mostlyai-engine 1.0.0__py3-none-any.whl

mostlyai/engine/domain.py

@@ -0,0 +1,171 @@
+# Copyright 2025 MOSTLY AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from enum import Enum
+
+from pydantic import BaseModel, Field, ConfigDict, field_validator
+
+
+class CustomBaseModel(BaseModel):
+    model_config = ConfigDict(protected_namespaces=(), populate_by_name=True)
+
+
+class ModelType(str, Enum):
+    """
+    The type of model.
+
+    - `TABULAR`: A generative AI model tailored towards tabular data, trained from scratch.
+    - `LANGUAGE`: A generative AI model build upon a (pre-trained) language model.
+    """
+
+    tabular = "TABULAR"
+    language = "LANGUAGE"
+
+
+class ModelEncodingType(str, Enum):
+    """
+    The encoding type used for model training and data generation.
+
+    - `AUTO`: Model chooses among available encoding types based on the column's data type.
+    - `TABULAR_CATEGORICAL`: Model samples from existing (non-rare) categories.
+    - `TABULAR_NUMERIC_AUTO`: Model chooses among 3 numeric encoding types based on the values.
+    - `TABULAR_NUMERIC_DISCRETE`: Model samples from existing discrete numerical values.
+    - `TABULAR_NUMERIC_BINNED`: Model samples from binned buckets, to then sample randomly within a bucket.
+    - `TABULAR_NUMERIC_DIGIT`: Model samples each digit of a numerical value.
+    - `TABULAR_CHARACTER`: Model samples each character of a string value.
+    - `TABULAR_DATETIME`: Model samples each part of a datetime value.
+    - `TABULAR_DATETIME_RELATIVE`: Model samples the relative difference between datetimes within a sequence.
+    - `TABULAR_LAT_LONG`: Model samples a latitude-longitude column. The format is "latitude,longitude".
+    - `LANGUAGE_TEXT`: Model will train a distinct LANGUAGE model for this column, to then generate free text.
+    """
+
+    auto = "AUTO"
+    tabular_categorical = "TABULAR_CATEGORICAL"
+    tabular_numeric_auto = "TABULAR_NUMERIC_AUTO"
+    tabular_numeric_discrete = "TABULAR_NUMERIC_DISCRETE"
+    tabular_numeric_binned = "TABULAR_NUMERIC_BINNED"
+    tabular_numeric_digit = "TABULAR_NUMERIC_DIGIT"
+    tabular_character = "TABULAR_CHARACTER"
+    tabular_datetime = "TABULAR_DATETIME"
+    tabular_datetime_relative = "TABULAR_DATETIME_RELATIVE"
+    tabular_lat_long = "TABULAR_LAT_LONG"
+    language_text = "LANGUAGE_TEXT"
+
+
+class ModelStateStrategy(str, Enum):
+    """
+    The strategy of how any existing model states and training progress are to be handled.
+
+    - `RESET`: Start training from scratch. Overwrite any existing model states and training progress.
+    - `REUSE`: Reuse any existing model states, but start progress from scratch. Used for fine-tuning existing models.
+    - `RESUME`: Reuse any existing model states and progress. Used for continuing an aborted training.
+    """
+
+    reset = "RESET"
+    reuse = "REUSE"
+    resume = "RESUME"
+
+
+class RareCategoryReplacementMethod(str, Enum):
+    """
+    Specifies how rare categories will be sampled.
+    Only applicable if value protection has been enabled.
+
+    - `CONSTANT`: Replace rare categories by a constant `_RARE_` token.
+    - `SAMPLE`: Replace rare categories by a sample from non-rare categories.
+    """
+
+    constant = "CONSTANT"
+    sample = "SAMPLE"
+
+
+class RebalancingConfig(CustomBaseModel):
+    """
+    Configure rebalancing.
+    """
+
+    column: str = Field(..., description="The name of the column to be rebalanced.")
+    probabilities: dict[str, float] = Field(
+        ...,
+        description="The target distribution of samples values. The keys are the categorical values, and the values "
+        "are the probabilities.",
+    )
+
+    @field_validator("probabilities", mode="after")
+    def validate_probabilities(cls, v):
+        if not all(0 <= v <= 1 for v in v.values()):
+            raise ValueError("the probabilities must be between 0 and 1")
+        if not sum(v.values()) <= 1:
+            raise ValueError("the sum of probabilities must be less than or equal to 1")
+        return v
+
+
+class ImputationConfig(CustomBaseModel):
+    """
+    Configure imputation. Imputed columns will suppress the sampling of NULL values.
+    """
+
+    columns: list[str] = Field(..., description="The names of the columns to be imputed.")
+
+
+class FairnessConfig(CustomBaseModel):
+    """
+    Configure a fairness objective for the table.
+
+    The generated synthetic data will maintain robust statistical parity between the target column and
+    the specified sensitive columns. All these columns must be categorical.
+    """
+
+    target_column: str = Field(..., alias="targetColumn")
+    sensitive_columns: list[str] = Field(..., alias="sensitiveColumns")
+
+
+class DifferentialPrivacyConfig(CustomBaseModel):
+    """
+    The differential privacy configuration for training the model.
+    If not provided, then no differential privacy will be applied.
+    """
+
+    max_epsilon: float | None = Field(
+        default=None,
+        alias="maxEpsilon",
+        description="Specifies the maximum allowable epsilon value. If the training process exceeds this threshold, "
+        "it will be terminated early. Only model checkpoints with epsilon values below this limit will be "
+        "retained. If not provided, the training will proceed without early termination based on "
+        "epsilon constraints.",
+        ge=0.0,
+        le=10000.0,
+    )
+    noise_multiplier: float = Field(
+        default=1.5,
+        alias="noiseMultiplier",
+        description="The ratio of the standard deviation of the Gaussian noise to the L2-sensitivity of the function "
+        "to which the noise is added (How much noise to add).\n",
+        ge=0.0,
+        le=10000.0,
+    )
+    max_grad_norm: float = Field(
+        default=1.0,
+        alias="maxGradNorm",
+        description="The maximum norm of the per-sample gradients for training the model with differential privacy.\n",
+        ge=0.0,
+        le=10000.0,
+    )
+    delta: float = Field(
+        default=1e-5,
+        description="The delta value for differential privacy. It is the probability of the privacy guarantee not "
+        "holding. The smaller the delta, the more confident you can be that the privacy guarantee holds.",
+        ge=0.0,
+        le=1.0,
+    )

