| 1 | # generate audio text map for WenetSpeech4TTS |
| 2 | # evaluate for vocab size |
| 3 | |
| 4 | import os |
| 5 | import sys |
| 6 | |
| 7 | |
| 8 | sys.path.append(os.getcwd()) |
| 9 | |
| 10 | import json |
| 11 | from concurrent.futures import ProcessPoolExecutor |
| 12 | from importlib.resources import files |
| 13 | |
| 14 | import torchaudio |
| 15 | from datasets import Dataset |
| 16 | from tqdm import tqdm |
| 17 | |
| 18 | from f5_tts.model.utils import convert_char_to_pinyin |
| 19 | |
| 20 | |
| 21 | def deal_with_sub_path_files(dataset_path, sub_path): |
| 22 | print(f"Dealing with: {sub_path}") |
| 23 | |
| 24 | text_dir = os.path.join(dataset_path, sub_path, "txts") |
| 25 | audio_dir = os.path.join(dataset_path, sub_path, "wavs") |
| 26 | text_files = os.listdir(text_dir) |
| 27 | |
| 28 | audio_paths, texts, durations = [], [], [] |
| 29 | for text_file in tqdm(text_files): |
| 30 | with open(os.path.join(text_dir, text_file), "r", encoding="utf-8") as file: |
| 31 | first_line = file.readline().split("\t") |
| 32 | audio_nm = first_line[0] |
| 33 | audio_path = os.path.join(audio_dir, audio_nm + ".wav") |
| 34 | text = first_line[1].strip() |
| 35 | |
| 36 | audio_paths.append(audio_path) |
| 37 | |
| 38 | if tokenizer == "pinyin": |
| 39 | texts.extend(convert_char_to_pinyin([text], polyphone=polyphone)) |
| 40 | elif tokenizer == "char": |
| 41 | texts.append(text) |
| 42 | |
| 43 | audio, sample_rate = torchaudio.load(audio_path) |
| 44 | durations.append(audio.shape[-1] / sample_rate) |
| 45 | |
| 46 | return audio_paths, texts, durations |
| 47 | |
| 48 | |
| 49 | def main(): |
| 50 | assert tokenizer in ["pinyin", "char"] |
| 51 | |
| 52 | audio_path_list, text_list, duration_list = [], [], [] |
| 53 | |
| 54 | executor = ProcessPoolExecutor(max_workers=max_workers) |
| 55 | futures = [] |
| 56 | for dataset_path in dataset_paths: |
| 57 | sub_items = os.listdir(dataset_path) |
| 58 | sub_paths = [item for item in sub_items if os.path.isdir(os.path.join(dataset_path, item))] |
| 59 | for sub_path in sub_paths: |
| 60 | futures.append(executor.submit(deal_with_sub_path_files, dataset_path, sub_path)) |
| 61 | for future in tqdm(futures, total=len(futures)): |
| 62 | audio_paths, texts, durations = future.result() |
| 63 | audio_path_list.extend(audio_paths) |
| 64 | text_list.extend(texts) |
| 65 | duration_list.extend(durations) |
| 66 | executor.shutdown() |
| 67 | |
| 68 | if not os.path.exists("data"): |
| 69 | os.makedirs("data") |
| 70 | |
| 71 | print(f"\nSaving to {save_dir} ...") |
| 72 | dataset = Dataset.from_dict({"audio_path": audio_path_list, "text": text_list, "duration": duration_list}) |
| 73 | dataset.save_to_disk(f"{save_dir}/raw", max_shard_size="2GB") # arrow format |
| 74 | |
| 75 | with open(f"{save_dir}/duration.json", "w", encoding="utf-8") as f: |
| 76 | json.dump( |
| 77 | {"duration": duration_list}, f, ensure_ascii=False |
| 78 | ) # dup a json separately saving duration in case for DynamicBatchSampler ease |
| 79 | |
| 80 | print("\nEvaluating vocab size (all characters and symbols / all phonemes) ...") |
| 81 | text_vocab_set = set() |
| 82 | for text in tqdm(text_list): |
| 83 | text_vocab_set.update(list(text)) |
| 84 | |
| 85 | # add alphabets and symbols (optional, if plan to ft on de/fr etc.) |
| 86 | if tokenizer == "pinyin": |
| 87 | text_vocab_set.update([chr(i) for i in range(32, 127)] + [chr(i) for i in range(192, 256)]) |
| 88 | |
| 89 | with open(f"{save_dir}/vocab.txt", "w") as f: |
| 90 | for vocab in sorted(text_vocab_set): |
| 91 | f.write(vocab + "\n") |
| 92 | print(f"\nFor {dataset_name}, sample count: {len(text_list)}") |
| 93 | print(f"For {dataset_name}, vocab size is: {len(text_vocab_set)}\n") |
| 94 | |
| 95 | |
| 96 | if __name__ == "__main__": |
| 97 | max_workers = 32 |
| 98 | |
| 99 | tokenizer = "pinyin" # "pinyin" | "char" |
| 100 | polyphone = True |
| 101 | dataset_choice = 1 # 1: Premium, 2: Standard, 3: Basic |
| 102 | |
| 103 | dataset_name = ( |
| 104 | ["WenetSpeech4TTS_Premium", "WenetSpeech4TTS_Standard", "WenetSpeech4TTS_Basic"][dataset_choice - 1] |
| 105 | + "_" |
| 106 | + tokenizer |
| 107 | ) |
| 108 | dataset_paths = [ |
| 109 | "<SOME_PATH>/WenetSpeech4TTS/Basic", |
| 110 | "<SOME_PATH>/WenetSpeech4TTS/Standard", |
| 111 | "<SOME_PATH>/WenetSpeech4TTS/Premium", |
| 112 | ][-dataset_choice:] |
| 113 | save_dir = str(files("f5_tts").joinpath("../../")) + f"/data/{dataset_name}" |
| 114 | print(f"\nChoose Dataset: {dataset_name}, will save to {save_dir}\n") |
| 115 | |
| 116 | main() |
| 117 | |
| 118 | # Results (if adding alphabets with accents and symbols): |
| 119 | # WenetSpeech4TTS Basic Standard Premium |
| 120 | # samples count 3932473 1941220 407494 |
| 121 | # pinyin vocab size 1349 1348 1344 (no polyphone) |
| 122 | # - - 1459 (polyphone) |
| 123 | # char vocab size 5264 5219 5042 |
| 124 | |
| 125 | # vocab size may be slightly different due to rjieba tokenizer and pypinyin (e.g. way of polyphoneme) |
| 126 | # please be careful if using pretrained model, make sure the vocab.txt is same |
| 127 |