| 1 | # Emilia Dataset: https://huggingface.co/datasets/amphion/Emilia-Dataset/tree/fc71e07 |
| 2 | # if use updated new version, i.e. WebDataset, feel free to modify / draft your own script |
| 3 | |
| 4 | # generate audio text map for Emilia ZH & EN |
| 5 | # evaluate for vocab size |
| 6 | |
| 7 | import os |
| 8 | import sys |
| 9 | |
| 10 | |
| 11 | sys.path.append(os.getcwd()) |
| 12 | |
| 13 | import json |
| 14 | from concurrent.futures import ProcessPoolExecutor |
| 15 | from importlib.resources import files |
| 16 | from pathlib import Path |
| 17 | |
| 18 | from datasets.arrow_writer import ArrowWriter |
| 19 | from tqdm import tqdm |
| 20 | |
| 21 | from f5_tts.model.utils import convert_char_to_pinyin, repetition_found |
| 22 | |
| 23 | |
| 24 | out_zh = { |
| 25 | "ZH_B00041_S06226", |
| 26 | "ZH_B00042_S09204", |
| 27 | "ZH_B00065_S09430", |
| 28 | "ZH_B00065_S09431", |
| 29 | "ZH_B00066_S09327", |
| 30 | "ZH_B00066_S09328", |
| 31 | } |
| 32 | zh_filters = ["い", "て"] |
| 33 | # seems synthesized audios, or heavily code-switched |
| 34 | out_en = { |
| 35 | "EN_B00013_S00913", |
| 36 | "EN_B00042_S00120", |
| 37 | "EN_B00055_S04111", |
| 38 | "EN_B00061_S00693", |
| 39 | "EN_B00061_S01494", |
| 40 | "EN_B00061_S03375", |
| 41 | "EN_B00059_S00092", |
| 42 | "EN_B00111_S04300", |
| 43 | "EN_B00100_S03759", |
| 44 | "EN_B00087_S03811", |
| 45 | "EN_B00059_S00950", |
| 46 | "EN_B00089_S00946", |
| 47 | "EN_B00078_S05127", |
| 48 | "EN_B00070_S04089", |
| 49 | "EN_B00074_S09659", |
| 50 | "EN_B00061_S06983", |
| 51 | "EN_B00061_S07060", |
| 52 | "EN_B00059_S08397", |
| 53 | "EN_B00082_S06192", |
| 54 | "EN_B00091_S01238", |
| 55 | "EN_B00089_S07349", |
| 56 | "EN_B00070_S04343", |
| 57 | "EN_B00061_S02400", |
| 58 | "EN_B00076_S01262", |
| 59 | "EN_B00068_S06467", |
| 60 | "EN_B00076_S02943", |
| 61 | "EN_B00064_S05954", |
| 62 | "EN_B00061_S05386", |
| 63 | "EN_B00066_S06544", |
| 64 | "EN_B00076_S06944", |
| 65 | "EN_B00072_S08620", |
| 66 | "EN_B00076_S07135", |
| 67 | "EN_B00076_S09127", |
| 68 | "EN_B00065_S00497", |
| 69 | "EN_B00059_S06227", |
| 70 | "EN_B00063_S02859", |
| 71 | "EN_B00075_S01547", |
| 72 | "EN_B00061_S08286", |
| 73 | "EN_B00079_S02901", |
| 74 | "EN_B00092_S03643", |
| 75 | "EN_B00096_S08653", |
| 76 | "EN_B00063_S04297", |
| 77 | "EN_B00063_S04614", |
| 78 | "EN_B00079_S04698", |
| 79 | "EN_B00104_S01666", |
| 80 | "EN_B00061_S09504", |
| 81 | "EN_B00061_S09694", |
| 82 | "EN_B00065_S05444", |
| 83 | "EN_B00063_S06860", |
| 84 | "EN_B00065_S05725", |
| 85 | "EN_B00069_S07628", |
| 86 | "EN_B00083_S03875", |
| 87 | "EN_B00071_S07665", |
| 88 | "EN_B00071_S07665", |
| 89 | "EN_B00062_S04187", |
| 90 | "EN_B00065_S09873", |
| 91 | "EN_B00065_S09922", |
| 92 | "EN_B00084_S02463", |
| 93 | "EN_B00067_S05066", |
| 94 | "EN_B00106_S08060", |
| 95 | "EN_B00073_S06399", |
| 96 | "EN_B00073_S09236", |
| 97 | "EN_B00087_S00432", |
| 98 | "EN_B00085_S05618", |
| 99 | "EN_B00064_S01262", |
| 100 | "EN_B00072_S01739", |
| 101 | "EN_B00059_S03913", |
| 102 | "EN_B00069_S04036", |
| 103 | "EN_B00067_S05623", |
| 104 | "EN_B00060_S05389", |
| 105 | "EN_B00060_S07290", |
| 106 | "EN_B00062_S08995", |
| 107 | } |
| 108 | en_filters = ["ا", "い", "て"] |
| 109 | |
| 110 | |
| 111 | def deal_with_audio_dir(audio_dir): |
| 112 | audio_jsonl = audio_dir.with_suffix(".jsonl") |
| 113 | sub_result, durations = [], [] |
| 114 | vocab_set = set() |
| 115 | bad_case_zh = 0 |
| 116 | bad_case_en = 0 |
| 117 | with open(audio_jsonl, "r") as f: |
| 118 | lines = f.readlines() |
| 119 | for line in tqdm(lines, desc=f"{audio_jsonl.stem}"): |
| 120 | obj = json.loads(line) |
| 121 | text = obj["text"] |
| 122 | if obj["language"] == "zh": |
| 123 | if obj["wav"].split("/")[1] in out_zh or any(f in text for f in zh_filters) or repetition_found(text): |
| 124 | bad_case_zh += 1 |
| 125 | continue |
| 126 | else: |
| 127 | text = text.translate( |
| 128 | str.maketrans({",": "，", "!": "！", "?": "？"}) |
| 129 | ) # not "。" cuz much code-switched |
| 130 | if obj["language"] == "en": |
| 131 | if ( |
| 132 | obj["wav"].split("/")[1] in out_en |
| 133 | or any(f in text for f in en_filters) |
| 134 | or repetition_found(text, length=4) |
| 135 | ): |
| 136 | bad_case_en += 1 |
| 137 | continue |
| 138 | if tokenizer == "pinyin": |
| 139 | text = convert_char_to_pinyin([text], polyphone=polyphone)[0] |
| 140 | duration = obj["duration"] |
| 141 | sub_result.append({"audio_path": str(audio_dir.parent / obj["wav"]), "text": text, "duration": duration}) |
| 142 | durations.append(duration) |
| 143 | vocab_set.update(list(text)) |
| 144 | return sub_result, durations, vocab_set, bad_case_zh, bad_case_en |
| 145 | |
| 146 | |
| 147 | def main(): |
| 148 | assert tokenizer in ["pinyin", "char"] |
| 149 | result = [] |
| 150 | duration_list = [] |
| 151 | text_vocab_set = set() |
| 152 | total_bad_case_zh = 0 |
| 153 | total_bad_case_en = 0 |
| 154 | |
| 155 | # process raw data |
| 156 | executor = ProcessPoolExecutor(max_workers=max_workers) |
| 157 | futures = [] |
| 158 | for lang in langs: |
| 159 | dataset_path = Path(os.path.join(dataset_dir, lang)) |
| 160 | [ |
| 161 | futures.append(executor.submit(deal_with_audio_dir, audio_dir)) |
| 162 | for audio_dir in dataset_path.iterdir() |
| 163 | if audio_dir.is_dir() |
| 164 | ] |
| 165 | for futures in tqdm(futures, total=len(futures)): |
| 166 | sub_result, durations, vocab_set, bad_case_zh, bad_case_en = futures.result() |
| 167 | result.extend(sub_result) |
| 168 | duration_list.extend(durations) |
| 169 | text_vocab_set.update(vocab_set) |
| 170 | total_bad_case_zh += bad_case_zh |
| 171 | total_bad_case_en += bad_case_en |
| 172 | executor.shutdown() |
| 173 | |
| 174 | # save preprocessed dataset to disk |
| 175 | if not os.path.exists(f"{save_dir}"): |
| 176 | os.makedirs(f"{save_dir}") |
| 177 | print(f"\nSaving to {save_dir} ...") |
| 178 | |
| 179 | # dataset = Dataset.from_dict({"audio_path": audio_path_list, "text": text_list, "duration": duration_list}) # oom |
| 180 | # dataset.save_to_disk(f"{save_dir}/raw", max_shard_size="2GB") |
| 181 | with ArrowWriter(path=f"{save_dir}/raw.arrow") as writer: |
| 182 | for line in tqdm(result, desc="Writing to raw.arrow ..."): |
| 183 | writer.write(line) |
| 184 | writer.finalize() |
| 185 | |
| 186 | # dup a json separately saving duration in case for DynamicBatchSampler ease |
| 187 | with open(f"{save_dir}/duration.json", "w", encoding="utf-8") as f: |
| 188 | json.dump({"duration": duration_list}, f, ensure_ascii=False) |
| 189 | |
| 190 | # vocab map, i.e. tokenizer |
| 191 | # add alphabets and symbols (optional, if plan to ft on de/fr etc.) |
| 192 | # if tokenizer == "pinyin": |
| 193 | # text_vocab_set.update([chr(i) for i in range(32, 127)] + [chr(i) for i in range(192, 256)]) |
| 194 | with open(f"{save_dir}/vocab.txt", "w") as f: |
| 195 | for vocab in sorted(text_vocab_set): |
| 196 | f.write(vocab + "\n") |
| 197 | |
| 198 | print(f"\nFor {dataset_name}, sample count: {len(result)}") |
| 199 | print(f"For {dataset_name}, vocab size is: {len(text_vocab_set)}") |
| 200 | print(f"For {dataset_name}, total {sum(duration_list) / 3600:.2f} hours") |
| 201 | if "ZH" in langs: |
| 202 | print(f"Bad zh transcription case: {total_bad_case_zh}") |
| 203 | if "EN" in langs: |
| 204 | print(f"Bad en transcription case: {total_bad_case_en}\n") |
| 205 | |
| 206 | |
| 207 | if __name__ == "__main__": |
| 208 | max_workers = 32 |
| 209 | |
| 210 | tokenizer = "pinyin" # "pinyin" | "char" |
| 211 | polyphone = True |
| 212 | |
| 213 | langs = ["ZH", "EN"] |
| 214 | dataset_dir = "<SOME_PATH>/Emilia_Dataset/raw" |
| 215 | dataset_name = f"Emilia_{'_'.join(langs)}_{tokenizer}" |
| 216 | save_dir = str(files("f5_tts").joinpath("../../")) + f"/data/{dataset_name}" |
| 217 | print(f"\nPrepare for {dataset_name}, will save to {save_dir}\n") |
| 218 | |
| 219 | main() |
| 220 | |
| 221 | # Emilia ZH & EN |
| 222 | # samples count 37837916 (after removal) |
| 223 | # pinyin vocab size 2543 (polyphone) |
| 224 | # total duration 95281.87 (hours) |
| 225 | # bad zh asr cnt 230435 (samples) |
| 226 | # bad eh asr cnt 37217 (samples) |
| 227 | |
| 228 | # vocab size may be slightly different due to rjieba tokenizer and pypinyin (e.g. way of polyphoneme) |
| 229 | # please be careful if using pretrained model, make sure the vocab.txt is same |
| 230 |