import argparse

import torch

from openmind import pipeline, is_torch_npu_available

from openmind_hub import snapshot_download





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():

    args = parse_args()

    model_path = args.model_name_or_path



    if is_torch_npu_available():

        device = "npu:0"

    else:

        device = "cpu"



    generator = pipeline('text-generation', model=model_path, device=device)



    output = generator("Hello, I'm a language model,", max_length=30, num_return_sequences=5, num_beams=5)



    print(f">>>output={output}", flush=True)





if __name__ == "__main__":

    main()