from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteriaList, StopStringTextMatchCriteria, StopStringCriteria
from time import time
import torch

N_RUNS = 100
MAX_NEW_TOKENS = 100
MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
STOP_STRINGS = ["Potato", "Carrots", "Onions", "Garlic", "Tomatoes", "Lettuce", "Cucumbers"]
BATCH_SIZE = 1

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto", dtype=torch.bfloat16)
inputs = tokenizer(["The quick brown"] * BATCH_SIZE, return_tensors="pt").to(model.device)

# warmup
for i in range(10):
    gen_out = model.generate(**inputs, do_sample=False, max_new_tokens=MAX_NEW_TOKENS, min_new_tokens=MAX_NEW_TOKENS)
    assert gen_out.shape[1] == MAX_NEW_TOKENS + inputs.input_ids.shape[1]

# ------------------------------------------------
# No stopping criteria
# ------------------------------------------------
all_times = []
for i in range(N_RUNS):
    start_time = time()
    gen_out = model.generate(**inputs, do_sample=False, max_new_tokens=MAX_NEW_TOKENS, min_new_tokens=MAX_NEW_TOKENS)
    assert gen_out.shape[1] == MAX_NEW_TOKENS + inputs.input_ids.shape[1]
    end_time = time()
    all_times.append(end_time - start_time)

avg_time = sum(all_times) / N_RUNS
print(f"[No stopping criteria] Average generation time: {avg_time} seconds")


# ------------------------------------------------
# StopStringCriteria
# ------------------------------------------------
# IMPORTANT NOTE: initializing this for the first time is slow, >1s. Prior to this PR, this initialization was
# done inside `generate` when `stop_strings` is set
init_start_time = time()
custom_stopping_criteria = StoppingCriteriaList([StopStringCriteria(tokenizer, STOP_STRINGS)])
init_end_time = time()
print(f"[StopStringCriteria] first init time: {init_end_time - init_start_time} seconds")

init_start_time = time()
custom_stopping_criteria = StoppingCriteriaList([StopStringCriteria(tokenizer, STOP_STRINGS)])
init_end_time = time()
print(f"[StopStringCriteria] second init time: {init_end_time - init_start_time} seconds")

all_times = []
for i in range(N_RUNS):
    start_time = time()
    gen_out = model.generate(
        **inputs,
        do_sample=False,
        stopping_criteria=custom_stopping_criteria,
        max_new_tokens=MAX_NEW_TOKENS,
        min_new_tokens=MAX_NEW_TOKENS,
    )
    assert gen_out.shape[1] == MAX_NEW_TOKENS + inputs.input_ids.shape[1]
    end_time = time()
    all_times.append(end_time - start_time)

avg_time = sum(all_times) / N_RUNS
print(f"[StopStringCriteria] Average generation time: {avg_time} seconds")

# ------------------------------------------------
# StopStringTextMatchCriteria
# ------------------------------------------------
init_start_time = time()
custom_stopping_criteria = StoppingCriteriaList([StopStringTextMatchCriteria(tokenizer, STOP_STRINGS)])
init_end_time = time()
print(f"[StopStringTextMatchCriteria] first init time: {init_end_time - init_start_time} seconds")

init_start_time = time()
custom_stopping_criteria = StoppingCriteriaList([StopStringTextMatchCriteria(tokenizer, STOP_STRINGS)])
init_end_time = time()
print(f"[StopStringTextMatchCriteria] second init time: {init_end_time - init_start_time} seconds")

all_times = []
for i in range(N_RUNS):
    start_time = time()
    gen_out = model.generate(
        **inputs,
        do_sample=False,
        stopping_criteria=custom_stopping_criteria,
        max_new_tokens=MAX_NEW_TOKENS,
        min_new_tokens=MAX_NEW_TOKENS,
    )
    assert gen_out.shape[1] == MAX_NEW_TOKENS + inputs.input_ids.shape[1]
    end_time = time()
    all_times.append(end_time - start_time)

avg_time = sum(all_times) / N_RUNS
print(f"[StopStringTextMatchCriteria] Average generation time: {avg_time} seconds")


# ------------------------------------------------
# generate with stop strings (using `StopStringTextMatchCriteria` under the hood)
# ------------------------------------------------
all_times = []
for i in range(N_RUNS):
    start_time = time()
    gen_out = model.generate(
        **inputs,
        do_sample=False,
        stop_strings=STOP_STRINGS,
        tokenizer=tokenizer,
        max_new_tokens=MAX_NEW_TOKENS,
        min_new_tokens=MAX_NEW_TOKENS,
    )
    assert gen_out.shape[1] == MAX_NEW_TOKENS + inputs.input_ids.shape[1]
    end_time = time()
    all_times.append(end_time - start_time)

avg_time = sum(all_times) / N_RUNS
print(f"[Default `stop_strings` criteria] Average generation time: {avg_time} seconds")