Spaces:
Running
Running
| """Capabilibara Hugging Face Space App (HCAI-Lab). | |
| Capability Provenance in Language Models: A Case Study in Social Reasoning (COLM 2026). | |
| Influence data is loaded from HCAI-Lab/dolma3-influence-heatmaps (bin-level TrackStar | |
| scores over the 576-bin WebOrganizer taxonomy) and z-score standardized within each | |
| benchmark, following the paper's aggregation (Section 3.3). | |
| """ | |
| import numpy as np | |
| import pandas as pd | |
| import gradio as gr | |
| import plotly.graph_objects as go | |
| from huggingface_hub import hf_hub_download | |
| # WebOrganizer taxonomy (Wettig et al. 2025): CSV label -> paper display name | |
| # (paper Tables 2 and 3), in the canonical grid order used by the dataset. | |
| TOPIC_LABELS = { | |
| "adult_content": "Adult", | |
| "art_and_design": "Art & Design", | |
| "crime_and_law": "Crime & Law", | |
| "education_and_jobs": "Education & Jobs", | |
| "electronics_and_hardware": "Hardware", | |
| "entertainment": "Entertainment", | |
| "fashion_and_beauty": "Fashion & Beauty", | |
| "finance_and_business": "Finance & Business", | |
| "food_and_dining": "Food & Dining", | |
| "games": "Games", | |
| "health": "Health", | |
| "history_and_geography": "History", | |
| "home_and_hobbies": "Home & Hobbies", | |
| "industrial": "Industrial", | |
| "literature": "Literature", | |
| "politics": "Politics", | |
| "religion": "Religion", | |
| "science_math_and_technology": "Science & Technology", | |
| "social_life": "Social Life", | |
| "software": "Software", | |
| "software_development": "Software Development", | |
| "sports_and_fitness": "Sports & Fitness", | |
| "transportation": "Transportation", | |
| "travel_and_tourism": "Travel", | |
| } | |
| FORMAT_LABELS = { | |
| "about_org": "About (Org.)", | |
| "about_pers": "About (Personal)", | |
| "academic_writing": "Academic Writing", | |
| "audio_transcript": "Audio Transcript", | |
| "comment_section": "Comment Section", | |
| "content_listing": "Content Listing", | |
| "creative_writing": "Creative Writing", | |
| "customer_support": "Customer Support", | |
| "documentation": "Documentation", | |
| "faq": "FAQ", | |
| "knowledge_article": "Knowledge Article", | |
| "legal_notices": "Legal Notices", | |
| "listicle": "Listicle", | |
| "news_article": "News Article", | |
| "news_org": "News (Org.)", | |
| "nonfiction_writing": "Nonfiction Writing", | |
| "personal_blog": "Personal Blog", | |
| "product_page": "Product Page", | |
| "q_a_forum": "Q&A Forum", | |
| "spam_ads": "Spam / Ads", | |
| "structured_data": "Structured Data", | |
| "truncated": "Truncated", | |
| "tutorial": "Tutorial", | |
| "user_review": "User Review", | |
| } | |
| TOPICS = list(TOPIC_LABELS.values()) | |
| FORMATS = list(FORMAT_LABELS.values()) | |
| # Benchmarks with released bin-level scores in dolma3-influence-heatmaps. | |
| BENCHMARKS = { | |
| "SocialIQA (social reasoning)": "socialiqa", | |
| "ARC-Challenge (STEM reasoning)": "arc_challenge", | |
| "MMLU Social Sciences (social knowledge)": "mmlu_social_science", | |
| "MMLU STEM (STEM knowledge)": "mmlu_stem", | |
| } | |
| CONTRAST = "Contrastive: SocialIQA − ARC-Challenge (Δz)" | |
| Z_CAP = 2.5 # color scale cap used in the paper's heatmaps (|z| = 2.5) | |
| def load_bin_scores(): | |
| """Download bin-level influence scores from the Hub and z-standardize per benchmark.""" | |
| z_matrices, doc_counts = {}, None | |
| for name, key in BENCHMARKS.items(): | |
| path = hf_hub_download( | |
| repo_id="HCAI-Lab/dolma3-influence-heatmaps", | |
| filename=f"influence_bin_scores/queries_{key}_bin_scores.csv", | |
| repo_type="dataset", | |
| ) | |
| df = pd.read_csv(path) | |
| df["z"] = (df.mean_score - df.mean_score.mean()) / df.mean_score.std() | |
| grid = df.pivot(index="topic_label", columns="format_label", values="z") | |
| grid = grid.loc[list(TOPIC_LABELS), list(FORMAT_LABELS)] | |
| z_matrices[name] = grid.to_numpy() | |
| if doc_counts is None: | |
| counts = df.pivot(index="topic_label", columns="format_label", values="doc_count") | |
| doc_counts = counts.loc[list(TOPIC_LABELS), list(FORMAT_LABELS)].to_numpy() | |
| return z_matrices, doc_counts | |
| Z_MATRICES, DOC_COUNTS = load_bin_scores() | |
| Z_MATRICES[CONTRAST] = ( | |
| Z_MATRICES["SocialIQA (social reasoning)"] - Z_MATRICES["ARC-Challenge (STEM reasoning)"] | |
| ) | |
| def create_heatmap(metric_choice: str): | |
| """Signed-influence heatmap over the 576-bin grid (RdBu, capped at |z| = 2.5).""" | |
| z_data = Z_MATRICES[metric_choice] | |
| if metric_choice == CONTRAST: | |
| title = "Contrastive provenance Δz: SocialIQA − ARC-Challenge (positive = social-tilted bin)" | |
| colorbar_title = "Δz" | |
| else: | |
| title = f"Signed mean influence (z) — {metric_choice}, Dolma3 working set" | |
| colorbar_title = "z" | |
| fig = go.Figure( | |
| go.Heatmap( | |
| z=z_data, | |
| x=FORMATS, | |
| y=TOPICS, | |
| zmin=-Z_CAP, | |
| zmax=Z_CAP, | |
| colorscale="RdBu_r", | |
| customdata=DOC_COUNTS, | |
| hovertemplate=( | |
| "<b>%{y} × %{x}</b><br>" | |
| f"{colorbar_title} = %{{z:.2f}}<br>" | |
| "working-set docs = %{customdata:,}<extra></extra>" | |
| ), | |
| colorbar=dict(title=colorbar_title), | |
| ) | |
| ) | |
| fig.update_layout( | |
| title=title, | |
| font_family="Inter, sans-serif", | |
| margin=dict(l=40, r=40, t=60, b=40), | |
| height=680, | |
| xaxis=dict(title="Corpus Format (24)", tickangle=-45), | |
| yaxis=dict(title="Corpus Topic (24)", autorange="reversed"), | |
| ) | |
| return fig | |
| def get_bin_details(topic: str, format_type: str): | |
| """Real per-bin stats for a topic-format cell across all benchmarks.""" | |
| t_idx = TOPICS.index(topic) | |
| f_idx = FORMATS.index(format_type) | |
| rows = "\n".join( | |
| f" | {name} | `{Z_MATRICES[name][t_idx, f_idx]:+.2f}` |" for name in BENCHMARKS | |
| ) | |
| diff = Z_MATRICES[CONTRAST][t_idx, f_idx] | |
| docs = int(DOC_COUNTS[t_idx, f_idx]) | |
| verdict = "🔥 Tilts toward social reasoning (SocialIQA)" if diff > 1.0 else ( | |
| "⚡ Tilts toward STEM reasoning (ARC-Challenge)" if diff < -1.0 else "⚖️ No strong social/STEM tilt" | |
| ) | |
| return f""" | |
| ### Bin: `{topic}` × `{format_type}` | |
| | Benchmark | Signed influence (z) | | |
| |---|---| | |
| {rows} | |
| | **Contrast (SocialIQA − ARC-Challenge)** | **`{diff:+.2f}`** | | |
| **Provenance diagnosis:** {verdict} | |
| *Working-set documents sampled in this bin: {docs:,} (target 10,000/bin; scarce bins fill lower).* | |
| """ | |
| # Build Gradio Interface | |
| theme = gr.themes.Soft( | |
| primary_hue="indigo", | |
| secondary_hue="blue", | |
| neutral_hue="slate" | |
| ) | |
| with gr.Blocks(theme=theme, title="Capabilibara — Capability Provenance in Language Models (COLM 2026)") as demo: | |
| gr.Markdown( | |
| """ | |
| # 🦫 Capabilibara: Capability Provenance in Language Models | |
| ### *A Case Study in Social Reasoning* (COLM 2026) | |
| **Hugging Face Space by [HCAI-Lab](https://huggingface.co/HCAI-Lab)** | [arXiv Paper](https://arxiv.org/abs/2606.19625) | [Project Website](https://eilab.gatech.edu/capabilibara/) | |
| --- | |
| This interactive Space explores training-data attribution across **576 corpus bins** in Dolma3 (24 Topics × 24 Formats WebOrganizer taxonomy), | |
| validating model capability origins using gradient-based influence (TrackStar via Bergson) and selective unlearning. | |
| Heatmap data is loaded live from [dolma3-influence-heatmaps](https://huggingface.co/datasets/HCAI-Lab/dolma3-influence-heatmaps). | |
| """ | |
| ) | |
| with gr.Tabs(): | |
| with gr.Tab("🗺️ 576-Bin Matrix Explorer"): | |
| gr.Markdown( | |
| """ | |
| ### WebOrganizer 24×24 Topic-by-Format Signed Influence | |
| Bin-level mean TrackStar influence, z-scored within each benchmark (paper §3.3); color capped at |z| = 2.5 as in the paper's Figures 2 and 12. | |
| Red bins are supportive, blue bins suppressive. Note the paper's headline pattern: **Literature × Customer Support** and interpersonal | |
| formats (Customer Support, FAQ, Q&A Forum) are strongly positive for SocialIQA only, while Documentation-like bins drive the comparison benchmarks. | |
| """ | |
| ) | |
| with gr.Row(): | |
| metric_dropdown = gr.Dropdown( | |
| choices=[CONTRAST] + list(BENCHMARKS), | |
| value=CONTRAST, | |
| label="Select Benchmark Influence Metric" | |
| ) | |
| heatmap_plot = gr.Plot(label="Corpus Influence Heatmap") | |
| metric_dropdown.change(fn=create_heatmap, inputs=metric_dropdown, outputs=heatmap_plot) | |
| demo.load(fn=create_heatmap, inputs=metric_dropdown, outputs=heatmap_plot) | |
| gr.Markdown("---") | |
| gr.Markdown("### Bin Inspector") | |
| with gr.Row(): | |
| topic_select = gr.Dropdown(choices=TOPICS, value="Literature", label="Select Topic (Y-axis)") | |
| format_select = gr.Dropdown(choices=FORMATS, value="Customer Support", label="Select Format (X-axis)") | |
| bin_output = gr.Markdown() | |
| topic_select.change(fn=get_bin_details, inputs=[topic_select, format_select], outputs=bin_output) | |
| format_select.change(fn=get_bin_details, inputs=[topic_select, format_select], outputs=bin_output) | |
| demo.load(fn=get_bin_details, inputs=[topic_select, format_select], outputs=bin_output) | |
| with gr.Tab("📊 Results & Unlearning Validation"): | |
| gr.Markdown( | |
| """ | |
| ## Headline Study Scale & Key Results | |
| | Metric | Value | Detail | | |
| |---|---|---| | |
| | **Corpus Bins** | `576` | WebOrganizer 24×24 topic-format matrix | | |
| | **Working Set** | `5,678,621` | Stratified unique Dolma3 documents (~10.5B tokens, 10,000/bin) | | |
| | **Source Corpus** | `~1.26B` | Unique documents in the de-duplicated Dolma3 6T mix | | |
| | **Base Model** | `OLMo3-7B` | Gradient index from Base; query gradients from Instruct | | |
| | **Unlearning Effect** | `+1.6 pts` | SocialIQA median paired damage, influence-targeted vs. random in-topic (Wilcoxon BH-adjusted $p \\approx 10^{-5}$, Table 49) | | |
| | **Attribution Compute** | `~37,000` | H200-equivalent GPU-hours across the full pipeline | | |
| ### Key Findings | |
| 1. **SocialIQA is the provenance-structure outlier**: its 576-bin profile correlates with the comparison benchmarks at only r ≤ 0.21, | |
| versus r = 0.76–0.86 among the three. Its support comes from interpersonal formats (Customer Support, FAQ, Q&A Forum) and the | |
| Literature and Social Life topics, while the comparison benchmarks concentrate in Documentation and Academic Writing. | |
| 2. **Unlearning validation**: forgetting the top-200 influence-ranked documents per topic damages SocialIQA more than 5×-larger | |
| random in-topic controls (median +1.6 accuracy points); effects on the comparison benchmarks are weaker, null, or reversed — | |
| attribution identifies capability-relevant regions, not just topic membership. | |
| """ | |
| ) | |
| with gr.Tab("📦 Data Artifacts"): | |
| gr.Markdown( | |
| """ | |
| ## Data Artifacts | |
| Influence matrices, sampling manifests, cross-probe statistics, unlearning checkpoints, and Hub artifacts. | |
| Each group below backs a figure or analysis in the [main paper](https://arxiv.org/abs/2606.19625); all artifacts are hosted under | |
| [HCAI-Lab](https://huggingface.co/HCAI-Lab) on the Hugging Face Hub, with code at [eilab-gt/capabilibara](https://github.com/eilab-gt/capabilibara). | |
| All artifacts are public, per the paper's open-source release. | |
| ### Influence Matrices — Figure 2 & Figure 12 (signed z-score heatmaps) | |
| Bin-level signed mean influence (576 × 4 matrix) behind the marginal panels of Figure 2 and the full signed heatmaps of Figure 12. | |
| - [dolma3-trackstar-influence-scores](https://huggingface.co/datasets/HCAI-Lab/dolma3-trackstar-influence-scores) — per-document TrackStar (Bergson) influence scores: OLMo3-7B Base gradient index × Instruct query gradients | |
| - [dolma3-influence-heatmaps](https://huggingface.co/datasets/HCAI-Lab/dolma3-influence-heatmaps) — rendered 24×24 topic-by-format signed heatmaps per benchmark | |
| - [dolma3-data-attribution-index](https://huggingface.co/datasets/HCAI-Lab/dolma3-data-attribution-index) — corpus gradient index used for attribution | |
| - [dolma3-attribution-job-archive](https://huggingface.co/datasets/HCAI-Lab/dolma3-attribution-job-archive) — archived attribution job outputs (~37K H200-equiv. GPU-hours) | |
| ### Sampling Manifests — Figure 1, stage 1 (corpus binning & stratification) | |
| The stratified working set of Figure 1: 5,678,621 documents (~10.5B tokens; 10,000/bin) drawn from ~1.26B unique de-duplicated Dolma3 documents. | |
| - [dolma3-6t-corpus-manifest](https://huggingface.co/datasets/HCAI-Lab/dolma3-6t-corpus-manifest) · [dolma3-olmo3-corpus-manifest](https://huggingface.co/datasets/HCAI-Lab/dolma3-olmo3-corpus-manifest) — corpus shard manifests | |
| - [dolma3-6t-unique](https://huggingface.co/datasets/HCAI-Lab/dolma3-6t-unique) · [dolma3-6t-bloom-index](https://huggingface.co/datasets/HCAI-Lab/dolma3-6t-bloom-index) — de-duplication state and bloom index | |
| - Stratified per-bin samples: [500](https://huggingface.co/datasets/HCAI-Lab/dolma3-6t-sample-500-docs) · [1k](https://huggingface.co/datasets/HCAI-Lab/dolma3-6t-sample-1000-docs) · [5k](https://huggingface.co/datasets/HCAI-Lab/dolma3-6t-sample-5000-docs) · [10k](https://huggingface.co/datasets/HCAI-Lab/dolma3-6t-sample-10000-docs) (working set) · [50k](https://huggingface.co/datasets/HCAI-Lab/dolma3-6t-sample-50000-docs) · [100k](https://huggingface.co/datasets/HCAI-Lab/dolma3-6t-sample-100000-docs) docs/bin | |
| - [dolma3-6t-preconditioner-100k](https://huggingface.co/datasets/HCAI-Lab/dolma3-6t-preconditioner-100k) — TrackStar preconditioner sample | |
| ### Cross-Probe Statistics — Table 1 benchmarks & Figure 34 (probe-profile correlations) | |
| The 2×2 benchmark queries (SocialIQA, ARC-Challenge, MMLU Social Sciences, MMLU STEM) plus the nine held-out social probes whose 576-bin | |
| profile correlations show SocialIQA as the structural outlier (r ≤ 0.21 vs. r = 0.76–0.86 among the comparison benchmarks). | |
| - Attribution query sets: [base-query-data](https://huggingface.co/datasets/HCAI-Lab/base-query-data) · [instruct-query-data](https://huggingface.co/datasets/HCAI-Lab/instruct-query-data) · [instruct-cot-query-data](https://huggingface.co/datasets/HCAI-Lab/instruct-cot-query-data) | |
| - OLMES evaluations: [olmo3-7b-base](https://huggingface.co/datasets/HCAI-Lab/olmes-eval-olmo3-7b-base) · [olmo3-7b-instruct-base](https://huggingface.co/datasets/HCAI-Lab/olmes-eval-olmo3-7b-instruct-base) · [olmo3-7b-instruct-cot](https://huggingface.co/datasets/HCAI-Lab/olmes-eval-olmo3-7b-instruct-cot) · [olmo3-7b-thinking](https://huggingface.co/datasets/HCAI-Lab/olmes-eval-olmo3-7b-thinking) | |
| - Held-out social probe statistics: [soc91-stats](https://huggingface.co/datasets/HCAI-Lab/soc91-stats) · [soc91-labels](https://huggingface.co/datasets/HCAI-Lab/soc91-labels) · [soc139-quality-sidecars](https://huggingface.co/datasets/HCAI-Lab/soc139-quality-sidecars) · [tombench-en](https://huggingface.co/datasets/HCAI-Lab/tombench-en) (MIT; evaluation only, per ToMBench authors) | |
| ### Unlearning Checkpoints — §4.3, Figure 43 & Table 49 (causal validation) | |
| NGDiff rank-8 LoRA unlearning on OLMo3-7B Base: top-200 influence-selected documents per topic vs. 1,000 random in-topic controls, | |
| three seeds per condition (SocialIQA paired difference significant at BH-adjusted p ≈ 10⁻⁵). | |
| - [unlearning-checkpoints](https://huggingface.co/HCAI-Lab/unlearning-checkpoints) — merged unlearned checkpoints for all conditions and seeds | |
| ### Hub Artifacts — cross-ecosystem replication (Appendices L & M) | |
| The causal test repeated end-to-end on Comma v0.1 7B-2T (Common Pile) and DCLM-Baseline-7B. | |
| - Comma OLMES evals: [comma-7b-1t](https://huggingface.co/datasets/HCAI-Lab/olmes-eval-comma-7b-1t) · [comma-7b-2t](https://huggingface.co/datasets/HCAI-Lab/olmes-eval-comma-7b-2t) | |
| - DCLM: [dclm-olmes-eval](https://huggingface.co/datasets/HCAI-Lab/dclm-olmes-eval) · [dclm-baseline-manifest](https://huggingface.co/datasets/HCAI-Lab/dclm-baseline-manifest) · [dclm-baseline-labels](https://huggingface.co/datasets/HCAI-Lab/dclm-baseline-labels) · [dclm-baseline-working-sample](https://huggingface.co/datasets/HCAI-Lab/dclm-baseline-working-sample) | |
| - [dolma3-corpus-explorer](https://huggingface.co/spaces/HCAI-Lab/dolma3-corpus-explorer) — companion Space for browsing the Dolma3 working set | |
| """ | |
| ) | |
| with gr.Tab("📜 Citation"): | |
| gr.Markdown( | |
| """ | |
| ### Cite This Work | |
| ```bibtex | |
| @inproceedings{matlin2026capabilityprovenance, | |
| title = {Capability Provenance in Language Models: A Case Study in Social Reasoning}, | |
| author = {Glenn Matlin and Chandreyi Chakraborty and Saehee Eom and Mika Okamoto and | |
| Rayan Castilla and Louis Jaburi and Alvin Deng and Taywon Min and | |
| Lucia Quirke and Stella Biderman and Mark Riedl}, | |
| booktitle = {Proceedings of the Conference on Language Modeling (COLM 2026)}, | |
| year = {2026}, | |
| eprint = {2606.19625}, | |
| archivePrefix = {arXiv}, | |
| primaryClass = {cs.CL}, | |
| url = {https://arxiv.org/abs/2606.19625} | |
| } | |
| ``` | |
| """ | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |