wehub-resource-sync
129 lines
4.5 KiB
Python
129 lines
4.5 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import argparse
|
|
import editdistance
|
|
import re
|
|
import shutil
|
|
import soundfile as sf
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
from examples.speech_to_text.data_utils import load_tsv_to_dicts
|
|
|
|
|
|
def preprocess_text(text):
|
|
text = "|".join(re.sub(r"[^A-Z' ]", " ", text.upper()).split())
|
|
text = " ".join(text)
|
|
return text
|
|
|
|
|
|
def prepare_w2v_data(
|
|
dict_dir, sample_rate, label, audio_paths, texts, split, data_dir
|
|
):
|
|
data_dir.mkdir(parents=True, exist_ok=True)
|
|
shutil.copyfile(
|
|
dict_dir / f"dict.{label}.txt",
|
|
data_dir / f"dict.{label}.txt"
|
|
)
|
|
with open(data_dir / f"{split}.tsv", "w") as f:
|
|
f.write("/\n")
|
|
for audio_path in audio_paths:
|
|
wav, sr = sf.read(audio_path)
|
|
assert sr == sample_rate, f"{sr} != sample_rate"
|
|
nsample = len(wav)
|
|
f.write(f"{audio_path}\t{nsample}\n")
|
|
with open(data_dir / f"{split}.{label}", "w") as f:
|
|
for text in texts:
|
|
text = preprocess_text(text)
|
|
f.write(f"{text}\n")
|
|
|
|
|
|
def run_asr(asr_dir, split, w2v_ckpt, w2v_label, res_dir):
|
|
"""
|
|
results will be saved at
|
|
{res_dir}/{ref,hypo}.word-{w2v_ckpt.filename}-{split}.txt
|
|
"""
|
|
cmd = ["python", "-m", "examples.speech_recognition.infer"]
|
|
cmd += [str(asr_dir.resolve())]
|
|
cmd += ["--task", "audio_finetuning", "--nbest", "1", "--quiet"]
|
|
cmd += ["--w2l-decoder", "viterbi", "--criterion", "ctc"]
|
|
cmd += ["--post-process", "letter", "--max-tokens", "4000000"]
|
|
cmd += ["--path", str(w2v_ckpt.resolve()), "--labels", w2v_label]
|
|
cmd += ["--gen-subset", split, "--results-path", str(res_dir.resolve())]
|
|
|
|
print(f"running cmd:\n{' '.join(cmd)}")
|
|
subprocess.run(cmd, check=True)
|
|
|
|
|
|
def compute_error_rate(hyp_wrd_path, ref_wrd_path, unit="word"):
|
|
"""each line is "<text> (None-<index>)" """
|
|
tokenize_line = {
|
|
"word": lambda x: re.sub(r" \(.*\)$", "", x.rstrip()).split(),
|
|
"char": lambda x: list(re.sub(r" \(.*\)$", "", x.rstrip()))
|
|
}.get(unit)
|
|
if tokenize_line is None:
|
|
raise ValueError(f"{unit} not supported")
|
|
|
|
inds = [int(re.sub(r"\D*(\d*)\D*", r"\1", line))
|
|
for line in open(hyp_wrd_path)]
|
|
hyps = [tokenize_line(line) for line in open(hyp_wrd_path)]
|
|
refs = [tokenize_line(line) for line in open(ref_wrd_path)]
|
|
assert(len(hyps) == len(refs))
|
|
err_rates = [
|
|
editdistance.eval(hyp, ref) / len(ref) for hyp, ref in zip(hyps, refs)
|
|
]
|
|
ind_to_err_rates = {i: e for i, e in zip(inds, err_rates)}
|
|
return ind_to_err_rates
|
|
|
|
|
|
def main(args):
|
|
samples = load_tsv_to_dicts(args.raw_manifest)
|
|
ids = [
|
|
sample[args.id_header] if args.id_header else "" for sample in samples
|
|
]
|
|
audio_paths = [sample[args.audio_header] for sample in samples]
|
|
texts = [sample[args.text_header] for sample in samples]
|
|
|
|
prepare_w2v_data(
|
|
args.w2v_dict_dir,
|
|
args.w2v_sample_rate,
|
|
args.w2v_label,
|
|
audio_paths,
|
|
texts,
|
|
args.split,
|
|
args.asr_dir
|
|
)
|
|
run_asr(args.asr_dir, args.split, args.w2v_ckpt, args.w2v_label, args.asr_dir)
|
|
ind_to_err_rates = compute_error_rate(
|
|
args.asr_dir / f"hypo.word-{args.w2v_ckpt.name}-{args.split}.txt",
|
|
args.asr_dir / f"ref.word-{args.w2v_ckpt.name}-{args.split}.txt",
|
|
args.err_unit,
|
|
)
|
|
|
|
uer_path = args.asr_dir / f"uer_{args.err_unit}.{args.split}.tsv"
|
|
with open(uer_path, "w") as f:
|
|
f.write("id\taudio\tuer\n")
|
|
for ind, (id_, audio_path) in enumerate(zip(ids, audio_paths)):
|
|
f.write(f"{id_}\t{audio_path}\t{ind_to_err_rates[ind]:.4f}\n")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--raw-manifest", required=True, type=Path)
|
|
parser.add_argument("--asr-dir", required=True, type=Path)
|
|
parser.add_argument("--id-header", default="id", type=str)
|
|
parser.add_argument("--audio-header", default="audio", type=str)
|
|
parser.add_argument("--text-header", default="src_text", type=str)
|
|
parser.add_argument("--split", default="raw", type=str)
|
|
parser.add_argument("--w2v-ckpt", required=True, type=Path)
|
|
parser.add_argument("--w2v-dict-dir", required=True, type=Path)
|
|
parser.add_argument("--w2v-sample-rate", default=16000, type=int)
|
|
parser.add_argument("--w2v-label", default="ltr", type=str)
|
|
parser.add_argument("--err-unit", default="word", type=str)
|
|
args = parser.parse_args()
|
|
|
|
main(args)
|