#!/bin/bash
#SBATCH --job-name=ar-tok
#SBATCH --partition=dev
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=96
#SBATCH --gpus-per-node=0
#SBATCH --time=24:00:00
#SBATCH --exclude=ip-10-0-136-14,ip-10-0-136-94,ip-10-0-136-253,ip-10-0-138-178,ip-10-0-139-113,ip-10-0-140-247
#SBATCH --output=logs/tok_%x_%j.out
#SBATCH --error=logs/tok_%x_%j.err

set -euo pipefail

# --- env ---
cd /fsx/ajay-stu/AR
source .venv/bin/activate
export PYTHONPATH=$PWD:${PYTHONPATH:-}
export TOKENIZERS_PARALLELISM=false
unset HF_TOKEN; export HF_TOKEN_PATH=/fsx/ajay-stu/hf_cache/token
export HF_HOME=/fsx/ajay-stu/ar-data/hf_home

# --- required inputs (pass via --export) ---
: "${SOURCE:?SOURCE must be set (fineweb|stack|finemath)}"
: "${TARGET:?TARGET must be set (token count, e.g. 130000000000)}"

OUT_DIR=${OUT_DIR:-/fsx/ajay-stu/ar-data/bin}

echo "=== tokenize source=$SOURCE target=$TARGET out=$OUT_DIR ==="
echo "=== free space before: $(df -h /fsx | tail -1) ==="

python -m ar.data.prepare_data \
    --source "$SOURCE" \
    --out_dir "$OUT_DIR" \
    --target_tokens "$TARGET" \
    --skip_docs "${SKIP_DOCS:-0}" --data_files "${DATA_FILES:-}"

# python -m ar.data.prepare_data \
#     --source "$SOURCE" \
#     --out_dir "$OUT_DIR" \
#     --target_tokens "$TARGET"

echo "=== free space after: $(df -h /fsx | tail -1) ==="