add handler

handler.py

@@ -0,0 +1,156 @@
+import os
+import torch
+from typing import Dict, List, Any
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
+
+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 = 4096
+        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
+        
+        # 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
+            
+        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
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Process the input data and generate a response using the Phi-4 model.
+        
+        Args:
+            data (Dict[str, Any]): Input data containing the prompt and generation parameters
+            
+        Returns:
+            Dict[str, Any]: Model response
+        """
+        # 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)
+    
+    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}
+    
+    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()
+
+# For local testing
+if __name__ == "__main__":
+    # Example usage
+    handler = EndpointHandler()
+    result = handler({"inputs": "What are the major features of Phi-4?"})
+    print(result)