Update handler.py

handler.py

@@ -194,15 +194,18 @@ def _to_block_list(content: Any) -> List[Dict[str, str]]:
     if content is None:
         return [{"type": "text", "text": ""}]
     if isinstance(content, list):
-        # Assume already blocks
-        return content
+        # Check if it's already in block format
+        if all(isinstance(item, dict) and "type" in item for item in content):
+            return content
+        # Otherwise treat as raw list, stringify
+        return [{"type": "text", "text": str(content)}]
     # Fallback: stringify
     return [{"type": "text", "text": str(content)}]
 
 
 class EndpointHandler:
     """
-    HF Inference Endpoint handler that mirrors OpenAI API format:
+    HF Inference Endpoint handler that accepts pure OpenAI API format:
     - Accepts OpenAI-style requests with 'messages', 'model', etc.
     - Returns OpenAI-style responses with 'choices', 'usage', etc.
     - Loads Unsloth 4-bit Gemma-3-4B IT
@@ -252,11 +255,10 @@ class EndpointHandler:
             role = m.get("role", "user")
             content = m.get("content", "")
             
-            # Convert content to block list format
+            # Convert content to block list format for internal processing
             if isinstance(content, str):
                 content = _to_block_list(content)
             elif isinstance(content, list):
-                # Already in block format or needs conversion
                 content = _to_block_list(content)
             else:
                 content = _to_block_list(str(content))
@@ -286,12 +288,14 @@ class EndpointHandler:
         
         Input format (OpenAI API):
         {
-            "messages": [{"role": "user", "content": "Hello"}],
-            "model": "gpt-3.5-turbo",  # optional
+            "model": "Ephraimmm/PIDGIN_gemma-3",
+            "messages": [
+                {"role": "user", "content": "What is deep learning?"}
+            ],
             "temperature": 0.7,         # optional
             "max_tokens": 256,          # optional
             "top_p": 1.0,              # optional
-            ...
+            "stream": false            # optional
         }
         
         Output format (OpenAI API):
@@ -299,7 +303,7 @@ class EndpointHandler:
             "id": "chatcmpl-...",
             "object": "chat.completion",
             "created": 1234567890,
-            "model": "gemma-3-4b-pidgin",
+            "model": "Ephraimmm/PIDGIN_gemma-3",
             "choices": [{
                 "index": 0,
                 "message": {
@@ -315,72 +319,48 @@ class EndpointHandler:
             }
         }
         """
-        logger.info(f"Incoming keys: {list(data.keys())}")
+        logger.info(f"Incoming request: {data}")
         
-        # Support both OpenAI format ('messages') and legacy format ('inputs')
-        if "messages" in data:
-            messages = data["messages"]
-            if not isinstance(messages, list):
-                raise ValueError("`messages` must be a list of message objects.")
-        elif "inputs" in data:
-            # Legacy support: convert inputs to messages format
-            prompt = data["inputs"]
-            if isinstance(prompt, str):
-                messages = [{"role": "user", "content": prompt}]
-            elif isinstance(prompt, dict):
-                messages = [prompt]
-            elif isinstance(prompt, list):
-                messages = prompt
-            else:
-                raise ValueError("Invalid `inputs` format.")
-        else:
-            raise ValueError("Missing `messages` or `inputs` in request body.")
-
-        # Extract parameters (support both OpenAI and custom parameter names)
-        params: Dict[str, Any] = data.get("parameters", {})
+        # Validate required fields
+        if "messages" not in data:
+            raise ValueError("Missing required field: 'messages'")
         
-        # Map OpenAI parameters to model parameters
-        max_new_tokens = int(
-            params.get("max_new_tokens") or 
-            data.get("max_tokens") or 
-            params.get("max_tokens") or 
-            256
-        )
-        temperature = float(
-            params.get("temperature") or 
-            data.get("temperature") or 
-            0.1
-        )
-        top_p = float(
-            params.get("top_p") or 
-            data.get("top_p") or 
-            1.0
-        )
-        top_k: Optional[int] = (
-            params.get("top_k") or 
-            data.get("top_k")
-        )
-        use_cache = bool(params.get("use_cache", False))
-        repetition_penalty = float(params.get("repetition_penalty", 1.0))
-        eos_token_id = params.get("eos_token_id", self.tokenizer.eos_token_id)
-        pad_token_id = params.get("pad_token_id", self.tokenizer.pad_token_id)
-
+        messages = data["messages"]
+        if not isinstance(messages, list):
+            raise ValueError("'messages' must be a list of message objects")
+        
+        # Extract model name (optional)
+        model_name = data.get("model", "Ephraimmm/PIDGIN_gemma-3")
+        
+        # Extract OpenAI-style parameters with defaults
+        max_tokens = int(data.get("max_tokens", 256))
+        temperature = float(data.get("temperature", 0.1))
+        top_p = float(data.get("top_p", 1.0))
+        top_k = data.get("top_k", None)
+        stream = data.get("stream", False)  # Not implemented yet
+        presence_penalty = data.get("presence_penalty", 0.0)
+        frequency_penalty = data.get("frequency_penalty", 0.0)
+        
+        # Build model inputs
         inputs = self._build_model_inputs(messages)
 
