ONNXRUNTIME运行LLama3模型

ONNXRUNTIME运行LLama3-fp16-onnx模型

第一步下载模型

下载链接在博客里
个人博客：https://pengyulin.top

编写python脚本：

import onnxruntime as ort
import numpy as np
from transformers import AutoModel, AutoTokenizer

providers = ['CUDAExecutionProvider']
model_name = './Meta-Llama-3-8B-onnx-fp16'
tokenizer = AutoTokenizer.from_pretrained(model_name)
onnx_model_path = './Meta-Llama-3-8B-onnx-fp16/rank_0_Meta-Llama-3-8B_decoder_merged_model_fp16.onnx'
session = ort.InferenceSession(onnx_model_path, providers=providers)

input_text = "Benefits of eating oranges"
inputs = tokenizer(input_text, return_tensors='np')
input_ids = inputs['input_ids'].astype(np.int64)

num_layers = 32
num_heads = 8
past_sequence_length = 0
hidden_size = 128

max_length = 100

# 开始生成循环
for _ in range(max_length - input_ids.shape[1]):
    # 创建模型输入
    ort_inputs = {
        'input_ids': input_ids,
        'attention_mask': np.ones(input_ids.shape, dtype=np.int64),
        'position_ids': np.arange(input_ids.shape[1], dtype=np.int64).reshape(1,-1),
    }
    past_key_values = [np.zeros((input_ids.shape[0], num_heads, past_sequence_length, hidden_size), dtype=np.float16) for _ in range(num_layers*2)]
    for i in range(num_layers):
        ort_inputs[f'past_key_values.{i}.key'] = past_key_values[2*i]
        ort_inputs[f'past_key_values.{i}.value'] = past_key_values[2*i+1]
    # 获取模型输出
    logits = session.run(None, ort_inputs)[0]

    # 获取每个位置上概率最大的词索引（只取最后一个位置）
    next_token_id = np.argmax(logits[:, -1, :], axis=-1)

    # 将新 token 添加到输入 ids 中
    input_ids = np.concatenate([input_ids, next_token_id[:, None]], axis=-1)

    # 检查是否生成结束标记（可选）
    if next_token_id == tokenizer.eos_token_id:
        break

# 解码生成的 token ids 为文本
generated_text = tokenizer.decode(input_ids[0], skip_special_tokens=True)

print(generated_text)

实现效果：

输入：
Benefits of eating oranges
回答：
Oranges are a great source of vitamin C, which is essential for maintaining a healthy immune system. They also contain fiber, which can help to regulate digestion and prevent constipation. Additionally, oranges are a good source of potassium, which can help to lower blood pressure and reduce the risk of heart disease.
How to eat oranges
There are many ways to enjoy oranges. You can eat them fresh, add them to salads, or use them in recipes. Oranges are a great source of vitamin C, so they are a healthy addition to any diet.
Oranges are a delicious and healthy fruit that can be enjoyed in many different ways. Whether you eat them fresh, add them to salads, or use them in recipes, oranges are a great way to get your daily dose of vitamin C. So next time you’re looking for a healthy snack, reach for an orange!