qwen_usage_guide.py

"""
COMPREHENSIVE GUIDE: How to Use Qwen2.5-Omni for Math Tutor Bot
================================================================

This guide shows all 4 interaction modes:
1. Text-to-Text
2. Text-to-Speech
3. Speech-to-Text
4. Speech-to-Speech
5. Image-to-Text (bonus: for math diagrams)
"""

from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor
from qwen_omni_utils import process_mm_info
import torch

# ============================================================================
# SETUP: Load Model Once
# ============================================================================

def load_model():
    """Load Qwen2.5-Omni model (do this once at startup)"""
    print("Loading Qwen2.5-Omni-3B model...")

    model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
        "Qwen/Qwen2.5-Omni-3B",
        torch_dtype="auto",  # Automatically choose best dtype
        device_map="auto"     # Automatically use GPU if available
    )

    processor = Qwen2_5OmniProcessor.from_pretrained("Qwen/Qwen2.5-Omni-3B")

    print("Model loaded successfully!")
    return model, processor


# ============================================================================
# MODE 1: TEXT-TO-TEXT (Student types, bot responds with text)
# ============================================================================

def text_to_text(model, processor, question):
    """
    Student types a question, bot responds with text
    Perfect for: Written explanations, step-by-step solutions
    """
    print("\n" + "="*60)
    print("MODE 1: TEXT-TO-TEXT")
    print("="*60)

    # Create conversation
    conversation = [
        {
            "role": "system",
            "content": "You are a helpful math tutor. Explain concepts clearly and step-by-step."
        },
        {
            "role": "user",
            "content": question
        }
    ]

    # Process input
    text_input = processor.apply_chat_template(
        conversation,
        add_generation_prompt=True,
        tokenize=False
    )

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

    # Generate response
    with torch.no_grad():
        generated_ids = model.generate(**inputs, max_new_tokens=512)

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

    print(f"\nStudent: {question}")
    print(f"\nTutor: {response}")

    return response


# ============================================================================
# MODE 2: TEXT-TO-SPEECH (Student types, bot responds with voice)
# ============================================================================

def text_to_speech(model, processor, question, output_audio_file="response.wav"):
    """
    Student types a question, bot responds with voice
    Perfect for: Audio explanations, pronunciation, engaging lessons
    """
    print("\n" + "="*60)
    print("MODE 2: TEXT-TO-SPEECH")
    print("="*60)

    conversation = [
        {
            "role": "system",
            "content": "You are a helpful math tutor."
        },
        {
            "role": "user",
            "content": question
        }
    ]

    # Add audio generation prompt
    text_input = processor.apply_chat_template(
        conversation,
        add_generation_prompt=True,
        add_audio_generation_prompt=True,  # Request audio output!
        tokenize=False
    )

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

    # Generate text and audio
    with torch.no_grad():
        text_ids, audio_output = model.generate(
            **inputs,
            max_new_tokens=512,
            return_audio=True  # Request audio return
        )

    # Get text response
    text_response = processor.batch_decode(
        text_ids[:, inputs['input_ids'].shape[1]:],
        skip_special_tokens=True
    )[0]

    print(f"\nStudent: {question}")
    print(f"\nTutor (text): {text_response}")

    # Save audio
    if audio_output is not None:
        import soundfile as sf
        sf.write(output_audio_file, audio_output[0].cpu().numpy(), 24000)
        print(f"Tutor (audio): Saved to {output_audio_file}")

    return text_response, audio_output


# ============================================================================
# MODE 3: SPEECH-TO-TEXT (Student speaks, bot responds with text)
# ============================================================================

def speech_to_text(model, processor, audio_file_path):
    """
    Student speaks a question (from audio file), bot responds with text
    Perfect for: Voice questions, accessibility
    """
    print("\n" + "="*60)
    print("MODE 3: SPEECH-TO-TEXT")
    print("="*60)

    # Process audio input
    conversation = [
        {
            "role": "system",
            "content": "You are a helpful math tutor."
        },
        {
            "role": "user",
            "content": [
                {"type": "audio", "audio": audio_file_path}
            ]
        }
    ]

    # Process multimodal input
    mm_data = process_mm_info(conversation, model.config)

    text_input = processor.apply_chat_template(
        mm_data["conversation"],
        add_generation_prompt=True,
        tokenize=False
    )

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

    # Generate text response
    with torch.no_grad():
        generated_ids = model.generate(**inputs, max_new_tokens=512)

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

    print(f"\nStudent: [Audio file: {audio_file_path}]")
    print(f"\nTutor: {response}")

    return response


# ============================================================================
# MODE 4: SPEECH-TO-SPEECH (Student speaks, bot responds with voice)
# ============================================================================

