#!/bin/bash
# export WANDB_MODE=disabled
# export CUDA_DEVICE_ORDER=PCI_BUS_ID
# export CUDA_VISIBLE_DEVICES=1,2,3
set -e
set -u

training_data=../../../data/train/MATH_Hard.jsonl
split=train

model_dirs=(
    /path/to/Qwen2.5-7B
    /path/to/Qwen2.5-3B
    /path/to/Qwen2.5-1.5B
)

training_configs=(
    ../args/Qwen2.5-7B-SFT.json
    ../args/Qwen2.5-3B-SFT.json
    ../args/Qwen2.5-1.5B-SFT.json
)


for i in "${!model_dirs[@]}"; do
    model_dir=${model_dirs[$i]}
    training_config=${training_configs[$i]}
    python ../sft.py --model_dir $model_dir \
                    --training_data $training_data \
                    --split $split \
                    --training_config $training_config \
                    --input_key problem \
                    --output_key solution \
                    --model_type qwen2.5-base \
                    --add_system_prompt 
    echo "Done training for model dir: $model_dir"
done

echo "Done training"
