import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

name = "google/gemma-2-9b-it"
tokenizer_name = name
llm_model_name = name
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
llm_model = AutoModelForCausalLM.from_pretrained(
        llm_model_name,
        device_map="auto",
        torch_dtype=torch.bfloat16,
        # attn_implementation="flash_attention_2",
    )
llm_model.eval()

def chatWithLLM(model: AutoModelForCausalLM, tokenizer: AutoTokenizer):
    messages = [[
        {"role": "user", "content": "laugh " * (idx + 1) + " How many laugh are there?"},
    ] for idx in range(5)]
    input_ids = tokenizer.apply_chat_template(
        messages, 
        padding=True, 
        add_generation_prompt=True, 
        return_tensors="pt", 
        return_dict=True
    ).to(model.device)
    
    outputs = model.generate(
        **input_ids,
        negative_prompt_attention_mask = input_ids["attention_mask"],
        do_sample=False,
        max_new_tokens=500,
        temperature=0.1,
    )

    # this is ugly code to isolate input message, but not related to the bug I guess
    response = tokenizer.batch_decode(outputs, skip_special_tokens=True)
    input_msg = tokenizer.batch_decode(input_ids["input_ids"], skip_special_tokens=True)
    for idx, im in enumerate(input_msg):
        response[idx] = response[idx][len(im):]
    
    return response

print(chatWithLLM(model=llm_model, tokenizer=tokenizer))
# ['', '', '', '', 'There are **5** laughs.  😄 \n']

messages = []
for idx in range(5):
    messages.append("laugh " * (idx + 1) + " How many laugh are there?")
print(messages)
input_ids = tokenizer(
    messages, 
    padding=True, 
    return_tensors="pt", 
).to(model.device)