wehub-resource-sync
81 lines
2.3 KiB
Bash
Executable File
81 lines
2.3 KiB
Bash
Executable File
#!/bin/bash
|
||
# Batch runner for generate_query_pairs.py
|
||
#
|
||
# Divides the page pool into non-overlapping batches and runs one at a time,
|
||
# skipping any batch whose output file already exists.
|
||
#
|
||
# Usage:
|
||
# bash run_generate_batches.sh <tiles_dir> <start> <end> [model] [extra args...]
|
||
#
|
||
# Examples:
|
||
# # Generate batches 0–99 with default Gemini Pro:
|
||
# bash run_generate_batches.sh /opt/dlami/nvme/kiwix_tiles 0 99
|
||
#
|
||
# # Use Flash Lite for cheaper/faster generation:
|
||
# bash run_generate_batches.sh /opt/dlami/nvme/kiwix_tiles 0 99 gemini-3.1-flash-lite-preview
|
||
#
|
||
# # Filter by page chunk count:
|
||
# bash run_generate_batches.sh /opt/dlami/nvme/kiwix_tiles 0 50 gemini-3.1-flash-lite-preview \
|
||
# --postfilter-min-page-chunks 21 --postfilter-max-page-chunks 50
|
||
|
||
set -euo pipefail
|
||
|
||
TILES_DIR=${1:?Usage: run_generate_batches.sh <tiles_dir> <start> <end> [model] [extra args...]}
|
||
START=${2:-0}
|
||
END=${3:-9}
|
||
MODEL=${4:-gemini-3.1-pro-preview}
|
||
if [ "$#" -ge 4 ]; then
|
||
shift 4
|
||
else
|
||
shift "$#"
|
||
fi
|
||
EXTRA_ARGS=("$@")
|
||
|
||
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
|
||
OUTPUT_DIR="$SCRIPT_DIR/batches"
|
||
mkdir -p "$OUTPUT_DIR"
|
||
|
||
TOTAL_BATCHES=1000
|
||
PAGES_PER_BATCH=2000
|
||
|
||
echo "Tiles dir: $TILES_DIR"
|
||
echo "Batches: $START to $END (of $TOTAL_BATCHES total)"
|
||
echo "Model: $MODEL"
|
||
echo "Extra args: ${EXTRA_ARGS[*]:-none}"
|
||
echo "Output dir: $OUTPUT_DIR"
|
||
echo "Started: $(date)"
|
||
echo ""
|
||
|
||
for i in $(seq "$START" "$END"); do
|
||
BATCH=$(printf "%03d" "$i")
|
||
OUT="$OUTPUT_DIR/batch_${BATCH}.jsonl"
|
||
LOG="$OUTPUT_DIR/batch_${BATCH}.log"
|
||
|
||
if [ -f "$OUT" ]; then
|
||
echo "[$BATCH] exists ($(wc -l < "$OUT") pairs), skipping."
|
||
continue
|
||
fi
|
||
|
||
echo "[$BATCH] start $(date)"
|
||
uv run python "$SCRIPT_DIR/generate_query_pairs.py" \
|
||
--tiles-dir "$TILES_DIR" \
|
||
--batch-index "$i" \
|
||
--total-batches "$TOTAL_BATCHES" \
|
||
--num-pages "$PAGES_PER_BATCH" \
|
||
--output "$OUT" \
|
||
--seed "$i" \
|
||
--model "$MODEL" \
|
||
"${EXTRA_ARGS[@]}" \
|
||
2>&1 | tee "$LOG"
|
||
|
||
echo "[$BATCH] done $(wc -l < "$OUT" 2>/dev/null || echo 0) pairs | $(date)"
|
||
echo ""
|
||
done
|
||
|
||
echo "All done | $(date)"
|
||
|
||
# Merge all batch files
|
||
MERGED="$OUTPUT_DIR/all_pairs.jsonl"
|
||
cat "$OUTPUT_DIR"/batch_*.jsonl > "$MERGED" 2>/dev/null || true
|
||
echo "Merged: $(wc -l < "$MERGED" 2>/dev/null || echo 0) pairs total → $MERGED"
|