from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
import torch
from threading import Thread

# === MODEL and TOKENIZER ===
model_id = "Bertug1911/BrtGPT-1"

tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)
model.eval().to("cuda" if torch.cuda.is_available() else "cpu")

# === CLEAN ===
def clean(text):
    return text.replace(" ", "").replace("Ġ", " ").replace("Ċ", "\n")

# === PAST ===
chat_history = []

# === LOOP: QA ===
while True:
    try:
        # Soru al
        question = input("\n🧠 Soru: ")

        if question.strip().lower() in ["q", "quit", "exit"]:
            print("Çıkılıyor...")
            break

        # Mesaj geçmişine kullanıcı mesajını ekle
        chat_history.append({"role": "user", "content": question})

        # Mesajı chat formatında hazırla (önceki geçmiş + yeni mesaj)
        inputs = tokenizer.apply_chat_template(
            chat_history,
            add_generation_prompt=True,
            return_tensors="pt"
        ).to(model.device)

        # Streamer tanımla
        streamer = TextIteratorStreamer(
            tokenizer,
            skip_prompt=True,
            skip_special_tokens=True
        )

        # Üretim thread'i
        def generate():
            model.generate(
                input_ids=inputs,
                streamer=streamer,
                max_new_tokens=256,
                do_sample=True,
                top_k=10,
                temperature=0.1,
            )

        thread = Thread(target=generate)
        thread.start()

        # Cevap tokenlarını temizleyerek yazdır
        print("🤖 Cevap:", end=" ", flush=True)
        full_answer = ""
        for token in streamer:
            cleaned = clean(token)
            full_answer += cleaned
            print(cleaned, end="", flush=True)

        # Cevabı da geçmişe assistant olarak ekle
        chat_history.append({"role": "assistant", "content": full_answer})

    except KeyboardInterrupt:
        print("\nManuel olarak durduruldu.")
        break