Update handler.py

handler.py

@@ -1,28 +1,41 @@
 #tope_version
 # handler.py
+# handler.py
 from typing import Any, Dict, List, Union, Optional
- 
 import torch
 from huggingface_inference_toolkit.logging import logger
 from unsloth import FastLanguageModel
 from peft import PeftModel
  
-Prompt = Union[str, List[Dict[str, str]]]
+# OpenAI-style prompt:
+#   - str
+#   - [{"role": "...", "content": "..." | [{"type":"text","text":"..."}]}]
+Prompt = Union[str, List[Dict[str, Any]], Dict[str, Any]]
+ 
+ 
+def _to_block_list(content: Any) -> List[Dict[str, str]]:
+    """Ensure content is a list[{'type':'text','text':...}]."""
+    if content is None:
+        return [{"type": "text", "text": ""}]
+    if isinstance(content, list):
+        # Assume already blocks
+        return content
+    # Fallback: stringify
+    return [{"type": "text", "text": str(content)}]
  
  
 class EndpointHandler:
     """
-    Hugging Face Inference Endpoint handler.
-    - Loads Unsloth 4-bit base
+    HF Inference Endpoint handler that mirrors your Gradio setup:
+    - Loads Unsloth 4-bit Gemma-3-4B IT
     - Attaches LoRA adapter
-    - Uses tokenizer.apply_chat_template like your Gradio app
-    - Always injects the Pidgin-English system instruction
+    - Uses tokenizer.apply_chat_template
+    - Always replies in Nigerian Pidgin
     """
  
     def __init__(self, model_dir: str, **_: Any):
         logger.info(f"Initializing with model_dir={model_dir}")
  
-        # --- Base model (4-bit) ---
         base_model_name = "unsloth/gemma-3-4b-it-unsloth-bnb-4bit"
         model, tokenizer = FastLanguageModel.from_pretrained(
             model_name=base_model_name,
@@ -31,45 +44,62 @@ class EndpointHandler:
             load_in_4bit=True,
         )
  
-        # --- LoRA adapter ---
         lora_repo = "Ephraimmm/PIDGIN_gemma-3"
-        model = PeftModel.from_pretrained(model, lora_repo)  # default adapter
+        model = PeftModel.from_pretrained(model, lora_repo)
         FastLanguageModel.for_inference(model)
  
         self.model = model.eval()
         self.tokenizer = tokenizer
