Update handler.py

handler.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Any, Optional
+from typing import Dict, List, Any
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import logging
@@ -11,113 +11,123 @@ class EndpointHandler:
     def __init__(self, path: str = ""):
         """
         Initialize the handler for Qwen2.5-Coder-7B-Instruct-Omni1.1
-        Optimized for Isaac Sim robotics code generation
+        Simple and robust implementation
         """
         logger.info(f"Loading model from {path}")
         
-        # Load tokenizer with proper configuration
-        self.tokenizer = AutoTokenizer.from_pretrained(
-            path,
-            trust_remote_code=True,
-            use_fast=False  # Use slow tokenizer to avoid tokenizer.json issues
-        )
-        
-        # Set pad token if not present
-        if self.tokenizer.pad_token is None:
-            self.tokenizer.pad_token = self.tokenizer.eos_token
-        
-        # Load model with optimizations for inference
-        self.model = AutoModelForCausalLM.from_pretrained(
-            path,
-            torch_dtype=torch.float16,
-            device_map="auto",
-            trust_remote_code=True,
-            low_cpu_mem_usage=True,
-            attn_implementation="flash_attention_2" if torch.cuda.is_available() else "eager"
-        )
-        
-        # Set model to evaluation mode
-        self.model.eval()
-        
-        logger.info("Model loaded successfully")
+        try:
+            # Load tokenizer - most robust approach
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                path,
+                trust_remote_code=True,
+                use_fast=False,
+                padding_side="left"
+            )
+            
+            # Ensure we have proper tokens
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+                
+            if self.tokenizer.chat_template is None:
+                # Set a basic chat template for Qwen
+                self.tokenizer.chat_template = "{% for message in messages %}{% if message['role'] == 'user' %}<|im_start|>user\n{{ message['content'] }}<|im_end|>\n{% elif message['role'] == 'assistant' %}<|im_start|>assistant\n{{ message['content'] }}<|im_end|>\n{% endif %}{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}"
+            
+            # Load model
+            self.model = AutoModelForCausalLM.from_pretrained(
+                path,
+                torch_dtype=torch.float16,
+                device_map="auto",
+                trust_remote_code=True,
+                low_cpu_mem_usage=True,
+                use_cache=True
+            )
+            
+            self.model.eval()
+            logger.info("Model and tokenizer loaded successfully")
+            
+        except Exception as e:
+            logger.error(f"Error loading model: {str(e)}")
+            raise e
     
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
         Handle inference requests
-        
-        Expected input format:
-        {
-            "inputs": "Your prompt here",
-            "parameters": {
-                "max_new_tokens": 512,
-                "temperature": 0.7,
-                "top_p": 0.9,
-                "do_sample": true,
-                "repetition_penalty": 1.1
-            }
-        }
         """
         try:
             # Extract inputs and parameters
             inputs = data.get("inputs", "")
             parameters = data.get("parameters", {})
             
-            # Default generation parameters optimized for code generation
+            if not inputs:
+                return [{"error": "No input provided", "generated_text": ""}]
+            
+            # Generation parameters
             max_new_tokens = parameters.get("max_new_tokens", 512)
             temperature = parameters.get("temperature", 0.7)
             top_p = parameters.get("top_p", 0.9)
             do_sample = parameters.get("do_sample", True)
             repetition_penalty = parameters.get("repetition_penalty", 1.1)
             
-            # Format input with proper chat template for Qwen2.5
-            if not inputs.startswith("<|im_start|>"):
-                formatted_input = f"<|im_start|>user\n{inputs}<|im_end|>\n<|im_start|>assistant"
-            else:
+            # Prepare input - handle both raw text and pre-formatted chat
+            if inputs.startswith("<|im_start|>"):
+                # Already formatted
                 formatted_input = inputs
+            else:
+                # Format as chat
+                formatted_input = f"<|im_start|>user\n{inputs}<|im_end|>\n<|im_start|>assistant\n"
             
-            # Tokenize input
+            # Tokenize
             input_ids = self.tokenizer.encode(
                 formatted_input,
                 return_tensors="pt",
+                add_special_tokens=False,
                 truncation=True,
-                max_length=2048  # Leave room for generation
-            ).to(self.model.device)
+                max_length=4096 - max_new_tokens  # Leave room for generation
+            )
+            
+            if input_ids.size(1) == 0:
+                return [{"error": "Input tokenization failed", "generated_text": ""}]
+            
+            input_ids = input_ids.to(self.model.device)
             
-            # Generate response
+            # Generate
             with torch.no_grad():
-                output_ids = self.model.generate(
+                outputs = self.model.generate(
                     input_ids,
                     max_new_tokens=max_new_tokens,
                     temperature=temperature,
                     top_p=top_p,
                     do_sample=do_sample,
                     repetition_penalty=repetition_penalty,
-                    pad_token_id=self.tokenizer.pad_token_id,
+                    pad_token_id=self.tokenizer.eos_token_id,
                     eos_token_id=self.tokenizer.eos_token_id,
-                    use_cache=True
+                    use_cache=True,
+                    num_return_sequences=1
                 )
             
-            # Decode only the new tokens (response)
-            response_ids = output_ids[0][input_ids.shape[1]:]
-            response_text = self.tokenizer.decode(
-                response_ids, 
+            # Decode response (only new tokens)
+            generated_ids = outputs[0][input_ids.size(1):]
+            response = self.tokenizer.decode(
+                generated_ids,
                 skip_special_tokens=True,
                 clean_up_tokenization_spaces=True
             )
             
             # Clean up response
-            response_text = response_text.strip()
+            response = response.strip()
+            
+            # Remove any remaining special tokens manually
+            response = response.replace("<|im_end|>", "").strip()
             
-            # Return in expected format
             return [{
-                "generated_text": response_text,
-                "generated_tokens": len(response_ids),
-                "finish_reason": "stop" if self.tokenizer.eos_token_id in response_ids else "length"
+                "generated_text": response,
+                "generated_tokens": len(generated_ids),
+                "finish_reason": "eos_token" if self.tokenizer.eos_token_id in generated_ids else "length"
             }]
             
         except Exception as e:
-            logger.error(f"Error during inference: {str(e)}")
+            logger.error(f"Generation error: {str(e)}")
             return [{
-                "error": f"Inference failed: {str(e)}",
+                "error": f"Generation failed: {str(e)}",
                 "generated_text": ""
-            }]
+            }]