import gradio as gr
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from transformers_interpret import SequenceClassificationExplainer

# ✅ Load model from Hugging Face
MODEL_NAME = "Ahsamkk/urdu-sentiment-model"  # ← replace with your actual repo name
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
model.eval()

# Ensure label mapping
model.config.id2label = {0: 'Negative', 1: 'Positive'}
model.config.label2id = {'Negative': 0, 'Positive': 1}

def analyze_sentiment(text):
    # Predict sentiment
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = model(**inputs)
        probs = F.softmax(outputs.logits, dim=1)[0]
        confidence = float(torch.max(probs))
        pred = torch.argmax(probs).item()
        sentiment = model.config.id2label[pred]

    # Explainability
    explainer = SequenceClassificationExplainer(model, tokenizer)
    word_attributions = explainer(text)

    top_influences = sorted(
        [{"word": w, "score": float(s)} for w, s in word_attributions],
        key=lambda x: abs(x["score"]),
        reverse=True
    )[:3]

    # Create readable output
    influences_text = "\n".join(
        [f"{i+1}. {inf['word']} ({inf['score']:.3f})" for i, inf in enumerate(top_influences)]
    )

    result_text = f"**Sentiment:** {sentiment}\n**Confidence:** {confidence:.3f}\n\n**Top Influences:**\n{influences_text}"
    return result_text

# ✅ Create Gradio interface
demo = gr.Interface(
    fn=analyze_sentiment,
    inputs=gr.Textbox(label="اپنا جملہ درج کریں (Enter Urdu sentence)"),
    outputs="text",
    title="Urdu Sentiment Analysis with Explainable AI",
    description="This model predicts sentiment and shows the top words influencing the prediction."
)

demo.launch()