chore: import upstream snapshot with attribution

scripts/dataset_processing/get_aishell_data.py

@@ -0,0 +1,177 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# USAGE: python get_aishell_data.py --data_root=<where to put data>
+
+import argparse
+import json
+import logging
+import os
+import subprocess
+import tarfile
+import urllib.request
+
+from tqdm import tqdm
+
+from nemo.utils.tar_utils import safe_extract
+
+parser = argparse.ArgumentParser(description="Aishell Data download")
+parser.add_argument("--data_root", required=True, default=None, type=str)
+args = parser.parse_args()
+
+URL = {"data_aishell": "http://www.openslr.org/resources/33/data_aishell.tgz"}
+
+
+def __retrieve_with_progress(source: str, filename: str):
+    """
+    Downloads source to destination
+    Displays progress bar
+    Args:
+        source: url of resource
+        destination: local filepath
+    Returns:
+    """
+    with open(filename, "wb") as f:
+        response = urllib.request.urlopen(source)
+        total = response.length
+
+        if total is None:
+            f.write(response.content)
+        else:
+            with tqdm(total=total, unit="B", unit_scale=True, unit_divisor=1024) as pbar:
+                for data in response:
+                    f.write(data)
+                    pbar.update(len(data))
+
+
+def __maybe_download_file(destination: str, source: str):
+    """
+    Downloads source to destination if it doesn't exist.
+    If exists, skips download
+    Args:
+        destination: local filepath
+        source: url of resource
+
+    Returns:
+
+    """
+    source = URL[source]
+    if not os.path.exists(destination):
+        logging.info("{0} does not exist. Downloading ...".format(destination))
+        __retrieve_with_progress(source, filename=destination + ".tmp")
+        os.rename(destination + ".tmp", destination)
+        logging.info("Downloaded {0}.".format(destination))
+    else:
+        logging.info("Destination {0} exists. Skipping.".format(destination))
+    return destination
+
+
+def __extract_all_files(filepath: str, data_root: str, data_dir: str):
+    if not os.path.exists(data_dir):
+        extract_file(filepath, data_root)
+        audio_dir = os.path.join(data_dir, "wav")
+        for subfolder, _, filelist in os.walk(audio_dir):
+            for ftar in filelist:
+                extract_file(os.path.join(subfolder, ftar), subfolder)
+    else:
+        logging.info("Skipping extracting. Data already there %s" % data_dir)
+
+
+def extract_file(filepath: str, data_dir: str):
+    try:
+        with tarfile.open(filepath) as tar:
+            safe_extract(tar, data_dir)
+    except Exception:
+        logging.info("Not extracting. Maybe already there?")
+
+
+def __process_data(data_folder: str, dst_folder: str):
+    """
+    To generate manifest
+    Args:
+        data_folder: source with wav files
+        dst_folder: where manifest files will be stored
+    Returns:
+
+    """
+
+    if not os.path.exists(dst_folder):
+        os.makedirs(dst_folder)
+
+    transcript_file = os.path.join(data_folder, "transcript", "aishell_transcript_v0.8.txt")
+    transcript_dict = {}
+    with open(transcript_file, "r", encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            audio_id, text = line.split(" ", 1)
+            # remove white space
+            text = text.replace(" ", "")
+            transcript_dict[audio_id] = text
+
+    data_types = ["train", "dev", "test"]
+    vocab_count = {}
+    for dt in data_types:
+        json_lines = []
+        audio_dir = os.path.join(data_folder, "wav", dt)
+        for sub_folder, _, file_list in os.walk(audio_dir):
+            for fname in file_list:
+                audio_path = os.path.join(sub_folder, fname)
+                audio_id = fname.strip(".wav")
+                if audio_id not in transcript_dict:
+                    continue
+                text = transcript_dict[audio_id]
+                for li in text:
+                    vocab_count[li] = vocab_count.get(li, 0) + 1
+                duration = subprocess.check_output(["soxi", "-D", audio_path])
+                duration = float(duration)
+                json_lines.append(
+                    json.dumps(
+                        {
+                            "audio_filepath": os.path.abspath(audio_path),
+                            "duration": duration,
+                            "text": text,
+                        },
+                        ensure_ascii=False,
+                    )
+                )
+
+        manifest_path = os.path.join(dst_folder, dt + ".json")
+        with open(manifest_path, "w", encoding="utf-8") as fout:
+            for line in json_lines:
+                fout.write(line + "\n")
+
+    vocab = sorted(vocab_count.items(), key=lambda k: k[1], reverse=True)
+    vocab_file = os.path.join(dst_folder, "vocab.txt")
+    with open(vocab_file, "w", encoding="utf-8") as f:
+        for v, c in vocab:
+            f.write(v + "\n")
+
+
+def main():
+    data_root = args.data_root
+    data_set = "data_aishell"
+    logging.info("\n\nWorking on: {0}".format(data_set))
+    file_path = os.path.join(data_root, data_set + ".tgz")
+    logging.info("Getting {0}".format(data_set))
+    __maybe_download_file(file_path, data_set)
+    logging.info("Extracting {0}".format(data_set))
+    data_folder = os.path.join(data_root, data_set)
+    __extract_all_files(file_path, data_root, data_folder)
+    logging.info("Processing {0}".format(data_set))
+    __process_data(data_folder, data_folder)
+    logging.info("Done!")
+
+
+if __name__ == "__main__":
+    main()