from fastapi import FastAPI, Request
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

print("COSMOS Llama Chatbot is starting...")

model_id = "ytu-ce-cosmos/Turkish-Llama-8b-DPO-v0.1"

print("Model loading started")
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
print("Model loading completed")

# bu mesaj değiştirilebilir ve chatbotun başlangıç mesajı olarak kullanılabilir
initial_message = [
    {"role": "system", "content": "Sen bir yapay zeka asistanısın. Kullanıcı sana bir görev verecek. Amacın görevi olabildiğince sadık bir şekilde tamamlamak."}
    # Görevi yerine getirirken adım adım düşün ve adımlarını gerekçelendir.
]

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Selected device:", device)

app = FastAPI()


@app.get('/')
def home():
    return {"hello": "Bitfumes"}


@app.post('/ask')
async def ask(request: Request):
    data = await request.json()
    prompt = data.get("prompt")
    if not prompt:
        return {"error": "Prompt is missing"}

    messages = initial_message.copy()
    messages.append({"role": "user", "content": prompt})

    # Tokenizer ve input
    input_ids = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt"
    ).to(model.device)

    terminators = [
        tokenizer.eos_token_id,
        tokenizer.convert_tokens_to_ids("<|eot_id|>")
    ]

    # Model generate sırasında gradient hesaplamayı kapatıyoruz
    with torch.no_grad():
        outputs = model.generate(
            input_ids,
            max_new_tokens=256,
            eos_token_id=terminators,
            do_sample=True,
            temperature=0.6,
            top_p=0.9,
        )
    
    # Cevabı çıkar ve tensorları GPU’dan kopar
    response = outputs[0][input_ids.shape[-1]:].detach()
    answer = tokenizer.decode(response, skip_special_tokens=True)

    # GPU’daki cache’i temizle
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    return {"answer": answer}