social-data-attribution-w2/README.md

# Data Attribution for Social Reasoning in Pre-Training Corpus

This repo is a workspace for running Bergson's quick start workflow on datasets like Dolma and SocialIQA.

## Install

1. Bootstrap the required local dependencies from the main checkout:

```bash
bash scripts/bootstrap_local_deps.sh
```

This script is idempotent. It clones `bergson` and `olmes` when missing, checks out the pinned commits, and applies the tracked local patches.

2. Install build dependencies:
```bash
sudo apt-get update && sudo apt-get install -y build-essential g++
```

3. Then setup the Python environment via `uv`:

```bash
uv pip install setuptools pybind11
uv sync --no-build-isolation
```

This project currently targets Python 3.12. The pinned `ai2-olmes` checkout still requires `<3.13`.

For a fresh git worktree, run `bash scripts/bootstrap/setup_worktree.sh` before the first `uv` command. It bootstraps the shared local dependencies from the main checkout, links them into the worktree, copies gitignored files (`settings.local.json`, `autopilot/state.json`, `uv.lock`) from the main checkout, installs `pybind11` for local builds, and then runs `uv sync --frozen`. The script uses uv's default cache (`~/.cache/uv`) for fast CoW clones on APFS; override with `UV_CACHE_DIR` if needed.

After bootstrap, use `uv run --no-sync ...` for routine verification commands in the worktree:

```bash
uv run --no-sync pytest -q
uv run --no-sync ruff check .
uv run --no-sync ruff format --check .
```

This avoids a second sync/build step during `uv run`. Plain `uv run ...` may try to resolve the build-system requirements for `setuptools` and `wheel`, which can fail in sandboxed or offline sessions.

If dependencies changed, rerun `bash scripts/bootstrap/setup_worktree.sh` or `uv sync --no-build-isolation` before using `--no-sync` again. If `uv run --no-sync ...` still fails, fall back to `./.venv/bin/python -m pytest ...` or `./.venv/bin/ruff ...`.

### Optional pre-push checks

Install the repo-local pre-push hook when you want push-time quality checks:

```bash
bash scripts/bootstrap/install_git_hooks.sh
```

The hook runs only before `git push`. It checks the commits being pushed with
`git diff --check`, then runs `ruff check` and `ruff format --check` on pushed
Python files. It does not run on every commit.

## R2 credentials

R2 hosts the raw 6T `.jsonl.zst` source corpus and the original per-shard sidecars. Many workflows can stay HF-only since the sidecars are now mirrored to `HCAI-Lab/soc91-labels`, `HCAI-Lab/soc91-stats`, and `HCAI-Lab/soc139-quality-sidecars` (see `docs/data_home/README.md`). You only need R2 for the raw shards or when you want per-shard sidecar granularity.

R2 scripts expect `R2_ACCESS_KEY_ID` and `R2_SECRET_ACCESS_KEY` in the environment.

If the credentials live in 1Password, use the wrapper script:

```bash
bash scripts/bootstrap/with_r2_credentials.sh \
  uv run python scripts/validation/analyze_sidecar_labels.py \
  --output-dir ~/scratch/soc91/eda \
  --max-shards 5
```

The wrapper defaults to the 1Password item `Research` / `Cloudflare R2 API Credentials`.
Override the lookup with `OP_R2_VAULT`, `OP_R2_ITEM`, `OP_R2_ACCESS_KEY_FIELD`, and
`OP_R2_SECRET_ACCESS_KEY_FIELD` if needed.

## Phoenix sidecar EDA

The sidecar EDA CLI now supports a fixed-manifest workflow for Phoenix CPU jobs.

Create a benchmark manifest from completed sidecars:

```bash
bash scripts/bootstrap/with_r2_credentials.sh \
  uv run python scripts/validation/analyze_sidecar_labels.py \
  --mode manifest \
  --output-dir ~/scratch/soc91/eda_benchmark \
  --manifest-path ~/scratch/soc91/eda_benchmark/completed_sidecars_manifest.jsonl \
  --max-shards 1000 \
  --chunk-count 32
```

Run a single benchmark job on Phoenix:

```bash
sbatch scripts/slurm/eda/sidecar_eda_benchmark_phoenix.sbatch
```

Run a chunked Phoenix array from an existing manifest:

