Create handler.py

handler.py

@@ -0,0 +1,126 @@
+# handler.py
+from __future__ import annotations
+
+from typing import Any, Dict, List, Union
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+Json = Dict[str, Any]
+Messages = List[Dict[str, str]]  # [{"role":"user|assistant|system", "content":"..."}]
+
+
+def _is_messages(x: Any) -> bool:
+    return (
+        isinstance(x, list)
+        and len(x) > 0
+        and all(isinstance(m, dict) and "role" in m and "content" in m for m in x)
+    )
+
+
+class EndpointHandler:
+    """
+    Hugging Face Inference Endpoints custom handler.
+    Expects:
+      - request body is a dict
+      - always contains `inputs`
+      - may contain `parameters` for generation
+    """
+
+    def __init__(self, model_dir: str):
+        self.model_dir = model_dir
+
+        # Pick dtype/device
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        if self.device == "cuda":
+            # bfloat16 is usually safe on A100/H100; if your instance doesn't support bf16, change to float16
+            self.dtype = torch.bfloat16
+        else:
+            self.dtype = torch.float32
+
+        # IMPORTANT: trust_remote_code=True because repo contains AsteriskForCausalLM.py + auto_map
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_dir,
+            trust_remote_code=True,
+            use_fast=True,
+        )
+
+        # Make sure pad token exists (your config uses pad_token_id=2 which equals eos_token_id in many llama-like models)
+        if self.tokenizer.pad_token_id is None and self.tokenizer.eos_token_id is not None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_dir,
+            trust_remote_code=True,
+            torch_dtype=self.dtype,
+            device_map="auto" if self.device == "cuda" else None,
+        )
+
+        if self.device != "cuda":
+            self.model.to(self.device)
+
+        self.model.eval()
+
+    @torch.inference_mode()
+    def __call__(self, data: Json) -> Union[Json, List[Json]]:
+        inputs = data.get("inputs", "")
+        params = data.get("parameters", {}) or {}
+
+        # Generation defaults (can be overridden via `parameters`)
+        max_new_tokens = int(params.get("max_new_tokens", 256))
+        temperature = float(params.get("temperature", 0.7))
+        top_p = float(params.get("top_p", 0.95))
+        top_k = int(params.get("top_k", 0))
+        repetition_penalty = float(params.get("repetition_penalty", 1.0))
+
+        do_sample = bool(params.get("do_sample", temperature > 0))
+        num_beams = int(params.get("num_beams", 1))
+
+        def _one(item: Any) -> Json:
+            # Accept:
+            # 1) string prompt
+            # 2) messages list: [{"role":"user","content":"..."}]
+            # 3) dict {"messages":[...]} (common chat style)
+            if isinstance(item, dict) and "messages" in item:
+                item = item["messages"]
+
+            if _is_messages(item):
+                # Chat template path exists in repo; tokenizer.apply_chat_template will use it if configured
+                input_ids = self.tokenizer.apply_chat_template(
+                    item,
+                    return_tensors="pt",
+                    add_generation_prompt=True,
+                )
+            else:
+                if not isinstance(item, str):
+                    item = str(item)
+                enc = self.tokenizer(item, return_tensors="pt")
+                input_ids = enc["input_ids"]
+
+            input_ids = input_ids.to(self.model.device)
+            input_len = input_ids.shape[-1]
+
+            gen_ids = self.model.generate(
+                input_ids=input_ids,
+                max_new_tokens=max_new_tokens,
+                do_sample=do_sample,
+                temperature=temperature if do_sample else None,
+                top_p=top_p if do_sample else None,
+                top_k=top_k if do_sample and top_k > 0 else None,
+                num_beams=num_beams,
+                repetition_penalty=repetition_penalty,
+                pad_token_id=self.tokenizer.pad_token_id,
+                eos_token_id=self.tokenizer.eos_token_id,
+            )
+
+            # Only return newly generated tokens
+            new_tokens = gen_ids[0, input_len:]
+            text = self.tokenizer.decode(new_tokens, skip_special_tokens=True)
+            return {"generated_text": text}
+
+        # Batch support
+        if isinstance(inputs, list) and not _is_messages(inputs):
+            return [_one(x) for x in inputs]
+        else:
+            return _one(inputs)