from openmind import omdatasets, pipeline, is_torch_npu_available, AutoTokenizer
import argparse
def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        help="Path to model",
        default="ChongqingAscend/distilroberta-base-climate-detector",
    )
    args = parser.parse_args()
    return args

args = parse_args()
model_path = args.model_name_or_path
prompt_text = "שלום, קוראים לי"
if is_torch_npu_available():
    device = "npu:0"
else:
    device = "cpu"

text_generator = pipeline(task="text-generation", model=model_path, tokenizer=model_path,device=device)

result = text_generator(prompt_text, num_return_sequences=1, batch_size=1, do_sample=True, top_k=40, top_p=0.92, temperature = 1, repetition_penalty=5.0, max_length = 192)


print("result = " + str(result))