File size: 750 Bytes
ff4ec71 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 |
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"]
]
|