```bash
sbatch --array=0-31 scripts/slurm/eda/sidecar_eda_array_phoenix.sbatch \
  --export=MANIFEST_PATH=~/scratch/soc91/eda_benchmark/completed_sidecars_manifest.jsonl,CHUNK_COUNT=32,OUTDIR=~/scratch/soc91/sidecar_eda_runs/current

sbatch --dependency=afterok:<array_job_id> scripts/slurm/eda/sidecar_eda_merge_phoenix.sbatch \
  --export=OUTDIR=~/scratch/soc91/sidecar_eda_runs/current
```

Worker jobs write `state.json`, `benchmark.json`, and the SQLite cache under each worker directory.
The merge job reads those worker artifacts and writes the final CSV and JSON outputs in `OUTDIR`.
The current Phoenix-safe default is `CHUNK_COUNT=32`, which stays below the `embers` submit limit we hit in practice.
Using a fixed `CHUNK_COUNT` keeps shard-to-worker assignment stable across reruns as the completed set grows,
which lets unchanged shards reuse the same worker cache instead of shifting between chunk directories.

To launch the full Phoenix path in one step:

```bash
scripts/slurm/eda/launch_sidecar_eda_phoenix.sh
```

Optional overrides:

- `CHUNK_COUNT=32`
- `RUN_ID=soc101_full_hash_YYYYMMDD_HHMMSS`
- `RUN_ROOT=$HOME/scratch/soc91/sidecar_eda_runs`
- `MAX_SHARDS=1000` for a bounded dry run
- `PYTHONEXEC=/path/to/python`

## Modal sidecar EDA

The current SOC-101 Modal entrypoint is `scripts/modal/sidecar_eda.py`.
It writes run artifacts under `soc91-stats/sidecar-eda/<run_id>/` in R2, maps
worker chunks in parallel, and merges the final outputs in place.

Run a bounded check:

```bash
python3 -m modal run scripts/modal/sidecar_eda.py --max-shards 1000 --chunk-count 32
```

Run the full current completed-sidecar set:

```bash
python3 -m modal run scripts/modal/sidecar_eda.py --chunk-count 128
```

## Bergson + Transformers version note (OLMo3)

This repo currently vendors Bergson from `./bergson` and pins that checkout through `scripts/bootstrap_local_deps.sh`.
The current pinned upstream release is `v0.9.0`.

Historical note:

- Older Bergson releases in the `v0.4.x` line pinned `transformers<4.56.0`.
- This repo carried a local override so OLMo3-compatible Transformers versions could be used.
- Upstream removed that transformers pin before `v0.9.0`, so that local Bergson patch is no longer the default path here.

The vendored workflow stays in place so bootstrap and worktree setup can reproduce the exact Bergson release used by the repo.

⚠️ Upstream docs and code are not perfectly aligned on every CLI example. Treat the installed `bergson --help` output as the source of truth when updating or running repo scripts.

## Bergson quick start (smoke)

This follows the Bergson quick start tutorial.

⚠️ The exact CLI flags depend on the installed Bergson version. Start by running `bergson --help` in the current project environment and use that output as the source of truth.

Build an index on a small dataset:

```bash
uv run bergson build runs/quickstart \
  --model EleutherAI/pythia-14m \
  --dataset NeelNanda/pile-10k \
  --truncation
```

Query the index interactively:

```bash
uv run bergson query runs/quickstart
```

## Pool sample dataset

The primary corpus for attribution experiments is a 175B-token random sample drawn from the full ~9T-token `allenai/dolma3_pool`. This avoids the confounding effects of AllenAI's RegMix upsampling in the curated mix.

**Dataset:** https://huggingface.co/datasets/HCAI-Lab/archive-dolma3-pool-150b

| Property | Value |
|----------|-------|
| Source | `allenai/dolma3_pool` (~9T tokens, 2732 shards) |
| Total tokens | 175,000,463,223 (175.0B) |
| Total docs | 163,560,678 |
| Format | `.jsonl.zst` (100 shards) + `.parquet` manifests |
| Sampling | Uniform random, seed=42, 5 jobs x 20 workers |
| Storage | ~568 GB |

