#!/bin/bash
# One-time: fresh data for the longer Stage 1 (~300B tokens).
#   English: fineweb-edu groups 007-009 -> new files fw007-fw009 (20B each)
#   Code: 5 languages continue after the documents already used (appended)
#   Math: finemath-3plus continues after the documents already used (appended)
# Run: nohup ./new_data.sh > logs/new_data.log 2>&1 &
cd /fsx/ajay-stu/AR
source .venv/bin/activate
D=/fsx/ajay-stu/ar-data/bin
P=sample/350BT

# 1) files we append to must end on a full document; if not, add one end-of-text id
#    (adding is safe while training reads the file; cutting it would not be)
python - <<'PY'
import numpy as np
EOT = 151645
for name in ["stack_python", "stack_javascript", "stack_java", "stack_cpp", "stack_typescript", "finemath3"]:
    p = f"/fsx/ajay-stu/ar-data/bin/{name}.bin"
    last = int(np.memmap(p, dtype=np.uint32, mode="r")[-1])
    if last != EOT:
        with open(p, "ab") as f:
            np.array([EOT], dtype=np.uint32).tofile(f)
        print(f"{name}: added end-of-text after a half-written document")
    else:
        print(f"{name}: ends on a full document")
PY

# 2) the Python job crashed before printing its skip value: count its documents,
#    plus 5% for the ~1% the code filter dropped, so the new part never overlaps
K=$(python -c "
import numpy as np
a = np.memmap('$D/stack_python.bin', dtype=np.uint32, mode='r'); n = 500_000_000
print(sum(int((a[i:i + n] == 151645).sum()) for i in range(0, len(a), n)))")
PYSKIP=$(( 20000 + K * 105 / 100 ))
echo "$(date -u) stack_python: $K documents in file, skip $PYSKIP"

# $1 source  $2 out dir  $3 tokens to add  $4 docs to skip  $5 data files (optional)
sub() {
  [ -z "$4" ] && { echo "$(date -u) $1: no skip value, NOT submitted"; return; }
  echo "$(date -u) $1 -> $2: add $3 tokens, skip $4 ${5:+files $5}"
  sbatch --export="ALL,SOURCE=$1,TARGET=$3,OUT_DIR=$2,SKIP_DOCS=$4,DATA_FILES=${5:-}" tokenize.sbatch
  sleep 300                                              # Hugging Face rate limit
}

for g in 007 008 009; do sub fineweb $D/fw$g 20000000000 0 "$P/${g}_*.parquet"; done
sub stack_python     $D 10000000000 "$PYSKIP"
sub stack_javascript $D  4000000000 3917625
sub stack_java       $D  4000000000 3062240
sub stack_cpp        $D  3000000000 978851
sub stack_typescript $D  2000000000 3330360
sub finemath3        $D 16000000000 4601977
echo "$(date -u) all 9 jobs submitted"