import numpy as np
import gradio as gr
from sentence_transformers import SentenceTransformer, util

# Function to load the selected model
def load_model(model_name):
    return SentenceTransformer(model_name)

# Function to compute similarity and classify relationship
def predict(model_name, original_sentence_input, *sentences_to_compare):
    model = load_model(model_name)
    result = {
        "Model Name": model_name,
        "Original Sentence": original_sentence_input,
        "Sentences to Compare": sentences_to_compare,
        "Similarity Scores": {}
    }
    
    if original_sentence_input and sentences_to_compare:
        sentences = [original_sentence_input] + list(sentences_to_compare)
        embeddings = model.encode(sentences)
        original_embedding = embeddings[0]

        for i, sentence in enumerate(sentences_to_compare, start=1):
            similarity_score = util.pytorch_cos_sim(original_embedding, embeddings[i]).item()
            result["Similarity Scores"][f"Sentence {i}"] = similarity_score
    
    return result

# Define inputs and outputs for Gradio interface
model_name = gr.Textbox(value="sartifyllc/African-Cross-Lingua-Embeddings-Model", label="Model Name")
original_sentence_input = gr.Textbox(lines=2, placeholder="Enter the original sentence here...", label="Original Sentence")
sentence_to_compare_inputs = gr.Textbox(lines=2, placeholder="Enter the sentence you want to compare...", label="Sentence to Compare", elem_id="sentence_to_compare")

inputs = [model_name, original_sentence_input, sentence_to_compare_inputs]
outputs = gr.JSON(label="Detailed Similarity Scores")

# Create Gradio interface
gr.Interface(
    fn=predict,
    title="African Cross-Lingua Embeddings Model's Demo",
    description="Compute the semantic similarity across various sentences among any African Languages using African-Cross-Lingua-Embeddings-Model.",
    inputs=inputs,
    outputs=outputs,
    cache_examples=False,
    allow_flagging="never"
).launch(debug=True, share=True)