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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mostlyai-engine
-Version: 1.5.2
+Version: 1.5.4
 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.4}/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.4"
 
 # 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.4}/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:
@@ -867,6 +864,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 +872,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.4}/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.4}/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.4}/pyproject.toml

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