if past_key_value is not None and use_cache:
    cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
    key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)

from transformers import AutoTokenizer, AutoModelForCausalLM
from utils import *
import random
import torch

def run_with_custom_order(model, tokenizer, device, prompts, order):
    original_layers = model.model.layers
    layer_dict = {i: layer for i, layer in enumerate(original_layers)}
    shuffled_layers = torch.nn.ModuleList([layer_dict[i] for i in order])
    model.model.layers = shuffled_layers

    for prompt in prompts:
        inputs = tokenizer(prompt, return_tensors="pt").to(device)
        input_length = inputs["input_ids"].shape[1]
        outputs = model.generate(**inputs, max_new_tokens=20, pad_token_id=tokenizer.eos_token_id, use_cache=False)
        generated_text = tokenizer.decode(outputs[0, input_length:], skip_special_tokens=True)
        print(f"{generated_text}")

    model.model.layers = original_layers
    pass

def main():
    device = (
        "cuda" if torch.cuda.is_available() else 
        "mps" if torch.backends.mps.is_available() else 
        "cpu"
    )    
    llm_name = "microsoft/Phi-3-mini-4k-instruct"
    tokenizer = AutoTokenizer.from_pretrained(llm_name)
    model = AutoModelForCausalLM.from_pretrained(llm_name, torch_dtype=torch.bfloat16, trust_remote_code=True)
    model.to(device)
    model.config.pad_token_id = tokenizer.eos_token_id
    num_layers = len(model.model.layers)

    # Load questions and answers
    dataset_type = "mmlu"  # Change this to "default" for capitals dataset
    num_questions = 1
    questions, _ = load_qa(dataset_type, num_questions)
    prompts = [f"<|user|>\n{question}\nChoose A, B, C, or D:<|end|>\n<|assistant|>" for question in questions]
    
    order = list(range(num_layers))
    run_with_custom_order(model, tokenizer, device, prompts, order)

    random.shuffle(order)
    run_with_custom_order(model, tokenizer, device, prompts, order)

    pass

if __name__ == "__main__":
    main()