import gradio as gr
import shap
import torch
import numpy as np
import matplotlib.pyplot as plt
from transformers import pipeline

# Load lightweight model
classifier = pipeline(
    "sentiment-analysis",
    model="distilbert-base-uncased-finetuned-sst-2-english"
)

# Create explainer
explainer = shap.Explainer(classifier)

def analyze(text):
    if not text.strip():
        return "Please enter text", None

    # Prediction
    result = classifier(text)[0]
    label = result["label"]
    score = result["score"]

    # SHAP values
    shap_values = explainer([text])

    tokens = shap_values[0].data
    values = shap_values[0].values

    # Create bar plot
    plt.figure()
    plt.barh(tokens, values)
    plt.xlabel("SHAP Value")
    plt.title("Word Contribution to Sentiment")

    return f"Prediction: {label} (Confidence: {score:.2f})", plt.gcf()

with gr.Blocks() as demo:
    gr.Markdown("# Sentiment Analysis with SHAP")

    inp = gr.Textbox(lines=4, placeholder="Enter text here...")
    prediction = gr.Textbox(label="Prediction")
    shap_plot = gr.Plot(label="SHAP Explanation")

    btn = gr.Button("Analyze")
    btn.click(analyze, inp, [prediction, shap_plot])

demo.launch()