encoder/src/model/transaction_fm.py

"""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)