get mask token from tokenizer (#2592)

microsoft · Dec 10, 2022 · 2076bf2 · 2076bf2
1 parent 35eabb0
commit 2076bf2
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/benchmarks/inference/bert-bench.py b/benchmarks/inference/bert-bench.py
@@ -58,7 +58,9 @@ def print_latency(latency_set, title, warmup=3):
 if dtype == torch.half:
     pipe.model.half()
 
-br = pipe("Hello I'm a [MASK] model")
+mask = pipe.tokenizer.mask_token
+
+br = pipe(f"Hello I'm a {mask} model")
 if args.deepspeed:
     pipe.model = deepspeed.init_inference(pipe.model,
                                           dtype=dtype,
@@ -74,7 +76,7 @@ def print_latency(latency_set, title, warmup=3):
 for i in range(args.trials):
     torch.cuda.synchronize()
     start = time.time()
-    r = pipe("Hello I'm a [MASK] model")
+    r = pipe(f"Hello I'm a {mask} model")
     torch.cuda.synchronize()
     end = time.time()
     responses.append(r)