import gradio as gr
from transformers import AutoModelForSequenceClassification, DistilBertTokenizerFast, pipeline

MODEL_ID = "Krish623/sentiment-model"
tokenizer = DistilBertTokenizerFast.from_pretrained(MODEL_ID)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
classifier = pipeline("text-classification", model=model, tokenizer=tokenizer, top_k=None)

def predict(text):
    results = classifier(text)
    scores = results[0] if isinstance(results[0], list) else results
    best = max(scores, key=lambda x: x["score"])
    return {"label": best["label"], "score": best["score"]}

# ✅ USE BLOCKS (IMPORTANT)
with gr.Blocks() as demo:
    inp = gr.Textbox(label="Enter text")
    out = gr.JSON()
    btn = gr.Button("Predict")

    # 👇 THIS LINE FIXES EVERYTHING
    btn.click(fn=predict, inputs=inp, outputs=out, api_name="/predict")

demo.launch()