Update handler.py

handler.py

@@ -1,28 +1,28 @@
+#tope_version
 # handler.py
-
-from typing import Any, Dict, List, Union
-
+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:
     """
-    Custom handler for Hugging Face Inference Endpoints.
-    Loads an Unsloth 4-bit base and attaches a LoRA adapter.
+    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):
-        """
-        Called once at startup. `model_dir` is the downloaded repo path.
-        """
-        logger.info(f"Initializing EndpointHandler with model_dir = {model_dir}")
-
-        # 1) Load Unsloth 4-bit base
+        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,
@@ -30,125 +30,289 @@ class EndpointHandler:
             dtype=torch.float16,
             load_in_4bit=True,
         )
-
-        # 2) Attach LoRA adapter
+ 
+        # --- LoRA adapter ---
         lora_repo = "Ephraimmm/PIDGIN_gemma-3"
         model = PeftModel.from_pretrained(model, lora_repo)  # default adapter
-        FastLanguageModel.for_inference(model)  # optimize for generation
-
-        # 3) Finalize
+        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"))
+        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"
         )
-        logger.info(f"Model ready. Device: {self.device}, EOS: {self.tokenizer.eos_token_id}")
-
-    def _build_input_ids(self, prompt: Prompt):
+ 
+    # -----------------------------
+    # 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):
         """
-        Accepts either a single string, or OpenAI-style chat messages:
-        [{"role":"system|user|assistant","content":"..."}]
-        Always injects system prompt for Pidgin English.
+        Use tokenizer.apply_chat_template exactly like your Gradio code:
+          add_generation_prompt=True, return_tensors='pt', tokenize=True.
         """
-        system_instruction = {
-            "role": "system",
-            "content": "Always reply in Pidgin English, answer in context of question."
-        }
-
-        if isinstance(prompt, list):
-            # Ensure system prompt is first
-            if not any(m.get("role") == "system" for m in prompt):
-                prompt = [system_instruction] + prompt
-
-            if hasattr(self.tokenizer, "apply_chat_template"):
-                text = self.tokenizer.apply_chat_template(
-                    prompt, tokenize=False, add_generation_prompt=True
-                )
-            else:
-                # Fallback concat
-                text = "\n".join(
-                    f"{m.get('role','user')}: {m.get('content','')}"
-                    for m in prompt
-                ) + "\nassistant:"
-        else:
-            # If it's just a raw string, prepend system message
-            text = f"System: {system_instruction['content']}\nUser: {prompt}\nAssistant:"
-
-        return self.tokenizer(text, return_tensors="pt").to(self.device), text
-
+        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]:
         """
-        Called per request. Expects {"inputs": <str|messages>, "parameters": {...}}.
+        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"EndpointHandler received data keys: {list(data.keys())}")
-
+        logger.info(f"Incoming keys: {list(data.keys())}")
         if "inputs" not in data:
-            raise ValueError("Request JSON must contain an `inputs` field (str or chat messages).")
-
+            raise ValueError("Missing `inputs` in request body.")
+ 
         prompt: Prompt = data["inputs"]
-        params = data.get("parameters", {}) or {}
-
-        max_new_tokens = int(params.get("max_new_tokens", 64))
-        temperature = float(params.get("temperature", 1))
-        top_k = int(params.get("top_k", 64))
-        top_p = float(params.get("top_p", 0.95))
-        do_sample = bool(params.get("do_sample", True))
-        repetition_penalty = float(params.get("repetition_penalty", 1.1))
+        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)
-
-        inputs, rendered_prompt = self._build_input_ids(prompt)
-
-        with torch.no_grad():
-            output_ids = self.model.generate(
-                **inputs,
-                max_new_tokens=max_new_tokens,
-                temperature=temperature,
-                top_k=top_k,
-                top_p=top_p,
-                do_sample=do_sample,
-                repetition_penalty=repetition_penalty,
-                eos_token_id=eos_token_id,
-            )
-
-        full_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
-
-        # Best-effort continuation extraction
-        continuation = (
-            full_text[len(rendered_prompt):]
-            if isinstance(prompt, str) and full_text.startswith(rendered_prompt)
-            else full_text
+        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": [full_text],
+            "reply": [reply_text],
             "usage": {
-                "prompt_tokens": int(inputs["input_ids"].shape[-1]),
-                "generated_tokens": int((output_ids.shape[-1] - inputs["input_ids"].shape[-1])),
+                "prompt_tokens": int(input_len),
+                "generated_tokens": int(generated_ids.shape[-1]),
             },
         }
-
-
+ 
+ 
 # Local quick test
 if __name__ == "__main__":
     h = EndpointHandler(model_dir=".")
