"""Orchestrator: encoder + projector + multi-surface LFM wrapper + 3 heads. Sibling of `transaction_fm.build_transaction_fm`. Reuses the encoder and projector from the existing demo so feature-token handling is identical. Replaces: - `LfmPseudoTokenBackbone` → `LfmMultiSurfaceBackbone` (loads Lfm2ForCausalLM, exposes LM head + text path). - The four classification heads → one probability head + one attribution head; the LM head from the backbone provides reasoning generation. The forward signatures are NOT the parent's MultiHeadModel contract (`(B, T, F) int -> dict[head_name, logits]`). This model takes a `MixedModalityBatch` and dispatches based on `batch.head_target`. This makes per-batch homogeneous-head sampling explicit at the type level: a batch is tagged for one head, and the model only computes that head's loss. Per-surface LoRA is implemented by loading a different LoRA artifact per surface (each `LfmMultiSurfaceBackbone` instance carries one adapter). A multi-surface inference server keeps one backbone in memory with a LoRA registry that hot-swaps via PEFT's `set_adapter`. That's a serving concern; this module trains one surface at a time. """ from __future__ import annotations from pathlib import Path from typing import Any import torch import torch.nn as nn import torch.nn.functional as F from src.data.schema import SchemaConfig from encoder.src.data.mixed_modality import ( MixedModalityBatch, build_combined_attention_mask, build_lm_target_layout, ) from encoder.src.model.heads.attribution_head import ( AttributionHead, AttributionHeadConfig, ) from encoder.src.model.heads.fraud_pattern_head import ( FraudPatternHead, FraudPatternHeadConfig, ) from encoder.src.model.heads.multi_treatment_probability_head import ( MultiTreatmentProbabilityHead, MultiTreatmentProbabilityHeadConfig, ) from encoder.src.model.heads.probability_head import ( ProbabilityHead, ProbabilityHeadConfig, ) from encoder.src.model.lfm_multisurface_wrapper import ( LfmMultiSurfaceBackbone, build_multisurface_lora_config, ) from encoder.src.model.projection_adapter import ProjectionAdapter from encoder.src.model.transaction_encoder import TransactionEncoder # Index of the "headline" class for the probability head. By convention # class 2 ({unlikely=0, ambiguous=1, likely=2}) is the high-severity # class whose softmax probability gets surfaced in the UI as the score. HEADLINE_CLASS_INDEX = 2 # LM-loss ignore index. Matches PyTorch CrossEntropyLoss default and # the value used in `build_lm_target_layout`. LM_IGNORE_INDEX = -100 class TransactionMultiSurfaceModel(nn.Module): """Encoder + projector + multi-surface LFM backbone + 3 heads. Forward dispatches based on `batch.head_target`: - "probability" → run encoder + projector + backbone (no LM head) + probability head; loss = CE(logits, labels). - "attribution" → run encoder + projector + backbone (no LM head) + attribution head; loss = weighted BCE. - "lm" → run encoder + projector + backbone (WITH LM head) + LM next-token loss on the text positions. All three paths share the same encoder + projector + backbone weights. The only difference is which head's loss is computed, enforcing the per-batch homogeneous-head sampling doctrine. Args: encoder: TransactionEncoder (compress mode). projector: ProjectionAdapter mapping (B, 64, d_encoder) → (B, 64, d_lfm). backbone: LfmMultiSurfaceBackbone. prob_head: ProbabilityHead. attr_head: AttributionHead. """ def __init__( self, encoder: TransactionEncoder, projector: ProjectionAdapter, backbone: LfmMultiSurfaceBackbone, prob_head: ProbabilityHead | MultiTreatmentProbabilityHead | FraudPatternHead, attr_head: AttributionHead, ) -> None: super().__init__() self.encoder = encoder self.projector = projector self.backbone = backbone self.prob_head = prob_head self.attr_head = attr_head # ---- shared forward components ---- def _build_combined_embeds( self, feature_ids: torch.Tensor, text_input_ids: torch.Tensor, disputed_idx: torch.Tensor | None = None, ) -> torch.Tensor: """Encoder + projector + concat with SEP + text embeddings. Args: feature_ids: (B, 64, 15) int64 tx feature ids. text_input_ids: (B, T_txt) int64 text tokens. disputed_idx: optional (B,) int64. When present, the encoder adds its learnable disputed marker at that position so the backbone knows which transaction is being asked about. Returns: (B, 64 + 1 + T_txt, d_lfm) combined embedding sequence, cast to the backbone's dtype. """ # (B, 64, d_encoder) tx_encoded = self.encoder(feature_ids, disputed_idx=disputed_idx) # (B, 64, d_lfm) tx_pseudo = self.projector(tx_encoded) # Cast to backbone dtype (bf16 on GPU, fp32 on CPU). target_dtype = next(self.backbone.base.parameters()).dtype if tx_pseudo.dtype != target_dtype: tx_pseudo = tx_pseudo.to(target_dtype) batch_size = feature_ids.shape[0] device = feature_ids.device # (B, T_txt, d_lfm) text_embeds = self.backbone.embed_text(text_input_ids).to(target_dtype) # (B, 1, d_lfm) sep_embeds = self.backbone.embed_sep(batch_size, device).to(target_dtype) # (B, 64 + 1 + T_txt, d_lfm) return torch.cat([tx_pseudo, sep_embeds, text_embeds], dim=1) # ---- per-head forward + loss ---- def forward_probability( self, batch: MixedModalityBatch, ) -> dict[str, torch.Tensor]: """Probability head path. Computes CE loss against labels_probability.""" if batch.labels_probability is None: raise ValueError( "Batch tagged for probability head has no labels_probability.", ) combined = self._build_combined_embeds( batch.feature_ids, batch.text_input_ids, disputed_idx=batch.disputed_idx, ) attn_mask = build_combined_attention_mask( batch_size=batch.batch_size, num_tx_positions=batch.num_tx_positions, text_attention_mask=batch.text_attention_mask, device=combined.device, ) outputs = self.backbone.forward_mixed( combined_embeds=combined, attention_mask=attn_mask, compute_lm_logits=False, ) # (B, 3) logits = self.prob_head(outputs["hidden_states"]) loss = self.prob_head.compute_loss(logits, batch.labels_probability) return {"loss": loss, "logits": logits} def forward_attribution( self, batch: MixedModalityBatch, ) -> dict[str, torch.Tensor]: """Attribution head path. Computes BCE loss per tx position.""" if batch.labels_attribution is None: raise ValueError( "Batch tagged for attribution head has no labels_attribution.", ) combined = self._build_combined_embeds( batch.feature_ids, batch.text_input_ids, disputed_idx=batch.disputed_idx, ) attn_mask = build_combined_attention_mask( batch_size=batch.batch_size, num_tx_positions=batch.num_tx_positions, text_attention_mask=batch.text_attention_mask, device=combined.device, ) outputs = self.backbone.forward_mixed( combined_embeds=combined, attention_mask=attn_mask, compute_lm_logits=False, ) # (B, 64) logits = self.attr_head(outputs["hidden_states"]) loss = self.attr_head.compute_loss(logits, batch.labels_attribution) return {"loss": loss, "logits": logits} def forward_lm( self, batch: MixedModalityBatch, ) -> dict[str, torch.Tensor]: """LM head path. Teacher-forced next-token CE on text positions. The text portion of `combined` is what we condition AND predict. The LM target layout puts the text token ids at positions [num_tx_positions + 1, T_total) and -100 elsewhere. Standard next-token shift: logits[..., :-1, :] predict targets[..., 1:]. """ combined = self._build_combined_embeds( batch.feature_ids, batch.text_input_ids, disputed_idx=batch.disputed_idx, ) attn_mask = build_combined_attention_mask( batch_size=batch.batch_size, num_tx_positions=batch.num_tx_positions, text_attention_mask=batch.text_attention_mask, device=combined.device, ) outputs = self.backbone.forward_mixed( combined_embeds=combined, attention_mask=attn_mask, compute_lm_logits=True, ) # (B, T_total, vocab_size) lm_logits = outputs["lm_logits"] # Targets aligned to combined sequence; ignore_index at non-text # positions and at text padding. targets = build_lm_target_layout( batch_size=batch.batch_size, num_tx_positions=batch.num_tx_positions, text_input_ids=batch.text_input_ids, text_attention_mask=batch.text_attention_mask, ignore_index=LM_IGNORE_INDEX, ) # Shift for next-token prediction: predict pos t+1 from pos t. # (B, T_total - 1, vocab_size) shift_logits = lm_logits[..., :-1, :].contiguous() # (B, T_total - 1) shift_targets = targets[..., 1:].contiguous() # Flatten and compute CE. ignore_index handles non-text and padding. loss = F.cross_entropy( shift_logits.view(-1, shift_logits.size(-1)), shift_targets.view(-1), ignore_index=LM_IGNORE_INDEX, ) return {"loss": loss, "lm_logits": lm_logits} def forward( self, batch: MixedModalityBatch, ) -> dict[str, torch.Tensor]: """Dispatch based on `batch.head_target`. The single entry point. The surface_trainer's batch sampler tags each batch with one head target; this forward enforces the doctrine that each minibatch updates exactly one head's loss. Contract gate: every dispute-legitimacy batch must carry `disputed_idx`. The label, attribution, and reasoning are all conditional on which transaction is being disputed; a batch without disputed_idx makes the task unidentifiable and the model silently degenerates to memorization. Fail loudly here so the bug is caught at the first forward pass rather than after a full training run. """ if batch.disputed_idx is None: raise ValueError( "MixedModalityBatch is missing disputed_idx. The dispute " "legitimacy task's labels are conditional on disputed_idx; " "the model cannot learn without it. Check the dataset " "collate function.", ) head = batch.head_target if head == "probability": return self.forward_probability(batch) if head == "attribution": return self.forward_attribution(batch) if head == "lm": return self.forward_lm(batch) raise ValueError( f"Unknown head_target: {head!r}. " f"Expected one of {{'probability', 'attribution', 'lm'}}.", ) # ---- inference helpers ---- @torch.no_grad() def predict( self, batch: MixedModalityBatch, ) -> dict[str, torch.Tensor]: """Compute all three heads' outputs in one forward (eval-time). Used by the demo UI: one inference returns the probability score, the top-k attributed transactions, and (if requested) the generated reasoning. No loss computed. """ combined = self._build_combined_embeds( batch.feature_ids, batch.text_input_ids, disputed_idx=batch.disputed_idx, ) attn_mask = build_combined_attention_mask( batch_size=batch.batch_size, num_tx_positions=batch.num_tx_positions, text_attention_mask=batch.text_attention_mask, device=combined.device, ) outputs = self.backbone.forward_mixed( combined_embeds=combined, attention_mask=attn_mask, compute_lm_logits=False, ) hidden = outputs["hidden_states"] prob_logits = self.prob_head(hidden) attr_logits = self.attr_head(hidden) return { "probability_logits": prob_logits, "probability_score": self.prob_head.score(prob_logits), "probability_band": self.prob_head.predict_band(prob_logits), "attribution_logits": attr_logits, "attribution_probabilities": self.attr_head.probabilities(attr_logits), "attribution_top_k": self.attr_head.top_k_positions(attr_logits, k=8), } # ---- introspection ---- def num_parameters(self) -> dict[str, int]: """Parameter breakdown across components.""" return { "encoder": sum(p.numel() for p in self.encoder.parameters()), "projector": sum(p.numel() for p in self.projector.parameters()), "backbone_total": self.backbone.total_parameters(), "backbone_trainable": self.backbone.trainable_parameters(), "prob_head": self.prob_head.num_parameters(), "attr_head": self.attr_head.num_parameters(), "trainable_total": sum( p.numel() for p in self.parameters() if p.requires_grad ), "total": sum(p.numel() for p in self.parameters()), } def build_transaction_multisurface( schema: SchemaConfig, model_path: str | Path, encoder_cfg: dict[str, Any] | None = None, projector_cfg: dict[str, Any] | None = None, head_cfg: dict[str, Any] | None = None, lora_cfg: dict[str, Any] | None = None, dtype: torch.dtype = torch.bfloat16, device_map: str | None = "auto", ) -> TransactionMultiSurfaceModel: """Construct the multi-surface model from configs. Args: schema: SchemaConfig from the parent's data module. model_path: HF-format directory or repo ID for LFM2.5. encoder_cfg: optional overrides for the TransactionEncoder (d_feat, d_encoder, mlp_hidden). projector_cfg: optional overrides for the ProjectionAdapter (hidden, use_layernorm). head_cfg: dict with `probability` and `attribution` sub-configs. See the YAML configs for shape. lora_cfg: per-surface LoRA configuration. {r, alpha, dropout, target_modules}. dtype: bfloat16 (GPU) or float32 (CPU). device_map: "auto" for GPU, None for CPU. Returns: TransactionMultiSurfaceModel ready for training or inference. """ encoder_cfg = encoder_cfg or {} projector_cfg = projector_cfg or {} head_cfg = head_cfg or {} lora_cfg = lora_cfg or {} # 1) Backbone — loads weights, applies LoRA. We build this first # because the projector + heads need d_lfm. lora = None if lora_cfg.get("enabled", True): lora = build_multisurface_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"), ) backbone = LfmMultiSurfaceBackbone( model_path, lora=lora, dtype=dtype, device_map=device_map, freeze_base=True, # encoder pattern; LoRA + heads carry the gradient ) # 2) Encoder + projector — compress mode (one pseudo-token per tx). 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), enable_collections_markers=encoder_cfg.get( "enable_collections_markers", False, ), enable_fraud_markers=encoder_cfg.get( "enable_fraud_markers", False, ), ) projector = ProjectionAdapter( d_encoder=encoder_cfg.get("d_encoder", 256), d_lfm=backbone.d_lfm, hidden=projector_cfg.get("hidden", 2 * backbone.d_lfm), use_layernorm=projector_cfg.get("use_layernorm", True), ) # 3) Heads. prob_cfg_dict = head_cfg.get("probability", {}) head_type = prob_cfg_dict.get("type", "probability") prob_head: ProbabilityHead | MultiTreatmentProbabilityHead | FraudPatternHead if head_type == "multi_treatment_probability": prob_head = MultiTreatmentProbabilityHead( MultiTreatmentProbabilityHeadConfig( name=prob_cfg_dict.get("name", "collections_treatment"), num_treatments=prob_cfg_dict.get("num_treatments", 4), num_bands=prob_cfg_dict.get("num_bands", 3), hidden_dim=backbone.d_lfm, mlp_hidden=prob_cfg_dict.get("mlp_hidden", 512), dropout=prob_cfg_dict.get("dropout", 0.1), num_tx_positions=prob_cfg_dict.get("num_tx_positions", 64), band_class_weights=prob_cfg_dict.get("band_class_weights"), ) ) elif head_type == "fraud_pattern": prob_head = FraudPatternHead( FraudPatternHeadConfig( name=prob_cfg_dict.get("name", "fraud_pattern"), num_stages=prob_cfg_dict.get("num_stages", 5), num_types=prob_cfg_dict.get("num_types", 4), hidden_dim=backbone.d_lfm, mlp_hidden=prob_cfg_dict.get("mlp_hidden", 256), dropout=prob_cfg_dict.get("dropout", 0.1), num_tx_positions=prob_cfg_dict.get("num_tx_positions", 64), stage_class_weights=prob_cfg_dict.get("stage_class_weights"), type_class_weights=prob_cfg_dict.get("type_class_weights"), ) ) elif head_type == "probability": prob_head = ProbabilityHead( ProbabilityHeadConfig( name=prob_cfg_dict.get("name", "dispute_legitimacy"), num_classes=prob_cfg_dict.get("num_classes", 3), hidden_dim=backbone.d_lfm, mlp_hidden=prob_cfg_dict.get("mlp_hidden", 256), dropout=prob_cfg_dict.get("dropout", 0.1), num_tx_positions=prob_cfg_dict.get("num_tx_positions", 64), ) ) else: raise ValueError( f"Unknown probability head type: {head_type!r}. Expected one of " f"{{'probability', 'multi_treatment_probability', 'fraud_pattern'}}.", ) attr_cfg_dict = head_cfg.get("attribution", {}) attr_head = AttributionHead( AttributionHeadConfig( name=attr_cfg_dict.get("name", "behavioral_attribution"), hidden_dim=backbone.d_lfm, mlp_hidden=attr_cfg_dict.get("mlp_hidden", 64), dropout=attr_cfg_dict.get("dropout", 0.1), num_tx_positions=attr_cfg_dict.get("num_tx_positions", 64), pos_weight=attr_cfg_dict.get("pos_weight", 5.0), ) ) return TransactionMultiSurfaceModel( encoder=encoder, projector=projector, backbone=backbone, prob_head=prob_head, attr_head=attr_head, )