#!/bin/bash
#SBATCH --account=bcmz-delta-gpu
#SBATCH --partition=gpuA40x4
#SBATCH --time=48:00:00
#SBATCH --nodes=1
#SBATCH --gpus-per-node=1
#SBATCH --ntasks=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=1
#SBATCH --mem=200G
#SBATCH -J job-delta-batch
#SBATCH -o job-delta-batch.o%j
#SBATCH -e job-delta-batch.e%j

# Load env variables
module load cuda
conda deactivate
conda activate llm


nvidia-smi
# python3 main.py --b -fa --mode ars --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model meta-llama/Llama-2-13b-chat-hf --dataset openai_humaneval 
# python3 main.py --b -fa --mode ars --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model meta-llama/Llama-2-13b-chat-hf --dataset xsum 
# python3 main.py --b -fa --mode ars --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model meta-llama/Llama-2-13b-chat-hf --dataset gsm8k 
# python3 main.py --b -fa --mode ars --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model meta-llama/Llama-2-13b-chat-hf --dataset finance-alpaca 

for gamma in 1 2 3 4 5 6 7 8
do 
    echo "Running with gamma: $gamma"
    # # SPS
    # python3 main.py --b -fa --gamma $gamma --mode sps --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model meta-llama/Llama-2-13b-chat-hf --dataset openai_humaneval 
    # python3 main.py --b -fa --gamma $gamma --mode sps --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model meta-llama/Llama-2-13b-chat-hf --dataset xsum 
    # python3 main.py --b -fa --gamma $gamma --mode sps --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model meta-llama/Llama-2-13b-chat-hf --dataset gsm8k 
    # python3 main.py --b -fa --gamma $gamma --mode sps --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model meta-llama/Llama-2-13b-chat-hf --dataset finance-alpaca 

    # # USPS
    # python3 main.py --b -fa --gamma $gamma --mode usps --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model meta-llama/Llama-2-13b-chat-hf --dataset openai_humaneval 
    python3 main.py --b --gamma $gamma --mode usps --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model meta-llama/Llama-2-13b-chat-hf --dataset xsum 
    # python3 main.py --b -fa --gamma $gamma --mode usps --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model meta-llama/Llama-2-13b-chat-hf --dataset gsm8k 
    # python3 main.py --b -fa --gamma $gamma --mode usps --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model meta-llama/Llama-2-13b-chat-hf --dataset finance-alpaca 


    # DSPS
    # python3 main.py --b -fa --gamma $gamma --mode dsps --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model meta-llama/Llama-2-13b-chat-hf --dataset openai_humaneval 
    # python3 main.py --b -fa --gamma $gamma --mode dsps --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model meta-llama/Llama-2-13b-chat-hf --dataset xsum 
    # python3 main.py --b -fa --gamma $gamma --mode dsps --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model meta-llama/Llama-2-13b-chat-hf --dataset gsm8k 
    # python3 main.py --b -fa --gamma $gamma --mode dsps --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model meta-llama/Llama-2-13b-chat-hf --dataset finance-alpaca 

    # # DHSPS
    # python3 main.py --b -fa --gamma $gamma --mode dhsps --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model meta-llama/Llama-2-13b-chat-hf --dataset openai_humaneval 
    # python3 main.py --b -fa --gamma $gamma --mode dhsps --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model meta-llama/Llama-2-13b-chat-hf --dataset xsum 
    # python3 main.py --b -fa --gamma $gamma --mode dhsps --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model meta-llama/Llama-2-13b-chat-hf --dataset gsm8k 
    # python3 main.py --b -fa --gamma $gamma --mode dhsps --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model meta-llama/Llama-2-13b-chat-hf --dataset finance-alpaca
    
    # python3 main.py --b -fa --gamma $gamma --mode pp --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model meta-llama/Llama-2-13b-chat-hf --dataset openai_humaneval 
    # python3 main.py --b --gamma $gamma --mode pp --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model meta-llama/Llama-2-13b-chat-hf --dataset xsum 
    # python3 main.py --b -fa --gamma $gamma --mode pp --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model meta-llama/Llama-2-13b-chat-hf --dataset gsm8k 
    # python3 main.py --b -fa --gamma $gamma --mode pp --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model meta-llama/Llama-2-13b-chat-hf --dataset finance-alpaca 
done


draft_model="google/gemma-1.1-2b-it"
target_model="google/gemma-1.1-7b-it"

# draft_model="facebook/opt-125m"
# target_model="facebook/opt-13b" 

