Update handler.py

handler.py

@@ -1,179 +1,597 @@
 #tope_version
 # handler.py
 # handler.py
-from typing import Any, Dict, List, Union, Optional
+# 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 TRIAL
+# handler.py
+import os
+import json
+import time
 import torch
+from typing import Any, Dict, List, Optional, Iterator
+
+# Note: these imports match your environment
 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':...}]."""
+
+# Try to import transformers' streaming helper (best-effort)
+try:
+    from transformers import TextIteratorStreamer
+    _HAS_STREAMER = True
+except Exception:
+    _HAS_STREAMER = False
+
+
+# ----- Config / constants -----
+BASE_MODEL = "unsloth/gemma-3-4b-it-unsloth-bnb-4bit"
+LORA_REPO = "Ephraimmm/PIDGIN_gemma-3"
+
+# Fixed system prompt (ALWAYS injected)
+PIDGIN_SYSTEM = (
+    "You be AI assistant always answer strictly for Nigerian pidgin. "
+    "Keep grammar natural like Naija street pidgin. "
+    "No switch go Queen English."
+)
+
+
+# ----- Helper utilities -----
+def _to_text_blocks(content: Any) -> List[Dict[str, str]]:
+    """Ensure content is a list of {'type':'text','text':...} blocks."""
     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)}]
- 
- 
+
+
+def _normalize_openai_messages(messages: Any) -> List[Dict[str, Any]]:
+    """
+    Normalize OpenAI-style messages to a list of dicts with block-style content.
+    Also ALWAYS prepend PIDGIN_SYSTEM as the first system message.
+    """
+    if isinstance(messages, str):
+        msgs = [
+            {"role": "system", "content": _to_text_blocks(PIDGIN_SYSTEM)},
+            {"role": "user", "content": _to_text_blocks(messages)},
+        ]
+        return msgs
+
+    if isinstance(messages, dict):
+        messages = [messages]
+
+    if not isinstance(messages, list):
+        raise ValueError("`messages` must be a string, dict, or list of message dicts.")
+
+    norm = []
+    for m in messages:
+        role = m.get("role", "user")
+        content = m.get("content", "")
+        # Accept both raw string content or block-style content
+        if isinstance(content, str):
+            content_blocks = _to_text_blocks(content)
+        else:
+            content_blocks = _to_text_blocks(content)
+        norm.append({"role": role, "content": content_blocks})
+
+    # ALWAYS prepend the pidgin system message
+    norm = [{"role": "system", "content": _to_text_blocks(PIDGIN_SYSTEM)}] + norm
+    return norm
+
+
+# ----- Endpoint handler -----
 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
