(
tokenizer_path: str,
dataset_path: Union[str, List[str]],
output_prefix: str,
language: str = None,
mode: str = "pretrain",
discard_overlong: bool = False,
sliding_stride: int = 200,
num_workers: int = 32,
seq_len: int = 2048,
)
| 79 | |
| 80 | |
| 81 | def main( |
| 82 | tokenizer_path: str, |
| 83 | dataset_path: Union[str, List[str]], |
| 84 | output_prefix: str, |
| 85 | language: str = None, |
| 86 | mode: str = "pretrain", |
| 87 | discard_overlong: bool = False, |
| 88 | sliding_stride: int = 200, |
| 89 | num_workers: int = 32, |
| 90 | seq_len: int = 2048, |
| 91 | ): |
| 92 | DATA_KEYS = ["input_ids", "attention_mask", "labels"] |
| 93 | |
| 94 | # create output dir |
| 95 | os.makedirs(os.path.dirname(output_prefix), exist_ok=True) |
| 96 | |
| 97 | tokenizer = CodeGeeXTokenizer(tokenizer_path=tokenizer_path) |
| 98 | pad_token_id = tokenizer.eos_token_id |
| 99 | |
| 100 | dataset = load_pretrain_dataset(dataset_path) |
| 101 | prompt_dataset = generate_prompt_samples(dataset, language=language, mode=mode) |
| 102 | |
| 103 | if num_workers == 0: |
| 104 | num_workers = multiprocessing.cpu_count() |
| 105 | pool = multiprocessing.Pool(num_workers) |
| 106 | output_bin_files = {} |
| 107 | output_idx_files = {} |
| 108 | builders = {} |
| 109 | |
| 110 | for key in DATA_KEYS: |
| 111 | output_bin_files[key] = "{}_{}.bin".format(output_prefix, key) |
| 112 | output_idx_files[key] = "{}_{}.idx".format(output_prefix, key) |
| 113 | builders[key] = make_mmap_builder( |
| 114 | output_bin_files[key], |
| 115 | vocab_size=None, # magic number, should change it |
| 116 | ) |
| 117 | |
| 118 | # NOTE that we use seq_len + 1 instead of seq_len, since the input tokens will be shifted by one. |
| 119 | processor = PromptDatasetProcessor( |
| 120 | tokenize=tokenizer.encode_code, |
| 121 | pad_token=pad_token_id, |
| 122 | max_seq_len=seq_len + 1, |
| 123 | discard_overlong=discard_overlong, |
| 124 | sliding_stride=sliding_stride, |
| 125 | eod_token=pad_token_id) |
| 126 | |
| 127 | processor.start_time = perf_counter() |
| 128 | doc_iter = pool.imap_unordered(processor.process_sample_strict, |
| 129 | prompt_dataset, |
| 130 | chunksize=20) |
| 131 | |
| 132 | for doc_idx, docs in tqdm(enumerate(doc_iter, start=1)): |
| 133 | processor.doc_processed += 1 |
| 134 | for doc in docs: |
| 135 | processor.doc_generated += 1 |
| 136 | for key in DATA_KEYS: |
| 137 | builders[key].add_item(torch.IntTensor(doc[key])) |
| 138 |
nothing calls this directly
no test coverage detected