amaye15/tabfm-gguf

tabfm-rs

Pure Rust converter and inference engine for google/tabfm-1.0.0-pytorch, a tabular foundation model (in-context learner, TabPFN-style) — not a time-series forecaster, unlike every other model in this repo. TabFM does zero-shot classification (up to 10 classes) and regression on structured tables with mixed numeric/categorical columns, with no fine-tuning.

Pre-converted GGUF files are available at amaye15/tabfm-gguf.

License note: weights are tabfm-non-commercial-v1.0 (non-commercial use only). Source code is Apache-2.0 via google-research/tabfm.

Scope

This converts the neural network (TabFM.forward) faithfully — checkpoint-loaded RoPE frequencies, the Fourier cell embedding, the dual column/row attention stages, and the in-context-learning stage all match the PyTorch reference to float32 precision (see Infer below for the low-level, single-forward-pass entry point).

On top of that, ensemble-predict (below) natively re-implements the HuggingFace TabFMClassifier/TabFMRegressor sklearn wrapper's preprocessing and ensembling: 3-stage feature scaling (all 5 normalization methods), categorical encoding, n_estimators-member ensembling with a bit-compatible port of CPython's random.Random (Mersenne Twister) so ensemble composition matches the real wrapper exactly for the same random_state, and the opt-in output-calibration (Platt/vector scaling) and NNLS ensemble-weighting paths, fit natively via an in-fold procedure — no Python or scipy involved at inference time.

Precision, two tiers: the default (always-on) ensembling + feature-scaling path matches the real TabFMClassifier/TabFMRegressor to ~1e-4–1e-6 max abs error (validated via scripts/compare_python.py --ensemble). The opt-in calibration/NNLS path (off by default in the wrapper itself) currently matches to ~0.02–0.09 max abs error — looser, because its out-of-fold cross-validation splits rows via our own seeded shuffle rather than sklearn's KFold(shuffle=True) (which draws from NumPy's legacy RandomState, a related but distinct, unported RNG family), so the two sides fit calibration/NNLS on slightly different fold memberships. This is a known, documented gap, not a correctness bug — the underlying NNLS solver is separately validated bit-exact against scipy.optimize.nnls.

Batch scope is B=1 per call (one table per infer/ensemble-predict invocation).

Build

cargo build --release

Convert

Downloads one variant (classification or regression) from HuggingFace and writes a GGUF file:

# F16 (recommended — good precision/size trade-off)
./target/release/tabfm-rs convert --task classification --dtype f16 --output gguf/tabfm-classification-f16.gguf
./target/release/tabfm-rs convert --task regression     --dtype f16 --output gguf/tabfm-regression-f16.gguf

To convert both variants, all three dtypes at once:

./scripts/convert_all.sh

HuggingFace token (optional for public models):

HF_TOKEN=hf_... ./scripts/convert_all.sh

Inspect tensors

./target/release/tabfm-rs inspect-tensors models/tabfm-classification/classification_model.safetensors

Infer

x is [T][H]: rows (training rows first, then query rows) x padded feature columns, all pre-encoded to floats (categorical columns ordinal/label-encoded by the caller). y is [T] labels; any finite placeholder at query-row positions is fine (masked internally). train_size is how many leading rows are training rows. cat_mask (optional, default all-false) marks which columns are categorical. d (optional, default H) is the actual unpadded feature count.

echo '{"x": [[0,1.1,-0.3],[1,0.4,0.9],[0,-1.2,0.2],[1,2.0,-0.5]], "y": [0,1,0,0], "train_size": 3, "cat_mask": [true,false,false], "d": 3}' \
  | ./target/release/tabfm-rs infer \
      --gguf gguf/tabfm-classification-f16.gguf \
      --config models/tabfm-classification/classification_config.json

Classification output includes raw logits and a plain softmax (no temperature/calibration):

{
  "task": "classification",
  "logits": [[...], [...], [...], [...]],
  "probabilities": [[...], [...], [...], [...]]
}

Regression output is the model's raw decoded scalar per row (no target un-normalization — that's part of the sklearn wrapper's preprocessing, replicated by ensemble-predict below, not this low-level command):

{ "task": "regression", "predictions": [0.12, -0.44, 0.03, 0.88], "raw": [[...], ...] }

Only rows >= train_size are meaningful predictions.

Ensemble predict

The full sklearn-wrapper-equivalent entry point: takes raw (unencoded, unscaled) train/test tables and handles categorical encoding, feature scaling, and ensembling internally — the counterpart to TabFMClassifier(model).fit(X_train, y_train).predict_proba(X_test) / TabFMRegressor(...).predict(X_test).

