encoder/src/demo/copilot_app_collections.py

"""Gradio app for the Collections Co-Pilot demo.

Mirrors the Dispute Co-Pilot layout (encoder/src/demo/copilot_app.py)
adapted to the Collections surface:
  - Cast strip + context quote + Analyze button.
  - Timeline of 64 dots, with the "now" position starred and the
    attribution glow on top-k contributing transactions.
  - Treatment scoreboard: 4 horizontal bars (one per option) showing
    P(likely_respond) and the model's recommendation pill.
  - Reasoning panel streams the analyst-facing rationale.

CLI:
    python -m encoder.src.demo.copilot_app_collections \\
        --checkpoint encoder/experiments/collections_v1/demo_checkpoint.pt \\
        --model-config encoder/configs/model_collections.yaml \\
        --schema data/schema.yaml \\
        --histories data/synthetic/token_ids.npy \\
        --cast encoder/data/collections_cast.json \\
        --port 7863
"""

from __future__ import annotations

import argparse
from pathlib import Path

import gradio as gr
import torch

from encoder.src.demo.copilot_inference_collections import (
    CollectionsCastMember,
    CollectionsCopilotModel,
)
from encoder.src.demo.copilot_render_collections import (
    render_cast_strip,
    render_context,
    render_header,
    render_reasoning,
    render_timeline,
    render_treatment_grid,
)


# Same container tuning as the dispute app.
_CONTAINER_WIDTH_PX = 1100


def _build_tab(model: CollectionsCopilotModel) -> None:
    """Build the Collections surface into the current Gradio context."""
    cast = model.cast

    # --- state ---
    selected_idx = gr.State(value=0)

    # --- cast strip + selection buttons ---
    cast_html = gr.HTML(render_cast_strip(cast, 0))

    with gr.Row():
        cast_buttons: list[gr.Button] = []
        for i, m in enumerate(cast):
            short = " ".join(m.display_name.split(" ")[:2])
            btn = gr.Button(
                value=short,
                variant="primary" if i == 0 else "secondary",
                scale=1,
            )
            cast_buttons.append(btn)

    # --- context quote ---
    context_html = gr.HTML(render_context(cast[0]))

    # --- analyze button ---
    with gr.Row():
        gr.HTML("<div style='flex: 1'></div>")
        analyze_btn = gr.Button(
            value="Analyze",
            variant="primary",
            size="lg",
            scale=0,
            min_width=180,
        )
        gr.HTML("<div style='flex: 1'></div>")

    # --- timeline ---
    timeline_html = gr.HTML(
        render_timeline(context_idx=cast[0].context_idx),
    )

    # --- scoreboard + reasoning side-by-side ---
    with gr.Row(equal_height=True):
        with gr.Column(scale=3):
            scoreboard_html = gr.HTML(render_treatment_grid(None))
        with gr.Column(scale=2):
            reasoning_html = gr.HTML(render_reasoning(None))

    # --- event wiring ---

    def _select(idx: int) -> tuple:
        member = cast[idx]
        button_updates = tuple(
            gr.update(variant="primary" if i == idx else "secondary")
            for i in range(len(cast))
        )
        return (
            idx,
            render_cast_strip(cast, idx),
            render_context(member),
            render_timeline(context_idx=member.context_idx),
            render_treatment_grid(None),
            render_reasoning(None),
        ) + button_updates

    for i, btn in enumerate(cast_buttons):
        btn.click(
            fn=lambda i=i: _select(i),
            inputs=None,
            outputs=[
                selected_idx, cast_html, context_html,
                timeline_html, scoreboard_html, reasoning_html,
                *cast_buttons,
            ],
        )

    def _analyze(idx: int):
        member: CollectionsCastMember = cast[idx]
        result = model.predict(member, top_k=5)
        timeline = render_timeline(
            context_idx=member.context_idx,
            top_k_positions=result.top_k_positions,
            attribution_probs=result.attribution_probs,
        )
        scoreboard = render_treatment_grid(result)
        yield timeline, scoreboard, render_reasoning("")
        for partial in model.stream_reasoning(member, result, chunk_chars=6):
            yield timeline, scoreboard, render_reasoning(partial)

    analyze_btn.click(
        fn=_analyze,
        inputs=[selected_idx],
        outputs=[timeline_html, scoreboard_html, reasoning_html],
    )


def _build_ui(model: CollectionsCopilotModel) -> gr.Blocks:
    """Standalone Blocks UI: Collections-specific header + tab content."""
    with gr.Blocks(title="Collections Co-Pilot — Liquid AI") as demo:
        gr.HTML(render_header())
        _build_tab(model)
    return demo


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Collections Co-Pilot Gradio demo",
    )
    parser.add_argument(
        "--checkpoint",
        type=Path,
        default=Path("encoder/experiments/collections_v1/demo_checkpoint.pt"),
    )
    parser.add_argument(
        "--model-config",
        type=Path,
        default=Path("encoder/configs/model_collections.yaml"),
    )
    parser.add_argument("--schema", type=Path, default=Path("data/schema.yaml"))
    parser.add_argument(
        "--histories",
        type=Path,
        default=Path("data/synthetic/token_ids.npy"),
    )
    parser.add_argument(
        "--cast",
        type=Path,
        default=Path("encoder/data/collections_cast.json"),
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cpu",
        choices=["cpu", "cuda", "mps"],
    )
    parser.add_argument("--port", type=int, default=7863)
    parser.add_argument("--share", action="store_true")
    args = parser.parse_args()

    device = torch.device(args.device)
    print(f"Loading CollectionsCopilotModel on {device} ...")
    model = CollectionsCopilotModel.from_paths(
        checkpoint_path=args.checkpoint,
        model_config_path=args.model_config,
        schema_path=args.schema,
        histories_path=args.histories,
        cast_path=args.cast,
        device=device,
    )
    print(f"  cast size: {len(model.cast)}")
    print(
        f"  histories: {model.histories.shape} dtype={model.histories.dtype}"
    )

    demo = _build_ui(model)
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=args.port,
        share=args.share,
        theme=gr.themes.Default(
            font=["Inter", "system-ui", "sans-serif"],
            font_mono=["JetBrains Mono", "ui-monospace", "monospace"],
        ),
        css=f"""
        .gradio-container {{
            max-width: {_CONTAINER_WIDTH_PX}px !important;
            background: #fafafa !important;
        }}
        """,
    )


if __name__ == "__main__":
    main()