fathom-code / viz /app.py

23f2002275

feat(C): demo materials - architecture diagram, demo script, blog draft, viz reward-pie panel, README polish (preflight green)

fb74a9b 4 months ago

Raw

History Blame Contribute Delete

8.29 kB

	"""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.")