import gradio as gr
from sentence_transformers import SentenceTransformer, util

# Load the pre-trained sentence transformer model
model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')

def inference(text1, text2):
    # Encode the input sentences into sentence embeddings
    embeddings1 = model.encode(text1, convert_to_tensor=True)
    embeddings2 = model.encode(text2, convert_to_tensor=True)

    # Calculate the cosine similarity between the two sentence embeddings
    similarity_score = util.pytorch_cos_sim(embeddings1, embeddings2).item()

    return round(similarity_score, 2)

with gr.Blocks() as demo:
    gr.Markdown(
    """
    # Sentence Similarity Calculator
    Start typing below to see the output.
    """)
    txt = gr.Textbox(label="Input 1", lines=2)
    txt_2 = gr.Textbox(label="Input 2")
    txt_3 = gr.Textbox(value="", label="Output")
    btn = gr.Button(value="Submit")
    btn.click(inference, inputs=[txt, txt_2], outputs=[txt_3])

demo.launch()