mostlyai/engine/encoding.py

@@ -0,0 +1,42 @@
+# Copyright 2025 MOSTLY AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pathlib import Path
+
+from mostlyai.engine._common import ProgressCallback
+from mostlyai.engine.domain import ModelType
+from mostlyai.engine._workspace import resolve_model_type
+
+
+def encode(*, workspace_dir: str | Path = "engine-ws", update_progress: ProgressCallback | None = None) -> None:
+    """
+    Encodes data in the workspace that has already been split and analyzed.
+
+    Creates the following folder structure within the `workspace_dir`:
+
+    - `OriginalData/encoded-data`: Encoded data for training, stored as parquet files.
+
+    Args:
+        workspace_dir: Directory path for workspace.
+        update_progress: Callback for progress updates.
+    """
+    model_type = resolve_model_type(workspace_dir)
+    if model_type == ModelType.tabular:
+        from mostlyai.engine._tabular.encoding import encode as encode_tabular
+
+        return encode_tabular(workspace_dir=workspace_dir, update_progress=update_progress)
+    else:
+        from mostlyai.engine._language.encoding import encode as encode_language
+
+        return encode_language(workspace_dir=workspace_dir, update_progress=update_progress)

mostlyai/engine/generation.py

@@ -0,0 +1,110 @@
+# Copyright 2025 MOSTLY AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+from pathlib import Path
+import pandas as pd
+
+from mostlyai.engine._common import ProgressCallback
+from mostlyai.engine.domain import (
+    ModelType,
+    RareCategoryReplacementMethod,
+    RebalancingConfig,
+    ImputationConfig,
+    FairnessConfig,
+)
+from mostlyai.engine._workspace import resolve_model_type
+
+
+def generate(
+    *,
+    ctx_data: pd.DataFrame | None = None,
+    seed_data: pd.DataFrame | None = None,
+    sample_size: int | None = None,
+    batch_size: int | None = None,
+    sampling_temperature: float = 1.0,
+    sampling_top_p: float = 1.0,
+    device: str | None = None,
+    rare_category_replacement_method: RareCategoryReplacementMethod | str | None = None,
+    rebalancing: RebalancingConfig | dict | None = None,
+    imputation: ImputationConfig | dict | None = None,
+    fairness: FairnessConfig | dict | None = None,
+    workspace_dir: str | Path = "engine-ws",
+    update_progress: ProgressCallback | None = None,
+) -> None:
+    """
+    Generates synthetic data from a trained model.
+
+    Creates the following folder structure within the `workspace_dir`:
+
+    - `SyntheticData`: Generated synthetic data, stored as parquet files.
+
+    Args:
+        ctx_data: Context data to be used for generation.
+        seed_data: Seed data to condition generation on fixed target columns.
+        sample_size: Number of samples to generate. Defaults to number of original samples.
+        batch_size: Batch size for generation. If None, determined automatically.
+        sampling_temperature: Sampling temperature. Higher values increase randomness.
+        sampling_top_p: Nucleus sampling probability threshold.
+        device: Device to run generation on ('cuda' or 'cpu'). Defaults to 'cuda' if available, else 'cpu'.
+        rare_category_replacement_method: Method for handling rare categories. Only applicable for tabular models.
+        rebalancing: Configuration for rebalancing column distributions. Only applicable for tabular models.
+        imputation: List of columns to impute missing values. Only applicable for tabular models.
+        fairness: Configuration for fairness constraints. Only applicable for tabular models.
+        workspace_dir: Directory path for workspace.
+        update_progress: Callback for progress updates.
+    """
+    model_type = resolve_model_type(workspace_dir)
+    if model_type == ModelType.tabular:
+        from mostlyai.engine._tabular.generation import generate as generate_tabular
+
+        return generate_tabular(
+            ctx_data=ctx_data,
+            seed_data=seed_data,
+            sample_size=sample_size,
+            batch_size=batch_size,
+            sampling_temperature=sampling_temperature,
+            sampling_top_p=sampling_top_p,
+            rare_category_replacement_method=inspect.signature(generate_tabular)
+            .parameters["rare_category_replacement_method"]
+            .default,
+            rebalancing=rebalancing,
+            imputation=imputation,
+            fairness=fairness,
+            device=device,
+            workspace_dir=workspace_dir,
+            update_progress=update_progress,
+        )
+    else:
+        from mostlyai.engine._language.generation import generate as generate_language
+
+        if imputation is not None:
+            raise ValueError("imputation is not supported for language models")
+        if fairness is not None:
+            raise ValueError("fairness is not supported for language models")
+        if rebalancing is not None:
+            raise ValueError("rebalancing is not supported for language models")
+        if rare_category_replacement_method is not None:
+            raise ValueError("rare_category_replacement_method is not supported for language models")
+        return generate_language(
+            ctx_data=ctx_data,
+            seed_data=seed_data,
+            sample_size=sample_size,
+            batch_size=batch_size,
+            sampling_temperature=sampling_temperature,
+            sampling_top_p=sampling_top_p,
+            device=device,
+            workspace_dir=workspace_dir,
+            update_progress=update_progress,
+        )

