import gradio as gr
from transformers.utils import logging
logging.set_verbosity_error()

from sentence_transformers import SentenceTransformer
from sentence_transformers import util

def compare(sentences1, sentences2):
    sentences1 = sentences1.splitlines()
    sentences2 = sentences2.splitlines()
    embeddings1 = model.encode(sentences1, convert_to_tensor=True)
    embeddings2 = model.encode(sentences2, convert_to_tensor=True)
    cosine_scores = util.cos_sim(embeddings1, embeddings2)

    output = ""
    for i in range(len(sentences1)):
        output += "Score: {:.4f} \t\t {} \t\t {}\n".format(cosine_scores[i][i], sentences1[i], sentences2[i])
        
    return output

model = SentenceTransformer("all-MiniLM-L6-v2")
demo = gr.Interface(
    compare,
    [
        gr.Textbox(
            label="Text",
            info="Initial text",
            lines=3,
            value="I like cats\nTea puts me to sleep\nThe quick brown fox jumped over the lazy dogs.",
        ),
        gr.Textbox(
            label="Compare Text",
            info="Text to compare",
            lines=3,
            value="I love kittens\nCoffee wakes me up\nThe fast brown fox jumps over lazy dogs.",
        ),
    ], outputs="text")
demo.launch()
gr.close_all()