diff --git a/belle_sft.py b/belle_sft.py index c783072..50aab61 100644 --- a/belle_sft.py +++ b/belle_sft.py @@ -32,6 +32,6 @@ if __name__ == "__main__": dataset = load_dataset("BelleGroup/train_3.5M_CN") process_dataset( dataset_dict=dataset, - output_subdir="belle_sft", + output_subdir="belle-sft", process_func=process_func, ) \ No newline at end of file diff --git a/sft_file.py b/sft_file.py index e69de29..66579ef 100644 --- a/sft_file.py +++ b/sft_file.py @@ -0,0 +1,15 @@ +from utils import dump_pkl_files, fetch_files +from tokenizer import BpeTokenizer +import os + +if __name__ == "__main__": + tokenizer = BpeTokenizer("tokenizer.json") + base_dir = [ + os.path.join("dataset", "belle-sft"), + ] + base_out_dir = "pkl_output" + files = [] + for dir_path in base_dir: + files.extend(fetch_files(dir_path)) + + dump_pkl_files(tokenizer, files, base_out_dir, packing_size=2048) \ No newline at end of file diff --git a/utils.py b/utils.py index ebb1f6a..f4701ab 100644 --- a/utils.py +++ b/utils.py @@ -38,7 +38,7 @@ def dump_pkl_files( ): def process_line(line: str) -> Tensor: line = json.loads(line)[key] - processed_line = encoder(line) + processed_line = encoder(line) if encoder else line ids = tokenizer.encode(processed_line) arrow = torch.tensor(ids, dtype=torch.int32) return arrow