from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

# 创建 FastAPI 应用
app = FastAPI()

# 加载模型和分词器
model_name = "IDEA-CCNL/Erlangshen-Roberta-330M-Sentiment"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# 定义请求数据格式
class TextRequest(BaseModel):
    text: str

# 定义情感分析 API
@app.post("/predict")
def predict_sentiment(request: TextRequest):
    inputs = tokenizer(request.text, return_tensors="pt", truncation=True, padding=True)
    outputs = model(**inputs)
    prediction = torch.argmax(outputs.logits, dim=1).item()
    return {"sentiment": ["负面", "中性", "正面"][prediction]}