democm / app.py
xingshining
Add application file
dda5ad5
Raw
History Blame Contribute Delete
780 Bytes
import gradio as gr
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
# 加载模型和tokenizer
model_name = "distilbert-base-uncased-finetuned-sst-2-english"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
def classify_text(text):
inputs = tokenizer(text, return_tensors="pt")
outputs = model(**inputs)
prediction = torch.nn.functional.softmax(outputs.logits, dim=-1)
return {"positive": prediction[0][1].item(), "negative": prediction[0][0].item()}
iface = gr.Interface(fn=classify_text, inputs="text", outputs="label")
iface.title = "Sentiment Analysis"
iface.description = "A simple sentiment analysis model using DistilBERT."
iface.launch()