Spaces:
Running
Running
File size: 5,760 Bytes
9831ced a9e6751 9831ced a9e6751 9831ced a9e6751 9831ced | 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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | import numpy as np
import pandas as pd
import threading
from models.base_model import BaseForecastModel, validate_history_and_horizon
from config import MOIRAI_CHECKPOINT, RANDOM_SEED
from features.feature_pipeline import validate_features
_RNG_LOCK = threading.Lock()
_MODULE_CACHE = {}
_MODULE_LOCK = threading.Lock()
def _step_from_history(history: pd.Series) -> pd.Timedelta:
if isinstance(history.index, pd.DatetimeIndex) and len(history.index) >= 2:
diffs = history.index.to_series().diff().dropna()
if not diffs.empty:
step = diffs.median()
if step > pd.Timedelta(0):
if not diffs.eq(step).all():
raise ValueError('Moirai needs a continuous history; irregular timestamps cannot be relabelled.')
return step
return pd.Timedelta(minutes=1)
class MoiraiModel(BaseForecastModel):
name = "Moirai"
def __init__(
self,
checkpoint: str = MOIRAI_CHECKPOINT,
context_length: int = 256,
num_samples: int = 20,
):
self.checkpoint = checkpoint
self.context_length = context_length
self.num_samples = num_samples
self._module = None
self.last_diagnostics = {}
def _get_module(self):
from uni2ts.model.moirai import MoiraiModule
if self._module is None:
with _MODULE_LOCK:
self._module = _MODULE_CACHE.get(self.checkpoint)
if self._module is None:
self._module = MoiraiModule.from_pretrained(self.checkpoint)
self._module.eval()
_MODULE_CACHE[self.checkpoint] = self._module
return self._module
def predict(
self, history: pd.Series, horizon: int = 1, features: pd.DataFrame = None
) -> list:
self.last_diagnostics = {}
validate_history_and_horizon(history, horizon, self.name)
try:
import torch
except ImportError as error:
raise ImportError('Moirai requires the optional foundation-model packages. Install requirements-foundation.txt.') from error
with _RNG_LOCK, torch.random.fork_rng(devices=[]), torch.inference_mode():
torch.manual_seed(RANDOM_SEED)
return self._predict(history, horizon, features)
def _predict(self, history, horizon, features):
values = np.asarray(history, dtype=np.float32)
if len(values) < 10:
raise ValueError("Moirai needs at least 10 candles of history.")
validate_history_and_horizon(values, horizon, self.name)
try:
from uni2ts.model.moirai import MoiraiForecast
from gluonts.dataset.pandas import PandasDataset
except ImportError as e:
raise ImportError(
f"Moirai needs the 'uni2ts' and 'gluonts' packages. Install with:\n pip install uni2ts torch\n(original error: {e})"
)
if features is not None:
validate_features(
features,
expected_len=len(values),
model_name="Moirai",
check_nan=False,
)
# Bound the context for the small checkpoint so a large dashboard
# window cannot cause an avoidable CPU-memory spike.
context_length = min(max(int(self.context_length), 32), 512)
context = values[-context_length:] if len(values) > context_length else values
step = _step_from_history(history)
end_ts = pd.Timestamp(history.index[-1]) if isinstance(history.index, pd.DatetimeIndex) else pd.Timestamp('2000-01-01')
if end_ts.tzinfo is not None:
end_ts = end_ts.tz_convert('UTC').tz_localize(None)
idx = pd.date_range(end=end_ts, periods=len(context), freq=step)
data = {"target": context}
feat_cols = []
if features is not None:
feat_context = features.iloc[-len(context) :].reset_index(drop=True)
validate_features(
feat_context, expected_len=len(context), model_name="Moirai"
)
for col in feat_context.columns:
data[col] = feat_context[col].to_numpy(dtype=np.float32)
feat_cols = list(feat_context.columns)
series_df = pd.DataFrame(data, index=idx)
dataset = PandasDataset(
series_df,
target="target",
past_feat_dynamic_real=feat_cols if feat_cols else None,
)
module = self._get_module()
forecaster = MoiraiForecast(
module=module,
prediction_length=horizon,
context_length=len(context),
patch_size="auto",
num_samples=self.num_samples,
target_dim=1,
feat_dynamic_real_dim=0,
past_feat_dynamic_real_dim=len(feat_cols),
)
predictor = forecaster.create_predictor(batch_size=1)
forecast = next(iter(predictor.predict(dataset)))
samples = np.asarray(forecast.samples, dtype=float)
if samples.ndim == 3 and samples.shape[-1] == 1:
samples = samples[..., 0]
if samples.ndim != 2 or samples.shape != (self.num_samples, horizon):
raise ValueError('Moirai returned an unexpected univariate sample/horizon shape.')
out = np.median(samples, axis=0).tolist()
if not all(np.isfinite(x) and x > 0 for x in out):
raise ValueError(
"Moirai produced a non-finite forecast (NaN/inf) for this window."
)
self.last_diagnostics = {'checkpoint': self.checkpoint, 'context_candles': len(context), 'point_statistic': 'sample median', 'seed': RANDOM_SEED, 'fallback': False, 'weights_loaded': True}
return out
|