social-data-attribution-w2/scripts/generate_6t_figures.py

"""Generate paper figures for the 6T working sample (SOC-135).

Downloads a working sample manifest from HuggingFace, converts it to
EDA-format CSVs, runs the existing WebOrganizer report pipeline, and
generates the Table C1 LaTeX fragment.

Usage:
    uv run python scripts/generate_6t_figures.py
    uv run python scripts/generate_6t_figures.py --sample-name sample_5000_docs
    uv run python scripts/generate_6t_figures.py --local-manifest /path/to/manifest.parquet
"""

from __future__ import annotations

import argparse
import json
import logging
from pathlib import Path

logger = logging.getLogger(__name__)

DEFAULT_SAMPLE_NAME = "sample_10000_docs"
DEFAULT_OUTPUT_DIR = Path("artifacts/paper_figures_6t")
HF_ORG = "HCAI-Lab"
HF_REPO_PREFIX = "dolma3_6T"


def _download_sample_artifacts(
    sample_name: str,
    cache_dir: Path,
) -> tuple[Path, Path, Path]:
    from huggingface_hub import hf_hub_download

    repo_id = f"{HF_ORG}/{HF_REPO_PREFIX}_{sample_name}"
    logger.info("Downloading artifacts from %s", repo_id)

    manifest_path = Path(
        hf_hub_download(
            repo_id,
            "working_sample_manifest.parquet",
            repo_type="dataset",
            cache_dir=str(cache_dir),
        )
    )
    contract_path = Path(
        hf_hub_download(
            repo_id,
            "sample_contract.json",
            repo_type="dataset",
            cache_dir=str(cache_dir),
        )
    )
    bin_summary_path = Path(
        hf_hub_download(
            repo_id,
            "bin_summary.csv",
            repo_type="dataset",
            cache_dir=str(cache_dir),
        )
    )
    return manifest_path, contract_path, bin_summary_path


def _write_comparison_note(output_dir: Path, contract: dict) -> Path:
    sample_docs = contract["WORKING_SAMPLE_REALIZED_DOC_COUNT"]
    sample_tokens = contract["WORKING_SAMPLE_REALIZED_TOKEN_TOTAL"]
    underfilled = contract["WORKING_SAMPLE_UNDERFILLED_BIN_COUNT"]
    covered = contract["WORKING_SAMPLE_COVERED_BIN_COUNT"]
    total_bins = contract.get("WORKING_SAMPLE_TOTAL_BIN_COUNT", 576)

    note = f"""# Comparison note: pool-sample vs 6T working sample

## Previous corpus (SOC-12 / SOC-13)

The earlier paper figures used a 175B-token random sample from the full
~9T-token `allenai/dolma3_pool` (163M docs, 100 shards).  EDA aggregates
in `artifacts/dolma_eda/` reflect that population (~87M docs after
enrichment coverage).

## Current corpus (SOC-135)

The paper uses the locked 10K docs/bin stratified working sample drawn
from the deduplicated 6T Dolma3 training mix (~1.26B unique docs).
Sample designated as final by SOC-148 (2026-03-30).

- Documents: {sample_docs:,}
- Tokens: {sample_tokens:,}
- Bins covered: {covered}/{total_bins}
- Underfilled bins: {underfilled}
- Seed: 42

## Key differences from the old pool-sample figures

1. The 6T population is deduplicated; the old 9T pool was not.
2. The working sample is stratified across 576 topic x format bins
   rather than being a uniform random draw.
3. Marginal distributions differ because stratification flattens
   the natural skew toward high-mass bins.
4. The heatmap shows more uniform coverage by design, with only
   {underfilled} underfilled bins.
"""
    path = output_dir / "comparison_note.md"
    path.write_text(note)
    logger.info("Wrote %s", path)
    return path


def _write_handoff_note(output_dir: Path) -> Path:
    note = """# Handoff note for SOC-44 and SOC-49

## SOC-44: Paper writing (Section 5 / Experimental Setup)

The following assets are in `artifacts/paper_figures_6t/` and use the
locked 10K docs/bin working sample (SOC-148, 2026-03-30):

- `table_c1.tex`: Two-row corpus summary (population + working sample).
  Include via `\\input{table_c1.tex}` in the appendix or Section 5.
- `fig_topic_token_count.pdf`: Topic distribution bar chart (Figure C1).
- `fig_format_token_count.pdf`: Format distribution bar chart (Figure C2).
- `fig_heatmap_token_count.pdf`: 24x24 joint topic-format heatmap (Figure C3).
- `table_concentration.tex`: Concentration metrics (Gini, effective bins, top-K share).

To regenerate after any pipeline changes:

```
uv run python scripts/generate_6t_figures.py
```

## SOC-49: Final paper review

Before submission, verify that:
1. All figure values match the locked sample contract (sample_contract.json).
2. The comparison note accurately reflects the pool-sample to 6T shift.
"""
    path = output_dir / "handoff_note.md"
    path.write_text(note)
    logger.info("Wrote %s", path)
    return path


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Generate paper figures for the 6T working sample (SOC-135).",
    )
    parser.add_argument(
        "--sample-name",
        default=DEFAULT_SAMPLE_NAME,
        help="HuggingFace sample name (default: %(default)s).",
    )
    group = parser.add_mutually_exclusive_group()
    group.add_argument(
        "--local-manifest",
        type=Path,
        help="Path to a local manifest parquet (skips HF download).",
    )
    group.add_argument(
        "--local-dir",
        type=Path,
        help="Local directory containing manifest.parquet, sample_contract.json, bin_summary.csv.",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=DEFAULT_OUTPUT_DIR,
        help="Output directory (default: %(default)s).",
    )
    parser.add_argument(
        "--format",
        choices=["pdf", "png", "html", "both", "all"],
        default="all",
        dest="output_format",
    )
    parser.add_argument("--verbose", action="store_true")
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(levelname)s %(name)s: %(message)s",
    )

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    eda_dir = output_dir / "eda"
    cache_dir = output_dir / ".hf_cache"

    if args.local_dir:
        manifest_path = args.local_dir / "working_sample_manifest.parquet"
        contract_path = args.local_dir / "sample_contract.json"
    elif args.local_manifest:
        manifest_path = args.local_manifest
        contract_path = manifest_path.parent / "sample_contract.json"
    else:
        manifest_path, contract_path, _ = _download_sample_artifacts(
            args.sample_name, cache_dir
        )

    if not manifest_path.exists():
        logger.error("Manifest not found: %s", manifest_path)
        return 1
    if not contract_path.exists():
        logger.error("Contract not found: %s", contract_path)
        return 1

    from dolma.distribution_report.manifest_bridge import manifest_to_eda_csvs
    from dolma.distribution_report.runner import run_report
    from dolma.distribution_report.table_c1 import write_table_c1

    logger.info("Converting manifest to EDA CSVs")
    manifest_to_eda_csvs(manifest_path, eda_dir)

    logger.info("Running WebOrganizer report pipeline")
    run_report(
        eda_dir=eda_dir,
        output_dir=output_dir,
        output_format=args.output_format,
        run_label=f"6T-working-sample-{args.sample_name}",
        use_dummy=True,
        representative_manifest=None,
        stratified_manifest=None,
    )

    logger.info("Generating Table C1")
    write_table_c1(contract_path, output_dir)

    contract = json.loads(contract_path.read_text())
    _write_comparison_note(output_dir, contract)
    _write_handoff_note(output_dir)

    logger.info("All outputs written to %s", output_dir)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())