Create reorganize_and_archive.py

reorganize_and_archive.py

@@ -0,0 +1,122 @@
+#!/usr/bin/env python3
+import os
+import logging
+import pandas
+from tqdm import tqdm
+import tarfile
+import logging
+import os
+from pathlib import Path
+import csv
+from functools import partial
+from multiprocessing import Pool
+
+
+logging.basicConfig(format='%(asctime)s %(levelname)s: %(message)s', level=logging.INFO)
+
+
+files_per_archive = 40_000
+
+
+def make_archive(archive_index_with_files, output_dir, lang, split):
+    archive_index, files = archive_index_with_files
+    archive_dir = f"{lang}_{split}_{archive_index}"
+    archive_path = os.path.join(output_dir, f"{archive_dir}.tar")
+    with tarfile.open(archive_path, "w") as tar:
+        for file in files:
+            _, filename = os.path.split(file)
+            tar.add(file, arcname=os.path.join(archive_dir, filename))
+
+
+def extract_archive(archive_path, target_dir):
+    with tarfile.open(archive_path, 'r:gz') as f:
+        f.extractall(path=target_dir)
+
+
+def main():
+    # langs = ["ab", "ar", "as", "ast"]
+    # langs = ["az"]
+    # langs = ["ba"]
+    # langs = ["bas", "be", "bg", "bn", "br"]
+    # langs = ["ca", "ckb", "en", "eo", "eu", "fr"]
+    # langs = ["hy-AM"]
+    # langs = ["rw", "th", "yue", "zh-CN"]
+    # langs = ["zh-HK", "zh-TW", "cv"]
+    # langs = ["cs", "da", "dv", "nl", "myv", "et", "fi", "fy-NL", "gl", ]
+    # langs = ["ka", "el", "gn", "cnh", "ha", "mrj", "hi", "ig", "id", "ia", "ga-IE", "it"]
+    # langs = ["ja", "kab", ]
+    # langs = ["kk", "kmr", "ky", "lg", "fa", "es"]
+    # langs = ["cy", "vi", "uz", "ug"]
+    # langs = ["ur", "uk", "tw", "tr", "tok", "ti"] +\
+    #         ["tig", "tt", "ta", "nan-tw", "sv-SE", "sw", "hsb", "sl", "sk", "sr"] + \
+    #         ["sc", "skr", "sat", "sah", "ru", ]
+    # langs = ["rm-sursilv", "ro", "pa-IN", "pt"] + \
+    #         ["pl", "or", "ne-NP", "mn", "mdf", "mhr", "mr", "mt", "ml", "mk", "lt", "nn-NO", "rm-vallader", "lv"]
+    # langs = ['hu', 'vot']
+    for lang in tqdm(langs, desc="languages"):
+        extract = True
+        logging.info(f"Language: {lang.upper()}, files per tar: {files_per_archive}")
+        orig_archive_path = f"/home/polina_huggingface_co/data/cv11/{lang}.tar.gz"
+        if extract:
+            logging.info("Extracting original archive...")
+            extract_archive(orig_archive_path, target_dir="/home/polina_huggingface_co/data/cv11/")
+            logging.info("Extracted.")
+
+        clip_path = f"/home/polina_huggingface_co/data/cv11/cv-corpus-11.0-2022-09-21/{lang}/clips"
+        # data_size = sum(f.stat().st_size for f in Path(clip_path).glob('*.mp3') if f.is_file())
+        # data_size_gb = data_size / 1024 ** 3
+        # num_procs = min(int(data_size_gb) + 2, 28)
+
+        splits = ("test", "dev", "train", "other", "invalidated")
+
+        for split in splits:
+            meta_path = f"/home/polina_huggingface_co/data/cv11/cv-corpus-11.0-2022-09-21/{lang}/{split}.tsv"
+            new_meta_dir = f"repos/common_voice_11_0/transcript/{lang}/"
+            Path(new_meta_dir).mkdir(parents=True, exist_ok=True)
+
+            data = pandas.read_csv(meta_path, sep='\t', quoting=csv.QUOTE_NONE, low_memory=False)
+            copy_command = f"cp {meta_path} {new_meta_dir}"
+            os.system(copy_command)
+
+            all_files = [os.path.join(clip_path, filename) for filename in list(data["path"])]
+
+            num_files = len(all_files)
+            if num_files == 0:
+                continue
+
+            logging.info(f"split: {split.upper()}, num_files: {num_files}")
+
+            new_clip_path = f"repos/common_voice_11_0/audio/{lang}/{split}"
+            Path(new_clip_path).mkdir(parents=True, exist_ok=True)
+
+            file_groups = [
+                (arch_index_in_dir, all_files[start_index:start_index + files_per_archive])
+                for arch_index_in_dir, start_index in enumerate(range(0, num_files, files_per_archive))
+            ]
+
+            n_file_groups = len(file_groups)
+            num_procs = max(1, min(n_file_groups, 26))
+            logging.info(f"N groups: {n_file_groups}, num procs: {num_procs}")
+
+            if n_file_groups > 1:
+                pool = Pool(num_procs)
+                pool.map(
+                    partial(
+                        make_archive,
+                        output_dir=new_clip_path,
+                        lang=lang,
+                        split=split,
+                    ),
+                    tqdm(file_groups, desc=f"Taring {split} subset...", position=0),
+                )
+            else:
+                make_archive(
+                    file_groups[0],
+                    output_dir=new_clip_path,
+                    lang=lang,
+                    split=split,
+                )
+
+
+if __name__ == "__main__":
+    main()