diff --git a/scripts/process_old_dataset.py b/scripts/process_dataset.py
similarity index 100%
rename from scripts/process_old_dataset.py
rename to scripts/process_dataset.py
diff --git a/scripts/transcribe_dataset.py b/scripts/transcribe_dataset.py
new file mode 100644
index 0000000..ca8dcfa
--- /dev/null
+++ b/scripts/transcribe_dataset.py
@@ -0,0 +1,79 @@
+import os
+import json
+import torch
+import torchaudio
+import whisperx
+
+from tqdm.auto import tqdm
+from pathlib import Path
+
+device = "cuda"
+batch_size = 16
+dtype = "float16"
+model_size = "large-v2"
+
+input_audio = "voice"
+output_dataset = "metadata"
+skip_existing = True
+
+model = whisperx.load_model(model_size, device, compute_type=dtype)
+
+align_model, align_model_metadata, align_model_language = (None, None, None)
+
+def pad(num, zeroes):
+ return str(num).zfill(zeroes+1)
+
+for dataset_name in os.listdir(f'./{input_audio}/'):
+ if not os.path.isdir(f'./{input_audio}/{dataset_name}/'):
+ print("Is not dir:", f'./{input_audio}/{dataset_name}/')
+ continue
+
+ for speaker_id in tqdm(os.listdir(f'./{input_audio}/{dataset_name}/'), desc="Processing speaker"):
+ if not os.path.isdir(f'./{input_audio}/{dataset_name}/{speaker_id}'):
+ print("Is not dir:", f'./{input_audio}/{dataset_name}/{speaker_id}')
+ continue
+
+ outpath = Path(f'./{output_dataset}/{dataset_name}/{speaker_id}/whisper.json')
+
+ if outpath.exists():
+ metadata = json.loads(open(outpath, 'r', encoding='utf-8').read())
+ else:
+ os.makedirs(f'./{output_dataset}/{dataset_name}/{speaker_id}/', exist_ok=True)
+ metadata = {}
+
+ for filename in tqdm(os.listdir(f'./{input_audio}/{dataset_name}/{speaker_id}/'), desc=f"Processing speaker: {speaker_id}"):
+
+ if skip_existing and filename in metadata:
+ continue
+
+ inpath = f'./{input_audio}/{dataset_name}/{speaker_id}/{filename}'
+
+ metadata[filename] = {
+ "segments": [],
+ "language": "",
+ "text": [],
+ }
+
+ audio = whisperx.load_audio(inpath)
+ result = model.transcribe(audio, batch_size=batch_size)
+ language = result["language"]
+
+ if align_model_language != language:
+ tqdm.write(f'Loading language: {language}')
+ align_model, align_model_metadata = whisperx.load_align_model(language_code=language, device=device)
+ align_model_language = language
+
+ result = whisperx.align(result["segments"], align_model, align_model_metadata, audio, device, return_char_alignments=False)
+
+ metadata[filename]["segments"] = result["segments"]
+ metadata[filename]["language"] = language
+
+ text = []
+ for segment in result["segments"]:
+ id = len(text)
+ text.append( segment["text"] )
+ metadata[filename]["segments"][id]["id"] = id
+
+ metadata[filename]["text"] = " ".join(text).strip()
+
+ open(outpath, 'w', encoding='utf-8').write(json.dumps(metadata))
\ No newline at end of file
diff --git a/vall_e/data.py b/vall_e/data.py
index 434b3d3..cf512be 100755
--- a/vall_e/data.py
+++ b/vall_e/data.py
@@ -177,10 +177,8 @@ def _get_phones(path, language="en"):
content = metadata["phonemes"]
else:
content = open(_get_phone_path(path), "r", encoding="utf-8").read().split(" ")
- content = _cleanup_phones( content )
return "".join(content)
- #return [""] + [ " " if not p else p for p in content ] + [""]
def _interleaved_reorder(l, fn):
groups = defaultdict(list)
@@ -431,8 +429,6 @@ class Dataset(_Dataset):
text = cfg.hdf5[key]["text"][:]
resps = cfg.hdf5[key]["audio"][:, :]
-
- text = np.array( _cleanup_phones( text, targets=[ self.phone_symmap[" "] ] ) )
text = torch.from_numpy(text).to(self.text_dtype)
resps = torch.from_numpy(resps).to(torch.int16)
@@ -455,12 +451,13 @@ class Dataset(_Dataset):
txt = cfg.hdf5[key]["text"][:]
qnt = cfg.hdf5[key]["audio"][:, :]
- txt = np.array( _cleanup_phones( txt, targets=[ self.phone_symmap[" "] ] ) )
+ txt = np.array( txt )
txt = torch.from_numpy(txt).to(self.text_dtype)
qnt = torch.from_numpy(qnt).to(torch.int16)
else:
- txt = torch.tensor([*map(self.phone_symmap.get, _get_phones(sampled_path))]).to(self.text_dtype)
+ #txt = torch.tensor([*map(self.phone_symmap.get, _get_phones(sampled_path))]).to(self.text_dtype)
+ txt = torch.tensor(tokenize(_get_phones(sampled_path))).to(self.text_dtype)
qnt = _load_quants(sampled_path)
# [original text] [new text]