Update README.md

README.md

@@ -8,18 +8,24 @@ base_model:
 - unsloth/Qwen2.5-Coder-1.5B-Instruct
 pipeline_tag: text-generation
 ---
+# Qwen25_Coder_MultipleChoice
 
-# Using tuandunghcmut/Qwen25_Coder_MultipleChoice
-The project "Knowledge Distillation About YAML-Based Structured Multi-Step Reasoning from a Teacher Model GPT-4o to a Small LLM: Qwen2.5 Coder 1.5B-Instruct" focuses on distilling structured multi-step reasoning from GPT-4o into a smaller model.
+*   This project focuses on distilling YAML-based structured multi-step reasoning capabilities from the GPT-4o teacher model into the smaller Qwen2.5 Coder 1.5B-Instruct LLM.
 
-This document provides everything you need to get started with tuandunghcmut/Qwen25_Coder_MultipleChoice, a model designed for multiple-choice coding questions.
+*   This document provides guidance on getting started with `tuandunghcmut/Qwen25_Coder_MultipleChoice`, a model fine-tuned for multiple-choice coding questions.
 
-I plan to refactor the project into a well-structured GitHub repository, expand the dataset, and re-train it later with distributed training for better scalability.
-## Installation and Setup
+*   Future plans include refactoring the project into a well-structured GitHub repository, expanding the dataset, and retraining the model using distributed training for improved scalability.
 
