social-data-attribution-w2/scripts/slurm/dedup/extract_ids.sbatch

#!/usr/bin/env bash
#SBATCH --job-name=extract_ids
#SBATCH --array=0-497%15
#SBATCH --cpus-per-task=4
#SBATCH --mem=16G
#SBATCH --time=04:00:00
#SBATCH -C localNVMe
#SBATCH --output=logs/dedup/%x_%A_%a.out
#SBATCH --error=logs/dedup/%x_%A_%a.err

set -euo pipefail

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

if [ -f .venv/bin/activate ]; then
    source .venv/bin/activate
fi

export PYTHONPATH="${REPO_DIR}/src${PYTHONPATH:+:$PYTHONPATH}"
export HF_TOKEN="${HF_TOKEN:-$(cat ~/.hf_token 2>/dev/null || true)}"

python3 scripts/slurm/dedup/extract_ids.py \
    --dataset "${DATASET_ID:-allenai/dolma3_mix-6T-1025-7B}" \
    --manifest "${SHARD_MANIFEST:-scripts/slurm/shard_manifest_6t.txt}" \
    --output-dir "${IDS_OUT_DIR:-/storage/ice-shared/cs7634/staff/TDA/soc-90/ids}" \
    --task-id "${SLURM_ARRAY_TASK_ID}" \
    --task-count "${TASK_COUNT:-498}" \
    --shards-per-task "${SHARDS_PER_TASK:-132}" \
    --startup-delay "${STARTUP_DELAY:-120}"