from fastapi import FastAPI
from pydantic import BaseModel
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer, Qwen2ForCausalLM
import os
import torch

app = FastAPI()

print("Downloading model...")

repo_path = snapshot_download(
    repo_id="dwarrrrrrrr/chatbot_breast_cancer"
)

MODEL_PATH = os.path.join(
    repo_path,
    "merged_model"
)

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)

print("Loading model...")
model = Qwen2ForCausalLM.from_pretrained(
    MODEL_PATH,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
)

model.eval()

print("Model loaded!")


class ChatRequest(BaseModel):
    message: str


@app.get("/")
def root():
    return {"status": "ok"}


@app.post("/chat")
async def chat(req: ChatRequest):

    messages = [
        {
            "role": "system",
            "content": "Kamu adalah asisten kesehatan virtual yang menjawab dalam Bahasa Indonesia."
        },
        {
            "role": "user",
            "content": req.message
        }
    ]

    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

    inputs = tokenizer(
        [text],
        return_tensors="pt"
    )

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=200,
            temperature=0.7,
            do_sample=True,
            top_p=0.9,
            pad_token_id=tokenizer.eos_token_id
        )

    input_len = inputs["input_ids"].shape[1]

    generated = outputs[0][input_len:]

    response = tokenizer.decode(
        generated,
        skip_special_tokens=True
    )

    return {
        "response": response.strip()
    }