mostlyai-engine 1.5.4__tar.gz → 1.5.6__tar.gz

{mostlyai_engine-1.5.4 → mostlyai_engine-1.5.6}/PKG-INFO

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

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

{mostlyai_engine-1.5.4 → mostlyai_engine-1.5.6}/mostlyai/engine/_common.py

@@ -56,6 +56,11 @@ SLEN_SUB_COLUMN_PREFIX = f"{POSITIONAL_COLUMN}{PREFIX_SUB_COLUMN}slen_"  # seque
 SDEC_SUB_COLUMN_PREFIX = f"{POSITIONAL_COLUMN}{PREFIX_SUB_COLUMN}sdec_"  # sequence index decile
 TABLE_COLUMN_INFIX = "::"  # this should be consistent as in mostly-data and mostlyai-qa
 
+# the latest version of the model uses SIDX/SLEN/RIDX positional column
+DEFAULT_HAS_SLEN = True
+DEFAULT_HAS_RIDX = True
+DEFAULT_HAS_SDEC = False
+
 ANALYZE_MIN_MAX_TOP_N = 1000  # the number of min/max values to be kept from each partition
 
 # the minimal number of min/max values to trigger the reduction; if less, the min/max will be reduced to None
@@ -318,6 +323,7 @@ def get_argn_name(
 def get_cardinalities(
     stats: dict, has_slen: bool | None = None, has_ridx: bool | None = None, has_sdec: bool | None = None
 ) -> dict[str, int]:
+    # the latest version of the model uses SIDX/SLEN/RIDX positional column (applies to sequential model only)
     cardinalities: dict[str, int] = {}
 
     if stats.get("is_sequential", False):
@@ -538,12 +544,11 @@ def decode_positional_column(df_encoded: pd.DataFrame, max_seq_len: int, prefix:
 
 
 def get_positional_cardinalities(
-    max_seq_len: int, has_slen: bool | None, has_ridx: bool | None, has_sdec: bool | None
+    max_seq_len: int, has_slen: bool | None = None, has_ridx: bool | None = None, has_sdec: bool | None = None
 ) -> dict[str, int]:
-    # the latest version of the model uses SIDX/SLEN/RIDX positional column
-    has_slen = has_slen if has_slen is not None else True
-    has_ridx = has_ridx if has_ridx is not None else True
-    has_sdec = has_sdec if has_sdec is not None else False
+    has_slen = has_slen if has_slen is not None else DEFAULT_HAS_SLEN
+    has_ridx = has_ridx if has_ridx is not None else DEFAULT_HAS_RIDX
+    has_sdec = has_sdec if has_sdec is not None else DEFAULT_HAS_SDEC
 
     if max_seq_len < SIDX_RIDX_DIGIT_ENCODING_THRESHOLD:
         # encode positional columns as numeric_discrete

{mostlyai_engine-1.5.4 → mostlyai_engine-1.5.6}/mostlyai/engine/_tabular/argn.py

@@ -226,15 +226,12 @@ class Embedders(nn.Module):
             ),
             None,
         )
-        last_ridx_sub_col = next(
-            (sub_col for sub_col in reversed(self.cardinalities) if sub_col.startswith(RIDX_SUB_COLUMN_PREFIX)), None
-        )
         for sub_col, dim_input in self.cardinalities.items():
             dim_output = _embedding_heuristic(id=self.id(sub_col), model_size=model_size, dim_input=dim_input)
             embedder = nn.Embedding(num_embeddings=dim_input, embedding_dim=dim_output, device=device)
-            # the embeddings of the last slen and ridx sub columns are never used
+            # the embeddings of the last SLEN sub column are never used
             # so we explicitly freeze them to make opacus not complain about "per sample gradient is not initialized"
-            if sub_col in [last_ridx_sub_col, last_slen_sub_col]:
+            if sub_col == last_slen_sub_col:
                 embedder.weight.requires_grad = False
             self.add(sub_column=sub_col, embedder=embedder)
             self.dims.append(dim_output)
@@ -874,6 +871,7 @@ class FlatModel(nn.Module):
         model_size: ModelSizeOrUnits,
         column_order: list[str] | None,
         device: torch.device,
