# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/01_text_similarity.ipynb (unless otherwise specified).

__all__ = ['get_sentence_encoding', 'cos_sim', 'get_sentence_similarity']

# Cell
from fastai.text import *
from sklearn.metrics.pairwise import cosine_similarity

# Cell
def get_sentence_encoding(input: str, tok_type: str = 'classical_unigram', lm_type: str = 'classical_lm'):
    "Returns a sentence enoding for `input` sentence using fastai `UMLFit` model"
    tok = Tokenizer(tok_type)
    token_ids = tok.numericalize(input)
    # get learner
    defaults.device = torch.device('cpu')
    learn = load_learner(lm_type)
    encoder = learn.model[0]
    encoder.reset()
    kk0 = encoder(Tensor([token_ids]).to(torch.int64))
    return np.array(kk0[0][-1][0][-1])

# Cell
def cos_sim(v1, v2):
    "Return consine similarity between Tenosr `v1` and `v2`"
    return F.cosine_similarity(Tensor(v1).unsqueeze(0), Tensor(v2).unsqueeze(0)).mean().item()

def get_sentence_similarity(sent1: str, sent2: str, cmp: Callable = cos_sim, **kwargs):
    "Return sentence similarity between `sent1` and sent2` using `cmp` comparison method"
    encoding1 = get_sentence_encoding(sen1, **kwargs)
    encoding2 = get_sentence_encoding(sen2, **kwargs)
    return cmp(encoding1, encoding2)