
export CUDA_VISIBLE_DEVICES=0

models=(
    
    meta-llama/Llama-3.1-8B-Instruct
    mistralai/Mistral-7B-Instruct-v0.3
    Qwen/Qwen2.5-7B-Instruct

)
max_capacity_prompts=(
    128
    256
    512
    1024
    2048
)
methods=(
    "fullkv"
    "compresskv"
    "cakekv"
    "snapkv"
    "pyramidkv"
    "streamingllm"
    "headkv"
    "adakv"
    "headcompresskv"
    "adacompresskv"
)

for model in "${models[@]}"
do
    head_idx_path=""
    layer_utility_path=""

    for method in "${methods[@]}"
    do
        if [[ "$method" == "fullkv" ]]; then
            echo "Running model=$model method=$method"
            python -m longbench.pred_compresskv \
                --model $model \
                --method $method 
        else
            for budget in "${max_capacity_prompts[@]}"
            do
                echo "Running model=$model method=$method budget=$budget "
                if [[ "$method" == "cakekv" ]]; then
                    tau1=""
                    tau2=""
                    if [[ "$budget" == "512" || "$budget" == "1024" || "$budget" == "2048" ]]; then
                        if [[ "$model" == *Qwen2.5-7B-Instruct* ]]; then
                            tau1="--tau1 1.5"
                            tau2="--tau2 0.6"
                        elif [[ "$model" == *Mistral-Nemo-Instruct-2407* ]]; then
                            tau1="--tau1 1.0"
                            tau2="--tau2 1.0"
                        elif [[ "$model" == *"Llama-3.1-8B-Instruct"* ]]; then
                            tau1="--tau1 1.6"
                            tau2="--tau2 0.4"
                        elif [[ "$model" == *"Mistral-7B-Instruct-v"* ]]; then
                            tau1="--tau1 0.8"
                            tau2="--tau2 0.5"
                        fi
                    elif [[ "$budget" == "128" || "$budget" == "256" ]]; then
                        if [[ "$model" == *Qwen2.5-7B-Instruct* ]]; then
                            tau1="--tau1 0.4"
                            tau2="--tau2 0.6"
                        elif [[ "$model" == *Mistral-Nemo-Instruct-2407* ]]; then
                            tau1="--tau1 1.0"
                            tau2="--tau2 1.0"
                        elif [[ "$model" == *"Llama-3.1-8B-Instruct"* ]]; then
                            tau1="--tau1 1.6"
                            tau2="--tau2 0.6"
                        elif [[ "$model" == *"Mistral-7B-Instruct-v"* ]]; then
                            tau1="--tau1 1.0"
                            tau2="--tau2 0.8"
                        fi
                    fi
                    python -m longbench.pred_compresskv \
                        --model $model \
                        --max_capacity_prompt $budget \
                        --method $method \
                        $tau1 $tau2
                elif [[ "$method" == "headkv" ]] || [[ "$method" == "headcompresskv" ]]; then
                    if [[ "$model" == *"Llama-3.1-8B-Instruct"* ]]; then
                        head_score_path="scores/Llama-3.1-8B-Instruct_retrieval_reasoning_heads.json"
                        head_idx_path="scores/Llama-3.1-8B-Instruct_head_idx.json"
                    elif [[ "$model" == *"Mistral-7B-Instruct-v"* ]]; then
                        head_score_path="scores/Mistral-7B-Instruct-v0.3_retrieval_reasoning_heads.json"
                        head_idx_path="scores/Mistral-7B-Instruct-v0.3_head_idx.json"
                    elif [[ "$model" == *"Qwen2.5-7B-Instruct"* ]]; then
                        head_score_path="scores/Qwen2.5-7B-Instruct_retrieval_reasoning_heads.json"
                        head_idx_path="scores/Qwen2.5-7B-Instruct_head_idx.json"
                    else
                        echo "Unknown model: $model"
                        exit 1
                    fi
                    if [[ "$method" == "headcompresskv" ]]; then
                        python -m longbench.pred_compresskv \
                            --model $model \
                            --max_capacity_prompt $budget \
                            --method $method \
                            --normalize \
                            --beta 1.2 \
                            --temp 1.0 \
                            --head_choice reason \
                            --floor_alpha 0.2 \
                            --gqa_support \
                            --head_score_path "$head_score_path" \
                            --importance_head_path "$head_idx_path"
                    else
                        
                        python -m longbench.pred_compresskv \
                            --model $model \
                            --max_capacity_prompt $budget \
                            --method $method \
                            --normalize \
                            --beta 1.2 \
                            --temp 1.0 \
                            --head_choice reason \
                            --floor_alpha 0.2 \
                            --gqa_support \
                            --head_score_path "$head_score_path"
                    fi
                elif [[ "$method" == "adakv" ]]; then
                    python -m longbench.pred_compresskv \
                        --model $model \
                        --max_capacity_prompt $budget \
                        --method $method \
                        --floor_alpha 0.2 \
                        --gqa_support
                elif [[ "$method" == "adacompresskv" ]] || [[ "$method" == "compresskv" ]]; then
                    if [[ "$model" == *"Llama-3.1-8B-Instruct"* ]]; then
                        head_idx_path="scores/Llama-3.1-8B-Instruct_head_idx.json"
                        layer_utility_path="scores/Llama-3.1-8B-Instruct_layer_score.jsonl"
                    elif [[ "$model" == *"Mistral-7B-Instruct-v"* ]]; then
                        head_idx_path="scores/Mistral-7B-Instruct-v0.3_head_idx.json"
                        layer_utility_path="scores/Mistral-7B-Instruct-v0.3_layer_score.jsonl"
                    elif [[ "$model" == *"Qwen2.5-7B-Instruct"* ]]; then
                        head_idx_path="scores/Qwen2.5-7B-Instruct_head_idx.json"
                        layer_utility_path="scores/Qwen2.5-7B-Instruct_layer_score.jsonl"
                    else
                        echo "Unknown model: $model"
                        exit 1
                    fi
                    if [[ "$method" == "adacompresskv" ]]; then
                        python -m longbench.pred_compresskv \
                        --model $model \
                        --max_capacity_prompt $budget \
                        --method $method \
                        --floor_alpha 0.2 \
                        --gqa_support \
                        --layer_importance_score_path "$layer_utility_path" \
                        --importance_head_path "$head_idx_path"
                    else 
                        python -m longbench.pred_compresskv \
                            --model $model \
                            --max_capacity_prompt $budget \
                            --method $method \
                            --layer_importance_score_path "$layer_utility_path" \
                            --importance_head_path "$head_idx_path" 
                    fi
                else
                    python -m longbench.pred_compresskv \
                        --model $model \
                        --max_capacity_prompt $budget \
                        --method $method \
                        
                fi
            done
        fi
    done
done

