Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import numpy as np | |
| from typing import Iterator, List, Optional | |
| from gluonts.dataset import Dataset as GluonDataset | |
| from gluonts.model import Forecast | |
| from gluonts.model.forecast import QuantileForecast | |
| from gluonts.model.predictor import RepresentablePredictor | |
| from tsfm_bench.eval.predictors import _ForecastConfig | |
| class MyUserModel(RepresentablePredictor): | |
| """A mock user-submitted pre-trained TSFM (e.g., an autoregressive or slightly perturbed model).""" | |
| def __init__( | |
| self, | |
| prediction_length: int, | |
| checkpoint_path: str | None = None, | |
| quantile_levels: Optional[List[float]] = None, | |
| ): | |
| super().__init__(prediction_length=prediction_length) | |
| self.config = _ForecastConfig.from_quantiles(quantile_levels) | |
| self.checkpoint_path = checkpoint_path | |
| # In a real model, we would load model weights here: | |
| # self.model = load_model_from_checkpoint(checkpoint_path) | |
| def predict(self, dataset: GluonDataset, **kwargs) -> Iterator[Forecast]: | |
| for entry in dataset: | |
| target = np.asarray(entry["target"], dtype=np.float64) | |
| univariate = target.ndim == 1 | |
| if univariate: | |
| working = target.reshape(1, -1) | |
| else: | |
| working = target | |
| # Autoregressive/perturbed simulation: | |
| # We take the last few points and project them forward with a small dampening factor | |
| # and a tiny bit of deterministic shift (to make it different from SeasonalNaive) | |
| horizon = self.prediction_length | |
| last_val = working[:, -1:] | |
| # Simple linear trend projection with noise dampening | |
| predictions = [] | |
| for h in range(1, horizon + 1): | |
| val = last_val * (0.985 ** h) + (np.sin(h) * 0.005 * last_val) | |
| predictions.append(val) | |
| mean = np.concatenate(predictions, axis=1) | |
| if univariate: | |
| mean = mean[0] | |
| arrays = [mean] | |
| if self.config.quantile_levels: | |
| for q in self.config.quantile_levels: | |
| # Deterministic quantiles for simulation | |
| if q <= 0.5: | |
| arrays.append(mean * (1 - (0.5 - q) * 0.08)) | |
| else: | |
| arrays.append(mean * (1 + (q - 0.5) * 0.08)) | |
| forecast_arrays = np.stack(arrays, axis=0) | |
| yield QuantileForecast( | |
| forecast_arrays=forecast_arrays, | |
| forecast_keys=self.config.forecast_keys, | |
| start_date=entry["start"], | |
| item_id=entry["item_id"], | |
| ) | |