modern-fm 0.2.1__tar.gz

modern_fm-0.2.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Masaya Kawamata
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

modern_fm-0.2.1/PKG-INFO

@@ -0,0 +1,117 @@
+Metadata-Version: 2.4
+Name: modern-fm
+Version: 0.2.1
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Rust
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Dist: numpy>=1.24
+Requires-Dist: scipy>=1.10
+Requires-Dist: pytest>=7 ; extra == 'dev'
+Requires-Dist: ruff>=0.4 ; extra == 'dev'
+Provides-Extra: dev
+License-File: LICENSE
+Summary: Fast, sklearn-compatible Factorization Machines and Field-aware Factorization Machines
+Keywords: factorization-machines,ffm,ctr,recommender,tabular
+Author: Masaya Kawamata
+License: MIT
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
+Project-URL: Homepage, https://github.com/Matapanino/modern_fm
+Project-URL: Issues, https://github.com/Matapanino/modern_fm/issues
+Project-URL: Repository, https://github.com/Matapanino/modern_fm
+
+# modern_fm
+
+Fast, sklearn-compatible Factorization Machines (FM) and Field-aware
+Factorization Machines (FFM) for Python.
+
+**Status: v0.2 (Beta).** A Rust CPU backend (parity-tested against pure-NumPy
+reference implementations) drives sklearn-style estimators — `FMClassifier`
+(binary + multiclass softmax), `FMRegressor`, and `FFMClassifier`
+(binary + multiclass softmax) — with the SGD / AdaGrad / Adam /
+**FTRL-Proximal** optimizers, **mini-batch**
+gradient averaging (`batch_size`), **multi-core training** via `rayon`
+(`n_jobs`), plus `sample_weight`/`class_weight`, `label_smoothing`, early
+stopping, a `CategoricalEncoder`, and `save_model`/`load_model`. FTRL's L1
+(`l1_linear`/`l1_factors`) yields exact-zero weights. See `docs/roadmap.md` for
+remaining niche gaps (FTRL + early stopping, multiclass + early-stopping for FFM).
+
+## Installation
+
+```bash
+pip install modern-fm        # once published; prebuilt wheels, no Rust needed
+```
+
+Until the first PyPI release, install from source (requires a Rust toolchain;
+see Development below).
+
+## Usage
+
+```python
+from modern_fm import FMClassifier, FFMClassifier
+
+model = FMClassifier(
+    n_factors=16,
+    optimizer="adagrad",
+    learning_rate=0.05,
+    max_iter=100,
+    batch_size=256,        # mini-batch gradient averaging (1 = per-row SGD)
+    n_jobs=-1,             # train batches across all CPU cores
+    l2_linear=1e-5,
+    l2_factors=1e-5,
+    random_state=42,
+)
+model.fit(X_train, y_train)
+proba = model.predict_proba(X_test)
+
+# FTRL-Proximal with L1 for sparse linear weights (classic CTR setup)
+sparse = FMClassifier(optimizer="ftrl", l1_linear=1.0, batch_size=256, random_state=42)
+sparse.fit(X_train, y_train)
+
+ffm = FFMClassifier(n_factors=8, n_jobs=-1, random_state=42)
+ffm.fit(X_train, y_train, field_ids=field_ids)
+```
+
+`FMRegressor`, multiclass `FMClassifier` (just pass a target with >2 classes),
+early stopping (`early_stopping=True` or `eval_set=(X_val, y_val)`), and the
+`CategoricalEncoder` are demonstrated in `examples/basic_usage.py`.
+`benchmarks/bench_synthetic.py` reports fit time and predict throughput against
+the NumPy reference floor.
+
+## Development
+
+Requires Python >= 3.10 and a recent Rust toolchain (1.74+; `rustup update`).
+
+```bash
+python3 -m venv .venv
+.venv/bin/pip install -e ".[dev]"   # builds the Rust extension via maturin
+.venv/bin/pytest -q
+.venv/bin/ruff check .
+```
+
+`pip install -e .` compiles `rust/` and installs the extension as
+`modern_fm._rust` (maturin mixed layout, config in `pyproject.toml`).
+After editing Rust code, re-run `pip install -e .` to rebuild. Rust-only
+checks:
+
+```bash
+cd rust
+PYO3_PYTHON=$PWD/../.venv/bin/python3 cargo test
+PYO3_PYTHON=$PWD/../.venv/bin/python3 cargo clippy
+```
+
+Without the extension built, the package still works: `modern_fm._backend`
+falls back to the pure-NumPy reference implementations, and the parity tests
+in `tests/test_rust_parity.py` are skipped.
+
+Design documents live in `docs/` — start with `docs/requirements.md` and
+`docs/math_spec.md`. The roadmap is in `docs/roadmap.md`.
+

