mostlyai-engine 2.7.0__tar.gz → 2.7.2__tar.gz

{mostlyai_engine-2.7.0 → mostlyai_engine-2.7.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: mostlyai-engine
-Version: 2.7.0
+Version: 2.7.2
 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-2.7.0 → mostlyai_engine-2.7.2}/mostlyai/engine/__init__.py

@@ -34,7 +34,7 @@ __all__ = [
     "TabularARGN",
     "LanguageModel",
 ]
-__version__ = "2.7.0"
+__version__ = "2.7.2"
 
 # suppress specific warning related to os.fork() in multi-threaded processes
 warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*multi-threaded.*fork.*")

{mostlyai_engine-2.7.0 → mostlyai_engine-2.7.2}/mostlyai/engine/_encoding_types/language/numeric.py

@@ -43,8 +43,8 @@ def analyze_language_numeric(values: pd.Series, root_keys: pd.Series, _: pd.Seri
 
     # determine max scale
     def count_scale(num: float) -> int:
-        # represent number as fixed point string, remove trailing zeros and decimal point
-        num = format(num, "f").rstrip("0").rstrip(".")
+        # preserve fractional precision without scientific notation or trailing zeros
+        num = np.format_float_positional(num, unique=True, trim="-")
         if "." in num:
             # in case of decimal, return number of digits after decimal point
             return len(num.split(".")[1])
@@ -133,11 +133,13 @@ def encode_language_numeric(values: pd.Series, stats: dict, _: pd.Series | None
 def decode_language_numeric(x: pd.Series, stats: dict[str, str]) -> pd.Series:
     x = pd.to_numeric(x, errors="coerce")
     x = x.round(stats["max_scale"])
+    # Nullable pandas numeric dtypes expose their corresponding NumPy dtype.
+    np_dtype = np.dtype(getattr(x.dtype, "numpy_dtype", x.dtype))
     if stats["min"] is not None:
-        reduced_min = np.dtype(x.dtype).type(stats["min"])
+        reduced_min = np_dtype.type(stats["min"])
         x.loc[x < reduced_min] = reduced_min
     if stats["max"] is not None:
-        reduced_max = np.dtype(x.dtype).type(stats["max"])
+        reduced_max = np_dtype.type(stats["max"])
         x.loc[x > reduced_max] = reduced_max
     dtype = "Int64" if stats["max_scale"] == 0 else float
     return x.astype(dtype)

{mostlyai_engine-2.7.0 → mostlyai_engine-2.7.2}/mostlyai/engine/_encoding_types/tabular/numeric.py

@@ -77,6 +77,7 @@ NUMERIC_DISCRETE_NULL_TOKEN = CATEGORICAL_NULL_TOKEN
 # maximum and minimum precision that is being considered
 NUMERIC_DIGIT_MAX_DECIMAL = 18
 NUMERIC_DIGIT_MIN_DECIMAL = -8
+NUMERIC_DIGIT_MAX_SCALE = 20
 
 
 def _type_safe_numeric_series(numeric_array: np.ndarray | list, pd_dtype: str) -> pd.Series:
@@ -136,7 +137,11 @@ def split_sub_columns_digit(
         # rely on `np.format_float_positional` to determine string representation of absolute values
         values_str = (
             values.abs()
-            .apply(lambda x: np.format_float_positional(x, unique=True, pad_left=50, pad_right=20, precision=20))
+            .apply(
+                lambda x: np.format_float_positional(
+                    x, unique=True, pad_left=50, pad_right=NUMERIC_DIGIT_MAX_SCALE, precision=NUMERIC_DIGIT_MAX_SCALE
+                )
+            )
             # convert to string[pyarrow] for faster processing
             .astype("string[pyarrow]")
             # replace nan with pd.NA for faster processing
@@ -199,7 +204,13 @@ def analyze_numeric(
     max_n = max_values.sort_values(ascending=False).head(ANALYZE_MIN_MAX_TOP_N).astype("float").tolist()
 
     # split values into digits; used for digit numeric encoding, plus to determine precision
-    df_split = split_sub_columns_digit(values)
+    # Extend the default digit window for fine fractions, up to the formatter's precision.
+    max_scale = max(
+        (len(np.format_float_positional(v, unique=True, trim="-").partition(".")[2]) for v in non_na_values),
+        default=0,
+    )
+    min_decimal = min(NUMERIC_DIGIT_MIN_DECIMAL, -min(max_scale, NUMERIC_DIGIT_MAX_SCALE))
+    df_split = split_sub_columns_digit(values, min_decimal=min_decimal)
     is_not_nan = df_split["nan"] == 0
     has_nan = sum(df_split["nan"]) > 0
     has_neg = sum(df_split["neg"]) > 0
@@ -238,9 +249,9 @@ def analyze_reduce_numeric(
     # check if there are negative values
     has_neg = any([j["has_neg"] for j in stats_list])
     # determine precision to apply rounding of sampled values during generation
-    keys = stats_list[0]["max_digits"].keys()
-    min_digits = {k: min([j["min_digits"][k] for j in stats_list]) for k in keys}
-    max_digits = {k: max([j["max_digits"][k] for j in stats_list]) for k in keys}
+    keys = sorted({k for stats in stats_list for k in stats["max_digits"]}, key=lambda k: int(k[1:]), reverse=True)
+    min_digits = {k: min([j["min_digits"].get(k, 0) for j in stats_list]) for k in keys}
+    max_digits = {k: max([j["max_digits"].get(k, 0) for j in stats_list]) for k in keys}
     non_zero_prec = [k for k in keys if max_digits[k] > 0 and k.startswith("E")]
     min_decimal = min([int(k[1:]) for k in non_zero_prec]) if len(non_zero_prec) > 0 else 0
 

{mostlyai_engine-2.7.0 → mostlyai_engine-2.7.2}/mostlyai/engine/_language/lstm.py

@@ -140,13 +140,16 @@ class LSTMFromScratchLMHeadModel(PreTrainedModel, GenerationMixin):
         )
 
     def prepare_inputs_for_generation(
-        self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs
+        self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None, **kwargs
     ) -> dict[str, torch.Tensor]:
         """
         This function is mandatory so that the model is able to use the Hugging Face `.generate()` method.
         Since `.generate()` works with left-padded sequences but the model is trained with right-padded sequences,
         we need to convert the padding side here to make it work properly.
         """
+        # Transformers may omit an all-ones mask for an unpadded batch.
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
         lengths = attention_mask.sum(dim=1)
         return {
             "input_ids": self.left_to_right_padding(input_ids, lengths),

{mostlyai_engine-2.7.0 → mostlyai_engine-2.7.2}/mostlyai/engine/_language/xgrammar_utils.py

@@ -32,14 +32,17 @@ JSON_NULL = "null"
 
 
 def prepend_grammar_root_with_space(grammar: str) -> str:
-    # XGrammar always starts with "{" when enforcing JSON Schema
-    # training is done on strings like ` {...} {...}`
-    # later, generation is done with prompt like ` {...}`
-    # so the first output token must be a space
-    start_of_grammar = 'root ::= "{"'
-    start_of_grammar_with_space = 'root ::= " {"'
-    assert start_of_grammar in grammar
-    return grammar.replace(start_of_grammar, start_of_grammar_with_space)
+    # JSON-schema grammars start at "{". Training strings look like ` {...} {...}`
+    # and generation prompts look like ` {...}`, so the first output token must be a space.
+    # xgrammar <=0.2.3 emits `root ::= "{"`. xgrammar 0.2.7 emits `root ::= (("{"`.
+    replacements = (
+        ('root ::= "{"', 'root ::= " {"'),
+        ('root ::= (("{"', 'root ::= ((" {"'),
+    )
+    for start_of_grammar, start_of_grammar_with_space in replacements:
+        if start_of_grammar in grammar:
+            return grammar.replace(start_of_grammar, start_of_grammar_with_space)
+    raise AssertionError("unrecognized xgrammar root rule")
 
 
 def ensure_seed_can_be_tokenized(seed_data: pd.DataFrame, tokenizer: PreTrainedTokenizerBase) -> pd.DataFrame:

{mostlyai_engine-2.7.0 → mostlyai_engine-2.7.2}/mostlyai/engine/_tabular/encoding.py

@@ -19,7 +19,7 @@ from pathlib import Path
 
 import numpy as np
 import pandas as pd
-from joblib import Parallel, cpu_count, delayed, parallel_config
+from joblib import Parallel, cpu_count, delayed
 
 from mostlyai.engine._common import (
     ARGN_COLUMN,
@@ -238,9 +238,7 @@ def encode_df(
             )
         )
     if delayed_encodes:
-        with parallel_config("loky", n_jobs=n_jobs):
-            df_columns.extend(Parallel()(delayed_encodes))
-
+        df_columns.extend(Parallel(n_jobs=n_jobs)(delayed_encodes))
     df = pd.concat(df_columns, axis=1) if df_columns else pd.DataFrame()
     return df, ctx_primary_key, tgt_context_key
 

{mostlyai_engine-2.7.0 → mostlyai_engine-2.7.2}/mostlyai/engine/_tabular/training.py

@@ -154,6 +154,8 @@ class BatchCollator:
         self.use_nested_ctxseq = use_nested_ctxseq
 
     def __call__(self, batch: list[dict]) -> dict[str, torch.Tensor]:
+        if not batch:
+            return {}
         batch = pd.DataFrame(batch)
         if self.is_sequential and self.max_sequence_window:
             batch = self._slice_sequences(batch, self.max_sequence_window)
@@ -678,6 +680,9 @@ def train(
                 max_grad_norm=dp_config.get("max_grad_norm"),
                 poisson_sampling=True,
             )
+            # Our dictionary collator represents empty batches as {}, including the first
+            # batch, which Opacus cannot infer from the raw dictionary dataset.
+            trn_dataloader.collate_fn = batch_collator
             # this further wraps the dataloader with batch_sampler=BatchSplittingSampler to achieve gradient accumulation
             # it will split the sampled logical batches into smaller sub-batches with batch_size
             trn_dataloader = wrap_data_loader(
@@ -712,22 +717,31 @@ def train(
                 except StopIteration:
                     trn_data_iter = iter(trn_dataloader)
                     step_data = next(trn_data_iter)
-                # forward pass + calculate sample losses
-                step_losses = _calculate_sample_losses(argn, step_data)
-                # FIXME in sequential case, this is an approximation, it should be divided by total sum of masks in the
-                #  entire batch to get the average loss per sample. Less importantly the final sample may be smaller
-                #  than the batch size in both flat and sequential case.
-                # calculate total step loss
-                step_loss = torch.mean(step_losses) / (1 if with_dp else gradient_accumulation_steps)
                 if with_dp:
-                    # opacus handles the gradient accumulation internally
                     optimizer.zero_grad(set_to_none=True)
-                # backward pass
-                with warnings.catch_warnings():
-                    warnings.filterwarnings("ignore", category=FutureWarning, message="Using a non-full backward hook*")
-                    if with_dp:
-                        warnings.filterwarnings("ignore", category=UserWarning, message="Full backward hook is firing*")
-                    step_loss.backward()
+                if with_dp and not step_data:
+                    # Poisson sampling can produce empty logical batches. Bypass the
+                    # RNN forward pass, but retain Opacus noise and accounting below.
+                    for parameter in optimizer.params:
+                        parameter.grad_sample = parameter.new_empty((0, *parameter.shape))
+                    step_losses = torch.empty(0, device=device)
+                else:
+                    # forward pass + calculate sample losses
+                    step_losses = _calculate_sample_losses(argn, step_data)
+                    # FIXME in sequential case, this is an approximation, it should be divided by total sum of masks in the
+                    #  entire batch to get the average loss per sample. Less importantly the final sample may be smaller
+                    #  than the batch size in both flat and sequential case.
+                    step_loss = torch.mean(step_losses) / (1 if with_dp else gradient_accumulation_steps)
+                    # backward pass
+                    with warnings.catch_warnings():
+                        warnings.filterwarnings(
+                            "ignore", category=FutureWarning, message="Using a non-full backward hook*"
+                        )
+                        if with_dp:
+                            warnings.filterwarnings(
+                                "ignore", category=UserWarning, message="Full backward hook is firing*"
+                            )
+                        step_loss.backward()
                 accumulated_steps += 1
                 # explicitly count the number of processed samples as the actual batch size can vary when DP is on
                 samples += step_losses.shape[0]

{mostlyai_engine-2.7.0 → mostlyai_engine-2.7.2}/mostlyai/engine/analysis.py

@@ -17,6 +17,7 @@ Provides analysis functionality of the engine
 """
 
 import logging
+import os
 import time
 from collections.abc import Iterable
 from pathlib import Path
@@ -24,7 +25,7 @@ from typing import Any, Literal
 
 import numpy as np
 import pandas as pd
-from joblib import Parallel, cpu_count, delayed, parallel_config
+from joblib import Parallel, cpu_count, delayed
 
 from mostlyai.engine._common import (
     ANALYZE_REDUCE_MIN_MAX_N,
@@ -251,18 +252,21 @@ def _analyze_partition(
     else:
         ctx_root_keys = ctx_primary_keys.rename("__rkey")
 
+    # Unique row IDs suffice for root-level counts; reuse them across columns.
+    tgt_root_keys = pd.Series(np.arange(len(tgt_df)), index=tgt_df.index, name="root_keys")
+
     # analyze all target columns
-    with parallel_config("loky", n_jobs=n_jobs):
-        results = Parallel()(
-            delayed(_analyze_col)(
-                values=tgt_df[column],
-                encoding_type=encoding_type,
-                context_keys=tgt_context_keys,
-            )
-            for column, encoding_type in tgt_encoding_types.items()
+    results = Parallel(n_jobs=n_jobs)(
+        delayed(_analyze_col)(
+            values=tgt_df[column],
+            root_keys=tgt_root_keys,
+            parent_pid=os.getpid(),
+            encoding_type=encoding_type,
+            context_keys=tgt_context_keys,
         )
-        tgt_column_stats = {column: stats for column, stats in zip(tgt_encoding_types.keys(), results)}
-
+        for column, encoding_type in tgt_encoding_types.items()
+    )
+    tgt_column_stats = {column: stats for column, stats in zip(tgt_encoding_types.keys(), results)}
     # collect target sequence length stats
     tgt_seq_len = _analyze_seq_len(
         tgt_context_keys=tgt_context_keys,
@@ -293,17 +297,16 @@ def _analyze_partition(
 
         # analyze all context columns
         assert isinstance(ctx_encoding_types, dict)
-        with parallel_config("loky", n_jobs=n_jobs):
-            results = Parallel()(
-                delayed(_analyze_col)(
-                    values=ctx_df[column],
-                    encoding_type=encoding_type,
-                    root_keys=ctx_root_keys,
-                )
-                for column, encoding_type in ctx_encoding_types.items()
+        results = Parallel(n_jobs=n_jobs)(
+            delayed(_analyze_col)(
+                values=ctx_df[column],
+                parent_pid=os.getpid(),
+                encoding_type=encoding_type,
+                root_keys=ctx_root_keys,
             )
-            ctx_column_stats = {column: stats for column, stats in zip(ctx_encoding_types.keys(), results)}
-
+            for column, encoding_type in ctx_encoding_types.items()
+        )
+        ctx_column_stats = {column: stats for column, stats in zip(ctx_encoding_types.keys(), results)}
         # persist context stats
         assert isinstance(ctx_stats_path, Path) and ctx_stats_path.exists()
         ctx_stats_file = ctx_stats_path / f"part.{partition_id}.json"
@@ -516,8 +519,10 @@ def _analyze_col(
     encoding_type: ModelEncodingType,
     root_keys: pd.Series | None = None,
     context_keys: pd.Series | None = None,
+    parent_pid: int | None = None,
 ) -> dict:
-    set_random_state(worker=True)
+    if os.getpid() != parent_pid:
+        set_random_state(worker=True)
 
     stats: dict = {"encoding_type": encoding_type}
 
@@ -526,7 +531,7 @@ def _analyze_col(
         return stats
 
     if root_keys is None:
-        root_keys = pd.Series([str(i) for i in range(len(values))], name="root_keys")
+        root_keys = pd.Series(np.arange(len(values)), index=values.index, name="root_keys")
 
     if is_sequential(values):
         # analyze sequential column

{mostlyai_engine-2.7.0 → mostlyai_engine-2.7.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mostlyai-engine"
-version = "2.7.0"
+version = "2.7.2"
 description = "Synthetic Data Engine"
 requires-python = ">=3.11,<3.14"
 readme = "README.md"