set -eux
LLM_RECIPES_DIR=/code/llm-recipes

CHECKPOINTS=(
    /work/models/additiona_trained/mistral-llm-recipes-mistral-ja-v1/iter_0019200/model.pt
    /work/models/additiona_trained/mistral-llm-recipes-mistral-ja-v1/iter_0019400/model.pt
    /work/models/additiona_trained/mistral-llm-recipes-mistral-ja-v1/iter_0019600/model.pt
    /work/models/additiona_trained/mistral-llm-recipes-mistral-ja-v1/iter_0019800/model.pt
    /work/models/additiona_trained/mistral-llm-recipes-mistral-ja-v1/iter_0020000/model.pt
)
# Concatenate the checkpoints
INPUTS=${CHECKPOINTS[@]}

OUTPUT_DIR=/work/models/additiona_trained/mistral-llm-recipes-mistral-ja-v1/averaged/averaged-19200-20000.pt
mkdir -p $(dirname $OUTPUT_DIR)

python $LLM_RECIPES_DIR/tools/merge_checkpoints.py \
       --inputs $INPUTS \
       --output $OUTPUT_DIR