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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mostlyai-engine
-Version: 2.2.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.2.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.2.0 → mostlyai_engine-2.3.0}/mostlyai/engine/__init__.py

@@ -34,7 +34,7 @@ __all__ = [
     "TabularARGN",
     "LanguageModel",
 ]
-__version__ = "2.2.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.2.0 → mostlyai_engine-2.3.0}/mostlyai/engine/_tabular/common.py

@@ -19,10 +19,32 @@ from pathlib import Path
 import pandas as pd
 import torch
 
-from mostlyai.engine._common import CTXFLT, CTXSEQ
+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")
 
 
@@ -165,7 +187,6 @@ def prepare_context_inputs(
         - encoded_dataframe: Encoded context DataFrame (for extracting keys if needed)
         - encoded_primary_key: Name of encoded primary key column (None if not provided)
     """
-    from mostlyai.engine._tabular.encoding import encode_df, pad_ctx_sequences
 
     # Encode context data
     ctx_encoded, ctx_primary_key_encoded, _ = encode_df(df=ctx_data, stats=ctx_stats, ctx_primary_key=ctx_primary_key)
@@ -198,3 +219,97 @@ def prepare_context_inputs(
 
     # 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

{mostlyai_engine-2.2.0 → mostlyai_engine-2.3.0}/mostlyai/engine/_tabular/generation.py

@@ -67,10 +67,8 @@ from mostlyai.engine._encoding_types.tabular.lat_long import decode_latlong
 from mostlyai.engine._encoding_types.tabular.numeric import (
     NUMERIC_BINNED_NULL_TOKEN,
     NUMERIC_BINNED_SUB_COL_SUFFIX,
-    NUMERIC_BINNED_UNKNOWN_TOKEN,
     NUMERIC_DISCRETE_NULL_TOKEN,
     NUMERIC_DISCRETE_SUB_COL_SUFFIX,
-    NUMERIC_DISCRETE_UNKNOWN_TOKEN,
     decode_numeric,
 )
 from mostlyai.engine._memory import get_available_ram_for_heuristics, get_available_vram_for_heuristics
@@ -80,10 +78,13 @@ from mostlyai.engine._tabular.argn import (
     SequentialModel,
 )
 from mostlyai.engine._tabular.common import (
+    check_column_order,
     create_and_load_model,
+    fix_rare_token_probs,
     load_model_artifacts,
     prepare_context_inputs,
     resolve_device,
+    translate_fixed_probs,
 )
 from mostlyai.engine._tabular.encoding import encode_df
 from mostlyai.engine._tabular.fairness import FairnessTransforms, get_fairness_transforms
@@ -454,43 +455,6 @@ def _generation_batch_size_heuristic(mem_available_gb: float, ctx_stats: dict, t
 #########################
 
 
-def _fix_rare_token_probs(
-    stats: dict,
-    rare_category_replacement_method: RareCategoryReplacementMethod | None = None,
-) -> dict[str, dict[str, CodeProbabilities]]:
-    # 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 _fix_imputation_probs(
     stats: dict,
     imputation: ImputationConfig | None = None,
@@ -573,23 +537,6 @@ def _fix_rebalancing_probs(
     return mask
 
 
-def _translate_fixed_probs(
-    fixed_probs: dict[str, dict[str, CodeProbabilities]], stats: dict
-) -> dict[str, CodeProbabilities]:
-    # translate fixed probs to ARGN conventions
-    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
-
-
 def _deepmerge(*dictionaries: dict, merged: dict | None = None) -> dict:
     merged = merged or {}
     for dictionary in dictionaries:
@@ -903,17 +850,12 @@ def generate(
         _LOG.info(f"{trn_column_order=}")
 
         if not enable_flexible_generation:
-            # check if resolved column order is the same as the one from training
-            if gen_column_order != trn_column_order:
-                raise ValueError(
-                    "The column order for generation does not match the column order from training, due to seed, rebalancing, fairness or imputation configs. "
-                    "A change in column order is only permitted for models that were trained with `enable_flexible_generation=True`."
-                )
+            check_column_order(gen_column_order, trn_column_order)
         _LOG.info(f"{rare_category_replacement_method=}")
-        rare_token_fixed_probs = _fix_rare_token_probs(tgt_stats, rare_category_replacement_method)
+        rare_token_fixed_probs = fix_rare_token_probs(tgt_stats, rare_category_replacement_method)
         imputation_fixed_probs = _fix_imputation_probs(tgt_stats, imputation)
         rebalancing_fixed_probs = _fix_rebalancing_probs(tgt_stats, rebalancing)
-        fixed_probs = _translate_fixed_probs(
+        fixed_probs = translate_fixed_probs(
             fixed_probs=_deepmerge(
                 rare_token_fixed_probs,
                 imputation_fixed_probs,

{mostlyai_engine-2.2.0 → mostlyai_engine-2.3.0}/mostlyai/engine/_tabular/interface.py

@@ -38,6 +38,7 @@ from mostlyai.engine._common import (
     median_fn,
     mode_fn,
 )
+from mostlyai.engine._tabular.probability import log_prob as _log_prob
 from mostlyai.engine._tabular.probability import predict_proba as _predict_proba
 from mostlyai.engine._workspace import Workspace
 from mostlyai.engine.analysis import analyze
@@ -730,3 +731,54 @@ class TabularARGN(BaseEstimator):
         )
 
         return probs_df
+
+    def log_prob(
+        self,
+        X: pd.DataFrame,
+        ctx_data: pd.DataFrame | None = None,
+        **kwargs,
+    ) -> pd.DataFrame:
+        """
+        Compute log probability of observations.
+
+        This method computes P(observation | model) - the likelihood of the observation
+        under the trained model. For autoregressive models:
+
+        log P(x1, x2, ..., xn) = log P(x1) + log P(x2|x1) + ... + log P(xn|x1,...,xn-1)
+
+        Each term is the probability the model assigns to the actual observed value.
+
+        Args:
+            X: DataFrame with all columns containing observed values.
+            ctx_data: Context data for generation. If None, uses the context data from training.
+            **kwargs: Additional generation parameters (device, etc.).
+
+        Returns:
+            pd.DataFrame of shape (n_samples, 1) with log probability per row.
+            Column name is "log_prob". Values are <= 0 (log probabilities).
+            More negative values indicate less likely samples.
+
+        Raises:
+            ValueError: If model is not fitted.
+        """
+        if not self._fitted:
+            raise ValueError("Model must be fitted before computing log_prob. Call fit() first.")
+
+        X_df = ensure_dataframe(X, columns=self._feature_names)
+
+        workspace = Workspace(self.workspace_dir)
+        device = kwargs.get("device", self.device)
+
+        # Prepare ctx_data
+        if ctx_data is None:
+            ctx_data = self.ctx_data
+        ctx_data_df = ensure_dataframe(ctx_data) if ctx_data is not None else None
+
+        log_probs = _log_prob(
+            workspace=workspace,
+            data=X_df,
+            ctx_data=ctx_data_df,
+            device=device,
+        )
+
+        return log_probs

{mostlyai_engine-2.2.0 → mostlyai_engine-2.3.0}/mostlyai/engine/_tabular/probability.py

@@ -43,22 +43,90 @@ from mostlyai.engine._encoding_types.tabular.numeric import (
 )
 from mostlyai.engine._tabular.argn import ModelSize
 from mostlyai.engine._tabular.common import (
+    check_column_order,
     create_and_load_model,
+    fix_rare_token_probs,
     load_model_artifacts,
     prepare_context_inputs,
     resolve_device,
+    translate_fixed_probs,
 )
 from mostlyai.engine._tabular.encoding import encode_df
-from mostlyai.engine._tabular.generation import _fix_rare_token_probs, _translate_fixed_probs
+from mostlyai.engine._tabular.training import _calculate_sample_losses
 from mostlyai.engine._workspace import Workspace
 from mostlyai.engine.domain import ModelEncodingType, RareCategoryReplacementMethod
 
 _LOG = logging.getLogger(__name__)
 
 
-##########################
-### PROBABILITY UTILS ###
-##########################
+def _initialize_model(
+    *,
+    workspace: Workspace,
+    rare_category_replacement_method: RareCategoryReplacementMethod | str | None = None,
+    device: torch.device | str | None = None,
+    allow_sequential: bool = True,
+) -> tuple[torch.nn.Module, dict, dict, dict, list[str], dict, torch.device, bool]:
+    """
+    Initialize model and artifacts for probability computation.
+
+    Args:
+        workspace: Workspace containing model and stats
+        rare_category_replacement_method: How to handle rare categories (None to skip fixed_probs)
+        device: Device for computation
+        allow_sequential: Whether to allow sequential models
+
+    Returns:
+        Tuple of (model, tgt_stats, ctx_stats, tgt_cardinalities, all_columns, fixed_probs, device, enable_flexible_generation)
+
+    Raises:
+        ValueError: If model is sequential and allow_sequential is False
+    """
+    # Load model artifacts
+    model_config, tgt_stats, ctx_stats, is_sequential = load_model_artifacts(workspace)
+
+    # Check model type
+    if is_sequential and not allow_sequential:
+        raise ValueError("Sequential models are not supported for this operation")
+
+    # Get cardinalities and config
+    tgt_cardinalities = get_cardinalities(tgt_stats)
+    ctx_cardinalities = get_cardinalities(ctx_stats)
+    enable_flexible_generation = model_config.get("enable_flexible_generation", True)
+
+    # Resolve device
+    device = resolve_device(device)
+    _LOG.info(f"Using device: {device}")
+
+    # Get all columns in training order
+    all_columns = get_columns_from_cardinalities(tgt_cardinalities)
+
+    # Create model
+    model_units = model_config.get("model_units") or ModelSize.M
+    ctx_seq_len_median = ctx_stats.get("sequence_len_median")
+
+    model = create_and_load_model(
+        workspace=workspace,
+        is_sequential=is_sequential,
+        tgt_cardinalities=tgt_cardinalities,
+        ctx_cardinalities=ctx_cardinalities,
+        model_units=model_units,
+        ctx_seq_len_median=ctx_seq_len_median,
+        column_order=all_columns,
+        device=device,
+    )
+
+    # Prepare fixed_probs to suppress rare tokens (if replacement method provided)
+    if rare_category_replacement_method is not None:
+        _LOG.info(f"{rare_category_replacement_method=}")
+        rare_token_fixed_probs = fix_rare_token_probs(tgt_stats, rare_category_replacement_method)
+        fixed_probs = translate_fixed_probs(
+            fixed_probs=rare_token_fixed_probs,
+            stats=tgt_stats,
+        )
+    else:
+        fixed_probs = {}
+
+    return model, tgt_stats, ctx_stats, tgt_cardinalities, all_columns, fixed_probs, device, enable_flexible_generation
 
 
 def _get_column_metadata(target_column: str, target_stats: dict) -> list[dict]:
@@ -181,6 +249,7 @@ def _generate_marginal_probs(
     ctx_data: pd.DataFrame | None = None,
     ctx_stats: dict | None = None,
     fixed_probs: dict | None = None,
+    enable_flexible_generation: bool = True,
 ) -> np.ndarray:
     """
     Generate P(target | seed_features, context).
@@ -195,6 +264,8 @@ def _generate_marginal_probs(
         device: Device for computation
         ctx_data: Optional context data
         ctx_stats: Optional context statistics (required if ctx_data provided)
+        fixed_probs: Optional fixed probabilities for rare token handling
+        enable_flexible_generation: Whether flexible generation is enabled
 
     Returns:
         DataFrame of shape (n_samples, cardinality) with probabilities and column names
@@ -229,6 +300,10 @@ def _generate_marginal_probs(
     # Determine column order: seed columns + target (in training order)
     gen_column_order = [col for col in all_columns if col in seed_columns or col == target_column]
 
+    # Check column order when flexible generation is disabled
+    if not enable_flexible_generation:
+        check_column_order(gen_column_order, all_columns)
+
     # Prepare context inputs if provided
     if ctx_data is not None and ctx_stats is not None:
         x, _, _ = prepare_context_inputs(ctx_data=ctx_data, ctx_stats=ctx_stats, device=device)
@@ -292,45 +367,14 @@ def predict_proba(
     """
     _LOG.info(f"PREDICT_PROBA started for targets: {target_columns}")
 
-    # Load model artifacts
-    model_config, tgt_stats, ctx_stats, is_sequential = load_model_artifacts(workspace)
-
-    # Check model type
-    if is_sequential:
-        raise ValueError("predict_proba is not yet supported for sequential models")
-
-    # Get cardinalities
-    tgt_cardinalities = get_cardinalities(tgt_stats)
-    ctx_cardinalities = get_cardinalities(ctx_stats)
-
-    # Resolve device
-    device = resolve_device(device)
-    _LOG.info(f"Using device: {device}")
-
-    # Get all columns in training order
-    all_columns = get_columns_from_cardinalities(tgt_cardinalities)
-
-    # Create model (will override column_order in forward pass for optimization)
-    model_units = model_config.get("model_units") or ModelSize.M
-    ctx_seq_len_median = ctx_stats.get("sequence_len_median")
-
-    model = create_and_load_model(
-        workspace=workspace,
-        is_sequential=is_sequential,
-        tgt_cardinalities=tgt_cardinalities,
-        ctx_cardinalities=ctx_cardinalities,
-        model_units=model_units,
-        ctx_seq_len_median=ctx_seq_len_median,
-        column_order=all_columns,  # Will override in forward pass
-        device=device,
-    )
-
-    # Prepare fixed_probs to suppress rare tokens
-    _LOG.info(f"{rare_category_replacement_method=}")
-    rare_token_fixed_probs = _fix_rare_token_probs(tgt_stats, rare_category_replacement_method)
-    fixed_probs = _translate_fixed_probs(
-        fixed_probs=rare_token_fixed_probs,
-        stats=tgt_stats,
+    # Initialize model
+    model, tgt_stats, ctx_stats, tgt_cardinalities, all_columns, fixed_probs, device, enable_flexible_generation = (
+        _initialize_model(
+            workspace=workspace,
+            rare_category_replacement_method=rare_category_replacement_method,
+            device=device,
+            allow_sequential=False,
+        )
     )
 
     # Encode seed data (features to condition on) - common for both single and multi-target
@@ -352,12 +396,13 @@ def predict_proba(
             target_cardinality = sum(col_stats["cardinalities"].values())
             total_cardinality *= target_cardinality
 
-        _LOG.warning(
-            f"Computing joint probabilities for {len(target_columns)} targets "
-            f"results in {total_cardinality:,} total probability values per sample. "
-            f"Computation complexity grows exponentially with the number of targets. "
-            f"Consider computing probabilities for targets separately if this takes too long."
-        )
+        if total_cardinality > 100:
+            _LOG.warning(
+                f"Computing joint probabilities for {len(target_columns)} targets "
+                f"results in {total_cardinality:,} total probability values per sample. "
+                f"Computation complexity grows exponentially with the number of targets. "
+                f"Consider computing probabilities for targets separately if this takes too long."
+            )
 
     # Initialize with first target: P(col1)
     first_target_df = _generate_marginal_probs(
@@ -371,6 +416,7 @@ def predict_proba(
         ctx_data=ctx_data,
         ctx_stats=ctx_stats,
         fixed_probs=fixed_probs,
+        enable_flexible_generation=enable_flexible_generation,
     )
     _LOG.info(f"Generated P({target_columns[0]}) with shape {first_target_df.shape}")
 
@@ -452,6 +498,7 @@ def predict_proba(
             ctx_data=batched_ctx_data,
             ctx_stats=ctx_stats,
             fixed_probs=fixed_probs,
+            enable_flexible_generation=enable_flexible_generation,
         )  # DataFrame: (n_samples * num_combos, current_card)
 
         # Extract probabilities for each combo and compute joint probabilities
@@ -479,3 +526,72 @@ def predict_proba(
 
     _LOG.info(f"PREDICT_PROBA finished: returned probabilities for {len(probs_df)} samples")
     return probs_df
+
+
+@torch.no_grad()
+def log_prob(
+    *,
+    workspace: Workspace,
+    data: pd.DataFrame,
+    ctx_data: pd.DataFrame | None = None,
+    device: torch.device | str | None = None,
+) -> pd.DataFrame:
+    """
+    Compute log probability of full observations.
+
+    This function computes P(observation | model) - the likelihood of the observation
+    under the trained model. For autoregressive models:
+
+    log P(x1, x2, ..., xn) = log P(x1) + log P(x2|x1) + ... + log P(xn|x1,...,xn-1)
+
+    Each term is the probability the model assigns to the actual observed value.
+
+    Supports all encoding types including multi-sub-column encodings like numeric_digit,
+    and both flat and sequential models.
+
+    Args:
+        workspace: Workspace object containing model and stats
+        data: DataFrame with ALL columns containing observed values
+        ctx_data: Optional context data (for models with context)
+        device: Device to run inference on ('cuda' or 'cpu'). Defaults to 'cuda' if available.
+
+    Returns:
+        pd.DataFrame of shape (n_samples, 1) with log probability per row.
+        Column name is "log_prob". Values are <= 0 (log probabilities).
+    """
+    _LOG.info("LOG_PROB started")
+
+    # Initialize model
+    model, tgt_stats, ctx_stats, _, all_columns, _, device, enable_flexible_generation = _initialize_model(
+        workspace=workspace,
+        device=device,
+    )
+
+    # Check column order of input data when flexible generation is disabled
+    if not enable_flexible_generation:
+        check_column_order(list(data.columns), all_columns)
+
+    # Encode full data to get observed codes for all columns
+    full_encoded, _, _ = encode_df(df=data, stats=tgt_stats)
+
+    n_samples = len(data)
+    _LOG.info(f"Computing log probabilities for {n_samples} samples")
+
+    # Build batch dict with ALL encoded values as tensors
+    batch_dict: dict[str, torch.Tensor] = {}
+    for col in full_encoded.columns:
+        batch_dict[col] = torch.tensor(full_encoded[col].values, dtype=torch.long, device=device).unsqueeze(-1)
+
+    # Add context data if provided
+    if ctx_data is not None and ctx_stats:
+        ctx_inputs, _, _ = prepare_context_inputs(ctx_data=ctx_data, ctx_stats=ctx_stats, device=device)
+        batch_dict.update(ctx_inputs)
+
+    # Use the training loss calculation directly - handles both flat and sequential models
+    losses = _calculate_sample_losses(model, batch_dict)
+
+    # Negate loss to get log probability (loss = -log_prob)
+    log_probs = -losses.cpu().numpy()
+
+    _LOG.info(f"LOG_PROB finished: computed log probs for {n_samples} samples")
+    return pd.DataFrame({"log_prob": log_probs})

{mostlyai_engine-2.2.0 → mostlyai_engine-2.3.0}/mostlyai/engine/splitting.py

@@ -19,6 +19,7 @@ Split original data for training and validation.
 import logging
 import time
 import warnings
+from collections.abc import Callable
 from pathlib import Path
 
 import numpy as np
@@ -73,7 +74,7 @@ def split(
     tgt_encoding_types: dict[str, str | ModelEncodingType] | None = None,
     ctx_encoding_types: dict[str, str | ModelEncodingType] | None = None,
     n_partitions: int = 1,
-    trn_val_split: float = 0.8,
+    trn_val_split: float | Callable[[pd.Series], tuple[pd.Series, pd.Series]] = 0.8,
     workspace_dir: str | Path = "engine-ws",
     update_progress: ProgressCallback | None = None,
 ) -> None:
@@ -99,7 +100,8 @@ def split(
         tgt_encoding_types: Encoding types for columns in the target data (excluding key columns).
         ctx_encoding_types: Encoding types for columns in the context data (excluding key columns).
         n_partitions: Number of partitions to split the data into.
-        trn_val_split: Fraction of data to use for training, with the remaining data used for validation.
+        trn_val_split: Fraction of data to use for training (0 < value < 1), or a callable
+            that takes keys as input and returns (trn_keys, val_keys) tuple.
         workspace_dir: Path to the workspace directory where files will be created.
         update_progress: A custom progress callback.
     """
@@ -189,13 +191,17 @@ def split(
             keys = tgt_data[tgt_context_key].drop_duplicates()
         else:
             keys = ctx_data[ctx_primary_key]
-        # shuffle keys
-        keys = keys.sample(frac=1)
-        # split randomly into trn and val
-        assert 0 < trn_val_split < 1, f"invalid trn_val_split: {trn_val_split}"
-        trn_cnt = round(trn_val_split * len(keys))
-        trn_keys = keys[:trn_cnt]
-        val_keys = keys[trn_cnt:]
+        # split into trn and val
+        if callable(trn_val_split):
+            trn_keys, val_keys = trn_val_split(keys)
+        else:
+            # shuffle keys
+            keys = keys.sample(frac=1)
+            # split randomly into trn and val
+            assert 0 < trn_val_split < 1, f"invalid trn_val_split: {trn_val_split}"
+            trn_cnt = round(trn_val_split * len(keys))
+            trn_keys = keys[:trn_cnt]
+            val_keys = keys[trn_cnt:]
 
         def save_partition(
             df: pd.DataFrame, path: Path, key: str, sel_trn_keys: np.ndarray, sel_val_keys: np.ndarray, idx: int

{mostlyai_engine-2.2.0 → mostlyai_engine-2.3.0}/pyproject.toml

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