Spaces:

cdotsanghvi
/

lfm2-transaction-encoder

Sleeping

lfm2-transaction-encoder / encoder /src /model /transaction_fm.py

add multi-head demo as 4th-6th tabs; restore Why Liquid + Integration

083b138 5 months ago

11.2 kB

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

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