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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mostlyai-engine
-Version: 1.5.2
+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.2 → 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.2"
+__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.2 → 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):
@@ -401,12 +407,11 @@ def get_sub_columns_lookup(
     return sub_cols_lookup
 
 
-class CtxSequenceLengthError(Exception):
-    """Error raised when the cols of the same table do not have the same stats value"""
-
-
 def get_ctx_sequence_length(ctx_stats: dict, key: str) -> dict[str, int]:
-    seq_stats: dict[str, int] = {}
+    """
+    Get the stats of sequence lengths from the first column_stats of each context table
+    """
+    ctxseq_stats: dict[str, int] = {}
 
     for column_stats in ctx_stats.get("columns", {}).values():
         if "seq_len" in column_stats:
@@ -414,12 +419,10 @@ def get_ctx_sequence_length(ctx_stats: dict, key: str) -> dict[str, int]:
                 argn_processor=column_stats[ARGN_PROCESSOR],
                 argn_table=column_stats[ARGN_TABLE],
             )
-            cur_value = seq_stats.get(table)
-            if cur_value and cur_value != column_stats["seq_len"][key]:
-                raise CtxSequenceLengthError()
-            seq_stats[table] = column_stats["seq_len"][key]
+            if table not in ctxseq_stats:
+                ctxseq_stats[table] = column_stats["seq_len"][key]
 
-    return seq_stats
+    return ctxseq_stats
 
 
 def get_max_data_points_per_sample(stats: dict) -> int:
@@ -541,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
@@ -867,6 +869,7 @@ def dp_non_rare(value_counts: dict[str, int], epsilon: float, threshold: int = 5
     noisy_counts = np.clip(np.array(list(value_counts.values())) + noise, 0, None).astype(int)
     for i, cat in enumerate(value_counts):
         value_counts[cat] = noisy_counts[i]
+    # NOTE: total_counts can be 0 in the edge case when the column only has null values
     total_counts = sum(value_counts.values())
 
     # 2. Collect all categories whose noisy count >= threshold
@@ -874,7 +877,7 @@ def dp_non_rare(value_counts: dict[str, int], epsilon: float, threshold: int = 5
 
     # 3. Compute the non-rare ratio
     noisy_total_counts = sum(selected.values())
-    non_rare_ratio = noisy_total_counts / total_counts
+    non_rare_ratio = noisy_total_counts / total_counts if total_counts > 0 else 0
 
     return list(selected.keys()), non_rare_ratio
 

{mostlyai_engine-1.5.2 → mostlyai_engine-1.5.5}/mostlyai/engine/_language/common.py

@@ -97,10 +97,10 @@ def load_base_model_and_config(
     else:
         quantization_config = None
 
-    if device.type == "cuda" and device.index is not None:
-        device_map = str(device)
-    else:
+    if device.type == "cuda" and device.index is None:
         device_map = "auto"
+    else:  # device is `cpu` or `cuda:0` (when using single GPU on a multi-GPU instance)
+        device_map = str(device)
 
     if hasattr(config, "text_config") and hasattr(config, "vision_config"):
         config.text_config.use_cache = use_cache

{mostlyai_engine-1.5.2 → 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.2 → 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.2 → mostlyai_engine-1.5.5}/mostlyai/engine/analysis.py

@@ -339,13 +339,25 @@ def _analyze_reduce(
     stats_list = [read_json(file) for file in stats_files]
     stats: dict[str, Any] = {"columns": {}}
 
+    # check how many context tables have sequential context
+    if mode == "ctx":
+        ctxseq_stats = {}
+        ctxseq_tables = []
+        for column, column_stats in stats_list[0]["columns"].items():
+            if "seq_len" in column_stats:
+                table_name = column.split(TABLE_COLUMN_INFIX)[0]
+                if table_name not in ctxseq_tables:
+                    ctxseq_tables.append(table_name)
+        n_ctxseq_tables = len(ctxseq_tables)
+        _LOG.info(f"{n_ctxseq_tables = }")
+
     encoding_types = {
         column: column_stats.get("encoding_type") for column, column_stats in stats_list[0]["columns"].items()
     }
 
-    # ctx: distribute the privacy budget across all columns
+    # ctx: distribute the privacy budget across all columns + sequence lengths of n_ctxseq_tables
     # tgt: distribute the privacy budget across all columns + sequence length
-    n_dp_splits = len(encoding_types) if mode == "ctx" else len(encoding_types) + 1
+    n_dp_splits = len(encoding_types) + n_ctxseq_tables if mode == "ctx" else len(encoding_types) + 1
     _LOG.info(f"{value_protection = }")
     if value_protection_epsilon is not None and n_dp_splits > 0:
         _LOG.info(f"epsilon for analyzing each column and sequence length: {value_protection_epsilon / n_dp_splits}")
@@ -364,13 +376,13 @@ def _analyze_reduce(
             stats["columns"][column] = {"encoding_type": encoding_type}
             continue
 
-        analyze_reduce_column_args = {
-            "stats_list": column_stats_list,
+        value_protection_args = {
             "value_protection": value_protection,
             "value_protection_epsilon": value_protection_epsilon / n_dp_splits
             if value_protection_epsilon is not None
             else None,
         }
+        analyze_reduce_column_args = {"stats_list": column_stats_list} | value_protection_args
 
         match encoding_type:
             case ModelEncodingType.tabular_categorical:
@@ -413,9 +425,16 @@ def _analyze_reduce(
         if encoding_type in _VALUE_PROTECTION_ENCODING_TYPES:
             stats_col = {"value_protection": value_protection} | stats_col
 
-        is_flat_column = "seq_len" not in column_stats_list[0]
-        if not is_flat_column:
-            stats_col["seq_len"] = _analyze_reduce_seq_len([column_stats_list[0]["seq_len"]])
+        is_ctxseq_column = "seq_len" in column_stats_list[0]
+        if is_ctxseq_column:
+            table_name = column.split(TABLE_COLUMN_INFIX)[0]
+            # only get the lengths from the first column of a ctxseq table and reuse the stats later
+            if table_name not in ctxseq_stats:
+                ctxseq_stats[table_name] = _analyze_reduce_seq_len(
+                    stats_list=[column_stats_list[0]["seq_len"]], **value_protection_args
+                )
+                _LOG.info(f"analyzed sequence length for context table `{table_name}`")
+            stats_col["seq_len"] = ctxseq_stats[table_name]
 
         is_language_column = encoding_type in (
             ModelEncodingType.language_text,

{mostlyai_engine-1.5.2 → mostlyai_engine-1.5.5}/pyproject.toml

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