Add pre-generated prompts option for benchmark

benchmark/python/benchmark_e2e.py

@@ -83,6 +83,14 @@
         generator.generate_next_token()
     return tokenizer.decode(generator.get_sequence(0))
 
+# Use prompt length to get pre-defined prompt
+def get_prompt_by_length(prompt_length):
+    json_path = "prompts.json"
+    with open(json_path) as prompts_file:
+        content = prompts_file.read()
+        data = json.load(content)
+        return data[f"{prompt_length}"]
+
 def get_target_pip_package_version(target_pip_package_name_list):
     # get package name and version
     import pkg_resources
@@ -218,7 +226,7 @@
     # Get user arguments
     num_repetitions = args.repetitions
     temperature = 1.0
 
     # Get tokenizer, and model
     if args.verbose: print("Loading model... ")
     model=og.Model(f'{args.input_folder}')
@@ -232,6 +240,18 @@
         # use random tokens instead of generating a prompt using the model and then tokenizing it
         tokens = np.random.randint(100, size=(batch_size, prompt_length))
         prompt = [tokenizer.decode(tokens[0])] * batch_size
+    elif args.use_prompt_set:
+        prompt = [get_prompt_by_length(prompt_length)] * batch_size
+        tokens = tokenizer.encode_batch(prompt)
+
+        if len(tokens) > max_length:
+            # Shorten the inputs from (batch_size, tokenized_length) to (batch_size, requested_length)
+            tokens = tokens[:, :max_length]
+        elif len(tokens) < max_length:
+            # Lengthen the inputs from (batch_size, tokenized_length) to (batch_size, requested_length)
+            tokens_first_col = tokens[:, 0].unsqueeze(0).T
+            for _ in range(max_length - len(tokens)):
+                tokens = np.hstack((tokens_first_col, tokens))
     else:
         prompt = [generate_prompt(model, tokenizer, prompt_length, args.use_graph_capture)] * batch_size
         tokens = tokenizer.encode_batch(prompt)
@@ -424,6 +444,7 @@
     parser.add_argument('-mn', '--model_name', type=str, default='model_name', help='Model name defined by users')
     parser.add_argument('-pr', '--precision', type=str, default='fp16', help='Model precision for metrics info')
     parser.add_argument('--use_random_tokens', action='store_true', help='Use random tokens instead of generating a prompt')
+    parser.add_argument('--use_prompt_set', action='store_true', help='Use pre-generated prompt set instead of generating a prompt')
     args = parser.parse_args()
 
     # check max_lengths