mostlyai-engine 1.4.8__tar.gz → 1.5.0__tar.gz

{mostlyai_engine-1.4.8 → mostlyai_engine-1.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mostlyai-engine
-Version: 1.4.8
+Version: 1.5.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
@@ -28,7 +28,7 @@ Requires-Dist: accelerate>=1.5.0
 Requires-Dist: datasets>=3.0.0
 Requires-Dist: huggingface-hub[hf-xet]>=0.30.2
 Requires-Dist: joblib>=1.4.2
-Requires-Dist: json-repair<0.47.0,>=0.30.0
+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
@@ -40,7 +40,7 @@ Requires-Dist: tokenizers>=0.21.0
 Requires-Dist: torch<2.7.1,>=2.7.0
 Requires-Dist: torchaudio<2.7.1,>=2.7.0
 Requires-Dist: torchvision<0.22.1,>=0.22.0
-Requires-Dist: transformers>=4.51.0
+Requires-Dist: transformers<4.54.0,>=4.51.0
 Requires-Dist: xgrammar>=0.1.19
 Provides-Extra: gpu
 Requires-Dist: bitsandbytes==0.42.0; (sys_platform == 'darwin') and extra == 'gpu'

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

{mostlyai_engine-1.4.8 → mostlyai_engine-1.5.0}/mostlyai/engine/_common.py

@@ -48,11 +48,12 @@ ARGN_COLUMN = "argn_column"
 PREFIX_TABLE = ":"
 PREFIX_COLUMN = "/"
 PREFIX_SUB_COLUMN = "__"
-SLEN_SIDX_SDEC_COLUMN = f"{TGT}{PREFIX_TABLE}{PREFIX_COLUMN}"
-SLEN_SIDX_DIGIT_ENCODING_THRESHOLD = 100
-SLEN_SUB_COLUMN_PREFIX = f"{SLEN_SIDX_SDEC_COLUMN}{PREFIX_SUB_COLUMN}slen_"  # sequence length
-SIDX_SUB_COLUMN_PREFIX = f"{SLEN_SIDX_SDEC_COLUMN}{PREFIX_SUB_COLUMN}sidx_"  # sequence index
-SDEC_SUB_COLUMN_PREFIX = f"{SLEN_SIDX_SDEC_COLUMN}{PREFIX_SUB_COLUMN}sdec_"  # sequence index decile
+SIDX_RIDX_DIGIT_ENCODING_THRESHOLD = 100
+POSITIONAL_COLUMN = f"{TGT}{PREFIX_TABLE}{PREFIX_COLUMN}"
+SIDX_SUB_COLUMN_PREFIX = f"{POSITIONAL_COLUMN}{PREFIX_SUB_COLUMN}sidx_"  # sequence index
+RIDX_SUB_COLUMN_PREFIX = f"{POSITIONAL_COLUMN}{PREFIX_SUB_COLUMN}ridx_"  # reverse index
+SLEN_SUB_COLUMN_PREFIX = f"{POSITIONAL_COLUMN}{PREFIX_SUB_COLUMN}slen_"  # sequence length
+SDEC_SUB_COLUMN_PREFIX = f"{POSITIONAL_COLUMN}{PREFIX_SUB_COLUMN}sdec_"  # sequence index decile
 TABLE_COLUMN_INFIX = "::"  # this should be consistent as in mostly-data and mostlyai-qa
 
 ANALYZE_MIN_MAX_TOP_N = 1000  # the number of min/max values to be kept from each partition
@@ -314,11 +315,14 @@ def get_argn_name(
     return "".join(name)
 
 
-def get_cardinalities(stats: dict) -> dict[str, int]:
+def get_cardinalities(
+    stats: dict, has_slen: bool | None = None, has_ridx: bool | None = None, has_sdec: bool | None = None
+) -> dict[str, int]:
     cardinalities: dict[str, int] = {}
+
     if stats.get("is_sequential", False):
         max_seq_len = get_sequence_length_stats(stats)["max"]
-        cardinalities |= get_slen_sidx_sdec_cardinalities(max_seq_len)
+        cardinalities |= get_positional_cardinalities(max_seq_len, has_slen, has_ridx, has_sdec)
 
     for i, column in enumerate(stats.get("columns", [])):
         column_stats = stats["columns"][column]
@@ -512,72 +516,70 @@ def skip_if_error(func: Callable) -> Callable:
     return skip_if_error_wrapper
 
 
-def encode_slen_sidx_sdec(vals: pd.Series, max_seq_len: int, prefix: str = "") -> pd.DataFrame:
+def encode_positional_column(vals: pd.Series, max_seq_len: int, prefix: str = "") -> pd.DataFrame:
     assert is_integer_dtype(vals)
-    if max_seq_len < SLEN_SIDX_DIGIT_ENCODING_THRESHOLD or prefix == SDEC_SUB_COLUMN_PREFIX:
-        # encode slen and sidx as numeric_discrete
+    if max_seq_len < SIDX_RIDX_DIGIT_ENCODING_THRESHOLD:
+        # encode positional column as numeric_discrete
         df = pd.DataFrame({f"{prefix}cat": vals})
     else:
-        # encode as numeric_digit
+        # encode positional column as numeric_digit
         n_digits = len(str(max_seq_len))
         df = pd.DataFrame(vals.astype(str).str.pad(width=n_digits, fillchar="0").apply(list).tolist()).astype(int)
         df.columns = [f"{prefix}E{i}" for i in range(n_digits - 1, -1, -1)]
     return df
 
 
-def decode_slen_sidx_sdec(df_encoded: pd.DataFrame, max_seq_len: int, prefix: str = "") -> pd.Series:
-    if max_seq_len < SLEN_SIDX_DIGIT_ENCODING_THRESHOLD or prefix == SDEC_SUB_COLUMN_PREFIX:
-        # decode slen and sidx as numeric_discrete
+def decode_positional_column(df_encoded: pd.DataFrame, max_seq_len: int, prefix: str = "") -> pd.Series:
+    if max_seq_len < SIDX_RIDX_DIGIT_ENCODING_THRESHOLD:
+        # decode positional column as numeric_discrete
         vals = df_encoded[f"{prefix}cat"]
     else:
-        # decode slen and sidx as numeric_digit
+        # decode positional column as numeric_digit
         n_digits = len(str(max_seq_len))
         vals = sum([df_encoded[f"{prefix}E{d}"] * 10 ** int(d) for d in list(range(n_digits))])
     return vals
 
 
-def get_slen_sidx_sdec_cardinalities(max_seq_len) -> dict[str, int]:
-    if max_seq_len < SLEN_SIDX_DIGIT_ENCODING_THRESHOLD:
-        # encode slen and sidx as numeric_discrete
-        slen_cardinalities = {f"{SLEN_SUB_COLUMN_PREFIX}cat": max_seq_len + 1}
+def get_positional_cardinalities(
+    max_seq_len: int, has_slen: bool | None, has_ridx: bool | None, has_sdec: bool | None
+) -> dict[str, int]:
+    # the latest version of the model uses SIDX/SLEN/RIDX positional column
+    has_slen = has_slen if has_slen is not None else True
+    has_ridx = has_ridx if has_ridx is not None else True
+    has_sdec = has_sdec if has_sdec is not None else False
+
+    if max_seq_len < SIDX_RIDX_DIGIT_ENCODING_THRESHOLD:
+        # encode positional columns as numeric_discrete
         sidx_cardinalities = {f"{SIDX_SUB_COLUMN_PREFIX}cat": max_seq_len + 1}
+        slen_cardinalities = {f"{SLEN_SUB_COLUMN_PREFIX}cat": max_seq_len + 1}
+        ridx_cardinalities = {f"{RIDX_SUB_COLUMN_PREFIX}cat": max_seq_len + 1}
     else:
-        # encode slen and sidx as numeric_digit
+        # encode positional columns as numeric_digit
         digits = [int(digit) for digit in str(max_seq_len)]
-        slen_cardinalities = {}
         sidx_cardinalities = {}
+        slen_cardinalities = {}
+        ridx_cardinalities = {}
         for idx, digit in enumerate(digits):
             # cap cardinality of the most significant position
             # less significant positions allow any digit
             card = digit + 1 if idx == 0 else 10
             e_idx = len(digits) - idx - 1
-            slen_cardinalities[f"{SLEN_SUB_COLUMN_PREFIX}E{e_idx}"] = card
             sidx_cardinalities[f"{SIDX_SUB_COLUMN_PREFIX}E{e_idx}"] = card
-    # order is important: slen first, then sidx, as the former has highest priority
+            ridx_cardinalities[f"{RIDX_SUB_COLUMN_PREFIX}E{e_idx}"] = card
+            slen_cardinalities[f"{SLEN_SUB_COLUMN_PREFIX}E{e_idx}"] = card
     sdec_cardinalities = {f"{SDEC_SUB_COLUMN_PREFIX}cat": 10}
-    return slen_cardinalities | sidx_cardinalities | sdec_cardinalities
-
-
-def trim_sequences(syn: pd.DataFrame, tgt_context_key: str, seq_len_min: int, seq_len_max: int):
-    if syn.empty:
-        return syn
-
-    # use SIDX and SLEN to determine sequence length
-    syn[SIDX_SUB_COLUMN_PREFIX] = decode_slen_sidx_sdec(syn, seq_len_max, prefix=SIDX_SUB_COLUMN_PREFIX)
-    syn[SLEN_SUB_COLUMN_PREFIX] = decode_slen_sidx_sdec(syn, seq_len_max, prefix=SLEN_SUB_COLUMN_PREFIX)
-    # ensure that seq_len_min is respected
-    syn[SLEN_SUB_COLUMN_PREFIX] = np.maximum(seq_len_min, syn[SLEN_SUB_COLUMN_PREFIX])
-    syn = syn[syn[SIDX_SUB_COLUMN_PREFIX] < syn[SLEN_SUB_COLUMN_PREFIX]].reset_index(drop=True)
-    # discarded padded context rows, ie where context key has been set to None
-    syn = syn.dropna(subset=[tgt_context_key])
-    # discard SLEN and SIDX columns
-    syn.drop(
-        [c for c in syn.columns if c.startswith(SLEN_SIDX_SDEC_COLUMN)],
-        axis=1,
-        inplace=True,
-    )
-    syn.reset_index(drop=True, inplace=True)
-    return syn
+    match has_slen, has_ridx, has_sdec:
+        case True, True, False:
+            # SIDX/SLEN/RIDX model
+            return sidx_cardinalities | slen_cardinalities | ridx_cardinalities
+        case True, False, True:
+            # SLEN/SIDX/SDEC model
+            return slen_cardinalities | sidx_cardinalities | sdec_cardinalities
+        case True, False, False:
+            # SLEN/SIDX model
+            return slen_cardinalities | sidx_cardinalities
+        case _:
+            raise ValueError(f"Invalid positional encoding: {has_slen=}, {has_ridx=}, {has_sdec=}")
 
 
 def persist_data_part(df: pd.DataFrame, output_path: Path, infix: str):

{mostlyai_engine-1.4.8 → mostlyai_engine-1.5.0}/mostlyai/engine/_language/common.py

@@ -27,6 +27,7 @@ from transformers import (
     PretrainedConfig,
     PreTrainedModel,
 )
+from transformers.quantizers import AutoQuantizationConfig
 
 from mostlyai.engine._language.lstm import LSTMFromScratchConfig
 
@@ -84,7 +85,9 @@ def load_base_model_and_config(
     else:
         attn_implementation = None
         torch_dtype = torch.float32
-    if is_gpu_training and is_bitsandbytes_available:
+    if hasattr(config, "quantization_config"):
+        quantization_config = AutoQuantizationConfig.from_dict(config.quantization_config)
+    elif is_gpu_training and is_bitsandbytes_available:
         quantization_config = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_quant_type="nf4",
@@ -117,7 +120,7 @@ def load_base_model_and_config(
         quantization_config=quantization_config,
         torch_dtype=torch_dtype,
     )
-    if quantization_config:
+    if isinstance(quantization_config, BitsAndBytesConfig):
         # convert all non-kbit layers to float32
         model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=False)
     if is_gpu_training and model.supports_gradient_checkpointing:

{mostlyai_engine-1.4.8 → mostlyai_engine-1.5.0}/mostlyai/engine/_language/generation.py

@@ -14,7 +14,6 @@
 
 import contextlib
 import importlib
-import json
 import logging
 import os
 import platform
@@ -72,10 +71,10 @@ def decode_buffered_samples(
 
     def parse_json(x, columns: list[str]):
         try:
-            parsed_x = json_repair.loads(x)
+            parsed_x = json_repair.loads(x, stream_stable=True)
             if not isinstance(parsed_x, dict):
                 raise ValueError("parsed_x has to be a dictionary")
-        except (json.decoder.JSONDecodeError, ValueError):
+        except Exception:
             parsed_x = {}
         return [parsed_x.get(c, INVALID_VALUE) for c in columns]
 
@@ -183,9 +182,9 @@ def generate(
         enable_flexible_generation = model_configs.get("enable_flexible_generation", True)
         _LOG.info(f"{enable_flexible_generation=}")
 
-        # resolve potential conflict between sample_seed and sample_size
+        # resolve potential conflict between seed_data and sample_size
         if seed_data is not None:
-            assert sample_size is None, "either sample_seed or sample_size can be provided, not both"
+            assert sample_size is None, "either seed_data or sample_size can be provided, not both"
             sample_size = len(seed_data)
 
         if has_context:
@@ -222,7 +221,7 @@ def generate(
             ctx_primary_key = tgt_context_key = DUMMY_CONTEXT_KEY
             ctx_data = pd.DataFrame({ctx_primary_key: range(sample_size)})
 
-        # ensure sample_seed exists; ensure valid columns
+        # ensure seed_data exists; ensure valid columns
         if seed_data is None:
             # build dummy seed
             seed_data = pd.DataFrame(index=list(range(sample_size)))
@@ -230,7 +229,7 @@ def generate(
         _LOG.info(f"{seed_data.shape=}")
 
         if not enable_flexible_generation:
-            # validate sample_seed maintains the same column order as the one from training
+            # validate seed_data maintains the same column order as the one from training
             seed_columns = seed_data.columns.tolist()
             if seed_columns != tgt_text_columns[: len(seed_columns)]:
                 raise ValueError(
@@ -238,7 +237,7 @@ def generate(
                     "A change in column order is only permitted for models that were trained with `enable_flexible_generation=True`."
                 )
 
-        # sanity check: at this point sample seed and context data should have the same number of rows
+        # sanity check: at this point seed data and context data should have the same number of rows
         assert len(seed_data) == len(ctx_data)
 
         # early exit in case generation context is empty
@@ -306,14 +305,14 @@ def generate(
         samples_processed = 0
         while samples_processed < sample_size:
             encoded_ctx_batch = encoded_ctx_data.iloc[samples_processed : samples_processed + batch_size]
-            sample_seed_batch = seed_data.iloc[samples_processed : samples_processed + batch_size]
+            seed_data_batch = seed_data.iloc[samples_processed : samples_processed + batch_size]
             ctx_batch = ctx_data.iloc[samples_processed : samples_processed + batch_size]
             ctx_keys = ctx_batch[ctx_primary_key]
 
             if enforce_json_output and not initialize_logits_processors_once:
                 t0 = time.time()
                 schemas = create_schemas(
-                    seed_df=sample_seed_batch,
+                    seed_df=seed_data_batch,
                     stats=tgt_stats,
                     rare_category_replacement_method=rare_category_replacement_method,
                 )
@@ -328,7 +327,7 @@ def generate(
             total_tokenize_fn_time += metrics.tokenize_time
             total_generate_fn_time += metrics.generate_time
 
-            buffer.add((outputs, ctx_keys, sample_seed_batch))
+            buffer.add((outputs, ctx_keys, seed_data_batch))
             if buffer.is_full():
                 decoded_data = decode_buffered_samples(
                     buffer, engine.tokenizer, tgt_stats, tgt_context_key, max_new_tokens

{mostlyai_engine-1.4.8 → mostlyai_engine-1.5.0}/mostlyai/engine/_language/training.py

@@ -319,7 +319,10 @@ def train(
             )
         ):
             device = torch.device("cuda:0")
-            _LOG.info("device set to single gpu (cuda:0) because model is too small or differential privacy is enabled")
+            if torch.cuda.device_count() > 1:
+                _LOG.info(
+                    "device set to single gpu (cuda:0) because model is too small or differential privacy is enabled"
+                )
 
         if not with_dp:
             if device.type == "cuda":

{mostlyai_engine-1.4.8 → mostlyai_engine-1.5.0}/mostlyai/engine/_language/xgrammar_utils.py

@@ -42,7 +42,7 @@ def prepend_grammar_root_with_space(grammar: str) -> str:
     return grammar.replace(start_of_grammar, start_of_grammar_with_space)
 
 
-def ensure_seed_can_be_tokenized(sample_seed: pd.DataFrame, tokenizer: PreTrainedTokenizerBase) -> pd.DataFrame:
+def ensure_seed_can_be_tokenized(seed_data: pd.DataFrame, tokenizer: PreTrainedTokenizerBase) -> pd.DataFrame:
     def transform(x: str | pd._libs.missing.NAType) -> str:
         if pd.isna(x):
             null = tokenizer.decode(tokenizer.encode(JSON_NULL), skip_special_tokens=True)
@@ -53,7 +53,7 @@ def ensure_seed_can_be_tokenized(sample_seed: pd.DataFrame, tokenizer: PreTraine
         # skip tokens unseen during training
         return tokenizer.decode(tokenizer.encode(x), skip_special_tokens=True)
 
-    return sample_seed.astype(STRING).map(transform)
+    return seed_data.astype(STRING).map(transform)
 
 
 def create_schemas(

{mostlyai_engine-1.4.8 → mostlyai_engine-1.5.0}/mostlyai/engine/_tabular/argn.py

@@ -36,6 +36,8 @@ from torch import nn
 from mostlyai.engine._common import (
     CTXFLT,
     CTXSEQ,
+    RIDX_SUB_COLUMN_PREFIX,
+    SLEN_SUB_COLUMN_PREFIX,
     get_columns_from_cardinalities,
     get_sub_columns_from_cardinalities,
     get_sub_columns_lookup,
@@ -214,9 +216,19 @@ class Embedders(nn.Module):
         self.embedders = nn.ModuleDict()
 
         # embedding layers for each sub column defined in cardinalities
+        last_slen_sub_col = next(
+            (sub_col for sub_col in reversed(self.cardinalities) if sub_col.startswith(SLEN_SUB_COLUMN_PREFIX)), None
+        )
+        last_ridx_sub_col = next(
+            (sub_col for sub_col in reversed(self.cardinalities) if sub_col.startswith(RIDX_SUB_COLUMN_PREFIX)), None
+        )
         for sub_col, dim_input in self.cardinalities.items():
             dim_output = _embedding_heuristic(id=self.id(sub_col), model_size=model_size, dim_input=dim_input)
             embedder = nn.Embedding(num_embeddings=dim_input, embedding_dim=dim_output, device=device)
+            # the embeddings of the last slen and ridx sub columns are never used
+            # so we explicitly freeze them to make opacus not complain about "per sample gradient is not initialized"
+            if sub_col in [last_ridx_sub_col, last_slen_sub_col]:
+                embedder.weight.requires_grad = False
             self.add(sub_column=sub_col, embedder=embedder)
             self.dims.append(dim_output)
 
@@ -715,7 +727,7 @@ def _make_permutation_mask(
         # create mask in provided order
         order = torch.tensor([columns.index(c) for c in column_order], dtype=torch.int32)
     elif is_sequential and n_cols >= 1:
-        # create mask in random order, but keep SLEN/SIDX/SDEC column at first position
+        # create mask in random order, but keep positional columns at first position
         order = torch.randperm(n_cols - 1) + 1
         order = torch.cat((torch.zeros(1, dtype=torch.int32), order), dim=0)
     else:
@@ -1241,13 +1253,17 @@ class SequentialModel(nn.Module):
         if mode == "trn":
             # forward pass through sub column embedders
             tgt_embeds = self.embedders(x)
+            tgt_embeds_slen_ridx_masked = {
+                k: torch.zeros_like(v) if k.startswith((RIDX_SUB_COLUMN_PREFIX, SLEN_SUB_COLUMN_PREFIX)) else v
+                for k, v in tgt_embeds.items()
+            }
 
             # forward pass through column embedders
-            tgt_col_embeds = self.column_embedders(tgt_embeds)
+            tgt_col_embeds = self.column_embedders(tgt_embeds_slen_ridx_masked)
 
             # history
             # time shift: remove last time step; add zeros for first time step; add randoms for all others
-            embeddings = torch.cat(list(tgt_embeds.values()), dim=-1)
+            embeddings = torch.cat(list(tgt_embeds_slen_ridx_masked.values()), dim=-1)
             history_in = embeddings[:, :-1, :]
             history_in = nn.ConstantPad2d((0, 0, 1, 0), 0)(history_in)
             history, _ = self.history_compressor(history_in)
@@ -1274,9 +1290,17 @@ class SequentialModel(nn.Module):
                 masked_col_embeds = [torch.mul(col_mask[lookup.col_idx, :].int(), col_embeddings)]
 
                 # 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]
-                ]
+                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]
+                }
+                if sub_col.startswith(RIDX_SUB_COLUMN_PREFIX):
+                    # RIDX sub-columns should not see SLEN sub-columns
+                    prev_sub_col_embeds = {
+                        k: torch.zeros_like(v) if k.startswith(SLEN_SUB_COLUMN_PREFIX) else v
+                        for k, v in prev_sub_col_embeds.items()
+                    }
+                prev_sub_col_embeds = list(prev_sub_col_embeds.values())
 
                 # regressor
                 regressor_in = context_history + masked_col_embeds + prev_sub_col_embeds
@@ -1335,10 +1359,17 @@ class SequentialModel(nn.Module):
 
                 else:  # sample from distribution
                     # collect previous sub column embeddings for current column
-                    prev_sub_col_embeds = [
-                        tgt_embeds[sub_col]
+                    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]
-                    ]
+                    }
+                    if sub_col.startswith(RIDX_SUB_COLUMN_PREFIX):
+                        # RIDX sub-columns should not see SLEN sub-columns
+                        prev_sub_col_embeds = {
+                            k: torch.zeros_like(v) if k.startswith(SLEN_SUB_COLUMN_PREFIX) else v
+                            for k, v in prev_sub_col_embeds.items()
+                        }
+                    prev_sub_col_embeds = list(prev_sub_col_embeds.values())
 
                     # regressor
                     regressor_in = context_history + [col_embeddings] + prev_sub_col_embeds
@@ -1362,16 +1393,20 @@ class SequentialModel(nn.Module):
 
                 # update current sub column embedding
                 tgt_embeds[sub_col] = self.embedders.get(sub_col)(out)
+                tgt_embeds_slen_ridx_masked = {
+                    k: torch.zeros_like(v) if k.startswith((SLEN_SUB_COLUMN_PREFIX, RIDX_SUB_COLUMN_PREFIX)) else v
+                    for k, v in tgt_embeds.items()
+                }
 
                 # update current column embedding
                 if sub_col in self.tgt_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)
+                    col_embed_in = torch.cat([tgt_embeds_slen_ridx_masked[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)
 
             # update history and hidden state
-            history_in = torch.cat([v for v in tgt_embeds.values()], dim=-1)
+            history_in = torch.cat([v for v in tgt_embeds_slen_ridx_masked.values()], dim=-1)
             history, history_state = self.history_compressor(history_in, history_state=history_state)
 
             # order outputs according to tgt_sub_columns

{mostlyai_engine-1.4.8 → mostlyai_engine-1.5.0}/mostlyai/engine/_tabular/encoding.py

@@ -24,13 +24,13 @@ from mostlyai.engine._common import (
     ARGN_COLUMN,
     ARGN_PROCESSOR,
     ARGN_TABLE,
-    SDEC_SUB_COLUMN_PREFIX,
+    RIDX_SUB_COLUMN_PREFIX,
     SIDX_SUB_COLUMN_PREFIX,
     SLEN_SUB_COLUMN_PREFIX,
     TGT,
     ProgressCallback,
     ProgressCallbackWrapper,
-    encode_slen_sidx_sdec,
+    encode_positional_column,
     get_argn_name,
     get_sequence_length_stats,
     is_a_list,
@@ -129,33 +129,33 @@ def _encode_partition(
             n_jobs=n_jobs,
         )
         # pad each list with one extra item
-        df_ctx = pad_horizontally(df_ctx, padding_value=0, right=False)
+        df_ctx = pad_ctx_sequences(df_ctx)
 
     if is_sequential:
         assert isinstance(tgt_context_key, str)
         # trim sequences to (privacy-protected) max_len
         max_len = seq_len_stats["max"]
         df = df[df.groupby(tgt_context_key).cumcount() < max_len].reset_index(drop=True)
-        # enrich with sequence lengths and sequence indexes
-        df = _enrich_slen_sidx_sdec(df, tgt_context_key, max_len)
-        # flatten to list columns
-        df = flatten_frame(df, tgt_context_key)
+        # pad each list with one extra item
+        df = pad_tgt_sequences(df, context_key=tgt_context_key)
         # add empty records for IDs, that are present in context, but not in target; i.e., for zero-sequence records
         if has_context:
-            zero_seq_ids = list(set(df_ctx[ctx_primary_key]) - set(df[tgt_context_key]))
-            df_miss = pd.DataFrame({tgt_context_key: zero_seq_ids})
-            df_pads = pd.DataFrame({c: [[]] for c in df.columns if c != tgt_context_key})
-            df_miss = df_miss.merge(df_pads, how="cross")
-            df = pd.concat([df, df_miss], axis=0).reset_index(drop=True)
-        # pad each list with one extra item
-        df = pad_horizontally(df, padding_value=0, right=True)
+            zero_seq_ids = set(df_ctx[ctx_primary_key]) - set(df[tgt_context_key])
+            df_miss = pd.DataFrame(
+                [{tgt_context_key: i, **{c: 0 for c in df.columns if c != tgt_context_key}} for i in zero_seq_ids]
+            )
+            df = pd.concat([df, df_miss], ignore_index=True)
+        # enrich with positional columns
+        df = _enrich_positional_columns(df, tgt_context_key, max_len)
+        # flatten to list columns
+        df = flatten_frame(df, tgt_context_key)
     elif has_context:
         # add 0-rows for IDs, that are present in context, but not in target; i.e., for zero-sequence records
-        zero_seq_ids = list(set(df_ctx[ctx_primary_key]) - set(df[tgt_context_key]))
-        df_miss = pd.DataFrame({tgt_context_key: zero_seq_ids})
-        df_pads = pd.DataFrame({c: [0] for c in df.columns if c != tgt_context_key})
-        df_miss = df_miss.merge(df_pads, how="cross")
-        df = pd.concat([df, df_miss], axis=0).reset_index(drop=True)
+        zero_seq_ids = set(df_ctx[ctx_primary_key]) - set(df[tgt_context_key])
+        df_miss = pd.DataFrame(
+            [{tgt_context_key: i, **{c: 0 for c in df.columns if c != tgt_context_key}} for i in zero_seq_ids]
+        )
+        df = pd.concat([df, df_miss], ignore_index=True)
         # ensure that max 1 item is retained per context_id for flat mode
         df = df[df.groupby(tgt_context_key).cumcount() < 1]
 
@@ -380,29 +380,61 @@ def flatten_frame(df: pd.DataFrame, group_key: str) -> pd.DataFrame:
     return flattened_data
 
 
-def _enrich_slen_sidx_sdec(df: pd.DataFrame, context_key: str, max_seq_len: int) -> pd.DataFrame:
+def _enrich_positional_columns(df: pd.DataFrame, context_key: str, max_seq_len: int) -> pd.DataFrame:
     df = df.reset_index(drop=True)
-    slen = df.groupby(context_key)[context_key].transform("size")  # sequence length
     sidx = df.groupby(context_key).cumcount(ascending=True)  # sequence index
-    sdec = (10 * sidx / slen.clip(lower=1)).astype(int)  # sequence index decile
-    slen = encode_slen_sidx_sdec(slen, max_seq_len=max_seq_len, prefix=SLEN_SUB_COLUMN_PREFIX)
-    sidx = encode_slen_sidx_sdec(sidx, max_seq_len=max_seq_len, prefix=SIDX_SUB_COLUMN_PREFIX)
-    sdec = encode_slen_sidx_sdec(sdec, max_seq_len=max_seq_len, prefix=SDEC_SUB_COLUMN_PREFIX)
-    df = pd.concat([slen, sidx, sdec, df], axis=1)
+    slen = df.groupby(context_key)[context_key].transform("size") - 1  # sequence length; -1 to account for padding
+    ridx = df.groupby(context_key).cumcount(ascending=False)  # sequence remainder
+    sidx = encode_positional_column(sidx, max_seq_len=max_seq_len, prefix=SIDX_SUB_COLUMN_PREFIX)
+    slen = encode_positional_column(slen, max_seq_len=max_seq_len, prefix=SLEN_SUB_COLUMN_PREFIX)
+    ridx = encode_positional_column(ridx, max_seq_len=max_seq_len, prefix=RIDX_SUB_COLUMN_PREFIX)
+    df = pd.concat([sidx, slen, ridx, df], axis=1)
     return df
 
 
-def pad_horizontally(df: pd.DataFrame, padding_value: int, right=True) -> pd.DataFrame:
-    if df.shape[0] == 0:
-        return df
-    list_cols = [c for c in df.columns if is_a_list(df.loc[0, c])]
+def pad_tgt_sequences(df: pd.DataFrame, context_key: str, padding_value: int = 0) -> pd.DataFrame:
+    """
+    Pad one extra row to for each subject in the target data frame.
 
-    def pad_right(x):
-        return x + [padding_value] if len(x) == 0 else x
+    Args:
+        df: Exploded (unflattened) target data frame. Each event is a row.
+        context_key: Context key.
+        padding_value: Value to pad with for columns other than context_key.
 
-    def pad_left(x):
-        return [padding_value] + x if len(x) == 0 else x
+    Returns:
+        Padded target data frame.
+    """
+
+    def pad_row(x):
+        return pd.concat(
+            [
+                x,
+                pd.DataFrame(
+                    {col: [padding_value] for col in x.columns if col not in [context_key]}
+                    | {context_key: [x.iloc[0][context_key]]}
+                ),
+            ],
+            axis=0,
+        )
+
+    return df.groupby(context_key)[df.columns].apply(pad_row).reset_index(drop=True)
+
+
+def pad_ctx_sequences(df: pd.DataFrame, padding_value: int = 0) -> pd.DataFrame:
+    """
+    Pad one extra item to the context sequences with a given padding value.
+
+    Args:
+        df: Flattened context data frame.
+        padding_value: Value to pad with.
 
+    Returns:
+        Padded context data frame.
+    """
+    if df.shape[0] == 0:
+        return df
+    list_cols = [c for c in df.columns if is_a_list(df.loc[0, c])]
     for col in list_cols:
-        df[col] = df[col].apply(pad_right if right else pad_left)
+        # Note: only pad empty sequences to keep the backward compatibility
+        df[col] = df[col].apply(lambda x: x + [padding_value] if len(x) == 0 else x)
     return df