#!/bin/bash
#SBATCH -N 1
#SBATCH -p GPU-shared
#SBATCH -t 12:00:00
#SBATCH --gpus=v100-32:1

#### --gpus=type:n


# Load env variables
module load cuda
module load cudnn
module load nvhpc

conda activate llm


nvidia-smi
# python3 main.py --b --mode ars --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model codellama/CodeLlama-13b-Python-hf --dataset openai_humaneval 
# python3 main.py --b --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 --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 --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 7
do 
    echo "Running with gamma: $gamma"
    # SPS
    # python3 main.py --b --gamma $gamma --mode sps --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model codellama/CodeLlama-13-Python-hf --dataset openai_humaneval 
    # python3 main.py --b --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 --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 --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 --gamma $gamma --mode usps --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model codellama/CodeLlama-13b-Python-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 --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 --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 --gamma $gamma --mode dsps --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model codellama/CodeLlama-13b-Python-hf --dataset openai_humaneval 
    # python3 main.py --b --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 --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 --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 --gamma $gamma --mode dhsps --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model codellama/CodeLlama-13b-Python-hf --dataset openai_humaneval 
    # python3 main.py --b --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 --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 --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 --gamma $gamma --mode pp --draft-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --target-model codellama/CodeLlama-13b-Python-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 --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 --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-2b"
target_model="google/gemma-7b"

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

# 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="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"


# python3 main.py --b --mode ars --draft-model $draft_model --target-model $target_model --dataset openai_humaneval 
# python3 main.py --b --mode ars --draft-model $draft_model --target-model $target_model --dataset xsum 
# python3 main.py --b --mode ars --draft-model $draft_model --target-model $target_model --dataset gsm8k 
# python3 main.py --b --mode ars --draft-model $draft_model --target-model $target_model --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 --gamma $gamma --mode sps --draft-model $draft_model --target-model $target_model --dataset openai_humaneval 
    python3 main.py --b --gamma $gamma --mode sps --draft-model $draft_model --target-model $target_model --dataset xsum 
    python3 main.py --b --gamma $gamma --mode sps --draft-model $draft_model --target-model $target_model --dataset gsm8k 
    python3 main.py --b --gamma $gamma --mode sps --draft-model $draft_model --target-model $target_model --dataset finance-alpaca 

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


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

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

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