encoder/src/demo/copilot_inference.py

"""Inference plumbing for the Co-Pilot demo.

Loads the slim multi-surface checkpoint, exposes one `predict` entry
point per cast member, and a token-by-token `stream_reasoning` generator
for the LM head. Everything is CPU-only, float32, greedy decode, fixed
seed: a demo must produce byte-identical output on every replay.

Three things this module enforces (no autoreload, no surprises):

  1. The frozen LFM2.5-350M base weights are reloaded fresh from HF
     when the model is constructed. The slim checkpoint only carries
     the trainable params (encoder + projector + LoRA + heads); it is
     loaded with `strict=False`.
  2. The encoder is wrapped with `torch.inference_mode()` for every
     call so no autograd memory accumulates. The Gradio process runs
     many sequential inferences; we cannot leak.
  3. Every inference call seeds `torch.manual_seed(42)` before the
     forward. The model has dropout layers; deterministic seeding makes
     them no-ops in eval mode but also pins any future stochastic
     decode behavior.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Iterator

import numpy as np
import torch
import yaml

from src.data.schema import SchemaConfig, load_schema

from encoder.src.data.mixed_modality import (
    MixedModalityBatch,
    build_combined_attention_mask,
    tokenize_texts,
)
from encoder.src.data.synthetic_dispute_legitimacy import (
    FEATURE_COUNTRY,
    FEATURE_CUSTOMER_MERCHANT_COUNT,
    FEATURE_CVV,
    FEATURE_AVS,
    FEATURE_ENTRY_MODE,
    FEATURE_MERCHANT_ID,
    CVV_MATCH,
    ENTRY_CNP,
    LABEL_LIKELY,
    LABEL_UNLIKELY,
    LABEL_AMBIGUOUS,
    RESERVED_OFFSET,
    _approximate_amount_usd,
    _decode_country_label,
)
from encoder.src.data.synthetic_dispute_legitimacy import FEATURE_AMOUNT
from encoder.src.model.transaction_fm_multisurface import (
    TransactionMultiSurfaceModel,
    build_transaction_multisurface,
)


# Greedy decode + fixed seed in every entry point. Demos must replay
# byte-identical on demand.
DEMO_SEED = 42

# Label index → human label for the probability head's 3-class output.
LABEL_NAMES: dict[int, str] = {
    0: "unlikely friendly fraud",
    1: "ambiguous",
    2: "likely friendly fraud",
}

# Color band per class (used by the renderer; kept here so the band
# and the predicted label can never drift apart).
LABEL_COLORS: dict[int, str] = {
    0: "#22c55e",  # green — unlikely (legit dispute)
    1: "#eab308",  # yellow — ambiguous
    2: "#ef4444",  # red — likely friendly fraud
}


@dataclass
class CastMember:
    """One entry from `encoder/data/demo_cast.json`, used by the UI.

    Attributes correspond 1:1 with the JSON fields. The renderer reads
    `display_name`, `complaint_text`, `tone`, `expected_label`; the
    inference module reads `customer_idx`, `disputed_idx`,
    `complaint_text`.
    """

    pattern: str
    display_name: str
    customer_idx: int
    disputed_idx: int
    expected_label: int
    expected_label_name: str
    complaint_text: str
    tone: str
    description: str


@dataclass
class CopilotResult:
    """Result of one inference call against the multi-surface model.

    Attributes:
        score: float in [0, 1] — softmax probability of the predicted
            class. Caller renders this as the headline number.
        predicted_class: int in {0, 1, 2} — argmax of probability logits.
        predicted_label: human-readable label string.
        color: hex color for the band (mapped from predicted_class).
        attribution_probs: (64,) float — per-position contribution
            probabilities. Renderer thresholds / top-k's these.
        top_k_positions: (k,) int — indices of the top-contributing
            transactions, sorted by descending attribution probability.
        disputed_idx: the disputed transaction's index, passed through
            so the renderer can mark it with a star.
    """

    score: float
    predicted_class: int
    predicted_label: str
    color: str
    attribution_probs: np.ndarray
    top_k_positions: np.ndarray
    disputed_idx: int


class CopilotModel:
    """Encapsulates the loaded model, tokenizer, schema, histories,
    and demo cast. Built once at app startup, then used for every
    inference call.

    Why a class and not free functions: the model, tokenizer, and the
    in-memory histories array are ~700 MB on disk. We load them once
    and reuse, not per-request.
    """

    def __init__(
        self,
        model: TransactionMultiSurfaceModel,
        schema: SchemaConfig,
        histories: np.ndarray,
        cast: list[CastMember],
        device: torch.device,
    ) -> None:
        self.model = model
        self.schema = schema
        self.histories = histories
        self.cast = cast
        self.device = device
        # Tokenizer comes from the backbone wrapper.
        self.tokenizer = model.backbone.tokenizer

        self.model.eval()

    @classmethod
    def from_paths(
        cls,
        checkpoint_path: Path,
        model_config_path: Path,
        schema_path: Path,
        histories_path: Path,
        cast_path: Path,
        device: torch.device = torch.device("cpu"),
    ) -> "CopilotModel":
        """Construct the full inference stack from on-disk artifacts.

        Args:
            checkpoint_path: slim checkpoint (.pt) from
                `encoder/scripts/slim_checkpoint.py`.
            model_config_path: same YAML the trainer used, so the
                architecture matches what produced the checkpoint.
            schema_path: parent's `data/schema.yaml`.
            histories_path: parent's `data/synthetic/token_ids.npy`.
            cast_path: `encoder/data/demo_cast.json`.
            device: torch.device. CPU is the demo default; the
                "this runs on a laptop" story is the pitch.

        Returns:
            CopilotModel ready for `predict` and `stream_reasoning`.
        """
        # --- schema, histories, cast ---
        schema = load_schema(schema_path)
        histories = np.load(histories_path, mmap_mode="r")
        cast = _load_cast(cast_path)

        # --- model: build fresh, then overlay the slim checkpoint ---
        with model_config_path.open() as f:
            mcfg = yaml.safe_load(f)

        # CPU demo runs in fp32 regardless of what the training config
        # said, because bf16 matmul on CPU is slow and inconsistently
        # supported across PyTorch builds. The slim ckpt's tensors get
        # cast to fp32 at load time.
        dtype = torch.float32 if device.type == "cpu" else torch.bfloat16

        model = build_transaction_multisurface(
            schema=schema,
            model_path=mcfg["backbone"]["hf_path"],
            encoder_cfg=mcfg.get("encoder"),
            projector_cfg=mcfg.get("projector"),
            head_cfg=mcfg.get("heads"),
            lora_cfg=mcfg["backbone"].get("lora"),
            dtype=dtype,
            device_map=None if device.type == "cpu" else "auto",
        )

        # Load slim weights with strict=False so the frozen base weights
        # (absent from the slim ckpt) stay at whatever the fresh HF
        # download produced. Those base weights are identical to what
        # was in the original full checkpoint.
        ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
        if not ckpt.get("model_state_dict_slim"):
            raise ValueError(
                f"Expected a slim checkpoint at {checkpoint_path}. "
                f"Run encoder/scripts/slim_checkpoint.py first.",
            )
        state = {k: v.to(dtype) if v.is_floating_point() else v
                 for k, v in ckpt["model_state_dict"].items()}
        missing, unexpected = model.load_state_dict(state, strict=False)
        # `missing` is expected (the frozen base keys); `unexpected`
        # would indicate a real version skew and should error.
        if unexpected:
            raise RuntimeError(f"Unexpected state_dict keys: {unexpected[:5]} ...")

        model.to(device)
        return cls(model=model, schema=schema, histories=histories,
                   cast=cast, device=device)

    # ----- inference -----

    @torch.inference_mode()
    def predict(self, member: CastMember, top_k: int = 5) -> CopilotResult:
        """Run the probability + attribution heads for one cast member.

        The LM head is NOT run here — call `stream_reasoning` separately
        so the UI can render the score immediately, then stream the
        reasoning underneath.

        Args:
            member: which cast member to predict on.
            top_k: how many attribution positions to surface for the
                timeline glow.

        Returns:
            CopilotResult — score, predicted class, attribution probs.
        """
        torch.manual_seed(DEMO_SEED)
        batch = self._build_batch(member)
        out = self.model.predict(batch)

        # Probability head — argmax for the headline label; softmax of
        # the predicted class for the score number.
        prob_logits = out["probability_logits"][0]  # (3,)
        probs = torch.softmax(prob_logits.float(), dim=-1)  # (3,)
        predicted_class = int(torch.argmax(probs).item())
        score = float(probs[predicted_class].item())

        # Attribution — sigmoid per position, then top-k by raw logit
        # (which is monotonic with sigmoid).
        attr_logits = out["attribution_logits"][0]  # (64,)
        attr_probs = torch.sigmoid(attr_logits.float()).cpu().numpy()
        top_k_positions = (
            torch.topk(attr_logits.float(), k=top_k, dim=-1).indices.cpu().numpy()
        )

        return CopilotResult(
            score=score,
            predicted_class=predicted_class,
            predicted_label=LABEL_NAMES[predicted_class],
            color=LABEL_COLORS[predicted_class],
            attribution_probs=attr_probs,
            top_k_positions=top_k_positions,
            disputed_idx=member.disputed_idx,
        )

    def build_reasoning_text(
        self,
        member: CastMember,
        result: "CopilotResult",
    ) -> str:
        """Render the analyst-facing reasoning deterministically.

        Doctrine reference: liquid-models-architecture §11 / Ottoguard
        principle: "Structural validity is not the model's job; it is
        the decoder's." The model produces (probability, attribution).
        This function consumes those plus the customer's ground-truth
        feature tokens to render a coherent paragraph. The LM head is
        no longer trained for this surface (350M for generation is the
        anti-pattern per liquid-finetuning-playbook §10), and the
        template here mirrors the synthesizer's `_build_reasoning_text`
        contract so the demo output matches what the corpus advertised.

        Args:
            member: cast member being analyzed (for ground-truth signal).
            result: CopilotResult containing predicted_class, score,
                top_k_positions, attribution_probs.

        Returns:
            One paragraph of reasoning text. Includes the score, the
            top contributing positions, and a one-line recommendation.
        """
        history = np.asarray(self.histories[member.customer_idx])
        return _render_reasoning(
            history=history,
            disputed_idx=member.disputed_idx,
            predicted_class=result.predicted_class,
            score=result.score,
            top_k_positions=result.top_k_positions.tolist(),
        )

    def stream_reasoning(
        self,
        member: CastMember,
        result: "CopilotResult",
        chunk_chars: int = 6,
    ) -> Iterator[str]:
        """Stream the templated reasoning chunk-by-chunk for the UI animation.

        The text is deterministic (built once), then yielded incrementally
        so the UI can render the "model thinking" effect. The previous
        implementation called the 350M LM head; that approach produced
        bag-of-words output (350M is below the doctrine's generation
        tier) and has been removed.

        Args:
            member: cast member.
            result: model prediction result.
            chunk_chars: chars per yield. ~6 chars at ~50ms cadence
                produces a smooth word-by-word reveal.

        Yields:
            Cumulative substring at each step; the last yield is the
            full reasoning text.
        """
        text = self.build_reasoning_text(member, result)
        for i in range(chunk_chars, len(text) + chunk_chars, chunk_chars):
            yield text[: min(i, len(text))]

    # ----- internals -----

    def _build_batch(self, member: CastMember) -> MixedModalityBatch:
        """Assemble a batch-of-1 MixedModalityBatch for inference.

        Pulls the customer's 64-tx history out of the in-memory
        histories array, tokenizes the complaint text, and returns
        a batch with no labels (predict-only).
        """
        # (64, 15) int16 -> (1, 64, 15) int64 on device. Copy out of
        # the mmap-backed array so the tensor owns writable memory.
        history = np.asarray(self.histories[member.customer_idx]).copy()
        feature_ids = torch.from_numpy(history).long().unsqueeze(0)
        feature_ids = feature_ids.to(self.device)

        input_ids, attention_mask, lengths = tokenize_texts(
            self.tokenizer, [member.complaint_text], max_length=256,
        )
        input_ids = input_ids.to(self.device)
        attention_mask = attention_mask.to(self.device)
        lengths = lengths.to(self.device)

        # disputed_idx is required by the model contract — the encoder
        # adds its disputed marker at this position so the backbone,
        # attribution head, and LM head know which transaction is being
        # asked about.
        disputed_idx = torch.tensor(
            [member.disputed_idx], dtype=torch.long, device=self.device,
        )

        # Tag the batch "probability" so any forward() dispatch (if used
        # by other code paths) lands on the prob head; predict() bypasses
        # this and runs all heads.
        return MixedModalityBatch(
            feature_ids=feature_ids,
            text_input_ids=input_ids,
            text_attention_mask=attention_mask,
            text_lengths=lengths,
            head_target="probability",
            disputed_idx=disputed_idx,
        )


def _render_reasoning(
    history: np.ndarray,
    disputed_idx: int,
    predicted_class: int,
    score: float,
    top_k_positions: list[int],
) -> str:
    """Render analyst-facing reasoning from model output + ground-truth features.

    The model contributes: predicted_class, score, top_k_positions.
    The history contributes: ground-truth feature tokens at the
    disputed transaction and the top-k attribution positions.

    Why pull ground truth from the history at render time: the model's
    job is the verdict; the explanation should cite specific, verifiable
    facts about the customer's history rather than asking the LM head
    to fabricate prose. This mirrors the Ottoguard principle ("structural
    validity belongs to the decoder, not the model") and is what the
    synthesizer trained the corpus against.
    """
    disp_country_token = int(history[disputed_idx, FEATURE_COUNTRY])
    disp_country = _decode_country_label(disp_country_token)
    disp_amount = _approximate_amount_usd(int(history[disputed_idx, FEATURE_AMOUNT]))
    disp_cmc_raw = int(history[disputed_idx, FEATURE_CUSTOMER_MERCHANT_COUNT])
    disp_cmc = max(0, disp_cmc_raw - RESERVED_OFFSET)
    disp_merchant = int(history[disputed_idx, FEATURE_MERCHANT_ID])
    disp_cvv = int(history[disputed_idx, FEATURE_CVV])
    disp_avs = int(history[disputed_idx, FEATURE_AVS])

    # How many of the top-k attribution positions share the disputed merchant?
    # Used to ground "the model attended to your prior usage of this merchant."
    same_merchant_positions = [
        p for p in top_k_positions
        if int(history[p, FEATURE_MERCHANT_ID]) == disp_merchant and p != disputed_idx
    ]

    # Country-break detection from the customer's history mode.
    countries = history[:, FEATURE_COUNTRY]
    mode_country_token = int(np.bincount(countries).argmax())
    geography_break = disp_country_token != mode_country_token

    # CNP cluster preceding the disputed transaction (card-testing signal).
    window = history[max(0, disputed_idx - 5):disputed_idx, FEATURE_ENTRY_MODE]
    cnp_count = int(np.sum(window == ENTRY_CNP))

    positions_str = ", ".join(str(p) for p in top_k_positions[:5])
    score_str = f"{score:.2f}"

    if predicted_class == LABEL_LIKELY:
        merchant_clause = (
            f"The customer has substantial prior history with this merchant "
            f"(familiarity bucket {disp_cmc}/19)"
            if disp_cmc >= 5
            else "The customer's history shows a recurring pattern at this merchant"
        )
        return (
            f"Score {score_str} — likely friendly fraud. {merchant_clause}, "
            f"and the disputed {disp_amount} charge sits inside the customer's "
            f"established pattern. The model's top contributing transactions are "
            f"positions {positions_str}"
            + (
                f"; {len(same_merchant_positions)} of those are this same merchant."
                if same_merchant_positions
                else "."
            )
            + " Recommend evidence request rather than auto-refund."
        )

    if predicted_class == LABEL_UNLIKELY:
        anomaly_clauses: list[str] = []
        if disp_cmc == 0:
            anomaly_clauses.append("no prior history with this merchant")
        if geography_break:
            anomaly_clauses.append(
                f"transaction occurred in {disp_country}, outside the "
                f"customer's home country"
            )
        if cnp_count >= 2:
            anomaly_clauses.append(
                f"{cnp_count} card-not-present transactions in the 5-tx window "
                f"preceding the disputed charge — consistent with card-testing"
            )
        if disp_cvv != CVV_MATCH:
            anomaly_clauses.append("CVV did not match")
        if not anomaly_clauses:
            anomaly_clauses.append("the disputed transaction is anomalous relative to the customer's pattern")
        anomaly_clause = "; ".join(anomaly_clauses)
        return (
            f"Score {score_str} — unlikely friendly fraud (probable true unauthorized). "
            f"The disputed {disp_amount} charge shows {anomaly_clause}. Top contributing "
            f"transactions are positions {positions_str}. Recommend auto-resolve refund."
        )

    # Ambiguous
    signal_clauses: list[str] = []
    if 1 <= disp_cmc <= 7:
        signal_clauses.append(
            f"some prior interactions with this merchant (familiarity bucket {disp_cmc}/19)"
        )
    if geography_break:
        signal_clauses.append(f"the disputed charge is in {disp_country}, not the customer's home country")
    if disp_cmc == 0 and not geography_break:
        signal_clauses.append("the merchant is unfamiliar but every other signal is clean")
    if not signal_clauses:
        signal_clauses.append("the signals are split between legitimate and friendly-fraud patterns")
    signal_clause = "; ".join(signal_clauses)
    return (
        f"Score {score_str} — ambiguous. The {disp_amount} charge shows mixed signals: "
        f"{signal_clause}. Top contributing transactions are positions {positions_str}. "
        f"Recommend specialist review before resolution."
    )


def _load_cast(cast_path: Path) -> list[CastMember]:
    """Read demo_cast.json into typed CastMember objects.

    Errors clearly if a required field is missing — the cast file is
    hand-curated and a typo there should not silently break the demo.
    """
    with cast_path.open() as f:
        raw = json.load(f)
    cast: list[CastMember] = []
    for entry in raw["cast"]:
        cast.append(
            CastMember(
                pattern=entry["pattern"],
                display_name=entry["display_name"],
                customer_idx=int(entry["customer_idx"]),
                disputed_idx=int(entry["disputed_idx"]),
                expected_label=int(entry["expected_label"]),
                expected_label_name=entry["expected_label_name"],
                complaint_text=entry["complaint_text"],
                tone=entry["tone"],
                description=entry["description"],
            ),
        )
    return cast