Spaces:
Sleeping
Sleeping
| 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 | |