import torch
from transformers import AutoModelForCausalLM, LlamaForCausalLM, AutoTokenizer, AutoConfig, LlamaTokenizer, LlamaConfig
from transformers.modeling_utils import load_sharded_checkpoint
from accelerate import init_empty_weights, load_checkpoint_in_model, load_checkpoint_and_dispatch
from utils_hh.modify_llama import convert_kvcache_llama_heavy_recent, LlamaAttention_heavy_hitter

ENABLE_Heavy_Hitter_FUNCTIONS = {
    "llama": convert_kvcache_llama_heavy_recent,
}

model_name = 'meta-llama/Llama-2-7b-hf'
cache_dir = 'checkpoint/models--meta-llama--Llama-2-7b-hf/snapshots/8a0442e81540efaeb1a0fe3e95477b5e0edfd423'
heavy_ratio = 0.1
recent_ratio = 0.1
length = 64
seed = 42


torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device: {}".format(device))

tokenizer = LlamaTokenizer.from_pretrained(model_name)
config = LlamaConfig.from_pretrained(model_name)
config.heavy_ratio = heavy_ratio
config.recent_ratio = recent_ratio

# load model without weight
with init_empty_weights():
    model = LlamaForCausalLM(config)

model = convert_kvcache_llama_heavy_recent(model, config)

print(model) 

# load checkpoint into the empty weight

model = load_checkpoint_and_dispatch(
    model,
    checkpoint=cache_dir,
    device_map="auto",
    offload_folder=cache_dir,
    dtype=torch.float16,
    offload_state_dict=True, 
)

prompt_text = 'Hello.'
input_ids = tokenizer(prompt_text, add_special_tokens=False, return_tensors='pt').input_ids.to(model.device)

generate_ids_hh = model.generate(
    input_ids,
    max_new_tokens=100,
    use_cache=True,
    ) # error here