def extract_references_from_wets(wet_files, metadata_dir, out_dir, tmp_dir=None): """Extract references from WET files into sharded output files.""" # Setup output files shard_files = make_ref_shard_files(out_dir) num_refs = 0 total = len(wet_files) for i, wet_file in enumerate(wet_files): num_refs_in_wet = 0 tf.logging.info(f"Processing file {i} of {total}") # Read metadata file metadata_fname = os.path.join( metadata_dir, os.path.basename(wet_file)) + cc_utils.METADTA_SUFFIX with tf.gfile.Open(cc_utils.readahead(metadata_fname)) as f: wet_metadata = json.loads(f.read()) if not wet_metadata: # No references in this WET file continue if wet_file.startswith("http"): # download if not tmp_dir: tmp_dir = tempfile.gettempdir() record_gen = cc_utils.wet_records_from_url(wet_file, tmp_dir) else: # local record_gen = cc_utils.wet_records_from_file_obj( cc_utils.gzip_memfile(wet_file), take_ownership=True) jj = 0 for wet_record in record_gen: jj += 1 # these are shard ids that contain one or more Wikipedia # articles that cite this reference shard_ids = wet_metadata.get(wet_record.url) if not shard_ids: # URL not in dataset continue # Serialize and write out ex = _make_example_from_record(wet_record) ex_str = ex.SerializeToString() for shard_id in shard_ids: shard_files[shard_id].write(ex_str) shard_files[shard_id].flush() num_refs += 1 num_refs_in_wet += 1 tf.logging.info("Wrote out %d references for this WET", num_refs_in_wet) tf.logging.info("Wrote out %d references total", num_refs) # Cleanup for shard_file in shard_files: shard_file.close()
def extract_references_from_wets(wet_files, metadata_dir, out_dir, tmp_dir=None): """Extract references from WET files into sharded output files.""" # Setup output files shard_files = make_ref_shard_files(out_dir) num_refs = 0 for i, wet_file in enumerate(wet_files): num_refs_in_wet = 0 tf.logging.info("Processing file %d", i) # Read metadata file metadata_fname = os.path.join( metadata_dir, os.path.basename(wet_file)) + cc_utils.METADTA_SUFFIX with tf.gfile.Open(cc_utils.readahead(metadata_fname)) as f: wet_metadata = json.loads(f.read()) if not wet_metadata: # No references in this WET file continue if wet_file.startswith("http"): # download if not tmp_dir: tmp_dir = tempfile.gettempdir() record_gen = cc_utils.wet_records_from_url(wet_file, tmp_dir) else: # local record_gen = cc_utils.wet_records_from_file_obj( cc_utils.gzip_memfile(wet_file), take_ownership=True) for wet_record in record_gen: shard_ids = wet_metadata.get(wet_record.url) if not shard_ids: # URL not in dataset continue # Serialize and write out ex = _make_example_from_record(wet_record) ex_str = ex.SerializeToString() for shard_id in shard_ids: shard_files[shard_id].write(ex_str) num_refs += 1 num_refs_in_wet += 1 tf.logging.info("Wrote out %d references for this WET", num_refs_in_wet) tf.logging.info("Wrote out %d references total", num_refs) # Cleanup for shard_file in shard_files: shard_file.close()