-        self.device = getattr(self.model, "device",
-                              torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+        self.device = getattr(
+            self.model,
+            "device",
+            torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+        )
+        # Safety: some tokenizers lack pad_token_id
+        if getattr(self.tokenizer, "pad_token_id", None) is None:
+            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+ 
         logger.info(f"Device: {self.device} | eos_id: {self.tokenizer.eos_token_id}")
  
-        # System prompt copied from your Gradio app
         self._system_text = (
-            "You are a Nigerian assistant that speaks PIDGIN ENGLISH."
-            " when asked how far reply I de o, how you de"
+            "You are a Nigerian assistant that speaks PIDGIN ENGLISH. "
+            "When asked 'how far', reply 'I dey o, how you dey?'. "
+            "Always answer in Pidgin English."
         )
  
-    # -----------------------------
-    # Input preparation
-    # -----------------------------
-    def _ensure_messages(self, prompt: Prompt) -> List[Dict[str, str]]:
-        """Normalize to a list of messages and inject the system instruction."""
-        sys_msg = {"role": "system", "content": self._system_text}
+    def _normalize_messages(self, prompt: Prompt) -> List[Dict[str, Any]]:
+        """Return a list of messages with block-style contents only."""
+        # Accept single dict
+        if isinstance(prompt, dict):
+            prompt = [prompt]
+ 
+        # Accept raw string
         if isinstance(prompt, str):
-            return [sys_msg, {"role": "user", "content": prompt}]
-        if isinstance(prompt, list):
-            has_system = any(m.get("role") == "system" for m in prompt)
-            return ([sys_msg] + prompt) if not has_system else prompt
-        raise ValueError("`inputs` must be a string or a list of {role, content} messages.")
+            msgs = [
+                {"role": "system", "content": _to_block_list(self._system_text)},
+                {"role": "user", "content": _to_block_list(prompt)},
+            ]
+            return msgs
+ 
+        if not isinstance(prompt, list):
+            raise ValueError("`inputs` must be a string, a message dict, or a list of messages.")
+ 
+        # Normalize all contents to blocks
+        norm = []
+        for m in prompt:
+            role = m.get("role", "user")
+            content = _to_block_list(m.get("content", ""))
+            norm.append({"role": role, "content": content})
+ 
+        # Ensure a system message exists (prepend if missing)
+        if not any(m.get("role") == "system" for m in norm):
+            norm = [{"role": "system", "content": _to_block_list(self._system_text)}] + norm
+ 
+        return norm
  
     def _build_model_inputs(self, prompt: Prompt):
-        """
-        Use tokenizer.apply_chat_template exactly like your Gradio code:
-          add_generation_prompt=True, return_tensors='pt', tokenize=True.
-        """
-        messages = self._ensure_messages(prompt)
- 
-        # Some tokenizers support content as either str or structured content.
-        # We pass OpenAI-style messages; Unsloth/Gemma-3 chat templates handle it.
+        messages = self._normalize_messages(prompt)
+        # Use the same path as your Gradio script
         model_inputs = self.tokenizer.apply_chat_template(
             messages,
             add_generation_prompt=True,
@@ -77,17 +107,9 @@ class EndpointHandler:
             tokenize=True,
             return_dict=True,
         ).to(self.device)
- 
         return model_inputs
  
-    # -----------------------------
-    # Generation
-    # -----------------------------
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        """
-        Expects: {"inputs": <str|messages>, "parameters": {...}}
-        Mirrors your Gradio defaults (temperature=0.1, top_p=1.0, top_k=None, use_cache=False).
-        """
         logger.info(f"Incoming keys: {list(data.keys())}")
         if "inputs" not in data:
             raise ValueError("Missing `inputs` in request body.")
@@ -95,17 +117,16 @@ class EndpointHandler:
         prompt: Prompt = data["inputs"]
         params: Dict[str, Any] = (data.get("parameters") or {})
  
-        # Map your defaults
-        max_new_tokens = int(params.get("max_new_tokens", 256))  # safer default than 9000
+        # Defaults aligned with your Gradio code
+        max_new_tokens = int(params.get("max_new_tokens", 256))   # raise per request if needed
         temperature = float(params.get("temperature", 0.1))
         top_p = float(params.get("top_p", 1.0))
         top_k: Optional[int] = params.get("top_k", None)
         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", getattr(self.tokenizer, "pad_token_id", None))
+        pad_token_id = params.get("pad_token_id", self.tokenizer.pad_token_id)
  
-        # Prepare inputs via chat template (like your Gradio code)
         inputs = self._build_model_inputs(prompt)
  
         gen_kwargs = dict(
@@ -113,20 +134,19 @@ class EndpointHandler:
             max_new_tokens=max_new_tokens,
             temperature=temperature,
             top_p=top_p,
-            do_sample=True,                 # sampling on (like your streamer setup)
+            do_sample=True,                 # streaming-like sampling on
             use_cache=use_cache,
             repetition_penalty=repetition_penalty,
             eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
         )
-        if pad_token_id is not None:
-            gen_kwargs["pad_token_id"] = pad_token_id
         if top_k is not None:
             gen_kwargs["top_k"] = int(top_k)
  
         with torch.no_grad():
             output_ids = self.model.generate(**gen_kwargs)
  
-        # Return only the continuation (like what streaming would show)
+        # Return ONLY the continuation (like your streamer)
         input_len = inputs["input_ids"].shape[-1]
         generated_ids = output_ids[0, input_len:]
         reply_text = self.tokenizer.decode(generated_ids, skip_special_tokens=True)
@@ -140,25 +160,180 @@ class EndpointHandler:
         }
  
  
-# Local quick test
 if __name__ == "__main__":
     h = EndpointHandler(model_dir=".")
- 
     # Raw string