The sample is organized as `job_{0-4}/worker_{000-019}/train-00000.jsonl.zst` with per-worker parquet manifests containing `doc_id`, `shard_path`, `token_count`, `weborganizer_topic`, and `weborganizer_format`. WebOrganizer labels are null in the source pool data and are populated downstream by the enrichment pipeline (SOC-5).

The 25B token excess over the 150B target provides headroom for downstream cleaning (blank/short doc removal).

## Sidecar corpus manifest

The unified corpus manifest now supports a sidecar-join path for the deduplicated 6T workflow.
This path reads completed `soc91-labels/*.parquet` sidecars plus their `.stats.json` files,
streams the matching source shards, and writes a partitioned Parquet dataset with one row per
document and no text payload.

The sidecar-backed manifest retains:

- `doc_id`, `shard_path`, `source_family`
- `original_word_count`, `estimated_token_count`, `token_count`
- canonical sampling labels in `weborganizer_topic`, `weborganizer_format`, and `bin_id`
- URL and noURL top-1 label ids, label strings, and confidences
- URL and noURL probability vectors when they are present in the sidecar schema

Local example:

```bash
uv run data-attribution-corpus-manifest \
  --sidecar-dir /path/to/sidecars \
  --source-dir /path/to/source_shards \
  --output /path/to/corpus_manifest_dataset
```

R2 example:

```bash
uv run data-attribution-corpus-manifest \
  --sidecar-r2-prefix soc91-labels \
  --bucket soc127-dedup \
  --endpoint-url https://0934ab8e84ac8f4e81decaf3eb121337.r2.cloudflarestorage.com \
  --output /path/to/corpus_manifest_dataset
```

Reusable work-manifest snapshot:

```bash
uv run data-attribution-corpus-manifest \
  --sidecar-r2-prefix soc91-labels \
  --bucket soc127-dedup \
  --write-work-manifest runs/manifests/soc95_work_manifest.jsonl \
  --manifest-only \
  --output /tmp/soc95-manifest-snapshot
```

When a saved work manifest contains R2 sidecar keys, later runs may still pass
`--source-dir /path/to/local/source_root` to read source shards locally while
continuing to fetch sidecar parquet files from R2.

This writes:

- `data/*.parquet` shard partitions
- `meta/*.json` per-sidecar join metadata
- `manifest_summary.json`
- `bin_totals.parquet`
- `bin_totals.csv`
- `mismatch_report.jsonl`

Exact sidecar schema inventory for rerun planning:

```bash
uv run data-attribution-sidecar-schema-inventory \
  --sidecar-r2-prefix soc91-labels \
  --bucket soc127-dedup \
  --output-dir runs/reports/soc95_sidecar_schema_inventory \
  --chunk-count 32
```

This writes a reusable completed-sidecar snapshot, per-chunk JSONL outputs, a merged exact old/new schema inventory, and `old_schema_sidecar_ids.txt` plus `old_schema_sidecar_keys.txt` for rerun handoff.

For the Modal CPU execution path, see [docs/SOC95_MODAL_RUNBOOK.md](docs/SOC95_MODAL_RUNBOOK.md).

To draw sampling manifests from a SOC-95 dataset root, use the supported manifest sampler:

```bash
uv run data-attribution-manifest-sample \
  --input /path/to/corpus_manifest_dataset \
  --output-dir /path/to/sample_outputs
```

This loader keeps the base manifest unchanged, drops unsampleable rows such as null topic/bin rows for the sampling path only, and writes:

- `pool_bin_stats.parquet`
- `pool_stratified_manifest.parquet`
- `pool_representative_manifest.parquet`

## WebOrganizer report

I keep the reusable WebOrganizer report workflow in [docs/WEBORGANIZER_REPORT.md](docs/WEBORGANIZER_REPORT.md).
It is the path I use when I want the 24 x 24 topic x format figures, concentration table, note, and metadata for a new datapool.

Local example:

```bash
uv run data-attribution-weborganizer-report \
  --eda-dir artifacts/dolma_eda \
  --output-dir artifacts/weborganizer_report \
  --run-label pool-150b \
  --format all \
  --representative-manifest /path/to/representative.parquet \
  --stratified-manifest /path/to/stratified_dir
```

The paper-facing alias `data-attribution-paper-figures` still points to the same implementation for the current paper cycle.

## Dolma subset

For Dolma, the simplest way to keep the run bounded is to use a sliced split (example sizes shown):

