import glob import json import logging import math import os import tarfile from copy import deepcopy import git from datasets import Dataset, Image, Sequence, Value, concatenate_datasets, load_from_disk from pathos.multiprocessing import ProcessingPool as Pool from tqdm import tqdm logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", ) logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) def write_file(path_file, to_write): f = open(path_file, "w") f.truncate(0) f.write(to_write) f.close() def html_to_web_documents( dataset, dom_tree_simplificator, pre_extraction_simplificator, num_proc, html_column_name="html", url_column_name="url", ): def func_html_to_web_documents(example): html_str = example[html_column_name] page_url = example[url_column_name] general_metadata = {} if all( [ column_name in example for column_name in ["url", "warc_filename", "warc_record_offset", "warc_record_length"] ] ): general_metadata = { "url": example["url"], "warc_filename": example["warc_filename"], "warc_record_offset": example["warc_record_offset"], "warc_record_length": example["warc_record_length"], } try: selectolax_tree = dom_tree_simplificator(html_str, type_return="selectolax_tree") list_nodes = pre_extraction_simplificator(selectolax_tree, page_url=page_url) except Exception: print("EXCEPTION") example["texts"] = [] example["images"] = [] example["metadata"] = json.dumps([]) example["general_metadata"] = json.dumps([]) return example texts = [] images = [] metadata = [] for node in list_nodes: if node.tag == "-text": texts.append(node.text) images.append("") metadata.append(None) elif node.tag == "img": texts.append(None) images.append(node.media_info["src"]) metadata.append(node.media_info) example["texts"] = texts example["images"] = images example["metadata"] = json.dumps(metadata) example["general_metadata"] = json.dumps(general_metadata) return example logger.info("Starting extracting the documents") dataset = dataset.map(func_html_to_web_documents, num_proc=num_proc, remove_columns=dataset.column_names) logger.info("Finished extracting the documents") return dataset def get_image_urls(dataset, num_proc, path_save_file_image_urls): def func_get_image_urls(example): example["urls"] = [el for el in example["images"] if el] return example logger.info("Starting getting the urls of all images") image_urls = dataset.map(func_get_image_urls, remove_columns=dataset.column_names, num_proc=num_proc) image_urls = [sub_el for el in image_urls["urls"] for sub_el in el if sub_el] image_urls = list(set(image_urls)) write_file(path_file=path_save_file_image_urls, to_write="\n".join(image_urls)) logger.info("Finished getting the urls of all images") def download_images( path_save_file_image_urls, path_save_dir_downloaded_images, number_sample_per_shard, image_size, resize_mode, num_proc, thread_count, ): # Before calling this method, set up a DNS solver # https://github.com/rom1504/img2dataset#setting-up-a-bind9-resolver logger.info("Starting downloading the images") os.system( "img2dataset" f" --url_list={path_save_file_image_urls} --output_folder={path_save_dir_downloaded_images}" f" --processes_count={num_proc} --thread_count={thread_count}" f" --number_sample_per_shard={number_sample_per_shard} --image_size={image_size}" f" --resize_mode={resize_mode} --output_format=webdataset" ) logger.info("Finished downloading the images") def create_dataset_images_from_tar( tar_paths, path_save_dir_tmp_datasets_images, num_proc, path_save_file_map_url_idx, path_save_dir_dataset_images, ): def process_one_tar(args): (tar_path, idx_tar) = args with tarfile.open(tar_path) as tar_file: tar_members = tar_file.getmembers() name_to_url = {} name_to_img = {} url_to_img = {} for tar_member in tar_members: if tar_member.name.endswith(".jpg"): name = tar_member.name.replace(".jpg", "") tar_member_file = tar_file.extractfile(tar_member) img = tar_member_file.read() tar_member_file.close() name_to_img[name] = img elif tar_member.name.endswith(".json"): name = tar_member.name.replace(".json", "") tar_member_file = tar_file.extractfile(tar_member) json_val = json.loads(tar_member_file.read()) status = json_val["status"] url = json_val["url"] tar_member_file.close() if status == "success": # Should always happend with webdataset format, not with parquet name_to_url[name] = url for name in name_to_url: url_to_img[name_to_url[name]] = name_to_img[name] new_urls_indexed = list(url_to_img.keys()) new_datasets_images = Dataset.from_dict( {"url": list(url_to_img.keys()), "image": list(url_to_img.values())} ) # We need to save the new datasets and then reload them, since `from_dict` store the dataset # in the RAM and does not use the disk space new_datasets_images.save_to_disk(os.path.join(path_save_dir_tmp_datasets_images, str(idx_tar))) return new_urls_indexed logger.info("Starting creating the dataset of all images") args_pool = [(tar_path, idx_tar) for idx_tar, tar_path in enumerate(tar_paths)] pool = Pool(num_proc) urls_indexed = pool.map(process_one_tar, args_pool) urls_indexed = [sub_el for el in urls_indexed for sub_el in el] map_url_idx = {url: idx for idx, url in enumerate(urls_indexed)} with open(path_save_file_map_url_idx, "w") as f: json.dump(map_url_idx, f) datasets_images = [ load_from_disk(os.path.join(path_save_dir_tmp_datasets_images, str(idx_tar))) for idx_tar in range(len(tar_paths)) ] dataset_images = concatenate_datasets(datasets_images) dataset_images.save_to_disk(path_save_dir_dataset_images) logger.info("Finished creating the dataset of all images") return dataset_images def create_dataset_images( path_save_dir_downloaded_images, path_save_dir_tmp_datasets_images, num_proc, path_save_file_map_url_idx, path_save_dir_dataset_images, ): tar_paths = glob.glob(os.path.join(path_save_dir_downloaded_images, "*.tar")) dataset_images = create_dataset_images_from_tar( tar_paths=tar_paths, path_save_dir_tmp_datasets_images=path_save_dir_tmp_datasets_images, num_proc=num_proc, path_save_file_map_url_idx=path_save_file_map_url_idx, path_save_dir_dataset_images=path_save_dir_dataset_images, ) return dataset_images def urls_to_images(dataset, dataset_images, map_url_idx, num_proc, some_urls_are_already_retrieved=False): if some_urls_are_already_retrieved: if "images_urls" not in dataset.features or "images" not in dataset.features: raise ValueError( "If some urls are already retrieved, the dataset must contain the features 'images_urls' and 'images'" ) def retrieve_image(url): if url not in map_url_idx: return None image = {"path": None, "bytes": dataset_images[map_url_idx[url]]["image"]} return image def func_urls_to_images_urls_in_images_col(example): example["images_urls"] = deepcopy(example["images"]) num_urls = sum([(url is not None and url != "") for url in example["images_urls"]]) example["images"] = [retrieve_image(url) if url else None for url in example["images"]] num_found = sum([img is not None for img in example["images"]]) num_not_found = num_urls - num_found example["num_found"] = num_found example["num_not_found"] = num_not_found return example def func_urls_to_images_urls_in_images_urls_col(example): num_urls = sum([(url is not None and url != "") for url in example["images_urls"]]) example["images"] = [ img if img is not None else retrieve_image(url) if url else None for img, url in zip(example["images"], example["images_urls"]) ] num_found = sum([img is not None for img in example["images"]]) num_not_found = num_urls - num_found example["num_found"] = num_found example["num_not_found"] = num_not_found return example func_urls_to_images = ( func_urls_to_images_urls_in_images_urls_col if some_urls_are_already_retrieved else func_urls_to_images_urls_in_images_col ) logger.info("Starting replacing urls by images") new_features = deepcopy(dataset.features) new_features["images"] = Sequence(Image()) new_features["images_urls"] = Sequence(Value("string")) new_features["num_found"] = Value("int32") new_features["num_not_found"] = Value("int32") dataset = dataset.map( func_urls_to_images, features=new_features, num_proc=num_proc, load_from_cache_file=False, ) logger.info("Finished replacing urls by images") return dataset def save_split_sharded_already_splitted_dataset(dataset, path_save_dir_sharded_dataset, shard_size): def save_split_ds(split_dataset, split_name): num_shards = math.ceil(len(split_dataset) / shard_size) for idx in tqdm(range(num_shards)): shard = split_dataset.shard(num_shards=num_shards, index=idx, contiguous=True) shard.save_to_disk(os.path.join(path_save_dir_sharded_dataset, split_name, f"shard_{idx}")) os.makedirs(path_save_dir_sharded_dataset, exist_ok=True) f = open(os.path.join(path_save_dir_sharded_dataset, "dataset_dict.json"), "w") f.write('{"splits": ["train", "valid"]}') f.close() os.makedirs(os.path.join(path_save_dir_sharded_dataset, "train"), exist_ok=True) os.makedirs(os.path.join(path_save_dir_sharded_dataset, "valid"), exist_ok=True) logger.info("Starting sharding the dataset") train_dataset = dataset["train"] valid_dataset = dataset["valid"] save_split_ds(train_dataset, "train") save_split_ds(valid_dataset, "valid") logger.info("Finished sharding the dataset") def save_split_sharded_dataset(dataset, path_save_dir_sharded_dataset, shard_size): os.makedirs(path_save_dir_sharded_dataset, exist_ok=True) f = open(os.path.join(path_save_dir_sharded_dataset, "dataset_dict.json"), "w") f.write('{"splits": ["train", "valid"]}') f.close() os.makedirs(os.path.join(path_save_dir_sharded_dataset, "train"), exist_ok=True) os.makedirs(os.path.join(path_save_dir_sharded_dataset, "valid"), exist_ok=True) logger.info("Starting sharding the dataset") num_shards = math.ceil(len(dataset) / shard_size) for idx in tqdm(range(num_shards)): shard = dataset.shard(num_shards=num_shards, index=idx, contiguous=True) if idx < 2: shard.save_to_disk(os.path.join(path_save_dir_sharded_dataset, "valid", f"shard_{idx}")) else: shard.save_to_disk(os.path.join(path_save_dir_sharded_dataset, "train", f"shard_{idx}")) logger.info("Finished sharding the dataset") class CommonCrawlWebDocumentExtractor: def __init__( self, html_dataset, dom_tree_simplificator, pre_extraction_simplificator, path_save_dir_dataset, num_proc, path_save_file_image_urls, path_save_dir_downloaded_images, thread_count, number_sample_per_shard, image_size, resize_mode, path_save_dir_tmp_datasets_images, path_save_dir_dataset_images, path_save_file_map_url_idx, num_proc_urls_to_images, path_save_dir_sharded_dataset, shard_size, ): self.dataset = html_dataset self.dom_tree_simplificator = dom_tree_simplificator self.pre_extraction_simplificator = pre_extraction_simplificator self.path_save_dir_dataset = path_save_dir_dataset self.num_proc = num_proc self.path_save_file_image_urls = path_save_file_image_urls self.path_save_dir_downloaded_images = path_save_dir_downloaded_images self.thread_count = thread_count self.number_sample_per_shard = number_sample_per_shard self.image_size = image_size self.resize_mode = resize_mode self.path_save_dir_tmp_datasets_images = path_save_dir_tmp_datasets_images self.path_save_dir_dataset_images = path_save_dir_dataset_images self.path_save_file_map_url_idx = path_save_file_map_url_idx self.num_proc_urls_to_images = num_proc_urls_to_images self.path_save_dir_sharded_dataset = path_save_dir_sharded_dataset self.shard_size = shard_size def html_to_web_documents(self): self.dataset = html_to_web_documents( dataset=self.dataset, dom_tree_simplificator=self.dom_tree_simplificator, pre_extraction_simplificator=self.pre_extraction_simplificator, num_proc=self.num_proc, ) def get_image_urls(self): get_image_urls( dataset=self.dataset, num_proc=self.num_proc, path_save_file_image_urls=self.path_save_file_image_urls ) def download_images(self): download_images( path_save_file_image_urls=self.path_save_file_image_urls, path_save_dir_downloaded_images=self.path_save_dir_downloaded_images, number_sample_per_shard=self.number_sample_per_shard, image_size=self.image_size, resize_mode=self.resize_mode, num_proc=self.num_proc, thread_count=self.thread_count, ) def create_dataset_images(self): self.dataset_images = create_dataset_images( path_save_dir_downloaded_images=self.path_save_dir_downloaded_images, path_save_dir_tmp_datasets_images=self.path_save_dir_tmp_datasets_images, num_proc=self.num_proc, path_save_file_map_url_idx=self.path_save_file_map_url_idx, path_save_dir_dataset_images=self.path_save_dir_dataset_images, ) def urls_to_images(self, reload_files=False): with open(self.path_save_file_map_url_idx) as f: self.map_url_idx = json.load(f) # Useful when this method is called independently without # the previous ones, so we need to load some files if reload_files: logger.info("Starting reloading variables for the step urls_to_images") self.dataset = load_from_disk(self.path_save_dir_dataset) self.dataset_images = load_from_disk(self.path_save_dir_dataset_images) logger.info("Finished reloading variables for the step urls_to_images") else: try: _ = self.dataset _ = self.dataset_images _ = self.map_url_idx except Exception: print("Set `reload_files=True` if you're calling this method alone to define the missing variables") self.dataset = urls_to_images( dataset=self.dataset, dataset_images=self.dataset_images, map_url_idx=self.map_url_idx, num_proc=self.num_proc_urls_to_images, ) def save_dataset(self): logger.info("Starting saving the dataset") self.dataset.save_to_disk(self.path_save_dir_dataset, num_proc=self.num_proc) logger.info("Finished saving the dataset") def save_commit_hash(self): logger.info("Starting writing the commit hash") repo = git.Repo(search_parent_directories=True) sha = repo.head.object.hexsha write_file(os.path.join(self.path_save_dir_dataset, "commit_hash.txt"), sha) logger.info("Finished writing the commit hash") def save_split_sharded_dataset(self, reload_files=False): # Useful when this method is called independently without # the previous ones, so we need to load some files if reload_files: self.dataset = load_from_disk(self.path_save_dir_dataset) else: try: _ = self.dataset except Exception: print("Set `reload_files=True` if you're calling this method alone to define the missing variables") save_split_sharded_dataset( dataset=self.dataset, path_save_dir_sharded_dataset=self.path_save_dir_sharded_dataset, shard_size=self.shard_size, )