File size: 8,290 Bytes

"""FATHOM Streamlit Demo — DEM-03.

3-panel layout:
  Left:   Recursion tree visualization (D3 or plotly sunburst)
  Middle: Pareto frontier (accuracy vs token-cost, α sweep)
  Right:  W&B training curves embed (live iframe)

Run: streamlit run viz/app.py
"""
from __future__ import annotations

import json
import os
from pathlib import Path

import streamlit as st

st.set_page_config(
    page_title="FATHOM — RL-Trained Recursive Language Model",
    page_icon="🧠",
    layout="wide",
)

# ---------------------------------------------------------------------------
# Header
# ---------------------------------------------------------------------------
st.markdown("## 🧠 FATHOM Demo")
st.markdown(
    "_First RL-Trained Recursive Language Model — Meta × PyTorch × HF Hackathon_"
)
st.divider()

# ---------------------------------------------------------------------------
# Reward composition badge (DEM-03 C.4: visible-to-judges reward overview)
# ---------------------------------------------------------------------------
with st.expander("Reward composition (4 components, deterministic verifier)", expanded=True):
    cb1, cb2 = st.columns([1, 2], gap="medium")
    with cb1:
        st.markdown("**Format gate** (multiplier)")
        st.success("`<answer>...</answer>` required \u2014 if missing, reward = 0")
        st.caption("Source: `rewards/format_gate.py` \u2014 audited against attack #2 in REWARD_AUDIT.md")
    with cb2:
        try:
            import plotly.graph_objects as go  # type: ignore
            labels = ["correctness", "token_budget", "recursion_efficiency"]
            weights = [0.75, 0.20, 0.05]
            colors = ["#2ca02c", "#1f77b4", "#ff7f0e"]
            fig0 = go.Figure(go.Pie(
                labels=labels, values=weights, marker=dict(colors=colors),
                hole=0.4, textinfo="label+percent",
            ))
            fig0.update_layout(margin=dict(l=10, r=10, t=10, b=10), height=200, showlegend=False)
            st.plotly_chart(fig0, use_container_width=True)
        except ImportError:
            st.metric("correctness", 0.75)
            st.metric("token_budget", 0.20)
            st.metric("recursion_efficiency", 0.05)

st.divider()

# ---------------------------------------------------------------------------
# Sidebar: controls
# ---------------------------------------------------------------------------
with st.sidebar:
    st.header("Controls")
    alpha_val = st.slider("Token-budget α", 0.05, 1.0, 0.20, 0.05, key="alpha")
    max_depth = st.slider("Max recursion depth shown", 1, 3, 2, key="max_depth")
    env_url = st.text_input(
        "Env server URL",
        value=os.environ.get("FATHOM_SPACE_URL", "http://localhost:8001"),
        key="env_url",
    )
    st.divider()
    st.caption("Source: [github.com/fathom](https://github.com)")

# ---------------------------------------------------------------------------
# 3 columns
# ---------------------------------------------------------------------------
col_tree, col_pareto, col_wb = st.columns([1.2, 1.2, 1], gap="medium")

