Merge branch 'aligner/nemotron5' of https://github.com/NVIDIA/NeMo into aligner/nemotron5

Author: arendu

examples/nlp/language_modeling/megatron_mamba_eval.py

@@ -17,6 +17,7 @@
 import json
 import os
 import threading
+import time
 from functools import partial
 
 import torch
@@ -352,23 +353,31 @@ def main(cfg) -> None:
     prompts = load_prompts(cfg)
 
     # First method of running text generation, call model.generate method
-    response = model.generate(inputs=prompts, length_params=length_params, sampling_params=sampling_params)
+    for i in range(3):
+        st = time.perf_counter()
+        response = model.generate(inputs=prompts, length_params=length_params, sampling_params=sampling_params)
+        tdiff = time.perf_counter() - st
+        print(f"[Try{i} model.generate took {tdiff} seconds...")
 
-    print("***************************")
-    print(response)
-    print("***************************")
+    # print("***************************")
+    # print(response)
+    # print("***************************")
 
     # Second method of running text generation, call trainer.predict [recommended]
     bs = 2
     ds = RequestDataSet(prompts)
     request_dl = DataLoader(dataset=ds, batch_size=bs)
     config = OmegaConf.to_container(cfg.inference)
     model.set_inference_config(config)
-    response = trainer.predict(model, request_dl)
-
-    print("***************************")
-    print(response)
-    print("***************************")
+    for i in range(3):
+        st = time.perf_counter()
+        response = trainer.predict(model, request_dl)
+        tdiff = time.perf_counter() - st
+        print(f"[Try{i} trainer.predict took {tdiff} seconds...")
+
+    # print("***************************")
+    # print(response)
+    # print("***************************")
 
     # Third method of running text generation, use inference server
     if cfg.server:

nemo/collections/nlp/modules/common/text_generation_server.py

@@ -171,15 +171,17 @@ def chat_completion(self, data):
         if OmegaConf.select(self.model.cfg, "data.chat_prompt_tokens") is not None:
             special_tokens = self.model.cfg.data.chat_prompt_tokens
         else:
-            #raise RuntimeError(
+            # raise RuntimeError(
             #    "You don't have a model (model_config.yaml) which has chat_prompt_tokens, are you sure this is a Chat/Instruction model?"
-            #)
+            # )
             # (@adithyare) hacking in the special tokens to test non-chat models for debugging
-            special_tokens = {"system_turn_start": "<SPECIAL_10>",
-                              "turn_start": "<SPECIAL_11>",
-                              "label_start": "<SPECIAL_12>",
-                              "end_of_name": "\n",
-                              "end_of_turn": "\n"}
+            special_tokens = {
+                "system_turn_start": "<SPECIAL_10>",
+                "turn_start": "<SPECIAL_11>",
+                "label_start": "<SPECIAL_12>",
+                "end_of_name": "\n",
+                "end_of_turn": "\n",
+            }
         nemo_source = self.convert_messages(data['messages'])
         header, conversation, data_type, mask_role = _get_header_conversation_type_mask_role(
             nemo_source, special_tokens
@@ -432,7 +434,7 @@ def put(self):
         # (@adithyare) resolves a json byte conversion issue (taken from chat_completeion)
         for i in range(len(output['tokens'])):
             tokens = output['tokens'][i]
-            output['tokens'][i] = [t.decode('utf-8', errors='replace') if isinstance(t, bytes) else t for t in tokens]      
+            output['tokens'][i] = [t.decode('utf-8', errors='replace') if isinstance(t, bytes) else t for t in tokens]
 
         if not all_probs:
             del output['full_logprob']