autobooster

AutoML for tabular data with a focus on massive feature engineering. You tag the columns by type; the library generates thousands of features (everything combined with everything), filters out noise with LightGBM against shadow features, trains a zoo of models, builds an ensemble and saves the result as a single file.

pip install autobooster-ml     # everything: LightGBM, CatBoost, XGBoost, torch, transformers, Prophet, SHAP, ...
autobooster download           # optional: fetch the pretrained models now (~2.5 GB) instead of on first use

One install brings every dependency (on Linux, PyPI's torch wheel includes CUDA libraries and is several GB; for a CPU-only machine install torch first: pip install torch --index-url https://download.pytorch.org/whl/cpu). Pretrained models (text and image encoders, RAG retriever, the default LLM) are downloaded from HuggingFace on first use, or all at once with autobooster download / ab.download_models() (--only text image rag llm picks groups). The import name is always import autobooster.

import autobooster as ab

m = ab.AutoBooster(
    task="auto",                      # regression|binary|multiclass|multilabel|ranking|forecast|anomaly|clustering
    target="churn",
    columns={                         # anything not listed gets an inferred role
        "numeric": ["age", "balance"],
        "categorical": ["city", "tariff"],
        "text": ["last_ticket"],      # multilingual embeddings
        "datetime": ["signup_at"],
        "sequential": ["events"],     # list cells: ["view", "cart", "buy"] or [1.2, 3.4, ...]
        "image": ["photo"],           # path, bytes, PIL.Image, numpy array or {"bytes": ...} from HF datasets
    },
    time_budget="30m", top_k=0.2,
)
m.fit(df)
m.leaderboard()                       # CV and holdout metrics of every model and of the ensemble
m.feature_report()                    # selected features, importance, which generator produced them
m.explain(df)                         # mean |SHAP|
m.report("report.html", df)
m.predict(new_df); m.predict_proba(new_df)
m.transform(new_df)                   # the selected features, to feed your own models
m.save("churn.ab"); m = ab.load("churn.ab")

What happens inside fit

1. Schema and task. Column roles come from columns or are inferred; the task is inferred from the target. The preset adapts to the number of rows (folds, how many features to mix, SVM mode, kNN, thread count).
2. Holdout. 12% of the data is set aside first: the latest rows for temporal data, whole groups for group_col and ranking. The final honest score is computed on it.
3. Generation, stage 1 (from raw columns):
   • numbers: raw, signed log, rank, fractional part, roundness, NaN flags, row-wise z-statistics;
   • categories: count, frequency code, one-hot (up to 12 levels), out-of-fold target encoding (per class for multiclass), crosses of category pairs;
   • dates: hour, day of week, month, day of year, week, weekend, month start/end, sin/cos, holidays (holidays_country="US"), differences between dates;
   • text: paraphrase-multilingual-MiniLM-L12-v2 embeddings → PCA, mean-pool across text columns and cosine similarity between them, char TF-IDF (2–4 characters) → SVD, word TF-IDF (words and bigrams) → SVD, NMF topics (10 topics, the share of each in the document), an out-of-fold linear model on the embeddings, on char TF-IDF and on word TF-IDF, surface statistics. Embeddings are cached in ~/.cache/autobooster;
   • sequential cells: for numeric sequences the shape of the series (quantiles, slope, tail, autocorrelation, mean crossings, FFT: peak frequency, band energy, spectral entropy); for tokens n-gram TF-IDF → SVD, item2vec (Word2Vec), an out-of-fold linear model, and for short fixed-length sequences a token per position (DNA, event codes);
   • images: DINOv2-small embeddings (CLS + mean-pooled patches, CPU) → PCA, an out-of-fold linear model (linear probe), size, brightness, contrast, mean colours. Plus classic scikit-image descriptors: GLCM texture (contrast, homogeneity, energy, correlation), LBP histograms, Canny edge density, sharpness (variance of the Laplacian), entropy, saturation and colourfulness as separate features, and HOG + LBP + HSV histogram (1920 numbers) → PCA with its own out-of-fold linear model (image_descriptors=False turns them off). Everything is cached by image content;
   • recommender pairs (user × item): detected automatically among high-cardinality categories or given via mf_pairs=[("user_id", "item_id")]. Produces out-of-fold ALS matrix factorisation with biases (prediction, user bias, item bias, dot product) and SVD embeddings of the interaction matrix;
   • time_col (+group_col): lags, rolling mean/std/min/max, ewm, time since the previous event within the series. For forecast, target lags shifted by the horizon are added. The history tail is stored in the model, so predict on future rows continues the windows;
   • Prophet: yhat, trend, weekly, yearly, daily and seasonal per series (the largest series get their own model, the rest share one). On train the values are strictly out-of-fold in time. In regression and forecasting Prophet also takes part in the ensemble as a separate model;
   • geo=[("lat","lon")]: rotated coordinates, clusters, distance to centres, grid density;
   • fit(df, related={"orders": (orders_df, "user_id")}): aggregates of child tables.
4. Generation, stage 2. A quick LightGBM ranks the stage-1 features. The top N are mixed pairwise (+ − × ÷), and for groups by category the mean, std, deviation from the group mean and z-score are computed. The number of features is capped by the memory budget: rows × features × 4B ≤ memory_budget_gb.
5. Selection. Filters first: empty, constant, duplicate, |corr| > 0.98. Then LightGBM with about 500 trees, early stopping on folds and gain importance against shuffled shadow copies: features stronger than the 95th percentile of the noise survive, and top_k of them are kept. With more than 200k rows, generation and selection run on a sample, after which the pipeline is pruned to what is needed and recomputed on all rows.
6. Models. LightGBM, CatBoost (native categories), XGBoost, Ridge/Logistic, ElasticNet (regression), MLP (torch), ExtraTrees, RandomForest, SVM (exact RBF on small folds, Nyström on large ones) and kNN. Each model gets its own Optuna tuning on one fold within its share of the budget, then CV with out-of-fold predictions. Ranking uses lambdarank, YetiRank and rank:ndcg.
7. Ensemble. The best single model, greedy Caruana selection and stacking (Ridge/Logistic on top of the OOF predictions) compete; the one with the best OOF score wins, and the complex ones must win by a margin.
8. Holdout score, then refit of the pipeline and the ensemble members on all data with the iteration counts found in CV.

A heavily skewed non-negative regression target is automatically learned as log1p, while metrics are reported on the original scale. With metric="rmsle" models always learn log1p(y).

Missing values

Before the schema is built, missing-value markers are turned into NaN. By default (na_values="auto") these are "", "?", "NA", "N/A", "null", "NaN", "#N/A" and similar. You can pass your own: a list of strings and/or numbers (na_values=["?", -999]), a per-column dict ({"age": [0], "city": ["unknown"]}) or None. A column that was numeric but was read as strings because of "?" becomes numeric again.

Then, with imputation="model" (the default), every numeric or categorical column with gaps gets its own LightGBM. The model predicts the column from the other columns; gaps in the predictors themselves are handled natively by LightGBM. The model is validated on 20% of the known values and kept only if it beats the median or mode (2% lower MAE, 0.5 pp higher accuracy); otherwise the median or mode is used. A <col>__was_missing flag is added alongside, so the fact that a value was absent is not lost ("no garage" and "unknown" are different things). Columns with more than 60% missing and categories with more than 100 levels are not filled; their NaN is left to the models. The target and cv_group are never used as predictors, so the same models work at predict time.

m = ab.AutoBooster("binary", "y", na_values=["?", ""]).fit(df)
m.imputation_report()          # column, how many gaps, method, model MAE/accuracy vs median/mode
clean = m.impute(df)           # the original dataframe with the gaps filled
clean = ab.Imputer(na_values=["?"], exclude=["y"]).fit_transform(df)   # the same without training task models

imputation="simple" fills with the median or mode, imputation=None only turns markers into NaN.

Validation

Folds are chosen automatically: stratified for classes, GroupKFold by group_col (the query for ranking), an expanding window by time_col. If the test set consists of new groups (other days, customers, shops), pass cv_group="day". Each group then lands entirely in either train or validation/holdout, and the column itself is not used as a feature. Otherwise feature selection and the ensemble adapt to an interpolation that will not exist on the test set.

If you know the rows you will predict on (a Kaggle test set, for example), pass them without the target: m.fit(train, test_df=test). This is adversarial validation: a classifier learns to tell train from test, and features that differ between them (e.g. "day of month" when the test covers other days) are dropped before selection. The test target is never used.

Task types

task	what it needs	default metric
regression	target	rmse
binary / multiclass	target (any labels)	auc / mlogloss
multilabel	target=[...] 0/1 columns	mean AUC
ranking	target relevance, group_col = query id	ndcg@10
forecast	target, time_col, optional group_col, horizon	rmse
anomaly	no target (if given, AUC is computed against it)	—
clustering	no target, optional n_clusters; works on features of any type (text, images, sequences)	silhouette

horizon in forecast sets how many steps of the series ahead the forecast is made. Target lags are shifted by horizon, so rows further than the horizon from the end of the history have no lags. Set horizon to the longest forecast distance you need.

Presets and budget

preset="auto" is chosen by the number of rows: up to 30k, 30–200k and over 200k. "fast" is for rough runs, "best" mixes more features and runs two shadow rounds. Single fields can be overridden, for example preset_overrides={"models": ("lightgbm", "catboost"), "memory_budget_gb": 4, "tune": False}.

time_budget ("30m", "2h" or seconds) is the wall-clock limit for the whole fit. Feature selection and stacking have deadlines, boosters and the MLP stop on a deadline callback, models that are not reached are skipped, and if no time is left for regenerating features, the final refit reuses the train-fitted pipeline. Feature generation itself is not interrupted, so with heavy text or image columns and a tight budget the run can overshoot slightly.

Saving

save writes a single joblib file with a manifest of the python, numpy, pandas, sklearn, lightgbm, catboost, xgboost, torch and other library versions. load warns on version mismatches, and ab.load(path, strict=True) raises instead. The sentence-transformers network is not saved by default, only its name: on load it comes from the HuggingFace cache. bundle_text_model=True puts the weights into the file, which is what you need for offline inference.

CLI

autobooster fit train.parquet --target y --schema schema.yaml --budget 1h -o model.ab
autobooster predict model.ab test.parquet --proba -o pred.csv
autobooster info model.ab

schema.yaml holds lists of columns per role; any AutoBooster argument can be passed in its options: section.

There are also sklearn wrappers: autobooster.sklearn.AutoBoosterRegressor and AutoBoosterClassifier.

Simple RAG (question answering over documents)

ab.RAG answers questions about your documents with a local LLM. Run any OpenAI-compatible server: Ollama (default, ollama pull llama3.1:8b), vLLM, llama.cpp server or LM Studio.

rag = ab.RAG().add(["docs/", "faq.md", "manual.pdf", "any raw text"])   # files, folders, strings, dicts
ans = rag.ask("How do I reset the router?")
print(ans)                    # answer citing passages as [1], [2]
ans.cited()                   # the passages it cited: .text, .source, .chunk.start
rag.search("reset router", k=5)     # retrieval only
rag.save("kb.rag"); rag = ab.RAG.load("kb.rag")

• Loading: txt / md / rst / html / csv / json / py files, whole folders, PDFs (with pypdf), raw strings, {"text": ..., "source": ..., **meta} dicts or a DataFrame (add(df, text_col="body")).
• Chunking: paragraph-aware chunks of ~900 characters with ~150 characters of overlap; long paragraphs are split by sentences. Identical chunks are indexed once, so add can be called again with updated folders.
• Retrieval: hybrid. intfloat/multilingual-e5-small embeddings (a question in one language finds passages in another) plus BM25 with English and Russian stopwords, fused by reciprocal rank fusion. A multilingual cross-encoder (cross-encoder/mmarco-mMiniLMv2-L12-H384-v1) re-ranks the top 30 candidates; it fixes most cross-language misses (rerank_model=None turns it off). embed_model=None, rerank_model=None gives BM25 only, without torch.
• Generation: the prompt tells the model to answer only from the numbered passages, cite them and say "don't know" otherwise, in the language of the question. llm="auto" takes the first model the server lists; pass a model name, a different llm_url (vLLM: http://localhost:8000/v1), or any callable messages -> str.
• Speed: an 8B model on CPU processes the prompt at a few tokens per second, so the context size dominates latency. Lower k= / max_context_chars= or raise timeout= (600 s by default) on slow machines.

Frozen LLM + trained head (text classification / regression)

ab.LLMHead turns a small HuggingFace LLM into a text classifier or regressor without fine-tuning: the model is loaded without its next-token head, frozen, and its hidden states are pooled into one vector per text; only a small head is trained on top. It is opt-in: nothing in AutoBooster uses an LLM unless you ask for it.

clf = ab.LLMHead().fit(texts, labels)                 # spam / not spam, sentiment, topic, ...
clf.predict(["WIN a free iPhone, click here"])        # -> array(["spam"])
clf.predict_proba(new_texts); clf.score(test_texts, test_labels)
reg = ab.LLMHead(task="regression").fit(descriptions, prices)
clf.save("spam.llmh"); clf = ab.LLMHead.load("spam.llmh")   # a few MB: the LLM is reloaded by name
X = clf.embed(texts)                                   # the frozen features themselves, for any other model

• Model: Qwen/Qwen2.5-0.5B by default (multilingual, ~1 GB, CPU is fine). Any HF model name works: HuggingFaceTB/SmolLM2-135M is 4x faster (English), Qwen/Qwen2.5-1.5B is stronger and slower.
• Pooling: pooling="mean" (default), "last" (last token) or "mean+last"; layers="auto" concatenates the layer at 2/3 of the depth with the last one, or pass indices (layers=[6, -1]). template="Message: {text}\nIs it spam?" wraps every text in a prompt, which mostly helps "last" pooling.
• Head: "linear" (default; logistic regression / ridge with the strength chosen by cross-validation), "mlp" (a small torch MLP with early stopping) or any sklearn estimator (head=LGBMClassifier()).
• Task: task="auto" treats non-numeric targets and integers with at most 20 values as classes, the rest as regression. DataFrame inputs with several text columns are joined as "col: value" lines.
• Speed: the forward pass dominates. On a laptop CPU (i5-1335U, max_length=128) SmolLM2-135M encodes ~3 texts/s and Qwen2.5-0.5B ~1 text/s; a GPU or a desktop CPU is many times faster. Lower max_length= for long texts. Embeddings are cached on disk by text hash, so another head on the same texts is instant.
• Quality (20 Newsgroups, 4 topics, 500 training texts, 300 test): TF-IDF + logistic regression 0.747, LLMHead("HuggingFaceTB/SmolLM2-135M") 0.847. Frozen LLM features shine with few labels; with tens of thousands of labelled texts TF-IDF and fine-tuned models catch up.

Inside AutoBooster the same encoder is available as extra features for text columns: ab.AutoBooster(..., llm_model="Qwen/Qwen2.5-0.5B") adds PCA components and out-of-fold linear predictions of the frozen LLM features (llm__<col>__pca*, llm__<col>__oof*), which then compete in feature selection.

Benchmarks

benchmarks/run.py runs Kaggle datasets from public mirrors (OpenML, UCI, GroupLens, HuggingFace) and compares the result with a reference top score. Results are in BENCHMARKS.md.

.venv/bin/python benchmarks/run.py --all --budget 12m --cluster

Development

python -m venv .venv && .venv/bin/pip install -e ".[dev]"
.venv/bin/pytest -m "not slow"     # fast tests, ~4 min
.venv/bin/pytest -m slow           # full model zoo + text + images (downloads the embedding models)

License

MIT, see LICENSE.