| 1 | # put in src/f5_tts/train/datasets/prepare_emilia_v2.py |
| 2 | # prepares Emilia dataset with the new format w/ Emilia-YODAS |
| 3 | |
| 4 | import json |
| 5 | import os |
| 6 | from concurrent.futures import ProcessPoolExecutor |
| 7 | from importlib.resources import files |
| 8 | from pathlib import Path |
| 9 | |
| 10 | from datasets.arrow_writer import ArrowWriter |
| 11 | from tqdm import tqdm |
| 12 | |
| 13 | from f5_tts.model.utils import repetition_found |
| 14 | |
| 15 | |
| 16 | # Define filters for exclusion |
| 17 | out_en = set() |
| 18 | en_filters = ["ا", "い", "て"] |
| 19 | |
| 20 | |
| 21 | def process_audio_directory(audio_dir): |
| 22 | sub_result, durations, vocab_set = [], [], set() |
| 23 | bad_case_en = 0 |
| 24 | |
| 25 | for file in audio_dir.iterdir(): |
| 26 | if file.suffix == ".json": |
| 27 | with open(file, "r") as f: |
| 28 | obj = json.load(f) |
| 29 | text = obj["text"] |
| 30 | if any(f in text for f in en_filters) or repetition_found(text, length=4): |
| 31 | bad_case_en += 1 |
| 32 | continue |
| 33 | |
| 34 | duration = obj["duration"] |
| 35 | audio_file = file.with_suffix(".mp3") |
| 36 | if audio_file.exists(): |
| 37 | sub_result.append({"audio_path": str(audio_file), "text": text, "duration": duration}) |
| 38 | durations.append(duration) |
| 39 | vocab_set.update(list(text)) |
| 40 | |
| 41 | return sub_result, durations, vocab_set, bad_case_en |
| 42 | |
| 43 | |
| 44 | def main(): |
| 45 | assert tokenizer in ["pinyin", "char"] |
| 46 | result, duration_list, text_vocab_set = [], [], set() |
| 47 | total_bad_case_en = 0 |
| 48 | |
| 49 | executor = ProcessPoolExecutor(max_workers=max_workers) |
| 50 | futures = [] |
| 51 | dataset_path = Path(dataset_dir) |
| 52 | for sub_dir in dataset_path.iterdir(): |
| 53 | if sub_dir.is_dir(): |
| 54 | futures.append(executor.submit(process_audio_directory, sub_dir)) |
| 55 | |
| 56 | for future in tqdm(futures, total=len(futures)): |
| 57 | sub_result, durations, vocab_set, bad_case_en = future.result() |
| 58 | result.extend(sub_result) |
| 59 | duration_list.extend(durations) |
| 60 | text_vocab_set.update(vocab_set) |
| 61 | total_bad_case_en += bad_case_en |
| 62 | |
| 63 | executor.shutdown() |
| 64 | |
| 65 | if not os.path.exists(f"{save_dir}"): |
| 66 | os.makedirs(f"{save_dir}") |
| 67 | |
| 68 | with ArrowWriter(path=f"{save_dir}/raw.arrow") as writer: |
| 69 | for line in tqdm(result, desc="Writing to raw.arrow ..."): |
| 70 | writer.write(line) |
| 71 | writer.finalize() |
| 72 | |
| 73 | with open(f"{save_dir}/duration.json", "w", encoding="utf-8") as f: |
| 74 | json.dump({"duration": duration_list}, f, ensure_ascii=False) |
| 75 | |
| 76 | with open(f"{save_dir}/vocab.txt", "w") as f: |
| 77 | for vocab in sorted(text_vocab_set): |
| 78 | f.write(vocab + "\n") |
| 79 | |
| 80 | print(f"For {dataset_name}, sample count: {len(result)}") |
| 81 | print(f"For {dataset_name}, vocab size is: {len(text_vocab_set)}") |
| 82 | print(f"For {dataset_name}, total {sum(duration_list) / 3600:.2f} hours") |
| 83 | print(f"Bad en transcription case: {total_bad_case_en}\n") |
| 84 | |
| 85 | |
| 86 | if __name__ == "__main__": |
| 87 | max_workers = 32 |
| 88 | tokenizer = "char" |
| 89 | dataset_dir = "/home/ubuntu/emilia-dataset/Emilia-YODAS/EN" |
| 90 | dataset_name = f"Emilia_EN_{tokenizer}" |
| 91 | # save_dir = os.path.expanduser(f"~/F5-TTS/data/{dataset_name}") |
| 92 | save_dir = str(files("f5_tts").joinpath("../../")) + f"/data/{dataset_name}" |
| 93 | |
| 94 | print(f"Prepare for {dataset_name}, will save to {save_dir}\n") |
| 95 | main() |
| 96 |