# ── Column 1: Recursion tree ────────────────────────────────────────────────
with col_tree:
    st.subheader("Recursion Tree")
    st.caption("One episode from the trained model")

    # Placeholder D3 tree — replaced with live env call post-training
    sample_tree = {
        "name": "root: 200K doc",
        "children": [
            {
                "name": "llm(chunk_0-50K)",
                "children": [{"name": "REPL: grep → 'azure'"}],
            },
            {"name": "REPL: count_tokens → 200K"},
            {"name": "→ <answer>azure</answer>"},
        ],
    }
    tree_json = json.dumps(sample_tree)
    d3_html = f"""
    <html>
    <head>
    <script src="https://cdn.jsdelivr.net/npm/d3@7"></script>
    <style>
      body {{ font-family: monospace; font-size: 12px; }}
      .node circle {{ fill: #6366f1; stroke: #312e81; stroke-width: 1.5px; }}
      .node text {{ fill: #1e1b4b; }}
      .link {{ fill: none; stroke: #a5b4fc; stroke-width: 1.5px; }}
    </style>
    </head>
    <body>
    <div id="tree"></div>
    <script>
    const data = {tree_json};
    const width = 340, height = 240;
    const svg = d3.select("#tree").append("svg").attr("width", width).attr("height", height);
    const g = svg.append("g").attr("transform", "translate(40,20)");
    const tree = d3.tree().size([height-40, width-80]);
    const root = d3.hierarchy(data);
    tree(root);
    g.selectAll(".link").data(root.links()).enter().append("path")
      .attr("class","link")
      .attr("d", d3.linkHorizontal().x(d=>d.y).y(d=>d.x));
    const node = g.selectAll(".node").data(root.descendants()).enter()
      .append("g").attr("class","node")
      .attr("transform", d=>`translate(${{d.y}},${{d.x}})`);
    node.append("circle").attr("r", 5);
    node.append("text").attr("dy","0.35em").attr("x", d=>d.children?-8:8)
      .attr("text-anchor", d=>d.children?"end":"start")
      .text(d=>d.data.name.slice(0,28));
    </script>
    </body></html>
    """
    st.components.v1.html(d3_html, height=260)

# ── Column 2: Pareto frontier ────────────────────────────────────────────────
with col_pareto:
    st.subheader("Pareto Frontier")
    st.caption(f"Accuracy vs token cost (α = {alpha_val:.2f})")

    try:
        import plotly.graph_objects as go  # type: ignore

        # Placeholder data — replaced with logged eval results post-training
        alpha_values = [0.05, 0.10, 0.20, 0.50, 1.00]
        accuracy = [0.62, 0.61, 0.58, 0.52, 0.44]
        token_cost = [1.00, 0.95, 0.85, 0.65, 0.48]

        fig = go.Figure()
        fig.add_trace(go.Scatter(
            x=token_cost,
            y=accuracy,
            mode="lines+markers",
            marker=dict(size=10, color="#6366f1"),
            line=dict(color="#a5b4fc", width=2),
            text=[f"α={a}" for a in alpha_values],
            textposition="top center",
        ))
        # Highlight current alpha
        idx = min(range(len(alpha_values)), key=lambda i: abs(alpha_values[i] - alpha_val))
        fig.add_trace(go.Scatter(
            x=[token_cost[idx]],
            y=[accuracy[idx]],
            mode="markers",
            marker=dict(size=16, color="#ef4444", symbol="star"),
            name=f"Current α={alpha_val:.2f}",
        ))
        fig.update_layout(
            xaxis_title="Normalized Token Cost",
            yaxis_title="Accuracy",
            margin=dict(l=20, r=10, t=20, b=40),
            height=240,
            showlegend=False,
        )
        st.plotly_chart(fig, use_container_width=True)
    except ImportError:
        st.info("plotly not installed — run `pip install plotly`")

# ── Column 3: W&B training curves ────────────────────────────────────────────
with col_wb:
    st.subheader("Training Curves")
    wb_url = os.environ.get("WANDB_RUN_URL", "")
    if wb_url:
        st.components.v1.iframe(wb_url, height=240, scrolling=True)
    else:
        st.info("Set `WANDB_RUN_URL` env var to embed live training curves.")
        st.caption("W&B logged metrics: composite, format_pass, correctness, token_budget, recursion_eff")
        # Placeholder sparkline
        try:
            import plotly.graph_objects as go  # type: ignore
            steps = list(range(0, 401, 50))
            fake_reward = [0.10, 0.18, 0.28, 0.38, 0.45, 0.52, 0.57, 0.60, 0.62]
            fig2 = go.Figure(go.Scatter(x=steps, y=fake_reward, mode="lines+markers",
                                        line=dict(color="#6366f1", width=2)))
            fig2.update_layout(
                xaxis_title="GRPO step",
                yaxis_title="Composite reward",
                margin=dict(l=20, r=10, t=10, b=40),
                height=240,
            )
            st.plotly_chart(fig2, use_container_width=True)
        except ImportError:
            pass

st.divider()
st.caption("FATHOM Phase 1 skeleton — full curves appear after GRPO training completes.")