test_model.py

"""
Quick test to verify the Qwen2.5-Omni model is working
"""

from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor
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"
)
processor = Qwen2_5OmniProcessor.from_pretrained("Qwen/Qwen2.5-Omni-3B")

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

# Test question
conversation = [
    {"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
text_input = processor.apply_chat_template(
    conversation,
    add_generation_prompt=True,
    tokenize=False
)

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

with torch.no_grad():
    generated_ids = model.generate(**inputs, max_new_tokens=100)

response = processor.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("  python app.py")
print("\nThen open: http://localhost:5000")