answerer-api / mapper.py
DaniilAlpha's picture
Update mapper.py
a2e05fb
raw
history blame
553 Bytes
from sentence_transformers import SentenceTransformer, util
class Mapper:
def __init__(self, repo: str, model: str):
self.__model = SentenceTransformer(f"{repo}/{model}")
def __call__(self, query: str, data: list[str]):
query_emb = self.__model.encode(query)
data_emb = self.__model.encode(data)
scores: list[float] = util.dot_score(query_emb, data_emb)[0].cpu().tolist()
data_score_pairs = list(zip(range(0, len(data)), scores))
return sorted(
data_score_pairs,
key=lambda x: x[1],
reverse=True,
)