import argparse
import torch
from openmind import is_torch_npu_available, AutoTokenizer, AutoModelForCausalLM
import time

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        help="Path to model",
        default=None,
    )

    args = parser.parse_args()
    return args


def main():
    start_time = time.time()  # 记录开始时间
    
    args = parse_args()
    if args.model_name_or_path:
        model_path = args.model_name_or_path
    else:
        model_path = "BioMedGPT-LM-7B"

    if is_torch_npu_available():
        device = "npu:0"
    else:
        device = "cpu"
#     device='npu:0'
    device = 'cpu'
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForCausalLM.from_pretrained(model_path).to(device)

    input_ids = tokenizer("Gra", return_tensors='pt').to(model.device)["input_ids"]
    output = model.generate(input_ids, max_new_tokens=48, do_sample=True, temperature=0.7)
    print(tokenizer.decode(output[0]))

    end_time = time.time()  # 记录结束时间
    elapsed_time = end_time - start_time  # 计算差值
    print(f"{device}:Program finished in {elapsed_time:.2f} seconds.")  # 打印运行时间


if __name__ == "__main__":
    main()