modern_fm-0.2.1/README.md

@@ -0,0 +1,86 @@
+# modern_fm
+
+Fast, sklearn-compatible Factorization Machines (FM) and Field-aware
+Factorization Machines (FFM) for Python.
+
+**Status: v0.2 (Beta).** A Rust CPU backend (parity-tested against pure-NumPy
+reference implementations) drives sklearn-style estimators — `FMClassifier`
+(binary + multiclass softmax), `FMRegressor`, and `FFMClassifier`
+(binary + multiclass softmax) — with the SGD / AdaGrad / Adam /
+**FTRL-Proximal** optimizers, **mini-batch**
+gradient averaging (`batch_size`), **multi-core training** via `rayon`
+(`n_jobs`), plus `sample_weight`/`class_weight`, `label_smoothing`, early
+stopping, a `CategoricalEncoder`, and `save_model`/`load_model`. FTRL's L1
+(`l1_linear`/`l1_factors`) yields exact-zero weights. See `docs/roadmap.md` for
+remaining niche gaps (FTRL + early stopping, multiclass + early-stopping for FFM).
+
+## Installation
+
+```bash
+pip install modern-fm        # once published; prebuilt wheels, no Rust needed
+```
+
+Until the first PyPI release, install from source (requires a Rust toolchain;
+see Development below).
+
+## Usage
+
+```python
+from modern_fm import FMClassifier, FFMClassifier
+
+model = FMClassifier(
+    n_factors=16,
+    optimizer="adagrad",
+    learning_rate=0.05,
+    max_iter=100,
+    batch_size=256,        # mini-batch gradient averaging (1 = per-row SGD)
+    n_jobs=-1,             # train batches across all CPU cores
+    l2_linear=1e-5,
+    l2_factors=1e-5,
+    random_state=42,
+)
+model.fit(X_train, y_train)
+proba = model.predict_proba(X_test)
+
+# FTRL-Proximal with L1 for sparse linear weights (classic CTR setup)
+sparse = FMClassifier(optimizer="ftrl", l1_linear=1.0, batch_size=256, random_state=42)
+sparse.fit(X_train, y_train)
+
+ffm = FFMClassifier(n_factors=8, n_jobs=-1, random_state=42)
+ffm.fit(X_train, y_train, field_ids=field_ids)
+```
+
+`FMRegressor`, multiclass `FMClassifier` (just pass a target with >2 classes),
+early stopping (`early_stopping=True` or `eval_set=(X_val, y_val)`), and the
+`CategoricalEncoder` are demonstrated in `examples/basic_usage.py`.
+`benchmarks/bench_synthetic.py` reports fit time and predict throughput against
+the NumPy reference floor.
+
+## Development
+
+Requires Python >= 3.10 and a recent Rust toolchain (1.74+; `rustup update`).
+
+```bash
+python3 -m venv .venv
+.venv/bin/pip install -e ".[dev]"   # builds the Rust extension via maturin
+.venv/bin/pytest -q
+.venv/bin/ruff check .
+```
+
+`pip install -e .` compiles `rust/` and installs the extension as
+`modern_fm._rust` (maturin mixed layout, config in `pyproject.toml`).
+After editing Rust code, re-run `pip install -e .` to rebuild. Rust-only
+checks:
+
+```bash
+cd rust
+PYO3_PYTHON=$PWD/../.venv/bin/python3 cargo test
+PYO3_PYTHON=$PWD/../.venv/bin/python3 cargo clippy
+```
+
+Without the extension built, the package still works: `modern_fm._backend`
+falls back to the pure-NumPy reference implementations, and the parity tests
+in `tests/test_rust_parity.py` are skipped.
+
+Design documents live in `docs/` — start with `docs/requirements.md` and
+`docs/math_spec.md`. The roadmap is in `docs/roadmap.md`.

