Spaces:
Sleeping
Sleeping
File size: 2,704 Bytes
e317359 | 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 | 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"],
)
|