fang8/Multi-Agent_Financial_QA_Optimization_Engine
0
1"""2Multi-Agent Financial QA Optimization Engine3Streamlit Dashboard — refactored for production clarity4"""5import json6import numpy as np7import streamlit as st8import plotly.graph_objects as go9import plotly.express as px10import pandas as pd11 12REPORT_PATH = "./experiments/report.json"13 14st.set_page_config(15 page_title="Multi-Agent Financial QA Optimization Engine",16 layout="wide",17 initial_sidebar_state="collapsed",18)19 20# ── Minimal shared style ──────────────────────────────────────────────────────21st.markdown("""22<style>23 .metric-card {24 background: #F8FAFC; border: 1px solid #E2E8F0;25 border-radius: 8px; padding: 14px 18px; margin-bottom: 0;26 }27 .metric-label { font-size: 11px; color: #64748B; text-transform: uppercase;28 letter-spacing: .06em; margin-bottom: 4px; }29 .metric-value { font-size: 26px; font-weight: 700; color: #0F172A; line-height: 1.1; }30 .metric-delta { font-size: 12px; margin-top: 4px; }31 .delta-pos { color: #16A34A; } .delta-neg { color: #DC2626; }32 .section-caption { color: #64748B; font-size: 13px; margin-top: -10px; margin-bottom: 12px; }33</style>34""", unsafe_allow_html=True)35 36# ── Data ──────────────────────────────────────────────────────────────────────37@st.cache_data38def load_report():39 try:40 with open(REPORT_PATH) as f:41 return json.load(f)42 except FileNotFoundError:43 # ── MOCK DATA (swap for real report.json) ────────────────────────────44 return {45 "meta": {46 "dataset": "TAT-QA",47 "total_n": 400,48 "n_per_type": 100,49 "models": ["openai-gpt-oss-20b", "openai-gpt-oss-120b"],50 },51 "runs": [52 # 120B53 {"model": "openai-gpt-oss-120b", "strategy": "zero_shot",54 "overall_f1": 0.740, "overall_em": 0.550, "weighted_f1": 0.692,55 "total_cost_usd": 0.087, "avg_latency_s": 2.69, "cost_per_correct": 0.0039,56 "by_type": {"arithmetic": {"f1": 0.500}, "span": {"f1": 0.811},57 "multi-span": {"f1": 0.649}, "count": {"f1": 1.000}},58 "fail_by_type": {"arithmetic": 50, "span": 40, "multi-span": 90, "count": 0}},59 {"model": "openai-gpt-oss-120b", "strategy": "chain_of_thought",60 "overall_f1": 0.661, "overall_em": 0.425, "weighted_f1": 0.615,61 "total_cost_usd": 0.100, "avg_latency_s": 4.72, "cost_per_correct": 0.0058,62 "by_type": {"arithmetic": {"f1": 0.433}, "span": {"f1": 0.689},63 "multi-span": {"f1": 0.520}, "count": {"f1": 1.000}},64 "fail_by_type": {"arithmetic": 70, "span": 70, "multi-span": 90, "count": 0}},65 {"model": "openai-gpt-oss-120b", "strategy": "router",66 "overall_f1": 0.716, "overall_em": 0.475, "weighted_f1": 0.660,67 "total_cost_usd": 0.130, "avg_latency_s": 5.10, "cost_per_correct": 0.0066,68 "by_type": {"arithmetic": {"f1": 0.433}, "span": {"f1": 0.768},69 "multi-span": {"f1": 0.664}, "count": {"f1": 1.000}},70 "fail_by_type": {"arithmetic": 70, "span": 50, "multi-span": 90, "count": 0}},71 # 20B72 {"model": "openai-gpt-oss-20b", "strategy": "zero_shot",73 "overall_f1": 0.680, "overall_em": 0.490, "weighted_f1": 0.630,74 "total_cost_usd": 0.041, "avg_latency_s": 1.82, "cost_per_correct": 0.0022,75 "by_type": {"arithmetic": {"f1": 0.440}, "span": {"f1": 0.760},76 "multi-span": {"f1": 0.580}, "count": {"f1": 0.940}},77 "fail_by_type": {"arithmetic": 60, "span": 45, "multi-span": 92, "count": 6}},78 {"model": "openai-gpt-oss-20b", "strategy": "chain_of_thought",79 "overall_f1": 0.610, "overall_em": 0.380, "weighted_f1": 0.560,80 "total_cost_usd": 0.052, "avg_latency_s": 3.50, "cost_per_correct": 0.0031,81 "by_type": {"arithmetic": {"f1": 0.380}, "span": {"f1": 0.640},82 "multi-span": {"f1": 0.460}, "count": {"f1": 0.960}},83 "fail_by_type": {"arithmetic": 75, "span": 72, "multi-span": 93, "count": 4}},84 {"model": "openai-gpt-oss-20b", "strategy": "router",85 "overall_f1": 0.650, "overall_em": 0.420, "weighted_f1": 0.600,86 "total_cost_usd": 0.068, "avg_latency_s": 3.90, "cost_per_correct": 0.0038,87 "by_type": {"arithmetic": {"f1": 0.400}, "span": {"f1": 0.710},88 "multi-span": {"f1": 0.610}, "count": {"f1": 0.950}},89 "fail_by_type": {"arithmetic": 72, "span": 52, "multi-span": 91, "count": 5}},90 ],91 "failure_patterns": [92 {"name": "Scale/Unit Contamination", "affected": "arithmetic",93 "frequency": "high",94 "description": "Model appends 'million' or '%' to numeric answers copied from table headers",95 "fix": "Post-process: strip units from numeric predictions"},96 {"name": "Span Truncation", "affected": "span", "frequency": "medium",97 "description": "Model paraphrases instead of exact-quoting source text",98 "fix": "Prompt: instruct model to copy verbatim from source"},99 {"name": "Multi-span Format Mismatch", "affected": "multi-span",100 "frequency": "high",101 "description": "Correct values but wrong order, case, or delimiter → EM=0",102 "fix": "Normalize: sort spans, lowercase, strip punctuation before eval"},103 {"name": "Count Unit Leakage", "affected": "count", "frequency": "low",104 "description": "Model adds unit to count answers ('2 years' vs '2')",105 "fix": "Post-process: extract first numeric token from count predictions"},106 ],107 }108 109report = load_report()110runs = report["runs"]111meta = report["meta"]112df = pd.DataFrame(runs)113 114COLORS = {115 "zero_shot": "#2563EB",116 "chain_of_thought": "#DC2626",117 "router": "#16A34A",118}119LABELS = {120 "zero_shot": "Zero-shot",121 "chain_of_thought": "Chain-of-Thought",122 "router": "Router Agent",123}124TYPES = ["arithmetic", "span", "multi-span", "count"]125 126# ── Router Decision Distribution — from real log files ───────────────────────127# Aggregated from router_openai-gpt-oss-120b_log.json and router_openai-gpt-oss-20b_log.json128# Format: per answer_type, count of questions routed to each sub-strategy129_LOG_120B = [130 ("arithmetic", "chain_of_thought", 8), ("arithmetic", "zero_shot", 2),131 ("span", "zero_shot", 5), ("span", "chain_of_thought", 5),132 ("multi-span", "zero_shot", 9), ("multi-span", "chain_of_thought", 1),133 ("count", "zero_shot", 1), ("count", "chain_of_thought", 9),134]135_LOG_20B = [136 ("arithmetic", "chain_of_thought", 8), ("arithmetic", "zero_shot", 2),137 ("span", "zero_shot", 7), ("span", "chain_of_thought", 3),138 ("multi-span", "zero_shot", 10), ("multi-span", "chain_of_thought", 0),139 ("count", "zero_shot", 2), ("count", "chain_of_thought", 8),140]141 142def _build_router_df(log):143 rows = {}144 for atype, strategy, cnt in log:145 if atype not in rows:146 rows[atype] = {"Zero-shot": 0, "Chain-of-Thought": 0}147 key = "Zero-shot" if strategy == "zero_shot" else "Chain-of-Thought"148 rows[atype][key] += cnt149 return pd.DataFrame(rows).T.rename_axis("Answer Type").reset_index()150 151ROUTER_DFS = {152 "openai-gpt-oss-120b": _build_router_df(_LOG_120B),153 "openai-gpt-oss-20b": _build_router_df(_LOG_20B),154}155 156# ── Title ─────────────────────────────────────────────────────────────────────157st.title("Multi-Agent Financial QA Optimization Engine")158st.caption(159 f"Dataset: {meta['dataset']} · "160 f"Sample: {meta['total_n']} questions ({meta['n_per_type']} per type) · "161 f"Models: {', '.join(meta['models'])}"162)163st.divider()164 165# ══════════════════════════════════════════════════════════════════════════════166# Section 0 — Performance Summary167# ══════════════════════════════════════════════════════════════════════════════168model_tab_label = st.radio(169 "Model scope", ["120B Model", "20B Model"], horizontal=True, label_visibility="collapsed"170)171sel_model = "openai-gpt-oss-120b" if "120B" in model_tab_label else "openai-gpt-oss-20b"172df_m = df[df["model"] == sel_model].copy()173 174best = df_m.loc[df_m["overall_f1"].idxmax()]175cot = df_m[df_m["strategy"] == "chain_of_thought"].iloc[0]176delta_f1 = best["overall_f1"] - cot["overall_f1"]177delta_lat = (best["avg_latency_s"] - cot["avg_latency_s"]) / cot["avg_latency_s"] * 100178delta_cost = (best["total_cost_usd"] - cot["total_cost_usd"]) / cot["total_cost_usd"] * 100179 180st.subheader("Performance Summary")181 182def card(label, value, delta_html=""):183 return f"""184 <div class="metric-card">185 <div class="metric-label">{label}</div>186 <div class="metric-value">{value}</div>187 {'<div class="metric-delta">' + delta_html + '</div>' if delta_html else ''}188 </div>"""189 190def delta_html(val, fmt, invert=False):191 pos = val > 0192 if invert: pos = not pos193 cls = "delta-pos" if pos else "delta-neg"194 sign = "+" if val > 0 else ""195 return f'<span class="{cls}">{sign}{fmt.format(val)} vs CoT</span>'196 197c = st.columns(5)198c[0].markdown(card("Best Strategy", LABELS[best["strategy"]]), unsafe_allow_html=True)199c[1].markdown(card("F1 Score", f"{best['overall_f1']:.3f}",200 delta_html(delta_f1, "{:.3f}")), unsafe_allow_html=True)201c[2].markdown(card("Avg Latency", f"{best['avg_latency_s']:.2f}s",202 delta_html(delta_lat, "{:.1f}%", invert=True)), unsafe_allow_html=True)203c[3].markdown(card("Cost / Correct", f"${best['cost_per_correct']:.4f}",204 delta_html(delta_cost, "{:.1f}%", invert=True)), unsafe_allow_html=True)205c[4].markdown(card("Risk-Weighted F1", f"{best['weighted_f1']:.3f}",206 '<span style="color:#64748B;font-size:11px">arith ×2 penalty</span>'),207 unsafe_allow_html=True)208 209st.divider()210 211# ══════════════════════════════════════════════════════════════════════════════212# Section 1 — Strategy Comparison + Pareto Frontier213# ══════════════════════════════════════════════════════════════════════════════214st.subheader("① Strategy Comparison & Pareto Frontier")215st.markdown('<p class="section-caption">Pareto-optimal frontier: maximize F1 while minimizing latency and cost.</p>',216 unsafe_allow_html=True)217 218col1, col2 = st.columns(2)219 220# 1a F1 & EM grouped bar221with col1:222 fig = go.Figure()223 for metric, opacity, suffix in [("overall_f1", 1.0, "F1"), ("overall_em", 0.35, "EM")]:224 fig.add_trace(go.Bar(225 name=suffix,226 x=[LABELS[s] for s in df_m["strategy"]],227 y=df_m[metric],228 marker_color=[COLORS[s] for s in df_m["strategy"]],229 opacity=opacity,230 text=[f"{v:.3f}" for v in df_m[metric]],231 textposition="outside",232 ))233 fig.update_layout(234 title=f"F1 & EM by Strategy ({sel_model.split('-')[-1].upper()})",235 barmode="group", yaxis_range=[0, 0.95], height=360,236 yaxis_title="Score", legend=dict(orientation="h", y=1.12),237 plot_bgcolor="white", paper_bgcolor="white",238 )239 st.plotly_chart(fig, use_container_width=True)240 241# 1b Pareto frontier242with col2:243 fig2 = go.Figure()244 245 # Sort by latency to draw frontier line246 df_sorted = df_m.sort_values("avg_latency_s")247 # Pareto frontier: keep points where no other point has lower latency AND lower cost248 pareto_pts = []249 min_cost = float("inf")250 for _, row in df_sorted.iterrows():251 if row["total_cost_usd"] < min_cost:252 min_cost = row["total_cost_usd"]253 pareto_pts.append(row)254 df_pareto = pd.DataFrame(pareto_pts)255 256 # Background scatter: all strategies (both models dimmed)257 for _, row in df.iterrows():258 is_sel = row["model"] == sel_model259 fig2.add_trace(go.Scatter(260 x=[row["avg_latency_s"]], y=[row["total_cost_usd"]],261 mode="markers",262 marker=dict(263 size=row["overall_f1"] * 55,264 color=COLORS[row["strategy"]],265 opacity=0.85 if is_sel else 0.20,266 line=dict(width=1.5 if is_sel else 0, color="white"),267 ),268 name=f"{LABELS[row['strategy']]} ({row['model'].split('-')[-1].upper()})",269 text=f"{LABELS[row['strategy']]}<br>F1={row['overall_f1']:.3f}<br>"270 f"${row['total_cost_usd']:.3f} | {row['avg_latency_s']:.2f}s",271 hoverinfo="text",272 showlegend=is_sel,273 ))274 275 # Pareto frontier line276 if len(df_pareto) > 1:277 fig2.add_trace(go.Scatter(278 x=df_pareto["avg_latency_s"], y=df_pareto["total_cost_usd"],279 mode="lines",280 line=dict(color="#F59E0B", width=2, dash="dot"),281 name="Pareto Frontier",282 ))283 284 fig2.update_layout(285 title="Pareto Frontier: Latency × Cost (bubble = F1)",286 xaxis_title="Avg Latency (s)", yaxis_title="Total Cost (USD)",287 height=360, legend=dict(orientation="h", y=-0.25, font_size=11),288 plot_bgcolor="white", paper_bgcolor="white",289 )290 st.plotly_chart(fig2, use_container_width=True)291 292st.divider()293 294# ══════════════════════════════════════════════════════════════════════════════295# Section 2 — Per-type Breakdown296# ══════════════════════════════════════════════════════════════════════════════297st.subheader("② Per-type Performance Breakdown")298st.markdown('<p class="section-caption">Arithmetic hardest (50% fail). Multi-span dominated by format mismatch (90% fail rate).</p>',299 unsafe_allow_html=True)300 301col3, col4 = st.columns(2)302 303with col3:304 strats = df_m["strategy"].tolist()305 matrix = []306 for s in strats:307 row_data = df_m[df_m["strategy"] == s].iloc[0]["by_type"]308 matrix.append([row_data.get(t, {}).get("f1", 0) for t in TYPES])309 310 fig3 = go.Figure(go.Heatmap(311 z=matrix, x=TYPES,312 y=[LABELS[s] for s in strats],313 colorscale="Blues",314 text=[[f"{v:.3f}" for v in row] for row in matrix],315 texttemplate="%{text}", showscale=True, zmin=0, zmax=1,316 ))317 fig3.update_layout(title=f"F1 Heatmap by Answer Type ({sel_model.split('-')[-1].upper()})",318 height=280, plot_bgcolor="white", paper_bgcolor="white")319 st.plotly_chart(fig3, use_container_width=True)320 321with col4:322 fail_data = []323 for s in strats:324 fail_by_type = df_m[df_m["strategy"] == s].iloc[0]["fail_by_type"]325 n = meta["n_per_type"]326 for t in TYPES:327 fail_data.append({328 "Strategy": LABELS[s], "Type": t,329 "Failure Rate": fail_by_type.get(t, 0) / n,330 })331 df_fail = pd.DataFrame(fail_data)332 fig4 = px.bar(333 df_fail, x="Type", y="Failure Rate", color="Strategy", barmode="group",334 color_discrete_map={LABELS[s]: COLORS[s] for s in COLORS},335 text_auto=".0%",336 title=f"Failure Rate by Answer Type ({sel_model.split('-')[-1].upper()})",337 )338 fig4.update_layout(height=280, yaxis_range=[0, 1.15],339 plot_bgcolor="white", paper_bgcolor="white",340 legend=dict(orientation="h", y=1.15))341 st.plotly_chart(fig4, use_container_width=True)342 343st.divider()344 345# ══════════════════════════════════════════════════════════════════════════════346# Section 3 — Router Decision Distribution ← NEW347# ══════════════════════════════════════════════════════════════════════════════348st.subheader("③ Router Decision Distribution")349st.markdown('<p class="section-caption">'350 'How the Router Agent dispatches each question type across sub-strategies. '351 'CoT dominates arithmetic; Zero-shot handles span efficiently.</p>',352 unsafe_allow_html=True)353 354col5, col6 = st.columns([1.1, 0.9])355 356with col5:357 rd = ROUTER_DFS[sel_model].set_index("Answer Type")358 rd_pct = rd.div(rd.sum(axis=1), axis=0) * 100359 360 fig5 = go.Figure()361 strat_color_map = {"Zero-shot": COLORS["zero_shot"], "Chain-of-Thought": COLORS["chain_of_thought"]}362 for strat in rd_pct.columns:363 fig5.add_trace(go.Bar(364 name=strat,365 x=rd_pct.index,366 y=rd_pct[strat],367 marker_color=strat_color_map[strat],368 text=[f"{v:.0f}%" for v in rd_pct[strat]],369 textposition="inside",370 customdata=rd[strat],371 hovertemplate="%{x}<br>%{customdata} questions → %{y:.0f}%<extra>" + strat + "</extra>",372 ))373 fig5.update_layout(374 title="Router Sub-Strategy Dispatch (% of questions per type)",375 barmode="stack", yaxis_title="Dispatch %", yaxis_range=[0, 110],376 height=320, legend=dict(orientation="h", y=1.15),377 plot_bgcolor="white", paper_bgcolor="white",378 )379 st.plotly_chart(fig5, use_container_width=True)380 381with col6:382 # Heatmap version — shows absolute counts383 z_vals = rd.values.tolist()384 fig6 = go.Figure(go.Heatmap(385 z=z_vals,386 x=rd.columns.tolist(),387 y=rd.index.tolist(),388 colorscale="Greens",389 text=[[str(v) for v in row] for row in z_vals],390 texttemplate="%{text}",391 showscale=True,392 ))393 fig6.update_layout(394 title="Dispatch Count Heatmap",395 height=320,396 plot_bgcolor="white", paper_bgcolor="white",397 )398 st.plotly_chart(fig6, use_container_width=True)399 400st.divider()401 402# ══════════════════════════════════════════════════════════════════════════════403# Section 4 — Failure Taxonomy404# ══════════════════════════════════════════════════════════════════════════════405st.subheader("④ Failure Taxonomy")406st.markdown('<p class="section-caption">4 systematic failure patterns identified across all strategies and models.</p>',407 unsafe_allow_html=True)408 409patterns = report["failure_patterns"]410freq_colors = {"high": "#FEE2E2", "medium": "#FEF9C3", "low": "#DCFCE7"}411freq_border = {"high": "#FCA5A5", "medium": "#FDE68A", "low": "#86EFAC"}412 413cols = st.columns(len(patterns))414for col, p in zip(cols, patterns):415 bg = freq_colors.get(p["frequency"], "#F3F4F6")416 bdr = freq_border.get(p["frequency"], "#E5E7EB")417 with col:418 st.markdown(f"""419 <div style="background:{bg}; border:1px solid {bdr}; padding:14px 16px;420 border-radius:8px; height:230px; position:relative;">421 <div style="font-weight:700; font-size:14px; margin-bottom:4px">{p['name']}</div>422 <div style="font-size:11px; color:#6B7280; margin-bottom:8px">423 Affects: <b>{p['affected']}</b> · Frequency: <b>{p['frequency']}</b>424 </div>425 <hr style="margin:8px 0; border-color:{bdr}">426 <div style="font-size:12px; color:#374151; line-height:1.5">{p['description']}</div>427 <div style="font-size:12px; margin-top:10px; color:#1D4ED8">428 💡 <i>{p['fix']}</i>429 </div>430 </div>""", unsafe_allow_html=True)431 432st.divider()433 434# ══════════════════════════════════════════════════════════════════════════════435# Section 5 — Cost-Performance Trade-off436# ══════════════════════════════════════════════════════════════════════════════437st.subheader("⑤ Cost-Performance Trade-off")438st.markdown('<p class="section-caption">'439 'Risk-Weighted F1 penalises arithmetic errors 2× — higher stakes in financial reporting.</p>',440 unsafe_allow_html=True)441 442col7, col8 = st.columns(2)443 444with col7:445 fig7 = go.Figure()446 for metric, opacity, name in [447 ("overall_f1", 1.0, "Standard F1"),448 ("weighted_f1", 0.45, "Weighted F1 (arith ×2)"),449 ]:450 fig7.add_trace(go.Bar(451 name=name,452 x=[LABELS[s] for s in df_m["strategy"]],453 y=df_m[metric],454 marker_color=[COLORS[s] for s in df_m["strategy"]],455 opacity=opacity,456 text=[f"{v:.3f}" for v in df_m[metric]],457 textposition="outside",458 ))459 fig7.update_layout(460 title=f"Standard vs Risk-Weighted F1 ({sel_model.split('-')[-1].upper()})",461 barmode="group", yaxis_range=[0, 0.95], height=340,462 legend=dict(orientation="h", y=1.12),463 plot_bgcolor="white", paper_bgcolor="white",464 )465 st.plotly_chart(fig7, use_container_width=True)466 467with col8:468 fig8 = go.Figure(go.Bar(469 x=[LABELS[s] for s in df_m["strategy"]],470 y=df_m["cost_per_correct"],471 marker_color=[COLORS[s] for s in df_m["strategy"]],472 text=[f"${v:.4f}" for v in df_m["cost_per_correct"]],473 textposition="outside",474 ))475 fig8.update_layout(476 title=f"Cost per Correct Answer ({sel_model.split('-')[-1].upper()})",477 yaxis_title="USD per correct answer",478 height=340,479 plot_bgcolor="white", paper_bgcolor="white",480 )481 st.plotly_chart(fig8, use_container_width=True)482 483# ── Footer ────────────────────────────────────────────────────────────────────484st.divider()485st.caption(486 "Multi-Agent Financial QA Optimization Engine · "487 "TAT-QA dataset (ACL 2021) · "488 "Strategies: Zero-shot, Chain-of-Thought, Router Agent · "489 "Tracked with MLflow"490)