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_display, original_sentence_input, sentence_1=None, sentence_2=None, sentence_3=None, sentence_4=None, sentence_5=None):
    model_name = "sartifyllc/African-Cross-Lingua-Embeddings-Model"
    model = load_model(model_name)
    result = {
        "Model Name": model_name,
        "Original Sentence": original_sentence_input,
        "Sentences to Compare": {
            "Sentence 1": sentence_1,
            "Sentence 2": sentence_2,
            "Sentence 3": sentence_3,
            "Sentence 4": sentence_4,
            "Sentence 5": sentence_5
        },
        "Similarity Scores": {}
    }
    
    if not sentence_1 or not sentence_2 or not sentence_3:
        return "Please provide a minimum of three sentences for comparison.", {}
    if not original_sentence_input:
        return "Please provide the original sentence.", {}
    
    sentences = [original_sentence_input, sentence_1, sentence_2, sentence_3]
    embeddings = model.encode(sentences)
    similarities = util.cos_sim(embeddings[0], embeddings[1:])
    similarity_scores = {
        "Sentence 1": float(similarities[0, 0]),
        "Sentence 2": float(similarities[0, 1]),
        "Sentence 3": float(similarities[0, 2]),
        "Sentence 4": float(similarities[0, 3]),
         "Sentence 5": float(similarities[0, 4]),
    }
    result["Similarity Scores"] = similarity_scores
    
    return result

model_name_display = gr.Markdown(value="**Model Name**: sartifyllc/African-Cross-Lingua-Embeddings-Model")
original_sentence_input = gr.Textbox(lines=2, placeholder="Enter the original sentence here...", label="Original Sentence")

sentence_1 = gr.Textbox(lines=2, placeholder="Enter the sentence to compare here...", label="Sentence 1")
sentence_2 = gr.Textbox(lines=2, placeholder="Enter the sentence to compare here...", label="Sentence 2")
sentence_3 = gr.Textbox(lines=2, placeholder="Enter the sentence to compare here...", label="Sentence 3")
sentence_4 = gr.Textbox(lines=2, placeholder="Enter the sentence to compare here...", label="Sentence 4")
sentence_5 = gr.Textbox(lines=2, placeholder="Enter the sentence to compare here...", label="Sentence 5")


inputs = [model_name_display, original_sentence_input, sentence_1, sentence_2, sentence_3]
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,
    article="Author: Innocent Charles. Model from Hugging Face Hub (sartify.com): [sartifyllc/African-Cross-Lingua-Embeddings-Model](https://huggingface.co/sartifyllc/African-Cross-Lingua-Embeddings-Model)",
  examples = [
    [
        "sartifyllc/African-Cross-Lingua-Embeddings-Model",
        "Jua linawaka sana leo.",
        "Òrùlé jẹ́ tí ó ti máa ń tan-ìmólẹ̀ lónìí.",
        "Ran na haske sosai yau.",
        "The sun is shining brightly today.",
        "napenda sana jua",
        "schooling is kinda boring, I don't like it"
    ],
    [
        "sartifyllc/African-Cross-Lingua-Embeddings-Model",
        "Mbuzi anaruka juu ya uzio.",
        "Àgbò ohun ti ń kọjú wá sórí orí-ọ̀kẹ́.",
        "Kura mai sauri tana tsalle kan kare mai barci.",
        "The goat is jumping over the fence.",
        "Àgbò ohun ti ń kọjú wá sórí orí-ọ̀kẹ́.",
        "The cat is sleeping under the table."
    ],
    [
        "sartifyllc/African-Cross-Lingua-Embeddings-Model",
        "Ninapenda kujifunza lugha mpya.",
        "Mo nífẹ̀ẹ́ láti kọ́ èdè tuntun.",
        "Ina son koyon sababbin harsuna.",
        "I love learning new languages.",
        "Ina son koyon sababbin harsuna.",
        "botu neyle ki lo no"
    ]
]

).launch(debug=True, share=True)