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