import gradio as gr
import torch
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity

model = SentenceTransformer(
    "sentence-transformers/sentence-t5-base",
    device="cuda" if torch.cuda.is_available() else "cpu"
)


def get_metrics(vec1, vec2):
    sim = float(cosine_similarity(vec1, vec2)[0][0])

    scs = abs((sim) ** 3)
    m = {
        "cosine_similarity": round(sim, 4),
        "scs": round(scs, 4)
    }
    return m


def compute(text1, text2):

    texts = [text1, text2]

    embeddings = model.encode(
        texts,
        show_progress_bar=False,
        convert_to_numpy=True,
        normalize_embeddings=True,
    )

    return get_metrics(embeddings[0].reshape(1, -1), embeddings[1].reshape(1, -1))


with gr.Blocks() as demo:
    with gr.Row():
        text1 = gr.Textbox(label="Enter Text 1")
        text2 = gr.Textbox(label="Enter Text 2")

    with gr.Column():
        submit_btn = gr.Button("Submit")

        output = gr.JSON(
            label="Score",
        )

    # # callback ---
    submit_btn.click(
        fn=compute,
        inputs=[text1, text2],
        outputs=output
    )


demo.launch()