added dataset transcription helper script (now I don't ever have to touch ai-voice-cloning) (to-do: unify scripts into the module)
This commit is contained in:
parent
b251669536
commit
ffc334cf58
79
scripts/transcribe_dataset.py
Normal file
79
scripts/transcribe_dataset.py
Normal file
|
@ -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))
|
|
@ -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 ["<s>"] + [ " " if not p else p for p in content ] + ["</s>"]
|
||||
|
||||
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)
|
||||
|
||||
# <s>[original text] [new text]</s>
|
||||
|
|
Loading…
Reference in New Issue
Block a user