# Get Transcription, WER and PPM """ TODO: [DONE]: Automatic generating Config """ import yaml import argparse import sys from pathlib import Path sys.path.append("./src") import lightning_module from UV import plot_UV, get_speech_interval from transformers import pipeline from rich.progress import track from rich import print as rprint import numpy as np import jiwer import pdb import torch.nn as nn import torch import torchaudio import gradio as gr from sys import flags from random import sample from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC # root_path = Path(__file__).parents[1] class ChangeSampleRate(nn.Module): def __init__(self, input_rate: int, output_rate: int): super().__init__() self.output_rate = output_rate self.input_rate = input_rate def forward(self, wav: torch.tensor) -> torch.tensor: # Only accepts 1-channel waveform input wav = wav.view(wav.size(0), -1) new_length = wav.size(-1) * self.output_rate // self.input_rate indices = torch.arange(new_length) * ( self.input_rate / self.output_rate ) round_down = wav[:, indices.long()] round_up = wav[:, (indices.long() + 1).clamp(max=wav.size(-1) - 1)] output = round_down * (1.0 - indices.fmod(1.0)).unsqueeze( 0 ) + round_up * indices.fmod(1.0).unsqueeze(0) return output model = lightning_module.BaselineLightningModule.load_from_checkpoint( "./src/epoch=3-step=7459.ckpt" ).eval() def calc_wer(audio_path, ref, ASR_pipeline): wav, sr = torchaudio.load(audio_path) osr = 16_000 batch = wav.unsqueeze(0).repeat(10, 1, 1) csr = ChangeSampleRate(sr, osr) out_wavs = csr(wav) # ASR trans = ASR_pipeline(audio_path)["text"] # WER wer = jiwer.wer( ref, trans, truth_transform=transformation, hypothesis_transform=transformation, ) return trans, wer if __name__ == "__main__": # Argparse parser = argparse.ArgumentParser( prog="get_ref_PPM", description="Generate Phoneme per Minute (and Voice/Unvoice plot)", epilog="", ) parser.add_argument( "--tag", type=str, default=None, required=False, help="ID tag for output *.csv", ) parser.add_argument("--ref_txt", type=str, required=True, help="Reference TXT") parser.add_argument( "--ref_wavs", type=str, required=True, help="Reference WAVs" ) parser.add_argument( "--metadata", type=str, required=False, help="metadata.csv including wav_id and reference", ) parser.add_argument( "--model", type=str, default='whisper-medium-FT', choices=['wav2vec+ctc', 'whipser-medium-FT', 'whipser-large-v2'], help="ASR engine for evaluation:\n ver1: wav2vec+ctc \n ver2: whipser-medium(Fined-tuned)\n ver3: whipser-large-v2", ) parser.add_argument( "--output_dir", type=str, required=True, help="Output Directory for *.csv", ) parser.add_argument( "--to_config", choices=["True", "False"], default="False", help="Generating Config from .txt and wavs/*wav", ) args = parser.parse_args() refs = np.loadtxt(args.ref_txt, delimiter="\n", dtype="str") refs_ids = [x.split()[0] for x in refs] refs_txt = [" ".join(x.split()[1:]) for x in refs] ref_wavs = [str(x) for x in sorted(Path(args.ref_wavs).glob("**/*.wav"))] # pdb.set_trace() try: len(refs) == len(ref_wavs) except ValueError: print("Error: Text and Wavs don't match") exit() # ASR part if args.model== "whisper-medium-FT": ASR_pipeline = pipeline("automatic-speech-recognition", model="KevinGeng/whipser_medium_en_PAL300_step25") elif args.model == "wav2vec+ctc": ASR_pipeline = pipeline("automatic-speech-recognition") elif args.model == "whisper-large-v2": ASR_pipeline = pipeline("automatic-speech-recognition", model="openai/whisper-large-v2") # pdb.set_trace() # WER part transformation = jiwer.Compose( [ jiwer.ToLowerCase(), jiwer.RemoveWhiteSpace(replace_by_space=True), jiwer.RemoveMultipleSpaces(), jiwer.ReduceToListOfListOfWords(word_delimiter=" "), ] ) # WPM part processor = Wav2Vec2Processor.from_pretrained( "facebook/wav2vec2-xlsr-53-espeak-cv-ft" ) phoneme_model = Wav2Vec2ForCTC.from_pretrained( "facebook/wav2vec2-xlsr-53-espeak-cv-ft" ) # phoneme_model = pipeline(model="facebook/wav2vec2-xlsr-53-espeak-cv-ft") description = """ MOS prediction demo using UTMOS-strong w/o phoneme encoder model, \ which is trained on the main track dataset. This demo only accepts .wav format. Best at 16 kHz sampling rate. Paper is available [here](https://arxiv.org/abs/2204.02152) Add ASR based on wav2vec-960, currently only English available. Add WER interface. """ referance_id = gr.Textbox( value="ID", placeholder="Utter ID", label="Reference_ID" ) referance_textbox = gr.Textbox( value="", placeholder="Input reference here", label="Reference" ) # Set up interface result = [] result.append("id,ref,hyp,wer") for id, x, y in track( zip(refs_ids, ref_wavs, refs_txt), total=len(refs_ids), description="Loading references information", ): trans, wer = calc_wer(x, y, ASR_pipeline=ASR_pipeline) record = ",".join( [ id, str(y), str(trans), str(wer) ] ) result.append(record) # Output if args.tag == None: args.tag = Path(args.ref_wavs).stem # Make output_dir # pdb.set_trace() Path.mkdir(Path(args.output_dir), exist_ok=True) # pdb.set_trace() with open("%s/%s.csv" % (args.output_dir, args.tag), "w") as f: print("\n".join(result), file=f) # Generating config if args.to_config == "True": config_dict = { "exp_id": args.tag, "ref_txt": args.ref_txt, "ref_feature": "%s/%s.csv" % (args.output_dir, args.tag), "ref_wavs": args.ref_wavs, "thre": { "minppm": 100, "maxppm": 100, "WER": 0.1, "AUTOMOS": 4.0, }, "auth": {"username": None, "password": None}, } with open("./config/%s.yaml" % args.tag, "w") as config_f: rprint("Dumping as config ./config/%s.yaml" % args.tag) rprint(config_dict) yaml.dump(config_dict, stream=config_f) rprint("Change parameter ./config/%s.yaml if necessary" % args.tag) print("Reference Dumping Finished")