add backend code

backend/inference.py

@@ -0,0 +1,25 @@
+import torch
+from backend.utils import load_model, load_embeddings, load_texts
+
+# Search
+def query_search(query: str, n_answers: int, model_name: str):
+    model = load_model(model_name)
+
+    # Creating embeddings
+    # query_emb = model.encode(query, convert_to_tensor=True)[None, :]
+    query_emb = model.encode(query, convert_to_tensor=True)
+
+    print("loading embedding")
+    corpus_emb = load_embeddings()
+    corpus_texts = load_texts()
+
+    # Getting hits
+    hits = torch.nn.functional.cosine_similarity(
+        query_emb[None, :], corpus_emb, dim=1, eps=1e-8
+    )
+
+    corpus_texts["Similarity"] = hits.tolist()
+
+    return corpus_texts.sort_values(by="Similarity", ascending=False).head(n_answers)[
+        ["Description", "Code"]
+    ]

backend/utils.py

@@ -0,0 +1,28 @@
+import pandas as pd
+from sentence_transformers import SentenceTransformer
+import streamlit as st
+import torch
+
+
+@st.cache(allow_output_mutation=True)
+def load_model(model_name):
+    # Lazy downloading
+    model = SentenceTransformer(model_name)
+    return model
+
+
+@st.cache(allow_output_mutation=True)
+def load_embeddings():
+    # embedding pre-generated
+    corpus_emb = torch.load(
+        "./embeddings/descriptions_emb_100000_examples.pt",
+        map_location=torch.device("cpu"),
+    )
+    return corpus_emb
+
+
+@st.cache(allow_output_mutation=True)
+def load_texts():
+    # texts database pre-generated
+    corpus_texts = pd.read_csv("./data/codesearchnet_100000_examples.csv")
+    return corpus_texts