social-data-attribution-w2/scripts/slurm/analysis/rq4_lexical_array.sbatch

#!/usr/bin/env bash
#SBATCH --job-name=rq4-lex-worker
#SBATCH --partition=ice-cpu
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=2
#SBATCH --mem=6G
#SBATCH --time=04:00:00
#SBATCH --array=0-31%16
#SBATCH --output=/storage/ice-shared/cs7634/staff/TDA/logs/rq4_lex_%A_%a.out
#SBATCH --error=/storage/ice-shared/cs7634/staff/TDA/logs/rq4_lex_%A_%a.err

set -euo pipefail

REPO="${REPO:-/storage/ice-shared/cs7634/staff/TDA/code/rq4_lexical}"
MANIFEST="${MANIFEST:-/storage/ice-shared/cs7634/staff/TDA/soc149/working_sample_manifest.parquet}"
SHARDS_DIR="${SHARDS_DIR:-/storage/ice-shared/cs7634/staff/TDA/tracstar/shards_10k/sample_10000_docs}"
SHARD_GLOB="${SHARD_GLOB:-shard_*.jsonl}"
CHUNK_COUNT="${CHUNK_COUNT:-32}"
OUTDIR="${OUTDIR:-/storage/ice-shared/cs7634/staff/TDA/outputs/rq4_lexical_workers}"
WRITE_PER_DOC="${WRITE_PER_DOC:-1}"

WORKER_DIR="$OUTDIR/$SLURM_ARRAY_TASK_ID"
mkdir -p "$WORKER_DIR"

echo "RQ4 lexical worker $SLURM_ARRAY_TASK_ID/$CHUNK_COUNT @ $(date)"
echo "Node: $(hostname)"
echo "REPO=$REPO"
echo "MANIFEST=$MANIFEST"
echo "SHARDS_DIR=$SHARDS_DIR"
echo "WORKER_DIR=$WORKER_DIR"
echo "WRITE_PER_DOC=$WRITE_PER_DOC"

PER_DOC_FLAG=""
if [[ "$WRITE_PER_DOC" = "1" ]]; then
  PER_DOC_FLAG="--write-per-doc"
fi

PYTHON_BIN="${PYTHON_BIN:-$REPO/.venv/bin/python}"

"$PYTHON_BIN" "$REPO/scripts/analysis/rq4_lexical_toolkit_profiles.py" \
  --mode worker \
  --manifest "$MANIFEST" \
  --shards-dir "$SHARDS_DIR" \
  --shard-glob "$SHARD_GLOB" \
  --chunk-index "$SLURM_ARRAY_TASK_ID" \
  --chunk-count "$CHUNK_COUNT" \
  --worker-output-dir "$WORKER_DIR" \
  $PER_DOC_FLAG

echo "Worker $SLURM_ARRAY_TASK_ID finished @ $(date)"