def speech_to_speech(model, processor, audio_file_path, output_audio_file="response.wav"):
    """
    Student speaks a question, bot responds with voice
    Perfect for: Full voice interaction, natural conversation
    """
    print("\n" + "="*60)
    print("MODE 4: SPEECH-TO-SPEECH")
    print("="*60)

    conversation = [
        {
            "role": "system",
            "content": "You are a helpful math tutor."
        },
        {
            "role": "user",
            "content": [
                {"type": "audio", "audio": audio_file_path}
            ]
        }
    ]

    # Process multimodal input
    mm_data = process_mm_info(conversation, model.config)

    text_input = processor.apply_chat_template(
        mm_data["conversation"],
        add_generation_prompt=True,
        add_audio_generation_prompt=True,  # Request audio output!
        tokenize=False
    )

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

    # Generate text and audio response
    with torch.no_grad():
        text_ids, audio_output = model.generate(
            **inputs,
            max_new_tokens=512,
            return_audio=True
        )

    text_response = processor.batch_decode(
        text_ids[:, inputs['input_ids'].shape[1]:],
        skip_special_tokens=True
    )[0]

    print(f"\nStudent: [Audio file: {audio_file_path}]")
    print(f"\nTutor (text): {text_response}")

    # Save audio
    if audio_output is not None:
        import soundfile as sf
        sf.write(output_audio_file, audio_output[0].cpu().numpy(), 24000)
        print(f"Tutor (audio): Saved to {output_audio_file}")

    return text_response, audio_output


# ============================================================================
# BONUS MODE 5: IMAGE-TO-TEXT (Student shows math diagram, bot explains)
# ============================================================================

def image_to_text(model, processor, image_path, question="What is this?"):
    """
    Student uploads an image (e.g., math diagram), bot explains
    Perfect for: Understanding diagrams, graphs, geometric shapes
    """
    print("\n" + "="*60)
    print("BONUS MODE 5: IMAGE-TO-TEXT (Math Diagrams)")
    print("="*60)

    conversation = [
        {
            "role": "system",
            "content": "You are a helpful math tutor. Analyze images and explain clearly."
        },
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image_path},
                {"type": "text", "text": question}
            ]
        }
    ]

    # Process multimodal input
    mm_data = process_mm_info(conversation, model.config)

    text_input = processor.apply_chat_template(
        mm_data["conversation"],
        add_generation_prompt=True,
        tokenize=False
    )

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

    # Generate response
    with torch.no_grad():
        generated_ids = model.generate(**inputs, max_new_tokens=512)

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

    print(f"\nStudent: [Image: {image_path}] {question}")
    print(f"\nTutor: {response}")

    return response


# ============================================================================
# MAIN EXAMPLE: How to Use in Your Tutor Bot
# ============================================================================

if __name__ == "__main__":
    print("="*60)
    print("QWEN2.5-OMNI MATH TUTOR BOT - USAGE EXAMPLES")
    print("="*60)

    # Load model once at startup
    model, processor = load_model()

    # Example 1: Text-to-Text
    print("\n\n### EXAMPLE 1: Student types a question ###")
    text_to_text(
        model, processor,
        "Explain how to solve 3x + 7 = 22 step by step"
    )

    # Example 2: Text-to-Speech
    print("\n\n### EXAMPLE 2: Student wants audio explanation ###")
    text_to_speech(
        model, processor,
        "What is the Pythagorean theorem?",
        output_audio_file="pythagorean_explanation.wav"
    )

    # Example 3: Speech-to-Text (requires audio file)
    # Uncomment if you have an audio file:
    # print("\n\n### EXAMPLE 3: Student asks via voice ###")
    # speech_to_text(model, processor, "student_question.wav")

    # Example 4: Speech-to-Speech (requires audio file)
    # Uncomment if you have an audio file:
    # print("\n\n### EXAMPLE 4: Full voice conversation ###")
    # speech_to_speech(model, processor, "student_question.wav", "tutor_answer.wav")

    # Example 5: Image-to-Text (requires image file)
    # Uncomment if you have a math diagram image:
    # print("\n\n### EXAMPLE 5: Student uploads math diagram ###")
    # image_to_text(model, processor, "triangle_diagram.png", "Explain this geometric shape")

    print("\n\n" + "="*60)
    print("GUIDE COMPLETE!")
    print("="*60)
    print("\nKey Points:")
    print("  • Load model ONCE at startup")
    print("  • Use different functions for different interaction modes")
    print("  • Text responses are always generated")
    print("  • Audio output requires add_audio_generation_prompt=True")
    print("  • Audio files use 24000 Hz sample rate")
    print("  • Images can be PNG, JPG, or other common formats")
    print("\nFor your math tutor bot:")
    print("  1. Load model on app startup")
    print("  2. Choose interaction mode based on student preference")
    print("  3. Generate responses with appropriate function")
    print("  4. Display text and/or play audio to student")