from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import uvicorn
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

app = FastAPI(title="情感分析API", version="1.0.0")

# 全局模型变量
model = None
tokenizer = None

class SentimentRequest(BaseModel):
    text: str
    max_length: int = 512

class SentimentResponse(BaseModel):
    text: str
    sentiment: str
    confidence: float
    processing_time: float

@app.on_event("startup")
async def load_model():
    """启动时加载模型"""
    global model, tokenizer
    
    print("🚀 加载情感分析模型...")
    
    model_path = "./output/sft_sentiment_qwen_merged"
    
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        torch_dtype=torch.float16,
        device_map="auto"
    )
    
    print("✅ 模型加载完成")

@app.post("/predict", response_model=SentimentResponse)
async def predict_sentiment(request: SentimentRequest):
    """预测文本情感"""
    import time
    
    start_time = time.time()
    
    try:
        # 构造输入
        prompt = f"""<|im_start|>user
请判断以下电影评论的情感倾向，只回答'正面'或'负面'：

评论：{request.text}
情感：<|im_end|>
<|im_start|>assistant
"""
        
        # 编码
        inputs = tokenizer(
            prompt, 
            return_tensors="pt",
            max_length=request.max_length,
            truncation=True
        ).to(model.device)
        
        # 推理
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=5,
                temperature=0.1,
                do_sample=False,
                pad_token_id=tokenizer.eos_token_id
            )
        
        # 解码
        response = tokenizer.decode(
            outputs[0][inputs.input_ids.shape[1]:],
            skip_special_tokens=True
        ).strip()
        
        # 解析结果
        if "正面" in response:
            sentiment = "positive"
            confidence = 0.9  # 简化的置信度
        elif "负面" in response:
            sentiment = "negative" 
            confidence = 0.9
        else:
            sentiment = "neutral"
            confidence = 0.5
        
        processing_time = time.time() - start_time
        
        return SentimentResponse(
            text=request.text,
            sentiment=sentiment,
            confidence=confidence,
            processing_time=processing_time
        )
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health")
async def health_check():
    """健康检查"""
    return {"status": "healthy", "model_loaded": model is not None}

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