from openmind import AutoModelForCausalLM, AutoTokenizer
from openmind import is_torch_npu_available, pipeline
import torch
import argparse
import torch.nn.functional as F
import time

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model_name_or_path",
        "-m",
        type=str,
        help="Path to model",
        default="huangjingwang/Promt-generator",
    )
    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"  
    #device = "cpu"  
    start_time = time.time()
    tokenizer = AutoTokenizer.from_pretrained(model_path )
    model = AutoModelForCausalLM.from_pretrained(model_path ).to(device)
    
    prompt = "a red car"
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    outputs = model.generate(**inputs)
    generated_prompt = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    print(f"生成结果： {generated_prompt }")
    end_time = time.time()
    print(f"硬件环境：{device},推理执行时间：{end_time - start_time}秒")


if __name__ == "__main__":
    main()

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("UnfilteredAI/Promt-generator")
model = AutoModelForCausalLM.from_pretrained("UnfilteredAI/Promt-generator")

prompt = "a red car"
inputs = tokenizer(prompt, return_tensors="pt")
outputs = model.generate(**inputs)
generated_prompt = tokenizer.decode(outputs[0], skip_special_tokens=True)

print(generated_prompt)