make truncation of token lists optional in example script (#317)

Signed-off-by: Max de Bayser <[email protected]>

Author: maxdebayser

examples/offline_inference/long_context.py

@@ -47,8 +47,14 @@
 parser.add_argument("--num-prompts", "-n", type=int, default=8)
 parser.add_argument("--compare-with-cpu",
                     action=argparse.BooleanOptionalAction)
+parser.add_argument("--trunc_print_len",
+                    "--trunc-print-len",
+                    type=int,
+                    required=False)
 args = parser.parse_args()
 
+trunc = args.trunc_print_len
+
 max_num_seqs = args.max_num_seqs  # defines the max batch size
 assert args.max_prompt_len < args.max_model_len
 
@@ -144,8 +150,8 @@ def round_up(t):
 print("Time elapsed for all prompts is %.2f sec" % (time.time() - t0))
 print("===============")
 for output, prompt in zip(outputs, prompts):
-    generated_text = output.outputs[0].text[:100]
-    prompt = prompt[:100]
+    generated_text = output.outputs[0].text[:trunc]
+    prompt = prompt[:trunc]
     print(f"\nPrompt:\n {prompt!r}")
     print(f"\nGenerated text (truncated):\n {generated_text!r}\n")
     print("-----------------------------------")
@@ -177,9 +183,9 @@ def round_up(t):
             any_differ = True
             spyre_output = outputs[i].outputs[0].text
             print(f"Results for prompt {i} differ on cpu")
-            print(f"\nPrompt:\n {prompt[:100]!r}")
-            print(f"\nSpyre generated text:\n {spyre_output[:100]!r}\n")
-            print(f"\nCPU generated text:\n {hf_generated_text[:100]!r}\n")
+            print(f"\nPrompt:\n {prompt[:trunc]!r}")
+            print(f"\nSpyre generated text:\n {spyre_output[:trunc]!r}\n")
+            print(f"\nCPU generated text:\n {hf_generated_text[:trunc]!r}\n")
             print("-----------------------------------")
 
     if not any_differ: