import torch
from unsloth import FastLanguageModel

MAX_SEQ_LENGTH = 2048
DTYPE = None
LOAD_IN_4BIT = True

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="zig_assistant_lora",
    max_seq_length=MAX_SEQ_LENGTH,
    dtype=DTYPE,
    load_in_4bit=LOAD_IN_4BIT,
)

FastLanguageModel.for_inference(model)

test_messages = [
    {
        "role": "system",
        "content": "You are an expert Zig programming assistant targeting Zig 0.17.0. Provide correct, idiomatic, memory-safe where applicable, and well-tested solutions with concise explanations.",
    },
    {
        "role": "user",
        "content": "Implement a generic stack data structure in Zig 0.17.0 using an explicit allocator.",
    },
]

inputs = tokenizer.apply_chat_template(
    test_messages, tokenize=True, add_generation_prompt=True, return_tensors="pt"
).to("cuda")

outputs = model.generate(
    input_ids=inputs, max_new_tokens=1500, use_cache=True, temperature=0.2, top_p=0.9
)

response = tokenizer.decode(outputs[0][inputs.shape[1] :], skip_special_tokens=True)
print("=== FINE-TUNED MODEL TEST OUTPUT ===")
print(response)