Update handler.py

handler.py

@@ -3,195 +3,347 @@
 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:
 
-    """
 
+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."
+        }
 
-        Called once at startup. `model_dir` is the downloaded repo path.
+        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())}")
 
-        logger.info(f"Initializing EndpointHandler with model_dir = {model_dir}")
+        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)
+
+
+
+# from typing import Any, Dict, List, Union
+
+# import torch
+
+# from huggingface_inference_toolkit.logging import logger
+
+# from unsloth import FastLanguageModel
+
+# from peft import PeftModel
  
-        # 1) Load Unsloth 4-bit base
+# Prompt = Union[str, List[Dict[str, str]]]
+ 
+ 
+# class EndpointHandler:
 
-        base_model_name = "unsloth/gemma-3-4b-it-unsloth-bnb-4bit"
+#     """
 
-        model, tokenizer = FastLanguageModel.from_pretrained(
+#     Custom handler for Hugging Face Inference Endpoints.
 
-            model_name      = base_model_name,
+#     Loads an Unsloth 4-bit base and attaches a LoRA adapter.
 
-            max_seq_length  = 2048,
+#     """
+ 
+#     def __init__(self, model_dir: str, **_: Any):
 
-            dtype           = torch.float16,
+#         """
 
-            load_in_4bit    = True,
+#         Called once at startup. `model_dir` is the downloaded repo path.
 
-        )
+#         """
+
+#         logger.info(f"Initializing EndpointHandler with model_dir = {model_dir}")
  
-        # 2) Attach LoRA adapter
+#         # 1) Load Unsloth 4-bit base
 
-        lora_repo = "Ephraimmm/PIDGIN_gemma-3"
+#         base_model_name = "unsloth/gemma-3-4b-it-unsloth-bnb-4bit"
+
+#         model, tokenizer = FastLanguageModel.from_pretrained(
 
-        model = PeftModel.from_pretrained(model, lora_repo)  # use default adapter name
+#             model_name      = base_model_name,
 
-        FastLanguageModel.for_inference(model)               # optimize for generation
+#             max_seq_length  = 2048,
+
+#             dtype           = torch.float16,
+
+#             load_in_4bit    = True,
+
+#         )
  
-        # 3) Finalize
+#         # 2) Attach LoRA adapter
 
-        self.model = model.eval()
+#         lora_repo = "Ephraimmm/PIDGIN_gemma-3"
 
-        self.tokenizer = tokenizer
+#         model = PeftModel.from_pretrained(model, lora_repo)  # use default adapter name
 
