modify the handler

handler.py

@@ -1,9 +1,17 @@
 import os
 import torch
-from typing import Dict, List, Any
+import logging
+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__)
+
 class EndpointHandler:
     def __init__(self, path=""):
         """
@@ -13,7 +21,7 @@ class EndpointHandler:
             path (str): Path to the model directory
         """
         # Set default parameters for inference
-        self.max_new_tokens = 4096
+        self.max_new_tokens = 1024  # Reduced from 4096 to avoid memory issues
         self.temperature = 0.7
         self.top_p = 0.9
         self.do_sample = True
@@ -22,41 +30,53 @@ class EndpointHandler:
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
         
-        # Load tokenizer
-        self.tokenizer = AutoTokenizer.from_pretrained(path)
-        
-        # 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)
-        
-        # Set model to evaluation mode
-        self.model.eval()
-        
-        print(f"Model loaded on {self.device} using {self.dtype}")
-
-    def format_prompt(self, prompt: str) -> str:
-        """
-        Format the user prompt for Phi-4 model.
-        
-        Args:
-            prompt (str): User input prompt
+        logger.info(f"Initializing model from {path} on {self.device}")
+        
+        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")
             
-        Returns:
-            str: Formatted prompt
-        """
-        # For Phi-4-mini-instruct, the prompt format is simple
-        # You may need to adjust this based on your specific fine-tuning
-        return prompt
+            # 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)
+            
+            # 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})")
+            
+        except Exception as e:
+            logger.error(f"Error during model initialization: {e}")
+            raise
 
-    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+    def __call__(self, data: Dict[str, Any]) -> Union[Dict[str, str], Generator]:
         """
         Process the input data and generate a response using the Phi-4 model.
         
@@ -64,89 +84,126 @@ class EndpointHandler:
             data (Dict[str, Any]): Input data containing the prompt and generation parameters
             
         Returns:
-            Dict[str, Any]: Model response
+            Dict[str, str] or Generator: Model response or stream
         """
-        # Extract input parameters with defaults
-        prompt = data.pop("inputs", "")
-        parameters = data.pop("parameters", {})
-        
-        # Get generation parameters with fallbacks to defaults
-        max_new_tokens = parameters.get("max_new_tokens", self.max_new_tokens)
-        temperature = parameters.get("temperature", self.temperature)
-        top_p = parameters.get("top_p", self.top_p)
-        do_sample = parameters.get("do_sample", self.do_sample)
-        stream = parameters.get("stream", False)
-        
-        # Format the prompt according to model requirements
-        formatted_prompt = self.format_prompt(prompt)
-        
-        # Tokenize the input
-        inputs = self.tokenizer(formatted_prompt, return_tensors="pt").to(self.device)
-        
-        # Handle streaming if requested
-        if stream:
-            return self._generate_stream(inputs, max_new_tokens, temperature, top_p, do_sample)
-        else:
-            return self._generate(inputs, max_new_tokens, temperature, top_p, do_sample)
+        try:
+            # Extract input parameters with defaults
+            if "inputs" not in data:
+                logger.warning("No 'inputs' field in request data")
+                return {"error": "Missing 'inputs' field in request"}
+                
+            prompt = data.get("inputs", "")
+            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)
+            stream = parameters.get("stream", False)
+            
+            # Tokenize the input safely
+            inputs = self.tokenizer(prompt, return_tensors="pt")
+            logger.info(f"Input tokens shape: {inputs.input_ids.shape}")
+            
+            # Move to device
+            inputs = {k: v.to(self.device) for k, v in inputs.items()}
+            
+            # Handle streaming if requested
+            if stream:
+                return self._generate_stream(inputs, max_new_tokens, temperature, top_p, do_sample)
+            else:
+                return self._generate(inputs, max_new_tokens, temperature, top_p, do_sample)
+                
+        except Exception as e:
+            logger.error(f"Error during generation: {e}")
+            return {"error": str(e)}
     
     def _generate(self, inputs, max_new_tokens, temperature, top_p, do_sample):
         """Generate text non-streaming mode"""