+        # Prepare generation parameters
         gen_kwargs = dict(
             **inputs,
-            max_new_tokens=max_new_tokens,
+            max_new_tokens=max_tokens,
             temperature=temperature,
             top_p=top_p,
-            do_sample=True,
-            use_cache=use_cache,
-            repetition_penalty=repetition_penalty,
-            eos_token_id=eos_token_id,
-            pad_token_id=pad_token_id,
+            do_sample=temperature > 0,  # Use sampling if temperature > 0
+            use_cache=True,
+            repetition_penalty=1.0 + frequency_penalty,  # Map frequency_penalty
+            eos_token_id=self.tokenizer.eos_token_id,
+            pad_token_id=self.tokenizer.pad_token_id,
         )
+        
         if top_k is not None:
             gen_kwargs["top_k"] = int(top_k)
 
+        # Generate response
         with torch.no_grad():
             output_ids = self.model.generate(**gen_kwargs)
 
@@ -389,17 +369,17 @@ class EndpointHandler:
         generated_ids = output_ids[0, input_len:]
         reply_text = self.tokenizer.decode(generated_ids, skip_special_tokens=True)
 
-        # Calculate tokens
+        # Calculate token usage
         prompt_tokens = int(input_len)
         completion_tokens = int(generated_ids.shape[-1])
         total_tokens = prompt_tokens + completion_tokens
 
-        # Return OpenAI-style response
+        # Return OpenAI-compatible response
         response = {
-            "id": f"chatcmpl-{int(time.time())}",
+            "id": f"chatcmpl-{int(time.time() * 1000)}",
             "object": "chat.completion",
             "created": int(time.time()),
-            "model": data.get("model", "gemma-3-4b-pidgin"),
+            "model": model_name,
             "choices": [
                 {
                     "index": 0,
@@ -423,31 +403,40 @@ class EndpointHandler:
 if __name__ == "__main__":
     h = EndpointHandler(model_dir=".")
     
-    # Test with OpenAI-style format
-    print("=== OpenAI Format Test ===")
+    # Test 1: Simple user message (OpenAI format)
+    print("=== Test 1: Simple User Message ===")
     response = h({
+        "model": "Ephraimmm/PIDGIN_gemma-3",
         "messages": [
-            {"role": "user", "content": "How far?"}
-        ],
-        "temperature": 0.7,
-        "max_tokens": 64
+            {"role": "user", "content": "What is deep learning?"}
+        ]
     })
     print(response)
+    print()
     
-    print("\n=== OpenAI Format with System Message ===")
+    # Test 2: With temperature and max_tokens
+    print("=== Test 2: With Parameters ===")
     response = h({
+        "model": "Ephraimmm/PIDGIN_gemma-3",
         "messages": [
-            {"role": "system", "content": "You are a helpful assistant."},
-            {"role": "user", "content": "Explain why sky dey blue."}
+            {"role": "user", "content": "How far?"}
         ],
+        "temperature": 0.7,
         "max_tokens": 64
     })
     print(response)
+    print()
     
-    # Legacy format still works
-    print("\n=== Legacy Format Test ===")
+    # Test 3: Multi-turn conversation
+    print("=== Test 3: Multi-turn Conversation ===")
     response = h({
-        "inputs": "How far?",
-        "parameters": {"max_new_tokens": 64}
+        "model": "Ephraimmm/PIDGIN_gemma-3",
+        "messages": [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": "Explain why the sky is blue"},
+            {"role": "assistant", "content": "The sky dey blue because..."},
+            {"role": "user", "content": "Tell me more"}
+        ],
+        "max_tokens": 100
     })
     print(response)