refactor(to_ids): 重构数据处理流程并支持自定义处理器

to_ids.py

@@ -2,17 +2,20 @@ from utils import dump_pkl_files, fetch_files
 from tokenizer import BpeTokenizer
 import os
 
+def processor(intput_str: str):
+    return f"{intput_str} <eos>"
 
 if __name__ == "__main__":
     tokenizer = BpeTokenizer("tokenizer.json")
     base_dir = [
-        os.path.join("dataset", "chinese-c4"),
+        # os.path.join("dataset", "chinese-c4"),
-        os.path.join("dataset", "english-fineweb")
+        # os.path.join("dataset", "english-fineweb"),
+        # os.path.join("dataset", "english-wiki"),
+        os.path.join("dataset", "chinese-wiki"),
     ]
     base_out_dir = "pkl_output"
-    
     files = []
     for dir_path in base_dir:
         files.extend(fetch_files(dir_path))
     
-    dump_pkl_files(tokenizer, files, base_out_dir)
+    dump_pkl_files(tokenizer, files, base_out_dir, processor)

utils.py

@@ -1,9 +1,10 @@
 from typing import List, Callable
 from datasets import load_dataset
 from tokenizer import BpeTokenizer
+from tqdm import tqdm
+from torch import Tensor
 import pickle as pkl
 import torch
-import tqdm
 import json
 import os
 import re
@@ -29,25 +30,28 @@ def dump_pkl_files(
     tokenizer: BpeTokenizer, 
     files: List[str], 
     base_out_dir: str, 
-    encder: Callable[[str], str]=None
+    encoder: Callable[[str], str]=None,
+    key: str='text',
 ):  
+    def process_line(line: str) -> Tensor:
+            line = json.loads(line)[key]
+            processed_line = encoder(line)
+            ids = tokenizer.encode(processed_line)
+            arrow = torch.tensor(ids, dtype=torch.int32)
+            return arrow
+    
     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)
-        
+        file_name = os.path.basename(file_path)
-        if not os.path.exists(out_file_path):
-            os.makedirs(os.path.dirname(out_file_path), exist_ok=True)
-        
         arrows = []
+        
+        os.makedirs(os.path.dirname(out_file_path), exist_ok=True)
         with open(file_path, "r") as f:    
             lines = f.readlines()
-            file_name = os.path.basename(file_path)
+        for line in tqdm(lines, desc=f"Processing {file_name}", leave=False):
-            for line in tqdm(lines, desc=f"Processing {file_name}", leave=False):
+            arrow = process_line(line)
-                line = json.loads(line)
+            arrows.append(arrow)
-                processed_line = encder(line)
-                ids = tokenizer.encode(processed_line)
-                arrow = torch.tensor(ids, dtype=torch.int32)
-                arrows.append(arrow)
         
         with open(out_file_path, "wb") as f:
             tensor = torch.cat(arrows)