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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mostlyai-engine
-Version: 1.5.4
+Version: 1.5.5
 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.5}/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.5"
 
 # 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.5}/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.5}/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)
@@ -1256,26 +1253,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 +1286,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 +1333,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 +1409,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.4 → mostlyai_engine-1.5.5}/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)
@@ -1023,7 +1028,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 +1107,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.5}/pyproject.toml

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