import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

model = AutoModelForSequenceClassification.from_pretrained("knightscode139/bert-base-cased-imdb-sentiment")
tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-cased")

def predict(text):
    inputs = tokenizer(text, return_tensors="pt", padding="max_length", truncation=True)
    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.nn.functional.softmax(outputs.logits[0], dim=0)
    
    return {"Negative": float(probs[0]), "Positive": float(probs[1])}

interface = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(lines=5, placeholder="Enter movie review..."),
    outputs=gr.Label(num_top_classes=2),
    title="IMDB Sentiment Classifier",
    description="Fine-tuned BERT on IMDB reviews - 92.8% test accuracy"
)

interface.launch()