sentiment-api / app.py
Krish623's picture
Update app.py
16e269b verified
Raw
History Blame Contribute Delete
879 Bytes
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()