from transformers import AutoTokenizer
from transformers import LlamaForCausalLM, AutoModelForCausalLM
from time import perf_counter
import torch
import argparse

def get_args():
    parser = argparse.ArgumentParser(description="Torch model Demo")
    
    parser.add_argument('--model-path', type=str, required=True,
                        help='Model Path, include config.ini and tokenizer files')
    parser.add_argument('--tokenizer-path', type=str, default=None)
    
    parser.add_argument("--prompt", type=str, required=False)
    parser.add_argument("--max-output-length", type=int, default=512)
    parser.add_argument("--warmups", type=int, default=10)
    parser.add_argument("--avgnums", type=int, default=10)
    args = parser.parse_args()

    print('\n=================== Arguments ===================')
    for k, v in vars(args).items():
        print(f' - {k.ljust(25, ".")}: {v}')
    print('=================================================')

    return args

def main():
    args = get_args()
    device = torch.device("cuda")
    
    prompt_template = "Human: {}\n\nAssistant:" # xverse
    # prompt_template = "<human>:{}\n<bot>:" # llama-ziya 13b

    prompt = prompt_template.format(args.prompt)
    
    model = AutoModelForCausalLM.from_pretrained(args.model_path, torch_dtype=torch.float16, trust_remote_code=True).eval().to(device)
    tokenizer = AutoTokenizer.from_pretrained(args.model_path, use_fast=False, trust_remote_code=True)

    test_batch_size = [1, 8, 16, 32, 64]
    print("test_batch_size: ", test_batch_size)

    for i, bs in enumerate(test_batch_size):

        prompts = [prompt] * bs

        # warmup gpu        
        for _ in range(args.warmups):
            input_ids = tokenizer(prompts, return_tensors="pt").input_ids.to(device)
            generate_ids = model.generate(
                        input_ids,
                        max_new_tokens=args.max_output_length, 
                        do_sample = False,
                        top_k = 30, 
                        top_p = 0.85, 
                        temperature = 1.0, 
                        repetition_penalty=1., 
                        eos_token_id=2, 
                        bos_token_id=1, 
                        pad_token_id=0)
            
            generate_ids = [output_ids[len(single_input_id):] for single_input_id, output_ids in zip(input_ids, generate_ids)]
            outputs = tokenizer.batch_decode(generate_ids)
            
        # test
        start = perf_counter()
        for _ in range(args.avgnums):
            input_ids = tokenizer(prompts, return_tensors="pt").input_ids.to(device)
            generate_ids = model.generate(
                        input_ids,
                        max_new_tokens=args.max_output_length, 
                        do_sample = False,
                        top_k = 30, 
                        top_p = 0.85, 
                        temperature = 1.0, 
                        repetition_penalty=1., 
                        eos_token_id=2, 
                        bos_token_id=1, 
                        pad_token_id=0)
            
            generate_ids = [output_ids[len(single_input_id):] for single_input_id, output_ids in zip(input_ids, generate_ids)]
            output_texts = tokenizer.batch_decode(generate_ids)
            
        end = perf_counter()
        cost = (end - start) / args.avgnums
        
        # 计算吞吐量
        input_output_texts = [prompt + ' ' + gtext for prompt, gtext in zip(prompts, output_texts)]
        tokens = 0
        input_tokens = len(tokenizer.encode(prompt))
        words = 0
        for text in input_output_texts:
            tokens += len(tokenizer.encode(text))
            words += len(text)
            
        avg_output_tokens = tokens / len(input_output_texts) - input_tokens
        print(
            f"\nBatch Size: {bs}, All tokens: {tokens}. Input tokens: {input_tokens}. Output tokens: {avg_output_tokens} Cost: {cost} seconds. Speed: {tokens/cost} tokens/s."
        )
        print(
            f"Batch Size: {bs}, All generated words: {words}. Cost: {cost} seconds. Speed: {words/cost} words/s."
        )

        if i == 0:
            for k in range(bs):
                print(
                    f"The {k} Sample, \n\t\tInputs: {prompts[k]}. \n\t\tOutputs: {output_texts[k].lstrip()}")
                if k > 2:
                    break

if __name__ == "__main__":
    main()