-    out1 = h({
-        "inputs": "How far?",
-        "parameters": {"max_new_tokens": 64}
-    })
-    print(">> STRING:", out1["reply"][0])
- 
-    # Chat-style messages
-    out2 = h({
+    print(h({"inputs": "How far?", "parameters": {"max_new_tokens": 64}}))
+    # Block-style content (your Gradio shape)
+    print(h({
         "inputs": [
-            {"role": "user", "content": "Explain why sky dey blue."}
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": "Explain why sky dey blue."}]
+            }
         ],
         "parameters": {"max_new_tokens": 64}
-    })
-    print(">> CHAT:", out2["reply"][0])
+    }))
+
+# from typing import Any, Dict, List, Union, Optional
+ 
+# import torch
+# from huggingface_inference_toolkit.logging import logger
+# from unsloth import FastLanguageModel
+# from peft import PeftModel
+ 
+# Prompt = Union[str, List[Dict[str, str]]]
+ 
+ 
+# class EndpointHandler:
+#     """
+#     Hugging Face Inference Endpoint handler.
+#     - Loads Unsloth 4-bit base
+#     - Attaches LoRA adapter
+#     - Uses tokenizer.apply_chat_template like your Gradio app
+#     - Always injects the Pidgin-English system instruction
+#     """
+ 
+#     def __init__(self, model_dir: str, **_: Any):
+#         logger.info(f"Initializing with model_dir={model_dir}")
+ 
+#         # --- Base model (4-bit) ---
+#         base_model_name = "unsloth/gemma-3-4b-it-unsloth-bnb-4bit"
+#         model, tokenizer = FastLanguageModel.from_pretrained(
+#             model_name=base_model_name,
+#             max_seq_length=2048,
+#             dtype=torch.float16,
+#             load_in_4bit=True,
+#         )
+ 
+#         # --- LoRA adapter ---
+#         lora_repo = "Ephraimmm/PIDGIN_gemma-3"
+#         model = PeftModel.from_pretrained(model, lora_repo)  # default adapter
+#         FastLanguageModel.for_inference(model)
+ 
+#         self.model = model.eval()
+#         self.tokenizer = tokenizer
+#         self.device = getattr(self.model, "device",
+#                               torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+#         logger.info(f"Device: {self.device} | eos_id: {self.tokenizer.eos_token_id}")
+ 
+#         # System prompt copied from your Gradio app
+#         self._system_text = (
+#             "You are a Nigerian assistant that speaks PIDGIN ENGLISH."
+#             " when asked how far reply I de o, how you de"
+#         )
+ 
+#     # -----------------------------
+#     # Input preparation
+#     # -----------------------------
+#     def _ensure_messages(self, prompt: Prompt) -> List[Dict[str, str]]:
+#         """Normalize to a list of messages and inject the system instruction."""
+#         sys_msg = {"role": "system", "content": self._system_text}
+#         if isinstance(prompt, str):
+#             return [sys_msg, {"role": "user", "content": prompt}]
+#         if isinstance(prompt, list):
+#             has_system = any(m.get("role") == "system" for m in prompt)
+#             return ([sys_msg] + prompt) if not has_system else prompt
+#         raise ValueError("`inputs` must be a string or a list of {role, content} messages.")
+ 
+#     def _build_model_inputs(self, prompt: Prompt):
+#         """
+#         Use tokenizer.apply_chat_template exactly like your Gradio code:
+#           add_generation_prompt=True, return_tensors='pt', tokenize=True.
+#         """
+#         messages = self._ensure_messages(prompt)
+ 
+#         # Some tokenizers support content as either str or structured content.
+#         # We pass OpenAI-style messages; Unsloth/Gemma-3 chat templates handle it.
+#         model_inputs = self.tokenizer.apply_chat_template(
+#             messages,
+#             add_generation_prompt=True,
+#             return_tensors="pt",
+#             tokenize=True,
+#             return_dict=True,
+#         ).to(self.device)
+ 
+#         return model_inputs
+ 
+#     # -----------------------------
+#     # Generation
+#     # -----------------------------
+#     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+#         """
+#         Expects: {"inputs": <str|messages>, "parameters": {...}}
+#         Mirrors your Gradio defaults (temperature=0.1, top_p=1.0, top_k=None, use_cache=False).
+#         """
+#         logger.info(f"Incoming keys: {list(data.keys())}")
+#         if "inputs" not in data:
+#             raise ValueError("Missing `inputs` in request body.")
+ 
+#         prompt: Prompt = data["inputs"]
+#         params: Dict[str, Any] = (data.get("parameters") or {})
+ 
+#         # Map your defaults
+#         max_new_tokens = int(params.get("max_new_tokens", 256))  # safer default than 9000
+#         temperature = float(params.get("temperature", 0.1))
+#         top_p = float(params.get("top_p", 1.0))
+#         top_k: Optional[int] = params.get("top_k", None)
+#         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", getattr(self.tokenizer, "pad_token_id", None))
+ 
+#         # Prepare inputs via chat template (like your Gradio code)
+#         inputs = self._build_model_inputs(prompt)
+ 
+#         gen_kwargs = dict(
+#             **inputs,
+#             max_new_tokens=max_new_tokens,
+#             temperature=temperature,
+#             top_p=top_p,
+#             do_sample=True,                 # sampling on (like your streamer setup)
+#             use_cache=use_cache,
+#             repetition_penalty=repetition_penalty,
+#             eos_token_id=eos_token_id,
+#         )
+#         if pad_token_id is not None:
+#             gen_kwargs["pad_token_id"] = pad_token_id
+#         if top_k is not None:
+#             gen_kwargs["top_k"] = int(top_k)
+ 
+#         with torch.no_grad():
+#             output_ids = self.model.generate(**gen_kwargs)
+ 
+#         # Return only the continuation (like what streaming would show)
+#         input_len = inputs["input_ids"].shape[-1]
+#         generated_ids = output_ids[0, input_len:]
+#         reply_text = self.tokenizer.decode(generated_ids, skip_special_tokens=True)
+ 
+#         return {
+#             "reply": [reply_text],
+#             "usage": {
+#                 "prompt_tokens": int(input_len),
+#                 "generated_tokens": int(generated_ids.shape[-1]),
+#             },
+#         }
+ 
+ 
+# # Local quick test
+# if __name__ == "__main__":
+#     h = EndpointHandler(model_dir=".")
+ 
+#     # Raw string
+#     out1 = h({
+#         "inputs": "How far?",
+#         "parameters": {"max_new_tokens": 64}
+#     })
+#     print(">> STRING:", out1["reply"][0])
+ 
+#     # Chat-style messages
+#     out2 = h({
+#         "inputs": [
+#             {"role": "user", "content": "Explain why sky dey blue."}
+#         ],
+#         "parameters": {"max_new_tokens": 64}
+#     })
+#     print(">> CHAT:", out2["reply"][0])