mostlyai-engine 1.3.1__tar.gz → 1.3.3__tar.gz

{mostlyai_engine-1.3.1 → mostlyai_engine-1.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mostlyai-engine
-Version: 1.3.1
+Version: 1.3.3
 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.3.1 → mostlyai_engine-1.3.3}/mostlyai/engine/__init__.py

@@ -22,7 +22,7 @@ from mostlyai.engine.generation import generate
 
 
 __all__ = ["split", "analyze", "encode", "train", "generate", "init_logging"]
-__version__ = "1.3.1"
+__version__ = "1.3.3"
 
 # suppress specific warning related to os.fork() in multi-threaded processes
 warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*multi-threaded.*fork.*")

{mostlyai_engine-1.3.1 → mostlyai_engine-1.3.3}/mostlyai/engine/_tabular/encoding.py

@@ -81,7 +81,7 @@ def encode(
                 output_path=workspace.encoded_data_path,
                 ctx_partition_file=ctx_pqt_partitions[i] if has_context else None,
                 ctx_stats=ctx_stats if has_context else None,
-                n_jobs=min(cpu_count() - 1, 16),
+                n_jobs=min(16, max(1, cpu_count() - 1)),
             )
             progress.update(completed=i, total=len(tgt_pqt_partitions) + 1)
     _LOG.info(f"ENCODE_TABULAR finished in {time.time() - t0:.2f}s")

{mostlyai_engine-1.3.1 → mostlyai_engine-1.3.3}/mostlyai/engine/_tabular/generation.py

