Update handler.py

handler.py

@@ -1,596 +1,181 @@
 #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 TRIAL
-# handler.py
-import os
-import json
-import time
+from typing import Any, Dict, List, Union, Optional
 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
-
-# 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."""
+ 
+# 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)}]
-
-
-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:
     """
-    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).
+    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 = "", **kwargs: Any):
-        logger.info(f"Initializing EndpointHandler, model_dir={model_dir}")
-
-        # Load base model (unsloth 4-bit) and tokenizer via FastLanguageModel
-        # returns (model, tokenizer)
+ 
+    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,
+            model_name=base_model_name,
             max_seq_length=2048,
             dtype=torch.float16,
             load_in_4bit=True,
         )
-
-        # Attach LoRA adapter (not merged)
-        model = PeftModel.from_pretrained(model, LORA_REPO)
-
-        # Prepare for inference (unsloth helper)
+ 
+        lora_repo = "Ephraimmm/PIDGIN_gemma-3"
+        model = PeftModel.from_pretrained(model, lora_repo)
         FastLanguageModel.for_inference(model)
-
+ 
         self.model = model.eval()
         self.tokenizer = tokenizer
-
-        # 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
+        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"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)
+ 
+        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
-
-    # 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)
+ 
+    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(
-            **model_inputs,
-            max_new_tokens=max_tokens,
+            **inputs,
+            max_new_tokens=max_new_tokens,
             temperature=temperature,
             top_p=top_p,
-            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,
+            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)
-
-        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)
-
-        input_len = model_inputs["input_ids"].shape[-1]
+ 
+        # Return ONLY the continuation (like your streamer)
+        input_len = inputs["input_ids"].shape[-1]
         generated_ids = output_ids[0, input_len:]
-        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",
-                }
-            ],
+        reply_text = self.tokenizer.decode(generated_ids, skip_special_tokens=True)
+ 
+        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 resp
-
-
-# If you want to smoke-test locally
+ 
+ 
 if __name__ == "__main__":
     h = EndpointHandler(model_dir=".")
-    req = {
-        "messages": [
-            {"role": "user", "content": "How far?"},
+    # 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,
-        "temperature": 0.2,
-        "stream": False,
-    }
-    print(json.dumps(h(req), indent=2)[:1000])
+        "parameters": {"max_new_tokens": 64}
+    }))
+
+
 
 
 # from typing import Any, Dict, List, Union, Optional