"""Orchestrator: encoder + projector + LFM2.5 wrapper → unified backbone. This module produces a `MultiHeadModel`-compatible object so the parent repo's `validate()`, `compute_losses`, and downstream-head infrastructure work without any rewriting. The encoder approach plugs into the parent harness as a different "backbone" — same downstream-head contract. The naming dance: - `EncoderBackbone` is what `MultiHeadModel.backbone` expects. It exposes `backbone_forward(token_ids) → (B, T, d_lfm)`. - `build_transaction_fm()` constructs the full `MultiHeadModel(backbone=EncoderBackbone, heads=…)` stack from configs. Head pool semantics with T=64, num_features=1: - `last_tx_mean` (num_features=1) ⇒ `hidden[:, -1:, :].mean(1)` = last pseudo-token = "tx 63's hidden state". Used by fraud (sequence-level). - `pre_last_tx` (num_features=1) ⇒ `hidden[:, -2, :]` = "tx 62's hidden state". Used by next_merchant / amount_range / mcc — predicts tx 63's features from the tx 62 representation (which has causally seen tx 0..62). The semantics line up with the parent's `T=960, num_features=15` layout: - Parent last_tx_mean reads positions 945..959 (tx 63's 15 features). - Parent pre_last_tx reads position 944 (end of tx 62). Encoder version reads the same logical positions at compressed sequence length. """ from __future__ import annotations from pathlib import Path from typing import Any import torch import torch.nn as nn from src.data.schema import SchemaConfig from src.model.task_heads import ( AnyHead, DownstreamHead, HeadConfig, MultiHeadModel, ) from encoder.src.model.encoder_heads import EncoderDownstreamHead, EncoderTiedEmbeddingHead from encoder.src.model.lfm_pseudo_token_wrapper import ( LfmPseudoTokenBackbone, build_lora_config, ) from encoder.src.model.projection_adapter import ProjectionAdapter from encoder.src.model.transaction_encoder import TransactionEncoder from encoder.src.model.transaction_encoder_nocompress import StructuredEncoder class EncoderBackbone(nn.Module): """Encoder (+ optional projector) + LFM2.5 wrapper as a single backbone. Implements the `backbone_forward(token_ids: Tensor) -> Tensor` contract that the parent's `MultiHeadModel` expects. Two modes, dispatched by the encoder type: - **Compress** (TransactionEncoder, default): encoder outputs (B, 64, d_encoder=256). Projector lifts to (B, 64, d_lfm=1024). Sequence length 64. - **Non-compress** (StructuredEncoder): encoder outputs (B, 960, d_lfm) directly. No projector. Sequence length 960. Matches the parent's structured-feature input shape. The projector is None in non-compress mode. The branch is selected at construction time by `build_transaction_fm`. """ def __init__( self, encoder: nn.Module, projector: ProjectionAdapter | None, lfm_wrapper: LfmPseudoTokenBackbone, ) -> None: super().__init__() self.encoder = encoder # nn.Module attribute, may be None in non-compress mode. self.projector = projector self.lfm = lfm_wrapper self.d_lfm = lfm_wrapper.d_lfm def backbone_forward(self, token_ids: torch.Tensor) -> torch.Tensor: # token_ids: (B, 64, 15) int64 # encoder output shape depends on mode: # compress → (B, 64, d_encoder) # nocompress → (B, 960, d_lfm) x = self.encoder(token_ids) if self.projector is not None: # Compress mode: lift d_encoder → d_lfm. x = self.projector(x) # x: (B, T, d_lfm) where T = 64 (compress) or 960 (nocompress) # Cast to the LFM base's dtype before injection. The LFM was loaded # in bf16 (GPU) or fp32 (CPU); upstream modules produce whatever # they want (default fp32). This cast bridges the boundary. target_dtype = next(self.lfm.base.parameters()).dtype if x.dtype != target_dtype: x = x.to(target_dtype) hidden = self.lfm(x) # → (B, T, d_lfm) in target_dtype return hidden # `MultiHeadModel.forward` calls `self.backbone.backbone_forward`, but # for debugging / standalone use we also expose `forward` as an alias. def forward(self, token_ids: torch.Tensor) -> torch.Tensor: return self.backbone_forward(token_ids) def build_heads_encoder( head_configs: dict[str, dict[str, Any]], hidden_dim: int, num_features: int, value_tables: nn.ModuleList | None = None, ) -> dict[str, AnyHead]: """Instantiate downstream heads from a config dict. `num_features` is the per-tx token count and shapes the pool strategy: - **compress mode**: `num_features=1` (one pseudo-token per tx) → `last_tx_mean` pools position -1, `pre_last_tx` pools position -2. - **nocompress mode**: `num_features=15` (matches parent's structured layout) → `last_tx_mean` pools last 15 positions, `pre_last_tx` pools position -(15+1) = -16 = end of tx 62. Two head implementations are dispatched: - **EncoderDownstreamHead** (default): fresh MLP from pool → output. - **EncoderTiedEmbeddingHead**: shares its classifier matrix with one of the encoder's per-feature value tables. Selected by `tied: true` on a head's config. Requires `value_tables` to be passed (only available in `nocompress` mode where the encoder is a `StructuredEncoder`); `compress` mode's TransactionEncoder uses d_feat=32 feature embeddings that can't be tied at d_lfm=1024. Heads with `tied: true` MUST have `target_type: "feature:N"` matching the value table they share weights with. The head's `output_dim` is ignored in tied mode (implicitly the feature's vocab_size). """ heads: dict[str, AnyHead] = {} for name, hcfg in head_configs.items(): hc = HeadConfig( name=name, output_dim=hcfg["output_dim"], loss_type=hcfg["loss"], pool_strategy=hcfg["pool"], target_type=hcfg["target"], weight=hcfg.get("weight", 1.0), mlp_hidden=hcfg.get("mlp_hidden", 128), dropout=hcfg.get("dropout", 0.1), ) if hcfg.get("tied", False): if value_tables is None: raise ValueError( f"Head '{name}' has tied=true but no value_tables were " f"passed. Tied heads are only supported in nocompress " f"mode (where the encoder is StructuredEncoder).", ) heads[name] = EncoderTiedEmbeddingHead( hc, hidden_dim, num_features=num_features, value_tables=value_tables, ) else: heads[name] = EncoderDownstreamHead( hc, hidden_dim, num_features=num_features, ) return heads def build_transaction_fm( schema: SchemaConfig, head_configs: dict[str, dict[str, Any]], model_path: str | Path, architecture_cfg: dict[str, Any] | None = None, encoder_cfg: dict[str, Any] | None = None, projector_cfg: dict[str, Any] | None = None, lora_cfg: dict[str, Any] | None = None, dtype: torch.dtype = torch.bfloat16, device_map: str | None = "auto", ) -> MultiHeadModel: """Construct the full encoder + LFM + heads stack. The architecture mode is selected by `architecture_cfg["mode"] ∈ {"compress", "nocompress"}`: - **compress** (default): TransactionEncoder → ProjectionAdapter → LFM. One pseudo-token per transaction; sequence length 64. - **nocompress**: StructuredEncoder → LFM directly. 15 pseudo-tokens per transaction (one per feature); sequence length 960. Mirrors parent's structured-feature input shape. Returns a `MultiHeadModel` so it's drop-in compatible with the parent's `validate()`, `compute_losses`, and training utilities regardless of mode. """ architecture_cfg = architecture_cfg or {} encoder_cfg = encoder_cfg or {} projector_cfg = projector_cfg or {} lora_cfg = lora_cfg or {} mode = architecture_cfg.get("mode", "compress") # Build the LFM wrapper first so we know d_lfm before constructing the # encoder/projector. (Compress mode needs d_lfm for the projector; # nocompress mode needs it for the StructuredEncoder.) lora = None if lora_cfg.get("enabled", True): # Allow config to override the default target module list. Passing # an explicit list (e.g. ["q_proj", ..., "in_proj"]) opts into # conv-LoRA. None falls through to build_lora_config's default # (attention names with conv.out_proj collision). lora = build_lora_config( r=lora_cfg.get("r", 16), alpha=lora_cfg.get("alpha", 32), dropout=lora_cfg.get("dropout", 0.05), target_modules=lora_cfg.get("target_modules"), strict_attention_only=lora_cfg.get("strict_attention_only", False), ) lfm_wrapper = LfmPseudoTokenBackbone( model_path, lora=lora, dtype=dtype, device_map=device_map, # `freeze_base=True` is the encoder-pattern default. The stage-2 # diagnostic (architecture.unfreeze_backbone=true) flips this so # the full 354M backbone is end-to-end trainable. freeze_base=not architecture_cfg.get("unfreeze_backbone", False), ) if mode == "compress": encoder = TransactionEncoder( schema, d_feat=encoder_cfg.get("d_feat", 32), d_encoder=encoder_cfg.get("d_encoder", 256), mlp_hidden=encoder_cfg.get("mlp_hidden", 384), ) projector: ProjectionAdapter | None = ProjectionAdapter( d_encoder=encoder_cfg.get("d_encoder", 256), d_lfm=lfm_wrapper.d_lfm, hidden=projector_cfg.get("hidden", 2 * lfm_wrapper.d_lfm), use_layernorm=projector_cfg.get("use_layernorm", True), ) num_features_for_heads = 1 elif mode == "nocompress": encoder = StructuredEncoder(schema, d_lfm=lfm_wrapper.d_lfm) projector = None # Matches parent: heads pool over 15-feature transaction stripes. num_features_for_heads = schema.num_features else: raise ValueError( f"Unknown architecture mode: {mode!r}. Expected 'compress' or 'nocompress'.", ) backbone = EncoderBackbone(encoder, projector, lfm_wrapper) # Tied heads need access to the encoder's per-feature value tables. # Only StructuredEncoder (nocompress mode) exposes them at d_lfm; the # compress mode's TransactionEncoder has d_feat=32 feature embeddings # that can't share weights with a d_lfm-dim classifier. value_tables_for_heads = ( encoder.value_tables if mode == "nocompress" else None ) heads = build_heads_encoder( head_configs, hidden_dim=lfm_wrapper.d_lfm, num_features=num_features_for_heads, value_tables=value_tables_for_heads, ) return MultiHeadModel(backbone=backbone, heads=heads)