update handler

README.md

@@ -0,0 +1,81 @@
+# Phi-4 Mini Inference Endpoint Handler
+
+This repository contains code for deploying the Phi-4 Mini model to a HuggingFace Inference Endpoint with an OpenAI-compatible API format.
+
+## Setup
+
+1. Install the required dependencies:
+   ```
+   pip install -r requirements.txt
+   ```
+
+2. Set the environment variable to your model path (optional if model is in the same directory):
+   ```
+   export MODEL_PATH=/path/to/your/model
+   ```
+
+## Usage
+
+When deploying to a HuggingFace Inference Endpoint, the `handler.py` file will be used to process requests. The endpoint accepts requests in an OpenAI-compatible format:
+
+```json
+{
+  "messages": [
+    {"role": "system", "content": "You are a helpful assistant."},
+    {"role": "user", "content": "Tell me about language models."}
+  ],
+  "max_tokens": 256,
+  "temperature": 0.7,
+  "top_p": 1.0,
+  "n": 1,
+  "stop": ["\n", "User:"]
+}
+```
+
+The endpoint returns responses in an OpenAI-compatible format:
+
+```json
+{
+  "id": "cmpl-12345",
+  "object": "chat.completion",
+  "created": 0,
+  "model": "phi4-mini-raw",
+  "choices": [
+    {
+      "index": 0,
+      "message": {
+        "role": "assistant",
+        "content": "Language models are computational systems designed to understand and generate human language..."
+      },
+      "finish_reason": "stop"
+    }
+  ],
+  "usage": {
+    "prompt_tokens": 42,
+    "completion_tokens": 156,
+    "total_tokens": 198
+  }
+}
+```
+
+## Local Testing
+
+To test the handler locally before deployment:
+
+```python
+from handler import EndpointHandler
+
+# Initialize the handler with your model path
+handler = EndpointHandler("./phi4-mini-raw")
+
+# Test with a sample request
+request = {
+  "messages": [
+    {"role": "system", "content": "You are a helpful assistant."},
+    {"role": "user", "content": "Hello, how are you?"}
+  ]
+}
+
+response = handler(request)
+print(response)
+``` 

handler.py

@@ -1,436 +1,166 @@
 import os
-import torch
-import logging
-import time
-import traceback
 import json
-from typing import Dict, List, Any, Union, Generator
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from threading import Thread
-
-# Set up logging
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
-)
-logger = logging.getLogger(__name__)
+import torch
+from transformers import pipeline, AutoTokenizer
+from typing import Dict, List, Any, Optional, Union
 
 class EndpointHandler:
     def __init__(self, path=""):
-        """
-        Initialize the model and tokenizer for Phi-4 inference.
-        
-        Args:
-            path (str): Path to the model directory
-        """
-        # Set default parameters for inference
-        self.max_new_tokens = 1024  # Keep at 1024 to avoid timeouts
-        self.temperature = 0.7
-        self.top_p = 0.9
-        self.do_sample = True
-        
-        # Determine if CUDA is available
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
+        # Initialize model and tokenizer
+        self.model_path = path if path else os.environ.get("MODEL_PATH", "")
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
         
