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

# Load your SentenceTransformer model fine-tuned for NLI
model = SentenceTransformer("Omartificial-Intelligence-Space/Arabic-Nli-Matryoshka")

# Define the labels for NLI
labels = ["contradiction", "entailment", "neutral"]

# Define the Matryoshka dimension
matryoshka_dim = 64

# Function to compute similarity and classify relationship
def predict(sentence1, sentence2):
    sentences = [sentence1, sentence2]
    embeddings = model.encode(sentences)
    
    # Shrink the embedding dimensions
    embeddings = embeddings[..., :matryoshka_dim]
    
    # Compute cosine similarity between the two sentences
    similarity_score = util.cos_sim(embeddings[0], embeddings[1])
    
    # Placeholder logic for NLI (needs to be replaced with actual model inference)
    # This is just an example; in reality, you need a classifier trained for NLI
    scores = np.random.rand(3)  # Replace this with actual model prediction logic
    scores = scores / scores.sum()  # Normalize to sum to 1
    
    label_probs = {labels[i]: float(scores[i]) for i in range(len(labels))}
    
    return similarity_score.item(), label_probs

# Define inputs and outputs for Gradio interface
inputs = [
    gr.Textbox(lines=2, placeholder="Enter the first sentence here...", label="Sentence 1"),
    gr.Textbox(lines=2, placeholder="Enter the second sentence here...", label="Sentence 2")
]

outputs = [
    gr.Textbox(label="Similarity Score"),
    gr.Label(num_top_classes=3, label="Label Probabilities")
]

examples = [
    ["يجلس شاب ذو شعر أشقر على الحائط يقرأ جريدة بينما تمر امرأة وفتاة شابة.", "ذكر شاب ينظر إلى جريدة بينما تمر إمرأتان بجانبه"],
    ["الشاب نائم بينما الأم تقود ابنتها إلى الحديقة", "ذكر شاب ينظر إلى جريدة بينما تمر إمرأتان بجانبه"]
]

# Create Gradio interface
gr.Interface(
    fn=predict,
    title="Arabic Sentence Similarity and NLI Classification",
    description="Compute the semantic similarity and classify the relationship between two Arabic sentences using a SentenceTransformer model.",
    inputs=inputs,
    examples=examples,
    outputs=outputs,
    cache_examples=False,
    article="Author: Your Name. Model from Hugging Face Hub: Omartificial-Intelligence-Space/Arabic-Nli-Matryoshka",
).launch(debug=True, share=True)