File size: 2,889 Bytes
b6c9fbe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | """
Custom Inference Endpoint handler for amazon/chronos-2, shaped to match the contract
server/services/forecastService.js already calls (see HF_ENDPOINT_URL there).
Deploy target: your OWN Hugging Face copy of amazon/chronos-2 (duplicated into your
namespace), with this file and requirements.txt added at the repo root. HF Inference
Endpoints auto-detects handler.py and switches the endpoint's task to "Custom" — see
hf-endpoint/chronos-2-handler/README.md for the exact steps.
Request shape sent by forecastService.js:
{ "inputs": [<daily prices...>], "parameters": { "prediction_length": 7, "num_samples": 20 } }
Response shape forecastService.js parses (server/services/forecastService.js:337-351):
{ "quantiles": { "0.1": [...], "0.5": [...], "0.9": [...] } }
Each array has `prediction_length` entries (one per forecasted day).
"""
from typing import Any, Dict
import pandas as pd
from chronos import Chronos2Pipeline
QUANTILE_LEVELS = [0.1, 0.5, 0.9]
class EndpointHandler:
def __init__(self, path: str = ""):
# path is the directory HF mounts the repo's model weights into.
self.pipeline = Chronos2Pipeline.from_pretrained(path, device_map="cpu")
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
inputs = data.get("inputs")
if not isinstance(inputs, list) or len(inputs) == 0:
return {"error": "inputs must be a non-empty list of numbers"}
parameters = data.get("parameters") or {}
prediction_length = int(parameters.get("prediction_length", 7))
context_df = pd.DataFrame({
"id": ["route"] * len(inputs),
# Chronos-2 only uses these to order/space the series; the collector logs one
# observation per day, so a daily index is the right granularity here.
"timestamp": pd.date_range(end=pd.Timestamp.today().normalize(), periods=len(inputs), freq="D"),
"target": [float(x) for x in inputs],
})
pred_df = self.pipeline.predict_df(
context_df,
prediction_length=prediction_length,
quantile_levels=QUANTILE_LEVELS,
id_column="id",
timestamp_column="timestamp",
target="target",
)
quantiles = {}
for level in QUANTILE_LEVELS:
key = str(level)
if key in pred_df.columns:
quantiles[key] = pred_df[key].tolist()
elif level == 0.5 and "predictions" in pred_df.columns:
# Some Chronos-2 releases name the point forecast column "predictions"
# rather than "0.5" even when 0.5 is in quantile_levels.
quantiles[key] = pred_df["predictions"].tolist()
if "0.5" not in quantiles:
return {"error": f"unexpected predict_df columns: {list(pred_df.columns)}"}
return {"quantiles": quantiles}
|