app.py

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