import torch
import os 

from transformers import AutoTokenizer, AutoModelForCausalLM


os.environ["CUDA_LAUNCH_BLOCKING"] = "1"

def test_flash_attention_2_continue_generate_with_position_ids():
    """
    Tests that the given attention implementation can work with packed sequences and infers the mask
    from position ids. This test requires the model to use new attention mask API which handles packing.
    """

    max_new_tokens = 2
    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B-Instruct")
    model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B-Instruct", device_map="auto", torch_dtype=torch.bfloat16).eval().cuda()
    model.set_attn_implementation("flash_attention_2")
    inputs_dict = tokenizer("Hello, how are you?", return_tensors="pt")

    # make sure that all models have enough positions for generation
    dummy_input_ids = inputs_dict["input_ids"]
    if hasattr(model.config, "max_position_embeddings"):
        model.config.max_position_embeddings = max_new_tokens + dummy_input_ids.shape[1] + 1

    dummy_input_ids = inputs_dict["input_ids"]
    dummy_position_ids = torch.arange(dummy_input_ids.shape[1], device="cuda")
    dummy_position_ids = dummy_position_ids.unsqueeze(0).repeat(dummy_input_ids.shape[0], 1)

    # Store cache for the input prompt
    output = model(dummy_input_ids.cuda(), position_ids=dummy_position_ids.cuda(), use_cache=True)


    # create new input_ids and position_ids to continue generation re-using the cache
    new_input_ids = output.logits[:, -1, :].float().argmax(-1)[:, None]
    past_length = dummy_input_ids.shape[1]
    position_ids = torch.arange(past_length, past_length + new_input_ids.shape[1], device="cuda")
    position_ids = position_ids.unsqueeze(0).repeat(new_input_ids.shape[0], 1)

    output = model(
        input_ids=new_input_ids,
        past_key_values=output.past_key_values,
        position_ids=position_ids,
        use_cache=True,
    )
    next_token_logits = output.logits[:, -1, :].float()


    generate_kwargs = {
        "pad_token_id": -1,
        "eos_token_id": -1,
        "forced_eos_token_id": None,
        "use_cache": True,
        "do_sample": False,
        "return_dict_in_generate": True,
        "output_logits": True,
        "max_new_tokens": max_new_tokens,
    }
    generation_out = model.generate(dummy_input_ids.cuda(), **generate_kwargs)
    next_token_logits_from_generate = generation_out.logits[-1]



    # acceptable numerical instability
    # print(next_token_logits_from_generate, next_token_logits)
    tol = torch.finfo(torch.bfloat16).eps
    torch.testing.assert_close(next_token_logits_from_generate, next_token_logits, rtol=tol, atol=tol)


test_flash_attention_2_continue_generate_with_position_ids()

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import torch
import os 

os.environ["CUDA_LAUNCH_BLOCKING"] = "1"

def test_generate_are_equivalent():
    """
    Tests that the given attention implementation can work with packed sequences and infers the mask
    from position ids. This test requires the model to use new attention mask API which handles packing.
    """

    max_new_tokens = 20
    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
    model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B", device_map="auto", torch_dtype=torch.bfloat16).eval().cuda()
    model.set_attn_implementation("flash_attention_2")
    inputs_dict = tokenizer("Hello, how are you?", return_tensors="pt")

    # Manual generation using for loop
    print("Manual generation:")
    position_ids = torch.arange(inputs_dict["input_ids"].shape[1], device="cuda")
    position_ids = position_ids.unsqueeze(0).repeat(inputs_dict["input_ids"].shape[0], 1)

    # Store cache for the input prompt
    output = model(inputs_dict["input_ids"].cuda(), position_ids=position_ids.cuda(), use_cache=True)
    past_key_values = output.past_key_values
    manual_tokens = inputs_dict["input_ids"].clone()

    # Generate tokens one by one
    for i in range(max_new_tokens):
        # Get next token
        new_input_ids = output.logits[:, -1, :].float().argmax(-1)[:, None]
        manual_tokens = torch.cat([manual_tokens, new_input_ids.cpu()], dim=1)

        # Print the newly generated token
        new_token_text = tokenizer.decode(new_input_ids[0].cpu().item())
        print(f"Token {i+1}: '{new_token_text}'")

        # Prepare for next iteration
        past_length = manual_tokens.shape[1] - 1
        position_ids = torch.arange(past_length, past_length + 1, device="cuda")
        position_ids = position_ids.unsqueeze(0).repeat(new_input_ids.shape[0], 1)

        output = model(
            input_ids=new_input_ids,
            past_key_values=past_key_values,
            position_ids=position_ids,
            use_cache=True,
        )
        past_key_values = output.past_key_values

    print("\nManual generation complete sequence:")
    print(f"'{tokenizer.decode(manual_tokens[0], skip_special_tokens=True)}'")

    # Generate using model.generate()
    print("\nGenerate method:")
    generate_kwargs = {
        "pad_token_id": tokenizer.eos_token_id if tokenizer.eos_token_id is not None else -1,
        "eos_token_id": -1,
        "forced_eos_token_id": None,
        "use_cache": True,
        "do_sample": False,
        "return_dict_in_generate": True,
        "output_logits": True,
        "max_new_tokens": max_new_tokens,
    }
    generation_out = model.generate(inputs_dict["input_ids"].cuda(), **generate_kwargs)

    # Print tokens from generate method
    generated_tokens = generation_out.sequences[0]
    original_length = inputs_dict["input_ids"].shape[1]

    for i in range(max_new_tokens):
        token_id = generated_tokens[original_length + i].cpu().item()
        token_text = tokenizer.decode(token_id)
        print(f"Token {i+1}: '{token_text}'")

    print("\nGenerate method complete sequence:")
    print(f"'{tokenizer.decode(generated_tokens, skip_special_tokens=True)}'")


test_generate_are_equivalent()