+    HuggingFace Inference Endpoint handler.
+
+    Expects OpenAI-style body:
+    {
+      "messages": [...],
+      "max_tokens": 256,
+      "temperature": 0.7,
+      "top_p": 0.9,
+      "stream": true|false,
+      ...
+    }
+
+    When stream=True this returns a generator yielding JSONL string chunks
+    in OpenAI-like streaming format (data: ...\n\n).
     """
- 
-    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"
+    def __init__(self, model_dir: str = "", **kwargs: Any):
+        logger.info(f"Initializing EndpointHandler, model_dir={model_dir}")
+
+        # Load base model (unsloth 4-bit) and tokenizer via FastLanguageModel
+        # returns (model, tokenizer)
         model, tokenizer = FastLanguageModel.from_pretrained(
-            model_name=base_model_name,
+            model_name=BASE_MODEL,
             max_seq_length=2048,
             dtype=torch.float16,
             load_in_4bit=True,
         )
- 
-        lora_repo = "Ephraimmm/PIDGIN_gemma-3"
-        model = PeftModel.from_pretrained(model, lora_repo)
+
+        # Attach LoRA adapter (not merged)
+        model = PeftModel.from_pretrained(model, LORA_REPO)
+
+        # Prepare for inference (unsloth helper)
         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
+
+        # device (some wrappers set it on the model)
+        self.device = getattr(self.model, "device", torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+
+        # safety: ensure pad token exists
         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)
+
+        logger.info(f"Loaded model. Device={self.device}, eos_id={self.tokenizer.eos_token_id}")
+
+    # Build tokenized inputs using tokenizer.apply_chat_template if available
+    def _build_inputs(self, messages: Any):
+        norm_msgs = _normalize_openai_messages(messages)
+        if hasattr(self.tokenizer, "apply_chat_template"):
+            # Mirror your existing pipeline: add_generation_prompt=True, tokenize=True/return_tensors
+            model_inputs = self.tokenizer.apply_chat_template(
+                norm_msgs,
+                add_generation_prompt=True,
+                return_tensors="pt",
+                tokenize=True,
+                return_dict=True,
+            ).to(self.device)
+        else:
+            # Fallback: naive formatting (Gemma chat format)
+            formatted = []
+            for m in norm_msgs:
+                role = m.get("role", "user")
+                text = "".join([b.get("text", "") for b in m.get("content", [])])
+                if role in ("system", "user"):
+                    formatted.append(f"<start_of_turn>user\n{text}<end_of_turn>\n")
+                else:
+                    formatted.append(f"<start_of_turn>model\n{text}<end_of_turn>\n")
+            formatted.append("<start_of_turn>model\n")
+            prompt = "".join(formatted)
+            model_inputs = self.tokenizer(
+                prompt,
+                return_tensors="pt",
+                truncation=True,
+                max_length=2048
+            ).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)
- 
+
+    # Helper to produce OpenAI-style id
+    def _make_id(self) -> str:
+        return f"chatcmpl-{os.urandom(12).hex()}"
+
+    # Streaming generator: best-effort using TextIteratorStreamer if available,
+    # otherwise falling back to chunking final text.
+    def _stream_generator(
+        self,
+        model_inputs: Dict[str, Any],
+        gen_kwargs: Dict[str, Any],
+        model_id: str,
+        index: int = 0,
+        chunk_size: int = 32,
+    ) -> Iterator[bytes]:
+        """
+        Yields bytes (JSONL lines) compatible with OpenAI streaming:
+          - data: <json>\n\n for each chunk
+          - data: [DONE]\n\n at end
+        Uses TextIteratorStreamer if available; otherwise generates full text and
+        yields sliced chunks (best-effort).
+        """
+        # Prepare some metadata
+        prompt_tokens = int(model_inputs["input_ids"].shape[1])
+        generated_tokens_so_far = 0
+        choices_metadata = {"index": index, "finish_reason": None}
+
+        # Try TextIteratorStreamer route (token-by-token streaming)
+        if _HAS_STREAMER:
+            try:
+                streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True)
+                gen_kwargs_stream = gen_kwargs.copy()
+                gen_kwargs_stream["streamer"] = streamer
+
+                # Kick off generation in a background thread to avoid blocking the generator;
+                # but since we are in a single-process environment, we'll call generate in another thread.
+                import threading
+                def _generate():
+                    with torch.no_grad():
+                        # model.generate will populate streamer as tokens arrive
+                        self.model.generate(**gen_kwargs_stream)
+
+                gen_thread = threading.Thread(target=_generate, daemon=True)
+                gen_thread.start()
+
+                # Iterate streamer for pieces of text
+                buffer = ""
+                for piece in streamer:
+                    if not piece:
+                        continue
+                    buffer += piece
+                    # when buffer reaches chunk_size, yield a chunk
+                    while len(buffer) >= chunk_size:
+                        chunk_text = buffer[:chunk_size]
+                        buffer = buffer[chunk_size:]
+                        generated_tokens_so_far += len(self.tokenizer.encode(chunk_text, add_special_tokens=False))
+                        data = {
+                            "id": model_id,
+                            "object": "chat.completion.chunk",
+                            "choices": [
+                                {
+                                    "delta": {"content": chunk_text},
+                                    "index": index,
+                                    "finish_reason": None,
+                                }
+                            ],
+                        }
+                        yield (f"data: {json.dumps(data)}\n\n").encode("utf-8")
+
+                # flush remaining buffer
+                if buffer:
+                    generated_tokens_so_far += len(self.tokenizer.encode(buffer, add_special_tokens=False))
+                    data = {
+                        "id": model_id,
+                        "object": "chat.completion.chunk",
+                        "choices": [
+                            {
+                                "delta": {"content": buffer},
+                                "index": index,
+                                "finish_reason": None,
+                            }
+                        ],
+                    }
+                    yield (f"data: {json.dumps(data)}\n\n").encode("utf-8")
+
+                # Wait for generation to finish
+                gen_thread.join(timeout=1.0)
+
+                # final "done" event with finish_reason and usage (best-effort)
+                completion_tokens = generated_tokens_so_far
+                total_tokens = prompt_tokens + completion_tokens
+                final_data = {
+                    "id": model_id,
+                    "object": "chat.completion",
+                    "choices": [
+                        {
+                            "index": index,
+                            "message": {
+                                "role": "assistant",
+                                "content": "",  # full content is not repeated in chunk streaming
+                            },
+                            "finish_reason": "stop",
+                        }
+                    ],
+                    "usage": {
+                        "prompt_tokens": prompt_tokens,
+                        "completion_tokens": completion_tokens,
+                        "total_tokens": total_tokens,
+                    },
+                }
+                # final chunk with DONE marker
+                yield (f"data: {json.dumps(final_data)}\n\n").encode("utf-8")
+                yield b"data: [DONE]\n\n"
+                return
+
+            except Exception as ex:
+                logger.warning(f"TextIteratorStreamer failed, falling back to chunking. Reason: {ex}")
+
+        # Fallback: non-streaming generate then chunk it
+        with torch.no_grad():
+            output_ids = self.model.generate(**gen_kwargs)
+        input_len = model_inputs["input_ids"].shape[-1]
+        generated_ids = output_ids[0, input_len:]
+        full_text = self.tokenizer.decode(generated_ids, skip_special_tokens=True)
+        # yield chunks
+        start = 0
+        while start < len(full_text):
+            chunk_text = full_text[start:start + chunk_size]
+            start += chunk_size
+            data = {
+                "id": model_id,
+                "object": "chat.completion.chunk",
+                "choices": [
+                    {
+                        "delta": {"content": chunk_text},
+                        "index": index,
+                        "finish_reason": None,
+                    }
+                ],
+            }
+            yield (f"data: {json.dumps(data)}\n\n").encode("utf-8")
+            time.sleep(0.01)  # small pause to give a streaming feel
+
+        # final metadata
+        completion_tokens = int(generated_ids.shape[-1])
+        total_tokens = prompt_tokens + completion_tokens
+        final_data = {
+            "id": model_id,
+            "object": "chat.completion",
+            "choices": [
+                {
+                    "index": index,
+                    "message": {
+                        "role": "assistant",
+                        "content": full_text,
+                    },
+                    "finish_reason": "stop",
+                }
+            ],
+            "usage": {
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": total_tokens,
+            },
+        }
+        yield (f"data: {json.dumps(final_data)}\n\n").encode("utf-8")
+        yield b"data: [DONE]\n\n"
+
+    # The core entrypoint used by HF Inference Endpoint
+    def __call__(self, data: Dict[str, Any]):
+        """
+        Expect OpenAI-style body. Returns either:
+          - dict (full OpenAI response) if stream=False
+          - generator yielding bytes (JSONL) if stream=True
+        """
+        # Validate
+        messages = data.get("messages")
+        if messages is None:
+            raise ValueError("Missing required field `messages` (OpenAI chat format).")
+
+        # read parameters (OpenAI names)
+        max_tokens = int(data.get("max_tokens", 256))
+        temperature = float(data.get("temperature", 0.7))
+        top_p = float(data.get("top_p", 0.9))
+        top_k = data.get("top_k", None)
+        stream = bool(data.get("stream", False))
+        do_sample = temperature > 0.0
+
+        # Build model inputs (normalizes messages and ALWAYS injects system prompt)
+        model_inputs = self._build_inputs(messages)
+
+        # generation kwargs (map OpenAI names to HF generate)
         gen_kwargs = dict(
-            **inputs,
-            max_new_tokens=max_new_tokens,
+            **model_inputs,
+            max_new_tokens=max_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,
+            do_sample=do_sample,
+            pad_token_id=self.tokenizer.pad_token_id or self.tokenizer.eos_token_id,
+            eos_token_id=self.tokenizer.eos_token_id,
         )
         if top_k is not None:
             gen_kwargs["top_k"] = int(top_k)
- 
+
+        model_id = self._make_id()
+
+        # STREAMING behavior: return generator of bytes (JSONL data:)
+        if stream:
+            return self._stream_generator(model_inputs=model_inputs, gen_kwargs=gen_kwargs, model_id=model_id)
+
+        # NON-STREAMING: generate and return full OpenAI-style completion dict
         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]
+
+        input_len = model_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],
+        generated_text = self.tokenizer.decode(generated_ids, skip_special_tokens=True)
+
+        prompt_tokens = int(input_len)
+        completion_tokens = int(generated_ids.shape[-1])
+        total_tokens = prompt_tokens + completion_tokens
+
+        resp = {
+            "id": model_id,
+            "object": "chat.completion",
+            "created": int(time.time()),
+            "model": LORA_REPO,
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {
+                        "role": "assistant",
+                        "content": generated_text.strip(),
+                    },
+                    "finish_reason": "stop",
+                }
+            ],
             "usage": {
-                "prompt_tokens": int(input_len),
-                "generated_tokens": int(generated_ids.shape[-1]),
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": total_tokens,
             },
         }
- 
- 
+        return resp
+
+
+# If you want to smoke-test locally
 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."}]
-            }
+    req = {
+        "messages": [
+            {"role": "user", "content": "How far?"},
         ],
-        "parameters": {"max_new_tokens": 64}
-    }))
+        "max_tokens": 64,
+        "temperature": 0.2,
+        "stream": False,
+    }
+    print(json.dumps(h(req), indent=2)[:1000])
+
 
 # from typing import Any, Dict, List, Union, Optional