diff --git a/chinese-wiki.py b/chinese-wiki.py index dca111f..2e08cc1 100644 --- a/chinese-wiki.py +++ b/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, ) \ No newline at end of file diff --git a/to_ids.py b/to_ids.py index 62aa857..ed4c038 100644 --- a/to_ids.py +++ b/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") diff --git a/utils.py b/utils.py index 656f3d1..9de8d0e 100644 --- a/utils.py +++ b/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}")