encoder/src/demo/copilot_app_unified.py

"""Unified Co-Pilot: 6 tabs over one Gradio Blocks.

Top-level layout:
    Tabs:
      - "Multi-Head Demo" — original encoder demo (4 task heads, V3 nocompress
                            + meanpool checkpoint, curated customer dropdown)
      - "Why Liquid" — architectural pitch (original demo content)
      - "Integration" — build-it-yourself guide (original demo content)
      - "Dispute Co-Pilot" — friendly-fraud classifier + attribution
      - "Collections Co-Pilot" — treatment-response scoreboard
      - "Fraud Co-Pilot" — pattern stage + type classifier

Each tab's content is a composable `_build_..._tab_contents(...)` helper
exported from the per-surface app module. Models are loaded once at
startup (4 model instances total — V3 multi-head + 3 surface-specific
multi-surface).

CLI:
    python -m encoder.src.demo.copilot_app_unified \\
        --multihead-checkpoint encoder/experiments/.../step_004999_slim.pt \\
        --multihead-config encoder/configs/model_nocompress.yaml \\
        --multihead-data-dir data/synthetic \\
        --dispute-checkpoint encoder/experiments/dispute_legitimacy_v7/demo_checkpoint.pt \\
        --collections-checkpoint encoder/experiments/collections_v3/demo_checkpoint.pt \\
        --fraud-checkpoint encoder/experiments/fraud_pattern_v1/demo_checkpoint.pt \\
        --cast-histories data/synthetic/cast_token_ids.npy \\
        --port 7860
"""

from __future__ import annotations

import argparse
from pathlib import Path

import gradio as gr
import torch

# Original multi-head demo (V3) imports
from src.data.schema import load_schema
from src.demo.app import DemoData
from src.demo.decode import TransactionDecoder
from src.demo.merchant_catalog import DemoMerchantCatalog

from encoder.src.demo.app import (
    _build_cold_start_tab_contents,
    _build_demo_tab_contents,
    _build_integration_tab_contents,
    _build_why_liquid_tab_contents,
    _build_theme,
    _CSS,
    _HEADER_HTML,
)
from encoder.src.demo.copilot_app import _build_tab as _build_dispute_tab
from encoder.src.demo.copilot_app_collections import (
    _build_tab as _build_collections_tab,
)
from encoder.src.demo.copilot_app_fraud_pattern import (
    _build_tab as _build_fraud_tab,
)
from encoder.src.demo.copilot_inference import CopilotModel
from encoder.src.demo.copilot_inference_collections import (
    CollectionsCopilotModel,
)
from encoder.src.demo.copilot_inference_fraud_pattern import (
    FraudPatternCopilotModel,
)
from encoder.src.demo.inference import EncoderDemoModel


def build_unified_ui(
    multihead_model: EncoderDemoModel,
    multihead_data: DemoData,
    multihead_decoder: TransactionDecoder,
    multihead_merchant_catalog: DemoMerchantCatalog,
    dispute_model: CopilotModel,
    collections_model: CollectionsCopilotModel,
    fraud_model: FraudPatternCopilotModel,
) -> gr.Blocks:
    """Compose the original 3-tab demo + 3 new Co-Pilot tabs into one Blocks.

    Uses the original demo's theme + CSS so the visual vocabulary stays
    consistent across tabs.
    """
    # Gradio 6.0 moved `css` and `theme` from Blocks() to launch(); the
    # caller passes them in via demo.queue().launch(theme=..., css=...).
    # We stash them on the returned Blocks instance for the caller.
    with gr.Blocks(title="Transaction Encoder — Liquid AI") as app:
        gr.HTML(_HEADER_HTML)
        with gr.Tabs():
            with gr.Tab("Multi-Head Demo"):
                _build_demo_tab_contents(
                    multihead_model,
                    multihead_data,
                    multihead_decoder,
                    multihead_merchant_catalog,
                    app,
                )
            with gr.Tab("Cold Start"):
                _build_cold_start_tab_contents()
            with gr.Tab("Why Liquid"):
                _build_why_liquid_tab_contents()
            with gr.Tab("Integration"):
                _build_integration_tab_contents()
            with gr.Tab("Dispute Co-Pilot"):
                _build_dispute_tab(dispute_model)
            with gr.Tab("Collections Co-Pilot"):
                _build_collections_tab(collections_model)
            with gr.Tab("Fraud Co-Pilot"):
                _build_fraud_tab(fraud_model)
    return app


def _load_multihead(
    checkpoint: Path,
    config: Path,
    schema: Path,
    data_dir: Path,
    dtype: torch.dtype,
    device: str,
) -> tuple[EncoderDemoModel, DemoData, TransactionDecoder, DemoMerchantCatalog]:
    """Load the original multi-head V3 model + curated test data."""
    print("[multihead] schema + data ...")
    schema_cfg = load_schema(schema)
    data = DemoData(data_dir, schema_cfg)
    merchant_catalog = DemoMerchantCatalog(schema_cfg)
    decoder = TransactionDecoder(schema_cfg, merchant_catalog)
    print("[multihead] loading EncoderDemoModel ...")
    model = EncoderDemoModel(
        model_config_path=config,
        schema_path=schema,
        checkpoint_path=checkpoint,
        dtype=dtype,
        device=device,
    )
    print(f"[multihead] checkpoint: {model.checkpoint_status}")
    return model, data, decoder, merchant_catalog


