AR/exps/get_txt_librilight.py

import argparse
import os
import time
import traceback
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import librosa
import numpy as np
import tqdm
import whisper
from soundstorm.s2.exps.hubert.feature_utils import get_shard_range
from soundstorm.utils import check_txt_file


def process_sentence(args,
                     fp: Path,
                     train_dump_dir: Path,
                     dev_dump_dir: Path,
                     test_dump_dir: Path,
                     VAD_dict):
    asr_model = whisper.load_model("tiny.en")
    utt_id = fp.stem
    sr = args.sr
    record = []
    train_txt_dir = train_dump_dir / "txt"
    train_txt_dir.mkdir(parents=True, exist_ok=True)

    dev_txt_dir = dev_dump_dir / "txt"
    dev_txt_dir.mkdir(parents=True, exist_ok=True)

    test_txt_dir = test_dump_dir / "txt"
    test_txt_dir.mkdir(parents=True, exist_ok=True)

    try:
        # get info for path
        wav_path_list = str(fp).strip().split('/')
        sub_dataset, spk_id, book_name = wav_path_list[-4], wav_path_list[
            -3], wav_path_list[-2]
        wav_name = wav_path_list[-1][:-5]
        assert wav_name == utt_id
        # key_name for big wav
        key_name = f'{wav_name}#{sub_dataset}#{spk_id}#{book_name}'
        # 判断 VAD 字典中不存在该条音频信息的情况
        if key_name not in VAD_dict.keys():
            print(key_name, 'not in VAD_dict !')
            return record
        wav = None
        sorted_split_VAD_dict = sorted(VAD_dict[key_name].items())
        len_dict = len(sorted_split_VAD_dict)
        for index, item in enumerate(sorted_split_VAD_dict):
            split_name, value = item
            start, end = value
            # train | dev | test
            if index == len_dict - 1:
                subset = 'test'
                txt_path = test_txt_dir / (split_name + ".txt")
            elif index == len_dict - 2:
                subset = 'dev'
                txt_path = dev_txt_dir / (split_name + ".txt")
            else:
                subset = 'train'
                txt_path = train_txt_dir / (split_name + ".txt")

            if os.path.exists(txt_path) and check_txt_file(txt_path):
                # print(txt_path, 'exits!')
                pass
            else:
                # 这里加判断保证在 sub wav 的循环中只 load 一次
                if wav is None:
                    # load big wav
                    # 在最外层 load 如果 sub wav 的特征都存在了就会白白消耗 load 的时间
                    wav, _ = librosa.load(str(fp), sr=sr)
                sub_wav = wav[int(start * sr):int(end * sr)]
                asr_result = asr_model.transcribe(sub_wav)["text"]
                with open(txt_path, 'w') as f:
                    f.write(asr_result)

            sub_record = {
                "utt_id": split_name,
                "txt_path": txt_path,
                "subset": subset
            }
            # recodrd 变成 List of Dict
            record.append(sub_record)
    except Exception:
        print("occur Exception")
        traceback.print_exc()
        # record 有可能是一个不完整的 List
        return record
    return record


def process_sentences(args,
                      fps: Path,
                      train_dump_dir: Path,
                      dev_dump_dir: Path,
                      test_dump_dir: Path,
                      VAD_dict,
                      nprocs: int=1):
    print("nprocs:", nprocs)
    if nprocs == 1:
        results = []
        for fp in tqdm.tqdm(fps, total=len(fps)):
            record = process_sentence(
                args=args,
                fp=fp,
                train_dump_dir=train_dump_dir,
                dev_dump_dir=dev_dump_dir,
                test_dump_dir=test_dump_dir,
                VAD_dict=VAD_dict)
            if record:
                results.append(record)
    else:
        with ThreadPoolExecutor(nprocs) as pool:
            futures = []
            with tqdm.tqdm(total=len(fps)) as progress:
                for fp in fps:
                    future = pool.submit(process_sentence, args, fp,
                                         train_dump_dir, dev_dump_dir,
                                         test_dump_dir, VAD_dict)
                    future.add_done_callback(lambda p: progress.update())
                    futures.append(future)

                results = []
                for ft in futures:
                    record = ft.result()
                    if record:
                        results.append(record)

    # torch.save() to a large `.pth` file
    txt_dict = dict()
    txt_dict['train'] = {}
    txt_dict['dev'] = {}
    txt_dict['test'] = {}
    # record 是 List of Dict, 一条大 wav 一个 record，一条小 wav 一个 sub_recored
    print(f"start to save {args.rank}_{args.nshard}.npy ...")
    save_start_time = time.time()
    for record in tqdm.tqdm(results, total=len(results), colour='green'):
        for sub_record in record:
            # 这里加 try, 因为 txt 文件可能损坏
            try:
                utt_id = sub_record["utt_id"]
                subset = sub_record["subset"]
                asr_result = check_txt_file(sub_record["txt_path"])
                if asr_result is not False:
                    txt_dict[subset][utt_id] = asr_result
                else:
                    print(f'asr result of {utt_id} is False')
            except Exception:
                print(f"{utt_id} occur Exception")
                traceback.print_exc()
                continue

    train_filename = train_dump_dir / f'txt_{args.rank}_{args.nshard}.npy'
    dev_filename = dev_dump_dir / f'txt_{args.rank}_{args.nshard}.npy'
    test_filename = test_dump_dir / f'txt_{args.rank}_{args.nshard}.npy'
    np.save(train_filename, txt_dict['train'])
    print(f"npy file '{train_filename}' write down")

    np.save(dev_filename, txt_dict['dev'])
    print(f"npy file '{dev_filename}' write down")

    np.save(test_filename, txt_dict['test'])
    print(f"npy file '{test_filename}' write down")
    print('time of save stage:', time.time() - save_start_time)


