import copy
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "left"

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B", attn_implementation="eager", dtype=torch.float32).cuda().eval()

sentences = ["short sentence", "a much much longer sentence that forces left padding."]
question = "\n\nQ: Is the sentence short or long? A:"


# (A) vanilla generation with batch_size = 1
tokenized_sentence = tokenizer(sentences[0], return_tensors="pt").to(model.device) # tokenize just the first sentence
tokenized_question = tokenizer(question, return_tensors="pt", add_special_tokens=False).to(model.device) # tokenize the question
new_input_ids = torch.cat([tokenized_sentence["input_ids"], tokenized_question["input_ids"]], dim=1)
new_attention_mask = torch.cat([tokenized_sentence["attention_mask"], tokenized_question["attention_mask"]], dim=1)
new_inputs = {"input_ids": new_input_ids, "attention_mask": new_attention_mask}
with torch.inference_mode():
    outputs = model.generate(**new_inputs, max_new_tokens=8, top_p=1.0, do_sample=False, pad_token_id=tokenizer.pad_token_id)
generated_ids = outputs[:, new_inputs["input_ids"].shape[1]:, ...]
generated_A = generated_ids[0]
print("------------------------------")
print("(A) Vanilla generation with batch_size = 1")
print("Batch_item 0:", tokenizer.decode(generated_A)) # correct
print("------------------------------")


# (B) vanilla generation with batch_size > 1
tokenized_sentence = tokenizer(sentences, padding=True, return_tensors="pt").to(model.device) # tokenize all sentences
B = tokenized_sentence["input_ids"].shape[0]
tokenized_question = tokenizer([question] * B, return_tensors="pt", add_special_tokens=False).to(model.device) # tokenize the question replicate per batch
new_input_ids = torch.cat([tokenized_sentence["input_ids"], tokenized_question["input_ids"]], dim=1)
new_attention_mask = torch.cat([tokenized_sentence["attention_mask"], tokenized_question["attention_mask"]], dim=1)
new_inputs = {"input_ids": new_input_ids, "attention_mask": new_attention_mask}
with torch.inference_mode():
    outputs = model.generate(**new_inputs, max_new_tokens=8, top_p=1.0, do_sample=False, pad_token_id=tokenizer.pad_token_id)
generated_ids = outputs[:, new_inputs["input_ids"].shape[1]:, ...]

print("------------------------------")
print("(B) Vanilla generation with batch_size > 1")
generated_B_0 = generated_ids[0]
print("Batch_item 0:", tokenizer.decode(generated_B_0)) # correct
generated_B_1 = generated_ids[1]
print("Batch_item 1:", tokenizer.decode(generated_B_1)) # correct
print("------------------------------")


# (C) prefill + continue generation with batch_size = 1
tokenized_sentence = tokenizer(sentences[0], return_tensors="pt").to(model.device) # tokenize just the first sentence
with torch.inference_mode():
    sentence_cache = model(**tokenized_sentence, use_cache=True).past_key_values
tokenized_question = tokenizer(question, return_tensors="pt", add_special_tokens=False).to(model.device) # tokenize the question
new_input_ids = torch.cat([tokenized_sentence["input_ids"], tokenized_question["input_ids"]], dim=1)
new_attention_mask = torch.cat([tokenized_sentence["attention_mask"], tokenized_question["attention_mask"]], dim=1)
new_inputs = {"input_ids": new_input_ids, "attention_mask": new_attention_mask}
past_key_values = copy.deepcopy(sentence_cache)
with torch.inference_mode():
    outputs = model.generate(**new_inputs, past_key_values=past_key_values, max_new_tokens=8, top_p=1.0, do_sample=False, pad_token_id=tokenizer.pad_token_id)
generated_ids = outputs[:, new_inputs["input_ids"].shape[1]:, ...]
generated_C = generated_ids[0]
print("------------------------------")
print("(C) Prefill + continue generation generation with batch_size = 1")
print("Batch_item 0:", tokenizer.decode(generated_C)) # correct
print("------------------------------")


# (D) prefill + continue generation with batch_size > 1
tokenized_sentence = tokenizer(sentences, padding=True, return_tensors="pt").to(model.device) # tokenize all sentences
with torch.inference_mode():
    sentence_cache = model(**tokenized_sentence, use_cache=True).past_key_values
B = tokenized_sentence["input_ids"].shape[0]
tokenized_question = tokenizer([question] * B, return_tensors="pt", add_special_tokens=False).to(model.device) # tokenize the question replicate per batch
new_input_ids = torch.cat([tokenized_sentence["input_ids"], tokenized_question["input_ids"]], dim=1)
new_attention_mask = torch.cat([tokenized_sentence["attention_mask"], tokenized_question["attention_mask"]], dim=1)
new_inputs = {"input_ids": new_input_ids, "attention_mask": new_attention_mask}
past_key_values = copy.deepcopy(sentence_cache)
with torch.inference_mode():
    outputs = model.generate(**new_inputs, past_key_values=past_key_values, max_new_tokens=8, top_p=1.0, do_sample=False, pad_token_id=tokenizer.pad_token_id)
generated_ids = outputs[:, new_inputs["input_ids"].shape[1]:, ...]
print("------------------------------")
print("(D) Prefill + continue generation generation with batch_size > 1")
generated_D_0 = generated_ids[0]
print("Batch_item 0:", tokenizer.decode(generated_D_0)) # <-- wrong: it has padded tokens and aren't handled correctly in generate
generated_D_1 = generated_ids[1]
print("Batch_item 1:", tokenizer.decode(generated_D_1)) # correct because it doesn't have padded tokens
print("------------------------------")