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

MODEL_PATH = "model"
TEMPERATURE = 0.5
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


def load_model_and_tokenizer(model_path):
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=False)
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float16,
        device_map="auto",
        trust_remote_code=False
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    return model, tokenizer


def generate_stream(model, tokenizer, prompt, temperature=0.4, max_new_tokens=512):
    """Generate text with streaming."""
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048)
    inputs = {k: v.to(model.device) for k, v in inputs.items()}

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

    generation_kwargs = dict(
        **inputs,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        do_sample=True,
        top_p=0.96,
        repetition_penalty=1.1,
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id,
        streamer=streamer,
    )

    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    for new_text in streamer:
        if "<|endoftext|>" in new_text:
            new_text = new_text.split("<|endoftext|>")[0]
            yield new_text
            break
        yield new_text

    thread.join()


def text_continuation(model, tokenizer, temperature):
    print(f"Instruct Mode (temp={temperature})")
    print("Enter your text and the model will continue it.")
    print("Commands: 'exit' or 'quit' — exit.")
    print("=" * 60)

    while True:
        try:
            user_input = input("\nInstruction: ").strip()
        except (KeyboardInterrupt, EOFError):
            print("\nGoodbye!")
            break

        if user_input.lower() in ["exit", "quit"]:
            print("Goodbye!")
            break

        if not user_input:
            continue

        print("Response: ", end="", flush=True)
        for token in generate_stream(model, tokenizer, "<|user|>"+user_input+"<|assistant|>", temperature=temperature):
            print(token, end="", flush=True)
        print()


if __name__ == "__main__":
    print(f"Loading model from {MODEL_PATH}...")
    model, tokenizer = load_model_and_tokenizer(MODEL_PATH)
    print(f"Model loaded. Device: {DEVICE}")

    text_continuation(model, tokenizer, temperature=TEMPERATURE)