+        with_dp: bool = False,
     ):
         super().__init__()
 
@@ -895,6 +893,7 @@ class FlatModel(nn.Module):
             ctx_cardinalities=self.ctx_cardinalities,
             ctxseq_len_median=self.ctxseq_len_median,
             device=device,
+            with_dp=with_dp,
         )
 
         # sub column embeddings
@@ -1256,26 +1255,28 @@ class SequentialModel(nn.Module):
         if context is None:
             context = self.context_compressor(x)
 
-        # SLEN is only masked for models with RIDX
         has_ridx = any(sub_col.startswith(RIDX_SUB_COLUMN_PREFIX) for sub_col in self.tgt_cardinalities)
-        masked_positional_columns = (
+
+        # SLEN and RIDX are masked for history
+        # NOTE: SLEN is not masked for models without RIDX (backwards compatibility)
+        history_masked_sub_cols = (
             (SLEN_SUB_COLUMN_PREFIX, RIDX_SUB_COLUMN_PREFIX) if has_ridx else (RIDX_SUB_COLUMN_PREFIX,)
         )
+        # SLEN is masked for column embeddings
+        # NOTE: SLEN is not masked for models without RIDX (backwards compatibility)
+        col_embeddings_masked_sub_cols = (SLEN_SUB_COLUMN_PREFIX,) if has_ridx else ()
 
         outputs = {}
         if mode == "trn":
             # forward pass through sub column embedders
             tgt_embeds = self.embedders(x)
-            tgt_embeds_positional_masked = {
-                k: torch.zeros_like(v) if k.startswith(masked_positional_columns) else v for k, v in tgt_embeds.items()
-            }
-
-            # forward pass through column embedders
-            tgt_col_embeds = self.column_embedders(tgt_embeds_positional_masked)
 
             # history
             # time shift: remove last time step; add zeros for first time step; add randoms for all others
-            embeddings = torch.cat(list(tgt_embeds_positional_masked.values()), dim=-1)
+            tgt_embeds_for_history = {
+                k: torch.zeros_like(v) if k.startswith(history_masked_sub_cols) else v for k, v in tgt_embeds.items()
+            }
+            embeddings = torch.cat(list(tgt_embeds_for_history.values()), dim=-1)
             history_in = embeddings[:, :-1, :]
             history_in = nn.ConstantPad2d((0, 0, 1, 0), 0)(history_in)
             history, _ = self.history_compressor(history_in)
@@ -1287,6 +1288,13 @@ class SequentialModel(nn.Module):
             flat_ctx = self._repeat_flat_context(flat_ctx, history.size(1))
             context_history = [torch.cat(flat_ctx + seq_ctx + [history], -1)]
 
+            # forward pass through column embedders
+            tgt_embeds_for_col_embeddings = {
+                k: torch.zeros_like(v) if k.startswith(col_embeddings_masked_sub_cols) else v
+                for k, v in tgt_embeds.items()
+            }
+            tgt_col_embeds = self.column_embedders(tgt_embeds_for_col_embeddings)
+
             # create batch-wise permutation mask
             col_embeddings = torch.cat(list(tgt_col_embeds.values()), dim=-1)
             col_mask = _make_permutation_mask(
@@ -1327,7 +1335,9 @@ class SequentialModel(nn.Module):
             return outputs, {}
 
         else:  # mode == "gen"
+            is_0th_step = False
             if history is None or history_state is None:
+                is_0th_step = True
                 # initialize history
                 history_in = torch.cat(
                     [
@@ -1401,24 +1411,32 @@ class SequentialModel(nn.Module):
                     )
 
                 # update timestep output
+                if is_0th_step and sub_col.startswith(RIDX_SUB_COLUMN_PREFIX):
+                    # overwrite output for RIDX sub-columns on 0th step with SLEN sub-columns
+                    slen_sub_col = sub_col.replace(RIDX_SUB_COLUMN_PREFIX, SLEN_SUB_COLUMN_PREFIX)
+                    out = outputs[slen_sub_col] if slen_sub_col in outputs else out
                 outputs[sub_col] = out
 
                 # update current sub column embedding
                 tgt_embeds[sub_col] = self.embedders.get(sub_col)(out)
-                tgt_embeds_positional_masked = {
-                    k: torch.zeros_like(v) if k.startswith(masked_positional_columns) else v
+                tgt_embeds_for_history = {
+                    k: torch.zeros_like(v) if k.startswith(history_masked_sub_cols) else v
+                    for k, v in tgt_embeds.items()
+                }
+                tgt_embeds_for_col_embeddings = {
+                    k: torch.zeros_like(v) if k.startswith(col_embeddings_masked_sub_cols) else v
                     for k, v in tgt_embeds.items()
                 }
 
                 # update current column embedding
                 if sub_col in self.tgt_last_sub_cols:
                     col_sub_cols = self.tgt_column_sub_columns[lookup.col_name]
-                    col_embed_in = torch.cat([tgt_embeds_positional_masked[sc] for sc in col_sub_cols], dim=-1)
+                    col_embed_in = torch.cat([tgt_embeds_for_col_embeddings[sc] for sc in col_sub_cols], dim=-1)
                     tgt_col_embeds[lookup.col_name] = self.column_embedders.get(lookup.col_name)(col_embed_in)
                     col_embeddings = torch.cat(list(tgt_col_embeds.values()), dim=-1)
 
             # update history and hidden state
-            history_in = torch.cat([v for v in tgt_embeds_positional_masked.values()], dim=-1)
+            history_in = torch.cat(list(tgt_embeds_for_history.values()), dim=-1)
             history, history_state = self.history_compressor(history_in, history_state=history_state)
 
             # order outputs according to tgt_sub_columns

mostlyai_engine-1.5.6/mostlyai/engine/_tabular/common.py

@@ -0,0 +1,37 @@
+# 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 time
+from pathlib import Path
+
+import torch
+
+_LOG = logging.getLogger(__name__)
+
+DPLSTM_SUFFIXES: tuple = ("ih.weight", "ih.bias", "hh.weight", "hh.bias")
+
+
+def load_model_weights(model: torch.nn.Module, path: Path, device: torch.device) -> None:
+    t0 = time.time()
+    incompatible_keys = model.load_state_dict(torch.load(f=path, map_location=device, weights_only=True), strict=False)
+    missing_keys = incompatible_keys.missing_keys
+    unexpected_keys = incompatible_keys.unexpected_keys
+    # for DP-trained models, we expect extra keys from the DPLSTM layers (which is fine to ignore because we use standard LSTM layers during generation)
+    # but if there're any other missing or unexpected keys, an error should be raised
+    if len(missing_keys) > 0 or any(not k.endswith(DPLSTM_SUFFIXES) for k in unexpected_keys):
+        raise RuntimeError(
+            f"failed to load model weights due to incompatibility: {missing_keys = }, {unexpected_keys = }"
+        )
+    _LOG.info(f"loaded model weights in {time.time() - t0:.2f}s")

{mostlyai_engine-1.5.4 → mostlyai_engine-1.5.6}/mostlyai/engine/_tabular/generation.py

@@ -30,6 +30,9 @@ from mostlyai.engine._common import (
     ARGN_TABLE,
     CTXFLT,
     CTXSEQ,
+    DEFAULT_HAS_RIDX,
+    DEFAULT_HAS_SDEC,
+    DEFAULT_HAS_SLEN,
     POSITIONAL_COLUMN,
     RIDX_SUB_COLUMN_PREFIX,
     SDEC_SUB_COLUMN_PREFIX,
@@ -695,6 +698,8 @@ def generate(
                 has_slen = any(SLEN_SUB_COLUMN_PREFIX in k for k in model_units.keys())
                 has_ridx = any(RIDX_SUB_COLUMN_PREFIX in k for k in model_units.keys())
                 has_sdec = any(SDEC_SUB_COLUMN_PREFIX in k for k in model_units.keys())
+            else:
+                has_slen, has_ridx, has_sdec = DEFAULT_HAS_SLEN, DEFAULT_HAS_RIDX, DEFAULT_HAS_SDEC
 
         tgt_cardinalities = get_cardinalities(tgt_stats, has_slen, has_ridx, has_sdec)
         ctx_cardinalities = get_cardinalities(ctx_stats)
@@ -874,11 +879,14 @@ def generate(
         no_of_model_params = get_no_of_model_parameters(model)
         _LOG.info(f"{no_of_model_params=}")
 
-        load_model_weights(
-            model=model,
-            path=workspace.model_tabular_weights_path,
-            device=device,
-        )
+        if workspace.model_tabular_weights_path.exists():
+            load_model_weights(
+                model=model,
+                path=workspace.model_tabular_weights_path,
+                device=device,
+            )
+        else:
+            _LOG.warning("model weights not found; generating data with an untrained model")
 
         model.to(device)
         model.eval()
@@ -1023,7 +1031,7 @@ def generate(
 
                     # fix SLEN by propagating sampled SLEN from first step
                     slen_vals = {}
-                    if seq_step > 0:
+                    if has_slen and seq_step > 0:
                         slen = out_df[SLEN_SUB_COLUMN_PREFIX]
                         slen = encode_positional_column(slen, max_seq_len=seq_len_max, prefix=SLEN_SUB_COLUMN_PREFIX)
                         slen_vals = {
@@ -1102,16 +1110,14 @@ def generate(
                     out_df[SIDX_SUB_COLUMN_PREFIX] = decode_positional_column(
                         out_df, seq_len_max, prefix=SIDX_SUB_COLUMN_PREFIX
                     )
-                    out_df[SLEN_SUB_COLUMN_PREFIX] = decode_positional_column(
-                        out_df, seq_len_max, prefix=SLEN_SUB_COLUMN_PREFIX
-                    )
-                    out_df[SLEN_SUB_COLUMN_PREFIX] = out_df[SLEN_SUB_COLUMN_PREFIX].clip(lower=seq_len_min)
+                    if has_slen:
+                        out_df[SLEN_SUB_COLUMN_PREFIX] = decode_positional_column(
+                            out_df, seq_len_max, prefix=SLEN_SUB_COLUMN_PREFIX
+                        )
+                        out_df[SLEN_SUB_COLUMN_PREFIX] = out_df[SLEN_SUB_COLUMN_PREFIX].clip(lower=seq_len_min)
                     if has_ridx:
-                        # set RIDX to SLEN for first step; beyond that decode RIDX
-                        out_df[RIDX_SUB_COLUMN_PREFIX] = (
-                            decode_positional_column(out_df, seq_len_max, prefix=RIDX_SUB_COLUMN_PREFIX)
-                            if seq_step > 0
-                            else out_df[SLEN_SUB_COLUMN_PREFIX]
+                        out_df[RIDX_SUB_COLUMN_PREFIX] = decode_positional_column(
+                            out_df, seq_len_max, prefix=RIDX_SUB_COLUMN_PREFIX
                         )
                         out_df[RIDX_SUB_COLUMN_PREFIX] = out_df[RIDX_SUB_COLUMN_PREFIX].clip(
                             lower=seq_len_min - seq_step, upper=seq_len_max

{mostlyai_engine-1.5.4 → mostlyai_engine-1.5.6}/mostlyai/engine/_tabular/training.py

@@ -462,6 +462,7 @@ def train(
                 model_size=model_size,
                 column_order=trn_column_order,
                 device=device,
+                with_dp=with_dp,
             )
         _LOG.info(f"model class: {argn.__class__.__name__}")
 

{mostlyai_engine-1.5.4 → mostlyai_engine-1.5.6}/pyproject.toml

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

mostlyai_engine-1.5.4/mostlyai/engine/_tabular/common.py

@@ -1,30 +0,0 @@
-# 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 time
-from pathlib import Path
-
-import torch
-
-_LOG = logging.getLogger(__name__)
-
-
-def load_model_weights(model: torch.nn.Module, path: Path, device: torch.device):
-    try:
-        t00 = time.time()
-        model.load_state_dict(torch.load(f=path, map_location=device, weights_only=True))
-        _LOG.info(f"loaded model weights in {time.time() - t00:.2f}s")
-    except Exception as e:
-        _LOG.warning(f"failed to load model weights: {e}")