refactor(utilsi): 优化项目结构

chinese-wiki.py

@@ -2,8 +2,7 @@ from utils import process_dataset
 
 if __name__ == "__main__":
     process_dataset(
-        dataset_name="wikimedia/wikipedia",
+        dataset_name="Blaze7451/Wiki-zh-20250601",
         output_subdir="chinese-wiki",
-        dataset_config="zh",
         max_chunk_size=5,
     )

to_ids.py

@@ -1,39 +1,7 @@
+from utils import process_files, fetch_files
 from tokenizer import BpeTokenizer
-from tqdm import tqdm
-from typing import List
-import json
-import pickle as pkl
-import torch
 import os
 
-def fetch_files(directory):
-    return [os.path.join(root, f) for root, _, files in os.walk(directory) for f in files]
-
-def convert_to_ids(tokenizer: BpeTokenizer, file_path, out_file_path):
-    arrows = []
-    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):
-            line = json.loads(line)
-            ids = tokenizer.encode(line["text"])
-            arrow = torch.tensor(ids, dtype=torch.int32)
-            arrows.append(arrow)
-    
-    with open(out_file_path, "wb") as f:
-        tensor = torch.cat(arrows)
-        pkl.dump(tensor, f)
-        
-def process_files(tokenizer: BpeTokenizer, files: List[str], base_out_dir):
-    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)
-        
-        if not os.path.exists(out_file_path):
-            os.makedirs(os.path.dirname(out_file_path), exist_ok=True)
-         
-        convert_to_ids(tokenizer, file_path, out_file_path)
-
 
 if __name__ == "__main__":
     tokenizer = BpeTokenizer("tokenizer.json")

utils.py

@@ -1,9 +1,43 @@
+from typing import List
 from datasets import load_dataset
+from tokenizer import BpeTokenizer
+import pickle as pkl
+import torch
+import tqdm
 import json
 import os
 import re
 
 
+def fetch_files(directory):
+    return [os.path.join(root, f) for root, _, files in os.walk(directory) for f in files]
+
+def convert_to_ids(tokenizer: BpeTokenizer, file_path, out_file_path):
+    arrows = []
+    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):
+            line = json.loads(line)
+            ids = tokenizer.encode(line["text"])
+            arrow = torch.tensor(ids, dtype=torch.int32)
+            arrows.append(arrow)
+    
+    with open(out_file_path, "wb") as f:
+        tensor = torch.cat(arrows)
+        pkl.dump(tensor, f)
+        
+def process_files(tokenizer: BpeTokenizer, files: List[str], base_out_dir):
+    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)
+        
+        if not os.path.exists(out_file_path):
+            os.makedirs(os.path.dirname(out_file_path), exist_ok=True)
+         
+        convert_to_ids(tokenizer, file_path, out_file_path)  
+
+
 def comprehensive_normalization(text):
     replacements = {
         '\u2018': "'", '\u2019': "'", '\u0060': "'",
@@ -51,6 +85,4 @@ def process_dataset(
                 json_line = {column_name : text}
                 f.write(json.dumps(json_line, ensure_ascii=False) + "\n")
 
-        print(f"Saved text chunk {i} to {output_path}")
-        
-  
+        print(f"Saved text chunk {i} to {output_path}")