#!/bin/bash

MODELS_TO_TRAIN=(
    "Qwen2.5-7B-Instruct"
    # "Llama-3.2-3B"
    # "Llama-3.2-1B"
    
)

export WANDB_API_KEY="62c57a07add7cf80060d09b29e313990bc2fada2"

for MODEL in "${MODELS_TO_TRAIN[@]}"; do
    echo "Starting training for model: ${MODEL}"
    
    
    # MODEL_NAME_OR_PATH="/aifs4su/yaodong/wenqi/projects/first-time_safety/output_models/${MODEL}-base/slice_end"
    # Second training phase
    MODEL_NAME_OR_PATH="/aifs4su/yaodong/wenqi/models/Qwen2.5-7B-Instruct"
    TRAIN_DATASETS="/aifs4su/yaodong/wenqi/projects/first-time_safety/data_annotation/data_output/safe-o1_0403/baseline_dataset"
    OUTPUT_DIR="/aifs4su/yaodong/wenqi/projects/first-time_safety/output_models/${MODEL}_safe_thinking"
    TRAIN_TEMPLATE="Safe_thinking"
    LOG_PROJECT="safe-o1"
    
    source ./setup.sh

    deepspeed \
        --master_port ${MASTER_PORT} \
        --module align_anything.trainers.text_to_text.sft \
        --model_name_or_path ${MODEL_NAME_OR_PATH} \
        --train_datasets ${TRAIN_DATASETS} \
        --train_split train \
        --train_template ${TRAIN_TEMPLATE} \
        --output_dir ${OUTPUT_DIR} \
        --log_project ${LOG_PROJECT} \
        --per_device_train_batch_size 4 \
        --per_device_eval_batch_size 4 \
        --gradient_accumulation_steps 2 \
        --learning_rate 2e-5 \
        --epochs 3 \
        --model_max_length 16384 \


    echo "Completed second phase training for ${MODEL}"
done

echo "All model training completed!"