modern_fm-0.2.1/pyproject.toml

@@ -0,0 +1,55 @@
+[build-system]
+requires = ["maturin>=1.7,<2"]
+build-backend = "maturin"
+
+[project]
+name = "modern-fm"
+version = "0.2.1"
+description = "Fast, sklearn-compatible Factorization Machines and Field-aware Factorization Machines"
+readme = "README.md"
+requires-python = ">=3.10"
+license = { text = "MIT" }
+authors = [{ name = "Masaya Kawamata" }]
+keywords = ["factorization-machines", "ffm", "ctr", "recommender", "tabular"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Programming Language :: Rust",
+    "Operating System :: OS Independent",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "numpy>=1.24",
+    "scipy>=1.10",
+]
+
+[project.optional-dependencies]
+dev = ["pytest>=7", "ruff>=0.4"]
+
+[project.urls]
+Homepage = "https://github.com/Matapanino/modern_fm"
+Repository = "https://github.com/Matapanino/modern_fm"
+Issues = "https://github.com/Matapanino/modern_fm/issues"
+
+[tool.maturin]
+manifest-path = "rust/Cargo.toml"
+python-source = "python"
+module-name = "modern_fm._rust"
+features = ["pyo3/extension-module"]
+include = [{ path = "LICENSE", format = "sdist" }]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[tool.ruff]
+line-length = 100
+target-version = "py310"
+
+[tool.ruff.lint]
+select = ["E", "F", "W", "I", "NPY", "UP"]

modern_fm-0.2.1/python/modern_fm/__init__.py

@@ -0,0 +1,17 @@
+"""modern_fm: fast, sklearn-compatible FM / FFM for Python."""
+
+from ._base import NotFittedError
+from .ffm import FFMClassifier
+from .fm import FMClassifier, FMRegressor
+from .preprocessing import CategoricalEncoder
+
+__version__ = "0.2.1"
+
+__all__ = [
+    "FMClassifier",
+    "FMRegressor",
+    "FFMClassifier",
+    "CategoricalEncoder",
+    "NotFittedError",
+    "__version__",
+]

modern_fm-0.2.1/python/modern_fm/_backend.py

