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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mostlyai-engine
-Version: 2.1.0
+Version: 2.2.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

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

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

@@ -0,0 +1,200 @@
+# 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 CTXFLT, CTXSEQ
+
+_LOG = logging.getLogger(__name__)
+
+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)
+    """
+    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)
+
+    # 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

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

@@ -78,10 +78,14 @@ from mostlyai.engine._tabular.argn import (
     FlatModel,
     ModelSize,
     SequentialModel,
-    get_no_of_model_parameters,
 )
-from mostlyai.engine._tabular.common import load_model_weights
-from mostlyai.engine._tabular.encoding import encode_df, pad_ctx_sequences
+from mostlyai.engine._tabular.common import (
+    create_and_load_model,
+    load_model_artifacts,
+    prepare_context_inputs,
+    resolve_device,
+)
+from mostlyai.engine._tabular.encoding import encode_df
 from mostlyai.engine._tabular.fairness import FairnessTransforms, get_fairness_transforms
 from mostlyai.engine._workspace import Workspace, ensure_workspace_dir, reset_dir
 from mostlyai.engine.domain import (
@@ -839,13 +843,10 @@ def generate(
         output_path = workspace.generated_data_path
         reset_dir(output_path)
 
-        model_configs = workspace.model_configs.read()
-        tgt_stats = workspace.tgt_stats.read()
-        is_sequential = tgt_stats["is_sequential"]
+        model_configs, tgt_stats, ctx_stats, is_sequential = load_model_artifacts(workspace)
         _LOG.info(f"{is_sequential=}")
         has_context = workspace.ctx_stats.path.exists()
         _LOG.info(f"{has_context=}")
-        ctx_stats = workspace.ctx_stats.read()
 
         # read model config
         model_units = model_configs.get("model_units") or ModelSize.M
@@ -871,11 +872,7 @@ def generate(
         _LOG.info(f"{len(ctx_sub_columns)=}")
 
         # resolve device
-        device = (
-            torch.device(device)
-            if device is not None
-            else (torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu"))
-        )
+        device = resolve_device(device)
         _LOG.info(f"{device=}")
 
         tgt_primary_key = tgt_stats.get("keys", {}).get("primary_key")
@@ -912,7 +909,6 @@ def generate(
                     "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`."
                 )
-
         _LOG.info(f"{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)
@@ -1016,43 +1012,18 @@ def generate(
         # init progress with total_count; +1 for the final decoding step
         progress.update(completed=0, total=no_of_batches * (seq_len_max + 1))
 
-        _LOG.info("create generative model")
-        model: FlatModel | SequentialModel
-        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=gen_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=gen_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; generating data with an untrained model")
-
-        model.to(device)
-        model.eval()
+        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=gen_column_order,
+            device=device,
+            seq_len_median=seq_len_median,
+            seq_len_max=seq_len_max,
+        )
 
         # calculate fairness transforms only once before batch generation
         fairness_transforms: FairnessTransforms | None = None
@@ -1121,13 +1092,11 @@ def generate(
             ctx_batch = ctx_batch.sort_values(ctx_primary_key).reset_index(drop=True)
             seed_batch = seed_batch.sort_values(tgt_context_key).reset_index(drop=True)
 
-            # encode ctx_batch
+            # encode ctx_batch and prepare tensor inputs
             _LOG.info(f"encode context {ctx_batch.shape}")
-            ctx_batch_encoded, ctx_primary_key_encoded, _ = encode_df(
-                df=ctx_batch, stats=ctx_stats, ctx_primary_key=ctx_primary_key
+            ctx_inputs, ctx_batch_encoded, ctx_primary_key_encoded = prepare_context_inputs(
+                ctx_data=ctx_batch, ctx_stats=ctx_stats, device=model.device, ctx_primary_key=ctx_primary_key
             )
-            # pad left context sequences to ensure non-empty sequences
-            ctx_batch_encoded = pad_ctx_sequences(ctx_batch_encoded)
             ctx_keys = ctx_batch_encoded[ctx_primary_key_encoded]
             ctx_keys.rename(tgt_context_key, inplace=True)
 
@@ -1149,30 +1118,12 @@ def generate(
                 syn = ctx_keys.to_frame().reset_index(drop=True)
                 buffer.add((syn, seed_batch))
             elif isinstance(model, SequentialModel):
-                ctxflt_inputs = {
-                    col: torch.unsqueeze(
-                        torch.as_tensor(ctx_batch_encoded[col].to_numpy(), device=model.device).type(torch.int),
-                        dim=-1,
-                    )
-                    for col in ctx_batch_encoded.columns
-                    if col.startswith(CTXFLT)
-                }
-                ctxseq_inputs = {
-                    col: torch.unsqueeze(
-                        torch.nested.as_nested_tensor(
-                            [torch.as_tensor(t, device=model.device).type(torch.int) for t in ctx_batch_encoded[col]],
-                            device=model.device,
-                        ),
-                        dim=-1,
-                    )
-                    for col in ctx_batch_encoded.columns
-                    if col.startswith(CTXSEQ)
-                }
+                # Use context inputs prepared earlier
                 seq_steps = model.tgt_seq_len_max
                 history = None
                 history_state = None
                 # process context just once for all sequence steps
-                context = model.context_compressor(ctxflt_inputs | ctxseq_inputs)
+                context = model.context_compressor(ctx_inputs)
                 # loop over sequence steps, and pass forward history to keep model state-less
                 out_df: pd.DataFrame | None = None
                 decode_prev_steps = {}
@@ -1358,26 +1309,8 @@ def generate(
                         persist_data_part(syn, output_path, f"{buffer.n_clears:06}.{0:06}")
                         buffer.clear()
             else:  # isinstance(model, FlatModel)
-                ctxflt_inputs = {
-                    col: torch.unsqueeze(
-                        torch.as_tensor(ctx_batch_encoded[col].to_numpy(), device=model.device).type(torch.int),
-                        dim=-1,
-                    )
-                    for col in ctx_batch_encoded.columns
-                    if col.startswith(CTXFLT)
-                }
-                ctxseq_inputs = {
-                    col: torch.unsqueeze(
-                        torch.nested.as_nested_tensor(
-                            [torch.as_tensor(t, device=model.device).type(torch.int) for t in ctx_batch_encoded[col]],
-                            device=model.device,
-                        ),
-                        dim=-1,
-                    )
-                    for col in ctx_batch_encoded.columns
-                    if col.startswith(CTXSEQ)
-                }
-                x = ctxflt_inputs | ctxseq_inputs
+                # Use context inputs prepared earlier
+                x = ctx_inputs
                 fixed_values = {
                     col: torch.as_tensor(seed_batch_encoded[col].to_numpy(), device=model.device).type(torch.int)
                     for col in seed_batch_encoded.columns