Add custom inference handler

b6c9fbe 2 months ago

2.89 kB

"""
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}

	"""
	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}