def _load_copilots(
    dispute_checkpoint: Path,
    dispute_config: Path,
    collections_checkpoint: Path,
    collections_config: Path,
    fraud_checkpoint: Path,
    fraud_config: Path,
    schema: Path,
    cast_histories: Path,
    dispute_cast: Path,
    collections_cast: Path,
    fraud_cast: Path,
    device: torch.device,
) -> tuple[CopilotModel, CollectionsCopilotModel, FraudPatternCopilotModel]:
    """Load the three Co-Pilot surfaces. Each has its own backbone copy."""
    print("[copilot 1/3] loading Dispute ...")
    dispute_model = CopilotModel.from_paths(
        checkpoint_path=dispute_checkpoint,
        model_config_path=dispute_config,
        schema_path=schema,
        histories_path=cast_histories,
        cast_path=dispute_cast,
        device=device,
    )
    print("[copilot 2/3] loading Collections ...")
    collections_model = CollectionsCopilotModel.from_paths(
        checkpoint_path=collections_checkpoint,
        model_config_path=collections_config,
        schema_path=schema,
        histories_path=cast_histories,
        cast_path=collections_cast,
        device=device,
    )
    print("[copilot 3/3] loading Fraud ...")
    fraud_model = FraudPatternCopilotModel.from_paths(
        checkpoint_path=fraud_checkpoint,
        model_config_path=fraud_config,
        schema_path=schema,
        histories_path=cast_histories,
        cast_path=fraud_cast,
        device=device,
    )
    return dispute_model, collections_model, fraud_model


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Unified Transaction Encoder Gradio app (6 tabs)",
    )
    # --- multi-head V3 ---
    parser.add_argument(
        "--multihead-checkpoint",
        type=Path,
        default=Path("encoder/experiments/nocompress_meanpool/"
                     "encoder_sft_20260519_144916/checkpoints/step_004999_slim.pt"),
    )
    parser.add_argument(
        "--multihead-config",
        type=Path,
        default=Path("encoder/configs/model_nocompress.yaml"),
    )
    parser.add_argument(
        "--multihead-data-dir",
        type=Path,
        default=Path("data/synthetic"),
    )
    # --- dispute ---
    parser.add_argument(
        "--dispute-checkpoint",
        type=Path,
        default=Path("encoder/experiments/dispute_legitimacy_v7/demo_checkpoint.pt"),
    )
    parser.add_argument(
        "--dispute-config",
        type=Path,
        default=Path("encoder/configs/model_dispute_legitimacy.yaml"),
    )
    parser.add_argument(
        "--dispute-cast",
        type=Path,
        default=Path("encoder/data/demo_cast.json"),
    )
    # --- collections ---
    parser.add_argument(
        "--collections-checkpoint",
        type=Path,
        default=Path("encoder/experiments/collections_v3/demo_checkpoint.pt"),
    )
    parser.add_argument(
        "--collections-config",
        type=Path,
        default=Path("encoder/configs/model_collections.yaml"),
    )
    parser.add_argument(
        "--collections-cast",
        type=Path,
        default=Path("encoder/data/collections_cast.json"),
    )
    # --- fraud ---
    parser.add_argument(
        "--fraud-checkpoint",
        type=Path,
        default=Path("encoder/experiments/fraud_pattern_v1/demo_checkpoint.pt"),
    )
    parser.add_argument(
        "--fraud-config",
        type=Path,
        default=Path("encoder/configs/model_fraud_pattern.yaml"),
    )
    parser.add_argument(
        "--fraud-cast",
        type=Path,
        default=Path("encoder/data/fraud_pattern_cast.json"),
    )
    # --- shared ---
    parser.add_argument(
        "--schema",
        type=Path,
        default=Path("data/schema.yaml"),
    )
    parser.add_argument(
        "--cast-histories",
        type=Path,
        default=Path("data/synthetic/token_ids.npy"),
        help="Histories file for the Co-Pilot tabs (subset of 18 cast customers).",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cpu",
        choices=["cpu", "cuda", "mps"],
    )
    parser.add_argument(
        "--dtype",
        type=str,
        default="float32",
        choices=["float32", "bfloat16"],
    )
    parser.add_argument("--port", type=int, default=7860)
    parser.add_argument("--share", action="store_true")
    args = parser.parse_args()

    device = torch.device(args.device)
    multihead_dtype = (
        torch.float32 if args.dtype == "float32" else torch.bfloat16
    )

    multihead_model, multihead_data, multihead_decoder, multihead_merchant = (
        _load_multihead(
            checkpoint=args.multihead_checkpoint,
            config=args.multihead_config,
            schema=args.schema,
            data_dir=args.multihead_data_dir,
            dtype=multihead_dtype,
            device=args.device,
        )
    )

    dispute_model, collections_model, fraud_model = _load_copilots(
        dispute_checkpoint=args.dispute_checkpoint,
        dispute_config=args.dispute_config,
        collections_checkpoint=args.collections_checkpoint,
        collections_config=args.collections_config,
        fraud_checkpoint=args.fraud_checkpoint,
        fraud_config=args.fraud_config,
        schema=args.schema,
        cast_histories=args.cast_histories,
        dispute_cast=args.dispute_cast,
        collections_cast=args.collections_cast,
        fraud_cast=args.fraud_cast,
        device=device,
    )

    print("all four models loaded.")

    demo = build_unified_ui(
        multihead_model=multihead_model,
        multihead_data=multihead_data,
        multihead_decoder=multihead_decoder,
        multihead_merchant_catalog=multihead_merchant,
        dispute_model=dispute_model,
        collections_model=collections_model,
        fraud_model=fraud_model,
    )
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=args.port,
        share=args.share,
        theme=_build_theme(),
        css=_CSS,
    )


if __name__ == "__main__":
    main()