Spaces:
Running
on
Zero
Running
on
Zero
File size: 8,284 Bytes
82bc972 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 |
from importlib.resources import path
import pathlib
import soundfile as sf
import numpy as np
import json
import multiprocessing
import argparse
import tqdm
import gzip
import time
import os
from tokenizer import TextTokenizer, tokenize_text
import glob
import sys
import os, random, numpy as np, socket
import json
import tqdm
def write_jsonl(data, fn):
with open(fn, "w") as file:
for entry in data:
file.write(json.dumps(entry, ensure_ascii=False) + "\n")
def read_jsonl(file_path):
cur_data = []
with open(file_path, 'r', encoding='utf-8-sig') as file:
for line in file:
cur_data.append(json.loads(line.strip()))
return cur_data
def save_audio(seq, fn):
output = seq
os.makedirs(os.path.dirname(fn), exist_ok=True)
sf.write(fn, output, samplerate=16000)
def save_text(text, fn):
os.makedirs(os.path.dirname(fn), exist_ok=True)
with open(fn, "w") as wwf:
wwf.writelines(text)
def phonemize_and_save(text, fn):
phn = tokenize_text(text_tokenizer, text)
os.makedirs(os.path.dirname(fn), exist_ok=True)
with open(fn, "w") as f:
f.write(' '.join(phn))
return set(phn)
def cut_sequence(task):
in_audio_fn, output_dir, metadata = task
if not os.path.isfile(in_audio_fn):
# print("missing: ", in_audio_fn)
return None
data, samplerate = sf.read(in_audio_fn)
assert len(data.shape) == 1
assert samplerate == 16000
all_phns = set()
for item in metadata:
out_fn = item['file_id']
out_audio_fn = os.path.join(output_dir, "audio", out_fn)
out_text_fn = os.path.join(output_dir, "audio", out_fn.replace(".flac", ".txt"))
out_phn_fn = os.path.join(output_dir, "phoneme", out_fn.replace(".flac", ".txt"))
save_audio(data[int(item['vad'][0]*samplerate):int(item['vad'][1]*samplerate)], out_audio_fn)
save_text(item['text'], out_text_fn)
phns = phonemize_and_save(item['text'], out_phn_fn)
all_phns.update(phns)
return all_phns
from collections import defaultdict
# Function to create a defaultdict recursively
def nested_defaultdict(levels, inner_type):
if levels <= 1:
return defaultdict(inner_type)
return defaultdict(lambda: nested_defaultdict(levels-1, inner_type))
def open_mani(fn):
print("load segmentation and transcription metadata...")
stime = time.time()
data = []
with gzip.open(fn, 'rt', encoding='utf-8') as f:
for line in f:
data.append(json.loads(line))
print(f"loading done, took {time.time() - stime:.4f} seconds")
return data
def cut(split,
audio_dir,
mani_dir,
output_dir,
n_process=32,
percent=0.5):
split2manifest = {
"train": [
"libriheavy_long_cuts_small.jsonl",
"libriheavy_long_cuts_medium.jsonl",
"libriheavy_long_cuts_large.jsonl",
"libriheavy_cuts_small.jsonl",
"libriheavy_cuts_medium.jsonl",
"libriheavy_cuts_large.jsonl",
],
"valid": [
"libriheavy_cuts_dev.jsonl",
"libriheavy_long_cuts_dev.jsonl"
],
"test": [
"libriheavy_cuts_test_clean.jsonl",
"libriheavy_cuts_test_other.jsonl",
"libriheavy_long_cuts_test_clean.jsonl",
"libriheavy_long_cuts_test_other.jsonl"
]
}
print("organize data by recording_id (i.e. the original big .flac file name)...")
stime = time.time()
organized_data = nested_defaultdict(4, list)
manifest_fn = os.path.join(output_dir, "manifest_mimi", split+".txt")
os.makedirs(os.path.join(output_dir, "manifest_mimi"), exist_ok=True)
with open(manifest_fn, "w") as wf:
for mani_fn in split2manifest[split]:
# data = open_mani(os.path.join(mani_dir, mani_fn))
data = read_jsonl(os.path.join(mani_dir, mani_fn))
for item in data:
file_id = item['supervisions'][0]['id'] + '.flac'
recording_id = item['recording']['id'] + '.flac'
sizeSplit, spk, book, flac = recording_id.split("/") # e.g. 'medium/100/emerald_city_librivox_64kb_mp3/emeraldcity_01_baum_64kb'
if os.path.isfile(os.path.join(audio_dir, recording_id)):
vad = (item['start'], item['start']+item['duration'])
text = item['supervisions'][0]['custom']['texts'][0]
file_id = file_id.replace(".flac", "") + f"_{vad[0]:.2f}_{vad[1]:.2f}.flac"
organized_data[sizeSplit][spk][book][recording_id].append({"file_id": file_id, "vad":vad, "text": text})
wf.writelines(f"{file_id}\t{item['duration']}\n")
# #### take only a subet of tasks
tasks = [(os.path.join(audio_dir, recording_id), output_dir, organized_data[sizeSplit][spk][book][recording_id], spk) for sizeSplit in organized_data for spk in organized_data[sizeSplit] for book in organized_data[sizeSplit][spk] for recording_id in organized_data[sizeSplit][spk][book]]
ntasks = len(tasks)
spk2tasks = defaultdict(list)
for task in tasks:
spk2tasks[task[3]].append(task)
# randomly shuffle each task list for each speaker
for spk in spk2tasks:
random.shuffle(spk2tasks[spk])
# take only 20% of the tasks, uniformly sampled from each speaker
# randomly pick a speaker, and then randomly pick a task from that speaker
tasks = []
while len(tasks) < ntasks * percent:
spk = random.choice(list(spk2tasks.keys()))
if len(spk2tasks[spk]) == 0:
continue
tasks.append(spk2tasks[spk].pop()[:-1])
print(f"take only {percent*100:.2f}% of the tasks, {len(tasks)} out of {ntasks} tasks")
#### take only a subet of tasks
print(f"organizing done, took {time.time() - stime:.4f} seconds")
print(f"Launching {n_process} processes")
phn_vocab = set()
cnt = 0
with multiprocessing.Pool(processes=n_process) as pool:
for phns in tqdm.tqdm(pool.imap_unordered(cut_sequence, tasks), total=len(tasks)):
cnt += 1
if phns != None:
phn_vocab.update(phns)
# save phn vocabulary
if split == "train":
vocab_fn = os.path.join(output_dir, "vocab.txt")
with open(vocab_fn, "w") as f:
for i, phn in enumerate(list(phn_vocab)):
if i < len(list(phn_vocab)) - 1:
f.write(f"{str(i)}\t{phn}\n")
else:
f.write(f"{str(i)}\t{phn}")
def parse_args():
parser = argparse.ArgumentParser(description="Cut a dataset in small "
"sequences using VAD files")
parser.add_argument('--split', type=str, default='train', choices=['train', 'valid', 'test'], help="train = libriheavy_cuts_{small,medium,large}.jsonl.gz, valid = libriheavy_cuts_dev_{clean,other}.jsonl.gz, test = libriheavy_cuts_test_{clean,other}.jsonl.gz")
parser.add_argument('--audio_dir', type=str, default="/data/scratch/pyp/datasets/librilight_example",
help="Path to the audio directory")
parser.add_argument('--manifest_dir', type=str, default="/data/scratch/pyp/datasets/librilight/libriheavy", help="path to the transcription file's dir, can be downloaded https://huggingface.co/datasets/pkufool/libriheavy/tree/main/v0.1")
parser.add_argument('--output_dir', type=str, default="/data/scratch/pyp/datasets/librilight/librilight_example_preprocessed",
help="Path to the output directory")
parser.add_argument('--n_workers', type=int, default=16,
help="Number of parallel worker processes")
parser.add_argument('--percent', type=float, default=0.5, help="take only this percent of the tasks, randomly sampled from each speaker")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
pathlib.Path(args.output_dir).mkdir(exist_ok=True, parents=True)
text_tokenizer = TextTokenizer()
cut(args.split, args.audio_dir, args.manifest_dir, args.output_dir, args.n_workers, args.percent) |