ONNX example

[lutzroeder]

Export

torch.onnx.export(model, context, 'gpt2.onnx',
    input_names=['input', 'temperature', 'top_k'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch', 1: 'tokens'}, 'output': {0: 'batch'}})

Inference

import sys
import onnxruntime
import tiktoken
prompt = 'The Eiffel tower is in' if len(sys.argv) <= 1 else sys.argv[1]
encoding = tiktoken.get_encoding('gpt2')
context = encoding.encode(prompt, allowed_special={'<|endoftext|>'})
session = onnxruntime.InferenceSession('gpt2.onnx')
print(prompt, end='', flush=True)
for _ in range(50): # max_tokens
    outputs = session.run(['output'], {'input': [context], 'temperature': [0.7], 'top_k': [40]})
    y = outputs[0][0]
    print(encoding.decode(y), end='', flush=True)
    context.extend(y)
print('')