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

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)"
cd "$ROOT_DIR"

export http_proxy="${http_proxy:-http://10.20.34.2:3128}"
export https_proxy="${https_proxy:-http://10.20.34.2:3128}"
export HTTP_PROXY="${HTTP_PROXY:-$http_proxy}"
export HTTPS_PROXY="${HTTPS_PROXY:-$https_proxy}"
export no_proxy="${no_proxy:-localhost,127.0.0.1,10.20.0.0/16,internal-cr,internal-cr:3443,10.20.34.1:6100}"
export NO_PROXY="${NO_PROXY:-$no_proxy}"
export HF_ENDPOINT="${HF_ENDPOINT:-https://hf-mirror.com}"

if [[ -z "${HF_TOKEN:-}" ]]; then
  echo "HF_TOKEN is required" >&2
  exit 2
fi

SEED="${SEED:-42}"
OUTPUT_ROOT="${OUTPUT_ROOT:-/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset}"
MAX_MBPS="${MAX_MBPS:-0}"

python dataset/pretrain/scripts/download_rebalanced_sources.py \
  --seed "$SEED" \
  --output-root "$OUTPUT_ROOT" \
  --max-mbps "$MAX_MBPS" \
  "$@"
