Create handler.py

handler.py

@@ -0,0 +1,187 @@
+from typing import Dict, List, Any
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+class EndpointHandler:
+    """
+    Custom handler for HuggingFace Inference Endpoints
+    Handles Nigerian Pidgin English text generation
+    """
+    
+    def __init__(self, path: str = ""):
+        """
+        Initialize the model and tokenizer
+        
+        Args:
+            path: Path to the model directory (provided by HuggingFace)
+        """
+        # Load model
+        self.model = AutoModelForCausalLM.from_pretrained(
+            path,
+            torch_dtype=torch.float16,
+            device_map="auto",
+            trust_remote_code=True
+        )
+        
+        # Load tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            path,
+            trust_remote_code=True
+        )
+        
+        # Set default system prompt
+        self.default_system_prompt = (
+            "You are a helpful assistant that speaks Nigerian Pidgin English. "
+            "Respond naturally in Pidgin."
+        )
+        
+        print("✓ Model and tokenizer loaded successfully")
+    
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        Handle inference requests
+        
+        Args:
+            data: Dictionary containing:
+                - inputs (str): The user's prompt/question
+                - parameters (dict, optional): Generation parameters
+                    - system_prompt (str): Custom system prompt
+                    - max_new_tokens (int): Max tokens to generate (default: 100)
+                    - temperature (float): Sampling temperature (default: 0.7)
+                    - top_p (float): Nucleus sampling (default: 0.9)
+                    - top_k (int): Top-k sampling (default: 50)
+                    - repetition_penalty (float): Penalty for repetition (default: 1.1)
+                    - do_sample (bool): Whether to use sampling (default: True)
+        
+        Returns:
+            List of dictionaries containing generated text
+        """
+        # Extract inputs
+        inputs_text = data.pop("inputs", data)
+        parameters = data.pop("parameters", {})
+        
+        # Get generation parameters with defaults
+        system_prompt = parameters.get(
+            "system_prompt", 
+            self.default_system_prompt
+        )
+        max_new_tokens = parameters.get("max_new_tokens", 100)
+        temperature = parameters.get("temperature", 0.7)
+        top_p = parameters.get("top_p", 0.9)
+        top_k = parameters.get("top_k", 50)
+        repetition_penalty = parameters.get("repetition_penalty", 1.1)
+        do_sample = parameters.get("do_sample", True)
+        return_full_text = parameters.get("return_full_text", False)
+        
+        # Construct full prompt with system prompt
+        if system_prompt:
+            full_prompt = f"{system_prompt}\n\nUser: {inputs_text}\nAssistant:"
+        else:
+            full_prompt = inputs_text
+        
+        # Tokenize
+        inputs = self.tokenizer(
+            full_prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=2048
+        ).to(self.model.device)
+        
+        # Generate
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                top_k=top_k,
+                repetition_penalty=repetition_penalty,
+                do_sample=do_sample,
+                pad_token_id=self.tokenizer.pad_token_id,
+                eos_token_id=self.tokenizer.eos_token_id,
+            )
+        
+        # Decode
+        generated_text = self.tokenizer.decode(
+            outputs[0],
+            skip_special_tokens=True
+        )
+        
+        # Extract only the assistant's response if system prompt was used
+        if not return_full_text:
+            if "Assistant:" in generated_text:
+                generated_text = generated_text.split("Assistant:")[-1].strip()
+            else:
+                # Remove the input prompt from output
+                generated_text = generated_text.replace(full_prompt, "").strip()
+        
+        # Return in HuggingFace expected format
+        return [{"generated_text": generated_text}]
+
+
+# For local testing (optional)
+if __name__ == "__main__":
+    # Test the handler locally
+    print("Testing handler locally...")
+    
+    # Initialize handler (use "." for current directory in local testing)
+    handler = EndpointHandler(path="Ephraimmm/pidgin_finetuned_model")
+    
+    # Test cases
+    test_cases = [
+        {
+            "name": "Simple greeting",
+            "data": {
+                "inputs": "How you dey?",
+                "parameters": {
+                    "max_new_tokens": 50,
+                    "temperature": 0.7
+                }
+            }
+        },
+        {
+            "name": "Without system prompt",
+            "data": {
+                "inputs": "Wetin you wan chop?",
+                "parameters": {
+                    "system_prompt": "",
+                    "max_new_tokens": 50
+                }
+            }
+        },
+        {
+            "name": "Custom system prompt",
+            "data": {
+                "inputs": "Tell me about Lagos",
+                "parameters": {
+                    "system_prompt": "You are a knowledgeable guide who speaks Pidgin.",
+                    "max_new_tokens": 100,
+                    "temperature": 0.8
+                }
+            }
+        },
+        {
+            "name": "Return full text",
+            "data": {
+                "inputs": "Wetin be your name?",
+                "parameters": {
+                    "max_new_tokens": 50,
+                    "return_full_text": True
+                }
+            }
+        }
+    ]
+    
+    # Run tests
+    for test in test_cases:
+        print(f"\n{'='*60}")
+        print(f"Test: {test['name']}")
+        print(f"{'='*60}")
+        print(f"Input: {test['data']['inputs']}")
+        
+        result = handler(test['data'])
+        print(f"Output: {result[0]['generated_text']}")
+    
+    print(f"\n{'='*60}")
+    print("✓ All tests completed!")