()
| 179 | |
| 180 | |
| 181 | def main(): |
| 182 | args = get_args() |
| 183 | args.tensor_model_parallel_size = 1 |
| 184 | args.rank = 0 |
| 185 | args.make_vocab_size_divisible_by = 128 |
| 186 | args.vocab_extra_ids = 0 |
| 187 | encoder = Encoder(args) |
| 188 | tokenizer = build_tokenizer(args) |
| 189 | print(f'Vocab size: {tokenizer.vocab_size}') |
| 190 | print(f'Output prefix: {args.output_prefix}') |
| 191 | |
| 192 | semaphore = Semaphore(10000 + args.workers) |
| 193 | |
| 194 | # use multiprocessing to iterate over input documents |
| 195 | file_list = os.listdir(args.input) |
| 196 | path_list = [os.path.join(args.input, file) for file in file_list] |
| 197 | fin = yield_from_files(path_list, semaphore) |
| 198 | |
| 199 | if args.workers > 1: |
| 200 | pool = multiprocessing.Pool(args.workers, |
| 201 | initializer=encoder.initializer) |
| 202 | encoded_docs = pool.imap(encoder.encode, fin, chunksize=25) |
| 203 | else: |
| 204 | encoder.initializer() |
| 205 | encoded_docs = (encoder.encode(doc) for doc in fin) |
| 206 | |
| 207 | output_bin_files = {} |
| 208 | output_idx_files = {} |
| 209 | builders = {} |
| 210 | for key in args.jsonl_keys: |
| 211 | output_bin_files[key] = '{}_{}_{}.bin'.format(args.output_prefix, key, |
| 212 | 'document') |
| 213 | output_idx_files[key] = '{}_{}_{}.idx'.format(args.output_prefix, key, |
| 214 | 'document') |
| 215 | builders[key] = indexed_dataset.MMapIndexedDatasetBuilder( |
| 216 | output_bin_files[key], |
| 217 | dtype=indexed_dataset.DType.optimal_dtype(tokenizer.vocab_size), |
| 218 | ) |
| 219 | |
| 220 | # actually do tokenization |
| 221 | proc_start = time.time() |
| 222 | total_bytes_processed = 0 |
| 223 | pbar = tqdm.tqdm() |
| 224 | for i, (doc, bytes_processed) in enumerate(encoded_docs, start=1): |
| 225 | total_bytes_processed += bytes_processed |
| 226 | |
| 227 | semaphore.release() |
| 228 | |
| 229 | # add each tokenized document / sentence |
| 230 | for key, sentences in doc.items(): |
| 231 | for sentence in sentences: |
| 232 | builders[key].add_item(torch.IntTensor(sentence)) |
| 233 | # separate with eos token |
| 234 | builders[key].end_document() |
| 235 | |
| 236 | # log progress |
| 237 | if i % args.log_interval == 0: |
| 238 | current = time.time() |
no test coverage detected