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-chinese")

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

# 定义一个 POST 请求处理函数
@app.post("/predict")
async def predict(request: TextRequest):
    result = sentiment_model(request.text)
    # 打印原始 result
    print("原始 result:", result)
    

    # 获取AI的概率并构建新的输出格式
    processed_result = [{
        "label": "AI" if result[0]["label"] == "ChatGPT" else "Human",
        "score": result[0]["score"] if result[0]["label"] == "ChatGPT" else 1 - result[0]["score"]
    }]
        
    # 打印处理后的 result
    print("处理后的 result:", processed_result)
    
    return {"result": processed_result}

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