from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline

# 创建 FastAPI 实例
app = FastAPI()

# 加载预训练模型
sentiment_model = pipeline("text-classification", model="Hello-SimpleAI/chatgpt-detector-roberta")

# 定义请求体的格式
class TextRequest(BaseModel):
    text: str

# 定义一个 POST 请求处理函数
@app.post("/predict")
async def predict(request: TextRequest):
    result = sentiment_model(request.text)
    return {"result": result}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)