from sentence_transformers import SentenceTransformer, util
from typing import Any, Optional


class ModelCompare:

    def __init__(self, model: Optional[Any] = None):

        if model is None:
            model = init_model()

        self.model = model

    def __call__(self, a: str, b: str):

        a_embed = self.model.encode(a)
        b_embed = self.model.encode(b)

        return util.cos_sim(a_embed, b_embed).numpy().squeeze()


def init_model(model_id: str = 'all-MiniLM-L6-v2'):

    embedding = SentenceTransformer(model_id)
    return embedding


def keyword_compare(a: str, b: str):

    return a in b


if __name__ == "__main__":

    model = init_model()
    comp = ModelCompare(model)

    print(comp("I am a dog", "I am a cat"))
    print(comp("I am a dog", "I am a dog"))
    print(comp("I am a dog", ["puppy", "kitten"]))