update handler.pu

handler.py

@@ -1,6 +1,8 @@
 import os
 import torch
 import logging
+import time
+import traceback
 from typing import Dict, List, Any, Union, Generator
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from threading import Thread
@@ -21,7 +23,7 @@ class EndpointHandler:
             path (str): Path to the model directory
         """
         # Set default parameters for inference
-        self.max_new_tokens = 2048  # Increased from 1024 to handle longer outputs
+        self.max_new_tokens = 1024  # Keep at 1024 to avoid timeouts
         self.temperature = 0.7
         self.top_p = 0.9
         self.do_sample = True
@@ -73,7 +75,8 @@ class EndpointHandler:
                 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}")
+            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):
@@ -103,7 +106,7 @@ class EndpointHandler:
         logger.info(f"Formatted prompt with {'system message and ' if system_message else ''}user message")
         return prompt
 
-    def __call__(self, data: Dict[str, Any]) -> Union[Dict[str, Any], Generator]:
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         """
         Process the input data and generate a response using the Phi-4 model.
         
@@ -111,33 +114,32 @@ class EndpointHandler:
             data (Dict[str, Any]): Input data containing the prompt and generation parameters
             
         Returns:
-            Dict[str, Any] or Generator: Model response
+            Dict[str, Any]: Model response
         """
+        start_time = time.time()
+        logger.info(f"Starting request processing: {data}")
+        
         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"}
+                return {"generated_text": "Error: Missing 'inputs' field in request"}
+            
+            # 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)
-                messages = [
-                    {"role": "system", "content": system_message} if system_message else None,
-                    {"role": "user", "content": user_message}
-                ]
-                messages = [m for m in messages if m is not None]  # Remove None values
                 
             # 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
-                system_message = None
-                user_message = ""
-                
                 for msg in messages:
                     if msg.get("role") == "system":
                         system_message = msg.get("content", "")
@@ -149,17 +151,17 @@ class EndpointHandler:
                 messages = data["inputs"]
                 
                 # Extract system and user messages for prompt formatting
-                system_message = None
-                user_message = ""
-                
                 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("Unsupported input format")
-                return {"error": "Unsupported input format. Expected string or messages object."}
+                logger.warning(f"Unsupported input format: {type(data['inputs'])}")
+                return {"generated_text": "Error: Unsupported input format. Expected string or messages object."}
+            
+            logger.info(f"Extracted user message: '{user_message}'")
+            logger.info(f"Extracted system message: '{system_message}'")
             
             # Format the prompt with system and user messages
             prompt = self.format_prompt_with_system(user_message, system_message)
@@ -169,72 +171,70 @@ class EndpointHandler:
             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), 4096)  # Increased to 4096
+            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)
             
-            # CRITICAL FIX: Use manual generation approach for Phi models with vocabulary mismatches
-            # This bypasses the token indexing issues
-            if stream:
-                return {"error": "Streaming temporarily disabled while fixing token indexing issues"}
+            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
-            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)
-            
-            # Format response in OpenAI-style format
-            if isinstance(response_text, dict) and "error" in response_text:
-                return response_text
-            else:
-                # OpenAI-style response format
-                openai_response = {
-                    "id": f"phi4-{int(torch.randint(10000, 99999, (1,)).item())}",
-                    "object": "chat.completion",
-                    "created": int(torch.cuda.current_stream().cuda_stream if torch.cuda.is_available() else 0),
-                    "model": "phi-4-mini",
-                    "choices": [
-                        {
-                            "index": 0,
-                            "message": {
-                                "role": "assistant",
-                                "content": response_text
-                            },
-                            "finish_reason": "stop"
-                        }
-                    ],
-                    "usage": {
-                        "prompt_tokens": len(input_ids[0]),
-                        "completion_tokens": len(self.tokenizer.encode(response_text)) if isinstance(response_text, str) else 0,
-                        "total_tokens": len(input_ids[0]) + (len(self.tokenizer.encode(response_text)) if isinstance(response_text, str) else 0)
-                    }
-                }
-                return openai_response
+            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'}")
+                
+                # Return response - keeping it simple for debug purposes
+                return {"generated_text": response_text}
+                
+            except RuntimeError as e:
+                logger.error(f"Runtime Error during generation: {str(e)}")
+                logger.error(traceback.format_exc())
+                return {"generated_text": f"Error during generation: {str(e)}"}
                 
         except Exception as e:
-            logger.error(f"Error during generation: {e}")
-            return {"error": str(e)}
+            logger.error(f"Unexpected error during request processing: {str(e)}")
+            logger.error(traceback.format_exc())
+            return {"generated_text": 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
-                max_steps = max_new_tokens  # Allow for full generation length
+                max_steps = min(max_new_tokens, 250)  # Limit to 250 tokens for reliability
                 current_ids = input_ids.clone()
                 
-                for _ in range(max_steps):
+                logger.info(f"Generating up to {max_steps} tokens")
+                
+                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,
@@ -275,111 +275,29 @@ class EndpointHandler:
                     
                     # 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")
                 else:
                     # Fallback if the expected format is not found
                     logger.warning("Could not find assistant tag in generated text")
                     response_text = generated_text
                     
-                logger.info(f"Generated {len(response_text)} characters")
                 return response_text
             
         except Exception as e:
             logger.error(f"Error in _safe_generate: {str(e)}")
-            return {"error": f"Generation error: {str(e)}. Please try a simpler input."}
-    
-    def _generate(self, inputs, max_new_tokens, temperature, top_p, do_sample):
-        """Generate text non-streaming mode"""
-        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}")
-                
-                # Fix: inputs is a dictionary, not an object with attributes
-                outputs = self.model.generate(
-                    inputs["input_ids"],
-                    attention_mask=inputs.get("attention_mask", 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 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"""
-        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.get("attention_mask", 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()
+            logger.error(traceback.format_exc())
+            return f"Generation error: {str(e)}. Please try a simpler input."
 
 # For local testing
 if __name__ == "__main__":