# draft_model="facebook/opt-125m"
# target_model="facebook/opt-30b" 

# draft_model="bigscience/bloom-560m"
# target_model="bigscience/bloom-7b1"

# draft_model="openlm-research/open_llama_3b"
# target_model="openlm-research/open_llama_7b"


# draft_model="stabilityai/stablelm-base-alpha-3b-v2"
# target_model="stabilityai/stablelm-base-alpha-7b-v2" 

draft_model="databricks/dolly-v2-3b"
target_model="databricks/dolly-v2-12b"

draft_model="EleutherAI/pythia-70m"

# draft_model="JackFram/llama-160m"
# target_model="lmsys/vicuna-13b-v1.5"


# draft_model="TinyLlama/TinyLlama-1.1B-Chat-v1.0"
# target_model="lmsys/vicuna-13b-v1.5"

# additional_args="-fa"
# additional_args="-q4"

# python3 main.py --b --mode ars --draft-model $draft_model --target-model $target_model --dataset openai_humaneval $additional_args
# python3 main.py --b --mode ars --draft-model $draft_model --target-model $target_model --dataset xsum $additional_args
# python3 main.py --b --mode ars --draft-model $draft_model --target-model $target_model --dataset gsm8k $additional_args
# python3 main.py --b --mode ars --draft-model $draft_model --target-model $target_model --dataset finance-alpaca $additional_args

for gamma in 1 2 3 4 5 6 7 8
do 
    echo "Running with gamma: $gamma"
    # SPS
    # python3 main.py --b $additional_args --gamma $gamma --mode sps --draft-model $draft_model --target-model $target_model --dataset openai_humaneval $additional_args
    # python3 main.py --b $additional_args --gamma $gamma --mode sps --draft-model $draft_model --target-model $target_model --dataset xsum 
    # python3 main.py --b $additional_args --gamma $gamma --mode sps --draft-model $draft_model --target-model $target_model --dataset gsm8k 
    # python3 main.py --b $additional_args --gamma $gamma --mode sps --draft-model $draft_model --target-model $target_model --dataset finance-alpaca 

    # # USPS
    # python3 main.py --b $additional_args --gamma $gamma --mode usps --draft-model $draft_model --target-model $target_model --dataset openai_humaneval 
    # python3 main.py --b $additional_args --gamma $gamma --mode usps --draft-model $draft_model --target-model $target_model --dataset xsum 
    # python3 main.py --b $additional_args --gamma $gamma --mode usps --draft-model $draft_model --target-model $target_model --dataset gsm8k 
    # python3 main.py --b $additional_args --gamma $gamma --mode usps --draft-model $draft_model --target-model $target_model --dataset finance-alpaca 

    # # DSPS
    # python3 main.py --b $additional_args --gamma $gamma --mode dsps --draft-model $draft_model --target-model $target_model --dataset openai_humaneval 
    # python3 main.py --b $additional_args --gamma $gamma --mode dsps --draft-model $draft_model --target-model $target_model --dataset xsum 
    # python3 main.py --b $additional_args --gamma $gamma --mode dsps --draft-model $draft_model --target-model $target_model --dataset gsm8k 
    # python3 main.py --b $additional_args --gamma $gamma --mode dsps --draft-model $draft_model --target-model $target_model --dataset finance-alpaca 

    # # DHSPS
    # python3 main.py --b $additional_args --gamma $gamma --mode dhsps --draft-model $draft_model --target-model $target_model --dataset openai_humaneval 
    # python3 main.py --b $additional_args --gamma $gamma --mode dhsps --draft-model $draft_model --target-model $target_model --dataset xsum 
    # python3 main.py --b $additional_args --gamma $gamma --mode dhsps --draft-model $draft_model --target-model $target_model --dataset gsm8k 
    # python3 main.py --b $additional_args --gamma $gamma --mode dhsps --draft-model $draft_model --target-model $target_model --dataset finance-alpaca 

    # # pp
    # python3 main.py --b $additional_args --gamma $gamma --mode pp --draft-model $draft_model --target-model $target_model --dataset openai_humaneval 
    # python3 main.py --b $additional_args --gamma $gamma --mode pp --draft-model $draft_model --target-model $target_model --dataset xsum 
    # python3 main.py --b $additional_args --gamma $gamma --mode pp --draft-model $draft_model --target-model $target_model --dataset gsm8k 
    # python3 main.py --b $additional_args --gamma $gamma --mode pp --draft-model $draft_model --target-model $target_model --dataset finance-alpaca 
done