import argparse

import torch
from openmind import pipeline, is_torch_npu_available
from openmind_hub import snapshot_download
import time

def parse_args():
    parser = argparse.ArgumentParser(description="Eval the model")
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        help="Path to the model",
       default="zhouhui/ArabianGPT-0.3B-QA",
        )
    args = parser.parse_args()
    return args


def main():
    args = parse_args()
    model_path = args.model_name_or_path
    if is_torch_npu_available():
        device = "npu:0"
    elif torch.cuda.is_available():
        device = "cuda:0"
    else:
        device = "cpu"
    #device = "cpu"
    start_time = time.time()
    qa = pipeline("question-answering", model=model_path, tokenizer=model_path, device=device)
    qa_input = {
            "question": "Why is model conversion important?",
            "context": "The option on convert models between FARM and openmind gives freedom to the user and let people easily switch between frameworks."
    }
    ans = qa(qa_input)
    print()
    print(ans)
    print()
    end_time = time.time()
    print(f"硬件环境：{device},推理执行时间：{end_time - start_time}秒")


if __name__ == "__main__":
    main()