-
-    # Test with plain string
-    resp1 = h({
-        "inputs": "Explain why sky dey blue.",
+ 
+    # Raw string
+    out1 = h({
+        "inputs": "How far?",
         "parameters": {"max_new_tokens": 64}
     })
-    print("Response (string):", resp1)
-
-    # Test with chat-style messages
-    resp2 = h({
+    print(">> STRING:", out1["reply"][0])
+ 
+    # Chat-style messages
+    out2 = h({
         "inputs": [
-            {"role": "user", "content": "How I go cook jollof rice?"}
+            {"role": "user", "content": "Explain why sky dey blue."}
         ],
         "parameters": {"max_new_tokens": 64}
     })
-    print("Response (chat):", resp2)
+    print(">> CHAT:", out2["reply"][0])
+
+
+
+# # handler.py
+
+# from typing import Any, Dict, List, Union
+
+# 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:
+#     """
+#     Custom handler for Hugging Face Inference Endpoints.
+#     Loads an Unsloth 4-bit base and attaches a LoRA adapter.
+#     """
+
+#     def __init__(self, model_dir: str, **_: Any):
+#         """
+#         Called once at startup. `model_dir` is the downloaded repo path.
+#         """
+#         logger.info(f"Initializing EndpointHandler with model_dir = {model_dir}")
+
+#         # 1) Load Unsloth 4-bit base
+#         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,
+#         )
+
+#         # 2) Attach LoRA adapter
+#         lora_repo = "Ephraimmm/PIDGIN_gemma-3"
+#         model = PeftModel.from_pretrained(model, lora_repo)  # default adapter
+#         FastLanguageModel.for_inference(model)  # optimize for generation
+
+#         # 3) Finalize
+#         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"Model ready. Device: {self.device}, EOS: {self.tokenizer.eos_token_id}")
+
+#     def _build_input_ids(self, prompt: Prompt):
+#         """
+#         Accepts either a single string, or OpenAI-style chat messages:
+#         [{"role":"system|user|assistant","content":"..."}]
+#         Always injects system prompt for Pidgin English.
+#         """
+#         system_instruction = {
+#             "role": "system",
+#             "content": "Always reply in Pidgin English, answer in context of question."
+#         }
+
+#         if isinstance(prompt, list):
+#             # Ensure system prompt is first
+#             if not any(m.get("role") == "system" for m in prompt):
+#                 prompt = [system_instruction] + prompt
+
+#             if hasattr(self.tokenizer, "apply_chat_template"):
+#                 text = self.tokenizer.apply_chat_template(
+#                     prompt, tokenize=False, add_generation_prompt=True
+#                 )
+#             else:
+#                 # Fallback concat
+#                 text = "\n".join(
+#                     f"{m.get('role','user')}: {m.get('content','')}"
+#                     for m in prompt
+#                 ) + "\nassistant:"
+#         else:
+#             # If it's just a raw string, prepend system message
+#             text = f"System: {system_instruction['content']}\nUser: {prompt}\nAssistant:"
+
+#         return self.tokenizer(text, return_tensors="pt").to(self.device), text
+
+#     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+#         """
+#         Called per request. Expects {"inputs": <str|messages>, "parameters": {...}}.
+#         """
+#         logger.info(f"EndpointHandler received data keys: {list(data.keys())}")
+
+#         if "inputs" not in data:
+#             raise ValueError("Request JSON must contain an `inputs` field (str or chat messages).")
+
+#         prompt: Prompt = data["inputs"]
+#         params = data.get("parameters", {}) or {}
+
+#         max_new_tokens = int(params.get("max_new_tokens", 64))
+#         temperature = float(params.get("temperature", 1))
+#         top_k = int(params.get("top_k", 64))
+#         top_p = float(params.get("top_p", 0.95))
+#         do_sample = bool(params.get("do_sample", True))
+#         repetition_penalty = float(params.get("repetition_penalty", 1.1))
+#         eos_token_id = params.get("eos_token_id", self.tokenizer.eos_token_id)
+
+#         inputs, rendered_prompt = self._build_input_ids(prompt)
+
+#         with torch.no_grad():
+#             output_ids = self.model.generate(
+#                 **inputs,
+#                 max_new_tokens=max_new_tokens,
+#                 temperature=temperature,
+#                 top_k=top_k,
+#                 top_p=top_p,
+#                 do_sample=do_sample,
+#                 repetition_penalty=repetition_penalty,
+#                 eos_token_id=eos_token_id,
+#             )
+
+#         full_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+#         # Best-effort continuation extraction
+#         continuation = (
+#             full_text[len(rendered_prompt):]
+#             if isinstance(prompt, str) and full_text.startswith(rendered_prompt)
+#             else full_text
+#         )
+
+#         return {
+#             "reply": [full_text],
+#             "usage": {
+#                 "prompt_tokens": int(inputs["input_ids"].shape[-1]),
+#                 "generated_tokens": int((output_ids.shape[-1] - inputs["input_ids"].shape[-1])),
+#             },
+#         }
+
+
+# # Local quick test
+# if __name__ == "__main__":
+#     h = EndpointHandler(model_dir=".")
+
+#     # Test with plain string
+#     resp1 = h({
+#         "inputs": "Explain why sky dey blue.",
+#         "parameters": {"max_new_tokens": 64}
+#     })
+#     print("Response (string):", resp1)
+
+#     # Test with chat-style messages
+#     resp2 = h({
+#         "inputs": [
+#             {"role": "user", "content": "How I go cook jollof rice?"}
+#         ],
+#         "parameters": {"max_new_tokens": 64}
+#     })
+#     print("Response (chat):", resp2)