social-data-attribution-w2/scripts/slurm/sampling/split_by_bin.sbatch

#!/usr/bin/env bash
#SBATCH --job-name=split_by_bin
#SBATCH --account=gts-mriedl3
#SBATCH --cpus-per-task=20
#SBATCH --mem=64G
#SBATCH --time=04:00:00
#SBATCH --array=0-4
#SBATCH --output=logs/split_by_bin/%A_job%a.out
#SBATCH --error=logs/split_by_bin/%A_job%a.err

# Split merged job files into per-bin (topic x format) zst files.
#
# 5 array tasks, one per job_id. Each task reads job_N.jsonl.zst and
# writes up to 576 bin files into by_bin/job_N/. A separate merge
# step (merge_bin_splits.py) combines them afterwards.
#
# Usage:
#   sbatch scripts/slurm/split_by_bin.sbatch

set -euo pipefail

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

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

INPUT_DIR="${INPUT_DIR:-stratified_data/merged_docs}"
OUTPUT_DIR="${OUTPUT_DIR:-stratified_data/by_bin}"
JOB_ID="$SLURM_ARRAY_TASK_ID"

echo "=============================================="
echo "  Split By Bin - Job $JOB_ID"
echo "=============================================="
echo "Input dir:    $INPUT_DIR"
echo "Output dir:   $OUTPUT_DIR"
echo "Start:        $(date)"
echo "Hostname:     $(hostname)"
echo "=============================================="

python3 scripts/sampling/split_by_bin.py \
    --input-dir "$INPUT_DIR" \
    --output-dir "$OUTPUT_DIR" \
    --job-id "$JOB_ID"

echo "=============================================="
echo "  Job $JOB_ID complete"
echo "=============================================="
echo "Done: $(date)"
echo "=============================================="