#!/usr/bin/env bash
set -euo pipefail

REPO_ROOT="${REPO_ROOT:-/mnt/beegfs/yi/laoyao_2b_moe}"
IMAGE="${IMAGE:-nvcr.io/nvidia/nemo:26.06}"
SOURCE_DATA_DIRS="${SOURCE_DATA_DIRS:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701:/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/train/logic_topup_proof_pile_17b_v1_20260701}"
WORK_DIR="${WORK_DIR:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_8192_v1}"
JSONL="${JSONL:-$WORK_DIR/text.jsonl}"
OUTPUT_PREFIX="${OUTPUT_PREFIX:-$WORK_DIR/laoyao_2b_moe_8192_text_document}"
TOKENIZER_MODEL="${TOKENIZER_MODEL:-$REPO_ROOT/tokenizer/glm5.2}"
WORKERS="${WORKERS:-16}"
MAX_DOCS="${MAX_DOCS:-0}"
MIN_FREE_GB="${MIN_FREE_GB:-1000}"
KEEP_JSONL="${KEEP_JSONL:-0}"

mkdir -p "$WORK_DIR"

available_gb="$(df -BG "$WORK_DIR" | awk 'NR==2 {gsub("G", "", $4); print $4}')"
if [[ "$MAX_DOCS" == "0" && "$available_gb" -lt "$MIN_FREE_GB" ]]; then
  cat >&2 <<EOF
Refusing full preprocessing: only ${available_gb}GB free at $WORK_DIR.
Full 200B-token Megatron indexed output is expected to require hundreds of GB,
and this script also stages a JSONL export. Set MIN_FREE_GB lower only if you
have confirmed an output filesystem with enough capacity, or run a bounded
probe with MAX_DOCS.
EOF
  exit 2
fi

input_args=()
IFS=':' read -r -a source_dirs <<< "$SOURCE_DATA_DIRS"
for source_dir in "${source_dirs[@]}"; do
  input_args+=(--input "${source_dir}/*.parquet")
done

docker run --rm --ipc=host --network=host \
  --ulimit memlock=-1 --ulimit stack=67108864 \
  -v /mnt/beegfs:/mnt/beegfs \
  -w "$REPO_ROOT" \
  "$IMAGE" \
  bash -lc '
    set -euo pipefail
    python3 dataset/pretrain/scripts/export_pretrain_parquet_text_jsonl.py \
      "$@" \
      --output '"$JSONL"' \
      --max-docs '"$MAX_DOCS"'
    python3 /opt/Megatron-Bridge/3rdparty/Megatron-LM/tools/preprocess_data.py \
      --input '"$JSONL"' \
      --json-keys text \
      --tokenizer-type HuggingFaceTokenizer \
      --tokenizer-model '"$TOKENIZER_MODEL"' \
      --append-eod \
      --output-prefix '"$OUTPUT_PREFIX"' \
      --workers '"$WORKERS"'
    if [[ "'"$KEEP_JSONL"'" == "0" ]]; then
      rm -f '"$JSONL"'
    fi
    ls -lh '"${OUTPUT_PREFIX}"'*
  ' bash "${input_args[@]}"

echo "Megatron indexed dataset prefix: $OUTPUT_PREFIX"
