(_)
| 71 | |
| 72 | |
| 73 | def main(_): |
| 74 | pool = multiprocessing.Pool(FLAGS.num_threads) |
| 75 | tf.logging.info("Using hub module %s", FLAGS.bert_hub_module_path) |
| 76 | tokenizer = bert_utils.get_tokenizer(FLAGS.bert_hub_module_path) |
| 77 | preprocessor = wiki_preprocessor.Preprocessor(get_sentence_splitter(), |
| 78 | FLAGS.max_block_length, |
| 79 | tokenizer) |
| 80 | mapper = functools.partial(create_block_info, preprocessor=preprocessor) |
| 81 | block_count = 0 |
| 82 | input_paths = tf.io.gfile.glob(FLAGS.input_pattern) |
| 83 | random.shuffle(input_paths) |
| 84 | tf.logging.info("Processing %d input files.", len(input_paths)) |
| 85 | |
| 86 | tf.io.gfile.makedirs(FLAGS.output_dir) |
| 87 | blocks_path = os.path.join(FLAGS.output_dir, "blocks.tfr") |
| 88 | examples_path = os.path.join(FLAGS.output_dir, "examples.tfr") |
| 89 | titles_path = os.path.join(FLAGS.output_dir, "titles.tfr") |
| 90 | |
| 91 | with tf.python_io.TFRecordWriter(blocks_path) as blocks_writer: |
| 92 | with tf.python_io.TFRecordWriter(examples_path) as examples_writer: |
| 93 | with tf.python_io.TFRecordWriter(titles_path) as titles_writer: |
| 94 | for block_info in pool.imap_unordered(mapper, input_paths): |
| 95 | for title, block, examples in block_info: |
| 96 | blocks_writer.write(block.encode("utf-8")) |
| 97 | examples_writer.write(examples) |
| 98 | titles_writer.write(title.encode("utf-8")) |
| 99 | block_count += 1 |
| 100 | if block_count % 10000 == 0: |
| 101 | tf.logging.info("Wrote %d blocks.", block_count) |
| 102 | tf.logging.info("Wrote %d blocks in total.", block_count) |
| 103 | |
| 104 | if __name__ == "__main__": |
| 105 | app.run(main) |
nothing calls this directly
no test coverage detected
searching dependent graphs…