#! /bin/bash
model_name="gpt2"

BASE_CODE_PATH="/home/naie/work/minillm"
export PYTHONPATH=${BASE_CODE_PATH}

big_model_path="/opt/dpcvol/datasets/8625883998351850434/ckpt/minillm/learngene/llama3-8b/stepwise_distill/train/gpt2_m-l-best_iter/bs8-lr5e-06-G2-N8-NN1-lm1-len512/pe4_rs0.5_nr256_ln_sr_tm0.2/5000/pytorch_model.bin"
small_model_path="/opt/dpcvol/datasets/8625883998351850434/ckpt/minillm/learngene/llama3-8b/stepwise_distill/train/gpt2_b-m-best_iter/bs8-lr5e-06-G2-N8-NN1-lm1-len512/pe4_rs0.5_nr256_ln_sr_tm0.2/5000/pytorch_model.bin"

# 138M
middle_model_path="/home/naie/work/minillm/configs/learngene/gpt2/gpt2-NEmbed_768_NHead_12_NLayer_14.json"
output_dir="/opt/dpcvol/datasets/8625883998351850434/ckpt/minillm/learngene/llama3-8b/Build_descendant_model/NEmbed_768_NHead_12_NLayer_14/sque_init/"
mkdir -p "$output_dir"

# for alpha in 0.1 0.2 0.3 0.4 0.6 0.5 0.6 0.7 0.8 1.0
#   do
#     mkdir -p "$output_dir"
#     python ${BASE_CODE_PATH}/learngene_init.py \
#       --big_model_path "$big_model_path" \
#       --small_model_path "$small_model_path" \
#       --model_base_config "$middle_model_path" \
#       --embeding_dim 768 \
#       --head_num 12 \
#       --layer_num 14 \
#       --alpha $alpha \
#       --output_type huggingface \
#       --output_dir "$output_dir"
#   done

# # 220M
# middle_model_path="/home/naie/work/minillm/configs/learngene/gpt2/gpt2-NEmbed_896_NHead_14_NLayer_18.json"
# output_dir="/opt/dpcvol/datasets/8625883998351850434/ckpt/minillm/learngene/llama3-8b/Build_descendant_model/NEmbed_896_NHead_14_NLayer_18/sque_init/"
# mkdir -p "$output_dir"

# for alpha in 0.1 0.2 0.3 0.4 0.6 0.5 0.6 0.7 0.8 1.0
#   do
#     mkdir -p "$output_dir"
#     python ${BASE_CODE_PATH}/learngene_init.py \
#       --big_model_path "$big_model_path" \
#       --small_model_path "$small_model_path" \
#       --model_base_config "$middle_model_path" \
#       --embeding_dim 896 \
#       --head_num 14 \
#       --layer_num 18 \
#       --alpha $alpha \
#       --output_type huggingface \
#       --output_dir "$output_dir"
#   done


# 277M
middle_model_path="/home/naie/work/minillm/configs/learngene/gpt2/gpt2-NEmbed_896_NHead_14_NLayer_24.json"
output_dir="/opt/dpcvol/datasets/8625883998351850434/ckpt/minillm/learngene/llama3-8b/Build_descendant_model/NEmbed_896_NHead_14_NLayer_24/sque_init/"
mkdir -p "$output_dir"

for alpha in 0.9 # 0.1 0.2 0.3 0.4 0.6 0.5 0.6 0.7 0.8 1.0
  do
    mkdir -p "$output_dir"
    python ${BASE_CODE_PATH}/learngene_init.py \
      --big_model_path "$big_model_path" \
      --small_model_path "$small_model_path" \
      --model_base_config "$middle_model_path" \
      --embeding_dim 896 \
      --head_num 14 \
      --layer_num 24 \
      --alpha $alpha \
      --output_type huggingface \
      --output_dir "$output_dir"
  done