-        logger.info(f"Initializing model from {path} on {self.device}")
+        # Create text generation pipeline
+        self.pipe = pipeline(
+            "text-generation",
+            model=self.model_path,
+            tokenizer=self.tokenizer,
+            torch_dtype=torch.float16,
+            device_map="auto",
+            return_full_text=False  # Only return the generated text, not the prompt
+        )
         
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Handle inference request in OpenAI-like format"""
         try:
-            # Load tokenizer - use original model ID as fallback
-            # This helps with common tokenizer mismatch issues
-            try:
-                self.tokenizer = AutoTokenizer.from_pretrained(path)
-                logger.info(f"Loaded tokenizer from local path")
-            except Exception as e:
-                logger.warning(f"Failed to load tokenizer from local path: {e}")
-                self.tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-4-mini-instruct")
-                logger.info("Loaded tokenizer from microsoft/Phi-4-mini-instruct")
-            
-            # Ensure tokenizer has EOS token set
-            if self.tokenizer.eos_token_id is None:
-                logger.warning("EOS token not set in tokenizer, using default")
-                self.tokenizer.eos_token_id = 199999  # Phi-4's default EOS token
-            
-            # Load model with appropriate settings
-            self.model = AutoModelForCausalLM.from_pretrained(
-                path,
-                torch_dtype=self.dtype,
-                device_map="auto" if self.device == "cuda" else None,
-                trust_remote_code=True
-            )
-            
-            # Move model to device if CPU
-            if self.device == "cpu":
-                self.model = self.model.to(self.device)
+            # Parse input data
+            inputs = self._parse_input(data)
             
-            # Set model to evaluation mode
-            self.model.eval()
-            
-            # Print diagnostic information
-            logger.info(f"Model loaded on {self.device} using {self.dtype}")
-            logger.info(f"Tokenizer vocabulary size: {len(self.tokenizer)}")
-            logger.info(f"Model vocabulary size: {self.model.config.vocab_size}")
-            logger.info(f"Model embedding size: {self.model.get_input_embeddings().weight.shape}")
-            
-            if len(self.tokenizer) != self.model.config.vocab_size:
-                logger.warning(f"Tokenizer vocab size ({len(self.tokenizer)}) doesn't match model vocab size ({self.model.config.vocab_size})")
+            # Generate response
+            outputs = self._generate(inputs)
             
+            # Format response in OpenAI-like format
+            return self._format_response(outputs, inputs)
         except Exception as e:
-            logger.error(f"Error during model initialization: {str(e)}")
-            logger.error(traceback.format_exc())
-            raise
-
-    def format_prompt_with_system(self, user_message, system_message=None):
-        """
-        Format the prompt with system and user messages according to Phi-4 format.
+            return {
+                "error": {
+                    "message": str(e),
+                    "type": "invalid_request_error",
+                    "code": 400
+                }
+            }
+    
+    def _parse_input(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Parse input data to extract generation parameters"""
+        # Extract messages
+        messages = data.get("messages", [])
+        if not messages:
+            raise ValueError("No messages provided")
         
-        Args:
-            user_message (str): The user's message
-            system_message (str, optional): The system message/instruction
-            
-        Returns:
-            str: Formatted prompt ready for the model
-        """
-        # Format using Phi-4's expected chat template:
-        # <|system|>
-        # {system_message}
-        # <|user|>
-        # {user_message}
-        # <|assistant|>
+        # Convert messages to prompt
+        prompt = self._convert_messages_to_prompt(messages)
         
-        if system_message:
-            prompt = f"<|system|>\n{system_message}\n<|user|>\n{user_message}\n<|assistant|>"
-        else:
-            # If no system message, just use user message with assistant tag
-            prompt = f"<|user|>\n{user_message}\n<|assistant|>"
+        # Extract generation parameters with defaults
+        generation_params = {
+            "max_tokens": data.get("max_tokens", 256),
+            "temperature": data.get("temperature", 0.7),
+            "top_p": data.get("top_p", 1.0),
+            "n": data.get("n", 1),
+            "stream": data.get("stream", False),
+            "stop": data.get("stop", None),
+            "presence_penalty": data.get("presence_penalty", 0.0),
+            "frequency_penalty": data.get("frequency_penalty", 0.0),
+        }
         
-        logger.info(f"Formatted prompt with {'system message and ' if system_message else ''}user message")
+        return {
+            "prompt": prompt,
+            "messages": messages,
+            "generation_params": generation_params
+        }
+    
+    def _convert_messages_to_prompt(self, messages: List[Dict[str, str]]) -> str:
+        """Convert list of messages to a prompt string"""
+        prompt = ""
+        for message in messages:
+            role = message.get("role", "")
+            content = message.get("content", "")
+            
+            if role == "system":
+                prompt += f"System: {content}\n\n"
+            elif role == "user":
+                prompt += f"User: {content}\n\n"
+            elif role == "assistant":
+                prompt += f"Assistant: {content}\n\n"
+            
+        # Add final assistant prompt
+        prompt += "Assistant: "
         return prompt
-
-    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        """
-        Process the input data and generate a response using the Phi-4 model.
+    
+    def _generate(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
+        """Generate response using the pipeline"""
+        prompt = inputs["prompt"]
+        params = inputs["generation_params"]
         
-        Args:
-            data (Dict[str, Any]): Input data containing the prompt and generation parameters
-            
-        Returns:
-            Dict[str, Any]: Model response
-        """
-        start_time = time.time()
-        logger.info(f"Starting request processing")
+        # Count input tokens
+        input_tokens = len(self.tokenizer.encode(prompt))
         
-        try:
-            # Extract input parameters with defaults
-            if "inputs" not in data:
-                logger.warning("No 'inputs' field in request data")
-                error_msg = "Missing 'inputs' field in request"
-                return self._format_error_response(error_msg)
-            
-            # Track user and system messages
-            user_message = ""
-            system_message = None
-            
-            # Handle different input formats
-            # 1. Direct string input
-            if isinstance(data["inputs"], str):
-                user_message = data["inputs"]
-                system_message = data.get("parameters", {}).get("system_message", None)
-                
-            # 2. Dict with messages format
-            elif isinstance(data["inputs"], dict) and "messages" in data["inputs"]:
-                messages = data["inputs"]["messages"]
-                
-                # Extract system and user messages for prompt formatting
-                for msg in messages:
-                    if msg.get("role") == "system":
-                        system_message = msg.get("content", "")
-                    elif msg.get("role") == "user":
-                        user_message = msg.get("content", "")
-                
-            # 3. Direct messages list format
-            elif isinstance(data["inputs"], list):
-                messages = data["inputs"]
-                
-                # Extract system and user messages for prompt formatting
-                for msg in messages:
-                    if msg.get("role") == "system":
-                        system_message = msg.get("content", "")
-                    elif msg.get("role") == "user":
-                        user_message = msg.get("content", "")
-            else:
-                logger.warning(f"Unsupported input format: {type(data['inputs'])}")
-                error_msg = "Unsupported input format. Expected string or messages object."
-                return self._format_error_response(error_msg)
-            
-            logger.info(f"Extracted user message length: {len(user_message)} characters")
-            if system_message:
-                logger.info(f"Extracted system message length: {len(system_message)} characters")
-            
-            # Format the prompt with system and user messages
-            prompt = self.format_prompt_with_system(user_message, system_message)
-            
-            parameters = data.get("parameters", {})
-            
-            logger.info(f"Processing input with {len(prompt)} characters")
-            
-            # Get generation parameters with fallbacks to defaults
-            max_new_tokens = min(parameters.get("max_new_tokens", self.max_new_tokens), 2048)
-            temperature = parameters.get("temperature", self.temperature)
-            top_p = parameters.get("top_p", self.top_p)
-            do_sample = parameters.get("do_sample", self.do_sample)
-            
-            logger.info(f"Generation parameters: max_new_tokens={max_new_tokens}, temperature={temperature}, top_p={top_p}, do_sample={do_sample}")
-            
-            # Manually implement generation to avoid token index errors
-            try:
-                input_ids = self.tokenizer.encode(prompt, return_tensors="pt").to(self.device)
-                logger.info(f"Input tokens shape: {input_ids.shape}")
-                
-                # Create attention mask
-                attention_mask = torch.ones_like(input_ids)
-                
-                # Perform safe generation with error handling for out-of-vocabulary issues
-                response_text = self._safe_generate(
-                    input_ids, 
-                    attention_mask, 
-                    max_new_tokens, 
-                    temperature, 
-                    top_p, 
-                    do_sample, 
-                    prompt
-                )
-                
-                logger.info(f"Response generation completed, text length: {len(response_text) if isinstance(response_text, str) else 'N/A'}")
-                
-                # Format and return response in OpenAI format
-                if isinstance(response_text, str):
-                    return self._format_openai_response(
-                        response_text, 
-                        input_ids.shape[1], 
-                        len(self.tokenizer.encode(response_text)) if response_text else 0
-                    )
-                else:
-                    return self._format_error_response(f"Error during generation: {response_text}")
-                
-            except RuntimeError as e:
-                logger.error(f"Runtime Error during generation: {str(e)}")
-                logger.error(traceback.format_exc())
-                return self._format_error_response(f"Error during generation: {str(e)}")
-                
-        except Exception as e:
-            logger.error(f"Unexpected error during request processing: {str(e)}")
-            logger.error(traceback.format_exc())
-            return self._format_error_response(f"Unexpected error: {str(e)}")
-        finally:
-            duration = time.time() - start_time
-            logger.info(f"Request processing completed in {duration:.2f} seconds")
-    
-    def _safe_generate(self, input_ids, attention_mask, max_new_tokens, temperature, top_p, do_sample, prompt):
-        """Safely generate text handling potential token index errors"""
-        try:
-            with torch.no_grad():
-                logger.info("Starting safe generation")
-                
-                # Get the input text to exclude from final output
-                input_text = prompt
-                logger.info(f"Input prompt length: {len(input_text)} characters")
-                
-                # Generate one token at a time to avoid index errors
-                # Increase from 250 to 500 to allow for longer completions
-                max_steps = min(max_new_tokens, 500)
-                current_ids = input_ids.clone()
-                
-                logger.info(f"Generating up to {max_steps} tokens")
-                
-                # Keep track of last 5 tokens to detect repetition
-                last_tokens = []
-                repetition_detected = False
-                
-                for i in range(max_steps):
-                    if i % 50 == 0:
-                        logger.info(f"Generated {i} tokens so far")
-                    
-                    # Get logits for next token
-                    outputs = self.model(
-                        input_ids=current_ids,
-                        attention_mask=attention_mask,
-                        return_dict=True
-                    )
-                    
-                    next_token_logits = outputs.logits[:, -1, :]
-                    
-                    # Apply temperature and sampling
-                    if temperature > 0:
-                        next_token_logits = next_token_logits / temperature
-                    
-                    if do_sample:
-                        # Apply top_p sampling
-                        sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
-                        cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
-                        
-                        # Remove tokens with cumulative probability above the threshold
-                        sorted_indices_to_remove = cumulative_probs > top_p
-                        # Shift the indices to the right to keep also the first token above the threshold
-                        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
-                        sorted_indices_to_remove[..., 0] = 0
-                        
-                        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
-                        next_token_logits[indices_to_remove] = -float('Inf')
-                        
-                        # Sample from the filtered distribution
-                        probs = torch.softmax(next_token_logits, dim=-1)
-                        next_token = torch.multinomial(probs, num_samples=1)
-                    else:
-                        # Take the token with highest probability
-                        next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
-                    
-                    # Add the predicted token to the sequence
-                    current_ids = torch.cat([current_ids, next_token], dim=-1)
-                    attention_mask = torch.cat([attention_mask, torch.ones_like(next_token)], dim=-1)
-                    
-                    # Add to last tokens list for repetition detection
-                    last_tokens.append(next_token.item())
-                    if len(last_tokens) > 5:
-                        last_tokens.pop(0)
-                    
-                    # Check for repetition (if we have at least 5 tokens)
-                    if len(last_tokens) >= 5:
-                        # Check if all last 5 tokens are the same
-                        if len(set(last_tokens)) == 1:
-                            logger.warning(f"Repetition detected after {i+1} tokens, stopping generation")
-                            repetition_detected = True
-                            break
-                    
-                    # Check if we've generated an EOS token
-                    if next_token[0, 0].item() == self.tokenizer.eos_token_id:
-                        logger.info(f"EOS token generated after {i+1} tokens")
-                        break
-                
-                # Decode the generated sequence
-                generated_text = self.tokenizer.decode(current_ids[0], skip_special_tokens=True)
-                logger.info(f"Decoded generated text: {len(generated_text)} characters")
-                
-                # Return only the newly generated text (after the assistant tag)
-                split_text = generated_text.split("<|assistant|>")
-                if len(split_text) > 1:
-                    response_text = split_text[1].strip()
-                    logger.info(f"Extracted assistant response: {len(response_text)} characters")
-                    
-                    # Check if the response text ends with a complete sentence
-                    if not repetition_detected and not response_text.endswith(('.', '!', '?', ':', ';', '"', "'", ')', ']', '}')):
-                        # Add an ellipsis to indicate truncation
-                        response_text += "..."
-                        logger.info("Added ellipsis to incomplete sentence")
-                    
-                else:
-                    # Fallback if the expected format is not found
-                    logger.warning("Could not find assistant tag in generated text")
-                    response_text = generated_text
-                    
-                return response_text
-            
-        except Exception as e:
-            logger.error(f"Error in _safe_generate: {str(e)}")
-            logger.error(traceback.format_exc())
-            return f"Generation error: {str(e)}. Please try a simpler input."
-    
-    def _format_openai_response(self, response_text, prompt_tokens, completion_tokens):
-        """Format the response in OpenAI-style format"""
-        try:
-            # Create a response ID
-            response_id = f"phi4-{int(time.time())}"
-            
-            # Build OpenAI-compatible response
-            openai_response = {
-                "id": response_id,
-                "object": "chat.completion",
-                "created": int(time.time()),
-                "model": "phi-4-mini",
-                "choices": [
-                    {
-                        "index": 0,
-                        "message": {
-                            "role": "assistant",
-                            "content": response_text
-                        },
-                        "finish_reason": "stop"
-                    }
-                ],
-                "usage": {
-                    "prompt_tokens": prompt_tokens,
-                    "completion_tokens": completion_tokens,
-                    "total_tokens": prompt_tokens + completion_tokens
-                }
-            }
-            
-            # For compatibility with Hugging Face UI, include the generated_text field
-            openai_response["generated_text"] = response_text
-            
-            logger.info(f"Formatted OpenAI-style response: {len(json.dumps(openai_response))} bytes")
-            return openai_response
-            
-        except Exception as e:
-            logger.error(f"Error formatting OpenAI response: {str(e)}")
-            # Fall back to simple response
-            return {"generated_text": response_text}
+        # Convert OpenAI-like parameters to pipeline parameters
+        generation_kwargs = {
+            "max_new_tokens": params["max_tokens"],
+            "temperature": params["temperature"],
+            "top_p": params["top_p"],
+            "num_return_sequences": params["n"],
+            "do_sample": params["temperature"] > 0,
+        }
+        
+        # Add stopping criteria if provided
+        if params["stop"]:
+            generation_kwargs["stopping_criteria"] = params["stop"]
+        
+        # Generate output using the pipeline
+        pipeline_outputs = self.pipe(
+            prompt,
+            **generation_kwargs
+        )
+        
+        # Extract generated texts
+        generated_texts = []
+        for output in pipeline_outputs:
+            gen_text = output["generated_text"]
+            
+            # Apply stop sequences if provided
+            if params["stop"]:
+                for stop in params["stop"]:
+                    if stop in gen_text:
+                        gen_text = gen_text[:gen_text.find(stop)]
+            
+            generated_texts.append(gen_text)
+        
+        # Count completion tokens
+        completion_tokens = [len(self.tokenizer.encode(text)) for text in generated_texts]
+        
+        return {
+            "generated_texts": generated_texts,
+            "prompt_tokens": input_tokens,
+            "completion_tokens": completion_tokens,
+        }
     
-    def _format_error_response(self, error_message):
-        """Format an error response in OpenAI-style format"""
-        try:
-            error_response = {
-                "id": f"phi4-error-{int(time.time())}",
-                "object": "chat.completion",
-                "created": int(time.time()),
-                "model": "phi-4-mini",
-                "choices": [
-                    {
-                        "index": 0,
-                        "message": {
-                            "role": "assistant",
-                            "content": f"Error: {error_message}"
-                        },
-                        "finish_reason": "error"
-                    }
-                ],
-                "usage": {
-                    "prompt_tokens": 0,
-                    "completion_tokens": 0,
-                    "total_tokens": 0
+    def _format_response(self, outputs: Dict[str, Any], inputs: Dict[str, Any]) -> Dict[str, Any]:
+        """Format response in OpenAI-like format"""
+        generated_texts = outputs["generated_texts"]
+        prompt_tokens = outputs["prompt_tokens"]
+        completion_tokens = outputs["completion_tokens"]
+        
+        choices = []
+        for i, text in enumerate(generated_texts):
+            choices.append({
+                "index": i,
+                "message": {
+                    "role": "assistant",
+                    "content": text
                 },
-                "error": {
-                    "message": error_message,
-                    "type": "invalid_request_error",
-                    "code": "error"
-                }
+                "finish_reason": "stop"
+            })
+        
+        return {
+            "id": f"cmpl-{hash(inputs['prompt']) % 10000}",
+            "object": "chat.completion",
+            "created": int(torch.cuda.current_device()) if torch.cuda.is_available() else 0,
+            "model": os.path.basename(self.model_path),
+            "choices": choices,
+            "usage": {
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": sum(completion_tokens),
+                "total_tokens": prompt_tokens + sum(completion_tokens)
             }
-            
-            # For compatibility with Hugging Face UI, include the generated_text field
-            error_response["generated_text"] = f"Error: {error_message}"
-            
-            logger.info(f"Formatted error response: {len(json.dumps(error_response))} bytes")
-            return error_response
-            
-        except Exception as e:
-            logger.error(f"Error formatting error response: {str(e)}")
-            # Fall back to simple error response
-            return {"generated_text": f"Error: {error_message}"}
-
-# For local testing
-if __name__ == "__main__":
-    # Example usage
-    handler = EndpointHandler()
-    
-    # Test with messages format
-    test_with_messages = {
-        "inputs": {
-            "messages": [
-                {"role": "system", "content": "You are an AI assistant that provides helpful, accurate, and concise information about AI models."},
-                {"role": "user", "content": "What are the major features of Phi-4?"}
-            ]
-        }
-    }
-    
-    # Run the test
-    result = handler(test_with_messages)
-    print(json.dumps(result, indent=2))
+        } 

requirements.txt

@@ -0,0 +1,4 @@
+torch>=2.0.0
+transformers>=4.30.0
+accelerate>=0.21.0
+sentencepiece>=0.1.99 

test_handler.py

@@ -0,0 +1,71 @@
+import json
+import argparse
+import sys
+from handler import EndpointHandler
+
+def test_inference(model_path=".", prompt=None, max_tokens=150, temperature=0.7):
+    """
+    Test the inference endpoint handler with a sample request.
+    
+    Args:
+        model_path: Path to the model directory
+        prompt: Custom prompt to use (optional)
+        max_tokens: Maximum number of tokens to generate
+        temperature: Temperature for generation
+    """
+    try:
+        print(f"Initializing handler with model path: {model_path}")
+        handler = EndpointHandler(model_path)
+        
+        # Default or custom prompt
+        if prompt is None:
+            messages = [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": "Explain quantum computing in simple terms."}
+            ]
+        else:
+            messages = [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": prompt}
+            ]
+        
+        # Sample request with OpenAI-like format
+        request = {
+            "messages": messages,
+            "max_tokens": max_tokens,
+            "temperature": temperature,
+            "top_p": 0.95
+        }
+        
+        print("Sending request to handler...")
+        print(f"Request: {json.dumps(request, indent=2)}")
+        
+        # Generate response
+        response = handler(request)
+        
+        # Print response in a readable format
+        print("\nResponse:")
+        print(json.dumps(response, indent=2))
+        
+        return response
+    
+    except Exception as e:
+        print(f"Error during inference: {str(e)}", file=sys.stderr)
+        import traceback
+        traceback.print_exc()
+        return {"error": str(e)}
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Test Phi-4 Mini inference")
+    parser.add_argument("--model_path", type=str, default=".", help="Path to the model directory")
+    parser.add_argument("--prompt", type=str, help="Custom prompt to use")
+    parser.add_argument("--max_tokens", type=int, default=150, help="Maximum number of tokens to generate")
+    parser.add_argument("--temperature", type=float, default=0.7, help="Temperature for generation")
+    
+    args = parser.parse_args()
+    test_inference(
+        model_path=args.model_path,
+        prompt=args.prompt,
+        max_tokens=args.max_tokens,
+        temperature=args.temperature
+    )