from sentence_transformers import SentenceTransformer


class model():
    def __init__(self):
        from sentence_transformers import SentenceTransformer
        # self.model = SentenceTransformer("./models/gemma/checkpoint-33246")
        self.model = SentenceTransformer("google/embeddinggemma-300m")
        self.model.load_adapter("./models/gemma/checkpoint-33246")
    

    def run(self, question:str, chunk:str)->int:
        query_embeddings = self.model.encode_query(question)
        document_embeddings = self.model.encode_document(chunk)
        similarities = self.model.similarity(query_embeddings, document_embeddings)
        return similarities