refactor(utilsi): 优化项目结构

This commit is contained in:
2025-07-14 14:41:42 +08:00
parent 1c132826c0
commit 7e53d35099
3 changed files with 37 additions and 38 deletions
+1 -33
View File
@@ -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")