source /mnt/shared-storage-user/caipengxiang/workspace/ChemBOMAS/train_regression/env_volcengine.sh

# suzuki_50
# arylation
# buchwald_Cc1ccc(Nc2ccc(C(F)(F)F)cc2)cc1.csv
# buchwald_Cc1ccc(Nc2ccccn2)cc1.csv
# buchwald_Cc1ccc(Nc2cccnc2)cc1.csv
# buchwald_CCc1ccc(Nc2ccc(C)cc2)cc1.csv
# buchwald_COc1ccc(Nc2ccc(C)cc2)cc1.csv

# grouped_exp

# 读取命令行参数
DATA_NAME="grouped_exp" # 存放数据集的名字

# Pretrained
BASE_MODEL_PATH="/mnt/shared-storage-user/caipengxiang/H200-share/models/share/step1_llama3_8b_0916_yearly_pistachio_ep3"

# LLaMA-3.1 8B
# BASE_MODEL_PATH="/mnt/shared-storage-user/caipengxiang/H200-ai4chem/Llama-3.1-8B"

# Qwen3-8B
# BASE_MODEL_PATH="/mnt/shared-storage-user/caipengxiang/H200-ai4chem/ChemBOMAS_Models/Qwen3-8B"

# GLM-4-9B
# BASE_MODEL_PATH="/mnt/shared-storage-user/caipengxiang/H200-ai4chem/ChemBOMAS_Models/GLM-4-9B"

latest="Now it is ${BASE_MODEL_PATH}"

# 检查所需路径是否都存在
PROJECT_DIR="/mnt/shared-storage-user/caipengxiang/workspace/ChemBOMAS"
SAVE_RESULTS_DIR="/mnt/shared-storage-user/caipengxiang/H200-ai4chem/Exp_AB_results"

NOTES="cluster_config/method_notes.txt"

required_paths=(
    "${BASE_MODEL_PATH}"
    "${PROJECT_DIR}/train_regression/data4regression"
    "yield_ft_ds_config.json"
    "${SAVE_RESULTS_DIR}"
    "${NOTES}"
)

echo $latest

for path in "${required_paths[@]}"; do
    if [ ! -e "$path" ]; then
        echo "Error: Required path not found: $path"
        exit 1
    fi
done

export TOKENIZERS_PARALLELISM=false 

deepspeed --num_nodes=${NNODES} \
            --num_gpus=${GPUS} \
            --master_addr=${MASTER_ADDR} \
            --master_port=${MASTER_PORT} \
            --node_rank=${RANK} \
            train_diff_base_model.py \
            --pretrained_model_path ${BASE_MODEL_PATH} \
            --lora_adapter_path "None" \
            --yield_predictor_path "None" \
            --num_epoch 100 \
            --lr 1e-4 \
            --data_path "${PROJECT_DIR}/train_regression/data4regression" \
            --data_name "${DATA_NAME}" \
            --per_device_train_batch_size 24 \
            --save_root "${SAVE_RESULTS_DIR}/train_regression/saved_models/${DATA_NAME}" \
            --gradient_accumulation_steps 1 \
            --use_lora 1 \
            --log_file "exp_train_${DATA_NAME}.log" \
            --deepspeed_config "yield_ft_ds_config.json" \
            --mlp_lr_multiplier 1 \
            --save_interval 30 \
            --save_lora_adapter 1 \
            --save_predictor 1 \
            --lora_adapter_save_path "${SAVE_RESULTS_DIR}/share/llama-3.1-8B/clustered/${DATA_NAME}" \
            --run_test 0 \
            --eval_save_ckpt 0 \
            --pooling_method "last_token" \
            --cluster_config "cluster_config/cluster_config.yaml" \
            --wandb_offline 1 \
            --load_by_torch 1 \
            --latest_method_notes "${latest}" \
            --use_cluster 0