diff --git a/examples/python/model-generate.py b/examples/python/model-generate.py index 9b5af011a..3e2b38d90 100644 --- a/examples/python/model-generate.py +++ b/examples/python/model-generate.py @@ -13,15 +13,30 @@ def main(args): prompts = args.prompts else: prompts = ["I like walking my cute dog", - "What is the best restaurant in town?", - "Hello, how are you today?"] + "What is the best restaurant in town?", + "Hello, how are you today?"] + + if args.chat_template: + if args.chat_template.count('{') != 1 or args.chat_template.count('}') != 1: + print("Error, chat template must have exactly one pair of curly braces, e.g. '<|user|>\n{input} <|end|>\n<|assistant|>'") + exit(1) + prompts[:] = [f'{args.chat_template.format(input=text)}' for text in prompts] + input_tokens = tokenizer.encode_batch(prompts) - if args.verbose: print("Prompt(s) encoded") + if args.verbose: print(f'Prompt(s) encoded: {prompts}') params = og.GeneratorParams(model) - params.set_search_options(max_length=args.max_length, top_p=args.top_p, top_k=args.top_k, temperature=args.temperature, repetition_penalty=args.repetition_penalty) - if args.cuda_graph_with_max_batch_size > 0: - params.try_use_cuda_graph_with_max_batch_size(args.cuda_graph_with_max_batch_size) + + search_options = {name:getattr(args, name) for name in ['do_sample', 'max_length', 'min_length', 'top_p', 'top_k', 'temperature', 'repetition_penalty'] if name in args} + + if (args.verbose): print(f'Args: {args}') + if (args.verbose): print(f'Search options: {search_options}') + + params.set_search_options(**search_options) + # Set the batch size for the CUDA graph to the number of prompts if the user didn't specify a batch size + params.try_use_cuda_graph_with_max_batch_size(len(prompts)) + if args.batch_size_for_cuda_graph: + params.try_use_cuda_graph_with_max_batch_size(args.batch_size_for_cuda_graph) params.input_ids = input_tokens if args.verbose: print("GeneratorParams created") @@ -37,19 +52,24 @@ def main(args): print() print() - print(f"Tokens: {len(output_tokens[0])} Time: {run_time:.2f} Tokens per second: {len(output_tokens[0])/run_time:.2f}") + total_tokens = sum(len(x) for x in output_tokens) + print(f"Tokens: {total_tokens} Time: {run_time:.2f} Tokens per second: {total_tokens/run_time:.2f}") print() if __name__ == "__main__": - parser = argparse.ArgumentParser(description="End-to-end token generation loop example for gen-ai") + parser = argparse.ArgumentParser(argument_default=argparse.SUPPRESS, description="End-to-end token generation loop example for gen-ai") parser.add_argument('-m', '--model', type=str, required=True, help='Onnx model folder path (must contain config.json and model.onnx)') - parser.add_argument('-pr', '--prompts', nargs='*', required=False, help='Input prompts to generate tokens from') - parser.add_argument('-l', '--max_length', type=int, default=512, help='Max number of tokens to generate after prompt') - parser.add_argument('-p', '--top_p', type=float, default=0.9, help='Top p probability to sample with') - parser.add_argument('-k', '--top_k', type=int, default=50, help='Top k tokens to sample from') - parser.add_argument('-t', '--temperature', type=float, default=1.0, help='Temperature to sample with') - parser.add_argument('-r', '--repetition_penalty', type=float, default=1.0, help='Repetition penalty to sample with') - parser.add_argument('-v', '--verbose', action='store_true', help='Print verbose output') - parser.add_argument('-c', '--cuda_graph_with_max_batch_size', type=int, default=0, help='Max batch size for CUDA graph') + parser.add_argument('-pr', '--prompts', nargs='*', required=False, help='Input prompts to generate tokens from. Provide this parameter multiple times to batch multiple prompts') + parser.add_argument('-i', '--min_length', type=int, help='Min number of tokens to generate including the prompt') + parser.add_argument('-l', '--max_length', type=int, help='Max number of tokens to generate including the prompt') + parser.add_argument('-ds', '--do_random_sampling', action='store_true', help='Do random sampling. When false, greedy or beam search are used to generate the output. Defaults to false') + parser.add_argument('-p', '--top_p', type=float, help='Top p probability to sample with') + parser.add_argument('-k', '--top_k', type=int, help='Top k tokens to sample from') + parser.add_argument('-t', '--temperature', type=float, help='Temperature to sample with') + parser.add_argument('-r', '--repetition_penalty', type=float, help='Repetition penalty to sample with') + parser.add_argument('-v', '--verbose', action='store_true', default=False, help='Print verbose output and timing information. Defaults to false') + parser.add_argument('-b', '--batch_size_for_cuda_graph', type=int, default=1, help='Max batch size for CUDA graph') + parser.add_argument('-c', '--chat_template', type=str, default='', help='Chat template to use for the prompt. User input will be injected into {input}. If not set, the prompt is used as is.') + args = parser.parse_args() main(args) \ No newline at end of file diff --git a/examples/python/model-qa.py b/examples/python/model-qa.py index 57ec9f6db..c1e628eb5 100644 --- a/examples/python/model-qa.py +++ b/examples/python/model-qa.py @@ -14,10 +14,15 @@ def main(args): tokenizer_stream = tokenizer.create_stream() if args.verbose: print("Tokenizer created") if args.verbose: print() + search_options = {name:getattr(args, name) for name in ['do_sample', 'max_length', 'min_length', 'top_p', 'top_k', 'temperature', 'repetition_penalty'] if name in args} - if args.chat_template.count('{') != 1 or args.chat_template.count('}') != 1: - print("Error, chat template must have exactly one pair of curly braces, e.g. '<|user|>\n{input} <|end|>\n<|assistant|>'") - exit(1) + + if args.verbose: print(search_options) + + if args.chat_template: + if args.chat_template.count('{') != 1 or args.chat_template.count('}') != 1: + print("Error, chat template must have exactly one pair of curly braces, e.g. '<|user|>\n{input} <|end|>\n<|assistant|>'") + exit(1) # Keep asking for input prompts in a loop while True: diff --git a/examples/python/phi3-qa.py b/examples/python/phi3-qa.py index 4f175c102..977a47589 100644 --- a/examples/python/phi3-qa.py +++ b/examples/python/phi3-qa.py @@ -15,6 +15,12 @@ def main(args): if args.verbose: print("Tokenizer created") if args.verbose: print() search_options = {name:getattr(args, name) for name in ['do_sample', 'max_length', 'min_length', 'top_p', 'top_k', 'temperature', 'repetition_penalty'] if name in args} + + # Set the max length to something sensible by default, unless it is specified by the user, + # since otherwise it will be set to the entire context length + if 'max_length' not in search_options: + search_options['max_length'] = 2048 + chat_template = '<|user|>\n{input} <|end|>\n<|assistant|>' # Keep asking for input prompts in a loop