encoder/src/model/heads/probability_head.py

"""Per-surface probability head: 3-class categorical with mean-pool.

This is the headline head for every multi-surface in ADR 0003. One
calibrated score per surface, expressed as a 3-class categorical
distribution {unlikely, ambiguous, likely}.

Why categorical instead of a single sigmoid scalar (the obvious choice):

    The `liquid-finetuning-playbook` Section 10 calls out that LFM2.5-
    350M numeric outputs cluster at attractor states (0.05, 0.45, 0.85,
    0.95). A continuous "friendly-fraud probability" trained on 350M
    will fight those attractors. The fix doctrine: use categorical
    labels (low / medium / high / critical) and let softmax produce a
    probability distribution. For dispute legitimacy:

        class 0 = unlikely      (low friendly-fraud probability)
        class 1 = ambiguous     (mid; analyst should review)
        class 2 = likely        (high friendly-fraud probability)

    The "score" displayed in the UI is `softmax(logits)[2]` — the
    probability assigned to the `likely` class. Softmax is naturally
    calibrated within {0, 1} without fighting attractor states because
    the model picks between three peaks, not one continuous range.

Why mean-pool over the tx positions only:

    LFM2.5's final layer is conv (kernel 3). Last-token pooling reads
    only a 3-token receptive field (`liquid-models-architecture`
    Section 3, 16). For sequence classification we mean-pool over all
    64 transaction pseudo-token positions, which inherits global
    context from the 6 attention layers at L2/L5/L8/L10/L12/L14.

    Importantly: we mean-pool over the 64 *transaction* positions, NOT
    the entire (tx + sep + text) combined sequence. The text tokens
    are conditioning information; the score belongs to the customer's
    behavioral signature, which lives in the transaction positions.

Shape contract:
    Input  hidden_states: (B, T_total, D)  where T_total = 64 + 1 + T_txt
    Output logits:        (B, num_classes)  default num_classes = 3
"""

from __future__ import annotations

from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F


@dataclass
class ProbabilityHeadConfig:
    """Probability head configuration.

    Attributes:
        name: Identifier for logging and config dispatch.
        num_classes: Categorical output dimension. Default 3 for the
            unlikely/ambiguous/likely scheme. Surfaces with a richer
            taxonomy (collections treatment, fraud pattern) override.
        hidden_dim: Backbone hidden size. 1024 for LFM2.5-350M, 2048
            for LFM2.5-1.2B.
        mlp_hidden: Width of the 2-layer MLP head. 256 is comfortable
            for 350M; 30 KB at fp32, ~7 KB at int8.
        dropout: Dropout between the two MLP layers. 0.1 default.
        num_tx_positions: Number of transaction pseudo-token positions
            at the start of the sequence. 64 in Phase 1. The head
            mean-pools over `hidden_states[:, :num_tx_positions, :]`.
    """

    name: str
    num_classes: int = 3
    hidden_dim: int = 1024
    mlp_hidden: int = 256
    dropout: float = 0.1
    num_tx_positions: int = 64


class ProbabilityHead(nn.Module):
    """Mean-pool over the transaction pseudo-tokens, then MLP, then softmax.

    Parameter count at default config (hidden_dim=1024, mlp_hidden=256,
    num_classes=3): ~263K params. Small relative to the LoRA adapter
    (~1M) and immaterial against the 350M backbone.

    The head is intentionally minimal. The work is done by the
    backbone + per-surface LoRA; the head is a thin calibration layer
    on top of the pooled representation.
    """

    def __init__(self, config: ProbabilityHeadConfig) -> None:
        super().__init__()
        self.config = config

        self.mlp = nn.Sequential(
            nn.Linear(config.hidden_dim, config.mlp_hidden),
            nn.ReLU(),
            nn.Dropout(config.dropout),
            nn.Linear(config.mlp_hidden, config.num_classes),
        )

    def pool(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """Mean-pool over the leading `num_tx_positions` positions.

        Args:
            hidden_states: (B, T_total, D) hidden states from the
                backbone after attending over [tx pseudo-tokens, SEP,
                text tokens]. T_total = num_tx_positions + 1 + T_txt.

        Returns:
            (B, D) pooled vector over the transaction positions only.
        """
        # (B, num_tx_positions, D) -> (B, D); cast to head dtype since
        # the backbone runs in bf16 but the small MLP head stays in
        # fp32 for numerical stability.
        head_dtype = next(self.mlp.parameters()).dtype
        tx_slice = hidden_states[:, : self.config.num_tx_positions, :].to(head_dtype)
        return tx_slice.mean(dim=1)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """Run the full head.

        Args:
            hidden_states: (B, T_total, D).

        Returns:
            (B, num_classes) raw logits. Caller applies softmax for
            probabilities or cross-entropy for loss.
        """
        # (B, D)
        pooled = self.pool(hidden_states)
        # (B, num_classes)
        return self.mlp(pooled)

    def compute_loss(
        self,
        logits: torch.Tensor,
        targets: torch.Tensor,
    ) -> torch.Tensor:
        """Cross-entropy over the categorical distribution.

        Args:
            logits:  (B, num_classes) raw logits from `forward`.
            targets: (B,) int64 class indices in [0, num_classes).

        Returns:
            Scalar CE loss.
        """
        return F.cross_entropy(logits, targets)

    @torch.no_grad()
    def score(self, logits: torch.Tensor) -> torch.Tensor:
        """Friendly-fraud probability for display in the UI.

        Returns softmax(logits)[..., -1] — the probability mass on the
        last class, which by convention is the "likely" / high-score
        class. For 3-class {unlikely, ambiguous, likely}, this is
        P(likely). Surfaces with different taxonomies can override
        which class index is the "headline" by passing it in.

        Args:
            logits: (B, num_classes).

        Returns:
            (B,) in [0, 1] — calibrated probability mass on the
            highest-severity class.
        """
        # (B, num_classes) -> (B,)
        return F.softmax(logits, dim=-1)[..., -1]

    @torch.no_grad()
    def predict_band(self, logits: torch.Tensor) -> torch.Tensor:
        """Predicted class index. Argmax over softmax.

        Args:
            logits: (B, num_classes).

        Returns:
            (B,) int64 class indices.
        """
        return logits.argmax(dim=-1)

    def num_parameters(self) -> int:
        return sum(p.numel() for p in self.parameters())