mostlyai-engine 1.4.2__tar.gz → 1.4.4__tar.gz

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/PKG-INFO

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

@@ -13,17 +13,16 @@
 # limitations under the License.
 import warnings
 
+from mostlyai.engine.analysis import analyze
+from mostlyai.engine.encoding import encode
+from mostlyai.engine.generation import generate
 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
-from mostlyai.engine.encoding import encode
-from mostlyai.engine.generation import generate
-
 
 __all__ = ["split", "analyze", "encode", "train", "generate", "init_logging", "set_random_state"]
-__version__ = "1.4.2"
+__version__ = "1.4.4"
 
 # suppress specific warning related to os.fork() in multi-threaded processes
 warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*multi-threaded.*fork.*")

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/_common.py

@@ -18,6 +18,7 @@ import json
 import logging
 import platform
 import time
+from collections.abc import Callable, Iterable
 from functools import wraps
 from pathlib import Path
 from typing import (
@@ -26,17 +27,14 @@ from typing import (
     NamedTuple,
     Protocol,
 )
-from collections.abc import Callable, Iterable
 
 import numpy as np
 import pandas as pd
 from pydantic import BaseModel
 
-
 from mostlyai.engine._dtypes import is_boolean_dtype, is_float_dtype, is_integer_dtype
 from mostlyai.engine.domain import ModelEncodingType
 
-
 _LOG = logging.getLogger(__name__)
 
 _LOG.info(f"running on Python ({platform.python_version()})")

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/_encoding_types/language/categorical.py

@@ -18,7 +18,7 @@ Categorical encoding for language models.
 
 import pandas as pd
 
-from mostlyai.engine._common import safe_convert_string, STRING
+from mostlyai.engine._common import STRING, safe_convert_string
 from mostlyai.engine._encoding_types.tabular.categorical import analyze_categorical, analyze_reduce_categorical
 
 CATEGORICAL_UNKNOWN_TOKEN = "_RARE_"

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/_encoding_types/language/text.py

@@ -14,7 +14,7 @@
 
 import pandas as pd
 
-from mostlyai.engine._common import safe_convert_string, STRING
+from mostlyai.engine._common import STRING, safe_convert_string
 
 
 def analyze_text(values: pd.Series, root_keys: pd.Series, _: pd.Series | None = None) -> dict:

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/_encoding_types/tabular/categorical.py

@@ -18,7 +18,7 @@ Categorical encoding maps each categorical value to its own integer code.
 
 import pandas as pd
 
-from mostlyai.engine._common import get_stochastic_rare_threshold, safe_convert_string, dp_non_rare
+from mostlyai.engine._common import dp_non_rare, get_stochastic_rare_threshold, safe_convert_string
 
 CATEGORICAL_UNKNOWN_TOKEN = "_RARE_"
 CATEGORICAL_NULL_TOKEN = "<<NULL>>"

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/_language/common.py

@@ -14,13 +14,13 @@
 
 import importlib
 import logging
-
-from peft import PeftConfig, prepare_model_for_kbit_training
-from transformers import AutoConfig, AutoModelForCausalLM, BitsAndBytesConfig, PreTrainedModel, PretrainedConfig
-from mostlyai.engine._language.lstm import LSTMFromScratchConfig
 from pathlib import Path
 
 import torch
+from peft import PeftConfig, prepare_model_for_kbit_training
+from transformers import AutoConfig, AutoModelForCausalLM, BitsAndBytesConfig, PretrainedConfig, PreTrainedModel
+
+from mostlyai.engine._language.lstm import LSTMFromScratchConfig
 
 _LOG = logging.getLogger(__name__)
 

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/_language/encoding.py

@@ -23,15 +23,15 @@ from pandas.core.dtypes.common import is_datetime64_any_dtype
 from tokenizers.pre_tokenizers import ByteLevel
 
 from mostlyai.engine._common import (
-    is_sequential,
+    TABLE_COLUMN_INFIX,
     ProgressCallback,
     ProgressCallbackWrapper,
-    TABLE_COLUMN_INFIX,
+    is_sequential,
 )
-from mostlyai.engine._workspace import ensure_workspace_dir, Workspace, reset_dir
 from mostlyai.engine._encoding_types.language.categorical import encode_language_categorical
-from mostlyai.engine._encoding_types.language.numeric import encode_language_numeric
 from mostlyai.engine._encoding_types.language.datetime import encode_language_datetime
+from mostlyai.engine._encoding_types.language.numeric import encode_language_numeric
+from mostlyai.engine._workspace import Workspace, ensure_workspace_dir, reset_dir
 
 _LOG = logging.getLogger(__name__)
 

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/_language/engine/hf_engine.py

@@ -20,11 +20,10 @@ from os import PathLike
 from pathlib import Path
 
 import torch
-import transformers
-import xgrammar as xgr
 from peft import PeftModel
 from pydantic import BaseModel
 from transformers import AutoTokenizer
+from xgrammar.contrib.hf import LogitsProcessor
 
 from mostlyai.engine._language.common import load_base_model_and_config
 from mostlyai.engine._language.engine.base import EngineMetrics, LanguageEngine
@@ -32,55 +31,6 @@ from mostlyai.engine._language.tokenizer_utils import tokenize_fn
 from mostlyai.engine._language.xgrammar_utils import create_compiled_grammars
 
 
-class XGrammarLogitsProcessor(transformers.LogitsProcessor):
-    """
-    Inspired by [LogitsProcessor](https://github.com/mlc-ai/xgrammar/blob/414473e7c029d0d9e2dfbeacb48afa946d0e3419/python/xgrammar/contrib/hf.py#L14).
-    HuggingFace's XGrammarLogitsProcessor cannot be reused. Logits processors must be initialized for each call to generate().
-    """
-
-    def __init__(self, compiled_grammars: list[xgr.CompiledGrammar]):
-        self.compiled_grammars = compiled_grammars
-        self.vocab_size = self.compiled_grammars[0].tokenizer_info.vocab_size
-        self.batch_size = len(compiled_grammars)
-
-        self.matchers: list[xgr.GrammarMatcher] | None = None
-        self.token_bitmask: torch.Tensor | None = None
-        self.prefilled = False
-
-    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
-        # lazily initialize GrammarMatchers and bitmask
-        if self.matchers is None:
-            self.matchers = [xgr.GrammarMatcher(self.compiled_grammars[i]) for i in range(self.batch_size)]
-            self.token_bitmask = xgr.allocate_token_bitmask(self.batch_size, self.vocab_size)
-
-        if input_ids.shape[0] != self.batch_size:
-            raise RuntimeError(
-                "Expect input_ids.shape[0] to be XGrammarLogitsProcessor.batch_size. "
-                + f"Got {input_ids.shape[0]} for the former, and {self.batch_size} for the latter."
-            )
-
-        if not self.prefilled:
-            # have not sampled a token yet
-            self.prefilled = True
-        else:
-            for i in range(self.batch_size):
-                if not self.matchers[i].is_terminated():
-                    sampled_token = input_ids[i][-1]
-                    assert self.matchers[i].accept_token(sampled_token)
-
-        for i in range(self.batch_size):
-            if not self.matchers[i].is_terminated():
-                self.matchers[i].fill_next_token_bitmask(self.token_bitmask, i)
-
-        device_type = scores.device.type
-        if device_type != "cuda":
-            scores = scores.to("cpu")
-        xgr.apply_token_bitmask_inplace(scores, self.token_bitmask.to(scores.device))
-        if device_type != "cuda":
-            scores = scores.to(device_type)
-        return scores
-
-
 class HuggingFaceEngine(LanguageEngine):
     def __init__(
         self, model_path: PathLike | str, device: torch.device, max_new_tokens: int, tokenizer_max_length: int
@@ -132,7 +82,7 @@ class HuggingFaceEngine(LanguageEngine):
             vocab_size=self._model_config.vocab_size,
             is_peft_adapter=self.is_peft_adapter,
         )
-        self._logits_processors = [XGrammarLogitsProcessor(list(compiled_grammars))]
+        self._logits_processors = [LogitsProcessor(list(compiled_grammars))]
 
     def generate(
         self, text: list[str], sampling_temperature: float, sampling_top_p: float

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/_language/generation.py

@@ -15,24 +15,23 @@
 import contextlib
 import importlib
 import json
+import logging
 import os
 import platform
-from typing import Any
-
-import json_repair
-import logging
 import time
 from pathlib import Path
+from typing import Any
 
+import json_repair
 import pandas as pd
 import torch
 from transformers import PreTrainedTokenizerBase
 
 from mostlyai.engine._common import (
-    persist_data_part,
     FixedSizeSampleBuffer,
     ProgressCallback,
     ProgressCallbackWrapper,
+    persist_data_part,
 )
 from mostlyai.engine._encoding_types.language.categorical import decode_language_categorical
 from mostlyai.engine._encoding_types.language.datetime import decode_language_datetime
@@ -40,9 +39,8 @@ from mostlyai.engine._encoding_types.language.numeric import decode_language_num
 from mostlyai.engine._encoding_types.language.text import decode_text
 from mostlyai.engine._language.common import MAX_LENGTH
 from mostlyai.engine._language.encoding import encode_df
-from mostlyai.engine._language.xgrammar_utils import create_schemas
-from mostlyai.engine._workspace import ensure_workspace_dir, Workspace, reset_dir
-from mostlyai.engine._language.xgrammar_utils import ensure_seed_can_be_tokenized
+from mostlyai.engine._language.xgrammar_utils import create_schemas, ensure_seed_can_be_tokenized
+from mostlyai.engine._workspace import Workspace, ensure_workspace_dir, reset_dir
 from mostlyai.engine.domain import ModelEncodingType, RareCategoryReplacementMethod
 
 INVALID_VALUE = "_INVALID_"  # when JSON parsing fails, the values of target columns will be set to this
@@ -181,6 +179,9 @@ def generate(
         tgt_text_columns = list(tgt_stats["columns"].keys())
         tgt_context_key = tgt_stats["keys"].get("context_key")
         has_context = workspace.ctx_stats.path.exists()
+        model_configs = workspace.model_configs.read()
+        enable_flexible_generation = model_configs.get("enable_flexible_generation", True)
+        _LOG.info(f"{enable_flexible_generation=}")
 
         # resolve potential conflict between sample_seed and sample_size
         if seed_data is not None:
@@ -228,6 +229,15 @@ def generate(
         seed_data = seed_data[[c for c in tgt_text_columns if c in seed_data.columns]]
         _LOG.info(f"{seed_data.shape=}")
 
+        if not enable_flexible_generation:
+            # validate sample_seed maintains the same column order as the one from training
+            seed_columns = seed_data.columns.tolist()
+            if seed_columns != tgt_text_columns[: len(seed_columns)]:
+                raise ValueError(
+                    "The order of columns in the seed data does not match the order of columns from training. "
+                    "A change in column order is only permitted for models that were trained with `enable_flexible_generation=True`."
+                )
+
         # sanity check: at this point sample seed and context data should have the same number of rows
         assert len(seed_data) == len(ctx_data)
 

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/_language/lstm.py

@@ -15,9 +15,9 @@
 import logging
 
 import torch
-from transformers import PreTrainedModel, PretrainedConfig, AutoConfig, AutoModelForCausalLM, AutoModel, GenerationMixin
-from transformers.modeling_outputs import CausalLMOutput
 import torch.nn as nn
+from transformers import AutoConfig, AutoModel, AutoModelForCausalLM, GenerationMixin, PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutput
 
 _LOG = logging.getLogger(__name__)
 

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/_language/tokenizer_utils.py

@@ -12,16 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from collections.abc import Iterator, Mapping
 from dataclasses import dataclass
 from typing import Any
-from collections.abc import Mapping, Iterator
 
-from transformers import DataCollatorForLanguageModeling, BatchEncoding, PreTrainedTokenizerFast, LlamaTokenizerFast
-from transformers.data.data_collator import pad_without_fast_tokenizer_warning, _torch_collate_batch
+from transformers import BatchEncoding, DataCollatorForLanguageModeling, LlamaTokenizerFast, PreTrainedTokenizerFast
+from transformers.data.data_collator import _torch_collate_batch, pad_without_fast_tokenizer_warning
 
 from mostlyai.engine.domain import ModelEncodingType
 
-
 #################
 ### TOKENIZER ###
 #################

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/_language/training.py

@@ -12,59 +12,58 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import gc
 import json
 import logging
 import math
+import os
 import time
-import gc
+import warnings
+from collections.abc import Callable
 from contextlib import nullcontext
 from functools import partial
+from importlib.metadata import version
 from pathlib import Path
-from collections.abc import Callable
-import warnings
 
-from importlib.metadata import version
-import pandas as pd
 import numpy as np
+import pandas as pd
 import torch
+from datasets import Dataset, DatasetDict, disable_progress_bar, load_dataset
+from opacus import GradSampleModule, PrivacyEngine
+from opacus.accountants import GaussianAccountant, PRVAccountant, RDPAccountant
 from opacus.grad_sample import register_grad_sampler
+from opacus.utils.batch_memory_manager import wrap_data_loader
+from peft import LoraConfig, PeftModel
 from torch import nn
 from torch.nn import CrossEntropyLoss
 from torch.optim.lr_scheduler import LRScheduler
-
-from opacus import PrivacyEngine, GradSampleModule
-from opacus.accountants import PRVAccountant, RDPAccountant, GaussianAccountant
-from opacus.utils.batch_memory_manager import wrap_data_loader
-
 from torch.utils.data import DataLoader
+from transformers import (
+    AutoTokenizer,
+    PreTrainedModel,
+)
 
-from mostlyai.engine._common import ProgressCallback, ProgressCallbackWrapper, TABLE_COLUMN_INFIX
+from mostlyai.engine._common import TABLE_COLUMN_INFIX, ProgressCallback, ProgressCallbackWrapper
 from mostlyai.engine._language.common import (
+    MAX_LENGTH,
     is_bf16_supported,
     load_base_model_and_config,
-    MAX_LENGTH,
 )
 from mostlyai.engine._language.encoding import row_to_json
+from mostlyai.engine._language.lstm import LSTMFromScratchConfig, LSTMFromScratchLMHeadModel
 from mostlyai.engine._language.tokenizer_utils import (
-    train_tokenizer,
     MostlyDataCollatorForLanguageModeling,
     tokenize_fn,
+    train_tokenizer,
 )
 from mostlyai.engine._training_utils import (
-    check_early_training_exit,
     EarlyStopper,
     ModelCheckpoint,
     ProgressMessage,
+    check_early_training_exit,
 )
-from mostlyai.engine.domain import ModelStateStrategy, DifferentialPrivacyConfig
 from mostlyai.engine._workspace import Workspace, ensure_workspace_dir
-from datasets import load_dataset, DatasetDict, Dataset, disable_progress_bar
-from transformers import (
-    AutoTokenizer,
-    PreTrainedModel,
-)
-from mostlyai.engine._language.lstm import LSTMFromScratchLMHeadModel, LSTMFromScratchConfig
-from peft import LoraConfig, PeftModel
+from mostlyai.engine.domain import DifferentialPrivacyConfig, ModelStateStrategy
 
 _LOG = logging.getLogger(__name__)
 
@@ -433,6 +432,10 @@ def train(
                 model_config = LSTMFromScratchConfig(vocab_size=len(tokenizer), with_dp=with_dp)
                 model = LSTMFromScratchLMHeadModel(model_config).to(device)
         else:
+            # use MOSTLY_HUGGING_FACE_TOKEN if available
+            if os.getenv("MOSTLY_HUGGING_FACE_TOKEN"):
+                os.environ["HF_TOKEN"] = os.environ["MOSTLY_HUGGING_FACE_TOKEN"]
+
             model, model_config = load_base_model_and_config(
                 model_id_or_path,
                 device=device,
@@ -471,6 +474,10 @@ def train(
                 )
                 model.add_adapter(peft_config)
 
+        # persist model configs
+        model_configs = {"enable_flexible_generation": enable_flexible_generation}
+        workspace.model_configs.write(model_configs)
+
         _LOG.info(f"model loading time: {time.time() - t0:.2f}s")
         model.train()
         no_of_model_params = model.num_parameters()

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/_tabular/argn.py

@@ -23,16 +23,15 @@ This module provides two models for tabular data:
 Each model can learn to generate data given some context data, which can either consist of scalar or sequential data.
 """
 
+import logging
 from enum import Enum
 from functools import partial
-import logging
 from typing import Any, Literal
 
 import numpy as np
 import torch
-from torch import nn
-
 from opacus.layers import DPLSTM
+from torch import nn
 
 from mostlyai.engine._common import (
     CTXFLT,

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/_tabular/common.py

@@ -14,9 +14,10 @@
 
 import logging
 import time
-import torch
 from pathlib import Path
 
+import torch
+
 _LOG = logging.getLogger(__name__)
 
 

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/_tabular/encoding.py

@@ -18,23 +18,23 @@ from pathlib import Path
 
 import numpy as np
 import pandas as pd
-from joblib import Parallel, delayed, parallel_config, cpu_count
+from joblib import Parallel, cpu_count, delayed, parallel_config
 
 from mostlyai.engine._common import (
     ARGN_COLUMN,
     ARGN_PROCESSOR,
     ARGN_TABLE,
+    SDEC_SUB_COLUMN_PREFIX,
     SIDX_SUB_COLUMN_PREFIX,
     SLEN_SUB_COLUMN_PREFIX,
     TGT,
+    ProgressCallback,
+    ProgressCallbackWrapper,
     encode_slen_sidx_sdec,
     get_argn_name,
     get_sequence_length_stats,
     is_a_list,
     is_sequential,
-    SDEC_SUB_COLUMN_PREFIX,
-    ProgressCallback,
-    ProgressCallbackWrapper,
 )
 from mostlyai.engine._encoding_types.tabular.categorical import encode_categorical
 from mostlyai.engine._encoding_types.tabular.character import encode_character
@@ -42,8 +42,8 @@ from mostlyai.engine._encoding_types.tabular.datetime import encode_datetime
 from mostlyai.engine._encoding_types.tabular.itt import encode_itt
 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.domain import ModelEncodingType
 from mostlyai.engine.random_state import set_random_state
 
 _LOG = logging.getLogger(__name__)

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/_tabular/fairness.py

@@ -13,24 +13,23 @@
 # limitations under the License.
 
 import logging
+from collections.abc import Callable
 from functools import partial
+from itertools import product
 from typing import Any, TypedDict
-from collections.abc import Callable
 
-import pandas as pd
 import numpy as np
-
-from itertools import product
+import pandas as pd
 import torch
 
 from mostlyai.engine._common import (
-    get_argn_name,
     ARGN_COLUMN,
     ARGN_PROCESSOR,
     ARGN_TABLE,
+    get_argn_name,
 )
 from mostlyai.engine._encoding_types.tabular.categorical import CATEGORICAL_SUB_COL_SUFFIX, CATEGORICAL_UNKNOWN_TOKEN
-from mostlyai.engine.domain import ModelEncodingType, FairnessConfig
+from mostlyai.engine.domain import FairnessConfig, ModelEncodingType
 
 _LOG = logging.getLogger(__name__)
 

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/_tabular/generation.py

@@ -24,33 +24,19 @@ import numpy as np
 import pandas as pd
 import torch
 
-from mostlyai.engine._memory import get_available_ram_for_heuristics, get_available_vram_for_heuristics
-from mostlyai.engine.domain import (
-    ModelEncodingType,
-    RareCategoryReplacementMethod,
-    RebalancingConfig,
-    ImputationConfig,
-    FairnessConfig,
-)
-from mostlyai.engine._tabular.common import load_model_weights
-
-from mostlyai.engine._tabular.fairness import get_fairness_transforms, FairnessTransforms
-
-from mostlyai.engine._tabular.argn import (
-    FlatModel,
-    ModelSize,
-    SequentialModel,
-    get_no_of_model_parameters,
-)
 from mostlyai.engine._common import (
     ARGN_COLUMN,
     ARGN_PROCESSOR,
     ARGN_TABLE,
     CTXFLT,
     CTXSEQ,
+    SDEC_SUB_COLUMN_PREFIX,
     SIDX_SUB_COLUMN_PREFIX,
     SLEN_SIDX_SDEC_COLUMN,
     SLEN_SUB_COLUMN_PREFIX,
+    FixedSizeSampleBuffer,
+    ProgressCallback,
+    ProgressCallbackWrapper,
     apply_encoding_type_dtypes,
     decode_slen_sidx_sdec,
     encode_slen_sidx_sdec,
@@ -60,16 +46,11 @@ from mostlyai.engine._common import (
     get_ctx_sequence_length,
     get_sequence_length_stats,
     get_sub_columns_from_cardinalities,
-    is_sequential,
-    trim_sequences,
     get_sub_columns_nested_from_cardinalities,
+    is_sequential,
     persist_data_part,
-    FixedSizeSampleBuffer,
-    SDEC_SUB_COLUMN_PREFIX,
-    ProgressCallback,
-    ProgressCallbackWrapper,
+    trim_sequences,
 )
-from mostlyai.engine._tabular.encoding import encode_df, pad_horizontally
 from mostlyai.engine._encoding_types.tabular.categorical import (
     CATEGORICAL_NULL_TOKEN,
     CATEGORICAL_SUB_COL_SUFFIX,
@@ -89,7 +70,24 @@ from mostlyai.engine._encoding_types.tabular.numeric import (
     NUMERIC_DISCRETE_UNKNOWN_TOKEN,
     decode_numeric,
 )
+from mostlyai.engine._memory import get_available_ram_for_heuristics, get_available_vram_for_heuristics
+from mostlyai.engine._tabular.argn import (
+    FlatModel,
+    ModelSize,
+    SequentialModel,
+    get_no_of_model_parameters,
+)
+from mostlyai.engine._tabular.common import load_model_weights
+from mostlyai.engine._tabular.encoding import encode_df, pad_horizontally
+from mostlyai.engine._tabular.fairness import FairnessTransforms, get_fairness_transforms
 from mostlyai.engine._workspace import Workspace, ensure_workspace_dir, reset_dir
+from mostlyai.engine.domain import (
+    FairnessConfig,
+    ImputationConfig,
+    ModelEncodingType,
+    RareCategoryReplacementMethod,
+    RebalancingConfig,
+)
 
 _LOG = logging.getLogger(__name__)
 
@@ -698,7 +696,7 @@ def generate(
         output_path = workspace.generated_data_path
         reset_dir(output_path)
 
-        model_configs = workspace.model_tabular_configs.read()
+        model_configs = workspace.model_configs.read()
         tgt_stats = workspace.tgt_stats.read()
         is_sequential = tgt_stats["is_sequential"]
         _LOG.info(f"{is_sequential=}")

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/_tabular/training.py

@@ -16,37 +16,31 @@ import logging
 import math
 import time
 import warnings
+from collections.abc import Callable
+from importlib.metadata import version
 from itertools import zip_longest
 from pathlib import Path
-from collections.abc import Callable
 
-from importlib.metadata import version
 import numpy as np
 import pandas as pd
 import torch
-from datasets import load_dataset, disable_progress_bar
-from torch.utils.data import DataLoader
-from torch.optim.lr_scheduler import LRScheduler
-
-from torch import nn
-
-from opacus import PrivacyEngine, GradSampleModule
-from opacus.accountants import PRVAccountant, RDPAccountant, GaussianAccountant
+from datasets import disable_progress_bar, load_dataset
+from opacus import GradSampleModule, PrivacyEngine
+from opacus.accountants import GaussianAccountant, PRVAccountant, RDPAccountant
 from opacus.utils.batch_memory_manager import wrap_data_loader
+from torch import nn
+from torch.optim.lr_scheduler import LRScheduler
+from torch.utils.data import DataLoader
 
-from mostlyai.engine._memory import get_available_ram_for_heuristics
-from mostlyai.engine.domain import ModelStateStrategy, DifferentialPrivacyConfig
-from mostlyai.engine._tabular.argn import (
-    FlatModel,
-    ModelSize,
-    SequentialModel,
-    get_model_units,
-    get_no_of_model_parameters,
-)
 from mostlyai.engine._common import (
     CTXFLT,
     CTXSEQ,
+    SDEC_SUB_COLUMN_PREFIX,
+    SIDX_SUB_COLUMN_PREFIX,
+    SLEN_SUB_COLUMN_PREFIX,
     TGT,
+    ProgressCallback,
+    ProgressCallbackWrapper,
     get_cardinalities,
     get_columns_from_cardinalities,
     get_ctx_sequence_length,
@@ -54,20 +48,24 @@ from mostlyai.engine._common import (
     get_sequence_length_stats,
     get_sub_columns_from_cardinalities,
     get_sub_columns_nested_from_cardinalities,
-    SIDX_SUB_COLUMN_PREFIX,
-    SLEN_SUB_COLUMN_PREFIX,
-    SDEC_SUB_COLUMN_PREFIX,
-    ProgressCallback,
-    ProgressCallbackWrapper,
+)
+from mostlyai.engine._memory import get_available_ram_for_heuristics
+from mostlyai.engine._tabular.argn import (
+    FlatModel,
+    ModelSize,
+    SequentialModel,
+    get_model_units,
+    get_no_of_model_parameters,
 )
 from mostlyai.engine._tabular.common import load_model_weights
 from mostlyai.engine._training_utils import (
-    check_early_training_exit,
     EarlyStopper,
     ModelCheckpoint,
     ProgressMessage,
+    check_early_training_exit,
 )
 from mostlyai.engine._workspace import Workspace, ensure_workspace_dir
+from mostlyai.engine.domain import DifferentialPrivacyConfig, ModelStateStrategy
 
 _LOG = logging.getLogger(__name__)
 
@@ -504,7 +502,7 @@ def train(
             "model_units": model_units,
             "enable_flexible_generation": enable_flexible_generation,
         }
-        workspace.model_tabular_configs.write(model_configs)
+        workspace.model_configs.write(model_configs)
 
         # heuristics for batch_size and for initial learn_rate
         mem_available_gb = get_available_ram_for_heuristics() / 1024**3

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/_workspace.py

@@ -14,10 +14,10 @@
 
 import logging
 import os
+from collections.abc import Callable
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
-from collections.abc import Callable
 
 from mostlyai.engine._common import read_json, write_json
 from mostlyai.engine.domain import ModelType
@@ -161,7 +161,7 @@ class Workspace:
         self.model_dp_accountant_path: Path = self.model_path / "dp-accountant.pt"
         self.model_tabular_weights_path: Path = self.model_path / "model-weights.pt"
         self.model_progress_messages_path: Path = self.model_path / "progress-messages.csv"
-        self.model_tabular_configs = make_json_path_desc(parts=model_data + ["model-configs.json"])
+        self.model_configs = make_json_path_desc(parts=model_data + ["model-configs.json"])
 
         # Generate-related
         generated_data = ["SyntheticData"]

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/analysis.py

@@ -18,13 +18,13 @@ Provides analysis functionality of the engine
 
 import logging
 import time
+from collections.abc import Iterable
 from pathlib import Path
 from typing import Any, Literal
-from collections.abc import Iterable
 
 import numpy as np
 import pandas as pd
-from joblib import Parallel, delayed, parallel_config, cpu_count
+from joblib import Parallel, cpu_count, delayed, parallel_config
 
 from mostlyai.engine._common import (
     ANALYZE_REDUCE_MIN_MAX_N,
@@ -33,22 +33,30 @@ from mostlyai.engine._common import (
     ARGN_TABLE,
     CTXFLT,
     CTXSEQ,
+    TABLE_COLUMN_INFIX,
     TGT,
+    ProgressCallback,
+    ProgressCallbackWrapper,
     dp_quantiles,
     get_stochastic_rare_threshold,
     is_a_list,
     is_sequential,
     read_json,
     write_json,
-    TABLE_COLUMN_INFIX,
-    ProgressCallback,
-    ProgressCallbackWrapper,
+)
+from mostlyai.engine._encoding_types.language.categorical import (
+    analyze_language_categorical,
+    analyze_reduce_language_categorical,
 )
 from mostlyai.engine._encoding_types.language.datetime import (
-    analyze_reduce_language_datetime,
     analyze_language_datetime,
+    analyze_reduce_language_datetime,
 )
 from mostlyai.engine._encoding_types.language.numeric import analyze_language_numeric, analyze_reduce_language_numeric
+from mostlyai.engine._encoding_types.language.text import (
+    analyze_reduce_text,
+    analyze_text,
+)
 from mostlyai.engine._encoding_types.tabular.categorical import (
     analyze_categorical,
     analyze_reduce_categorical,
@@ -70,22 +78,13 @@ from mostlyai.engine._encoding_types.tabular.numeric import (
     analyze_numeric,
     analyze_reduce_numeric,
 )
-from mostlyai.engine._encoding_types.language.text import (
-    analyze_text,
-    analyze_reduce_text,
-)
-from mostlyai.engine._encoding_types.language.categorical import (
-    analyze_language_categorical,
-    analyze_reduce_language_categorical,
-)
-from mostlyai.engine.domain import ModelEncodingType, DifferentialPrivacyConfig
-
 from mostlyai.engine._workspace import (
     PathDesc,
     Workspace,
     ensure_workspace_dir,
     reset_dir,
 )
+from mostlyai.engine.domain import DifferentialPrivacyConfig, ModelEncodingType
 from mostlyai.engine.random_state import set_random_state
 
 _LOG = logging.getLogger(__name__)

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/domain.py

@@ -14,7 +14,7 @@
 
 from enum import Enum
 
-from pydantic import BaseModel, Field, ConfigDict, field_validator
+from pydantic import BaseModel, ConfigDict, Field, field_validator
 
 
 class CustomBaseModel(BaseModel):

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/generation.py

@@ -13,17 +13,18 @@
 # limitations under the License.
 
 from pathlib import Path
+
 import pandas as pd
 
 from mostlyai.engine._common import ProgressCallback
+from mostlyai.engine._workspace import resolve_model_type
 from mostlyai.engine.domain import (
+    FairnessConfig,
+    ImputationConfig,
     ModelType,
     RareCategoryReplacementMethod,
     RebalancingConfig,
-    ImputationConfig,
-    FairnessConfig,
 )
-from mostlyai.engine._workspace import resolve_model_type
 
 
 def generate(

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/splitting.py

@@ -16,30 +16,28 @@
 Split original data for training and validation.
 """
 
+import logging
 import time
 import warnings
 from pathlib import Path
 
-from mostlyai.engine._common import ProgressCallbackWrapper, ProgressCallback
+import numpy as np
+import pandas as pd
+
+from mostlyai.engine._common import ProgressCallback, ProgressCallbackWrapper
 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,
 )
+from mostlyai.engine.domain import ModelEncodingType, ModelType
 
 warnings.simplefilter(action="ignore", category=UserWarning)
 

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/training.py

@@ -13,14 +13,14 @@
 # limitations under the License.
 
 import inspect
-from pathlib import Path
 from collections.abc import Callable
+from pathlib import Path
 
 import torch
 
 from mostlyai.engine._common import ProgressCallback
-from mostlyai.engine.domain import ModelType, ModelStateStrategy, DifferentialPrivacyConfig
 from mostlyai.engine._workspace import resolve_model_type
+from mostlyai.engine.domain import DifferentialPrivacyConfig, ModelStateStrategy, ModelType
 
 
 def train(

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mostlyai-engine"
-version = "1.4.2"
+version = "1.4.4"
 description = "Synthetic Data Engine"
 authors = [{ name = "MOSTLY AI", email = "dev@mostly.ai" }]
 requires-python = ">=3.10"
@@ -64,6 +64,7 @@ docs = [
     "mkdocs>=1.6",
     "mkdocstrings[crystal, python]>=0.29",
     "mkdocs-material>=9.0",
+    "mkdocs-llmstxt>=0.2",
     "griffe>=1.0",
     "pymdown-extensions>=10.0",
     "griffe-fieldz>=0.2",
@@ -94,3 +95,6 @@ build-backend = "hatchling.build"
 [tool.ruff]
 target-version = "py310"
 line-length = 120
+
+[tool.ruff.lint]
+extend-select = ["I"]

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/_memory.py

@@ -12,11 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 import os
-import psutil
 import re
-import logging
 
+import psutil
 import torch
 
 _LOG = logging.getLogger(__name__)

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/_training_utils.py

@@ -16,9 +16,9 @@ import abc
 import logging
 import time
 
-from opacus.accountants import IAccountant
 import pandas as pd
 import torch
+from opacus.accountants import IAccountant
 from pydantic import BaseModel, Field, field_validator
 
 from mostlyai.engine._workspace import Workspace

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/encoding.py

@@ -15,8 +15,8 @@
 from pathlib import Path
 
 from mostlyai.engine._common import ProgressCallback
-from mostlyai.engine.domain import ModelType
 from mostlyai.engine._workspace import resolve_model_type
+from mostlyai.engine.domain import ModelType
 
 
 def encode(

{mostlyai_engine-1.4.2 → mostlyai_engine-1.4.4}/mostlyai/engine/logging.py

@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import sys
 import logging
+import sys
 
 _LOG = logging.getLogger(__name__.rsplit(".", 1)[0])  # get the logger with the root module name (mostlyai.engine)