social-data-attribution-w2/scripts/slurm/sampling/sample_data.sh

#!/usr/bin/env bash
#SBATCH --job-name=sample_data
#SBATCH --partition=coc-gpu
#SBATCH --output=logs/download/%x_%j.out
#SBATCH --error=logs/download/%x_%j.err
#SBATCH --cpus-per-task=4
#SBATCH --mem=32G
#SBATCH --time=02:00:00

set -euo pipefail

mkdir -p logs/download
cd "${SLURM_SUBMIT_DIR:-$(pwd)}"

uv run python - <<'PY'
import json
import os
from datasets import load_dataset

SEED = int(os.environ.get("SEED", 42))
N = int(os.environ.get("N", 100_000))
BUFFER = int(os.environ.get("BUFFER", 500_000))
OUT = os.environ.get("OUT", "data/dolma3_mix_sample.jsonl")

os.makedirs(os.path.dirname(OUT) or ".", exist_ok=True)

ds = load_dataset(
    "allenai/dolma3_mix-150B-1025",
    split="train",
    streaming=True,
)
ds = ds.shuffle(seed=SEED, buffer_size=BUFFER)

written = 0
with open(OUT, "w", encoding="utf-8") as f:
    for ex in ds:
        if written >= N:
            break
        f.write(json.dumps(ex, ensure_ascii=False) + "\n")
        written += 1

print(f"wrote {OUT} ({written} examples, seed={SEED}, buffer={BUFFER})")
PY

echo "Done: $(date)"