Update handler.py

handler.py

@@ -1,194 +1,18 @@
 #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
- 
-# # 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:
-#     """
-#     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
-#     - Always replies in Nigerian Pidgin
-#     """
- 
-#     def __init__(self, model_dir: str, **_: Any):
-#         logger.info(f"Initializing with model_dir={model_dir}")
- 
-#         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_repo = "Ephraimmm/PIDGIN_gemma-3"
-#         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"),
-#         )
-#         # 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}")
- 
-#         self._system_text = (
-#             "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."
-#         )
- 
-#     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):
-#             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):
-#         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,
-#             return_tensors="pt",
-#             tokenize=True,
-#             return_dict=True,
-#         ).to(self.device)
-#         return model_inputs
- 
-#     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-#         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 {})
- 
-#         # 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", self.tokenizer.pad_token_id)
- 
-#         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,                 # 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 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 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)
- 
-#         return {
-#             "reply": [reply_text],
-#             "usage": {
-#                 "prompt_tokens": int(input_len),
-#                 "generated_tokens": int(generated_ids.shape[-1]),
-#             },
-#         }
- 
- 
-# if __name__ == "__main__":
-#     h = EndpointHandler(model_dir=".")
-#     # Raw string
-#     print(h({"inputs": "How far?", "parameters": {"max_new_tokens": 64}}))
-#     # Block-style content (your Gradio shape)
-#     print(h({
-#         "inputs": [
-#             {
-#                 "role": "user",
-#                 "content": [{"type": "text", "text": "Explain why sky dey blue."}]
-#             }
-#         ],
-#         "parameters": {"max_new_tokens": 64}
-#     }))
-
-#openai api
 from typing import Any, Dict, List, Union, Optional
 import torch
-import time
 from huggingface_inference_toolkit.logging import logger
 from unsloth import FastLanguageModel
 from peft import PeftModel
-
+ 
 # 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:
@@ -198,22 +22,20 @@ def _to_block_list(content: Any) -> List[Dict[str, str]]:
         return content
     # Fallback: stringify
     return [{"type": "text", "text": str(content)}]
-
-
+ 
+ 
 class EndpointHandler:
     """
-    HF Inference Endpoint handler that mirrors OpenAI API format:
-    - Accepts OpenAI-style requests with 'messages', 'model', etc.
-    - Returns OpenAI-style responses with 'choices', 'usage', etc.
+    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
     - Always replies in Nigerian Pidgin
     """
-
+ 
     def __init__(self, model_dir: str, **_: Any):
         logger.info(f"Initializing with model_dir={model_dir}")
-
+ 
         base_model_name = "unsloth/gemma-3-4b-it-unsloth-bnb-4bit"
         model, tokenizer = FastLanguageModel.from_pretrained(
             model_name=base_model_name,
@@ -221,11 +43,11 @@ class EndpointHandler:
             dtype=torch.float16,
             load_in_4bit=True,
         )
-
+ 
         lora_repo = "Ephraimmm/PIDGIN_gemma-3"
         model = PeftModel.from_pretrained(model, lora_repo)
         FastLanguageModel.for_inference(model)
-
+ 
         self.model = model.eval()
         self.tokenizer = tokenizer
         self.device = getattr(
@@ -236,143 +58,83 @@ class EndpointHandler:
         # 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}")
-
+ 
         self._system_text = (
             "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."
         )
-
-    def _normalize_messages(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
-        """Normalize OpenAI-style messages to block-style contents."""
+ 
+    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):
+            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 messages:
+        for m in prompt:
             role = m.get("role", "user")
-            content = m.get("content", "")
-            
-            # Convert content to block list format
-            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))
-            
+            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, messages: List[Dict[str, Any]]):
-        normalized_messages = self._normalize_messages(messages)
+ 
+    def _build_model_inputs(self, prompt: Prompt):
+        messages = self._normalize_messages(prompt)
+        # Use the same path as your Gradio script
         model_inputs = self.tokenizer.apply_chat_template(
-            normalized_messages,
+            messages,
             add_generation_prompt=True,
             return_tensors="pt",
             tokenize=True,
             return_dict=True,
         ).to(self.device)
         return model_inputs
-
+ 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        """
-        Process OpenAI-style requests and return OpenAI-style responses.
-        
-        Input format (OpenAI API):
-        {
-            "messages": [{"role": "user", "content": "Hello"}],
-            "model": "gpt-3.5-turbo",  # optional
-            "temperature": 0.7,         # optional
-            "max_tokens": 256,          # optional
-            "top_p": 1.0,              # optional
-            ...
-        }
-        
-        Output format (OpenAI API):
-        {
-            "id": "chatcmpl-...",
-            "object": "chat.completion",
-            "created": 1234567890,
-            "model": "gemma-3-4b-pidgin",
-            "choices": [{
-                "index": 0,
-                "message": {
-                    "role": "assistant",
-                    "content": "Response text"
-                },
-                "finish_reason": "stop"
-            }],
-            "usage": {
-                "prompt_tokens": 10,
-                "completion_tokens": 20,
-                "total_tokens": 30
-            }
-        }
-        """
         logger.info(f"Incoming keys: {list(data.keys())}")
