mapper.py

from sentence_transformers import SentenceTransformer, util

class Mapper:
  def __init__(self, model: str):
    self.__model = SentenceTransformer(f"models/{model}/pytorch_model.bin")

  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()
    return scores