"""
Qwen3-0.6B inference test
Model: /home/runner/workspace/model (cloned from HuggingFace)
Library: transformers (airllm is for large sharded models; Qwen3-0.6B fits in RAM directly)
"""
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_path = "/home/runner/workspace/model"

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(model_path)

print("Loading Qwen3-0.6B model...")
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    dtype=torch.float32,
    device_map="cpu",
)
model.eval()
print("Model loaded!\n")

def chat(prompt: str, max_new_tokens: int = 150) -> str:
    messages = [{"role": "user", "content": prompt}]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )
    inputs = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=0.7,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
        )
    return tokenizer.decode(
        outputs[0][inputs["input_ids"].shape[1]:],
        skip_special_tokens=True,
    )

# Test prompts
prompts = [
    "Hello! What can you do?",
    "What is 2 + 2? Explain briefly.",
    "Write a haiku about the ocean.",
]

for p in prompts:
    print(f"User: {p}")
    response = chat(p)
    print(f"Qwen3: {response}")
    print("-" * 50)

print("\n✓ Qwen3-0.6B is running smoothly!")