31 lines
946 B
Python
31 lines
946 B
Python
from modules.utils import dump_pkl_files, fetch_files
|
|
from modules.tokenizer import BpeTokenizer
|
|
import torch
|
|
import os
|
|
|
|
def get_processor(tokenizer: BpeTokenizer):
|
|
def processor(intput_dict: dict) -> dict:
|
|
segment = intput_dict["text"]
|
|
ids = tokenizer.encode(f"{segment} <eos>")
|
|
t_ids = torch.tensor(ids, dtype=torch.int32)
|
|
|
|
return {'sequence': t_ids}
|
|
|
|
return processor
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tokenizer = BpeTokenizer("tokenizer.json")
|
|
base_dir = [
|
|
os.path.join("dataset", "chinese-c4"),
|
|
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))
|
|
processor = get_processor(tokenizer)
|
|
dump_pkl_files(files, base_out_dir, processor, ["text"]) |