#!/bin/bash

MODEL_SIZE=${1:-"1B"}
# path
project_dir=/path/paper_code

if [ "$MODEL_SIZE" = "1B" ]; then
    export CUDA_VISIBLE_DEVICES=0
    gpu_count=1
    model_source="meta-llama/Llama-3.2-1B"
elif [ "$MODEL_SIZE" = "3B" ]; then
    export CUDA_VISIBLE_DEVICES=0,1
    gpu_count=2
    model_source="meta-llama/Llama-3.2-3B"
elif [ "$MODEL_SIZE" = "8B" ]; then
    export CUDA_VISIBLE_DEVICES=0,1,2,3
    gpu_count=4
    model_source="meta-llama/Llama-3.1-8B"
else
    exit 1
fi

# Token settings
token_select_strategy="full_token"
use_forgetting=False
# config
dataset="training_dataset/tulu_10k_sample.json"
rho=0.0
max_length=2048
port_number=29509
batch_per_gpu=2

# fine-tuning
echo "start finetuning..."
bash_files/finetune.sh "$model_source" "$dataset" "$max_length" "$batch_per_gpu" "$gpu_count" "$model_source" "$project_dir" "$rho" "$port_number" "$token_select_strategy" "$use_forgetting"