Create main.py

main.py

@@ -0,0 +1,24 @@
+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]}