-        with torch.no_grad():
-            outputs = self.model.generate(
-                **inputs,
-                max_new_tokens=max_new_tokens,
-                temperature=temperature,
-                top_p=top_p,
-                do_sample=do_sample,
-                pad_token_id=self.tokenizer.eos_token_id
-            )
-        
-        # Decode the generated text
-        generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        # Return only the newly generated text (without the prompt)
-        prompt_length = len(self.tokenizer.decode(inputs.input_ids[0], skip_special_tokens=True))
-        response_text = generated_text[prompt_length:]
-        
-        return {"generated_text": response_text}
+        try:
+            with torch.no_grad():
+                generation_config = {
+                    "max_new_tokens": max_new_tokens,
+                    "temperature": temperature,
+                    "top_p": top_p,
+                    "do_sample": do_sample,
+                    "pad_token_id": self.tokenizer.eos_token_id
+                }
+                
+                logger.info(f"Generating with config: {generation_config}")
+                
+                outputs = self.model.generate(
+                    inputs.input_ids,
+                    attention_mask=inputs.attention_mask if hasattr(inputs, 'attention_mask') else None,
+                    **generation_config
+                )
+            
+            # Decode the generated text
+            generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            
+            # Return only the newly generated text (without the prompt)
+            input_text = self.tokenizer.decode(inputs.input_ids[0], skip_special_tokens=True)
+            
+            if generated_text.startswith(input_text):
+                response_text = generated_text[len(input_text):]
+            else:
+                # Fallback if the decoded text doesn't start with the input
+                response_text = generated_text
+                
+            logger.info(f"Generated {len(response_text)} characters")
+            return {"generated_text": response_text}
+            
+        except Exception as e:
+            logger.error(f"Error in _generate: {e}")
+            return {"error": str(e)}
     
     def _generate_stream(self, inputs, max_new_tokens, temperature, top_p, do_sample):
         """Generate text in streaming mode"""
-        # Create a streamer object
-        streamer = TextIteratorStreamer(self.tokenizer, skip_special_tokens=True)
-        
-        # Set up generation in a separate thread
-        generation_kwargs = dict(
-            **inputs,
-            streamer=streamer,
-            max_new_tokens=max_new_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            do_sample=do_sample,
-            pad_token_id=self.tokenizer.eos_token_id
-        )
-        
-        thread = Thread(target=self.model.generate, kwargs=generation_kwargs)
-        thread.start()
-        
-        # Determine input text length to strip it from outputs
-        prompt_text = self.tokenizer.decode(inputs.input_ids[0], skip_special_tokens=True)
-        prompt_length = len(prompt_text)
-        
-        # Stream the output
-        def generate_stream():
-            # Skip the prompt part in the first chunk
-            first_chunk = True
-            for text in streamer:
-                if first_chunk:
-                    # Only yield new tokens, not the original prompt
-                    if len(text) > prompt_length:
-                        yield {"generated_text": text[prompt_length:]}
-                    first_chunk = False
-                else:
-                    yield {"generated_text": text}
-        
-        return generate_stream()
+        try:
+            # Create a streamer object
+            streamer = TextIteratorStreamer(self.tokenizer, skip_special_tokens=True)
+            
+            # Set up generation in a separate thread
+            generation_kwargs = {
+                "input_ids": inputs.input_ids,
+                "attention_mask": inputs.attention_mask if hasattr(inputs, 'attention_mask') else None,
+                "streamer": streamer,
+                "max_new_tokens": max_new_tokens,
+                "temperature": temperature,
+                "top_p": top_p,
+                "do_sample": do_sample,
+                "pad_token_id": self.tokenizer.eos_token_id
+            }
+            
+            thread = Thread(target=self.model.generate, kwargs=generation_kwargs)
+            thread.start()
+            
+            # Determine input text length to strip it from outputs
+            input_text = self.tokenizer.decode(inputs.input_ids[0], skip_special_tokens=True)
+            
+            # Stream the output
+            def generate_stream():
+                # Skip the prompt part in the first chunk
+                full_text = ""
+                for text in streamer:
+                    full_text += text
+                    # Only return the part after the prompt
+                    if full_text.startswith(input_text):
+                        current_response = full_text[len(input_text):]
+                    else:
+                        current_response = full_text
+                    yield {"generated_text": current_response}
+            
+            return generate_stream()
+            
+        except Exception as e:
+            logger.error(f"Error in _generate_stream: {e}")
+            def error_stream():
+                yield {"error": str(e)}
+            return error_stream()
 
 # For local testing
 if __name__ == "__main__":