ONNX 模型库
返回模型

说明文档

这是 Eleusis-12B 的 ONNX 优化版本。 如需了解该模型功能的更全面信息,请访问原模型仓库。

推理

环境要求

如果你使用的是仅 CPU 的机器:

pip install onnxruntime

如果你有可用的 NVIDIA GPU:

pip uninstall onnxruntime -y
pip install onnxruntime-gpu

请确保已安装 CUDA ToolkitcuDNN

import onnxruntime as ort
from transformers import AutoTokenizer
import numpy as np
import argparse

def generate_text(prompt, num_tokens, model_path, tokenizer_path):
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
    session = ort.InferenceSession(model_path, providers=providers)

    input_ids = tokenizer(prompt, return_tensors="np").input_ids

    for _ in range(num_tokens):
        # 创建注意力掩码和位置 ID
        attention_mask = np.ones_like(input_ids)
        position_ids = np.arange(input_ids.shape[1])[None, :]

        outputs = session.run(
            output_names=['logits'],
            input_feed={
                'input_ids': input_ids,
                'attention_mask': attention_mask,
                'position_ids': position_ids
            }
        )

        next_token = np.argmax(outputs[0][0, -1, :])

        input_ids = np.concatenate([input_ids, [[next_token]]], axis=1)

    return tokenizer.decode(input_ids[0], skip_special_tokens=True)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='使用 ONNX 模型生成文本')
    parser.add_argument('prompt', type=str, help='生成用的输入提示词')
    parser.add_argument('num_tokens', type=int, help='要生成的 token 数量')
    parser.add_argument('--model_path', type=str, default='model.onnx', 
                        help='ONNX 模型文件路径')
    parser.add_argument('--tokenizer_path', type=str, default='tokenizer',
                        help='分词器目录路径')

    args = parser.parse_args()

    result = generate_text(args.prompt, args.num_tokens, args.model_path, args.tokenizer_path)
    print(result)
python onnx_inference.py "Once upon a time" 512 --model_path /path/to/model.onnx --tokenizer_path /path/to/model/dir

这是一个示例脚本,并未经过充分优化。

PygmalionAI/Eleusis-12B-ONNX

作者 PygmalionAI

text-generation
↓ 1 ♥ 0

创建时间: 2025-01-27 23:12:25+00:00

更新时间: 2025-01-27 23:19:00+00:00

在 Hugging Face 上查看

文件 (10)

.gitattributes
README.md
config.json
generation_config.json
model.onnx ONNX
model.onnx_data
special_tokens_map.json
tekken.json
tokenizer.json
tokenizer_config.json