-### Prerequisites
+*   A demonstration notebook is available on Google Colab (click the badge below). Please note that the training code has been omitted from this notebook. It is intended solely for testing and inference using the latest checkpoint.
+    [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://drive.google.com/file/d/1Q4jtRjIkFWIAM82pAg4OBPCLjpQ8ndpI/view?usp=sharing)
 
-Make sure you have Python 3.8+ installed. Then install the required packages:
+*   Note: The Qwen2.5 Coder 1.5B-Instruct model might be too small for this task, and the current training dataset may be insufficient. Future iterations will explore using a larger model and more extensive data. However, the current model successfully adheres to the desired YAML format and demonstrates structured reasoning.
+
+*   The guide below provides an explanation of the code presented in the notebook.
+
+## Installation
+
+First, install the required dependencies:
 
 ```bash
 # Install core dependencies
@@ -35,175 +41,381 @@ pip install flash-attn --no-build-isolation
 pip install datasets pyyaml
 ```
 
-### Flash Attention Setup
+## Key Classes
 
-Flash Attention provides a significant speedup for transformer models. To use it with the Qwen model:
-
-1. Install Flash Attention as shown above
-2. Enable it when loading the model:
+The project provides several key classes for working with the model:
 
+### 1. QwenModelHandler
 ```python
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-# Enable Flash Attention during model loading
-model = AutoModelForCausalLM.from_pretrained(
-    "tuandunghcmut/Qwen25_Coder_MultipleChoice",
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
-    trust_remote_code=True,
-    use_flash_attention_2=True  # Enable Flash Attention
-)
+class QwenModelHandler:
+    """Handler for Qwen models with inference and saving capabilities using Unsloth"""
+    
+    def __init__(self, model_name="unsloth/Qwen2.5-7B", max_seq_length=768, 
+                 quantization=None, device_map="auto", cache_dir=None):
+        """
+        Initialize model and tokenizer using Unsloth
+        
+        Args:
+            model_name: Name or path of the model (preferably an unsloth model)
+            max_seq_length: Maximum sequence length for the model
+            quantization: Quantization type (None, '4bit', '8bit') - for compatibility
+            device_map: Device mapping strategy
+            cache_dir: Cache directory for models
+        """
 ```
 
-Flash Attention will provide:
-- 2-3x faster inference speed
-- Lower memory usage
-- Compatible with 4-bit quantization for even more efficiency
-
-### Environment Variables
-
-If you're using Hugging Face Hub models, you may want to set up your access token:
-
-```bash
-# Set environment variable for Hugging Face token
-export HF_TOKEN="your_huggingface_token_here"
+This class handles the core model operations:
+- Model loading and initialization
+- Text generation with streaming support
+- Perplexity calculation
+- Model saving and pushing to HuggingFace Hub
 
-# Or in Python
-import os
-os.environ["HF_TOKEN"] = "your_huggingface_token_here"
+### 2. PromptCreator
+```python
+class PromptCreator:
+    """Creates and formats prompts for multiple choice questions"""
+    
+    # Prompt types
+    BASIC = "basic"  # Simple answer-only format
+    YAML_REASONING = "yaml"  # YAML formatted reasoning
+    TEACHER_REASONED = "teacher"  # Same YAML format but using teacher completions
 ```
 
-### GPU Setup
+This class manages prompt creation with three modes:
+- Basic: Simple answer-only format
+- YAML Reasoning: Structured reasoning in YAML format
+- Teacher Reasoned: YAML format with teacher completions for training
 
-For optimal performance, you'll need a CUDA-compatible GPU. Check your installation:
+### 3. ResponseParser
+```python
+class ResponseParser:
+    """Parser for model responses with support for different formats"""
+    
+    # Parser modes
+    BASIC = "basic"  # Extract single letter answer
+    YAML = "yaml"    # Parse YAML formatted response with reasoning
+```
 
-```bash
-# Verify CUDA is available
-python -c "import torch; print('CUDA available:', torch.cuda.is_available())"
+This class handles response parsing:
+- Extracts answers from model responses
+- Parses YAML-formatted reasoning
+- Supports both basic and YAML formats
 
-# Print CUDA device info
-python -c "import torch; print('CUDA device count:', torch.cuda.device_count()); print('CUDA device name:', torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'No GPU')"
+### 4. MultipleChoiceTester
+```python
+class MultipleChoiceTester:
+    """Framework for testing Qwen models on multiple choice questions"""
+    
+    def __init__(self, model_handler, prompt_creator=None):
+        """
+        Initialize with model handler and prompt configuration
+        
+        Args:
+            model_handler: The QwenModelHandler instance
+            prompt_creator: Optional PromptCreator instance
+        """
 ```
 
-## Required Classes
+This class provides a complete testing framework:
+- Single example inference
+- Batch processing
+- Dataset evaluation
+- Performance metrics tracking
+- Results saving and visualization
 
-Below are the essential classes needed to work with the model. Copy these into your Python files to use them in your project.
+## Full Class Implementations
 
-### PromptCreator
+<details>
+<summary>Click to expand/collapse full class implementations</summary>
 
-This class formats prompts for multiple-choice questions:
+### 1. QwenModelHandler
+```python
+class QwenModelHandler:
+    """Handler for Qwen models with inference and saving capabilities using Unsloth"""
+    
+    def __init__(self, model_name="unsloth/Qwen2.5-7B", max_seq_length=768, 
+                 quantization=None, device_map="auto", cache_dir=None):
+        self.model_name = model_name
+        self.max_seq_length = max_seq_length
+        self.device_map = device_map
+        self.quantization = quantization
+        self.cache_dir = cache_dir
+        
+        # Convert quantization parameter to load_in_4bit parameter for Unsloth
+        self.load_in_4bit = quantization == "4bit"
+        
+        # Load tokenizer and model
+        self.tokenizer, self.model = self._load_model()
+        self.response_parser = ResponseParser()
+        
+    def _load_model(self):
+        """Load model and tokenizer with Unsloth for optimization"""
+        from unsloth import FastLanguageModel
+        import torch
+        
+        print(f"Loading {self.model_name} with Unsloth, max_seq_length={self.max_seq_length}")
+        
+        # Set dtype based on hardware
+        dtype = None  # None for auto detection
+        
+        # Load model and tokenizer with Unsloth
+        model, tokenizer = FastLanguageModel.from_pretrained(
+            model_name=self.model_name,
+            max_seq_length=self.max_seq_length,
+            dtype=dtype,
+            load_in_4bit=self.load_in_4bit,
+            cache_dir=self.cache_dir,
+        )
+        
+        return tokenizer, model
+      
+    def generate_with_streaming(self, prompt, temperature=0.7, max_tokens=1024, stream=True):
+        """Generate completion with optional streaming using Unsloth's optimized inference"""
+        # Enable faster inference
+        from unsloth import FastLanguageModel
+        FastLanguageModel.for_inference(self.model)
+        
+        # Format as chat
+        messages = [{"role": "user", "content": prompt}]
+        chat_text = self.tokenizer.apply_chat_template(
+            messages, 
+            tokenize=False, 
+            add_generation_prompt=True
+        )
+        
+        # Tokenize input
+        model_inputs = self.tokenizer([chat_text], return_tensors="pt").to(self.model.device)
+        
+        # Generate with streaming if requested
+        if stream:
+            from transformers import TextIteratorStreamer
+            import threading
+            
+            # Set up streamer
+            streamer = TextIteratorStreamer(
+                self.tokenizer,
+                skip_prompt=True,
+                skip_special_tokens=True
+            )
+            
+            # Start generation in a thread
+            generation_kwargs = {
+                "input_ids": model_inputs.input_ids,
+                "attention_mask": model_inputs.attention_mask,
+                "temperature": temperature,
+                "max_new_tokens": max_tokens,
+                "streamer": streamer,
+                "do_sample": temperature > 0.0,
+                "use_cache": True,
+                "min_p": 0.1 if temperature > 0.0 else None,
+            }
+            
+            thread = threading.Thread(target=self.model.generate, kwargs=generation_kwargs)
+            thread.start()
+            
+            return streamer
+        else:
+            # Generate without streaming
+            generated_ids = self.model.generate(
+                input_ids=model_inputs.input_ids,
+                attention_mask=model_inputs.attention_mask,
+                temperature=temperature,
+                max_new_tokens=max_tokens,
+                do_sample=temperature > 0.0,
+                use_cache=True,
+                min_p=0.1 if temperature > 0.0 else None,
+            )
+            
+            # Decode the generated text
+            generated_text = self.tokenizer.decode(
+                generated_ids[0][model_inputs.input_ids.shape[1]:],
+                skip_special_tokens=True
+            )
+            
+            return generated_text
+            
+    def calculate_perplexity(self, prompt, answer, temperature=0.0):
+        """Calculate perplexity for a prompt and answer pair"""
+        import torch
+        
+        # Format chat for perplexity calculation
+        messages = [
+            {"role": "user", "content": prompt},
+            {"role": "assistant", "content": answer}
+        ]
+        chat_text = self.tokenizer.apply_chat_template(
+            messages, 
+            tokenize=False
+        )
+        
+        # Tokenize the text
+        encodings = self.tokenizer(chat_text, return_tensors="pt").to(self.model.device)
+        
+        # Calculate loss
+        with torch.no_grad():
+            outputs = self.model(**encodings, labels=encodings.input_ids)
+            
+        # Get loss and calculate perplexity
+        neg_log_likelihood = outputs.loss.item()
+        perplexity = torch.exp(torch.tensor(neg_log_likelihood)).item()
+        
+        return perplexity
+      
+    def save_model(self, output_dir, save_method="lora"):
+        """Save model to disk using Unsloth's optimized methods"""
+        import os
+        
+        os.makedirs(output_dir, exist_ok=True)
+        
+        # Use Unsloth's saving methods
+        if save_method == "lora":
+            self.model.save_pretrained(output_dir)
+            self.tokenizer.save_pretrained(output_dir)
+        elif save_method == "merged_16bit":
+            self.model.save_pretrained_merged(output_dir, self.tokenizer, save_method="merged_16bit")
+        elif save_method == "merged_4bit":
+            self.model.save_pretrained_merged(output_dir, self.tokenizer, save_method="merged_4bit")
+        elif save_method == "gguf":
+            self.model.save_pretrained_gguf(output_dir, self.tokenizer, quantization_method="q4_k_m")
+        else:
+            raise ValueError(f"Unknown save method: {save_method}")
+            
+        print(f"Model saved to {output_dir} using method {save_method}")
+        return output_dir
+        
+    def push_to_hub(self, repo_id, token=None, save_method="lora", private=False):
+        """Push model to Hugging Face Hub using Unsloth's optimized methods"""
+        if save_method == "lora":
+            self.model.push_to_hub_merged(repo_id, self.tokenizer, save_method="lora", token=token)
+        elif save_method == "merged_16bit":
+            self.model.push_to_hub_merged(repo_id, self.tokenizer, save_method="merged_16bit", token=token)
+        elif save_method == "merged_4bit":
+            self.model.push_to_hub_merged(repo_id, self.tokenizer, save_method="merged_4bit", token=token)
+        elif save_method == "gguf":
+            self.model.push_to_hub_gguf(
+                repo_id, 
+                self.tokenizer, 
+                quantization_method=["q4_k_m", "q5_k_m"], 
+                token=token
+            )
+        else:
+            raise ValueError(f"Unknown save method: {save_method}")
+        
+        print(f"Model successfully pushed to: https://huggingface.co/{repo_id}")
+        return f"https://huggingface.co/{repo_id}"
+```
 
+### 2. PromptCreator
 ```python
 class PromptCreator:
-    """
-    Creates and formats prompts for multiple choice questions
-    Supports different prompt styles for training and inference
-    """
-
+    """Creates and formats prompts for multiple choice questions"""
+    
     # Prompt types
     BASIC = "basic"  # Simple answer-only format
     YAML_REASONING = "yaml"  # YAML formatted reasoning
-    TEACHER_REASONED = "teacher"  # Same YAML format as YAML_REASONING but using teacher completions for training
+    TEACHER_REASONED = "teacher"  # Same YAML format but using teacher completions
 
     def __init__(self, prompt_type=BASIC):
+        if prompt_type == self.TEACHER_REASONED:
+            prompt_type = self.YAML_REASONING
         self.prompt_type = prompt_type
-        # Initialize parser mode based on prompt type
-        if prompt_type == self.YAML_REASONING or prompt_type == self.TEACHER_REASONED:
-            self.parser_mode = "yaml"
-        else:
-            self.parser_mode = "basic"
+        self.original_type = prompt_type
 
     def format_choices(self, choices):
-        """Format choices with letter prefixes"""
-        return "\n".join([f"{chr(65 + i)}. {choice}" for i, choice in enumerate(choices)])
+        """Format choices as a lettered list"""
+        return "\n".join(
+            [f"{chr(65 + i)}. {choice}" for i, choice in enumerate(choices)]
+        )
 
     def get_max_letter(self, choices):
-        """Get the last valid letter based on choice count"""
+        """Get the maximum letter based on number of choices"""
         return chr(65 + len(choices) - 1)
 
     def create_inference_prompt(self, question, choices):
-        """Create a prompt for inference based on the configured prompt type"""
+        """Create a prompt for inference based on current prompt type"""
         formatted_choices = self.format_choices(choices)
         max_letter = self.get_max_letter(choices)
-        
-        if self.prompt_type == self.BASIC:
-            return self._create_basic_prompt(question, formatted_choices, max_letter)
-        elif self.prompt_type == self.YAML_REASONING or self.prompt_type == self.TEACHER_REASONED:
+
+        if self.prompt_type == self.YAML_REASONING:
             return self._create_yaml_prompt(question, formatted_choices, max_letter)
         else:
             return self._create_basic_prompt(question, formatted_choices, max_letter)
 
     def _create_basic_prompt(self, question, formatted_choices, max_letter):
-        """Create a basic prompt that just asks for an answer letter"""
+        """Create a basic prompt asking for just the answer letter"""
         return f"""
+QUESTION:
 {question}
 
+CHOICES:
 {formatted_choices}
 
-Select the correct answer from A through {max_letter}:
+Answer with a single letter from A through {max_letter} without any additional explanation or commentary.
 """
 
     def _create_yaml_prompt(self, question, formatted_choices, max_letter):
-        """Create a prompt with YAML formatted reasoning structure"""
+        """Create a prompt requesting YAML-formatted reasoning"""
         return f"""
+QUESTION:
 {question}
 
+CHOICES:
 {formatted_choices}
 
-Think through this step-by-step:
-- Understand what the question is asking
-- Analyze each option carefully
-- Reason about why each option might be correct or incorrect
-- Select the most appropriate answer
+Analyze this question step-by-step and provide a detailed explanation.
+Your response MUST be in YAML format as follows:
 
-Your response should be in YAML format:
 understanding: |
-  <your understanding of the question>
+  <your understanding of what the question is asking>
 analysis: |
   <your analysis of each option>
 reasoning: |
-  <your reasoning about the correct answer>
+  <your step-by-step reasoning process>
 conclusion: |
   <your final conclusion>
-answer: <single letter A through {max_letter} representing your final answer>
+answer: <single letter A through {max_letter}>
+
+The answer field MUST contain ONLY a single character letter.
 """
 
     def create_training_prompt(self, question, choices):
-        """Create a prompt for training based on the configured prompt type"""
+        """Create a prompt for training with the current prompt type"""
         formatted_choices = self.format_choices(choices)
         max_letter = self.get_max_letter(choices)
-        
-        if self.prompt_type == self.BASIC:
-            return self._create_basic_training_prompt(question, formatted_choices, max_letter)
-        elif self.prompt_type == self.YAML_REASONING or self.prompt_type == self.TEACHER_REASONED:
-            return self._create_yaml_training_prompt(question, formatted_choices, max_letter)
+
+        if self.prompt_type == self.YAML_REASONING:
+            return self._create_yaml_training_prompt(
+                question, formatted_choices, max_letter
+            )
         else:
-            return self._create_basic_training_prompt(question, formatted_choices, max_letter)
+            return self._create_basic_training_prompt(
+                question, formatted_choices, max_letter
+            )
 
     def _create_basic_training_prompt(self, question, formatted_choices, max_letter):
         """Create a basic training prompt"""
         return f"""
+QUESTION:
 {question}
 
+CHOICES:
 {formatted_choices}
 
-Select the correct answer from A through {max_letter}:
+The answer is a single letter (A, B, C, etc.). Only provide ONE character as your answer:
 """
 
     def _create_yaml_training_prompt(self, question, formatted_choices, max_letter):
-        """Create a training prompt with YAML formatted reasoning structure"""
+        """Create a YAML-formatted training prompt"""
         return f"""
+QUESTION:
 {question}
 
+CHOICES:
 {formatted_choices}
 
-Think through this step-by-step:
-- Understand what the question is asking
-- Analyze each option carefully
-- Reason about why each option might be correct or incorrect
-- Select the most appropriate answer
+Analyze this question step-by-step and provide a detailed explanation.
+Follow the YAML format in your response:
 
-Your response should be in YAML format:
 understanding: |
   <your understanding of the question>
 analysis: |
@@ -212,420 +424,823 @@ reasoning: |
   <your reasoning about the correct answer>
 conclusion: |
   <your final conclusion>
-answer: <single letter A through {max_letter} representing your final answer>
+answer: <single letter A through {max_letter}>
 """
 
     def set_prompt_type(self, prompt_type):
-        """Set the prompt type and update parser mode accordingly"""
+        """Set the prompt type"""
+        self.original_type = prompt_type
+        if prompt_type == self.TEACHER_REASONED:
+            pass
         self.prompt_type = prompt_type
-        if prompt_type == self.YAML_REASONING or prompt_type == self.TEACHER_REASONED:
-            self.parser_mode = "yaml"
-        else:
-            self.parser_mode = "basic"
-
+        return self
+        
     def is_teacher_mode(self):
-        """Check if prompt type is teacher mode"""
-        return self.prompt_type == self.TEACHER_REASONED
+        """Check if we're using teacher mode"""
+        return self.original_type == self.TEACHER_REASONED
 ```
 
-### ResponseParser
-
-This class extracts answers from model responses:
-
+### 3. ResponseParser
 ```python
 class ResponseParser:
-    """
-    Parser for model responses with support for different formats
-    Extracts answers and reasoning from model outputs
-    """
+    """Parser for model responses with support for different formats"""
     
     # Parser modes
-    BASIC = "basic"        # Extract single letter answer
-    YAML = "yaml"          # Parse YAML formatted response with reasoning
+    BASIC = "basic"  # Extract single letter answer
+    YAML = "yaml"    # Parse YAML formatted response with reasoning
     
     def __init__(self, parser_mode=BASIC):
-        """Initialize with parser mode (basic or yaml)"""
         self.parser_mode = parser_mode
-        
+    
     def parse(self, response_text):
-        """Parse the response text and extract answer and reasoning"""
+        """Parse the model's response according to the current mode"""
         if self.parser_mode == self.YAML:
             return self._parse_yaml_response(response_text)
         else:
             return self._parse_basic_response(response_text)
     
     def _parse_basic_response(self, response_text):
-        """
-        Parse a basic response to extract the answer letter
-        
-        Returns:
-            tuple: (answer_letter, None)
-        """
-        # Look for just the letter at the end of text
+        """Parse basic response looking for a letter answer"""
         import re
         
-        # Try to find the last occurrence of letters A-Z by themselves
-        matches = re.findall(r'\b([A-Z])\b', response_text)
-        if matches:
-            return matches[-1], None  # Return the last matching letter
-            
-        # Try to find "The answer is X" pattern
-        answer_match = re.search(r'[Tt]he answer is[:\s]+([A-Z])', response_text)
+        # Try to extract a single letter answer (A-Z)
+        answer_match = re.search(r"(?:^|\s)([A-Z])(?:\s|$|\.)", response_text)
         if answer_match:
-            return answer_match.group(1), None
-            
-        # If nothing else works, just get the last uppercase letter
-        uppercase_letters = re.findall(r'[A-Z]', response_text)
-        if uppercase_letters:
-            return uppercase_letters[-1], None
-            
-        return None, None  # No answer found
+            answer = answer_match.group(1)
+        else:
+            # Take first character if it's a letter
+            if response_text and response_text[0].isalpha():
+                answer = response_text[0].upper()
+            else:
+                answer = None
+        
+        # For basic mode, we don't extract detailed reasoning
+        reasoning = ""
+        
+        return answer, reasoning
     
     def _parse_yaml_response(self, response_text):
-        """
-        Parse a YAML formatted response to extract the answer and reasoning
-        
-        Returns:
-            tuple: (answer_letter, reasoning_dict)
-        """
+        """Parse YAML formatted response extracting answer and reasoning"""
         import re
         import yaml
         
-        # First try to extract just the answer field
-        answer_match = re.search(r'answer:\s*([A-Z])', response_text)
-        answer = answer_match.group(1) if answer_match else None
+        # First try to find answer in YAML format
+        yaml_match = re.search(r"answer:\s*([A-Z])", response_text)
+        if yaml_match:
+            answer = yaml_match.group(1)
+        else:
+            # Fall back to basic extraction if YAML parsing fails
+            answer_match = re.search(r"(?:^|\s)([A-Z])(?:\s|$|\.)", response_text)
+            if answer_match:
+                answer = answer_match.group(1)
+            elif response_text and response_text[0].isalpha():
+                answer = response_text[0].upper()
+            else:
+                answer = None
         
-        # Try to extract the entire YAML
-        try:
-            # Remove potential code block markers
-            yaml_text = response_text
-            if "```yaml" in yaml_text:
-                yaml_text = yaml_text.split("```yaml")[1]
-                if "```" in yaml_text:
-                    yaml_text = yaml_text.split("```")[0]
-            elif "```" in yaml_text:
-                # Assume the whole thing is a code block
-                parts = yaml_text.split("```")
-                if len(parts) >= 3:
-                    yaml_text = parts[1]
-                    
-            # Parse the YAML
-            parsed_yaml = yaml.safe_load(yaml_text)
-            
-            # If successful, use the answer from the YAML, and return the parsed structure
-            if isinstance(parsed_yaml, dict) and "answer" in parsed_yaml:
-                return parsed_yaml.get("answer"), parsed_yaml
-        except Exception:
-            # If YAML parsing fails, we already have the answer from regex
-            pass
+        # Try to parse reasoning from YAML format
+        reasoning = ""
+        if "reasoning:" in response_text:
+            yaml_content = yaml.safe_load("---\n" + response_text)
+            if isinstance(yaml_content, dict) and "reasoning" in yaml_content:
+                reasoning = yaml_content["reasoning"]
+                
+                # Add other YAML fields if available
+                if "understanding" in yaml_content:
+                    reasoning = f"Understanding: {yaml_content['understanding']}\n\n{reasoning}"
+                if "conclusion" in yaml_content:
+                    reasoning = f"{reasoning}\n\nConclusion: {yaml_content['conclusion']}"
+        else:
+            # Use the full response as reasoning if not in YAML format
+            reasoning = response_text
         
-        return answer, None
+        return answer, reasoning
     
     def set_parser_mode(self, parser_mode):
         """Set the parser mode"""
         self.parser_mode = parser_mode
+        return self
     
     @classmethod
     def from_prompt_type(cls, prompt_type):
-        """
-        Create a ResponseParser with the appropriate mode based on prompt type
-        
-        Args:
-            prompt_type: The prompt type (e.g., PromptCreator.YAML_REASONING)
-            
-        Returns:
-            ResponseParser: A parser configured for the prompt type
-        """
-        if prompt_type in ["yaml", "teacher"]:
-            return cls("yaml")
+        """Create a parser instance with mode matching the prompt type"""
+        if prompt_type == PromptCreator.YAML_REASONING or prompt_type == PromptCreator.TEACHER_REASONED:
+            return cls(parser_mode=cls.YAML)
         else:
-            return cls("basic")
+            return cls(parser_mode=cls.BASIC)
 ```
 
-### QwenModelHandler
-
-This class handles model loading and inference:
-
+### 4. MultipleChoiceTester
 ```python
-class QwenModelHandler:
-    def __init__(self, model_name="unsloth/Qwen2.5-7B", max_seq_length=768, 
-                 quantization=None, device_map="auto", cache_dir=None,
-                 use_flash_attention=True):
-        """
-        Initialize a handler for Qwen models
+class MultipleChoiceTester:
+    """Framework for testing Qwen models on multiple choice questions"""
+
+    def __init__(self, model_handler, prompt_creator=None):
+        self.model_handler = model_handler
+        self.prompt_creator = prompt_creator or PromptCreator(PromptCreator.BASIC)
+        self.response_parser = ResponseParser.from_prompt_type(self.prompt_creator.prompt_type)
+
+    def infer_example(self, example, temperature=0.7, max_tokens=1024, prompt_type=None, stream=False):
+        """Inference on a single example for visualization/demonstration"""
+        # Allow temporary override of prompt type
+        original_prompt_type = None
+        if prompt_type is not None:
+            original_prompt_type = self.prompt_creator.prompt_type
+            self.prompt_creator.set_prompt_type(prompt_type)
+            self.response_parser = ResponseParser.from_prompt_type(prompt_type)
         
-        Args:
-            model_name: Model identifier (local path or Hugging Face model ID)
-            max_seq_length: Maximum sequence length
-            quantization: Quantization method ("4bit", "8bit", or None)
-            device_map: Device mapping strategy
-            cache_dir: Directory to cache downloaded models
-            use_flash_attention: Whether to use Flash Attention 2 for faster inference
-        """
-        self.model_name = model_name
-        self.max_seq_length = max_seq_length
-        self.quantization = quantization
-        self.device_map = device_map
-        self.cache_dir = cache_dir
-        self.use_flash_attention = use_flash_attention
+        # Prepare data
+        question = example["question"]
+        
+        # Handle different formats of choices
+        if isinstance(example["choices"], list):
+            choices = example["choices"]
+        elif isinstance(example["choices"], str) and example["choices"].startswith("["):
+            import ast
+            choices = ast.literal_eval(example["choices"]) if "[" in example["choices"] else example["choices"].split(",")
+        else:
+            choices = str(example["choices"]).split(",")
         
-        self.model = None
-        self.tokenizer = None
+        # Generate the prompt using prompt creator
+        prompt = self.prompt_creator.create_inference_prompt(question, choices)
         
-        # Load the model and tokenizer
-        self._load_model()
+        # Start timing
+        start_time = time.time()
         
-    def _load_model(self):
-        """Load the model and tokenizer with appropriate settings"""
-        from transformers import AutoModelForCausalLM, AutoTokenizer
-        import torch
+        if stream:
+            # Use streaming generation
+            streamer = self.model_handler.generate_with_streaming(
+                prompt=prompt,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                stream=True
+            )
+            
+            # Collect output from streamer
+            raw_response = ""
+            print("Model response:")
+            for text_chunk in streamer:
+                print(text_chunk, end="", flush=True)
+                raw_response += text_chunk
+            print("\n")
+        else:
+            # Generate without streaming
+            raw_response = self.model_handler.generate_with_streaming(
+                prompt=prompt,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                stream=False
+            )
         
-        # Load tokenizer
-        self.tokenizer = AutoTokenizer.from_pretrained(
-            self.model_name,
-            trust_remote_code=True,
-            cache_dir=self.cache_dir
-        )
+        response_time = time.time() - start_time
+        
+        # Parse the response using the response parser
+        predicted_answer, reasoning = self.response_parser.parse(raw_response)
         
-        # Prepare model loading kwargs
-        model_kwargs = {
-            "trust_remote_code": True,
-            "cache_dir": self.cache_dir,
-            "device_map": self.device_map,
+        # Prepare results
+        result = {
+            "question": question,
+            "choices": choices,
+            "predicted_answer": predicted_answer,
+            "reasoning": reasoning,
+            "response_time": response_time,
+            "raw_response": raw_response,
+            "prompt_type": self.prompt_creator.prompt_type,
         }
         
-        # Add Flash Attention if requested and available
-        if self.use_flash_attention:
-            try:
-                import flash_attn
-                model_kwargs["use_flash_attention_2"] = True
-                print("Flash Attention 2 enabled!")
-            except ImportError:
-                print("Flash Attention not available. For faster inference, install with: pip install flash-attn")
-        
-        # Add quantization if specified
-        if self.quantization == "4bit":
-            try:
-                from transformers import BitsAndBytesConfig
-                model_kwargs["quantization_config"] = BitsAndBytesConfig(
-                    load_in_4bit=True,
-                    bnb_4bit_compute_dtype=torch.bfloat16
-                )
-            except ImportError:
-                print("bitsandbytes not available, loading without 4-bit quantization")
-        elif self.quantization == "8bit":
-            model_kwargs["load_in_8bit"] = True
-        else:
-            model_kwargs["torch_dtype"] = torch.bfloat16
+        # Add task_id if available
+        if "task_id" in example:
+            result["task_id"] = example["task_id"]
+            
+        # Calculate metrics if label is provided
+        if "answer" in example:
+            label = example["answer"]
+            result["correct_answer"] = label
+            result["is_correct"] = predicted_answer == label
+            
+            # Calculate perplexity if requested
+            if hasattr(self.model_handler, "calculate_perplexity"):
+                perplexity = self.model_handler.calculate_perplexity(prompt, raw_response)
+                result["perplexity"] = perplexity
         
-        # Load the model
-        self.model = AutoModelForCausalLM.from_pretrained(
-            self.model_name,
-            **model_kwargs
-        )
+        # Restore original prompt type if it was overridden
+        if original_prompt_type is not None:
+            self.prompt_creator.set_prompt_type(original_prompt_type)
+            self.response_parser = ResponseParser.from_prompt_type(original_prompt_type)
+            
+        return result
+
+    def infer_batch(self, examples, temperature=0.7, max_tokens=1024, prompt_type=None, batch_size=4):
+        """Inference on a batch of examples"""
+        # Allow temporary override of prompt type
+        original_prompt_type = None
+        if prompt_type is not None:
+            original_prompt_type = self.prompt_creator.prompt_type
+            self.prompt_creator.set_prompt_type(prompt_type)
+            self.response_parser = ResponseParser.from_prompt_type(prompt_type)
         
-    def generate_with_streaming(self, prompt, temperature=0.7, max_tokens=1024, stream=True):
-        """
-        Generate text from the model with optional streaming
+        # Prepare all prompts
+        prompts = []
+        metadata = []
         
-        Args:
-            prompt: Input text prompt
-            temperature: Temperature for sampling (0 for deterministic)
-            max_tokens: Maximum number of tokens to generate
-            stream: Whether to stream the output
+        for i, example in enumerate(examples):
+            # Extract data
+            question = example["question"]
             
-        Returns:
-            String containing the generated text, or generator if streaming
-        """
-        import torch
+            # Handle different formats of choices
+            if isinstance(example["choices"], list):
+                choices = example["choices"]
+            elif isinstance(example["choices"], str) and example["choices"].startswith("["):
+                import ast
+                choices = ast.literal_eval(example["choices"]) if "[" in example["choices"] else example["choices"].split(",")
+            else:
+                choices = str(example["choices"]).split(",")
+            
+            # Generate the prompt using prompt creator
+            prompt = self.prompt_creator.create_inference_prompt(question, choices)
+            prompts.append(prompt)
+            
+            # Store metadata for later
+            meta = {
+                "question": question,
+                "choices": choices,
+                "index": i,
+            }
+            
+            # Add label if available
+            if "answer" in example:
+                meta["label"] = example["answer"]
+                
+            if "task_id" in example:
+                meta["task_id"] = example["task_id"]
+            
+            metadata.append(meta)
         
-        # Tokenize prompt
-        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
-        input_ids = inputs.input_ids
-        attention_mask = inputs.attention_mask
-        
-        # Set generation parameters
-        generation_config = {
-            "max_new_tokens": max_tokens,
-            "temperature": temperature,
-            "do_sample": temperature > 0,
-            "top_p": 0.95 if temperature > 0 else 1.0,
-            "repetition_penalty": 1.1,
-            "pad_token_id": self.tokenizer.eos_token_id,
-        }
+        # Process in batches
+        results = []
+        correct_count = 0
+        total_count = 0
+        perplexities = []
         
-        # If not streaming, do normal generation
-        if not stream:
-            with torch.no_grad():
-                outputs = self.model.generate(
-                    input_ids=input_ids,
-                    attention_mask=attention_mask,
-                    **generation_config
+        for i in range(0, len(prompts), batch_size):
+            batch_prompts = prompts[i:i+batch_size]
+            batch_meta = metadata[i:i+batch_size]
+            
+            # Process batch
+            start_time = time.time()
+            batch_responses = []
+            
+            for prompt in batch_prompts:
+                response = self.model_handler.generate_with_streaming(
+                    prompt=prompt,
+                    temperature=temperature,
+                    max_tokens=max_tokens,
+                    stream=False
                 )
+                batch_responses.append(response)
             
-            # Decode the generated text (skip the prompt)
-            generated_text = self.tokenizer.decode(
-                outputs[0][input_ids.shape[1]:], 
-                skip_special_tokens=True
-            )
+            batch_time = time.time() - start_time
             
-            return generated_text
+            # Process each response in the batch
+            for j, (response, meta) in enumerate(zip(batch_responses, batch_meta)):
+                # Parse response
+                predicted_answer, reasoning = self.response_parser.parse(response)
+                
+                # Create result
+                result = {
+                    "question": meta["question"],
+                    "choices": meta["choices"],
+                    "predicted_answer": predicted_answer,
+                    "reasoning": reasoning,
+                    "raw_response": response,
+                    "prompt_type": self.prompt_creator.prompt_type,
+                    "response_time": batch_time / len(batch_prompts),
+                }
+                
+                # Add task_id if available
+                if "task_id" in meta:
+                    result["task_id"] = meta["task_id"]
+                
+                # Add metrics if label available
+                if "label" in meta:
+                    label = meta["label"]
+                    result["correct_answer"] = label
+                    result["is_correct"] = predicted_answer == label
+                    
+                    # Update counts for accuracy
+                    total_count += 1
+                    if result["is_correct"]:
+                        correct_count += 1
+                        
+                    # Calculate perplexity if possible
+                    if hasattr(self.model_handler, "calculate_perplexity"):
+                        prompt = batch_prompts[j]
+                        perplexity = self.model_handler.calculate_perplexity(prompt, response)
+                        result["perplexity"] = perplexity
+                        perplexities.append(perplexity)
+                        
+                results.append(result)
         
-        # If streaming, yield generated tokens one by one
-        else:
-            generated = []
-            
-            # Initialize generator
-            with torch.no_grad():
-                generated_ids = self.model.generate(
-                    input_ids=input_ids,
-                    attention_mask=attention_mask,
-                    **generation_config,
-                    streamer=None  # Would need a custom streamer here if available
-                )
+        # Calculate aggregate metrics
+        summary_metrics = {}
+        if total_count > 0:
+            summary_metrics["accuracy"] = correct_count / total_count
+            summary_metrics["correct_count"] = correct_count
+            summary_metrics["total_count"] = total_count
             
-            # Decode the entire sequence at once (not truly streaming, but simpler)
-            full_text = self.tokenizer.decode(
-                generated_ids[0][input_ids.shape[1]:], 
-                skip_special_tokens=True
+            if perplexities:
+                summary_metrics["avg_perplexity"] = sum(perplexities) / len(perplexities)
+                summary_metrics["min_perplexity"] = min(perplexities)
+                summary_metrics["max_perplexity"] = max(perplexities)
+        
+        # Restore original prompt type if it was overridden
+        if original_prompt_type is not None:
+            self.prompt_creator.set_prompt_type(original_prompt_type)
+            self.response_parser = ResponseParser.from_prompt_type(original_prompt_type)
+            
+        return results, summary_metrics
+
+    def evaluate_dataset(self, dataset, temperature=0.7, max_tokens=1024, num_examples=None, 
+                        verbose=True, prompt_type=None, batch_size=4, log_to_wandb=False):
+        """Inference on a whole dataset with metrics calculation"""
+        # Allow overriding the prompt type for this evaluation
+        original_prompt_type = self.prompt_creator.prompt_type
+        if prompt_type is not None:
+            self.prompt_creator.set_prompt_type(prompt_type)
+            self.response_parser = ResponseParser.from_prompt_type(prompt_type)
+            
+        # Select subset if specified
+        if num_examples is not None:
+            dataset = dataset.select(range(min(num_examples, len(dataset))))
+
+        results = []
+        correct_count = 0
+        total_count = 0
+        perplexities = []
+        
+        # Process examples in batches
+        for i in range(0, len(dataset), batch_size):
+            batch_examples = dataset[i:i+batch_size]
+            
+            if verbose:
+                batch_desc = f"Batch {i//batch_size + 1}/{(len(dataset) + batch_size - 1) // batch_size}"
+                print(f"\nProcessing {batch_desc} with {len(batch_examples)} examples...")
+            
+            # Infer batch
+            batch_results, batch_metrics = self.infer_batch(
+                examples=batch_examples,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                batch_size=batch_size
             )
             
-            return full_text
+            # Update metrics
+            results.extend(batch_results)
+            if "correct_count" in batch_metrics:
+                correct_count += batch_metrics["correct_count"]
+                total_count += batch_metrics["total_count"]
+                
+                if verbose:
+                    batch_accuracy = batch_metrics["accuracy"]
+                    overall_accuracy = correct_count / total_count
+                    print(f"Batch accuracy: {batch_accuracy:.2%}, Overall: {overall_accuracy:.2%} ({correct_count}/{total_count})")
+                    
+            # Collect perplexities
+            if "avg_perplexity" in batch_metrics:
+                for result in batch_results:
+                    if "perplexity" in result:
+                        perplexities.append(result["perplexity"])
+                        
+        # Calculate final accuracy
+        accuracy = correct_count / total_count if total_count > 0 else 0.0
+        
+        if verbose:
+            prompt_type_str = self.prompt_creator.prompt_type
+            print(f"\nFinal accuracy with {prompt_type_str} prompts: {accuracy:.2%} ({correct_count}/{total_count})")
+            if perplexities:
+                avg_perplexity = sum(perplexities) / len(perplexities)
+                print(f"Average perplexity: {avg_perplexity:.4f}")
+        
+        # Prepare comprehensive summary
+        summary = {
+            "accuracy": accuracy,
+            "correct_count": correct_count,
+            "total_count": total_count,
+            "prompt_type": self.prompt_creator.prompt_type,
+            "results": results,
+        }
+        
+        # Add perplexity metrics if available
+        if perplexities:
+            summary["avg_perplexity"] = sum(perplexities) / len(perplexities)
+            summary["min_perplexity"] = min(perplexities)
+            summary["max_perplexity"] = max(perplexities)
+            
+        # Log results to wandb if requested
+        if log_to_wandb and wandb.run is not None:
+            metrics = {
+                "test/accuracy": accuracy,
+                "test/correct_count": correct_count,
+                "test/total_count": total_count,
+            }
+            if perplexities:
+                metrics["test/avg_perplexity"] = summary["avg_perplexity"]
+                metrics["test/min_perplexity"] = summary["min_perplexity"]
+                metrics["test/max_perplexity"] = summary["max_perplexity"]
+            
+            wandb.log(metrics)
+            
+            # Create a table of results for visualization if task_id exists
+            if "task_id" in dataset.features:
+                columns = ["task_id", "question", "correct_answer", "predicted_answer", "is_correct"]
+                table = wandb.Table(columns=columns)
+                
+                for res in results[:min(100, len(results))]:
+                    table.add_data(
+                        res.get("task_id", "unknown"),
+                        res["question"][:100] + "...",
+                        res.get("correct_answer", ""),
+                        res.get("predicted_answer", ""),
+                        res.get("is_correct", False)
+                    )
+                
+                wandb.log({"test_samples": table})
+        
+        # Restore original prompt type
+        self.prompt_creator.set_prompt_type(original_prompt_type)
+        self.response_parser = ResponseParser.from_prompt_type(original_prompt_type)
+            
+        return summary
+
+    def save_results(self, results, output_dir="./results"):
+        """Save evaluation results to file"""
+        os.makedirs(output_dir, exist_ok=True)
+
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        results_file = os.path.join(output_dir, f"results_{timestamp}.json")
+
+        # Create serializable results
+        serializable_results = {
+            "accuracy": results.get("accuracy", 0.0),
+            "correct_count": results.get("correct_count", 0),
+            "total_count": results.get("total_count", 0),
+            "timestamp": timestamp,
+            "prompt_type": results.get("prompt_type", "unknown"),
+        }
+
+        # Add perplexity metrics if available
+        if "avg_perplexity" in results:
+            serializable_results["avg_perplexity"] = results["avg_perplexity"]
+            serializable_results["min_perplexity"] = results["min_perplexity"]
+            serializable_results["max_perplexity"] = results["max_perplexity"]
+
+        # Process individual results
+        serializable_results["individual_results"] = []
+        for result in results["results"]:
+            # Skip perplexity in individual results to save space
+            result_copy = result.copy()
+            if "perplexity" in result_copy:
+                del result_copy["perplexity"]
+
+            # Convert choices if needed
+            choices = result_copy["choices"]
+            if not isinstance(choices, list):
+                try:
+                    import ast
+                    result_copy["choices"] = ast.literal_eval(choices)
+                except (SyntaxError, ValueError):
+                    pass
+
+            serializable_results["individual_results"].append(result_copy)
+
+        # Save to file
+        with open(results_file, "w") as f:
+            import json
+            json.dump(serializable_results, f, indent=2)
+
+        print(f"Results saved to {results_file}")
+        return results_file
 ```
 
-## Hardware Requirements and Optimization
+</details>
 
-### Flash Attention Benefits
+## Quick Start
 
-Flash Attention is a highly optimized implementation of the attention mechanism that:
+Here's a simple example of how to use the model:
 
-1. **Speeds up inference by 2-3x** compared to standard attention
-2. **Reduces memory usage** by avoiding materializing large attention matrices
-3. **Works perfectly with 4-bit quantization** for even further optimization
-4. **Scales better with sequence length**, which is important for complex coding questions
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-For the best performance, make sure to:
-- Install Flash Attention (`pip install flash-attn`)
-- Enable it when loading the model (see QwenModelHandler class)
-- Use with CUDA-compatible NVIDIA GPUs
+# Load the model and tokenizer
+model_id = "tuandunghcmut/Qwen25_Coder_MultipleChoice"
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    trust_remote_code=True
+)
+
+# Example question
+question = "What is the correct way to open a file in Python for reading?"
+choices = [
+    "open('file.txt', 'r')",
+    "file.open('file.txt', 'read')",
+    "read('file.txt')",
+    "File.open('file.txt')"
+]
+
+# Format the prompt
+prompt = f"""
+QUESTION:
+{question}
+
+CHOICES:
+{chr(65 + i)}. {choice}
+for i, choice in enumerate(choices)}
+
+Answer with a single letter from A through {chr(65 + len(choices) - 1)} without any additional explanation or commentary.
+"""
 
-### Hardware Recommendations
+# Generate response
+inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+outputs = model.generate(**inputs, max_new_tokens=10)
+response = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-For optimal performance, we recommend:
+print(f"Model's answer: {response}")
+```
 
-- **GPU**: NVIDIA GPU with at least 8GB VRAM (16GB+ recommended for larger models)
-- **RAM**: 16GB+ system RAM
-- **Storage**: At least 10GB free disk space for model files
-- **CPU**: Modern multi-core processor (for preprocessing)
+## Advanced Usage
 
-### Reducing Memory Usage
+### Using the MultipleChoiceTester Framework
 
-If you're facing memory constraints:
+For more advanced usage, you can use the provided `MultipleChoiceTester` framework:
 
 ```python
-# Use 4-bit quantization with Flash Attention for optimal memory-efficiency
+from save import QwenModelHandler, MultipleChoiceTester, PromptCreator
+
+# Initialize the model handler
 model_handler = QwenModelHandler(
     model_name="tuandunghcmut/Qwen25_Coder_MultipleChoice",
+    max_seq_length=2048,
     quantization="4bit",
-    use_flash_attention=True
+    device_map="auto"
 )
 
-# Further optimize with unsloth
-try:
-    from unsloth.models import FastLanguageModel
-    FastLanguageModel.for_inference(model_handler.model)
-    print("Using unsloth for additional optimization")
-except ImportError:
-    print("unsloth not available")
+# Create a prompt creator with YAML reasoning format
+prompt_creator = PromptCreator(PromptCreator.YAML_REASONING)
+
+# Initialize the tester
+tester = MultipleChoiceTester(model_handler, prompt_creator=prompt_creator)
+
+# Example question
+example = {
+    "question": "What is the correct way to open a file in Python for reading?",
+    "choices": [
+        "open('file.txt', 'r')",
+        "file.open('file.txt', 'read')",
+        "read('file.txt')",
+        "File.open('file.txt')"
+    ],
+    "answer": "A"  # Optional ground truth
+}
+
+# Get prediction with reasoning
+result = tester.infer_example(example, temperature=0.0001, stream=True)
+print(f"Predicted answer: {result['predicted_answer']}")
+print("Reasoning:")
+print(result['reasoning'])
 ```
 
-## Usage Example
+### Batch Processing
 
-Here's how to use these classes with Flash Attention enabled:
+You can also process multiple questions in batches:
 
 ```python
-# 1. Load the model with Flash Attention and 4-bit quantization
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
+# List of examples
+examples = [
+    {
+        "question": "What is the correct way to open a file in Python for reading?",
+        "choices": ["open('file.txt', 'r')", "file.open('file.txt', 'read')", "read('file.txt')", "File.open('file.txt')"],
+        "answer": "A"
+    },
+    # Add more examples...
+]
 
-hub_model_id = "tuandunghcmut/Qwen25_Coder_MultipleChoice"
+# Process batch
+results, metrics = tester.infer_batch(examples, batch_size=4)
+print(f"Batch accuracy: {metrics['accuracy']:.2%}")
+```
 
-# Create model handler with Flash Attention and 4-bit quantization
+### Streaming Inference
+
+The model supports streaming inference, which provides real-time output as the model generates its response. This is particularly useful for interactive applications and when you want to see the reasoning process in real-time.
+
+#### Basic Streaming Usage
+
+Here's how to use streaming inference:
+
+```python
+# Initialize model handler and tester as before
 model_handler = QwenModelHandler(
-    model_name=hub_model_id,
-    max_seq_length=2048,
-    quantization="4bit",
-    use_flash_attention=True
+    model_name="tuandunghcmut/Qwen25_Coder_MultipleChoice",
+    max_seq_length=2048
+)
+tester = MultipleChoiceTester(model_handler)
+
+# Example with streaming
+example = {
+    "question": "Which Python method is used to remove whitespace from both ends of a string?",
+    "choices": [
+        "strip()",
+        "trim()",
+        "clean()",
+        "remove_whitespace()"
+    ],
+    "answer": "A"
+}
+
+# Enable streaming with stream=True
+result = tester.infer_example(
+    example,
+    temperature=0.0001,
+    max_tokens=1024,
+    stream=True  # Enable streaming
 )
 
-# Optional: Use unsloth for even faster inference
-try:
-    from unsloth.models import FastLanguageModel
-    FastLanguageModel.for_inference(model_handler.model)
-    print("Using unsloth for faster inference")
-except ImportError:
-    print("unsloth not available, using standard inference")
+# The output will be printed in real-time as the model generates it
+# You can also access the complete response after generation
+print("\nFinal result:")
+print(f"Predicted answer: {result['predicted_answer']}")
+print("Complete reasoning:")
+print(result['reasoning'])
+```
 
-# 2. Create prompt creator with YAML reasoning format
-prompt_creator = PromptCreator(PromptCreator.YAML_REASONING)
+#### Advanced Streaming Patterns
 
-# 3. Example question
-question = "Which of the following correctly defines a list comprehension in Python?"
-choices = [
-    "[x**2 for x in range(10)]",
-    "for(x in range(10)) { return x**2; }",
-    "map(lambda x: x**2, range(10))",
-    "[for x in range(10): x**2]"
-]
+##### 1. Custom Stream Processing
+
+You can process the streamed output in custom ways:
+
+```python
+def process_stream(streamer):
+    """Custom stream processing function"""
+    collected_text = ""
+    for chunk in streamer:
+        # Process each chunk as it arrives
+        collected_text += chunk
+        # You can do custom processing here
+        # For example, parse partial YAML, update UI, etc.
+        yield chunk, collected_text
+
+# Use custom stream processing
+result = tester.infer_example(
+    example,
+    temperature=0.0001,
+    stream=True
+)
+
+# Process the stream with custom logic
+for chunk, full_text in process_stream(result['stream']):
+    # Do something with each chunk
+    print(f"Chunk: {chunk}")
+    print(f"Full text so far: {full_text}")
+```
+
+##### 2. YAML Streaming with Real-time Parsing
+
+When using YAML reasoning format, you can parse the output as it streams:
+
+```python
+import yaml
+from io import StringIO
+
+def parse_yaml_stream(streamer):
+    """Parse YAML content as it streams"""
+    buffer = StringIO()
+    for chunk in streamer:
+        buffer.write(chunk)
+        try:
+            # Try to parse the current buffer as YAML
+            yaml_content = yaml.safe_load(buffer.getvalue())
+            if yaml_content:
+                yield chunk, yaml_content
+        except yaml.YAMLError:
+            # Not enough content for valid YAML yet
+            continue
+
+# Use YAML streaming with parsing
+result = tester.infer_example(
+    example,
+    temperature=0.0001,
+    prompt_type=PromptCreator.YAML_REASONING,
+    stream=True
+)
 
-# 4. Create prompt and generate answer
-prompt = prompt_creator.create_inference_prompt(question, choices)
-response = model_handler.generate_with_streaming(prompt, temperature=0.0, stream=False)
+# Process YAML content as it streams
+for chunk, yaml_content in parse_yaml_stream(result['stream']):
+    if isinstance(yaml_content, dict):
+        # Access YAML fields as they become available
+        if 'understanding' in yaml_content:
+            print(f"Understanding: {yaml_content['understanding']}")
+        if 'reasoning' in yaml_content:
+            print(f"Reasoning: {yaml_content['reasoning']}")
+        if 'answer' in yaml_content:
+            print(f"Answer: {yaml_content['answer']}")
+```
+
+##### 3. Streaming with Progress Tracking
 
-# 5. Parse the response
-parser = ResponseParser(prompt_creator.parser_mode)
-answer, reasoning = parser.parse(response)
+You can track generation progress and timing:
 
-print(f"Question: {question}")
-print(f"Answer: {answer}")
-if reasoning:
-    print(f"Reasoning: {reasoning}")
+```python
+import time
+
+def stream_with_progress(streamer):
+    """Stream with progress tracking"""
+    start_time = time.time()
+    tokens_generated = 0
+    
+    for chunk in streamer:
+        tokens_generated += len(chunk.split())
+        elapsed = time.time() - start_time
+        tokens_per_second = tokens_generated / elapsed if elapsed > 0 else 0
+        
+        yield {
+            'chunk': chunk,
+            'tokens': tokens_generated,
+            'tokens_per_second': tokens_per_second,
+            'elapsed': elapsed
+        }
+
+# Use streaming with progress tracking
+result = tester.infer_example(
+    example,
+    temperature=0.0001,
+    stream=True
+)
+
+for progress in stream_with_progress(result['stream']):
+    print(f"Generated {progress['tokens']} tokens "
+          f"({progress['tokens_per_second']:.2f} tokens/sec)")
+    print(f"Chunk: {progress['chunk']}")
 ```
 
-## Troubleshooting
+#### Implementation Details
+
+The streaming implementation uses Unsloth's optimized inference with the following key features:
+
+1. **Efficient Token Generation**
+   - Uses Unsloth's `FastLanguageModel` for optimized inference
+   - Implements streaming using `TextIteratorStreamer`
+   - Supports both greedy and temperature-based sampling
+
+2. **Memory Management**
+   - Streams tokens without storing the entire response in memory
+   - Efficiently handles long responses
+   - Supports batch processing with streaming
+
+3. **Performance Optimizations**
+   - Uses `use_cache=True` for faster generation
+   - Implements `min_p` sampling for better quality
+   - Supports 4-bit quantization for reduced memory usage
+
+4. **Error Handling**
+   - Gracefully handles streaming interruptions
+   - Provides partial results if generation is interrupted
+   - Maintains context for resumed generation
+
+The streaming output will show the model's reasoning process in real-time, including:
+- Understanding of the question
+- Analysis of each option
+- Step-by-step reasoning
+- Final conclusion
+- Answer selection
+
+This is particularly useful for:
+- Debugging model behavior
+- Creating interactive demos
+- Understanding the model's reasoning process
+- Providing immediate feedback to users
+- Building real-time applications
 
-### Common Issues
+## Model Features
 
-1. **Flash Attention Installation Issues**: If you encounter problems installing `flash-attn`:
-   ```bash
-   # Try with specific CUDA version (e.g., for CUDA 11.8)
-   pip install flash-attn==2.3.4+cu118 --no-build-isolation
-   
-   # For older GPUs
-   pip install flash-attn==2.3.4 --no-build-isolation
-   ```
+- **YAML-Based Reasoning**: The model provides structured reasoning in YAML format
+- **Multiple Prompt Types**: Supports both basic and YAML-formatted reasoning prompts
+- **Batch Processing**: Efficiently process multiple questions at once
+- **Performance Metrics**: Tracks accuracy, perplexity, and response times
+- **Streaming Support**: Real-time output streaming for interactive use
 
-2. **CUDA Out of Memory**: Try combining 4-bit quantization with Flash Attention.
-   ```python
-   model_handler = QwenModelHandler(
-       model_name=hub_model_id, 
-       quantization="4bit",
-       use_flash_attention=True
-   )
-   ```
+## Examples
 
-3. **Module Not Found Errors**: Make sure you've installed all required packages.
-   ```bash
-   pip install transformers torch unsloth datasets pyyaml bitsandbytes flash-attn
-   ```
+Check out the [example notebook](https://colab.research.google.com/drive/1YOUR_NOTEBOOK_ID) for more detailed usage examples and demonstrations.
 
-4. **Parsing Errors**: If the model isn't producing valid YAML responses, try adjusting the temperature:
-   ```python
-   response = model_handler.generate_with_streaming(prompt, temperature=0.0, stream=False)
-   ```
+## Contributing
 
-### Getting Help
+Contributions are welcome! Please feel free to submit a Pull Request.
 
-If you encounter issues, check the [model repository on Hugging Face](https://huggingface.co/tuandunghcmut/Qwen25_Coder_MultipleChoice) for updates and community discussions.
+## License
 
-This guide provides you with all the necessary code and optimization techniques to use the model effectively for multiple-choice coding questions.
+This project is licensed under the MIT License - see the LICENSE file for details.