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


model_id = "Bertug1911/BrtGPT-Conversation-1.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")


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



chat_history = [
    {
        "role": "system",
        "content": (
            "You are hepfull LLM."
        )
    }
]


while True:
    try:
        question = input("\nQuestion (write "q" for exit.): ")

        if question.strip().lower() in ["q", "quit", "exit"]:
            print("Exiting...")
            break

        chat_history.append({"role": "user", "content": question})

        inputs = tokenizer.apply_chat_template(
            chat_history,
            add_generation_prompt=True,
            return_tensors="pt"
        ).to(model.device)

        streamer = TextIteratorStreamer(
            tokenizer,
            skip_prompt=True,
            skip_special_tokens=True
        )

        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()

        print("🤖 Answer:", end=" ", flush=True)
        full_answer = ""
        for token in streamer:
            cleaned = clean(token)
            full_answer += cleaned
            print(cleaned, end="", flush=True)

        chat_history.append({"role": "assistant", "content": full_answer})

    except KeyboardInterrupt:
        print("\nStoped.")
        break