优化处理形式

2025-06-30 22:00:15 +08:00
parent 6238764047
commit b508402724
1 changed files with 8 additions and 5 deletions
@@ -12,17 +12,20 @@ def fetch_files(directory):
 def convert_to_ids(tokenizer: BpeTokenizer, file_path, out_file_path):
    arrows = []
    with open(file_path, "r") as f:
-        for line in f:
+        lines = f.readlines()
+        file_name = os.path.basename(file_path)
+        for line in tqdm(lines, desc=f"Processing {file_name}", leave=False):
            line = json.loads(line)
            ids = tokenizer.encode(line["text"])
-            arrows.append(ids)
+            arrow = torch.tensor(ids, dtype=torch.int32)
+            arrows.append(arrow)
    
    with open(out_file_path, "w") as f:
-        tensor = torch.tensor(arrows, dtype=torch.int32)
+        tensor = torch.cat(arrows)
        pkl.dump(tensor, f)
        
 def process_files(tokenizer: BpeTokenizer, files: List[str], base_out_dir):
-    for file_path in tqdm(files, desc="Processing files", total=len(files)):
+    for file_path in files:
        out_file_name = os.path.basename(file_path).replace(".jsonl", ".pkl")
        out_file_path = os.path.join(base_out_dir, out_file_name)