ViperEkura/DataPipeline
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

优化处理形式

Browse Source
This commit is contained in:
ViperEkura
2025-06-30 22:00:15 +08:00
parent 6238764047
commit b508402724
1 changed files with 8 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
to_ids.py
+8 -5
View File
@@ -12,23 +12,26 @@ def fetch_files(directory):
def convert_to_ids(tokenizer: BpeTokenizer, file_path, out_file_path): def convert_to_ids(tokenizer: BpeTokenizer, file_path, out_file_path):
arrows = [] arrows = []
with open(file_path, "r") as f: 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) line = json.loads(line)
ids = tokenizer.encode(line["text"]) 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: with open(out_file_path, "w") as f:
tensor = torch.tensor(arrows, dtype=torch.int32) tensor = torch.cat(arrows)
pkl.dump(tensor, f) pkl.dump(tensor, f)
def process_files(tokenizer: BpeTokenizer, files: List[str], base_out_dir): 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_name = os.path.basename(file_path).replace(".jsonl", ".pkl")
out_file_path = os.path.join(base_out_dir, out_file_name) out_file_path = os.path.join(base_out_dir, out_file_name)
if not os.path.exists(out_file_path): if not os.path.exists(out_file_path):
os.makedirs(os.path.dirname(out_file_path), exist_ok=True) os.makedirs(os.path.dirname(out_file_path), exist_ok=True)
convert_to_ids(tokenizer, file_path, out_file_path) convert_to_ids(tokenizer, file_path, out_file_path)