@@ -720,7 +720,7 @@ def generate(
         # read model config
         model_units = model_configs.get("model_units") or ModelSize.M
         _LOG.debug(f"{model_units=}")
-        enable_flexible_generation = model_configs.get("enable_flexible_generation")
+        enable_flexible_generation = model_configs.get("enable_flexible_generation", True)
         _LOG.info(f"{enable_flexible_generation=}")
 
         # resolve device
@@ -755,6 +755,22 @@ def generate(
             fairness=fairness,
         )
         _LOG.info(f"{gen_column_order=}")
+        if not enable_flexible_generation:
+            # check if resolved column order is the same as the one from training
+            trn_column_order = [
+                get_argn_name(
+                    argn_processor=tgt_stats["columns"][col][ARGN_PROCESSOR],
+                    argn_table=tgt_stats["columns"][col][ARGN_TABLE],
+                    argn_column=tgt_stats["columns"][col][ARGN_COLUMN],
+                )
+                for col in tgt_stats["columns"].keys()
+            ]
+            _LOG.info(f"{trn_column_order=}")
+            if gen_column_order != trn_column_order:
+                raise ValueError(
+                    "The column order for generation does not match the column order from training, due to seed, rebalancing, fairness or imputation configs. "
+                    "A change in column order is only permitted for models that were trained with `enable_flexible_generation=True`."
+                )
 
         _LOG.info(f"{rare_category_replacement_method=}")
         rare_token_fixed_probs = _fix_rare_token_probs(tgt_stats, rare_category_replacement_method)

{mostlyai_engine-1.3.1 → mostlyai_engine-1.3.3}/mostlyai/engine/analysis.py

@@ -162,7 +162,7 @@ def analyze(
                 ctx_encoding_types=ctx_encoding_types,
                 ctx_primary_key=ctx_primary_key if has_context else None,
                 ctx_root_key=ctx_root_key,
-                n_jobs=min(cpu_count() - 1, 16),
+                n_jobs=min(16, max(1, cpu_count() - 1)),
             )
             progress.update(completed=i, total=len(tgt_pqt_partitions) + 1)
 

{mostlyai_engine-1.3.1 → mostlyai_engine-1.3.3}/mostlyai/engine/training.py

@@ -26,8 +26,8 @@ from mostlyai.engine._workspace import resolve_model_type
 def train(
     *,
     model: str | None = None,
-    max_training_time: float = 14400.0,  # 10 days
-    max_epochs: float = 100.0,  # 100 epochs
+    max_training_time: float | None = 14400.0,  # 10 days
+    max_epochs: float | None = 100.0,  # 100 epochs
     batch_size: int | None = None,
     gradient_accumulation_steps: int | None = None,
     enable_flexible_generation: bool = True,
@@ -48,8 +48,8 @@ def train(
 
     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.
+        max_training_time: Maximum training time in minutes. If None, defaults to 10 days.
+        max_epochs: Maximum number of training epochs. If None, defaults to 100 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.
@@ -65,11 +65,12 @@ def train(
     if model_type == ModelType.tabular:
         from mostlyai.engine._tabular.training import train as train_tabular
 
+        args = inspect.signature(train_tabular).parameters
         train_tabular(
-            model=model if model else inspect.signature(train_tabular).parameters["model"].default,
+            model=model if model else args["model"].default,
             workspace_dir=workspace_dir,
-            max_training_time=max_training_time,
-            max_epochs=max_epochs,
+            max_training_time=max_training_time if max_training_time else args["max_training_time"].default,
+            max_epochs=max_epochs if max_epochs else args["max_epochs"].default,
             batch_size=batch_size,
             gradient_accumulation_steps=gradient_accumulation_steps,
             enable_flexible_generation=enable_flexible_generation,
@@ -78,9 +79,7 @@ def train(
             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,
+            max_sequence_window=max_sequence_window if max_sequence_window else args["max_sequence_window"].default,
         )
     else:
         from mostlyai.engine._language.training import train as train_language
@@ -88,11 +87,12 @@ def train(
         if max_sequence_window is not None:
             raise ValueError("max_sequence_window is not supported for language models")
 
+        args = inspect.signature(train_language).parameters
         train_language(
-            model=model if model else inspect.signature(train_language).parameters["model"].default,
+            model=model if model else args["model"].default,
             workspace_dir=workspace_dir,
-            max_training_time=max_training_time,
-            max_epochs=max_epochs,
+            max_training_time=max_training_time if max_training_time else args["max_training_time"].default,
+            max_epochs=max_epochs if max_epochs else args["max_epochs"].default,
             batch_size=batch_size,
             gradient_accumulation_steps=gradient_accumulation_steps,
             enable_flexible_generation=enable_flexible_generation,

{mostlyai_engine-1.3.1 → mostlyai_engine-1.3.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mostlyai-engine"
-version = "1.3.1"
+version = "1.3.3"
 description = "Synthetic Data Engine"
 authors = [{ name = "MOSTLY AI", email = "dev@mostly.ai" }]
 requires-python = ">=3.10"
@@ -54,20 +54,20 @@ gpu = [
 
 [dependency-groups]
 dev = [
-    "pytest>=7.1.2,<8",
-    "ruff==0.9.2",
-    "pre-commit>=4.0.0,<5",
-    "twine>=6.1.0,<7",
-    "ipykernel>=6.25.2",
+    "pytest>=8.0",
+    "ruff>=0.11",  # sync'ed with .pre-commit-config
+    "pre-commit>=4.0",
+    "twine>=6.1",
+    "ipykernel>=6.25",
 ]
 docs = [
-    "mkdocs>=1.5.3",
-    "mkdocstrings[crystal, python]>=0.27.0,<0.28",
-    "mkdocs-material>=9.0.0,<10",
-    "griffe>=1.0.0,<2",
-    "pymdown-extensions>=10.0.0,<11",
-    "griffe-fieldz>=0.2.0,<0.3",
-    "black>=22.1.0",
+    "mkdocs>=1.6",
+    "mkdocstrings[crystal, python]>=0.29",
+    "mkdocs-material>=9.0",
+    "griffe>=1.0",
+    "pymdown-extensions>=10.0",
+    "griffe-fieldz>=0.2",
+    "black>=25.0",
 ]
 
 [project.urls]