A Keras-based entity embedding encoder for tabular ML and GBM pipelines
Project description
Category Embedding
A Keras-based neural encoder for categorical variables in tabular machine learning.
It learns dense vector representations (embeddings) for categorical features and outputs a clean numeric DataFrame that integrates seamlessly with gradient-boosted tree models.
This library is designed for ML engineers who want the benefits of deep-learning-based embeddings without replacing their GBM models.
It is sklearn-compatible, deterministic, and production-ready.
🚀 Features
-
Multi-column categorical embeddings
Each categorical feature receives its own learned embedding matrix. -
Smart default embedding dimensions
Uses a simple, interpretable rule:- If
n_cat ≤ 10:dim = n_cat - 1 - Else:
dim = max(10, n_cat // 2) - Always capped at 30.
- If
-
Per-column embedding dimension overrides
Pass a list of integers to control embedding size manually. -
Hashing for unseen categories
Unseen values at inference time are deterministically mapped into valid embedding indices. -
Residual MLP architecture
LayerNorm + GELU + Dropout + skip connections for stable training. -
Automatic numeric feature scaling
Numeric columns are scaled using StandardScaler during training.
You can choose whether the transformed output returns scaled or raw numeric values viascaled_num_out. -
Optional log-scaling of regression targets
For regression tasks, the target can be optionally transformed usinglog(y + 1e-6)during training (and inverse-transformed at prediction time) by settinglog_target=True. -
Supports regression and binary classification
The neural head is used only for training/tuning; GBMs remain the final predictor. -
Optional external validation set
Enables clean early stopping and stable embedding learning. -
Sklearn-compatible API
Implementsfit,transform,predict, andget_feature_names_out. -
Outputs a pandas DataFrame
Perfect for LightGBM, XGBoost, CatBoost, or any sklearn model.
📦 Installation
pip install category-embedding
Requires:
- Python ≥ 3.9
- TensorFlow ≥ 2.12
- scikit-learn ≥ 1.2
- pandas ≥ 1.5
🔧 Quick Start
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from lightgbm import LGBMRegressor
from category_embedding import CategoryEmbedding
df = pd.DataFrame({
"country": ["DE", "FR", "DE", "US", "US"],
"device": ["mobile", "desktop", "tablet", "mobile", "desktop"],
"age": [25, 40, 31, 22, 35],
})
y = [10.5, 20.1, 15.3, 8.7, 18.0]
categorical = ["country", "device"]
numeric = ["age"]
# CategoryEmbedding acts as a transformer inside a sklearn pipeline
preprocess = ColumnTransformer(
transformers=[
("emb", CategoryEmbedding(
task="regression",
categorical_cols=categorical,
numeric_cols=numeric,
epochs=20,
batch_size=32,
scaled_num_out=True, # return scaled numeric features
), categorical + numeric)
],
remainder="drop",
)
model = Pipeline([
("prep", preprocess),
("lgbm", LGBMRegressor(n_estimators=300, learning_rate=0.05)),
])
model.fit(df, y)
preds = model.predict(df)
🧠 Why Use Category Embedding?
Traditional encoders struggle with:
- high-cardinality categorical features
- sparse interactions
- noisy or rare categories
Neural embeddings solve this by learning dense, continuous representations that capture similarity structure between categories.
This library gives you:
- the power of deep learning
- the simplicity and performance of GBMs
- a clean sklearn interface
- deterministic, production-ready behavior
- automatic numeric scaling and target log-transforming for stable training
⚙️ API Overview
CategoryEmbedding(...)
log_target : bool, default=False Whether to apply a log transformation to the target variable for regression tasks.
Parameters
| Parameter | Type | Default | Description |
|---|---|---|---|
| task | "regression" or "classification" | "regression" | Determines loss function. |
| log_target | bool | False | Whether to apply a log transformation to the target variable for regression tasks. |
| categorical_cols | list[str] or None | None | Names of categorical columns to embed. |
| numeric_cols | list[str] or None | None | Names of numeric columns to include as inputs. |
| embedding_dims | list[int] or None | None | Optional list of integers specifying the embedding dimension for each categorical column, in the same order as categorical_cols. If None, a per-column default rule is used (if n_cat <= 10: dim = n_cat - 1, else: dim = max(10, n_cat // 2), in all cases: dim <= 30. |
| hidden_units | int | 64 | Width of each residual MLP block. |
| n_blocks | int | 2 | Number of residual blocks. |
| dropout_rate | float | 0.2 | Dropout rate inside residual blocks and before output head. |
| l2_emb | float | 1e-6 | L2 regularization for embedding weights. |
| l2_dense | float | 1e-6 | L2 regularization for dense layers. |
| batch_size | int | 512 | Training batch size. |
| epochs | int | 30 | Maximum number of training epochs. |
| lr | float | 2e-3 | Learning rate for Adam optimizer. |
| random_state | int | 42 | TensorFlow random seed. |
| verbose | int | 1 | Verbosity passed to Keras .fit(). |
| patience | int | 4 | Early stopping patience. |
| reduce_lr_factor | float | 0.5 | LR reduction factor when validation loss plateaus. |
| reduce_lr_patience | int | 2 | Patience before reducing LR. |
| val_set | tuple(X_val, y_val) or None | None | Optional external validation set. |
| scaled_num_out | bool | True | If True, transform() returns scaled numeric columns; otherwise raw. |
Methods
.fit(X, y)
Trains the embedding model:
- learns embeddings
- fits numeric scaler
- applies log-scaling to regression targets
- trains the neural model
.transform(X)
Returns a DataFrame containing:
- learned embeddings
- numeric features (scaled or raw depending on
scaled_num_out)
.predict(X)
Uses the neural head for tuning/evaluation.
For regression: automatically applies inverse log-transform.
.get_feature_names_out()
Returns the names of all output columns.
🔑 Keywords
machine learning, deep learning, tabular data, categorical encoding, entity embeddings, category embeddings, neural encoder, keras, tensorflow, scikit-learn, sklearn transformer, feature engineering, gradient boosting, lightgbm, xgboost, embeddings for gbm, high-cardinality features, optuna tuning, ml pipelines
📄 License
This project is licensed under the MIT License. See the LICENSE file for details.
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file category_embedding-0.2.2.tar.gz.
File metadata
- Download URL: category_embedding-0.2.2.tar.gz
- Upload date:
- Size: 14.4 kB
- Tags: Source
- Uploaded using Trusted Publishing? Yes
- Uploaded via: twine/6.1.0 CPython/3.13.7
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
110dbd971ba4021821361fe57695a251642cf35c6c88e0815bfe3944f726615a
|
|
| MD5 |
8a6298ba28dbbf8af22a4ad80706aac4
|
|
| BLAKE2b-256 |
c7c309abe60001a9bcee56eba6d7a2c11f32550fba0eeb316d26592324fa878c
|
Provenance
The following attestation bundles were made for category_embedding-0.2.2.tar.gz:
Publisher:
publish.yml on Karabush/category-embedding
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
category_embedding-0.2.2.tar.gz -
Subject digest:
110dbd971ba4021821361fe57695a251642cf35c6c88e0815bfe3944f726615a - Sigstore transparency entry: 954806513
- Sigstore integration time:
-
Permalink:
Karabush/category-embedding@a302178e5c28700484a576cbe591d2c6c00bca47 -
Branch / Tag:
refs/tags/v0.2.2 - Owner: https://github.com/Karabush
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish.yml@a302178e5c28700484a576cbe591d2c6c00bca47 -
Trigger Event:
push
-
Statement type:
File details
Details for the file category_embedding-0.2.2-py3-none-any.whl.
File metadata
- Download URL: category_embedding-0.2.2-py3-none-any.whl
- Upload date:
- Size: 11.9 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? Yes
- Uploaded via: twine/6.1.0 CPython/3.13.7
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
318ea46afbc8613031388478bc9f1e7692683ae401ac6217ec0f972ac60021db
|
|
| MD5 |
167035808584b9aaad87f1c5a0d76fc9
|
|
| BLAKE2b-256 |
420b59828cee4171df8ad40e2d0ab055b6debe32665dacfd5418b2b287cbbf6d
|
Provenance
The following attestation bundles were made for category_embedding-0.2.2-py3-none-any.whl:
Publisher:
publish.yml on Karabush/category-embedding
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
category_embedding-0.2.2-py3-none-any.whl -
Subject digest:
318ea46afbc8613031388478bc9f1e7692683ae401ac6217ec0f972ac60021db - Sigstore transparency entry: 954806515
- Sigstore integration time:
-
Permalink:
Karabush/category-embedding@a302178e5c28700484a576cbe591d2c6c00bca47 -
Branch / Tag:
refs/tags/v0.2.2 - Owner: https://github.com/Karabush
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish.yml@a302178e5c28700484a576cbe591d2c6c00bca47 -
Trigger Event:
push
-
Statement type: