HCAI-Lab/w2-consensus-deepdive-unlearning-artifacts/social-data-attribution-w2 - Storage Bucket

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 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:

sudo apt-get update && sudo apt-get install -y build-essential g++

3. Then setup the Python environment via uv:

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:

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 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 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 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:

sbatch scripts/slurm/eda/sidecar_eda_benchmark_phoenix.sbatch

Run a chunked Phoenix array from an existing manifest:

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:

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:

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

Run the full current completed-sidecar set:

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:

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

Query the index interactively:

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:

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:

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:

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:

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.

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

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. 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:

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):

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.