echo '{"x_train": [[0.1,1.1,"red"],[0.4,0.9,"blue"],[-1.2,0.2,"red"],[2.0,-0.5,"blue"],[0.3,0.1,"red"],[1.1,-0.2,"blue"]],
       "y_train": [0,1,0,1,0,1],
       "x_test": [[0.2,0.9,"red"],[1.5,-0.4,"blue"]],
       "cat_mask": [false,false,true]}' \
  | ./target/release/tabfm-rs ensemble-predict \
      --gguf gguf/tabfm-classification-f16.gguf \
      --config models/tabfm-classification/classification_config.json

Defaults match the sklearn wrapper's own: n_estimators=32, norm_methods=["none","power"] (cycled per member), class_shift=true, outlier_threshold=4.0, softmax_temperature=0.9, average_logits=true, random_state=42. All 5 normalization methods are supported via norm_methods: "none", "power" (Yeo-Johnson), "quantile", "quantile_rtdl", "robust". Calibration (binary_calibration_method: "platt" for 2 classes, multiclass_calibration_method: "vector" for more) and enable_nnls are off by default, matching the wrapper; when turned on, they're fit natively via a num_folds_for_cv-fold (default 5) out-of-fold pass over x_train in the same call — no separate fit step.

{
  "task": "classification",
  "probabilities": [[0.1, 0.9], [0.8, 0.2]],
  "predicted_labels": ["1", "0"],
  "classes": ["0", "1"]
}

Regression output: {"task": "regression", "predictions": [0.42, -1.1]}.

Verify precision against the real sklearn wrapper

python3 scripts/compare_python.py --ensemble --task classification
python3 scripts/compare_python.py --ensemble --task regression
python3 scripts/compare_python.py --ensemble --task classification --norm-methods quantile,robust
python3 scripts/compare_python.py --ensemble --task classification --enable-nnls
python3 scripts/compare_python.py --ensemble --task classification --multiclass-calib vector

Python bindings

python -m venv .venv && source .venv/bin/activate
pip install maturin
maturin develop --features python

import tabfm_rs

model = tabfm_rs.TabFM("gguf/tabfm-classification-f16.gguf", "models/tabfm-classification/classification_config.json")

# Low-level, single-forward-pass API (pre-encoded/scaled x) — see Infer above.
out = model.predict(x, y, train_size=3, cat_mask=[True, False, False], d=3)

# Full sklearn-wrapper-equivalent pipeline (raw x, scaling + encoding + ensembling handled
# internally) — see Ensemble predict above. Same defaults as the CLI's ensemble-predict command.
result = model.ensemble_predict(
    x_train=[[0.1, 1.1, "red"], [0.4, 0.9, "blue"], [-1.2, 0.2, "red"],
             [2.0, -0.5, "blue"], [0.3, 0.1, "red"], [1.1, -0.2, "blue"]],
    y_train=[0, 1, 0, 1, 0, 1],
    x_test=[[0.2, 0.9, "red"], [1.5, -0.4, "blue"]],
    cat_mask=[False, False, True],
)
# {"task": "classification", "probabilities": [[...], [...]],
#  "predicted_labels": [...], "classes": [...]}

ensemble_predict accepts the same keyword arguments as the CLI's stdin JSON (n_estimators, norm_methods, class_shift, outlier_threshold, softmax_temperature, average_logits, random_state, binary_calibration_method, multiclass_calibration_method, num_folds_for_cv, enable_nnls, nnls_beta, calibration_lambda), all with the same defaults.

Architecture notes

• CellEmbedder: each cell's value is grouped with feature_group_size neighboring columns (index wraparound (h + 2^i - 1) % d), Fourier-expanded (sin/cos in float32), linearly projected, and summed over the group — plus a learned embedding of the row's label, added only at training-row positions.
• ColEmbedding: a SetTransformer (induced attention) mixes information across rows, per column, masked so only training rows are attendable keys.
• RowInteraction: self-attention across columns (with checkpoint-loaded RoPE), masked to the unpadded feature columns plus a prepended block of learned CLS tokens; the second pass collapses to just the CLS-token slice.
• ICLearning: 24 self-attention blocks over rows, re-injecting each row's label at training-row positions, masked so only training rows are attendable keys — the actual "in-context learning" step — followed by an MLP decoder to per-class logits or a scalar.
• Attention: q is pre-scaled by a learned per-dimension softplus'd scale before the QK matmul; attention itself runs at scale=1.0 (no additional 1/√d). RoPE (where used) is the interleaved-pair variant, checkpoint-loaded frequencies (never recomputed from a formula).
• Norms: RMSNorm throughout (float32), four per attention+FFN block (pre/post each).
• FFN: SwiGLU everywhere.