mostlyai-engine 1.5.1__tar.gz → 1.5.3__tar.gz

{mostlyai_engine-1.5.1 → mostlyai_engine-1.5.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mostlyai-engine
-Version: 1.5.1
+Version: 1.5.3
 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.1 → mostlyai_engine-1.5.3}/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.1"
+__version__ = "1.5.3"
 
 # suppress specific warning related to os.fork() in multi-threaded processes
 warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*multi-threaded.*fork.*")

{mostlyai_engine-1.5.1 → mostlyai_engine-1.5.3}/mostlyai/engine/_common.py

@@ -401,12 +401,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 +413,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:

{mostlyai_engine-1.5.1 → mostlyai_engine-1.5.3}/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.1 → mostlyai_engine-1.5.3}/mostlyai/engine/_tabular/argn.py

@@ -216,8 +216,15 @@ class Embedders(nn.Module):
         self.embedders = nn.ModuleDict()
 
         # embedding layers for each sub column defined in cardinalities
+        has_ridx = any(sub_col.startswith(RIDX_SUB_COLUMN_PREFIX) for sub_col in self.cardinalities)
         last_slen_sub_col = next(
-            (sub_col for sub_col in reversed(self.cardinalities) if sub_col.startswith(SLEN_SUB_COLUMN_PREFIX)), None
+            (
+                sub_col
+                for sub_col in reversed(self.cardinalities)
+                if sub_col.startswith(SLEN_SUB_COLUMN_PREFIX)
+                if has_ridx  # last SLEN sub column is dangling for model with RIDX only
+            ),
+            None,
         )
         last_ridx_sub_col = next(
             (sub_col for sub_col in reversed(self.cardinalities) if sub_col.startswith(RIDX_SUB_COLUMN_PREFIX)), None
@@ -1249,21 +1256,26 @@ 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_SUB_COLUMN_PREFIX, RIDX_SUB_COLUMN_PREFIX) if has_ridx else (RIDX_SUB_COLUMN_PREFIX,)
+        )
+
         outputs = {}
         if mode == "trn":
             # forward pass through sub column embedders
             tgt_embeds = self.embedders(x)
-            tgt_embeds_slen_ridx_masked = {
-                k: torch.zeros_like(v) if k.startswith((RIDX_SUB_COLUMN_PREFIX, SLEN_SUB_COLUMN_PREFIX)) else v
-                for k, v in tgt_embeds.items()
+            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_slen_ridx_masked)
+            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_slen_ridx_masked.values()), dim=-1)
+            embeddings = torch.cat(list(tgt_embeds_positional_masked.values()), dim=-1)
             history_in = embeddings[:, :-1, :]
             history_in = nn.ConstantPad2d((0, 0, 1, 0), 0)(history_in)
             history, _ = self.history_compressor(history_in)
@@ -1393,20 +1405,20 @@ class SequentialModel(nn.Module):
 
                 # update current sub column embedding
                 tgt_embeds[sub_col] = self.embedders.get(sub_col)(out)
-                tgt_embeds_slen_ridx_masked = {
-                    k: torch.zeros_like(v) if k.startswith((SLEN_SUB_COLUMN_PREFIX, RIDX_SUB_COLUMN_PREFIX)) else v
+                tgt_embeds_positional_masked = {
+                    k: torch.zeros_like(v) if k.startswith(masked_positional_columns) 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_slen_ridx_masked[sc] for sc in col_sub_cols], dim=-1)
+                    col_embed_in = torch.cat([tgt_embeds_positional_masked[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_slen_ridx_masked.values()], dim=-1)
+            history_in = torch.cat([v for v in tgt_embeds_positional_masked.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.1 → mostlyai_engine-1.5.3}/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.1 → mostlyai_engine-1.5.3}/pyproject.toml

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