Upload 3 files

Dockerfile

@@ -6,10 +6,11 @@ ENV PORT=7860
 ENV CUDA_VISIBLE_DEVICES=""
 ENV TRANSFORMERS_NO_ADVISORY_WARNINGS=1
 ENV HF_HUB_DISABLE_TELEMETRY=1
-ENV CHRONOS_BACKEND=baseline_cpu
+ENV CHRONOS_BACKEND=hf_cpu
 ENV CHRONOS_MAX_CONTEXT_LENGTH=512
 ENV CHRONOS_MAX_HORIZON_STEP=288
 ENV CHRONOS_MIN_REQUIRED_POINTS=32
+ENV CHRONOS_ALLOW_BASELINE_FALLBACK=false
 ENV UV_SYSTEM_PYTHON=1
 
 WORKDIR /app

requirements.txt

@@ -1,3 +1,6 @@
 fastapi==0.115.12
 uvicorn==0.34.0
 pydantic==2.11.3
+numpy>=2.2.0
+torch>=2.6.0
+chronos-forecasting>=2.0.0

service.py

@@ -1,137 +1,234 @@
-from __future__ import annotations
-
-import math
-import os
-from statistics import mean
-from typing import Any
-
-from schemas import HealthResponse, PredictRequest, PredictResponse, PredictionItem
-
-
-class ChronosService:
-    """CPU-first HF Space service wrapper for Chronos.
-
-    This scaffold is designed for HuggingFace free CPU Spaces and keeps the
-    serving contract aligned with `tsf-bridge`. The default backend is a
-    deterministic CPU baseline rather than real Chronos inference.
-    """
-
-    def __init__(self) -> None:
-        self.model_id = "chronos"
-        self.model_name = os.getenv(
-            "CHRONOS_MODEL_NAME",
-            "amazon/chronos-2",
-        )
-        self.backend = os.getenv("CHRONOS_BACKEND", "baseline_cpu").strip() or "baseline_cpu"
-        self.device = "cpu"
-        self.ready = True
-        self.max_context_length = int(os.getenv("CHRONOS_MAX_CONTEXT_LENGTH", "512"))
-        self.max_horizon_step = int(os.getenv("CHRONOS_MAX_HORIZON_STEP", "288"))
-        self.confidence_floor = float(os.getenv("CHRONOS_CONFIDENCE_FLOOR", "0.16"))
-        self.confidence_ceiling = float(os.getenv("CHRONOS_CONFIDENCE_CEILING", "0.80"))
-        self.min_required_points = int(os.getenv("CHRONOS_MIN_REQUIRED_POINTS", "32"))
-
-    def health(self) -> HealthResponse:
-        return HealthResponse(
-            status="ok",
-            model=self.model_name,
-            model_id=self.model_id,
-            backend=self.backend,
-            device=self.device,
-            ready=self.ready,
-            max_context_length=self.max_context_length,
-            max_horizon_step=self.max_horizon_step,
-        )
-
-    def predict(self, payload: PredictRequest) -> PredictResponse:
-        self._validate_request(payload)
-        closes = payload.close_prices[-payload.context_length :]
-        predictions = self._predict_with_baseline(closes, payload.horizons)
-        return PredictResponse(model_id=self.model_id, predictions=predictions)
-
-    def _validate_request(self, payload: PredictRequest) -> None:
-        if payload.context_length > self.max_context_length:
-            raise ValueError(
-                f"context_length {payload.context_length} exceeds "
-                f"CHRONOS_MAX_CONTEXT_LENGTH={self.max_context_length}"
-            )
-        if payload.context_length > len(payload.close_prices):
-            raise ValueError("context_length must not exceed len(close_prices)")
-        if len(payload.close_prices) < self.min_required_points:
-            raise ValueError(
-                f"at least {self.min_required_points} close prices are required "
-                "for CPU baseline stability"
-            )
-        if any(step > self.max_horizon_step for step in payload.horizons):
-            raise ValueError(
-                f"horizons contain values above CHRONOS_MAX_HORIZON_STEP={self.max_horizon_step}"
-            )
-
-    def _predict_with_baseline(
-        self, close_prices: list[float], horizons: list[int]
-    ) -> list[PredictionItem]:
-        last_price = close_prices[-1]
-        short_window = close_prices[-min(10, len(close_prices)) :]
-        mid_window = close_prices[-min(24, len(close_prices)) :]
-        long_window = close_prices[-min(64, len(close_prices)) :]
-
-        short_mean = mean(short_window)
-        mid_mean = mean(mid_window)
-        long_mean = mean(long_window)
-        momentum = 0.0 if short_mean == 0 else (last_price - short_mean) / short_mean
-        mean_reversion = 0.0 if long_mean == 0 else (mid_mean - long_mean) / long_mean
-        local_trend = self._slope(mid_window)
-
-        predictions: list[PredictionItem] = []
-        for step in horizons:
-            horizon_scale = min(1.0, math.log(step + 1.0) / 3.8)
-            expected_return = momentum * 0.35 + mean_reversion * 0.30 + local_trend * 0.35
-            expected_return *= horizon_scale
-
-            pred_price = max(0.00000001, last_price * (1.0 + expected_return))
-            confidence = self._confidence(close_prices, step, abs(expected_return))
-            predictions.append(
-                PredictionItem(
-                    step=step,
-                    pred_price=round(pred_price, 8),
-                    pred_confidence=round(confidence, 4),
-                )
-            )
-        return predictions
-
-    def _confidence(
-        self, close_prices: list[float], step: int, expected_move_abs: float
-    ) -> float:
-        if len(close_prices) < 3:
-            return self.confidence_floor
-
-        changes: list[float] = []
-        for previous, current in zip(close_prices[:-1], close_prices[1:]):
-            if previous <= 0:
-                continue
-            changes.append(abs((current - previous) / previous))
-
-        realized_vol = mean(changes[-min(64, len(changes)) :]) if changes else 0.0
-        stability = max(0.0, 1.0 - min(realized_vol * 18.0, 1.0))
-        horizon_decay = 1.0 / (1.0 + math.log(step + 1.0))
-        raw = 0.20 + min(expected_move_abs / (realized_vol + 1e-9), 2.0) * 0.18
-        raw += stability * 0.20 + horizon_decay * 0.22
-        return max(self.confidence_floor, min(self.confidence_ceiling, raw))
-
-    @staticmethod
-    def _slope(values: list[float]) -> float:
-        if len(values) < 2 or values[0] == 0:
-            return 0.0
-        return (values[-1] - values[0]) / values[0]
-
-    def describe_runtime(self) -> dict[str, Any]:
-        return {
-            "model_id": self.model_id,
-            "model_name": self.model_name,
-            "backend": self.backend,
-            "device": self.device,
-            "ready": self.ready,
-            "max_context_length": self.max_context_length,
-            "max_horizon_step": self.max_horizon_step,
-            "min_required_points": self.min_required_points,
-        }
+from __future__ import annotations
+
+import math
+import os
+from statistics import mean
+from typing import Any
+
+from schemas import HealthResponse, PredictRequest, PredictResponse, PredictionItem
+
+
+class ChronosService:
+    """HF Space service wrapper for Chronos-2."""
+
+    def __init__(self) -> None:
+        self.model_id = "chronos"
+        self.model_name = os.getenv(
+            "CHRONOS_MODEL_NAME",
+            "amazon/chronos-2",
+        )
+        self.backend = os.getenv("CHRONOS_BACKEND", "hf_cpu").strip() or "hf_cpu"
+        self.device = "cpu"
+        self.max_context_length = int(os.getenv("CHRONOS_MAX_CONTEXT_LENGTH", "512"))
+        self.max_horizon_step = int(os.getenv("CHRONOS_MAX_HORIZON_STEP", "288"))
+        self.confidence_floor = float(os.getenv("CHRONOS_CONFIDENCE_FLOOR", "0.16"))
+        self.confidence_ceiling = float(os.getenv("CHRONOS_CONFIDENCE_CEILING", "0.80"))
+        self.min_required_points = int(os.getenv("CHRONOS_MIN_REQUIRED_POINTS", "32"))
+        self.num_samples = int(os.getenv("CHRONOS_NUM_SAMPLES", "20"))
+        self.allow_baseline_fallback = os.getenv("CHRONOS_ALLOW_BASELINE_FALLBACK", "false").lower() == "true"
+
+        self.ready = False
+        self.load_error = ""
+        self._torch = None
+        self._pipeline = None
+
+        self._initialize_backend()
+
+    def health(self) -> HealthResponse:
+        return HealthResponse(
+            status="ok" if self.ready else "degraded",
+            model=self.model_name,
+            model_id=self.model_id,
+            backend=self.backend,
+            device=self.device,
+            ready=self.ready,
+            max_context_length=self.max_context_length,
+            max_horizon_step=self.max_horizon_step,
+        )
+
+    def predict(self, payload: PredictRequest) -> PredictResponse:
+        self._validate_request(payload)
+        closes = payload.close_prices[-payload.context_length :]
+
+        if self.backend == "hf_cpu":
+            if not self.ready:
+                raise RuntimeError(self.load_error or "chronos backend not ready")
+            predictions = self._predict_with_hf(closes, payload.horizons)
+        else:
+            predictions = self._predict_with_baseline(closes, payload.horizons)
+
+        return PredictResponse(model_id=self.model_id, predictions=predictions)
+
+    def _initialize_backend(self) -> None:
+        if self.backend == "baseline_cpu":
+            self.ready = True
+            return
+        if self.backend != "hf_cpu":
+            raise ValueError(f"unsupported CHRONOS_BACKEND={self.backend}")
+
+        try:
+            self._load_hf_model()
+            self.ready = True
+        except Exception as exc:
+            self.load_error = f"chronos hf load failed: {exc}"
+            if self.allow_baseline_fallback:
+                self.backend = "baseline_cpu"
+                self.ready = True
+            else:
+                self.ready = False
+
+    def _load_hf_model(self) -> None:
+        import torch
+        from chronos import Chronos2Pipeline
+
+        self._torch = torch
+        torch.set_num_threads(max(1, int(os.getenv("CHRONOS_TORCH_THREADS", "2"))))
+        self._pipeline = Chronos2Pipeline.from_pretrained(
+            self.model_name,
+            device_map="cpu",
+        )
+
+    def _predict_with_hf(
+        self, close_prices: list[float], horizons: list[int]
+    ) -> list[PredictionItem]:
+        assert self._torch is not None
+        assert self._pipeline is not None
+
+        torch = self._torch
+        context = torch.tensor(close_prices[-self.max_context_length :], dtype=torch.float32)
+        forecast = self._pipeline.predict(
+            context,
+            prediction_length=self.max_horizon_step,
+            num_samples=self.num_samples,
+        )
+
+        dense_mean, dense_conf = self._extract_forecast(forecast)
+        if len(dense_mean) < max(horizons):
+            raise RuntimeError(
+                f"Chronos output horizon {len(dense_mean)} is shorter than requested {max(horizons)}"
+            )
+
+        predictions: list[PredictionItem] = []
+        for step in horizons:
+            predictions.append(
+                PredictionItem(
+                    step=step,
+                    pred_price=round(max(0.00000001, float(dense_mean[step - 1])), 8),
+                    pred_confidence=round(dense_conf[step - 1], 4),
+                )
+            )
+        return predictions
+
+    def _extract_forecast(self, forecast: Any) -> tuple[list[float], list[float]]:
+        assert self._torch is not None
+        torch = self._torch
+
+        if hasattr(forecast, "detach"):
+            tensor = forecast.detach().cpu()
+        else:
+            tensor = torch.as_tensor(forecast)
+
+        if tensor.ndim == 3:
+            samples = tensor[0]
+            mean_forecast = samples.mean(dim=0).tolist()
+            std_forecast = samples.std(dim=0).tolist()
+        elif tensor.ndim == 2:
+            mean_forecast = tensor[0].tolist()
+            std_forecast = [0.0 for _ in mean_forecast]
+        else:
+            raise RuntimeError(f"unexpected Chronos forecast shape: {tuple(tensor.shape)}")
+
+        confidence: list[float] = []
+        for pred, std in zip(mean_forecast, std_forecast):
+            dispersion = abs(float(std)) / max(abs(float(pred)), 1e-6)
+            raw = 1.0 / (1.0 + dispersion)
+            confidence.append(max(self.confidence_floor, min(self.confidence_ceiling, raw)))
+        return [float(item) for item in mean_forecast], confidence
+
+    def _validate_request(self, payload: PredictRequest) -> None:
+        if payload.context_length > self.max_context_length:
+            raise ValueError(
+                f"context_length {payload.context_length} exceeds "
+                f"CHRONOS_MAX_CONTEXT_LENGTH={self.max_context_length}"
+            )
+        if payload.context_length > len(payload.close_prices):
+            raise ValueError("context_length must not exceed len(close_prices)")
+        if len(payload.close_prices) < self.min_required_points:
+            raise ValueError(
+                f"at least {self.min_required_points} close prices are required "
+                "for Chronos stability"
+            )
+        if any(step > self.max_horizon_step for step in payload.horizons):
+            raise ValueError(
+                f"horizons contain values above CHRONOS_MAX_HORIZON_STEP={self.max_horizon_step}"
+            )
+
+    def _predict_with_baseline(
+        self, close_prices: list[float], horizons: list[int]
+    ) -> list[PredictionItem]:
+        last_price = close_prices[-1]
+        short_window = close_prices[-min(10, len(close_prices)) :]
+        mid_window = close_prices[-min(24, len(close_prices)) :]
+        long_window = close_prices[-min(64, len(close_prices)) :]
+
+        short_mean = mean(short_window)
+        mid_mean = mean(mid_window)
+        long_mean = mean(long_window)
+        momentum = 0.0 if short_mean == 0 else (last_price - short_mean) / short_mean
+        mean_reversion = 0.0 if long_mean == 0 else (mid_mean - long_mean) / long_mean
+        local_trend = self._slope(mid_window)
+
+        predictions: list[PredictionItem] = []
+        for step in horizons:
+            horizon_scale = min(1.0, math.log(step + 1.0) / 3.8)
+            expected_return = momentum * 0.35 + mean_reversion * 0.30 + local_trend * 0.35
+            expected_return *= horizon_scale
+
+            pred_price = max(0.00000001, last_price * (1.0 + expected_return))
+            confidence = self._baseline_confidence(close_prices, step, abs(expected_return))
+            predictions.append(
+                PredictionItem(
+                    step=step,
+                    pred_price=round(pred_price, 8),
+                    pred_confidence=round(confidence, 4),
+                )
+            )
+        return predictions
+
+    def _baseline_confidence(
+        self, close_prices: list[float], step: int, expected_move_abs: float
+    ) -> float:
+        if len(close_prices) < 3:
+            return self.confidence_floor
+
+        changes: list[float] = []
+        for previous, current in zip(close_prices[:-1], close_prices[1:]):
+            if previous <= 0:
+                continue
+            changes.append(abs((current - previous) / previous))
+
+        realized_vol = mean(changes[-min(64, len(changes)) :]) if changes else 0.0
+        stability = max(0.0, 1.0 - min(realized_vol * 18.0, 1.0))
+        horizon_decay = 1.0 / (1.0 + math.log(step + 1.0))
+        raw = 0.20 + min(expected_move_abs / (realized_vol + 1e-9), 2.0) * 0.18
+        raw += stability * 0.20 + horizon_decay * 0.22
+        return max(self.confidence_floor, min(self.confidence_ceiling, raw))
+
+    @staticmethod
+    def _slope(values: list[float]) -> float:
+        if len(values) < 2 or values[0] == 0:
+            return 0.0
+        return (values[-1] - values[0]) / values[0]
+
+    def describe_runtime(self) -> dict[str, Any]:
+        return {
+            "model_id": self.model_id,
+            "model_name": self.model_name,
+            "backend": self.backend,
+            "device": self.device,
+            "ready": self.ready,
+            "load_error": self.load_error,
+            "max_context_length": self.max_context_length,
+            "max_horizon_step": self.max_horizon_step,
+            "min_required_points": self.min_required_points,
+        }