transformers==4.45.1
numpy==1.24.4
gguf==0.10.0
accelerate
openmind-hub
einops

from openmind import AutoModelForCausalLM, AutoTokenizer
from openmind import is_torch_npu_available
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",
        type=str,
        help="Path to model",
        default="./h2o-danube3-500m-chat-GGUFF",
    )
    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()
    filename = "h2o-danube3-500m-chat-Q4_K_M.gguf"
    tokenizer = AutoTokenizer.from_pretrained(model_path, gguf_file=filename, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_path, torch_dtype=torch.float16, gguf_file=filename, device_map=device
    )

    
    prompt = "Q: What is the largest animal?\nA:"
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
    input_ids = input_ids.to(model.device)
    generation_output = model.generate(input_ids=input_ids, max_new_tokens=32)

    print(tokenizer.decode(generation_output[0]))
    end_time = time.time()
    print(f"硬件环境：{device},推理执行时间：{end_time - start_time}秒")


if __name__ == "__main__":
    main()

<|prompt|>Why is drinking water so healthy?</s><|answer|>