social-data-attribution-w2/scripts/slurm/sampling/stratified_sample.sh

#!/usr/bin/env bash
#SBATCH --job-name=stratified_sample
#SBATCH --account=cse
#SBATCH --qos=coc-ice
#SBATCH --cpus-per-task=8
#SBATCH --mem=64G
#SBATCH --time=04:00:00
#SBATCH --output=logs/sample/%x_%j.out
#SBATCH --error=logs/sample/%x_%j.err

set -euo pipefail
mkdir -p logs/sample
cd "${SLURM_SUBMIT_DIR:-$(pwd)}"

INPUT_MANIFEST="${INPUT_MANIFEST:-}"
if [[ -n "$INPUT_MANIFEST" ]]; then
  TARGET_DOCS_PER_BIN="${TARGET_DOCS_PER_BIN:-100000}"
  uv run data-attribution-manifest-sample \
    --input "$INPUT_MANIFEST" \
    --target-docs-per-bin "$TARGET_DOCS_PER_BIN"
else
  DUMMY_DOCS="${DUMMY_DOCS:-10000000}"
  TARGET_DOCS_PER_BIN="${TARGET_DOCS_PER_BIN:-1000}"
  uv run data-attribution-manifest-sample \
    --dummy \
    --n-dummy-docs "$DUMMY_DOCS" \
    --target-docs-per-bin "$TARGET_DOCS_PER_BIN"
fi

echo "Done: $(date)"