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

#!/usr/bin/env bash
#SBATCH --job-name=clean_pool
#SBATCH --account=coc
#SBATCH --partition=ice-cpu
#SBATCH --cpus-per-task=8
#SBATCH --mem=64G
#SBATCH --time=18:00:00
#SBATCH --output=logs/clean_pool/%x_%j.out
#SBATCH --error=logs/clean_pool/%x_%j.err

# Clean the 150B pool sample: stream from HF (or local), drop blank/short docs.
#
# This script avoids `uv run` entirely — the olmes path dependency in
# pyproject.toml makes uv sync fail on compute nodes.  Instead we activate
# the pre-built .venv and set PYTHONPATH=src so Python finds the project
# packages directly.
#
# HF download uses an inline Python snippet instead of huggingface-cli,
# because system pip is too old to install a modern huggingface-hub[cli].

set -euo pipefail

REPO_DIR="${SLURM_SUBMIT_DIR:-$(cd "$(dirname "$0")/../../.." && pwd)}"
cd "$REPO_DIR"
mkdir -p logs/clean_pool

# ── Activate the pre-built venv ──────────────────────────────────────
if [ -f .venv/bin/activate ]; then
    source .venv/bin/activate
else
    echo "ERROR: .venv not found at $REPO_DIR/.venv"
    echo "Create it from an interactive session first:"
    echo "  python3 -m venv .venv && .venv/bin/activate && pip install zstandard huggingface_hub"
    exit 1
fi

export PYTHONPATH="${REPO_DIR}/src${PYTHONPATH:+:$PYTHONPATH}"

# ── Configuration ────────────────────────────────────────────────────
HF_REPO="${HF_REPO:-HCAI-Lab/archive-dolma3-pool-150b}"
HF_FILE="${HF_FILE:-archive-dolma3-pool-150b.jsonl.zst}"
MIN_WORDS="${MIN_WORDS:-3}"
OUTFILE="${POOL_OUTPUT:-$HOME/scratch/archive-dolma3-pool-150b/archive-dolma3-pool-150b-cleaned.jsonl.zst}"

# Set USE_LOCAL=1 to skip download and read from a local file.
USE_LOCAL="${USE_LOCAL:-0}"
INFILE="${POOL_INPUT:-$HOME/scratch/archive-dolma3-pool-150b/archive-dolma3-pool-150b.jsonl.zst}"

# HF authentication
export HF_TOKEN="${HF_TOKEN:-$(cat ~/.hf_token 2>/dev/null || true)}"

echo "=============================================="
echo "A-02  Pool Sample Cleaning"
echo "=============================================="
echo "Repo dir:   $REPO_DIR"
echo "Output:     $OUTFILE"
echo "Min words:  $MIN_WORDS"
echo "Start:      $(date)"

# ── Download (unless USE_LOCAL=1) ────────────────────────────────────
if [ "$USE_LOCAL" = "1" ]; then
    echo "Mode:       LOCAL ($INFILE)"
    if [ ! -f "$INFILE" ]; then
        echo "ERROR: Local input file not found: $INFILE"
        exit 1
    fi
else
    # Download to $TMPDIR (node-local, ~1.7 TB) to avoid scratch quota pressure.
    DOWNLOAD_DIR="${TMPDIR:-/tmp}/dolma3_pool_download"
    mkdir -p "$DOWNLOAD_DIR"

    echo "Mode:       HF DOWNLOAD → $DOWNLOAD_DIR"
    echo "Downloading $HF_REPO / $HF_FILE ..."

    python3 -c "
from huggingface_hub import hf_hub_download
import os
path = hf_hub_download(
    repo_id='${HF_REPO}',
    filename='${HF_FILE}',
    repo_type='dataset',
    local_dir='${DOWNLOAD_DIR}',
    token=os.environ.get('HF_TOKEN'),
)
print(f'Downloaded to: {path}')
"
    INFILE="${DOWNLOAD_DIR}/${HF_FILE}"
    echo "Download complete: $(ls -lh "$INFILE" | awk '{print $5}')"
fi

# ── Run cleaning ─────────────────────────────────────────────────────
echo "Cleaning:   $INFILE → $OUTFILE"
mkdir -p "$(dirname "$OUTFILE")"

python3 -m dolma.remove_white_spaces \
    --input "$INFILE" \
    --output "$OUTFILE" \
    --min-words "$MIN_WORDS" \
    --compress \
    --verbose

echo "=============================================="
echo "Done:       $(date)"
echo "Output:     $(ls -lh "$OUTFILE" | awk '{print $5}')"
echo "=============================================="