Create handler.py

handler.py

@@ -0,0 +1,159 @@
+# handler.py
+# Hugging Face Inference Toolkit custom handler for chat-style or plain text prompts.
+# Supports two input formats:
+# 1) HF standard: {"inputs": "your prompt", "parameters": {...}}
+# 2) Chat format: {"messages": [{"role":"system"|"user"|"assistant","content":"..."}], "parameters": {...}}
+
+import os
+import json
+import torch
+
+from typing import Any, Dict, List, Optional, Union
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    TextIteratorStreamer,
+)
+
+# Optional: respect a few env knobs (set in Endpoint settings)
+DEFAULT_MAX_NEW_TOKENS = int(os.getenv("MAX_NEW_TOKENS", "512"))
+DEFAULT_TEMPERATURE    = float(os.getenv("TEMPERATURE", "0.7"))
+DEFAULT_TOP_P          = float(os.getenv("TOP_P", "0.9"))
+DEFAULT_TOP_K          = int(os.getenv("TOP_K", "50"))
+DEFAULT_DO_SAMPLE      = os.getenv("DO_SAMPLE", "true").lower() in {"1","true","yes"}
+DEFAULT_REPETITION_PEN = float(os.getenv("REPETITION_PENALTY", "1.05"))
+
+class EndpointHandler:
+    """Hugging Face custom handler contract"""
+    def __init__(self, model_dir: str, *args, **kwargs):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        dtype = torch.bfloat16 if self.device == "cuda" else torch.float32
+
+        # Load model + tokenizer from model_dir (automatically provided by the endpoint)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_dir, use_fast=True)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_dir,
+            torch_dtype=dtype,
+            device_map="auto" if self.device == "cuda" else None,
+        )
+        self.model.eval()
+
+        # Some chat models need this to avoid warnings
+        if not self.tokenizer.pad_token_id:
+            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+
+    # -------- Utilities --------
+    def _to_prompt_from_messages(self, messages: List[Dict[str, str]]) -> str:
+        """
+        If tokenizer supports chat template, use it. Otherwise, build a simple prompt.
+        messages = [{"role":"system"|"user"|"assistant","content":"..."}]
+        """
+        if hasattr(self.tokenizer, "apply_chat_template"):
+            try:
+                return self.tokenizer.apply_chat_template(
+                    messages,
+                    tokenize=False,
+                    add_generation_prompt=True
+                )
+            except Exception:
+                pass  # fallback below
+
+        # Minimal fallback prompt
+        role_map = {"system": "System", "user": "User", "assistant": "Assistant"}
+        lines = []
+        for m in messages:
+            role = role_map.get(m.get("role","user"), "User")
+            content = m.get("content","")
+            lines.append(f"{role}: {content}")
+        lines.append("Assistant:")
+        return "\n".join(lines)
+
+    def _pack_inputs(
+        self,
+        payload: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """
+        Normalize inbound payload to a single string prompt + gen params.
+        Accepts:
+          {"inputs": "...", "parameters": {...}}
+        or
+          {"messages":[...], "parameters": {...}}
+        """
+        parameters = payload.get("parameters", {}) or {}
+
+        max_new_tokens   = int(parameters.get("max_new_tokens", DEFAULT_MAX_NEW_TOKENS))
+        temperature      = float(parameters.get("temperature", DEFAULT_TEMPERATURE))
+        top_p            = float(parameters.get("top_p", DEFAULT_TOP_P))
+        top_k            = int(parameters.get("top_k", DEFAULT_TOP_K))
+        do_sample        = bool(parameters.get("do_sample", DEFAULT_DO_SAMPLE))
+        repetition_pen   = float(parameters.get("repetition_penalty", DEFAULT_REPETITION_PEN))
+
+        if "messages" in payload:
+            prompt = self._to_prompt_from_messages(payload["messages"])
+        else:
+            prompt = payload.get("inputs", "")
+            if not isinstance(prompt, str):
+                # Some clients send list[str]
+                if isinstance(prompt, list) and prompt and isinstance(prompt[0], str):
+                    prompt = prompt[0]
+                else:
+                    prompt = str(prompt)
+
+        gen_kwargs = {
+            "max_new_tokens": max_new_tokens,
+            "temperature": temperature,
+            "top_p": top_p,
+            "top_k": top_k,
+            "do_sample": do_sample,
+            "repetition_penalty": repetition_pen,
+            "eos_token_id": self.tokenizer.eos_token_id,
+            "pad_token_id": self.tokenizer.pad_token_id,
+        }
+        return {"prompt": prompt, "gen_kwargs": gen_kwargs}
+
+    # -------- Main inference entry --------
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Return shape (recommended):
+          {"text": "..."}                      # single string
+        Optional:
+          {"generated_text": "..."}            # alt field some tools expect
+          {"usage": {"prompt_tokens":..., "completion_tokens":..., "total_tokens":...}}
+        """
+        try:
+            packed = self._pack_inputs(data)
+            prompt = packed["prompt"]
+            gen_kwargs = packed["gen_kwargs"]
+
+            if not prompt.strip():
+                return {"text": "Empty prompt."}
+
+            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+
+            with torch.no_grad():
+                output_ids = self.model.generate(
+                    **inputs,
+                    **gen_kwargs,
+                )
+
+            # Remove the prompt portion to get only the newly generated tokens
+            gen_ids = output_ids[0][inputs["input_ids"].shape[-1]:]
+            text = self.tokenizer.decode(gen_ids, skip_special_tokens=True)
+
+            # Basic usage metrics (approx)
+            prompt_tokens = int(inputs["input_ids"].numel())
+            completion_tokens = int(gen_ids.numel())
+            total_tokens = prompt_tokens + completion_tokens
+
+            return {
+                "text": text,
+                "generated_text": text,
+                "usage": {
+                    "prompt_tokens": prompt_tokens,
+                    "completion_tokens": completion_tokens,
+                    "total_tokens": total_tokens,
+                },
+            }
+        except Exception as e:
+            # Never crash the container: return a JSON error
+            return {"error": str(e)}