-        
-        # 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", {})
-        
-        # 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")
-        )
+        if "inputs" not in data:
+            raise ValueError("Missing `inputs` in request body.")
+ 
+        prompt: Prompt = data["inputs"]
+        params: Dict[str, Any] = (data.get("parameters") or {})
+ 
+        # 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", self.tokenizer.pad_token_id)
-
-        inputs = self._build_model_inputs(messages)
-
+ 
+        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,
+            do_sample=True,                 # streaming-like sampling on
             use_cache=use_cache,
             repetition_penalty=repetition_penalty,
             eos_token_id=eos_token_id,
@@ -380,588 +142,35 @@ class EndpointHandler:
         )
         if top_k is not None:
             gen_kwargs["top_k"] = int(top_k)
-
+ 
         with torch.no_grad():
             output_ids = self.model.generate(**gen_kwargs)
-
-        # Extract only the generated tokens (exclude prompt)
+ 
+        # 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)
-
-        # Calculate tokens
-        prompt_tokens = int(input_len)
-        completion_tokens = int(generated_ids.shape[-1])
-        total_tokens = prompt_tokens + completion_tokens
-
-        # Return OpenAI-style response
-        response = {
-            "id": f"chatcmpl-{int(time.time())}",
-            "object": "chat.completion",
-            "created": int(time.time()),
-            "model": data.get("model", "gemma-3-4b-pidgin"),
-            "choices": [
-                {
-                    "index": 0,
-                    "message": {
-                        "role": "assistant",
-                        "content": reply_text
-                    },
-                    "finish_reason": "stop"
-                }
-            ],
+ 
+        return {
+            "reply": [reply_text],
             "usage": {
-                "prompt_tokens": prompt_tokens,
-                "completion_tokens": completion_tokens,
-                "total_tokens": total_tokens
-            }
+                "prompt_tokens": int(input_len),
+                "generated_tokens": int(generated_ids.shape[-1]),
+            },
         }
-
-        return response
-
-
+ 
+ 
 if __name__ == "__main__":
     h = EndpointHandler(model_dir=".")
-    
-    # Test with OpenAI-style format
-    print("=== OpenAI Format Test ===")
-    response = h({
-        "messages": [
-            {"role": "user", "content": "How far?"}
-        ],
-        "temperature": 0.7,
-        "max_tokens": 64
-    })
-    print(response)
-    
-    print("\n=== OpenAI Format with System Message ===")
-    response = h({
-        "messages": [
-            {"role": "system", "content": "You are a helpful assistant."},
-            {"role": "user", "content": "Explain why sky dey blue."}
+    # Raw string
+    print(h({"inputs": "How far?", "parameters": {"max_new_tokens": 64}}))
+    # Block-style content (your Gradio shape)
+    print(h({
+        "inputs": [
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": "Explain why sky dey blue."}]
+            }
         ],
-        "max_tokens": 64
-    })
-    print(response)
-    
-    # Legacy format still works
-    print("\n=== Legacy Format Test ===")
-    response = h({
-        "inputs": "How far?",
         "parameters": {"max_new_tokens": 64}
-    })
-    print(response)
-
-
-
-# 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])
-
-
-
-# # 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)
-
-
-
-# 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)  # use default adapter name
-
-#         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":"..."}]
-
-#         """
-
-#         if isinstance(prompt, list):
-
-#             # If tokenizer has a chat template, use it.
-
-#             if hasattr(self.tokenizer, "apply_chat_template"):
-
-#                 text = self.tokenizer.apply_chat_template(
-
-#                     prompt, tokenize=False, add_generation_prompt=True
-
-#                 )
-
-#             else:
-
-#                 # Fallback: naive concat
-
-#                 text = "\n".join(f"{m.get('role','user')}: {m.get('content','')}" for m in prompt) + "\nassistant:"
-
-#         else:
-
-#             text = prompt
- 
-#         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],
-
-#             # "continuations": [continuation],
-
-#             "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=".")
-
-#     resp = h({
-
-#         "inputs": "Write in Nigerian Pidgin: How you dey?",
-
-#         "parameters": {"max_new_tokens": 64}
-
-#     })
-
-#     print("Response:", resp)
-
- 
+    }))