from sentence_transformers import SentenceTransformer

def return_similarity(query, docs, model, prompt_name):
    query_prompt_name = prompt_name
    query_embeddings = model.encode(query, prompt_name=query_prompt_name)
    doc_embeddings = model.encode(docs)

    similarities = model.similarity(query_embeddings, doc_embeddings)

    return similarities[0].tolist()


if __name__ == "__main__":
    model = SentenceTransformer("/workspace/models/stella_en_400M_v5", trust_remote_code=True).cuda()
