#!/usr/bin/env bash
cd /ROOT_DIR
# ---------------- EDIT THESE PATHS/VALUES ----------------
PYTHON=python
CLIP_SCRIPT=./clip_inference.py

JSONL_FILE=./scenery-output.jsonl
ROOT_DIR=./tipo/generated_outputs # for each model-result directories. 
NUM_FOLDERS=200 # Number of subfolders in the root_dir (splitted since too many files in single folder is bad)

# Number of GPUs (and hence workers) you want to use
NUM_WORKERS=4

# ---------------------------------------------------------

# Create logs directory if not existing
mkdir -p logs

echo "Running CLIP inference with $NUM_WORKERS workers..."

for (( i=0; i<$NUM_WORKERS; i++ )); do
  echo "Launching worker $i..."
  export CUDA_VISIBLE_DEVICES=$i

  # Launch clip_inference.py with:
  #  1) JSONL file,
  #  2) root_dir,
  #  3) num_folders,
  #  4) total num_workers,
  #  5) worker_index,
  #  6) optional "original"
  #  7) output file with worker index (clip_scores_tipo_worker_0.jsonl)
  #
  #
  echo "$PYTHON $CLIP_SCRIPT $JSONL_FILE $ROOT_DIR $NUM_FOLDERS $NUM_WORKERS $i False clip_result/clip_scores_tipo_worker_${i}.jsonl"
  $PYTHON "$CLIP_SCRIPT" \
    "$JSONL_FILE" \
    "$ROOT_DIR" \
    "$NUM_FOLDERS" \
    "$NUM_WORKERS" \
    "$i" \
    "False" \
    "clip_result/clip_scores_tipo_worker_${i}.jsonl" \
    >> "logs/clip_worker_${i}.log" 2>&1 &

done

# Wait for all background jobs to finish
wait

echo "All CLIP inference workers have finished."
