Add custom handler for inference endpoints

handler.py

@@ -0,0 +1,176 @@
+"""
+Custom Handler for MORBID-Actuarial v0.1.0 Conversational Model
+Hugging Face Inference Endpoints
+"""
+
+from typing import Dict, List, Any
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+class EndpointHandler:
+    def __init__(self, path: str = ""):
+        """
+        Initialize the handler with model and tokenizer
+        
+        Args:
+            path: Path to the model directory
+        """
+        # Load tokenizer and model
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            path,
+            torch_dtype=torch.float16,
+            device_map="auto",
+            low_cpu_mem_usage=True
+        )
+        
+        # Set padding token if not already set
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        
+        # System prompt for conversational behavior
+        self.system_prompt = """You are MORBID.AI, a friendly and conversational actuarial assistant. 
+You have expertise in:
+- Life expectancy and mortality statistics
+- Insurance and risk calculations
+- Financial mathematics (FM exam - 100% accuracy)
+- Probability theory (P exam - 100% accuracy)
+- Investment and financial markets (IFM exam - 93.3% accuracy)
+
+Be warm, helpful, and engaging. Respond naturally to greetings and casual conversation while maintaining your actuarial expertise.
+When users greet you, respond warmly. When they ask for help, be supportive and clear.
+Balance personality with precision when discussing technical topics."""
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        Process the inference request
+        
+        Args:
+            data: Dictionary containing the input data
+                - inputs (str or list): The input text(s)
+                - parameters (dict): Generation parameters
+        
+        Returns:
+            List of generated responses
+        """
+        # Extract inputs
+        inputs = data.get("inputs", "")
+        parameters = data.get("parameters", {})
+        
+        # Handle both string and list inputs
+        if isinstance(inputs, str):
+            inputs = [inputs]
+        elif not isinstance(inputs, list):
+            inputs = [str(inputs)]
+        
+        # Set default generation parameters
+        generation_params = {
+            "max_new_tokens": parameters.get("max_new_tokens", 200),
+            "temperature": parameters.get("temperature", 0.8),
+            "top_p": parameters.get("top_p", 0.95),
+            "do_sample": parameters.get("do_sample", True),
+            "repetition_penalty": parameters.get("repetition_penalty", 1.1),
+            "pad_token_id": self.tokenizer.pad_token_id,
+            "eos_token_id": self.tokenizer.eos_token_id,
+        }
+        
+        # Process each input
+        results = []
+        for input_text in inputs:
+            # Format the prompt with conversational context
+            prompt = self._format_prompt(input_text)
+            
+            # Tokenize
+            inputs_tokenized = self.tokenizer(
+                prompt,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=512
+            ).to(self.model.device)
+            
+            # Generate response
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    **inputs_tokenized,
+                    **generation_params
+                )
+            
+            # Decode the response
+            generated_text = self.tokenizer.decode(
+                outputs[0],
+                skip_special_tokens=True
+            )
+            
+            # Extract only the assistant's response
+            response = self._extract_response(generated_text, prompt)
+            
+            results.append({
+                "generated_text": response,
+                "conversation": {
+                    "user": input_text,
+                    "assistant": response
+                }
+            })
+        
+        return results
+    
+    def _format_prompt(self, user_input: str) -> str:
+        """
+        Format the user input into a conversational prompt
+        
+        Args:
+            user_input: The user's message
+            
+        Returns:
+            Formatted prompt string
+        """
+        # Check if it's a greeting or casual message
+        lower_input = user_input.lower().strip()
+        
+        # For very short inputs or greetings, add conversational context
+        if len(lower_input) <= 20 or any(greet in lower_input for greet in ["hi", "hello", "hey", "howdy"]):
+            return f"{self.system_prompt}\n\nHuman: {user_input}\nAssistant: "
+        
+        # For longer inputs, check if they're actuarial
+        actuarial_keywords = ["mortality", "life expectancy", "insurance", "premium", "annuity", 
+                             "probability", "risk", "actuarial", "death", "survival"]
+        
+        if any(keyword in lower_input for keyword in actuarial_keywords):
+            # Actuarial query - be precise but friendly
+            return f"As a conversational actuarial AI assistant, provide a helpful and accurate response.\n\nHuman: {user_input}\nAssistant: "
+        else:
+            # General conversation - be more casual
+            return f"{self.system_prompt}\n\nHuman: {user_input}\nAssistant: "
+    
+    def _extract_response(self, generated_text: str, prompt: str) -> str:
+        """
+        Extract only the assistant's response from the generated text
+        
+        Args:
+            generated_text: Full generated text including prompt
+            prompt: The original prompt
+            
+        Returns:
+            Just the assistant's response
+        """
+        # Remove the prompt from the beginning
+        if generated_text.startswith(prompt):
+            response = generated_text[len(prompt):].strip()
+        else:
+            # Try to find "Assistant:" marker
+            if "Assistant:" in generated_text:
+                response = generated_text.split("Assistant:")[-1].strip()
+            else:
+                response = generated_text.strip()
+        
+        # Clean up any remaining markers
+        if response.startswith(":"):
+            response = response[1:].strip()
+        
+        # Ensure we have a response
+        if not response:
+            response = "I'm here to help! Could you please rephrase your question?"
+        
+        return response