Skip to content

Add sklearn-compatible adapter (anomaly, forecasting, embedding) - #86

Open
DrLHS wants to merge 1 commit into
moment-timeseries-foundation-model:mainfrom
DrLHS:lhs/sklearn-adapter
Open

Add sklearn-compatible adapter (anomaly, forecasting, embedding)#86
DrLHS wants to merge 1 commit into
moment-timeseries-foundation-model:mainfrom
DrLHS:lhs/sklearn-adapter

Conversation

@DrLHS

@DrLHS DrLHS commented Jun 7, 2026

Copy link
Copy Markdown

This PR adds momentfm.sklearn_adapter, three sklearn estimator-API adapters that let MOMENT drop into sklearn.pipeline.Pipeline, ColumnTransformer, GridSearchCV, and any other tooling that follows the sklearn contract.

Adapters

  • MOMENTAnomalyDetector(TransformerMixin) — reconstruction-MSE per-sample anomaly score. transform(X) returns shape (n_samples,); score_samples(X) returns the negated score in the sklearn outlier-detector convention (higher = more normal).
  • MOMENTForecaster(RegressorMixin)predict(X) returns the forecast horizon as a flat 2D array of shape (n_samples, n_channels * forecast_horizon) so it slots into sklearn pipelines without a custom reshape step. Construct with MOMENTForecaster(forecast_horizon=H).
  • MOMENTEmbedder(TransformerMixin)transform(X) returns the encoder embedding as a 2D feature matrix (n_samples, embedding_dim) so downstream sklearn estimators (LogisticRegression, KMeans, t-SNE) can consume them directly.

Input shape

All three adapters accept either 2D (n_samples, n_channels * context_length) input (sklearn pipeline convention) or 3D (n_samples, n_channels, context_length) (MOMENT native). Shape coercion is centralised in a _to_3d helper with explicit error messages.

fit contract

fit is zero-shot by default — it loads the pretrained weights and sets the model to eval(). This matches MOMENT's intended zero-shot-foundation-model usage. Calling predict or transform before fit is equivalent to calling fit first.

Tests

tests/test_sklearn_adapter.py adds 18 tests:

  • 6 shape-coercion tests (2D passthrough, 2D → 3D, multichannel, rejection of 4D, rejection of wrong 2D width, rejection of wrong 3D shape)
  • 6 sklearn API conformance tests (`get_params` and `set_params` parametrised across all three adapters)
  • 1 `MOMENTForecaster` validation test (must set `forecast_horizon` before `fit`)
  • 5 integration tests covering the real forward pass for each adapter plus a sklearn-Pipeline integration test. Marked `skip` by default so CI doesn't need network access; flip the `skipif` to `False` to enable when weights are cached locally.

13 unit tests pass on the current branch; 5 integration tests correctly skip.

README

A new ## 🧩 scikit-learn Integration section just before the tutorials section shows the three usage idioms.

What this enables

  • Drop-in MOMENT for anyone with an existing sklearn-pipeline-based feature engineering or model stack
  • Reproducible grid search over MOMENT model sizes via standard sklearn `GridSearchCV`
  • Mixed pipelines that pre-process tabular features upstream and feed MOMENT downstream
  • A natural integration target for tabular ML practitioners who don't want to write a custom PyTorch loop

Files changed

  • `momentfm/sklearn_adapter.py` (new, ~250 lines)
  • `tests/test_sklearn_adapter.py` (new, 18 tests)
  • `README.md` (one new section, 35 lines)

Companion to #85 (CI) and #84 (release-cut request).

Three sklearn estimator-API adapters that let MOMENT drop into sklearn.pipeline.Pipeline, ColumnTransformer, GridSearchCV, and any other tooling that follows the sklearn contract.

- MOMENTAnomalyDetector (TransformerMixin): reconstruction-MSE per-sample anomaly score.
- MOMENTForecaster (RegressorMixin): predict(X) returns flat 2D forecast horizon.
- MOMENTEmbedder (TransformerMixin): transform(X) returns encoder embeddings.

All three adapters accept 2D (n_samples, n_channels * context_length) or 3D (n_samples, n_channels, context_length) input. fit() loads the pretrained weights zero-shot; predict/transform perform inference.

Tests: 13 sklearn-API unit tests pass; 5 end-to-end smoke tests are marked skip by default since they need downloaded weights.

README adds a scikit-learn Integration section.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant