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()