mostlyai-engine 1.4.0__tar.gz → 1.4.1__tar.gz

{mostlyai_engine-1.4.0 → mostlyai_engine-1.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mostlyai-engine
-Version: 1.4.0
+Version: 1.4.1
 Summary: Synthetic Data Engine
 Project-URL: homepage, https://github.com/mostly-ai/mostlyai-engine
 Project-URL: repository, https://github.com/mostly-ai/mostlyai-engine

{mostlyai_engine-1.4.0 → mostlyai_engine-1.4.1}/mostlyai/engine/__init__.py

@@ -14,6 +14,7 @@
 import warnings
 
 from mostlyai.engine.logging import init_logging
+from mostlyai.engine.random_state import set_random_state
 from mostlyai.engine.splitting import split
 from mostlyai.engine.analysis import analyze
 from mostlyai.engine.training import train
@@ -21,8 +22,8 @@ from mostlyai.engine.encoding import encode
 from mostlyai.engine.generation import generate
 
 
-__all__ = ["split", "analyze", "encode", "train", "generate", "init_logging"]
-__version__ = "1.4.0"
+__all__ = ["split", "analyze", "encode", "train", "generate", "init_logging", "set_random_state"]
+__version__ = "1.4.1"
 
 # suppress specific warning related to os.fork() in multi-threaded processes
 warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*multi-threaded.*fork.*")

{mostlyai_engine-1.4.0 → mostlyai_engine-1.4.1}/mostlyai/engine/_common.py

@@ -16,9 +16,7 @@ import inspect
 import itertools
 import json
 import logging
-import os
 import platform
-import struct
 import time
 from functools import wraps
 from pathlib import Path
@@ -883,34 +881,3 @@ def dp_non_rare(value_counts: dict[str, int], epsilon: float, threshold: int = 5
 
 def get_stochastic_rare_threshold(min_threshold: int = 5, noise_multiplier: float = 3) -> int:
     return min_threshold + int(noise_multiplier * np.random.uniform())
-
-
-def set_random_state(random_state: int | None = None, worker: bool = False):
-    def get_random_int_from_os() -> int:
-        # 32-bit, cryptographically secure random int from os
-        return int(struct.unpack("I", os.urandom(4))[0])
-
-    if worker:  # worker process
-        if "MOSTLYAI_ENGINE_SEED" in os.environ:
-            random_state = int(os.environ["MOSTLYAI_ENGINE_SEED"])
-        else:
-            # normally, the seed should have been set in the main process
-            # but if not (e.g. in tests), we fallback to generating a random seed here
-            random_state = get_random_int_from_os()
-    else:  # main process
-        if random_state is not None:
-            _LOG.info(f"Global random_state set to `{random_state}`")
-
-        if random_state is None:
-            random_state = get_random_int_from_os()
-
-        os.environ["MOSTLYAI_ENGINE_SEED"] = str(random_state)
-
-    import random
-    import numpy as np
-    import torch
-
-    random.seed(random_state)
-    np.random.seed(random_state)
-    torch.manual_seed(random_state)
-    torch.cuda.manual_seed_all(random_state)

{mostlyai_engine-1.4.0 → mostlyai_engine-1.4.1}/mostlyai/engine/_language/encoding.py

@@ -27,7 +27,6 @@ from mostlyai.engine._common import (
     ProgressCallback,
     ProgressCallbackWrapper,
     TABLE_COLUMN_INFIX,
-    set_random_state,
 )
 from mostlyai.engine._workspace import ensure_workspace_dir, Workspace, reset_dir
 from mostlyai.engine._encoding_types.language.categorical import encode_language_categorical
@@ -158,11 +157,9 @@ def _encode_partition(
 def encode(
     workspace_dir: str | Path | None = None,
     update_progress: ProgressCallback | None = None,
-    random_state: int | None = None,
 ) -> None:
     _LOG.info("ENCODE_LANGUAGE started")
     t0 = time.time()
-    set_random_state(random_state)
     with ProgressCallbackWrapper(update_progress) as progress:
         workspace_dir = ensure_workspace_dir(workspace_dir)
         workspace = Workspace(workspace_dir)

{mostlyai_engine-1.4.0 → mostlyai_engine-1.4.1}/mostlyai/engine/_language/generation.py

@@ -33,7 +33,6 @@ from mostlyai.engine._common import (
     FixedSizeSampleBuffer,
     ProgressCallback,
     ProgressCallbackWrapper,
-    set_random_state,
 )
 from mostlyai.engine._encoding_types.language.categorical import decode_language_categorical
 from mostlyai.engine._encoding_types.language.datetime import decode_language_datetime
@@ -150,11 +149,9 @@ def generate(
     device: torch.device | str | None = None,
     workspace_dir: str | Path = "engine-ws",
     update_progress: ProgressCallback | None = None,
-    random_state: int | None = None,
 ):
     _LOG.info("GENERATE_LANGUAGE started")
     t0_ = time.time()
-    set_random_state(random_state)
     os.environ["VLLM_LOGGING_LEVEL"] = "WARNING"
     os.environ["VLLM_NO_DEPRECATION_WARNING"] = "1"
 

{mostlyai_engine-1.4.0 → mostlyai_engine-1.4.1}/mostlyai/engine/_language/training.py

@@ -38,7 +38,7 @@ from opacus.utils.batch_memory_manager import wrap_data_loader
 
 from torch.utils.data import DataLoader
 
-from mostlyai.engine._common import ProgressCallback, ProgressCallbackWrapper, TABLE_COLUMN_INFIX, set_random_state
+from mostlyai.engine._common import ProgressCallback, ProgressCallbackWrapper, TABLE_COLUMN_INFIX
 from mostlyai.engine._language.common import (
     is_bf16_supported,
     load_base_model_and_config,
@@ -75,7 +75,7 @@ _LOG = logging.getLogger(__name__)
 
 
 def _physical_batch_size_heuristic(
-    no_of_records: int, no_of_model_params: int, max_tokens: int, device: torch.device
+    no_of_records: int, no_of_model_params: int, max_tokens: int, model_id: str, device: torch.device
 ) -> int:
     """
     Calculate the physical batch size that fits in memory.
@@ -84,16 +84,19 @@ def _physical_batch_size_heuristic(
         no_of_records (int): Number of records in the training dataset.
         no_of_model_params (int): Number of model parameters.
         max_tokens (int): Maximum number of tokens that are in the training dataset.
+        model_id (str): Model ID.
         device (torch.device): Device to run training on.
 
     Returns:
         Batch size (int)
     """
     min_batches = 8
-    max_batch_size = max(1, no_of_records // min_batches)
 
     if device.type == "cuda":
-        batch_size = 2**10  # 1024, max 10 reductions
+        if model_id == LSTMFromScratchConfig.model_id:
+            batch_size = 64  # empirically tuned for LSTM to have a better training dynamics
+        else:
+            batch_size = 2**10  # 1024, max 10 reductions
     else:
         if no_of_model_params < 10_000_000:
             batch_size = 32
@@ -101,7 +104,7 @@ def _physical_batch_size_heuristic(
             batch_size = 16 if max_tokens < 100 else 8
         else:
             batch_size = 8 if max_tokens < 100 else 4
-    batch_size = 2 ** int(np.log2(no_of_records / min_batches)) if no_of_records > 0 else 1
+    max_batch_size = 2 ** int(np.log2(no_of_records / min_batches)) if no_of_records > 0 else 1
     return int(np.clip(a=batch_size, a_min=1, a_max=max_batch_size))
 
 
@@ -281,11 +284,9 @@ def train(
     device: torch.device | str | None = None,
     workspace_dir: str | Path = "engine-ws",
     update_progress: ProgressCallback | None = None,
-    random_state: int | None = None,
 ):
     _LOG.info("TRAIN_LANGUAGE started")
     t0_ = time.time()
-    set_random_state(random_state)
     workspace_dir = ensure_workspace_dir(workspace_dir)
     workspace = Workspace(workspace_dir)
 
@@ -502,6 +503,7 @@ def train(
                 no_of_records=trn_cnt,
                 no_of_model_params=no_of_model_params,
                 max_tokens=max_tokens,
+                model_id=model.config.model_type,
                 device=device,
             )
         batch_size = max(1, min(batch_size, trn_cnt))

{mostlyai_engine-1.4.0 → mostlyai_engine-1.4.1}/mostlyai/engine/_tabular/encoding.py

@@ -35,7 +35,6 @@ from mostlyai.engine._common import (
     SDEC_SUB_COLUMN_PREFIX,
     ProgressCallback,
     ProgressCallbackWrapper,
-    set_random_state,
 )
 from mostlyai.engine._encoding_types.tabular.categorical import encode_categorical
 from mostlyai.engine._encoding_types.tabular.character import encode_character
@@ -45,6 +44,7 @@ from mostlyai.engine._encoding_types.tabular.lat_long import encode_latlong
 from mostlyai.engine._encoding_types.tabular.numeric import encode_numeric
 from mostlyai.engine.domain import ModelEncodingType
 from mostlyai.engine._workspace import Workspace, ensure_workspace_dir, reset_dir
+from mostlyai.engine.random_state import set_random_state
 
 _LOG = logging.getLogger(__name__)
 
@@ -52,11 +52,9 @@ _LOG = logging.getLogger(__name__)
 def encode(
     workspace_dir: str | Path | None = None,
     update_progress: ProgressCallback | None = None,
-    random_state: int | None = None,
 ) -> None:
     _LOG.info("ENCODE_TABULAR started")
     t0 = time.time()
-    set_random_state(random_state)
     with ProgressCallbackWrapper(update_progress) as progress:
         # build paths based on workspace dir
         workspace_dir = ensure_workspace_dir(workspace_dir)

{mostlyai_engine-1.4.0 → mostlyai_engine-1.4.1}/mostlyai/engine/_tabular/generation.py

@@ -61,7 +61,6 @@ from mostlyai.engine._common import (
     get_sequence_length_stats,
     get_sub_columns_from_cardinalities,
     is_sequential,
-    set_random_state,
     trim_sequences,
     get_sub_columns_nested_from_cardinalities,
     persist_data_part,
@@ -689,11 +688,9 @@ def generate(
     device: torch.device | str | None = None,
     workspace_dir: str | Path = "engine-ws",
     update_progress: ProgressCallback | None = None,
-    random_state: int | None = None,
 ) -> None:
     _LOG.info("GENERATE_TABULAR started")
     t0 = time.time()
-    set_random_state(random_state)
     with ProgressCallbackWrapper(update_progress) as progress:
         # build paths based on workspace dir
         workspace_dir = ensure_workspace_dir(workspace_dir)

{mostlyai_engine-1.4.0 → mostlyai_engine-1.4.1}/mostlyai/engine/_tabular/training.py

@@ -59,7 +59,6 @@ from mostlyai.engine._common import (
     SDEC_SUB_COLUMN_PREFIX,
     ProgressCallback,
     ProgressCallbackWrapper,
-    set_random_state,
 )
 from mostlyai.engine._tabular.common import load_model_weights
 from mostlyai.engine._training_utils import (
@@ -350,11 +349,9 @@ def train(
     device: torch.device | str | None = None,
     workspace_dir: str | Path = "engine-ws",
     update_progress: ProgressCallback | None = None,
-    random_state: int | None = None,
 ):
     _LOG.info("TRAIN_TABULAR started")
     t0 = time.time()
-    set_random_state(random_state)
     workspace_dir = ensure_workspace_dir(workspace_dir)
     workspace = Workspace(workspace_dir)
     with ProgressCallbackWrapper(

{mostlyai_engine-1.4.0 → mostlyai_engine-1.4.1}/mostlyai/engine/analysis.py

@@ -39,7 +39,6 @@ from mostlyai.engine._common import (
     is_a_list,
     is_sequential,
     read_json,
-    set_random_state,
     write_json,
     TABLE_COLUMN_INFIX,
     ProgressCallback,
@@ -87,6 +86,7 @@ from mostlyai.engine._workspace import (
     ensure_workspace_dir,
     reset_dir,
 )
+from mostlyai.engine.random_state import set_random_state
 
 _LOG = logging.getLogger(__name__)
 
@@ -109,7 +109,6 @@ def analyze(
     differential_privacy: DifferentialPrivacyConfig | None = None,
     workspace_dir: str | Path = "engine-ws",
     update_progress: ProgressCallback | None = None,
-    random_state: int | None = None,
 ) -> None:
     """
     Generates (privacy-safe) column-level statistics of the original data, that has been `split` into the workspace.
@@ -128,7 +127,6 @@ def analyze(
 
     _LOG.info("ANALYZE started")
     t0 = time.time()
-    set_random_state(random_state)
     with ProgressCallbackWrapper(update_progress) as progress:
         # build paths based on workspace dir
         workspace_dir = ensure_workspace_dir(workspace_dir)

{mostlyai_engine-1.4.0 → mostlyai_engine-1.4.1}/mostlyai/engine/encoding.py

@@ -23,7 +23,6 @@ def encode(
     *,
     workspace_dir: str | Path = "engine-ws",
     update_progress: ProgressCallback | None = None,
-    random_state: int | None = None,
 ) -> None:
     """
     Encodes data in the workspace that has already been split and analyzed.
@@ -35,14 +34,13 @@ def encode(
     Args:
         workspace_dir: Directory path for workspace.
         update_progress: Callback for progress updates.
-        random_state: Seed for the random number generators.
     """
     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, random_state=random_state)
+        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, random_state=random_state)
+        return encode_language(workspace_dir=workspace_dir, update_progress=update_progress)

{mostlyai_engine-1.4.0 → mostlyai_engine-1.4.1}/mostlyai/engine/generation.py

@@ -41,7 +41,6 @@ def generate(
     fairness: FairnessConfig | dict | None = None,
     workspace_dir: str | Path = "engine-ws",
     update_progress: ProgressCallback | None = None,
-    random_state: int | None = None,
 ) -> None:
     """
     Generates synthetic data from a trained model.
@@ -64,7 +63,6 @@ def generate(
         fairness: Configuration for fairness constraints. Only applicable for tabular models.
         workspace_dir: Directory path for workspace.
         update_progress: Callback for progress updates.
-        random_state: Seed for the random number generators.
     """
     model_type = resolve_model_type(workspace_dir)
     if model_type == ModelType.tabular:
@@ -84,7 +82,6 @@ def generate(
             device=device,
             workspace_dir=workspace_dir,
             update_progress=update_progress,
-            random_state=random_state,
         )
     else:
         from mostlyai.engine._language.generation import generate as generate_language
@@ -106,5 +103,4 @@ def generate(
             device=device,
             workspace_dir=workspace_dir,
             update_progress=update_progress,
-            random_state=random_state,
         )

mostlyai_engine-1.4.1/mostlyai/engine/random_state.py

@@ -0,0 +1,49 @@
+# 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 logging
+import os
+import random
+import struct
+
+import numpy as np
+import torch
+
+_LOG = logging.getLogger(__name__)
+
+
+def set_random_state(random_state: int | None = None, worker: bool = False):
+    def get_random_int_from_os() -> int:
+        # 32-bit, cryptographically secure random int from os
+        return int(struct.unpack("I", os.urandom(4))[0])
+
+    if worker:  # worker process
+        if "MOSTLYAI_ENGINE_SEED" in os.environ:
+            random_state = int(os.environ["MOSTLYAI_ENGINE_SEED"])
+        else:
+            # don't set seed for worker process if not set in main process
+            return
+    else:  # main process
+        if random_state is not None:
+            _LOG.info(f"Global random_state set to `{random_state}`")
+
+        if random_state is None:
+            random_state = get_random_int_from_os()
+
+        os.environ["MOSTLYAI_ENGINE_SEED"] = str(random_state)
+
+    random.seed(random_state)
+    np.random.seed(random_state)
+    torch.manual_seed(random_state)
+    torch.cuda.manual_seed_all(random_state)

{mostlyai_engine-1.4.0 → mostlyai_engine-1.4.1}/mostlyai/engine/splitting.py

@@ -20,7 +20,7 @@ import time
 import warnings
 from pathlib import Path
 
-from mostlyai.engine._common import ProgressCallbackWrapper, ProgressCallback, set_random_state
+from mostlyai.engine._common import ProgressCallbackWrapper, ProgressCallback
 from mostlyai.engine._dtypes import (
     is_date_dtype,
     is_float_dtype,
@@ -75,10 +75,9 @@ def split(
     tgt_encoding_types: dict[str, str | ModelEncodingType] | None = None,
     ctx_encoding_types: dict[str, str | ModelEncodingType] | None = None,
     n_partitions: int = 1,
-    trn_val_split: float = 0.9,
+    trn_val_split: float = 0.8,
     workspace_dir: str | Path = "engine-ws",
     update_progress: ProgressCallback | None = None,
-    random_state: int | None = None,
 ) -> None:
     """
     Splits the provided original data into training and validation sets, and stores these as partitioned Parquet files.
@@ -105,11 +104,9 @@ def split(
         trn_val_split: Fraction of data to use for training, with the remaining data used for validation.
         workspace_dir: Path to the workspace directory where files will be created.
         update_progress: A custom progress callback.
-        random_state: Seed for the random number generators.
     """
     _LOG.info("SPLIT started")
     t0 = time.time()
-    set_random_state(random_state)
     with ProgressCallbackWrapper(update_progress) as progress:
         # validate input
         if tgt_primary_key and tgt_primary_key not in tgt_data:

{mostlyai_engine-1.4.0 → mostlyai_engine-1.4.1}/mostlyai/engine/training.py

@@ -38,7 +38,6 @@ def train(
     workspace_dir: str | Path = "engine-ws",
     update_progress: ProgressCallback | None = None,
     upload_model_data_callback: Callable | None = None,
-    random_state: int | None = None,
 ) -> None:
     """
     Trains a model with optional early stopping and differential privacy.
@@ -61,7 +60,6 @@ def train(
         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.
-        random_state: Seed for the random number generators.
     """
     model_type = resolve_model_type(workspace_dir)
     if model_type == ModelType.tabular:
@@ -82,7 +80,6 @@ def train(
             model_state_strategy=model_state_strategy,
             device=device,
             max_sequence_window=max_sequence_window if max_sequence_window else args["max_sequence_window"].default,
-            random_state=random_state,
         )
     else:
         from mostlyai.engine._language.training import train as train_language
@@ -104,5 +101,4 @@ def train(
             upload_model_data_callback=upload_model_data_callback,
             model_state_strategy=model_state_strategy,
             device=device,
-            random_state=random_state,
         )

{mostlyai_engine-1.4.0 → mostlyai_engine-1.4.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mostlyai-engine"
-version = "1.4.0"
+version = "1.4.1"
 description = "Synthetic Data Engine"
 authors = [{ name = "MOSTLY AI", email = "dev@mostly.ai" }]
 requires-python = ">=3.10"