Update handler.py

handler.py

@@ -1,22 +1,127 @@
 from typing import Dict, List, Any
-from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import logging
 
-class EndpointHandler():
-    def __init__(self, path=""):
-        # Load the model and tokenizer from the specified path
-        self.model = AutoModelForCausalLM.from_pretrained(path)
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class EndpointHandler:
+    def __init__(self, path: str = ""):
+        """
+        Initialize the model and tokenizer when the endpoint starts.
+        
+        Args:
+            path (str): Path to the model files
+        """
+        logger.info(f"Loading model from {path}")
+        
+        # Load tokenizer and model
         self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            path,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            device_map="auto" if torch.cuda.is_available() else None,
+            trust_remote_code=True
+        )
+        
+        # Set pad token if it doesn't exist
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+            
+        logger.info("Model loaded successfully")
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-        # Extract input text from the request
-        inputs = data.pop("inputs", data)
-
-        # Tokenize input and generate text
-        input_ids = self.tokenizer.encode(inputs, return_tensors="pt")
-        output_ids = self.model.generate(input_ids)
-
-        # Decode the generated output
-        output_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        """
+        Process the inference request.
+        
+        Args:
+            data (Dict[str, Any]): Request data containing:
+                - inputs (str): The input text/prompt
+                - parameters (dict, optional): Generation parameters
+                    - max_new_tokens (int): Maximum tokens to generate (default: 256)
+                    - temperature (float): Sampling temperature (default: 0.7)
+                    - top_p (float): Top-p sampling (default: 0.9)
+                    - do_sample (bool): Whether to use sampling (default: True)
+                    - repetition_penalty (float): Repetition penalty (default: 1.1)
+                    - return_full_text (bool): Return full text including input (default: False)
+                
+        Returns:
+            List[Dict[str, Any]]: Generated text response
+        """
+        try:
+            # Extract inputs
+            inputs = data.get("inputs", "")
+            if not inputs:
+                return [{"error": "No input text provided"}]
+            
+            # Extract generation parameters
+            parameters = data.get("parameters", {})
+            max_new_tokens = parameters.get("max_new_tokens", 256)
+            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)
+            return_full_text = parameters.get("return_full_text", False)
+            
+            # Format the input as a chat message if it doesn't already contain instruction formatting
+            if not any(marker in inputs.lower() for marker in ["[inst]", "<s>", "### instruction", "user:", "assistant:"]):
+                formatted_input = f"[INST] {inputs} [/INST]"
+            else:
+                formatted_input = inputs
+            
+            # Tokenize input
+            input_ids = self.tokenizer.encode(
+                formatted_input, 
+                return_tensors="pt",
+                truncation=True,
+                max_length=2048  # Reasonable limit for input
+            )
+            
+            # Move to GPU if available
+            if torch.cuda.is_available():
+                input_ids = input_ids.cuda()
+            
+            # Generate response
+            with torch.no_grad():
+                output_ids = 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,
+                    eos_token_id=self.tokenizer.eos_token_id,
+                    use_cache=True
+                )
+            
+            # Decode the response
+            if return_full_text:
+                generated_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
+            else:
+                # Only return the newly generated tokens
+                new_tokens = output_ids[0][input_ids.shape[-1]:]
+                generated_text = self.tokenizer.decode(new_tokens, skip_special_tokens=True)
+            
+            # Clean up the response
+            generated_text = generated_text.strip()
+            
+            # Return in the expected format
+            return [{
+                "generated_text": generated_text,
+                "input_length": input_ids.shape[-1],
+                "output_length": len(output_ids[0]) - input_ids.shape[-1]
+            }]
+            
+        except Exception as e:
+            logger.error(f"Error during inference: {str(e)}")
+            return [{"error": f"Inference failed: {str(e)}"}]
 
-        # Return the generated text
-        return [{"generated_text": output_text}]
+    def __del__(self):
+        """Clean up resources when the handler is destroyed."""
+        if hasattr(self, 'model'):
+            del self.model
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()