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)
    print(f">>>output={output}", flush=True)

if __name__ == "__main__":
    main()

@misc{ParsGPT2,
  author = {Hooshvare Team},
  title = {ParsGPT2 the Persian version of GPT2},
  year = {2021},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/hooshvare/parsgpt}},
}