#! /bin/bash

#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=8
#SBATCH --mem=100G
#SBATCH -p gpu
#SBATCH --gres gpu:1
#SBATCH --partition=gpu_h100
#SBATCH --time=01-00:00:00

#SBATCH -o /scratch-shared/dwu18/cache/logs/out.emoon.%j.o
#SBATCH -o /scratch-shared/dwu18/cache/logs/out.emoon.%j.e

source activate py38cuda11
# source activate calibration

export HF_HUB_CACHE=/gpfs/work4/0/gus20642/dwu18/cache
export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python

################## MAIN ##################

LR=$1

SETTING=emoon_$LR
TEST_DATASET=enote_dataset

echo $LR

python -m llama_recipes.finetuning --use_peft --peft_method lora \
	--model_name meta-llama/Llama-2-7b-hf \
        --output_dir ./checkpoints/7B/emoon/${SETTING} \
        --dataset enote_dataset \
        --rule_names "ryjlzs0001,ryjlzs0002,ryjlzs0004,ryjlzs0005,ryjlxbs0001,ryjljws0003,ryjljws0004,ryjljws0006,ryjljws0007,ryjljws0008,ryjljws0010" \
        --batching_strategy padding \
        --num_epochs 2 \
        --lr $LR \
        --batch_size_training 1 \
        --val_batch_size 1 \
        --gradient_accumulation_steps 8 \
        --use_wandb

for EPOCH in 0; do
    BASE_SYS=results/emoon/${SETTING}-beam1/${EPOCH}
    python inference_formal.py --model_name meta-llama/Llama-2-7b-hf \
            --peft_model ./checkpoints/7B/emoon/${SETTING}/${EPOCH} \
            --dataset ${TEST_DATASET} \
            --val_batch_size 1 \
            --do_sample False \
            --output_dir ${BASE_SYS} \
            --rule_names "ryjlzs0001,ryjlzs0002" \
            --beam_size 1
done