@@ -0,0 +1,259 @@
+"""Backend dispatch: Rust extension when built, NumPy reference otherwise.
+
+Private module. The NumPy implementations in `_reference` remain the ground
+truth; the Rust extension (`modern_fm._rust`, built via maturin) is an
+optimized drop-in whose parity is enforced by tests/test_rust_parity.py.
+
+Both prediction and training are dispatched here (FM/FFM predict, FM binary and
+multiclass-softmax training, FFM training); training parity with the reference
+trainers is enforced by tests/test_rust_train_parity.py.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import scipy.sparse as sp
+
+from . import _reference, _reference_train
+
+try:
+    from . import _rust
+except ImportError:  # extension not built — pure-Python install
+    _rust = None
+
+
+def has_rust():
+    return _rust is not None
+
+
+def _prep_dense(X):
+    return np.ascontiguousarray(X, dtype=np.float64)
+
+
+def _prep_vec(a, dtype=np.float64):
+    return np.ascontiguousarray(a, dtype=dtype)
+
+
+def _prep_csr(X):
+    X = X.tocsr().astype(np.float64)
+    X.sum_duplicates()
+    return (
+        np.ascontiguousarray(X.indptr, dtype=np.int64),
+        np.ascontiguousarray(X.indices, dtype=np.int64),
+        np.ascontiguousarray(X.data, dtype=np.float64),
+        X.shape[1],
+    )
+
+
+def fm_predict_fast(X, w0, w, V):
+    """FM prediction (math_spec.md), Rust-accelerated when available."""
+    if _rust is None:
+        return _reference.fm_predict_fast(X, w0, w, V)
+    w = _prep_vec(w)
+    V = _prep_dense(V)
+    if sp.issparse(X):
+        indptr, indices, data, n_features = _prep_csr(X)
+        return _rust.fm_predict_fast_csr(indptr, indices, data, n_features, float(w0), w, V)
+    return _rust.fm_predict_fast_dense(_prep_dense(X), float(w0), w, V)
+
+
+def ffm_predict(X, field_ids, w0, w, V):
+    """FFM prediction (math_spec.md), Rust-accelerated when available."""
+    if _rust is None:
+        return _reference.ffm_predict(X, field_ids, w0, w, V)
+    field_ids = _prep_vec(field_ids, dtype=np.int64)
+    w = _prep_vec(w)
+    V = _prep_dense(V)
+    if sp.issparse(X):
+        indptr, indices, data, n_features = _prep_csr(X)
+        return _rust.ffm_predict_csr(
+            indptr, indices, data, n_features, field_ids, float(w0), w, V
+        )
+    return _rust.ffm_predict_dense(_prep_dense(X), field_ids, float(w0), w, V)
+
+
+def _prep_fit(X, y, params, row_orders):
+    """Common coercion for the Rust fit entry points.
+
+    Dense X is converted to CSR (exact zeros are skipped either way, matching
+    the reference). Returns fresh float64 copies of w and V that the Rust
+    kernel mutates in place; the caller's `params` are left untouched.
+    """
+    w0, w, V = params
+    w = np.array(w, dtype=np.float64, order="C", copy=True)
+    V = np.array(V, dtype=np.float64, order="C", copy=True)
+    y = _prep_vec(y)
+    row_orders = np.ascontiguousarray(row_orders, dtype=np.int64)
+    if row_orders.ndim == 1:
+        row_orders = row_orders[None, :]
+    Xc = X if sp.issparse(X) else sp.csr_matrix(np.asarray(X, dtype=np.float64))
+    return _prep_csr(Xc), y, float(w0), w, V, row_orders
+
+
+def _acc_arrays(state, w, V):
+    """AdaGrad accumulators (acc_w0, acc_w, acc_v) from `state` or fresh zeros.
+
+    `state` (a mutable list, for the epoch-driven early-stopping path) persists
+    the accumulators across calls; None means a single all-epochs run.
+    """
+    if state is None:
+        return 0.0, np.zeros(len(w)), np.zeros_like(V)
+    acc_w0, acc_w, acc_v = state
+    return float(acc_w0), _prep_vec(acc_w), _prep_dense(acc_v)
+
+
+def fm_fit(
+    X, y, params, *, loss, optimizer, learning_rate, l2_linear, l2_factors, row_orders,
+    l1_linear=0.0, l1_factors=0.0, beta_1=0.9, beta_2=0.999, epsilon=1e-8, ftrl_beta=1.0,
+    batch_size=1, n_jobs=1, sample_weight=None, state=None, adam_state=None,
+):
+    """Train an FM (docs/optimization_spec.md).
+
+    `params` = (w0, w, V) initial values (unchanged); returns new float64
+    (w0, w, V). `sample_weight` scales each row's gradient (None -> all ones).
+    `batch_size` averages each batch's gradient (batch_size=1 is per-row).
+    `n_jobs` (>= 1) splits each batch across that many rayon threads; n_jobs=1
+    is the serial path matching the reference. `beta_1`/`beta_2`/`epsilon` apply
+    only when optimizer == "adam". `state` carries AdaGrad accumulators in/out
+    for epoch-by-epoch training. `adam_state` does the same for Adam moments;
+    because the Rust kernel keeps Adam state internal, Adam + early stopping
+    (adam_state given) runs on the NumPy reference path. The reference fallback
+    is always serial (it is the n_jobs=1 ground truth).
+    """
+    if _rust is None or adam_state is not None:
+        return _reference_train.fm_fit_reference(
+            X, y, params, loss=loss, optimizer=optimizer, learning_rate=learning_rate,
+            l2_linear=l2_linear, l2_factors=l2_factors, l1_linear=l1_linear,
+            l1_factors=l1_factors, row_orders=row_orders, beta_1=beta_1, beta_2=beta_2,
+            epsilon=epsilon, ftrl_beta=ftrl_beta, batch_size=batch_size,
+            sample_weight=sample_weight, state=state, adam_state=adam_state,
+        )
+    (indptr, indices, data, n_features), y, w0, w, V, row_orders = _prep_fit(
+        X, y, params, row_orders
+    )
+    sw = np.ones(len(y)) if sample_weight is None else _prep_vec(sample_weight)
+    acc_w0, acc_w, acc_v = _acc_arrays(state, w, V)
+    w0, acc_w0 = _rust.fm_fit_csr(
+        indptr, indices, data, n_features, y, sw, w0, acc_w0, w, V, acc_w, acc_v,
+        loss, optimizer, learning_rate, l2_linear, l2_factors, beta_1, beta_2, epsilon,
+        row_orders, batch_size, n_jobs, l1_linear, l1_factors, ftrl_beta,
+    )
+    if state is not None:
+        state[0], state[1], state[2] = acc_w0, acc_w, acc_v
+    return w0, w, V
+
+
+def fm_fit_multiclass(
+    X, y, params, *, optimizer, learning_rate, l2_linear, l2_factors, row_orders,
+    label_smoothing=0.0, l1_linear=0.0, l1_factors=0.0, beta_1=0.9, beta_2=0.999, epsilon=1e-8,
+    ftrl_beta=1.0, batch_size=1, sample_weight=None, state=None, adam_state=None,
+):
+    """Train a multiclass (softmax) FM (optimization_spec.md).
+
+    `params` = (w0 (C,), w (C, n), V (C, n, k)) initial values (unchanged);
+    `y` holds integer class indices in [0, C). Returns new float64 (w0, w, V).
+    `batch_size` averages each batch's gradient (batch_size=1 is per-row).
+    `beta_1`/`beta_2`/`epsilon` apply only when optimizer == "adam". The Rust
+    multiclass kernel keeps optimizer state internal, so early stopping (`state`
+    or `adam_state` given) round-trips it across epochs on the NumPy reference
+    path. Otherwise Rust-accelerated when available, reference fallback otherwise.
+    """
+    if _rust is None or state is not None or adam_state is not None:
+        return _reference_train.fm_fit_multiclass_reference(
+            X, y, params, optimizer=optimizer, learning_rate=learning_rate,
+            l2_linear=l2_linear, l2_factors=l2_factors, l1_linear=l1_linear,
+            l1_factors=l1_factors, row_orders=row_orders, label_smoothing=label_smoothing,
+            beta_1=beta_1, beta_2=beta_2, epsilon=epsilon, ftrl_beta=ftrl_beta,
+            batch_size=batch_size, sample_weight=sample_weight, state=state,
+            adam_state=adam_state,
+        )
+    w0, w, V = params
+    w0 = np.array(w0, dtype=np.float64, order="C", copy=True)  # (C,), mutated in place
+    w = np.array(w, dtype=np.float64, order="C", copy=True)  # (C, n)
+    V = np.array(V, dtype=np.float64, order="C", copy=True)  # (C, n, k)
+    y = _prep_vec(y, dtype=np.int64)
+    row_orders = np.ascontiguousarray(row_orders, dtype=np.int64)
+    if row_orders.ndim == 1:
+        row_orders = row_orders[None, :]
+    Xc = X if sp.issparse(X) else sp.csr_matrix(np.asarray(X, dtype=np.float64))
+    indptr, indices, data, n_features = _prep_csr(Xc)
+    sw = np.ones(len(y)) if sample_weight is None else _prep_vec(sample_weight)
+    _rust.fm_fit_multiclass_csr(
+        indptr, indices, data, n_features, y, sw, w0, w, V,
+        optimizer, learning_rate, l2_linear, l2_factors, label_smoothing,
+        beta_1, beta_2, epsilon, row_orders, batch_size, l1_linear, l1_factors, ftrl_beta,
+    )
+    return w0, w, V
+
+
+def ffm_fit(
+    X, y, field_ids, params, *, optimizer, learning_rate, l2_linear, l2_factors, row_orders,
+    l1_linear=0.0, l1_factors=0.0, beta_1=0.9, beta_2=0.999, epsilon=1e-8, ftrl_beta=1.0,
+    batch_size=1, n_jobs=1, sample_weight=None, state=None, adam_state=None,
+):
+    """Train an FFM (logistic loss); see fm_fit. `batch_size` averages each
+    batch's gradient (batch_size=1 is per-row); `n_jobs` (>= 1) splits each batch
+    across that many rayon threads (n_jobs=1 matches the serial reference).
+    Adam + early stopping (adam_state given) runs on the NumPy reference path,
+    like fm_fit."""
+    if _rust is None or adam_state is not None:
+        return _reference_train.ffm_fit_reference(
+            X, y, field_ids, params, optimizer=optimizer, learning_rate=learning_rate,
+            l2_linear=l2_linear, l2_factors=l2_factors, l1_linear=l1_linear,
+            l1_factors=l1_factors, row_orders=row_orders, beta_1=beta_1, beta_2=beta_2,
+            epsilon=epsilon, ftrl_beta=ftrl_beta, batch_size=batch_size,
+            sample_weight=sample_weight, state=state, adam_state=adam_state,
+        )
+    field_ids = _prep_vec(field_ids, dtype=np.int64)
+    (indptr, indices, data, n_features), y, w0, w, V, row_orders = _prep_fit(
+        X, y, params, row_orders
+    )
+    sw = np.ones(len(y)) if sample_weight is None else _prep_vec(sample_weight)
+    acc_w0, acc_w, acc_v = _acc_arrays(state, w, V)
+    w0, acc_w0 = _rust.ffm_fit_csr(
+        indptr, indices, data, n_features, y, sw, field_ids, w0, acc_w0, w, V, acc_w, acc_v,
+        optimizer, learning_rate, l2_linear, l2_factors, beta_1, beta_2, epsilon, row_orders,
+        batch_size, n_jobs, l1_linear, l1_factors, ftrl_beta,
+    )
+    if state is not None:
+        state[0], state[1], state[2] = acc_w0, acc_w, acc_v
+    return w0, w, V
+
+
+def ffm_fit_multiclass(
+    X, y, field_ids, params, *, optimizer, learning_rate, l2_linear, l2_factors, row_orders,
+    label_smoothing=0.0, l1_linear=0.0, l1_factors=0.0, beta_1=0.9, beta_2=0.999, epsilon=1e-8,
+    ftrl_beta=1.0, batch_size=1, sample_weight=None,
+):
+    """Train a multiclass (softmax) FFM (one FFM per class, coupled by softmax).
+
+    `params` = (w0 (C,), w (C, n), V (C, n, n_fields, k)); `y` holds class indices
+    in [0, C). Serial (no n_jobs), like FM multiclass. Rust-accelerated when
+    available, NumPy reference fallback otherwise.
+    """
+    if _rust is None:
+        return _reference_train.ffm_fit_multiclass_reference(
+            X, y, field_ids, params, optimizer=optimizer, learning_rate=learning_rate,
+            l2_linear=l2_linear, l2_factors=l2_factors, l1_linear=l1_linear,
+            l1_factors=l1_factors, row_orders=row_orders, label_smoothing=label_smoothing,
+            beta_1=beta_1, beta_2=beta_2, epsilon=epsilon, ftrl_beta=ftrl_beta,
+            batch_size=batch_size, sample_weight=sample_weight,
+        )
+    w0, w, V = params
+    w0 = np.array(w0, dtype=np.float64, order="C", copy=True)  # (C,), mutated in place
+    w = np.array(w, dtype=np.float64, order="C", copy=True)  # (C, n)
+    V = np.array(V, dtype=np.float64, order="C", copy=True)  # (C, n, n_fields, k)
+    y = _prep_vec(y, dtype=np.int64)
+    field_ids = _prep_vec(field_ids, dtype=np.int64)
+    row_orders = np.ascontiguousarray(row_orders, dtype=np.int64)
+    if row_orders.ndim == 1:
+        row_orders = row_orders[None, :]
+    Xc = X if sp.issparse(X) else sp.csr_matrix(np.asarray(X, dtype=np.float64))
+    indptr, indices, data, n_features = _prep_csr(Xc)
+    sw = np.ones(len(y)) if sample_weight is None else _prep_vec(sample_weight)
+    _rust.ffm_fit_multiclass_csr(
+        indptr, indices, data, n_features, y, sw, field_ids, w0, w, V,
+        optimizer, learning_rate, l2_linear, l2_factors, label_smoothing,
+        beta_1, beta_2, epsilon, row_orders, batch_size, l1_linear, l1_factors, ftrl_beta,
+    )
+    return w0, w, V

modern_fm-0.2.1/python/modern_fm/_base.py

@@ -0,0 +1,87 @@
+"""Minimal sklearn-compatible parameter handling.
+
+Avoids a hard scikit-learn dependency in v0.1 while keeping the contract:
+__init__ stores constructor arguments verbatim; get_params/set_params
+round-trip them. Phase 3 may swap this for sklearn's BaseEstimator.
+"""
+
+from __future__ import annotations
+
+import inspect
+import pickle
+
+
+class NotFittedError(ValueError, AttributeError):
+    """Raised when a predict-like method is called before fit.
+
+    Inherits ValueError and AttributeError to match sklearn's exception of
+    the same name, so generic sklearn-style error handling keeps working.
+    """
+
+
+def check_is_fitted(estimator, attribute="w0_"):
+    if not hasattr(estimator, attribute):
+        raise NotFittedError(
+            f"This {type(estimator).__name__} instance is not fitted yet; "
+            "call 'fit' before using this method."
+        )
+
+
+class ParamsMixin:
+    @classmethod
+    def _param_names(cls):
+        sig = inspect.signature(cls.__init__)
+        return [name for name in sig.parameters if name != "self"]
+
+    def get_params(self, deep=True):
+        return {name: getattr(self, name) for name in self._param_names()}
+
+    def set_params(self, **params):
+        valid = set(self._param_names())
+        for key, value in params.items():
+            if key not in valid:
+                raise ValueError(
+                    f"Invalid parameter {key!r} for estimator {type(self).__name__}. "
+                    f"Valid parameters are: {sorted(valid)}."
+                )
+            setattr(self, key, value)
+        return self
+
+    def __repr__(self):
+        args = ", ".join(f"{k}={v!r}" for k, v in self.get_params().items())
+        return f"{type(self).__name__}({args})"
+
+
+class ModelIOMixin:
+    """save_model / load_model for fitted estimators.
+
+    Stores `{format_version, class, params, attrs}` via pickle, where `attrs`
+    are the learned trailing-underscore attributes; this is generic over the
+    estimator (binary, regression, multiclass) and round-trips constructor
+    params too. The estimators also pickle natively (plain attributes), so
+    `pickle.dumps(model)` works as an alternative.
+    """
+
+    _IO_VERSION = 1
+
+    def save_model(self, path):
+        check_is_fitted(self)
+        state = {
+            "format_version": self._IO_VERSION,
+            "class": type(self).__name__,
+            "params": self.get_params(),
+            "attrs": {k: getattr(self, k) for k in vars(self) if k.endswith("_")},
+        }
+        with open(path, "wb") as f:
+            pickle.dump(state, f)
+
+    @classmethod
+    def load_model(cls, path):
+        with open(path, "rb") as f:
+            state = pickle.load(f)
+        if state.get("class") != cls.__name__:
+            raise ValueError(f"{path!r} holds a {state.get('class')!r}, not a {cls.__name__!r}")
+        model = cls(**state["params"])
+        for key, value in state["attrs"].items():
+            setattr(model, key, value)
+        return model