-        self.device = getattr(self.model, "device", torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+#         FastLanguageModel.for_inference(model)               # optimize for generation
+ 
+#         # 3) Finalize
 
-        logger.info(f"Model ready. Device: {self.device}, EOS: {self.tokenizer.eos_token_id}")
+#         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):
+#     def _build_input_ids(self, prompt: Prompt):
 
-        """
+#         """
 
-        Accepts either a single string, or OpenAI-style chat messages:
+#         Accepts either a single string, or OpenAI-style chat messages:
 
-        [{"role":"system|user|assistant","content":"..."}]
+#         [{"role":"system|user|assistant","content":"..."}]
 
-        """
+#         """
 
-        if isinstance(prompt, list):
+#         if isinstance(prompt, list):
 
-            # If tokenizer has a chat template, use it.
+#             # If tokenizer has a chat template, use it.
 
-            if hasattr(self.tokenizer, "apply_chat_template"):
+#             if hasattr(self.tokenizer, "apply_chat_template"):
 
-                text = self.tokenizer.apply_chat_template(
+#                 text = self.tokenizer.apply_chat_template(
 
-                    prompt, tokenize=False, add_generation_prompt=True
+#                     prompt, tokenize=False, add_generation_prompt=True
 
-                )
+#                 )
 
-            else:
+#             else:
 
-                # Fallback: naive concat
+#                 # Fallback: naive concat
 
-                text = "\n".join(f"{m.get('role','user')}: {m.get('content','')}" for m in prompt) + "\nassistant:"
+#                 text = "\n".join(f"{m.get('role','user')}: {m.get('content','')}" for m in prompt) + "\nassistant:"
 
-        else:
+#         else:
 
-            text = prompt
+#             text = prompt
  
-        return self.tokenizer(text, return_tensors="pt").to(self.device), text
+#         return self.tokenizer(text, return_tensors="pt").to(self.device), text
  
-    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+#     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
 
-        """
+#         """
 
-        Called per request. Expects {"inputs": <str|messages>, "parameters": {...}}.
+#         Called per request. Expects {"inputs": <str|messages>, "parameters": {...}}.
 
-        """
+#         """
 
-        logger.info(f"EndpointHandler received data keys: {list(data.keys())}")
+#         logger.info(f"EndpointHandler received data keys: {list(data.keys())}")
  
-        if "inputs" not in data:
+#         if "inputs" not in data:
 
-            raise ValueError("Request JSON must contain an `inputs` field (str or chat messages).")
+#             raise ValueError("Request JSON must contain an `inputs` field (str or chat messages).")
  
-        prompt: Prompt = data["inputs"]
+#         prompt: Prompt = data["inputs"]
 
-        params = data.get("parameters", {}) or {}
+#         params = data.get("parameters", {}) or {}
  
-        max_new_tokens = int(params.get("max_new_tokens", 64))
+#         max_new_tokens = int(params.get("max_new_tokens", 64))
 
-        temperature    = float(params.get("temperature", 1))
+#         temperature    = float(params.get("temperature", 1))
 
-        top_k          = int(params.get("top_k", 64))
+#         top_k          = int(params.get("top_k", 64))
 
-        top_p          = float(params.get("top_p", 0.95))
+#         top_p          = float(params.get("top_p", 0.95))
 
-        do_sample      = bool(params.get("do_sample", True))
+#         do_sample      = bool(params.get("do_sample", True))
 
-        repetition_penalty = float(params.get("repetition_penalty", 1.1))
+#         repetition_penalty = float(params.get("repetition_penalty", 1.1))
 
-        eos_token_id   = params.get("eos_token_id", self.tokenizer.eos_token_id)
+#         eos_token_id   = params.get("eos_token_id", self.tokenizer.eos_token_id)
  
-        inputs, rendered_prompt = self._build_input_ids(prompt)
+#         inputs, rendered_prompt = self._build_input_ids(prompt)
  
-        with torch.no_grad():
+#         with torch.no_grad():
 
-            output_ids = self.model.generate(
+#             output_ids = self.model.generate(
 
-                **inputs,
+#                 **inputs,
 
-                max_new_tokens=max_new_tokens,
+#                 max_new_tokens=max_new_tokens,
 
-                temperature=temperature,
+#                 temperature=temperature,
 
-                top_k=top_k,
+#                 top_k=top_k,
 
-                top_p=top_p,
+#                 top_p=top_p,
 
-                do_sample=do_sample,
+#                 do_sample=do_sample,
 
-                repetition_penalty=repetition_penalty,
+#                 repetition_penalty=repetition_penalty,
 
-                eos_token_id=eos_token_id,
+#                 eos_token_id=eos_token_id,
 
-            )
+#             )
  
-        full_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
+#         full_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
  
-        # Best-effort continuation extraction
+#         # Best-effort continuation extraction
 
-        continuation = full_text[len(rendered_prompt):] if isinstance(prompt, str) and full_text.startswith(rendered_prompt) else full_text
+#         continuation = full_text[len(rendered_prompt):] if isinstance(prompt, str) and full_text.startswith(rendered_prompt) else full_text
  
-        return {
+#         return {
 
-            "reply": [full_text],
+#             "reply": [full_text],
 
-            # "continuations": [continuation],
+#             # "continuations": [continuation],
 
-            "usage": {
+#             "usage": {
 
-                "prompt_tokens": int(inputs["input_ids"].shape[-1]),
+#                 "prompt_tokens": int(inputs["input_ids"].shape[-1]),
 
-                "generated_tokens": int((output_ids.shape[-1] - inputs["input_ids"].shape[-1])),
+#                 "generated_tokens": int((output_ids.shape[-1] - inputs["input_ids"].shape[-1])),
 
-            },}
+#             },}
  
  
-# Local quick test
+# # Local quick test
 
-if __name__ == "__main__":
+# if __name__ == "__main__":
 
-    h = EndpointHandler(model_dir=".")
+#     h = EndpointHandler(model_dir=".")
 
-    resp = h({
+#     resp = h({
 
-        "inputs": "Write in Nigerian Pidgin: How you dey?",
+#         "inputs": "Write in Nigerian Pidgin: How you dey?",
 
-        "parameters": {"max_new_tokens": 64}
+#         "parameters": {"max_new_tokens": 64}
 
-    })
+#     })
 
-    print("Response:", resp)
+#     print("Response:", resp)