Merging with main and fixing merge conflict

tools/preprocess_data.py

@@ -122,8 +122,10 @@ def get_args():
                        choices=['lazy', 'cached', 'mmap'])
 
     group = parser.add_argument_group(title='runtime')
-    group.add_argument('--workers', type=int, default=1,
+    group.add_argument('--workers', type=int, required=True,
                        help='Number of worker processes to launch')
+    group.add_argument('--chunk-size', type=int, required=True,
+                       help='Chunk size assigned to each worker process')
     group.add_argument('--log-interval', type=int, default=100,
                        help='Interval between progress updates')
     args = parser.parse_args()
@@ -154,7 +156,7 @@ def main():
     encoder = Encoder(args)
     tokenizer = build_tokenizer(args)
     pool = multiprocessing.Pool(args.workers, initializer=encoder.initializer)
-    encoded_docs = pool.imap(encoder.encode, fin, 25)
+    encoded_docs = pool.imap(encoder.encode, fin, args.chunk_size)
     #encoded_docs = map(encoder.encode, fin)
 
     level = "document"

tools/run_text_generation_server.py

@@ -28,6 +28,7 @@ from megatron.model import GPTModel
 from megatron.training import get_model
 from megatron.text_generation_server import MegatronServer
 from megatron.text_generation import generate_and_post_process
+from megatron.text_generation import beam_search_and_post_process
 import torch
 
 def model_provider(pre_process=True, post_process=True):
@@ -82,3 +83,8 @@ if __name__ == "__main__":
                 generate_and_post_process(model)
             except ValueError as ve:
                 pass
+        elif choice[0].item() == 1:
+            try:
+                beam_search_and_post_process(model)
+            except ValueError as ve:
+                pass