fix bugs of data preprocessing with multiple json keys

NVIDIA · Dec 25, 2024 · 96f966d · 96f966d
1 parent 2da43ef
commit 96f966d
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/tools/preprocess_data.py b/tools/preprocess_data.py
@@ -11,8 +11,6 @@
 import time
 import gzip
 import glob
-import torch
-import numpy as np
 import multiprocessing
 try:
     import nltk
@@ -184,7 +182,8 @@ def process_json_file(self, file_name):
             self.print_processing_stats(i, proc_start, total_bytes_processed)
 
         fin.close()
-        builders[key].finalize(output_idx_files[key])
+        for key in builders.keys():
+            builders[key].finalize(output_idx_files[key])
 
 
 def get_args():