Upload service.py

service.py

@@ -124,8 +124,16 @@ class ChronosService:
 
     def _extract_forecast(self, forecast: Any) -> tuple[list[float], list[float]]:
         assert self._torch is not None
+        assert self._pipeline is not None
         torch = self._torch
 
+        if isinstance(forecast, (list, tuple)):
+            if not forecast:
+                raise RuntimeError("empty Chronos forecast output")
+            if len(forecast) != 1:
+                raise RuntimeError(f"unexpected Chronos batch size: {len(forecast)}")
+            forecast = forecast[0]
+
         if hasattr(forecast, "detach"):
             tensor = forecast.detach().cpu()
         else:
@@ -139,9 +147,27 @@ class ChronosService:
             mean_forecast = squeezed.tolist()
             std_forecast = [0.0 for _ in mean_forecast]
         else:
-            samples = squeezed.reshape(-1, squeezed.shape[-1])
-            mean_forecast = samples.mean(dim=0).tolist()
-            std_forecast = samples.std(dim=0).tolist()
+            if squeezed.ndim == 2:
+                quantile_tensor = squeezed
+            else:
+                quantile_tensor = squeezed.reshape(-1, squeezed.shape[-2], squeezed.shape[-1]).mean(dim=0)
+
+            quantiles = list(getattr(self._pipeline, "quantiles", []))
+            if not quantiles:
+                median_idx = quantile_tensor.shape[0] // 2
+                mean_forecast = quantile_tensor[median_idx].tolist()
+                std_forecast = [0.0 for _ in mean_forecast]
+            else:
+                median_idx = min(range(len(quantiles)), key=lambda idx: abs(quantiles[idx] - 0.5))
+                lower_idx = min(range(len(quantiles)), key=lambda idx: abs(quantiles[idx] - 0.1))
+                upper_idx = min(range(len(quantiles)), key=lambda idx: abs(quantiles[idx] - 0.9))
+                mean_forecast = quantile_tensor[median_idx].tolist()
+                lower_forecast = quantile_tensor[lower_idx].tolist()
+                upper_forecast = quantile_tensor[upper_idx].tolist()
+                std_forecast = [
+                    max(0.0, (float(upper) - float(lower)) / 2.0)
+                    for lower, upper in zip(lower_forecast, upper_forecast)
+                ]
 
         confidence: list[float] = []
         for pred, std in zip(mean_forecast, std_forecast):