from sentence_transformers import SentenceTransformer, util
import gradio as gr

# Load the model
model = SentenceTransformer('sartifyllc/AviLaBSE')

# Function to compute similarities
def compute_similarity(original_sentence, sentences_to_compare):
    # Encode the original sentence and the sentences to compare
    embeddings_original = model.encode([original_sentence])
    embeddings_sentences_to_compare = model.encode(sentences_to_compare)

    # Compute cosine similarities
    similarities = util.cos_sim(embeddings_original, embeddings_sentences_to_compare)

    # Prepare results
    results = []
    for i, sentence in enumerate(sentences_to_compare):
        results.append((sentence, similarities[0][i].item()))

    return results

# Interface using Gradio
def sentence_similarity(original_sentence, sentence1, sentence2, sentence3):
    sentences_to_compare = [sentence1, sentence2, sentence3]
    results = compute_similarity(original_sentence, sentences_to_compare)
    return results

iface = gr.Interface(
    fn=sentence_similarity,
    inputs=[
        gr.Textbox(label="Original Sentence"),
        gr.Textbox(label="Sentence 1"),
        gr.Textbox(label="Sentence 2"),
        gr.Textbox(label="Sentence 3")
    ],
    outputs=gr.Table(label="Similarity Scores", columns=["Sentence", "Similarity Score"]),
    title="Sentence Similarity Calculator",
    description="Enter an original sentence and three sentences to compare their similarity."
)

iface.launch()