time generate method

Signed-off-by: arendu <adithya.r@gmail.com>
NVIDIA · Nov 19, 2024 · 0278a01 · 0278a01
1 parent d27c4a5
commit 0278a01
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/nemo/collections/nlp/modules/common/text_generation_server.py b/nemo/collections/nlp/modules/common/text_generation_server.py
@@ -405,6 +405,7 @@ def put(self):
                     if neighbors is not None:
                         self.inference_strategy.update_neighbors(neighbors)
 
+            st = time.perf_counter()
             output = generate(
                 self.model,
                 sentences,
@@ -422,6 +423,8 @@ def put(self):
                 random_seed=random_seed,
                 **extra,
             )
+            tdiff = time.perf_counter() - st
+            logging.info(f"Chat server generation completed in {tdiff} seconds")
             for k in output:
                 if isinstance(output[k], torch.Tensor):
                     output[k] = output[k].tolist()