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"
)

save_dir="outputs/needle_result"

for model in "${models[@]}"
do
    if [[ "$model" == *"Llama-3.1-8B-Instruct"* ]]; then
        context_length_arg="--context_length 8000 16000 24000 32000 40000 48000 56000 64000 72000 80000 88000 96000 104000 112000 120000 128000"
        save_dir="outputs/needle_result/needle_result_llama3_128k"
    elif [[ "$model" == *"Mistral-7B-Instruct-v"* ]]; then
        context_length_arg="--context_length 2000 4000 6000 8000 10000 12000 14000 16000 18000 20000 22000 24000 26000 28000 30000 32000"
        save_dir="outputs/needle_result/needle_result_mistral_32k"
    elif [[ "$model" == *"Qwen2.5-7B-Instruct"* ]]; then
        context_length_arg="--context_length  8000 16000 24000 32000 40000 48000 56000 64000 72000 80000 88000 96000 104000 112000 120000 128000"
        save_dir="outputs/needle_result/needle_result_qwen2_128k"
    else
        echo "Unknown model: $model"
        exit 1
    fi

    for method in "${methods[@]}"
    do
        if [[ "$method" == "fullkv" ]]; then
            echo "Running model=$model method=$method"
            python -m needle.run_needle_in_haystack \
                --model_path $model \
                --method $method \
                --save_dir "$save_dir" \
                $context_length_arg
        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 needle.run_needle_in_haystack \
                        --model_path $model \
                        --max_capacity_prompt $budget \
                        --method $method \
                        --save_dir "$save_dir" \
                        $tau1 $tau2  \
                        $context_length_arg
                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 needle.run_needle_in_haystack \
                            --model_path $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" \
                            --save_dir "$save_dir" \
                            $context_length_arg
                    else
                        
                        python -m needle.run_needle_in_haystack \
                            --model_path $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" \
                            --save_dir "$save_dir" \
                            $context_length_arg
                    fi
                elif [[ "$method" == "adakv" ]]; then
                    python -m needle.run_needle_in_haystack \
                        --model_path $model \
                        --max_capacity_prompt $budget \
                        --method $method \
                        --floor_alpha 0.2 \
                        --gqa_support \
                        --save_dir "$save_dir" \
                        $context_length_arg
                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 needle.run_needle_in_haystack \
                        --model_path $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" \
                        --save_dir "$save_dir" \
                        $context_length_arg
                    else
                        python -m needle.run_needle_in_haystack \
                            --model_path $model \
                            --max_capacity_prompt $budget \
                            --method $method \
                            --layer_importance_score_path "$layer_utility_path" \
                            --importance_head_path "$head_idx_path" \
                            --save_dir "$save_dir" \
                            $context_length_arg
                    fi
                else
                    python -m needle.run_needle_in_haystack \
                        --model_path $model \
                        --max_capacity_prompt $budget \
                        --method $method \
                        --save_dir "$save_dir" \
                        $context_length_arg
                        
                fi
            done
        fi
    done
done