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

# Khởi tạo FastAPI
app = FastAPI(title="DeepSeek R1 Chatbot with Prompt Engineering")

# Định nghĩa cấu trúc dữ liệu đầu vào
class Message(BaseModel):
    text: str

# Chỉ định thư mục cache tùy chỉnh
cache_dir = "/app/cache"  # Thư mục có thể ghi trong container
os.makedirs(cache_dir, exist_ok=True)  # Tạo thư mục nếu chưa có
os.environ["HF_HOME"] = cache_dir  # Đặt biến môi trường cho Hugging Face
os.environ["TRANSFORMERS_CACHE"] = cache_dir  # Đặt cache cho transformers

# Tải mô hình và tokenizer
model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"
tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir)

# Kiểm tra GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# System Prompt
SYSTEM_PROMPT = """Bạn là một trợ lý AI thông minh, chuyên giải toán và trò chuyện bằng tiếng Việt. Hãy luôn trả lời tự nhiên, dễ hiểu và sử dụng tiếng Việt chuẩn. Khi giải bài toán, suy luận từng bước theo kiểu Chain of Thought và đặt đáp án cuối cùng trong \boxed{}. Dưới đây là một ví dụ:

Ví dụ:
- Câu hỏi: "Giải bài toán: 2x + 3 = 7"
- Trả lời: "Ta có phương trình: 2x + 3 = 7.  
  Bước 1: Lấy 2x + 3 - 3 = 7 - 3, suy ra 2x = 4.  
  Bước 2: Chia hai vế cho 2, suy ra x = 2.  
  Vậy đáp án là: \boxed{2}."
"""

# Hàm tạo phản hồi
def generate_response(user_text: str) -> str:
    full_prompt = f"{SYSTEM_PROMPT}\n\nCâu hỏi từ người dùng: {user_text}\nTrả lời:"
    inputs = tokenizer(full_prompt, return_tensors="pt").to(device)
    outputs = model.generate(
        inputs["input_ids"],
        max_new_tokens=150,
        temperature=0.7,
        top_p=0.9,
        do_sample=True
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response[len(full_prompt):].strip()

# Endpoint gốc
@app.get("/")
def read_root():
    return {"message": "Chào mừng đến với DeepSeek R1 Chatbot với Prompt Engineering!"}

# Endpoint chat
@app.post("/chat")
def chat(message: Message):
    try:
        response = generate_response(message.text)
        return {"response": response}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Lỗi: {str(e)}")

# Chạy cục bộ
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)