diff --git a/samples/dnn/gpt2_inference.py b/samples/dnn/gpt2_inference.py index 8f0ddcf7d7..f9a5a2b072 100644 --- a/samples/dnn/gpt2_inference.py +++ b/samples/dnn/gpt2_inference.py @@ -30,21 +30,17 @@ Run the script: python gpt2_inference.py --model= --prompt= ''' - - import numpy as np import tiktoken import argparse import cv2 as cv -from tqdm import tqdm def parse_args(): parser = argparse.ArgumentParser(description='Use this script to run GPT-2 inference in OpenCV', formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument('--model', type=str, required=True, help='Path to GPT-2 model ONNX model file.') parser.add_argument("--prompt", type=str, default="Hello, I'm a language model,", help="Prompt to start with.") - parser.add_argument("--max_seq_len", type=int, default=40, help="Number of tokens to continue.") - parser.add_argument("--batch_size", type=int, default=1, help="Number of batches.") + parser.add_argument("--max_seq_len", type=int, default=1024, help="Number of tokens to continue.") parser.add_argument("--seed", type=int, default=0, help="Random seed") return parser.parse_args() @@ -53,21 +49,27 @@ def stable_softmax(logits): return exp_logits / np.sum(exp_logits, axis=-1, keepdims=True) -def gpt2_inference(net, tokens, max_length, num_return_sequences=1): + +def gpt2_inference(net, tokens, max_length, tokenizer): print("Inferencing GPT-2 model...") x = np.array(tokens) - x = np.tile(x, (num_return_sequences, 1)).astype(np.int32) + x = np.tile(x, (1, 1)).astype(np.int32) pos = np.arange(0, len(x), dtype=np.int32) - counter = x.shape[1] - pbar = tqdm(total=max_length - counter, desc="Generating tokens") - while counter < max_length: + # warm up + net.setInputsNames(['input_ids', 'position_ids']) + net.setInput(x, 'input_ids') + net.setInput(pos, 'position_ids') + logits = net.forward() + + stop_tokens = (50256, ) ## could be extended to include more stop tokens + print("\n", tokenizer.decode(tokens), sep="", end="") + while 0 < max_length and x[:, -1] not in stop_tokens: net.setInputsNames(['input_ids', 'position_ids']) net.setInput(x, 'input_ids') net.setInput(pos, 'position_ids') - logits = net.forward() # logits is assumed to be (B, seq_length, vocab_size) and needs to be the last token's logits @@ -87,23 +89,27 @@ def gpt2_inference(net, tokens, max_length, num_return_sequences=1): sampled_indices = [np.random.choice(topk_indices[i], p=topk_probs[i]) for i in range(len(topk_probs))] sampled_indices = np.array(sampled_indices).reshape(-1, 1) + # Decode and print the new token + new_word = tokenizer.decode([sampled_indices[0, 0]]) + + ## clean the prints from the previous line + print(new_word, end='', flush=True) + # Append to the sequence x = np.concatenate((x, sampled_indices), axis=1) pos = np.arange(0, x.shape[1], dtype=np.int32) # shape (T) - counter += 1 - pbar.update(1) + max_length -= 1 - pbar.close() - print("Inference done!") - return x + print('\n') if __name__ == '__main__': args = parse_args() + print("Preparing GPT-2 model...") + np.random.seed(args.seed) max_length = args.max_seq_len - num_return_sequences = args.batch_size prompt = args.prompt net = cv.dnn.readNet(args.model) @@ -111,9 +117,4 @@ if __name__ == '__main__': enc = tiktoken.get_encoding('gpt2') tokens = enc.encode(prompt) - output = gpt2_inference(net, tokens, max_length, num_return_sequences) - - for i in range(num_return_sequences): - tokens = output[i].tolist() - decoded = enc.decode(tokens) - print(">>>>", decoded) + gpt2_inference(net, tokens, max_length, enc)