from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline, DistilBertForSequenceClassification, AutoTokenizer
import os
import uvicorn

# 设置环境变量
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'

# 创建 FastAPI 实例
app = FastAPI()

# 异步加载预训练模型
@app.on_event("startup")
async def load_model():
    global sentiment_model
    model = DistilBertForSequenceClassification.from_pretrained("ayatsuri/academic-ai-detector", from_tf=True)
    tokenizer = AutoTokenizer.from_pretrained("ayatsuri/academic-ai-detector")
    sentiment_model = pipeline("text-classification", model=model, tokenizer=tokenizer)

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

# 定义一个 POST 请求处理函数
@app.post("/predict")
async def predict(request: TextRequest):
    try:
        result = sentiment_model(request.text)
        # 遍历 result 列表，检查每个元素的 label
        for item in result:
            if item["label"] == "human-produced":
               item["score"] = 1 - item["score"]
        return {"result": result}
    except Exception as e:
        return {"error": str(e)}

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