mostlyai/engine/splitting.py

@@ -0,0 +1,257 @@
+# Copyright 2025 MOSTLY AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Split original data for training and validation.
+"""
+
+import time
+import warnings
+from pathlib import Path
+
+from mostlyai.engine._common import ProgressCallbackWrapper, ProgressCallback
+from mostlyai.engine._dtypes import (
+    is_date_dtype,
+    is_float_dtype,
+    is_integer_dtype,
+    is_timestamp_dtype,
+)
+
+import numpy as np
+import logging
+
+import pandas as pd
+
+from mostlyai.engine.domain import ModelType, ModelEncodingType
+from mostlyai.engine._workspace import (
+    PathDesc,
+    Workspace,
+    ensure_workspace_dir,
+    reset_dir,
+)
+
+warnings.simplefilter(action="ignore", category=UserWarning)
+
+_LOG = logging.getLogger(__name__)
+
+
+def _get_default_tabular_encoding_type(x: pd.Series) -> ModelEncodingType:
+    if is_integer_dtype(x) or is_float_dtype(x):
+        return ModelEncodingType.tabular_numeric_auto
+    elif is_date_dtype(x) or is_timestamp_dtype(x):
+        return ModelEncodingType.tabular_datetime
+    else:
+        return ModelEncodingType.tabular_categorical
+
+
+def _get_default_language_encoding_type(_: pd.Series) -> ModelEncodingType:
+    return ModelEncodingType.language_text
+
+
+def split(
+    tgt_data: pd.DataFrame,
+    *,
+    ctx_data: pd.DataFrame | None = None,
+    tgt_primary_key: str | None = None,
+    ctx_primary_key: str | None = None,
+    tgt_context_key: str | None = None,
+    model_type: str | ModelType | None = None,
+    tgt_encoding_types: dict[str, str | ModelEncodingType] | None = None,
+    ctx_encoding_types: dict[str, str | ModelEncodingType] | None = None,
+    n_partitions: int = 1,
+    workspace_dir: str | Path = "engine-ws",
+    update_progress: ProgressCallback | None = None,
+) -> None:
+    """
+    Splits the provided original data into training and validation sets, and stores these as partitioned Parquet files.
+    This is a simplified version of `mostlyai-data`, tailored towards single- and two-table use cases, while requiring
+    all data to be passed as DataFrames in memory.
+
+    Creates the following folder structure within the `workspace_dir`:
+
+      - `OriginalData/tgt-data`: Partitioned target data files.
+      - `OriginalData/tgt-meta`: Metadata files for target data.
+      - `OriginalData/ctx-data`: Partitioned context data files (if context is provided).
+      - `OriginalData/ctx-meta`: Metadata files for context data (if context is provided).
+
+    Args:
+        tgt_data: DataFrame containing the target data.
+        ctx_data: DataFrame containing the context data.
+        tgt_primary_key: Primary key column name in the target data.
+        ctx_primary_key: Primary key column name in the context data.
+        tgt_context_key: Context key column name in the target data.
+        model_type: Model type for the target data. If not provided, it will be inferred from the encoding types, or set to TABULAR by default.
+        tgt_encoding_types: Encoding types for columns in the target data (excluding key columns).
+        ctx_encoding_types: Encoding types for columns in the context data (excluding key columns).
+        n_partitions: Number of partitions to split the data into.
+        workspace_dir: Path to the workspace directory where files will be created.
+        update_progress: A custom progress callback.
+    """
+
+    _LOG.info("SPLIT started")
+    t0 = time.time()
+    with ProgressCallbackWrapper(update_progress) as progress:
+        # validate input
+        if tgt_primary_key and tgt_primary_key not in tgt_data:
+            raise Exception("tgt_primary_key not found in tgt_data")
+        if tgt_primary_key and tgt_data[tgt_primary_key].duplicated().any():
+            raise Exception("tgt_primary_key must be unique")
+        if ctx_data is not None:
+            if tgt_context_key is None:
+                raise Exception("tgt_context_key must be provided")
+            if tgt_context_key not in tgt_data:
+                raise Exception("tgt_context_key not found in tgt_data")
+            if ctx_primary_key is None:
+                raise Exception("ctx_primary_key must be provided")
+            if ctx_primary_key not in ctx_data:
+                raise Exception("ctx_primary_key not found in ctx_data")
+            if ctx_data[ctx_primary_key].duplicated().any():
+                raise Exception("ctx_primary_key must be unique")
+            if not tgt_data[tgt_context_key].isin(ctx_data[ctx_primary_key]).all():
+                raise Exception("tgt_primary_keys must match ctx_primary_keys")
+
+        # prepare data
+        if ctx_data is None and tgt_context_key is not None:
+            ctx_primary_key = ctx_primary_key or tgt_context_key
+            ctx_data = pd.DataFrame({ctx_primary_key: tgt_data[tgt_context_key].drop_duplicates()})
+
+        # prepare workspace directory
+        workspace_dir = ensure_workspace_dir(workspace_dir)
+        ws = Workspace(workspace_dir)
+        reset_dir(ws.tgt_data_path)
+        reset_dir(ws.tgt_meta_path)
+        if ctx_data is not None:
+            reset_dir(ws.ctx_data_path)
+            reset_dir(ws.ctx_meta_path)
+
+        # set model_type if not provided
+        if not model_type and not tgt_encoding_types:
+            # use TABULAR as default
+            model_type = ModelType.tabular
+        elif not model_type and tgt_encoding_types:
+            # consider ModelType of the first encoding type
+            first_enc_type = next(iter(tgt_encoding_types.values()))
+            model_type = ModelType.tabular if first_enc_type.startswith(ModelType.tabular.value) else ModelType.language
+        model_type = ModelType(model_type).value
+
+        # populate model_encoding_types with defaults if not specified
+        tgt_encoding_types = tgt_encoding_types or {}
+        tgt_cols = [c for c in tgt_data if c != tgt_primary_key and c != tgt_context_key]
+        for col in tgt_cols:
+            if col not in tgt_encoding_types:
+                if model_type == ModelType.language:
+                    tgt_encoding_types[col] = _get_default_language_encoding_type(tgt_data[col]).value
+                else:
+                    tgt_encoding_types[col] = _get_default_tabular_encoding_type(tgt_data[col]).value
+            else:
+                tgt_encoding_types[col] = ModelEncodingType(tgt_encoding_types[col]).value
+        ctx_encoding_types = ctx_encoding_types or {}
+        ctx_cols = [c for c in ctx_data if c != ctx_primary_key] if ctx_data is not None else []
+        for col in ctx_cols:
+            if col not in ctx_encoding_types:
+                ctx_encoding_types[col] = _get_default_tabular_encoding_type(ctx_data[col]).value
+            else:
+                ctx_encoding_types[col] = ModelEncodingType(ctx_encoding_types[col]).value
+        _LOG.info(f"{model_type=}")
+        _LOG.info(f"{tgt_encoding_types=}")
+        if ctx_encoding_types:
+            _LOG.info(f"{ctx_encoding_types=}")
+
+        # check if tgt_encoding_types are consistent with model_type
+        for col, enc_type in tgt_encoding_types.items():
+            if not enc_type.startswith(model_type):
+                raise Exception("Mismatch between model_type and model_encoding_types")
+
+        # split into trn/val
+        tmp_key = "__key"
+        if ctx_data is None and tgt_context_key is None:
+            # assign temporary key to enable consistent splitting logic; we drop that key later on again
+            tgt_data = tgt_data.copy()
+            tgt_data[tmp_key] = range(len(tgt_data))
+            tgt_context_key = tmp_key
+        # gather all keys
+        if ctx_data is None:
+            keys = tgt_data[tgt_context_key].drop_duplicates()
+        else:
+            keys = ctx_data[ctx_primary_key]
+        # shuffle keys and split
+        keys = keys.sample(frac=1)
+        trn_cnt = round(0.9 * len(keys))
+        trn_keys = keys[:trn_cnt]
+        val_keys = keys[trn_cnt:]
+
+        def save_partition(
+            df: pd.DataFrame, path: Path, key: str, sel_trn_keys: np.ndarray, sel_val_keys: np.ndarray, idx: int
+        ):
+            trn_df = df[df[key].isin(sel_trn_keys)].drop(columns=[tmp_key], errors="ignore")
+            val_df = df[df[key].isin(sel_val_keys)].drop(columns=[tmp_key], errors="ignore")
+            trn_df.to_parquet(path / f"part.{idx:06d}-trn.parquet", engine="pyarrow", index=False)
+            val_df.to_parquet(path / f"part.{idx:06d}-val.parquet", engine="pyarrow", index=False)
+
+        # save tgt-data
+        for i in range(n_partitions):
+            save_partition(
+                df=tgt_data,
+                path=ws.tgt_data_path,
+                key=tgt_context_key,
+                sel_trn_keys=trn_keys[i::n_partitions],
+                sel_val_keys=val_keys[i::n_partitions],
+                idx=i,
+            )
+            progress.update(completed=i, total=2 * n_partitions)
+        if tgt_context_key == tmp_key:
+            tgt_context_key = None
+        # save tgt-meta
+        _save_meta_encoding_types(tgt_encoding_types, ws.tgt_encoding_types)
+        _save_meta_keys(primary_key=tgt_primary_key, context_key=tgt_context_key, path=ws.tgt_keys)
+
+        if ctx_data is not None:
+            # save ctx-data
+            for i in range(n_partitions):
+                save_partition(
+                    df=ctx_data,
+                    path=ws.ctx_data_path,
+                    key=ctx_primary_key,
+                    sel_trn_keys=trn_keys[i::n_partitions],
+                    sel_val_keys=val_keys[i::n_partitions],
+                    idx=i,
+                )
+                progress.update(completed=n_partitions + i, total=2 * n_partitions)
+            # save ctx-meta
+            _save_meta_keys(primary_key=ctx_primary_key, path=ws.ctx_keys)
+            _save_meta_encoding_types(ctx_encoding_types, path=ws.ctx_encoding_types)
+
+    _LOG.info(f"SPLIT finished in {time.time() - t0:.2f}s")
+
+
+def _save_meta_keys(
+    path: PathDesc,
+    primary_key: str | None = None,
+    context_key: str | None = None,
+):
+    keys = {}
+    if primary_key is not None:
+        keys["primary_key"] = primary_key
+    if context_key is not None:
+        keys["context_key"] = context_key
+    path.write(keys)
+
+
+def _save_meta_encoding_types(encoding_types: dict[str, ModelEncodingType], path: PathDesc):
+    for col, enc_type in encoding_types.items():
+        try:
+            ModelEncodingType(enc_type)
+        except ValueError:
+            raise ValueError(f"unknown encoding type: {enc_type} on column {col}")
+    path.write(encoding_types)

mostlyai/engine/training.py

@@ -0,0 +1,104 @@
+# Copyright 2025 MOSTLY AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+from pathlib import Path
+from collections.abc import Callable
+
+import torch
+
+from mostlyai.engine._common import ProgressCallback
+from mostlyai.engine.domain import ModelType, ModelStateStrategy, DifferentialPrivacyConfig
+from mostlyai.engine._workspace import resolve_model_type
+
+
+def train(
+    *,
+    model: str | None = None,
+    max_training_time: float = 60.0,
+    max_epochs: float = 100.0,
+    batch_size: int | None = None,
+    gradient_accumulation_steps: int | None = None,
+    enable_flexible_generation: bool = True,
+    max_sequence_window: int | None = None,
+    differential_privacy: DifferentialPrivacyConfig | dict | None = None,
+    model_state_strategy: ModelStateStrategy = ModelStateStrategy.reset,
+    device: torch.device | str | None = None,
+    workspace_dir: str | Path = "engine-ws",
+    update_progress: ProgressCallback | None = None,
+    upload_model_data_callback: Callable | None = None,  # TODO
+) -> None:
+    """
+    Trains a model with optional early stopping and differential privacy.
+
+    Creates the following folder structure within the `workspace_dir`:
+
+    - `ModelStore`: Trained model checkpoints and logs.
+
+    Args:
+        model: The identifier of the model to train. If tabular, defaults to MOSTLY_AI/Medium. If language, defaults to MOSTLY_AI/LSTMFromScratch-3m.
+        max_training_time: Maximum training time in minutes.
+        max_epochs: Maximum number of training epochs.
+        batch_size: Per-device batch size for training and validation. If None, determined automatically.
+        gradient_accumulation_steps: Number of steps to accumulate gradients. If None, determined automatically.
+        enable_flexible_generation: Whether to enable flexible order generation. Defaults to True.
+        max_sequence_window: Maximum sequence window for tabular sequential models. Only applicable for tabular models.
+        differential_privacy: Configuration for differential privacy training. If None, DP is disabled.
+        model_state_strategy: Strategy for handling existing model state (reset/resume/reuse).
+        device: Device to run training on ('cuda' or 'cpu'). Defaults to 'cuda' if available, else 'cpu'.
+        workspace_dir: Directory path for workspace. Training outputs are stored in ModelStore subdirectory.
+        update_progress: Callback function to report training progress.
+        upload_model_data_callback: Callback function to upload model data during training.
+    """
+    model_type = resolve_model_type(workspace_dir)
+    if model_type == ModelType.tabular:
+        from mostlyai.engine._tabular.training import train as train_tabular
+
+        train_tabular(
+            model=model if model else inspect.signature(train_tabular).parameters["model"].default,
+            workspace_dir=workspace_dir,
+            max_training_time=max_training_time,
+            max_epochs=max_epochs,
+            batch_size=batch_size,
+            gradient_accumulation_steps=gradient_accumulation_steps,
+            enable_flexible_generation=enable_flexible_generation,
+            differential_privacy=differential_privacy,
+            update_progress=update_progress,
+            upload_model_data_callback=upload_model_data_callback,
+            model_state_strategy=model_state_strategy,
+            device=device,
+            max_sequence_window=max_sequence_window
+            if max_sequence_window
+            else inspect.signature(train_tabular).parameters["max_sequence_window"].default,
+        )
+    else:
+        from mostlyai.engine._language.training import train as train_language
+
+        if max_sequence_window is not None:
+            raise ValueError("max_sequence_window is not supported for language models")
+
+        train_language(
+            model=model if model else inspect.signature(train_language).parameters["model"].default,
+            workspace_dir=workspace_dir,
+            max_training_time=max_training_time,
+            max_epochs=max_epochs,
+            batch_size=batch_size,
+            gradient_accumulation_steps=gradient_accumulation_steps,
+            enable_flexible_generation=enable_flexible_generation,
+            differential_privacy=differential_privacy,
+            update_progress=update_progress,
+            upload_model_data_callback=upload_model_data_callback,
+            model_state_strategy=model_state_strategy,
+            device=device,
+        )