import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_PATH = "./outputs/qwen3_0.6b_fft"

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(
    MODEL_PATH,
    trust_remote_code=True
)

print("Loading model...")

dtype = (
    torch.bfloat16
    if torch.cuda.is_available() and torch.cuda.is_bf16_supported()
    else torch.float16
)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    dtype=dtype,
    trust_remote_code=True,
    device_map="auto"
)

model.eval()

print("=" * 60)
print("Qwen3 Full Fine-Tuned Model")
print("Type 'exit' to quit")
print("=" * 60)

while True:

    user_input = input("\nUser: ")

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

    messages = [
        {
            "role": "user",
            "content": user_input
        }
    ]

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

    inputs = tokenizer(
        text,
        return_tensors="pt"
    ).to(model.device)

    with torch.no_grad():

        outputs = model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.7,
            top_p=0.95,
            do_sample=True,
            repetition_penalty=1.1,
            eos_token_id=tokenizer.eos_token_id
        )

    response = tokenizer.decode(
        outputs[0][inputs["input_ids"].shape[1]:],
        skip_special_tokens=True
    )

    print("\nAssistant:", response)