test_model_fixed.py

"""
Test the Qwen2.5-Omni model with corrected conversation format
"""

from transformers import Qwen2_5OmniForConditionalGeneration, AutoTokenizer
import torch

print("="*60)
print("Testing Qwen2.5-Omni-3B Model")
print("="*60)

print("\nLoading model...")
model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
    "Qwen/Qwen2.5-Omni-3B",
    torch_dtype="auto",
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Omni-3B")

print(f"✓ Model loaded on device: {model.device}\n")

# Test question - use tokenizer's chat template
messages = [
    {"role": "system", "content": "You are a helpful math tutor."},
    {"role": "user", "content": "What is 5 + 3?"}
]

print("Test Question: What is 5 + 3?")
print("Generating response...\n")

# Generate response using tokenizer
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():
    output = model.generate(
        **inputs,
        max_new_tokens=100,
        do_sample=False
    )

# Handle tuple return (text_ids, audio) or just text_ids
if isinstance(output, tuple):
    generated_ids = output[0]
else:
    generated_ids = output

response = tokenizer.batch_decode(
    generated_ids[:, inputs['input_ids'].shape[1]:],
    skip_special_tokens=True
)[0]

print("="*60)
print("Model Response:")
print("="*60)
print(response)
print("\n" + "="*60)
print("✓ Model is working correctly!")
print("="*60)
print("\nYou can now run the web app with:")
print("  ~/anaconda3/envs/tutor/bin/python app.py")
print("\nThen open: http://localhost:5000")