mostlyai-engine 2.1.0__tar.gz → 2.3.0__tar.gz

{mostlyai_engine-2.1.0 → mostlyai_engine-2.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mostlyai-engine
-Version: 2.1.0
+Version: 2.3.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
@@ -88,7 +88,8 @@ Two model classes with these methods are available:
    * `argn.fit(data)`: Train a TabularARGN model
    * `argn.sample(n_samples)`: Generate samples
    * `argn.predict(target, n_draws, agg_fn)`: Predict a feature
-   * `argn.predict_proba(target, n_draws)`: Estimate probabilities
+   * `argn.predict_proba(target)`: Estimate probabilities
+   * `argn.log_prob(data)`: Compute log likelihood
    * `argn.impute(data)`: Fill missing values
 2. `LanguageModel()`: For semi-structured, flat textual tabular data.
    * `.fit(data)`: Train a Language model
@@ -191,9 +192,6 @@ from sklearn.metrics import accuracy_score, roc_auc_score
 # predict class labels for a categorical
 predictions = argn.predict(data_test, target="income", n_draws=100, agg_fn="mode")
 
-# predict class probabilities for a categorical
-probabilities = argn.predict_proba(data_test, target="income", n_draws=100)
-
 # evaluate performance
 accuracy = accuracy_score(data_test["income"], predictions)
 auc = roc_auc_score(data_test["income"], probabilities[:, 1])
@@ -215,6 +213,46 @@ mae = mean_absolute_error(data_test["age"], predictions)
 print(f"MAE: {mae:.1f} years")
 ```
 
+### Conditional Probabilities
+
+Assess any marginal conditional probability, for one or more target columns:
+
+```python
+# extract class probabilities for a categorical
+argn.predict_proba(
+    X=pd.DataFrame({
+        "age": [25, 30, 35],
+        "sex": ["Male", "Female", "Male"],
+    }),
+    target="income"
+)
+
+# extract bin probabilities for a numerical
+argn.predict_proba(
+    X=pd.DataFrame({
+        # "age": [25, 30, 35],
+        "sex": ["Male", "Female", "Male"],
+        "occupation": ["Craft-repair", "Craft-repair", "Craft-repair"]
+    }),
+    target="capital_gain"
+)
+
+# extract two-way marginals
+argn.predict_proba(
+    X=data_test[["age", "race"]],
+    target=["sex", "income"]
+)
+```
+
+### Log Probability
+
+Compute log likelihood of observations:
+
+```python
+# compute log probability for each observation
+log_probs = argn.log_prob(data_test)
+```
+
 ## TabularARGN for Sequential Data
 
 For sequential data (e.g., time series or event logs), specify the context key:

{mostlyai_engine-2.1.0 → mostlyai_engine-2.3.0}/README.md

@@ -37,7 +37,8 @@ Two model classes with these methods are available:
    * `argn.fit(data)`: Train a TabularARGN model
    * `argn.sample(n_samples)`: Generate samples
    * `argn.predict(target, n_draws, agg_fn)`: Predict a feature
-   * `argn.predict_proba(target, n_draws)`: Estimate probabilities
+   * `argn.predict_proba(target)`: Estimate probabilities
+   * `argn.log_prob(data)`: Compute log likelihood
    * `argn.impute(data)`: Fill missing values
 2. `LanguageModel()`: For semi-structured, flat textual tabular data.
    * `.fit(data)`: Train a Language model
@@ -140,9 +141,6 @@ from sklearn.metrics import accuracy_score, roc_auc_score
 # predict class labels for a categorical
 predictions = argn.predict(data_test, target="income", n_draws=100, agg_fn="mode")
 
-# predict class probabilities for a categorical
-probabilities = argn.predict_proba(data_test, target="income", n_draws=100)
-
 # evaluate performance
 accuracy = accuracy_score(data_test["income"], predictions)
 auc = roc_auc_score(data_test["income"], probabilities[:, 1])
@@ -164,6 +162,46 @@ mae = mean_absolute_error(data_test["age"], predictions)
 print(f"MAE: {mae:.1f} years")
 ```
 
+### Conditional Probabilities
+
+Assess any marginal conditional probability, for one or more target columns:
+
+```python
+# extract class probabilities for a categorical
+argn.predict_proba(
+    X=pd.DataFrame({
+        "age": [25, 30, 35],
+        "sex": ["Male", "Female", "Male"],
+    }),
+    target="income"
+)
+
+# extract bin probabilities for a numerical
+argn.predict_proba(
+    X=pd.DataFrame({
+        # "age": [25, 30, 35],
+        "sex": ["Male", "Female", "Male"],
+        "occupation": ["Craft-repair", "Craft-repair", "Craft-repair"]
+    }),
+    target="capital_gain"
+)
+
+# extract two-way marginals
+argn.predict_proba(
+    X=data_test[["age", "race"]],
+    target=["sex", "income"]
+)
+```
+
+### Log Probability
+
+Compute log likelihood of observations:
+
+```python
+# compute log probability for each observation
+log_probs = argn.log_prob(data_test)
+```
+
 ## TabularARGN for Sequential Data
 
 For sequential data (e.g., time series or event logs), specify the context key:

{mostlyai_engine-2.1.0 → mostlyai_engine-2.3.0}/mostlyai/engine/__init__.py

@@ -34,7 +34,7 @@ __all__ = [
     "TabularARGN",
     "LanguageModel",
 ]
-__version__ = "2.1.0"
+__version__ = "2.3.0"
 
 # suppress specific warning related to os.fork() in multi-threaded processes
 warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*multi-threaded.*fork.*")

{mostlyai_engine-2.1.0 → mostlyai_engine-2.3.0}/mostlyai/engine/_tabular/argn.py

@@ -972,10 +972,63 @@ class FlatModel(nn.Module):
 
         return context
 
+    def _initialize_generation(self, x, batch_size, effective_column_order):
+        """Initialize context, embeddings, and sub-column order for generation/probs mode."""
+        # forward pass through context compressor
+        context = self.context_compressor(x)
+        context = self._handle_context(context)
+
+        # initialize embeddings
+        tgt_embeds = self.embedders.zero_mask(batch_size)
+        tgt_col_embeds = self.column_embedders.zero_mask(batch_size)
+        col_embeddings = torch.cat(list(tgt_col_embeds.values()), dim=-1)
+
+        # determine sub-column order
+        column_order = effective_column_order or self.tgt_columns
+        sub_column_order = [sub_col for col in column_order for sub_col in self.tgt_column_sub_columns[col]]
+
+        return context, tgt_embeds, tgt_col_embeds, col_embeddings, sub_column_order
+
+    def _update_embeddings(self, sub_col, out, tgt_embeds, tgt_col_embeds, col_embeddings):
+        """Update sub-column and column embeddings after setting a value.
+
+        Returns updated col_embeddings if this sub-column completes a column,
+        otherwise returns the unchanged col_embeddings.
+        """
+        lookup = self.tgt_sub_columns_lookup[sub_col]
+
+        # update current sub column embedding
+        tgt_embeds[sub_col] = self.embedders.get(sub_col)(out)
+
+        # update current column embedding if this is the last sub-column
+        if sub_col in self.last_sub_cols:
+            col_sub_cols = self.tgt_column_sub_columns[lookup.col_name]
+            col_embed_in = torch.cat([tgt_embeds[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)
+
+        return col_embeddings
+
+    def _compute_logits(self, sub_col, context, col_embeddings, tgt_embeds):
+        """Compute logits for a sub-column given context and previous embeddings."""
+        lookup = self.tgt_sub_columns_lookup[sub_col]
+
+        # collect previous sub column embeddings for current column
+        prev_sub_col_embeds = [
+            tgt_embeds[sc] for sc in self.tgt_sub_columns[lookup.sub_col_offset : lookup.sub_col_cum]
+        ]
+
+        # regressor + predictor
+        regressor_in = context + [col_embeddings] + prev_sub_col_embeds
+        xs = self.regressors(regressor_in, sub_col)
+        xs = self.predictors(xs, sub_col)
+
+        return xs
+
     def forward(
         self,
         x,
-        mode: Literal["trn", "gen"],
+        mode: Literal["trn", "gen", "probs"],
         batch_size: int | None = None,
         fixed_probs=None,
         fixed_values=None,
@@ -1020,7 +1073,7 @@ class FlatModel(nn.Module):
 
                 # collect previous sub column embeddings for current column
                 prev_sub_col_embeds = [
-                    tgt_embeds[sub_col] for sub_col in self.tgt_sub_columns[lookup.sub_col_offset : lookup.sub_col_cum]
+                    tgt_embeds[sc] for sc in self.tgt_sub_columns[lookup.sub_col_offset : lookup.sub_col_cum]
                 ]
 
                 # regressor
@@ -1033,84 +1086,68 @@ class FlatModel(nn.Module):
                 # update output
                 outputs[sub_col] = xs
 
-        else:  # mode == "gen"
-            # forward pass through context compressor
-            context = self.context_compressor(x)
-            context = self._handle_context(context)
-
-            # initialize sub column embeddings
-            tgt_embeds = self.embedders.zero_mask(batch_size)
-
-            # initialize column embeddings
-            tgt_col_embeds = self.column_embedders.zero_mask(batch_size)
-
-            # concatenate column embeddings
-            col_embeddings = torch.cat(list(tgt_col_embeds.values()), dim=-1)
+            return outputs, {}
 
-            # take sub columns in the specified generation order
-            column_order = effective_column_order or self.tgt_columns
-            sub_column_order = [sub_col for col in column_order for sub_col in self.tgt_column_sub_columns[col]]
+        elif mode == "gen":
+            context, tgt_embeds, tgt_col_embeds, col_embeddings, sub_column_order = self._initialize_generation(
+                x, batch_size, effective_column_order
+            )
 
             for sub_col in sub_column_order:
-                lookup = self.tgt_sub_columns_lookup[sub_col]
-
-                # if sub column is fixed, skip sampling and use that value
+                # handle fixed values
                 if sub_col in fixed_values:
                     out = fixed_values[sub_col]
+                else:
+                    # compute probabilities and sample
+                    logits = self._compute_logits(sub_col, context, col_embeddings, tgt_embeds)
+                    probs_tensor = nn.Softmax(dim=-1)(logits)
 
-                else:  # sample from distribution
-                    # collect previous sub column embeddings for current column
-                    prev_sub_col_embeds = [
-                        tgt_embeds[sub_col]
-                        for sub_col in self.tgt_sub_columns[lookup.sub_col_offset : lookup.sub_col_cum]
-                    ]
-
-                    # regressor
-                    regressor_in = context + [col_embeddings] + prev_sub_col_embeds
-                    xs = self.regressors(regressor_in, sub_col)
-
-                    # predictor
-                    xs = self.predictors(xs, sub_col)
-
-                    # softmax to probs
-                    xs = nn.Softmax(dim=-1)(xs)
-
-                    # keep probabilities (used e.g. for fairness)
+                    # optionally keep probabilities
                     if sub_col in return_probs:
-                        probs[sub_col] = xs
+                        probs[sub_col] = probs_tensor
 
-                    # apply fairness transform when generating the target sub column
+                    # apply fairness transforms
                     if fairness_transforms:
-                        xs = apply_fairness_transforms(sub_col, xs, outputs, fairness_transforms)
+                        probs_tensor = apply_fairness_transforms(sub_col, probs_tensor, outputs, fairness_transforms)
 
                     # sample
                     out = torch.squeeze(
-                        _sample(
-                            probs=xs,
-                            temperature=temperature,
-                            top_p=top_p,
-                            fixed_probs=fixed_probs.get(sub_col),
-                        ),
+                        _sample(probs_tensor, temperature, top_p, fixed_probs.get(sub_col)),
                         dim=-1,
                     )
 
-                # update output
+                # update output and embeddings
                 outputs[sub_col] = out
+                col_embeddings = self._update_embeddings(sub_col, out, tgt_embeds, tgt_col_embeds, col_embeddings)
 
-                # update current sub column embedding
-                tgt_embeds[sub_col] = self.embedders.get(sub_col)(out)
+            # order outputs and return
+            outputs = {sub_col: outputs[sub_col] for sub_col in self.tgt_sub_columns}
+            return outputs, probs
 
-                # update current column embedding
-                if sub_col in self.last_sub_cols:
-                    col_sub_cols = self.tgt_column_sub_columns[lookup.col_name]
-                    col_embed_in = torch.cat([tgt_embeds[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)
+        elif mode == "probs":
+            context, tgt_embeds, tgt_col_embeds, col_embeddings, sub_column_order = self._initialize_generation(
+                x, batch_size, effective_column_order
+            )
 
-            # order outputs according to tgt_sub_columns
-            outputs = {sub_col: outputs[sub_col] for sub_col in self.tgt_sub_columns}
+            for sub_col in sub_column_order:
+                # handle fixed values
+                if sub_col in fixed_values:
+                    out = fixed_values[sub_col]
+                    # update embeddings to maintain correct autoregressive context
+                    col_embeddings = self._update_embeddings(sub_col, out, tgt_embeds, tgt_col_embeds, col_embeddings)
+                else:
+                    # compute probabilities without sampling
+                    logits = self._compute_logits(sub_col, context, col_embeddings, tgt_embeds)
+                    probs_tensor = nn.Softmax(dim=-1)(logits)
+
+                    # apply fixed_probs mask if provided
+                    if sub_col in fixed_probs:
+                        probs_tensor = _sampling_fixed_probs(probs_tensor, fixed_probs[sub_col])
+
+                    # store probabilities (no sampling, no embedding updates)
+                    probs[sub_col] = probs_tensor
 
-        return outputs, probs
+            return {}, probs
 
 
 class AttentionModule(nn.Module):
@@ -1334,8 +1371,7 @@ class SequentialModel(nn.Module):
 
                 # collect previous sub column embeddings for current column
                 prev_sub_col_embeds = {
-                    sub_col: tgt_embeds[sub_col]
-                    for sub_col in self.tgt_sub_columns[lookup.sub_col_offset : lookup.sub_col_cum]
+                    sc: tgt_embeds[sc] for sc in self.tgt_sub_columns[lookup.sub_col_offset : lookup.sub_col_cum]
                 }
                 if sub_col.startswith(RIDX_SUB_COLUMN_PREFIX):
                     # RIDX sub-columns should not see SLEN sub-columns
@@ -1405,8 +1441,7 @@ class SequentialModel(nn.Module):
                 else:  # sample from distribution
                     # collect previous sub column embeddings for current column
                     prev_sub_col_embeds = {
-                        sub_col: tgt_embeds[sub_col]
-                        for sub_col in self.tgt_sub_columns[lookup.sub_col_offset : lookup.sub_col_cum]
+                        sc: tgt_embeds[sc] for sc in self.tgt_sub_columns[lookup.sub_col_offset : lookup.sub_col_cum]
                     }
                     if sub_col.startswith(RIDX_SUB_COLUMN_PREFIX):
                         # RIDX sub-columns should not see SLEN sub-columns

mostlyai_engine-2.3.0/mostlyai/engine/_tabular/common.py

@@ -0,0 +1,315 @@
+# Copyright 2025 MOSTLY AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import time
+from pathlib import Path
+
+import pandas as pd
+import torch
+
+from mostlyai.engine._common import (
+    ARGN_COLUMN,
+    ARGN_PROCESSOR,
+    ARGN_TABLE,
+    CTXFLT,
+    CTXSEQ,
+    get_argn_name,
+)
+from mostlyai.engine._encoding_types.tabular.categorical import (
+    CATEGORICAL_SUB_COL_SUFFIX,
+    CATEGORICAL_UNKNOWN_TOKEN,
+)
+from mostlyai.engine._encoding_types.tabular.numeric import (
+    NUMERIC_BINNED_SUB_COL_SUFFIX,
+    NUMERIC_BINNED_UNKNOWN_TOKEN,
+    NUMERIC_DISCRETE_SUB_COL_SUFFIX,
+    NUMERIC_DISCRETE_UNKNOWN_TOKEN,
+)
+from mostlyai.engine._tabular.encoding import encode_df, pad_ctx_sequences
+from mostlyai.engine.domain import ModelEncodingType, RareCategoryReplacementMethod
+
+_LOG = logging.getLogger(__name__)
+
+# Type alias for fixed probabilities
+CodeProbabilities = dict[int, float]
+
+DPLSTM_SUFFIXES: tuple = ("ih.weight", "ih.bias", "hh.weight", "hh.bias")
+
+
+def load_model_weights(model: torch.nn.Module, path: Path, device: torch.device) -> None:
+    t0 = time.time()
+    incompatible_keys = model.load_state_dict(torch.load(f=path, map_location=device, weights_only=True), strict=False)
+    missing_keys = incompatible_keys.missing_keys
+    unexpected_keys = incompatible_keys.unexpected_keys
+    # for DP-trained models, we expect extra keys from the DPLSTM layers (which is fine to ignore because we use standard LSTM layers during generation)
+    # but if there're any other missing or unexpected keys, an error should be raised
+    if len(missing_keys) > 0 or any(not k.endswith(DPLSTM_SUFFIXES) for k in unexpected_keys):
+        raise RuntimeError(
+            f"failed to load model weights due to incompatibility: {missing_keys = }, {unexpected_keys = }"
+        )
+    _LOG.info(f"loaded model weights in {time.time() - t0:.2f}s")
+
+
+def load_model_artifacts(workspace):
+    """
+    Load model configurations and statistics from workspace.
+
+    Returns:
+        Tuple of (model_config, tgt_stats, ctx_stats, is_sequential)
+    """
+    model_config = workspace.model_configs.read()
+    tgt_stats = workspace.tgt_stats.read()
+    ctx_stats = workspace.ctx_stats.read()
+    is_sequential = tgt_stats["is_sequential"]
+    return model_config, tgt_stats, ctx_stats, is_sequential
+
+
+def resolve_device(device: torch.device | str | None) -> torch.device:
+    """
+    Resolve device to use for inference.
+
+    Args:
+        device: Device specification ('cuda', 'cpu', or None for auto-detect)
+
+    Returns:
+        torch.device instance
+    """
+    if device is None:
+        return torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    return torch.device(device)
+
+
+def create_and_load_model(
+    workspace,
+    is_sequential: bool,
+    tgt_cardinalities: dict,
+    ctx_cardinalities: dict,
+    model_units,
+    ctx_seq_len_median: int | None,
+    column_order: list[str],
+    device: torch.device,
+    seq_len_median: int | None = None,
+    seq_len_max: int | None = None,
+):
+    """
+    Create model, load weights, and prepare for inference.
+
+    Args:
+        workspace: Workspace containing model weights
+        is_sequential: Whether to create SequentialModel or FlatModel
+        tgt_cardinalities: Target column cardinalities
+        ctx_cardinalities: Context column cardinalities
+        model_units: Model size configuration
+        ctx_seq_len_median: Median context sequence length
+        column_order: Order of columns for generation
+        device: Device to load model on
+        seq_len_median: Median sequence length (for sequential models)
+        seq_len_max: Maximum sequence length (for sequential models)
+
+    Returns:
+        Initialized model ready for inference
+    """
+    from mostlyai.engine._tabular.argn import FlatModel, SequentialModel, get_no_of_model_parameters
+
+    _LOG.info("Creating generative model")
+
+    if is_sequential:
+        model = SequentialModel(
+            tgt_cardinalities=tgt_cardinalities,
+            tgt_seq_len_median=seq_len_median,
+            tgt_seq_len_max=seq_len_max,
+            ctx_cardinalities=ctx_cardinalities,
+            ctxseq_len_median=ctx_seq_len_median,
+            model_size=model_units,
+            column_order=column_order,
+            device=device,
+        )
+    else:
+        model = FlatModel(
+            tgt_cardinalities=tgt_cardinalities,
+            ctx_cardinalities=ctx_cardinalities,
+            ctxseq_len_median=ctx_seq_len_median,
+            model_size=model_units,
+            column_order=column_order,
+            device=device,
+        )
+
+    no_of_model_params = get_no_of_model_parameters(model)
+    _LOG.info(f"{no_of_model_params=}")
+
+    if workspace.model_tabular_weights_path.exists():
+        load_model_weights(
+            model=model,
+            path=workspace.model_tabular_weights_path,
+            device=device,
+        )
+    else:
+        _LOG.warning("Model weights not found; using untrained model")
+
+    model.to(device)
+    model.eval()
+
+    return model
+
+
+def prepare_context_inputs(
+    ctx_data: pd.DataFrame,
+    ctx_stats: dict,
+    device: torch.device | str,
+    ctx_primary_key: str | None = None,
+) -> tuple[dict[str, torch.Tensor], pd.DataFrame, str | None]:
+    """
+    Encode context data and prepare tensor inputs for model forward pass.
+
+    Handles both flat context (CTXFLT) and sequential context (CTXSEQ).
+
+    Args:
+        ctx_data: Context DataFrame to encode
+        ctx_stats: Context statistics from training
+        device: Device for tensor placement
+        ctx_primary_key: Optional primary key column for context
+
+    Returns:
+        Tuple of (context_tensors, encoded_dataframe, encoded_primary_key):
+        - context_tensors: Dict of CTXFLT/* and CTXSEQ/* tensors for model.context_compressor()
+        - encoded_dataframe: Encoded context DataFrame (for extracting keys if needed)
+        - encoded_primary_key: Name of encoded primary key column (None if not provided)
+    """
+
+    # Encode context data
+    ctx_encoded, ctx_primary_key_encoded, _ = encode_df(df=ctx_data, stats=ctx_stats, ctx_primary_key=ctx_primary_key)
+
+    # Pad empty sequences (required for model)
+    ctx_encoded = pad_ctx_sequences(ctx_encoded)
+
+    # Build flat context inputs (CTXFLT/*)
+    ctxflt_inputs = {
+        col: torch.unsqueeze(
+            torch.as_tensor(ctx_encoded[col].to_numpy(), device=device).type(torch.int),
+            dim=-1,
+        )
+        for col in ctx_encoded.columns
+        if col.startswith(CTXFLT)
+    }
+
+    # Build sequential context inputs (CTXSEQ/*)
+    ctxseq_inputs = {
+        col: torch.unsqueeze(
+            torch.nested.as_nested_tensor(
+                [torch.as_tensor(t, device=device).type(torch.int) for t in ctx_encoded[col]],
+                device=device,
+            ),
+            dim=-1,
+        )
+        for col in ctx_encoded.columns
+        if col.startswith(CTXSEQ)
+    }
+
+    # Merge and return with encoded dataframe
+    return (ctxflt_inputs | ctxseq_inputs), ctx_encoded, ctx_primary_key_encoded
+
+
+def check_column_order(
+    gen_column_order: list[str],
+    trn_column_order: list[str],
+) -> None:
+    """
+    Check if column order matches training order.
+
+    Args:
+        gen_column_order: Column order for the current operation
+        trn_column_order: Column order from training
+
+    Raises:
+        ValueError: If column order doesn't match training order
+    """
+    if gen_column_order != trn_column_order:
+        raise ValueError(
+            "Column order does not match training order. "
+            "A change in column order is only permitted for models that were trained with `enable_flexible_generation=True`."
+        )
+
+
+def fix_rare_token_probs(
+    stats: dict,
+    rare_category_replacement_method: RareCategoryReplacementMethod | None = None,
+) -> dict[str, dict[str, CodeProbabilities]]:
+    """
+    Create fixed probabilities to suppress rare tokens.
+
+    Args:
+        stats: Target statistics dict
+        rare_category_replacement_method: How to handle rare categories
+
+    Returns:
+        Dict of column -> sub_column -> code -> probability
+    """
+    # suppress rare token for categorical when no_of_rare_categories == 0
+    mask = {
+        col: {CATEGORICAL_SUB_COL_SUFFIX: {col_stats["codes"][CATEGORICAL_UNKNOWN_TOKEN]: 0.0}}
+        for col, col_stats in stats["columns"].items()
+        if col_stats["encoding_type"] == ModelEncodingType.tabular_categorical
+        if "codes" in col_stats
+        if col_stats.get("no_of_rare_categories", 0) == 0
+    }
+    # suppress rare token for categorical if RareCategoryReplacementMethod is sample
+    if rare_category_replacement_method == RareCategoryReplacementMethod.sample:
+        mask |= {
+            col: {CATEGORICAL_SUB_COL_SUFFIX: {col_stats["codes"][CATEGORICAL_UNKNOWN_TOKEN]: 0.0}}
+            for col, col_stats in stats["columns"].items()
+            if col_stats["encoding_type"] == ModelEncodingType.tabular_categorical
+            if "codes" in col_stats
+        }
+    # always suppress rare token for numeric_binned
+    mask |= {
+        col: {NUMERIC_BINNED_SUB_COL_SUFFIX: {col_stats["codes"][NUMERIC_BINNED_UNKNOWN_TOKEN]: 0.0}}
+        for col, col_stats in stats["columns"].items()
+        if col_stats["encoding_type"] == ModelEncodingType.tabular_numeric_binned
+        if "codes" in col_stats
+    }
+    # always suppress rare token for numeric_discrete
+    mask |= {
+        col: {NUMERIC_DISCRETE_SUB_COL_SUFFIX: {col_stats["codes"][NUMERIC_DISCRETE_UNKNOWN_TOKEN]: 0.0}}
+        for col, col_stats in stats["columns"].items()
+        if col_stats["encoding_type"] == ModelEncodingType.tabular_numeric_discrete
+        if "codes" in col_stats
+    }
+    return mask
+
+
+def translate_fixed_probs(
+    fixed_probs: dict[str, dict[str, CodeProbabilities]], stats: dict
+) -> dict[str, CodeProbabilities]:
+    """
+    Translate fixed probs to ARGN naming conventions.
+
+    Args:
+        fixed_probs: Dict of column -> sub_column -> code -> probability
+        stats: Target statistics dict
+
+    Returns:
+        Dict of ARGN sub_column name -> code -> probability
+    """
+    mask = {
+        get_argn_name(
+            argn_processor=stats["columns"][col][ARGN_PROCESSOR],
+            argn_table=stats["columns"][col][ARGN_TABLE],
+            argn_column=stats["columns"][col][ARGN_COLUMN],
+            argn_sub_column=sub_col,
+        ): sub_col_mask
+        for col, col_mask in fixed_probs.items()
+        for sub_col, sub_col_mask in col_mask.items()
+    }
+    return mask