#!/bin/bash
# Parse the model argument
while [[ $# -gt 0 ]]; do
    case $1 in
        --model)
            model="$2"
            shift 2
            ;;
        *)
            echo "Unknown argument: $1"
            exit 1
            ;;
    esac
done

# Set the model name based on the provided argument
case $model in
    phi3)
        model_name="microsoft/Phi-3-mini-4k-instruct"
        ;;
    llama3c)
        model_name="meta-llama/Meta-Llama-3-8B"
        ;;
    llama3)
        model_name="meta-llama/Meta-Llama-3-8B-Instruct"
        ;;
    llama7b)
        model_name="meta-llama/Llama-2-7b-chat-hf"
        ;;
    llama13b)
        model_name="meta-llama/Llama-2-13b-chat-hf"
        per_device_train_batch_size=1
        gradient_accumulation_steps=16
        ;;
    llama3_70b)
        model_name="meta-llama/Meta-Llama-3-70B-Instruct"
        per_device_train_batch_size=1
        gradient_accumulation_steps=16
        ;;
    mistral)
        model_name="mistralai/Mistral-7B-Instruct-v0.2"
        ;;
    *)
        echo "Invalid model: $model"
        exit 1
        ;;
esac

# Change directory and create the results directory
output_root_dir=./$model/
mkdir -p $output_root_dir

# Run the Python script with parameterized arguments and capture the output
logfile=$output_root_dir/finetune-lora.log
> $logfile
accelerate launch --config_file /datadrive/huggingface/accelerate/default_config.yaml \
    finetune-lora.py \
    --model_name_or_path $model_name  \
    --output_dir $output_root_dir/finetune-lora-small \
    --data_path chatdoc \
    --num_train_epochs 3 \
    --per_device_train_batch_size 8 \
    --lora_r 2 \
    --lora_alpha 4 \
    --learning_rate 1e-3 \
    --warmup_steps 500 \
    --eval_strategy 'epoch' \
    --save_strategy 'no' \
    --save_total_limit 3 \
    --fp16 True \
    --lr_scheduler_type 'cosine' \
    --gradient_accumulation_steps 4 \
    --remove_unused_columns False \
    | tee $logfile
# python alpaca-lora.py \
#     --output_dir /datadrive2/fingerprinting/lora/$model/alpaca-lora \
#     --model $model_name \
#     | tee $logfile