Views
No views yet
pip install onnxruntime1pip uninstall onnxruntime -y
2pip install onnxruntime-gpu1import onnxruntime as ort
2from transformers import AutoTokenizer
3import numpy as np
4import argparse
5
6def generate_text(prompt, num_tokens, model_path, tokenizer_path):
7 tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
8 providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
9 session = ort.InferenceSession(model_path, providers=providers)
10
11 input_ids = tokenizer(prompt, return_tensors="np").input_ids
12
13 for _ in range(num_tokens):
14 # Create attention mask and position ids
15 attention_mask = np.ones_like(input_ids)
16 position_ids = np.arange(input_ids.shape[1])[None, :]
17
18 outputs = session.run(
19 output_names=['logits'],
20 input_feed={
21 'input_ids': input_ids,
22 'attention_mask': attention_mask,
23 'position_ids': position_ids
24 }
25 )
26
27 next_token = np.argmax(outputs[0][0, -1, :])
28
29 input_ids = np.concatenate([input_ids, [[next_token]]], axis=1)
30
31 return tokenizer.decode(input_ids[0], skip_special_tokens=True)
32
33if __name__ == "__main__":
34 parser = argparse.ArgumentParser(description='Generate text using ONNX model')
35 parser.add_argument('prompt', type=str, help='Input prompt for generation')
36 parser.add_argument('num_tokens', type=int, help='Number of tokens to generate')
37 parser.add_argument('--model_path', type=str, default='model.onnx',
38 help='Path to ONNX model file')
39 parser.add_argument('--tokenizer_path', type=str, default='tokenizer',
40 help='Path to tokenizer directory')
41
42 args = parser.parse_args()
43
44 result = generate_text(args.prompt, args.num_tokens, args.model_path, args.tokenizer_path)
45 print(result)python onnx_inference.py "Once upon a time" 512 --model_path /path/to/model.onnx --tokenizer_path /path/to/model/dir