social-data-attribution-w2/scripts/slurm/upload/upload_stratified_docs.sbatch

#!/usr/bin/env bash
#SBATCH --job-name=upload_stratified
#SBATCH --account=gts-mriedl3
#SBATCH --cpus-per-task=4
#SBATCH --mem=8G
#SBATCH --time=04:00:00
#SBATCH --output=logs/upload_stratified_docs/%A.out
#SBATCH --error=logs/upload_stratified_docs/%A.err

# Delete per-worker artifacts from HCAI-Lab/archive-dolma3-pool-stratified,
# then upload any combined job_N.jsonl.zst files not yet present.
#
# Usage:
#   sbatch scripts/slurm/upload_stratified_docs.sbatch

set -euo pipefail

REPO_DIR="${SLURM_SUBMIT_DIR:-$PWD}"
cd "$REPO_DIR"
mkdir -p logs/upload_stratified_docs

if [ -f .venv/bin/activate ]; then
    source .venv/bin/activate
else
    echo "ERROR: .venv not found at $REPO_DIR/.venv"
    exit 1
fi

export HF_TOKEN="${HF_TOKEN:-$(cat ~/.hf_token 2>/dev/null || true)}"
export HF_HUB_DISABLE_XET=1
export PYTHONUNBUFFERED=1

INPUT_DIR="${INPUT_DIR:-stratified_data/merged_docs}"
REPO_ID="${REPO_ID:-HCAI-Lab/archive-dolma3-pool-stratified}"

echo "=============================================="
echo "  Upload Stratified Docs"
echo "=============================================="
echo "Input dir:    $INPUT_DIR"
echo "Repo:         $REPO_ID"
echo "Start:        $(date)"
echo "Hostname:     $(hostname)"
echo "=============================================="

python3 scripts/sampling/upload_stratified_docs.py \
    --input-dir "$INPUT_DIR" \
    --repo-id "$REPO_ID"

echo "=============================================="
echo "  Upload complete"
echo "=============================================="
echo "Done: $(date)"
echo "=============================================="