pdfchat / app.py
iabualhaol's picture
Update app.py
36ea7dc
raw
history blame contribute delete
No virus
664 Bytes
import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer
model_name = "distilbert-base-uncased-finetuned-sst-2-english"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
def predict_sentiment(text):
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
outputs = model(**inputs)
probs = outputs.logits.softmax(dim=1).detach().numpy()[0]
return {"Negative": float(probs[0]), "Positive": float(probs[1])}
iface = gr.Interface(fn=predict_sentiment, inputs="text", outputs="label")
iface.launch()