```bash
uv run bergson build runs/dolma_small \
  --model EleutherAI/pythia-14m \
  --dataset allenai/dolma3_mix-150B-1025 \
  --subset default \
  --split "train[:1024]" \
  --truncation
```

## SocialIQA dataset

I can construct prompt strings from `allenai/social_i_qa` on Hugging Face instead of relying on the committed v1.4 archive.
If we want a labeled split for "correct answer" completions, we should start with the train/validation splits that contain `correct`.
The SocialIQA manifest and export helpers now emit a `completion` field with the correct choice when labels are available so Bergson
builds can point `--prompt_column prompt` and `--completion_column completion` at the generated JSONL.

## Repo smoke: SocialIQA on Dolma (GPU only)

This is the bounded smoke path we run manually (not in CI). Defaults: Dolma slice `train[:4509]`, SocialIQA sample 100 `train` queries, model `allenai/Olmo-3-1025-7B`, projection_dim 16, bf16, truncation on, skip_preconditioners on. Outputs go to `runs/attributions/<run_id>` by default.

1. Precache (requires `HF_TOKEN`):  
   `uv run data-attribution-precache --dataset-repo-id allenai/dolma3_mix-150B-1025 --allow-full-download --snapshot`
2. Dolma index:  
   `uv run data-attribution-dolma-train --train-slice 4509 --model-id allenai/Olmo-3-1025-7B`
3. Dolma doc text sidecar (aligns doc_id to text/source_id):  
   `uv run data-attribution-dolma-metadata --train-slice 4509 --model-id allenai/Olmo-3-1025-7B --run-root runs/dolma_train_allenai-Olmo-3-1025-7B_train4509_proj16_bf16`
4. SocialIQA queries + query index:  
   `uv run data-attribution-socialiqa-queries --split train --sample-size 100 --model-id allenai/Olmo-3-1025-7B`
5. Attribution with doc metadata:  
   `uv run data-attribution run --query-gradients runs/socialiqa_allenai-Olmo-3-1025-7B_train_proj16_bf16 --index-path runs/dolma_train_allenai-Olmo-3-1025-7B_train4509_proj16_bf16 --metadata-path runs/dolma_train_allenai-Olmo-3-1025-7B_train4509_proj16_bf16/metadata.parquet --metadata-join-keys text source_id`

Optional analysis helper: join doc text and query prompts into the attribution parquet:  
`uv run data-attribution-enrich --attributions runs/attributions/<run_id>/attributions.parquet --doc-metadata runs/dolma_train_allenai-Olmo-3-1025-7B_train4509_proj16_bf16/metadata.parquet --query-index runs/socialiqa_allenai-Olmo-3-1025-7B_train_proj16_bf16`

See `notes/BERGSON_COOKBOOK.md` for details and troubleshooting.

## Plan

`PLAN.md` tracks the current approach. The current direction is to keep a quickstart-first workflow and avoid a custom runner until the basic Bergson pipeline is working end to end.

## Hugging Face auth (rate limits)

Some Hugging Face endpoints will rate limit unauthenticated requests.
I keep secrets in `.env.secret` (git-ignored) and load them in scripts that need them.
The repo defaults Hugging Face caches to `./.hf_cache`; smoke helpers auto-source `.env.secret` and set `HF_HOME`, `HF_DATASETS_CACHE`, `HUGGINGFACE_HUB_CACHE`, and `TRANSFORMERS_CACHE` to that path.

1. Create `.env.secret` by copying `.env.secret.example`.
2. Set `HF_TOKEN` to a read-only Hugging Face token.
3. Re-run `data-attribution-precache --snapshot --allow-full-download --dataset-repo-id allenai/dolma3_mix-150B-1025`.
4. If you hit 429 rate limits, retry with `--max-workers 1 --retry-attempts 10`.

## Precache allowlists

The precache CLI only accepts predefined model and dataset ids.

Allowed model ids:
- `allenai/Olmo-3-1025-7B`
- `allenai/Olmo-3-1125-32B`

Allowed dataset ids:
- `allenai/dolma3_mix-150B-1025` (subset: `default`)

## Ingest policy

`data-attribution ingest` only reads from local Dolma shards.
Provide `--cache-dir` or `--download-dir` pointing at the prepared shards.
Use `data-attribution-precache` to fetch shards before ingesting.