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

# 创建 FastAPI 实例
app = FastAPI()

# 加载预训练模型
sentiment_model = pipeline("text-classification", model="shahxeebhassan/bert_base_ai_content_detector")

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

# 定义一个 POST 请求处理函数
@app.post("/predict")
async def predict(request: TextRequest):
    result = sentiment_model(request.text)
    print("处理前的 result:", result)
    
    # 获取原始预测结果
    original = result[0]
    
    # 计算AI的概率
    ai_probability = original["score"] if original["label"] == "LABEL_1" else 1 - original["score"]
    
    processed_result = [{
        "label": "AI" if ai_probability > 0.5 else "Human",
        "score": ai_probability
    }]
    
    print("处理后的 result:", processed_result)
    return {"result": processed_result}
    
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)