def main():
    # parse config and args
    parser = argparse.ArgumentParser(
        description="Preprocess audio and then extract features for LibriLight.")

    parser.add_argument(
        "--data_dir", default=None, type=str, help="directory to dataset.")

    parser.add_argument(
        "--dump_dir",
        type=str,
        required=True,
        help="directory to dump feature files.")

    parser.add_argument(
        "--num-cpu", type=int, default=1, help="number of process.")

    parser.add_argument(
        '--sr', type=int, default=16000, help='sample rate of model')

    # For LibriLight dataset
    parser.add_argument(
        "--sub_dataset",
        default="small",
        type=str,
        help="name of sub dataset of LibriLight",
        choices=['small', 'medium', 'large', 'duplicate'], )
    parser.add_argument(
        "--VAD_path", type=str, default='./VAD/librilight_segment_dict.npy')
    parser.add_argument("--nshard", type=int, default=3)
    parser.add_argument("--rank", type=int, default=0)

    args = parser.parse_args()

    data_dir = Path(args.data_dir).expanduser()
    dump_dir = Path(args.dump_dir).expanduser()
    # use absolute path
    dump_dir = dump_dir.resolve()
    dump_dir.mkdir(parents=True, exist_ok=True)

    assert data_dir.is_dir()

    # sub_dataset here
    sub_dataset_dir = data_dir / args.sub_dataset
    # olny spk_id in list, sort by lexicographical order 
    speaker_list = sorted(os.listdir(sub_dataset_dir))
    start, end = get_shard_range(len(speaker_list), args.nshard, args.rank)
    # speaker_list for this rank
    speaker_list = speaker_list[start:end]

    all_wav_files = []

    for speaker in speaker_list:
        wav_files = sorted(list((sub_dataset_dir / speaker).rglob("*/*.flac")))
        # filter out ._*.flac
        wav_files = [
            file for file in wav_files if not file.name.startswith('._')
        ]
        all_wav_files += wav_files

    print(f"num of wav files in rank {args.rank}:", len(all_wav_files))
    # get VAD info
    VAD_dict = np.load(args.VAD_path, allow_pickle=True).item()

    sub_dataset_dump_dir = dump_dir / args.sub_dataset
    sub_dataset_dump_dir.mkdir(parents=True, exist_ok=True)
    train_dump_dir = sub_dataset_dump_dir / "train"
    train_dump_dir.mkdir(parents=True, exist_ok=True)
    dev_dump_dir = sub_dataset_dump_dir / "dev"
    dev_dump_dir.mkdir(parents=True, exist_ok=True)
    test_dump_dir = sub_dataset_dump_dir / "test"
    test_dump_dir.mkdir(parents=True, exist_ok=True)

    # 每条大 wav 分出一个 dev 一个 test，比例大概是 96:2:2
    if all_wav_files:
        process_sentences(
            args=args,
            fps=all_wav_files,
            train_dump_dir=train_dump_dir,
            dev_dump_dir=dev_dump_dir,
            test_dump_dir=test_dump_dir,
            VAD_dict=VAD_dict,
            nprocs=args.num_cpu)


if __name__ == "__main__":
    main()