Kamae bridges the gap between offline data processing and online model serving. Build preprocessing pipelines in Spark for big data workloads, then export them as Keras 3 models for low-latency inference. Multi-backend support allows numeric operations to run on TensorFlow, JAX, or PyTorch backends, while string and datetime operations require TensorFlow.

Training and serving often happen on different platforms. Spark for batch processing at scale, Keras for low-latency inference. Manually reimplementing preprocessing logic in both places creates:

• Training/serving skew: Subtle bugs from inconsistent implementations
• Development overhead: Writing and maintaining duplicate code
• Deployment friction: Changes require updates in multiple systems

Kamae solves this by generating the inference model directly from your Spark pipeline, guaranteeing consistency between training and serving.

pip install kamae

from pyspark.sql import SparkSession
from kamae.spark.estimators import StandardScaleEstimator, StringIndexEstimator
from kamae.spark.pipeline import KamaeSparkPipeline
from kamae.spark.transformers import LogTransformer, ArrayConcatenateTransformer

# Define preprocessing in Spark
spark = SparkSession.builder.getOrCreate()
data = spark.createDataFrame(
    [(1, 2, "a"), (4, 5, "b"), (7, 8, "c")],
    ["col1", "col2", "category"]
)

pipeline = KamaeSparkPipeline(stages=[
    LogTransformer(inputCol="col1", outputCol="log_col1", alpha=1, inputDtype="float"),
    ArrayConcatenateTransformer(inputCols=["log_col1", "col2"], outputCol="features", inputDtype="float"),
    StandardScaleEstimator(inputCol="features", outputCol="scaled_features"),
    StringIndexEstimator(inputCol="category", outputCol="category_indexed"),
])

fitted_pipeline = pipeline.fit(data)
fitted_pipeline.transform(data).show()  # Use in Spark

# Export for TensorFlow Serving
tf_input_schema = [
    {"name": "col1", "dtype": "int32", "shape": (None, 1)},
    {"name": "col2", "dtype": "int32", "shape": (None, 1)},
    {"name": "category", "dtype": "string", "shape": (None, 1)},
]
keras_model = fitted_pipeline.build_keras_model(tf_input_schema=tf_input_schema)
keras_model.save("./preprocessing_model.keras")

Spark Pipeline (Recommended): Build preprocessing pipelines using Spark transformers and estimators, fit on DataFrames, then export as Keras models. See examples for common patterns.

Direct Keras Layers: Import and compose Keras layers directly for non-tabular data or custom workflows. Browse available layers in the transformation table below.

Backend Selection: Set KERAS_BACKEND environment variable before importing keras:

import os
os.environ['KERAS_BACKEND'] = 'tensorflow'  # or 'jax' or 'torch'

Multi-backend layers (numeric operations) work on all backends. TensorFlow-only layers (string/datetime operations) require TensorFlow backend. See the Backend column in the transformation table below, or use the discovery API:

import kamae
# Get layers/transformers compatible with current backend
layers = kamae.get_compatible_layers()
transformers = kamae.get_compatible_transformers()

# Get layers/transformers compatible with specific backend
jax_layers = kamae.get_compatible_layers('jax')
torch_transformers = kamae.get_compatible_transformers('torch')

Note: TensorFlow is a required dependency for Kamae, as the package includes TensorFlow-only layers. JAX and PyTorch backends provide an alternative execution path for numeric operations only.

• Examples: Full working examples for common use cases
• Chaining models: Use Kamae preprocessing models as inputs to trainable models
• Type parity: Ensuring consistent dtypes between Spark and Keras
• Shape parity: Ensuring consistent shapes between Spark and Keras
• Testing inference: Validate model outputs with TensorFlow Serving
• Adding transformers: Contributing new transformations
• v3.0.0 Migration Guide: Upgrading from Kamae 2.x to 3.0.0 (Keras 3 multi-backend, sklearn removal, deprecated layer removals, hash indexer changes)

Requirements: Python 3.10 (for development), pipx (installation instructions)

make setup      # Install dependencies and pre-commit hooks
make all        # Run tests, formatting, and linting
make help       # See all available commands

The package supports Python 3.8-3.12 in production.

make run-example        # Run example pipeline
make test-tf-serving    # Test TensorFlow Serving inference
make test-end-to-end    # Run example + test serving

Create a branch from main and open a pull request. Follow the adding transformers guide for new transformers.

Code quality: Pre-commit hooks enforce formatting and linting. Install with uv run pre-commit install. PRs must pass all tests in tests/.

Versioning: Automated via semantic-release. Use conventional commit prefixes in PR titles: fix: (patch), feat: (minor), BREAKING CHANGE: (major).

Contact: Questions? Reach out to the Kamae team.