Create handler.py

handler.py

@@ -0,0 +1,82 @@
+"""
+HF Inference Endpoints handler for ChatTS-14B.
+
+Expected request JSON:
+    {
+        "inputs": {
+            "prompt": "Describe the trend of this series.",
+            "timeseries": [[0.1, 0.2, 0.3, ...]],          # list of float lists, one per <ts><ts/>
+            "max_new_tokens": 300
+        }
+    }
+
+The prompt MUST contain one `<ts><ts/>` placeholder per series in `timeseries`.
+
+Response:
+    {"generated_text": "..."}
+"""
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import torch
+from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer
+
+
+class EndpointHandler:
+    def __init__(self, path: str = "") -> None:
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        dtype = torch.float16 if self.device == "cuda" else torch.float32
+
+        self.model = AutoModelForCausalLM.from_pretrained(
+            path,
+            trust_remote_code=True,
+            torch_dtype=dtype,
+            device_map=0 if self.device == "cuda" else None,
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
+        self.processor = AutoProcessor.from_pretrained(
+            path, trust_remote_code=True, tokenizer=self.tokenizer
+        )
+        self.model.eval()
+
+    def __call__(self, data: dict[str, Any]) -> dict[str, Any]:
+        inputs = data.get("inputs", {})
+        if isinstance(inputs, str):
+            return {
+                "error": "ChatTS requires structured inputs. "
+                "Use {'inputs': {'prompt': str, 'timeseries': [[...]], 'max_new_tokens': int}}"
+            }
+
+        prompt: str = inputs["prompt"]
+        ts_lists = inputs["timeseries"]
+        max_new_tokens: int = int(inputs.get("max_new_tokens", 300))
+
+        ts_arrays = [np.asarray(t, dtype=np.float64) for t in ts_lists]
+
+        formatted = (
+            "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
+            f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
+        )
+
+        proc_inputs = self.processor(
+            text=[formatted],
+            timeseries=ts_arrays,
+            padding=True,
+            return_tensors="pt",
+        )
+        proc_inputs = {k: v.to(self.device) for k, v in proc_inputs.items()}
+
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **proc_inputs,
+                max_new_tokens=max_new_tokens,
+                do_sample=False,
+            )
+
+        generated = self.tokenizer.batch_decode(
+            outputs[:, proc_inputs["input_ids"].shape[1] :],
+            skip_special_tokens=True,
+        )
+        return {"generated_text": generated[0]}