mostlyai-engine 1.5.7__tar.gz → 1.6.0__tar.gz

{mostlyai_engine-1.5.7 → mostlyai_engine-1.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mostlyai-engine
-Version: 1.5.7
+Version: 1.6.0
 Summary: Synthetic Data Engine
 Project-URL: homepage, https://github.com/mostly-ai/mostlyai-engine
 Project-URL: repository, https://github.com/mostly-ai/mostlyai-engine
@@ -31,7 +31,7 @@ Requires-Dist: joblib>=1.4.2
 Requires-Dist: json-repair>=0.47.0
 Requires-Dist: numpy>=2.0.0
 Requires-Dist: opacus>=1.5.4
-Requires-Dist: pandas~=2.2.0
+Requires-Dist: pandas>=2.2.0
 Requires-Dist: peft>=0.12.0
 Requires-Dist: psutil<6,>=5.9.5
 Requires-Dist: pyarrow>=16.0.0
@@ -91,7 +91,7 @@ pip install -U 'mostlyai-engine[gpu]'
 On Linux, one can explicitly install the CPU-only variant of torch together with `mostlyai-engine`:
 
 ```bash
-pip install -U torch==2.7.0+cpu torchvision==0.22.0+cpu mostlyai-engine --extra-index-url https://download.pytorch.org/whl/cpu
+pip install -U torch==2.8.0+cpu torchvision==0.23.0+cpu mostlyai-engine --extra-index-url https://download.pytorch.org/whl/cpu
 ```
 
 ## Quick start

{mostlyai_engine-1.5.7 → mostlyai_engine-1.6.0}/README.md

@@ -41,7 +41,7 @@ pip install -U 'mostlyai-engine[gpu]'
 On Linux, one can explicitly install the CPU-only variant of torch together with `mostlyai-engine`:
 
 ```bash
-pip install -U torch==2.7.0+cpu torchvision==0.22.0+cpu mostlyai-engine --extra-index-url https://download.pytorch.org/whl/cpu
+pip install -U torch==2.8.0+cpu torchvision==0.23.0+cpu mostlyai-engine --extra-index-url https://download.pytorch.org/whl/cpu
 ```
 
 ## Quick start

{mostlyai_engine-1.5.7 → mostlyai_engine-1.6.0}/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.7"
+__version__ = "1.6.0"
 
 # suppress specific warning related to os.fork() in multi-threaded processes
 warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*multi-threaded.*fork.*")

{mostlyai_engine-1.5.7 → mostlyai_engine-1.6.0}/mostlyai/engine/_common.py

@@ -884,3 +884,86 @@ def dp_non_rare(value_counts: dict[str, int], epsilon: float, threshold: int = 5
 
 def get_stochastic_rare_threshold(min_threshold: int = 5, noise_multiplier: float = 3) -> int:
     return min_threshold + int(noise_multiplier * np.random.uniform())
+
+
+def get_empirical_probs_for_predictor_init(
+    first_encoded_part: Path, tgt_cardinalities: dict[str, int], is_sequential: bool, alpha: float = 1.0
+) -> dict[str, np.ndarray]:
+    """
+    Calculate empirical probabilities of each sub column from the first partition of encoded data.
+    The probabilities will be used for predictor layer initialization.
+
+    Args:
+        first_encoded_part: Path to the first partition of encoded data.
+        tgt_cardinalities: Mapping from column name to its cardinality.
+        is_sequential: Whether the model is sequential.
+        alpha: Laplace smoothing parameter. If smaller or equal to 0, no smoothing is applied.
+
+    Returns:
+        dict[str, np.ndarray]: Mapping from sub column name to its empirical probabilities.
+    """
+    df_part = pd.read_parquet(first_encoded_part)
+    # for sequential models, we will use the empirical probs of the first time step for weight initialization
+    if is_sequential:
+        for sub_col in df_part.columns:
+            df_part[sub_col] = df_part[sub_col].apply(lambda x: x[0] if isinstance(x, np.ndarray) else x)
+    # check which columns have a separate NaN sub column
+    has_nan_map = {
+        col: f"{col}{PREFIX_SUB_COLUMN}nan" in tgt_cardinalities
+        for col in get_columns_from_cardinalities(tgt_cardinalities)
+    }
+    probs_map: dict[str, np.ndarray] = {}
+    for sub_col, cardinality in tgt_cardinalities.items():
+        col, _ = sub_col.split(PREFIX_SUB_COLUMN)
+        nan_sub_col = f"{col}{PREFIX_SUB_COLUMN}nan"
+        if has_nan_map[col] is True and sub_col != nan_sub_col and (df_part[nan_sub_col] == 0).sum() > 0:
+            # exclude NaN rows from the calculation if
+            # - this column has a separate NaN sub column
+            # - the NaN sub column has at least one non-NaN row
+            # - this sub column is not the NaN sub column
+            df_part_sub_col = df_part.loc[df_part[nan_sub_col] == 0, sub_col]
+        else:
+            df_part_sub_col = df_part[sub_col]
+        # calculate empirical probabilities
+        vc = df_part_sub_col.value_counts()
+        if vc.empty:
+            # fallback to uniform distribution
+            probs_map[sub_col] = np.full(cardinality, 1.0 / cardinality)
+        else:
+            counts = np.zeros(cardinality)
+            for idx, count in vc.items():
+                counts[int(idx)] = float(count)
+            # apply Laplace smoothing
+            alpha = max(0.0, alpha)
+            total = counts.sum() + alpha * len(counts)
+            probs_map[sub_col] = (counts + alpha) / max(total, 1e-12)
+            probs_map[sub_col] = np.clip(probs_map[sub_col], a_min=1e-12, a_max=None)
+    return probs_map
+
+
+def impute_from_non_nan_distribution(values: pd.Series, column_stats: dict) -> tuple[pd.Series, pd.Series]:
+    """
+    Impute NaNs with values from the empirical distributions of non-NaN rows.
+    This is helpful especially in the low-data regime to avoid bias towards strong artificial patterns.
+    It is applied when encoding columns with the following encoding types:
+    - TABULAR_NUMERIC_DIGIT
+    - TABULAR_DATETIME
+    - TABULAR_LAT_LONG
+    - TABULAR_CHARACTER
+
+    Args:
+        values: Series of values before encoding.
+        column_stats: Column statistics.
+
+    Returns:
+        tuple[pd.Series, pd.Series]: The series with imputed values and the mask of NaNs.
+    """
+    nan_mask = values.isna()
+    vc = values.value_counts(normalize=True)
+    if vc.empty:
+        return values, nan_mask.astype(int)
+    probs = vc.values
+    categories = vc.index
+    # NOTE: an alternative will be to use the largest remainder method
+    values[nan_mask] = np.random.choice(categories, size=nan_mask.sum(), p=probs)
+    return values, nan_mask.astype(int)

{mostlyai_engine-1.5.7 → mostlyai_engine-1.6.0}/mostlyai/engine/_encoding_types/tabular/character.py

@@ -19,7 +19,12 @@ Character encoding splits any value into its characters, and encodes each positi
 import numpy as np
 import pandas as pd
 
-from mostlyai.engine._common import dp_non_rare, get_stochastic_rare_threshold, safe_convert_string
+from mostlyai.engine._common import (
+    dp_non_rare,
+    get_stochastic_rare_threshold,
+    impute_from_non_nan_distribution,
+    safe_convert_string,
+)
 
 UNKNOWN_TOKEN = "\0"
 MAX_LENGTH_CHARS = 50
@@ -92,15 +97,18 @@ def analyze_reduce_character(
 
 def encode_character(values: pd.Series, stats: dict, _: pd.Series | None = None) -> pd.DataFrame:
     values = safe_convert_string(values)
+    values, nan_mask = impute_from_non_nan_distribution(values, stats)
     max_string_length = stats["max_string_length"]
     df_split = split_sub_columns_character(values, max_string_length)
-    if not stats["has_nan"]:
-        df_split.drop(["nan"], axis=1, inplace=True)
     for idx in range(max_string_length):
         sub_col = f"P{idx}"
         np_codes = np.array(pd.Categorical(df_split[sub_col], categories=stats["codes"][sub_col]).codes)
         np.place(np_codes, np_codes == -1, 0)
         df_split[sub_col] = np_codes
+    if stats["has_nan"]:
+        df_split["nan"] = nan_mask
+    else:
+        df_split.drop(["nan"], axis=1, inplace=True)
     return df_split
 
 

{mostlyai_engine-1.5.7 → mostlyai_engine-1.6.0}/mostlyai/engine/_encoding_types/tabular/datetime.py

@@ -30,6 +30,7 @@ from mostlyai.engine._common import (
     compute_log_histogram,
     dp_approx_bounds,
     get_stochastic_rare_threshold,
+    impute_from_non_nan_distribution,
     safe_convert_datetime,
 )
 from mostlyai.engine._dtypes import is_date_dtype, is_timestamp_dtype
@@ -168,22 +169,26 @@ def encode_datetime(values: pd.Series, stats: dict, _: pd.Series | None = None)
     if stats["max"] is not None:
         reduced_max = pd.Series([stats["max"]], dtype=values.dtype).iloc[0]
         values.loc[values > reduced_max] = reduced_max
+    values, nan_mask = impute_from_non_nan_distribution(values, stats)
     # split to sub_columns
     df = split_sub_columns_datetime(values)
-    is_not_nan = df["nan"] == 0
     # encode values so that each datetime part ranges from 0 to `max_value-min_value`
     for key in DATETIME_PARTS:
         # subtract minimum value
-        df[key] = df[key].where(~is_not_nan, df[key] - stats["min_values"][key])
-        # ensure that any value is mapped onto valid value range
+        df[key] = df[key] - stats["min_values"][key]
+        # clamp to valid range
         df[key] = np.minimum(df[key], stats["max_values"][key] - stats["min_values"][key])
         df[key] = np.maximum(df[key], 0)
-    if not stats["has_nan"]:
-        df.drop(["nan"], inplace=True, axis=1)
+
     if not stats["has_time"]:
         df.drop(["hour", "minute", "second"], inplace=True, axis=1)
     if not stats["has_ms"]:
         df.drop(["ms_E2", "ms_E1", "ms_E0"], inplace=True, axis=1)
+
+    if stats["has_nan"]:
+        df["nan"] = nan_mask
+    else:
+        df.drop(["nan"], inplace=True, axis=1)
     return df
 
 

{mostlyai_engine-1.5.7 → mostlyai_engine-1.6.0}/mostlyai/engine/_encoding_types/tabular/lat_long.py

@@ -19,7 +19,12 @@ import numpy as np
 import pandas as pd
 from numpy.typing import NDArray
 
-from mostlyai.engine._common import dp_non_rare, get_stochastic_rare_threshold, safe_convert_string
+from mostlyai.engine._common import (
+    dp_non_rare,
+    get_stochastic_rare_threshold,
+    impute_from_non_nan_distribution,
+    safe_convert_string,
+)
 from mostlyai.engine._encoding_types.tabular.categorical import (
     CATEGORICAL_UNKNOWN_TOKEN,
     encode_categorical,
@@ -336,21 +341,29 @@ def analyze_reduce_latlong(
 
 def encode_latlong(
     values: pd.Series,
-    column_stats: dict,
+    stats: dict,
     context_keys: pd.Series | None = None,
 ) -> pd.DataFrame:
     values = safe_convert_string(values)
+    # convert invalid entries to NaNs before imputation
+    latitude_longitude = split_str_to_latlong(values)
+    invalid_entry_or_nan_mask = latitude_longitude.isna().any(axis=1)
+    values[invalid_entry_or_nan_mask] = np.nan
+    values, nan_mask = impute_from_non_nan_distribution(values, stats)
     # split to sub_columns
     quads = split_sub_columns_latlong(values)
     encoded_quads = pd.DataFrame()  # empty DF to include all the ModelEncodingType.tabular_categorical quads
-    for quad, value_counts in column_stats["quad_codes"].items():
-        stats = {"codes": value_counts}
-        encoded_quads[quad] = encode_categorical(quads[quad], stats)
+    for quad, value_counts in stats["quad_codes"].items():
+        quad_stats = {"codes": value_counts}
+        encoded_quads[quad] = encode_categorical(quads[quad], quad_stats)
+
+    encoded_quadtile = encode_character(quads["QUADTILE"], stats["quadtile_characters"])
 
-    encoded_quadtile = encode_character(quads["QUADTILE"], column_stats["quadtile_characters"])
-    if column_stats["has_nan"]:
-        encoded_quadtile["nan"] = quads["nan"]
     df = pd.concat([encoded_quads, encoded_quadtile], axis=1)
+    if stats["has_nan"]:
+        # FIXME: consider moving nan sub column to the beginning
+        df["nan"] = nan_mask
+
     return df
 
 

{mostlyai_engine-1.5.7 → mostlyai_engine-1.6.0}/mostlyai/engine/_encoding_types/tabular/numeric.py

@@ -36,6 +36,7 @@ from mostlyai.engine._common import (
     dp_non_rare,
     find_distinct_bins,
     get_stochastic_rare_threshold,
+    impute_from_non_nan_distribution,
     safe_convert_numeric,
 )
 from mostlyai.engine._dtypes import is_float_dtype, is_integer_dtype
@@ -401,29 +402,31 @@ def _encode_numeric_digit(values: pd.Series, stats: dict, _: pd.Series | None =
     # replace extreme values with min/max
     if stats["min"] is not None:
         reduced_min = _type_safe_numeric_series([stats["min"]], dtype).iloc[0]
-        values.loc[values < reduced_min] = reduced_min
+        values = values.where((values.isna()) | (values >= reduced_min), reduced_min)
     if stats["max"] is not None:
         reduced_max = _type_safe_numeric_series([stats["max"]], dtype).iloc[0]
-        values.loc[values > reduced_max] = reduced_max
+        values = values.where((values.isna()) | (values <= reduced_max), reduced_max)
+    values, nan_mask = impute_from_non_nan_distribution(values, stats)
     # split to sub_columns
     df = split_sub_columns_digit(values, stats["max_decimal"], stats["min_decimal"])
-    is_not_nan = df["nan"] == 0
+
     # normalize values to `[0, max_digit-min_digit]`
     for d in np.arange(stats["max_decimal"], stats["min_decimal"] - 1, -1):
         key = f"E{d}"
         # subtract minimum value
-        df[key] = df[key].where(~is_not_nan, df[key] - stats["min_digits"][key])
+        df[key] = df[key] - stats["min_digits"][key]
         # ensure that any value is mapped onto valid value range
         df[key] = np.minimum(df[key], stats["max_digits"][key] - stats["min_digits"][key])
         df[key] = np.maximum(df[key], 0)
-
     # ensure that encoded digits are mapped onto valid value range
     for d in np.arange(stats["max_decimal"], stats["min_decimal"] - 1, -1):
         df[f"E{d}"] = np.minimum(df[f"E{d}"], stats["max_digits"][f"E{d}"])
-    if not stats["has_nan"]:
-        df.drop("nan", inplace=True, axis=1)
     if not stats["has_neg"]:
         df.drop("neg", inplace=True, axis=1)
+    if stats["has_nan"]:
+        df["nan"] = nan_mask
+    else:
+        df.drop("nan", inplace=True, axis=1)
     return df
 
 

{mostlyai_engine-1.5.7 → mostlyai_engine-1.6.0}/mostlyai/engine/_language/__init__.py

@@ -12,6 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import warnings
+
 from mostlyai.engine._language.lstm import register_mostly_lstm_model
 
+# suppress xgrammar max_rollback_tokens deprecation warnings
+warnings.filterwarnings("ignore", message=".*max_rollback_tokens.*", category=DeprecationWarning)
+
 register_mostly_lstm_model()

{mostlyai_engine-1.5.7 → mostlyai_engine-1.6.0}/mostlyai/engine/_language/training.py

@@ -355,7 +355,7 @@ def train(
         max_training_time = max(0.0, max_training_time * 60)  # convert to seconds
         _LOG.info(f"{max_training_time=}s")
         max_epochs = max(0.0, max_epochs)
-        max_epochs_cap = math.ceil((trn_cnt + val_cnt) / 50)
+        max_epochs_cap = math.ceil((trn_cnt + val_cnt) / 25)
         if max_epochs_cap < max_epochs:
             _LOG.info(f"{max_epochs=} -> max_epochs={max_epochs_cap} due to small sample size")
             max_epochs = max_epochs_cap
@@ -425,7 +425,7 @@ def train(
         _LOG.info(f"{model_state_strategy=}")
         if model_state_strategy in [ModelStateStrategy.resume, ModelStateStrategy.reuse]:
             _LOG.info("load existing model weights")
-            torch.serialization.add_safe_globals([np.core.multiarray.scalar, np.dtype, np.dtypes.Float64DType])
+            torch.serialization.add_safe_globals([np._core.multiarray.scalar, np.dtype, np.dtypes.Float64DType])
             resume_from_last_checkpoint = True
             model_id_or_path = workspace.model_path
         else:  # ModelStateStrategy.reset
@@ -637,7 +637,9 @@ def train(
             # therefore, we choose RDP instead as it is more stable and provides comparable privacy guarantees
             dp_accountant = "rdp"  # hard-coded for now
             _LOG.info(f"{dp_config=}, {dp_accountant=}")
-            privacy_engine = PrivacyEngine(accountant=dp_accountant)
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", category=UserWarning, message=".*Secure RNG turned off*")
+                privacy_engine = PrivacyEngine(accountant=dp_accountant)
             if model_state_strategy == ModelStateStrategy.resume and workspace.model_dp_accountant_path.exists():
                 _LOG.info("restore DP accountant state")
                 torch.serialization.add_safe_globals([getattr, PRVAccountant, RDPAccountant, GaussianAccountant])
@@ -707,6 +709,7 @@ def train(
                     # FIXME approximation, should be divided by total sum of number of tokens in the batch
                     #  as in _calculate_per_label_losses, also the final sample may be smaller than the batch size.
                     if with_dp:
+                        warnings.filterwarnings("ignore", category=UserWarning, message="Full backward hook is firing*")
                         step_loss = outputs.loss
                         step_loss.backward()
                     else:

{mostlyai_engine-1.5.7 → mostlyai_engine-1.6.0}/mostlyai/engine/_tabular/argn.py

@@ -702,6 +702,7 @@ class Predictors(nn.Module):
         cardinalities: dict[str, int],
         regressors_dims: dict[str, int],
         device: torch.device,
+        empirical_probs: dict[str, np.ndarray] | None = None,
     ):
         super().__init__()
 
@@ -710,10 +711,23 @@ class Predictors(nn.Module):
         self.device = device
 
         self.predictors = nn.ModuleDict()
+        empirical_probs = empirical_probs or {}
+        if empirical_probs:
+            _LOG.info("initializing predictor bias with empirical log probabilities")
 
         for sub_col, dim_output in self.cardinalities.items():
             dim_input = self.regressors_dims[sub_col]
             self.predictors[sub_col] = nn.Linear(in_features=dim_input, out_features=dim_output, device=self.device)
+            if empirical_probs:
+                nn.init.xavier_uniform_(self.predictors[sub_col].weight)
+                with torch.no_grad():
+                    self.predictors[sub_col].bias.copy_(
+                        torch.as_tensor(
+                            np.log(empirical_probs[sub_col]),
+                            dtype=self.predictors[sub_col].bias.dtype,
+                            device=device,
+                        )
+                    )
 
     def forward(self, x: torch.Tensor, sub_col: str) -> torch.Tensor:
         return self.predictors[sub_col](x)
@@ -739,7 +753,7 @@ def _make_permutation_mask(
         order = torch.randperm(n_cols)
 
     # convert order into a binary mask consisting of 0s and 1s
-    idx = np.argsort(order)
+    idx = torch.argsort(order)
     ones = torch.ones(n_cols, n_cols, dtype=torch.int32, device=device)
     mask = torch.tril(ones, diagonal=-1)  # strict lower triangular matrix
     mask = mask[idx, :][:, idx].bool()  # re-order rows and columns
@@ -872,6 +886,7 @@ class FlatModel(nn.Module):
         column_order: list[str] | None,
         device: torch.device,
         with_dp: bool = False,
+        empirical_probs_for_predictor_init: dict[str, np.ndarray] | None = None,
     ):
         super().__init__()
 
@@ -926,6 +941,7 @@ class FlatModel(nn.Module):
             cardinalities=self.tgt_cardinalities,
             regressors_dims=self.regressors.dims_output,
             device=device,
+            empirical_probs=empirical_probs_for_predictor_init if not with_dp else None,
         )
 
     def _handle_context(
@@ -1155,6 +1171,7 @@ class SequentialModel(nn.Module):
         column_order: list[str] | None,
         device: torch.device,
         with_dp: bool = False,
+        empirical_probs_for_predictor_init: dict[str, np.ndarray] | None = None,
     ):
         super().__init__()
 
@@ -1228,6 +1245,7 @@ class SequentialModel(nn.Module):
             cardinalities=self.tgt_cardinalities,
             regressors_dims=self.regressors.dims_output,
             device=device,
+            empirical_probs=empirical_probs_for_predictor_init if not with_dp else None,
         )
 
     def _repeat_flat_context(self, flat_ctx: list[torch.Tensor], repetition: int) -> list[torch.Tensor]:

{mostlyai_engine-1.5.7 → mostlyai_engine-1.6.0}/mostlyai/engine/_tabular/fairness.py

@@ -182,8 +182,8 @@ def get_fairness_transforms(
         for sensitive_group, target_quantiles in conditional_target_quantiles.items():
             transforms[i][sensitive_group] = partial(
                 torch_interp,
-                xp=torch.as_tensor(target_quantiles[:, i], device=device),
-                fp=torch.as_tensor(marginal_target_quantiles[:, i], device=device),
+                xp=torch.as_tensor(target_quantiles[:, i], device=device).contiguous(),
+                fp=torch.as_tensor(marginal_target_quantiles[:, i], device=device).contiguous(),
             )
 
     _LOG.info("fairness: created transforms")
@@ -217,7 +217,7 @@ def apply_fairness_transforms(
                         sensitive_values == torch.as_tensor(group, device=device).tile((probs.shape[0], 1)),
                         dim=1,
                     ),
-                    transform(probs[:, i]),
+                    transform(probs[:, i].contiguous()),
                     probs[:, i],
                 )
         # normalize transformed probabilities to make sure they sum to 1

{mostlyai_engine-1.5.7 → mostlyai_engine-1.6.0}/mostlyai/engine/_tabular/training.py

@@ -44,6 +44,7 @@ from mostlyai.engine._common import (
     get_cardinalities,
     get_columns_from_cardinalities,
     get_ctx_sequence_length,
+    get_empirical_probs_for_predictor_init,
     get_max_data_points_per_sample,
     get_sequence_length_stats,
     get_sub_columns_from_cardinalities,
@@ -390,7 +391,7 @@ def train(
         max_training_time = max(0.0, max_training_time) * 60  # convert to seconds
         _LOG.info(f"{max_training_time=}s")
         max_epochs = max(0.0, max_epochs)
-        max_epochs_cap = math.ceil((trn_cnt + val_cnt) / 50)
+        max_epochs_cap = math.ceil((trn_cnt + val_cnt) / 25)
         if max_epochs_cap < max_epochs:
             _LOG.info(f"{max_epochs=} -> max_epochs={max_epochs_cap} due to small sample size")
             max_epochs = max_epochs_cap
@@ -437,6 +438,14 @@ def train(
         _LOG.info(f"{max_sequence_window=}")
         ctx_seq_len_median = get_ctx_sequence_length(ctx_stats, key="median")
 
+        empirical_probs_for_predictor_init = (
+            get_empirical_probs_for_predictor_init(
+                workspace.encoded_data_trn.fetch_all()[0], tgt_cardinalities, is_sequential
+            )
+            if not with_dp
+            else None
+        )
+
         # the line below fixes issue with growing epoch time for later epochs
         # https://discuss.pytorch.org/t/training-time-gets-slower-and-slower-on-cpu/145483
         torch.set_flush_denormal(True)
@@ -444,28 +453,24 @@ def train(
         _LOG.info("create training model")
         model_checkpoint = TabularModelCheckpoint(workspace=workspace)
         argn: SequentialModel | FlatModel
+        model_kwargs = {
+            "tgt_cardinalities": tgt_cardinalities,
+            "ctx_cardinalities": ctx_cardinalities,
+            "ctxseq_len_median": ctx_seq_len_median,
+            "model_size": model_size,
+            "column_order": trn_column_order,
+            "device": device,
+            "with_dp": with_dp,  # this flag decides whether the model is initialized with LSTM or DPLSTM layers
+            "empirical_probs_for_predictor_init": empirical_probs_for_predictor_init,
+        }
         if is_sequential:
             argn = SequentialModel(
-                tgt_cardinalities=tgt_cardinalities,
-                ctx_cardinalities=ctx_cardinalities,
+                **model_kwargs,
                 tgt_seq_len_median=tgt_seq_len_median,
                 tgt_seq_len_max=tgt_seq_len_max,
-                ctxseq_len_median=ctx_seq_len_median,
-                model_size=model_size,
-                column_order=trn_column_order,
-                device=device,
-                with_dp=with_dp,  # this flag decides whether the model is initialized with LSTM or DPLSTM layers
             )
         else:
-            argn = FlatModel(
-                tgt_cardinalities=tgt_cardinalities,
-                ctx_cardinalities=ctx_cardinalities,
-                ctxseq_len_median=ctx_seq_len_median,
-                model_size=model_size,
-                column_order=trn_column_order,
-                device=device,
-                with_dp=with_dp,
-            )
+            argn = FlatModel(**model_kwargs)
         _LOG.info(f"model class: {argn.__class__.__name__}")
 
         if isinstance(model_state_strategy, str):
@@ -476,7 +481,7 @@ def train(
         _LOG.info(f"{model_state_strategy=}")
         if model_state_strategy in [ModelStateStrategy.resume, ModelStateStrategy.reuse]:
             _LOG.info("load existing model weights")
-            torch.serialization.add_safe_globals([np.core.multiarray.scalar, np.dtype, np.dtypes.Float64DType])
+            torch.serialization.add_safe_globals([np._core.multiarray.scalar, np.dtype, np.dtypes.Float64DType])
             load_model_weights(model=argn, path=workspace.model_tabular_weights_path, device=device)
         else:  # ModelStateStrategy.reset
             _LOG.info("remove existing checkpoint files")
@@ -634,7 +639,9 @@ def train(
             # therefore, we choose RDP instead as it is more stable and provides comparable privacy guarantees
             dp_accountant = "rdp"  # hard-coded for now
             _LOG.info(f"{dp_config=}, {dp_accountant=}")
-            privacy_engine = PrivacyEngine(accountant=dp_accountant)
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", category=UserWarning, message=".*Secure RNG turned off*")
+                privacy_engine = PrivacyEngine(accountant=dp_accountant)
             if model_state_strategy == ModelStateStrategy.resume and workspace.model_dp_accountant_path.exists():
                 _LOG.info("restore DP accountant state")
                 torch.serialization.add_safe_globals([getattr, PRVAccountant, RDPAccountant, GaussianAccountant])
@@ -700,6 +707,8 @@ def train(
                 # backward pass
                 with warnings.catch_warnings():
                     warnings.filterwarnings("ignore", category=FutureWarning, message="Using a non-full backward hook*")
+                    if with_dp:
+                        warnings.filterwarnings("ignore", category=UserWarning, message="Full backward hook is firing*")
                     step_loss.backward()
                 accumulated_steps += 1
                 # explicitly count the number of processed samples as the actual batch size can vary when DP is on

{mostlyai_engine-1.5.7 → mostlyai_engine-1.6.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mostlyai-engine"
-version = "1.5.7"
+version = "1.6.0"
 description = "Synthetic Data Engine"
 authors = [{ name = "MOSTLY AI", email = "dev@mostly.ai" }]
 requires-python = ">=3.10"
@@ -27,7 +27,7 @@ classifiers = [
 dependencies = [
     "setuptools>=77.0.3",  # similar to vllm 0.10.1.1
     "numpy>=2.0.0",
-    "pandas~=2.2.0",
+    "pandas>=2.2.0",
     "pyarrow>=16.0.0",
     "joblib>=1.4.2",
     "psutil>=5.9.5,<6",  # upgrade when colab psutil is updated