reach-vb/asr-pyctcdecode

1import nltk2import librosa3import torch4import kenlm5import gradio as gr6from pyctcdecode import build_ctcdecoder7from transformers import Wav2Vec2Processor, AutoModelForCTC8 9nltk.download("punkt")10 11wav2vec2processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h") 12wav2vec2model = AutoModelForCTC.from_pretrained("facebook/wav2vec2-base-960h")13hubertprocessor = Wav2Vec2Processor.from_pretrained("facebook/hubert-large-ls960-ft") 14hubertmodel = AutoModelForCTC.from_pretrained("facebook/hubert-large-ls960-ft")15 16def return_processor_and_model(model_name):17    return Wav2Vec2Processor.from_pretrained(model_name), AutoModelForCTC.from_pretrained(model_name)18 19def load_and_fix_data(input_file):  20  speech, sample_rate = librosa.load(input_file)21  if len(speech.shape) > 1: 22      speech = speech[:,0] + speech[:,1]23  if sample_rate !=16000:24    speech = librosa.resample(speech, sample_rate,16000)25  return speech26 27def fix_transcription_casing(input_sentence):28  sentences = nltk.sent_tokenize(input_sentence)29  return (' '.join([s.replace(s[0],s[0].capitalize(),1) for s in sentences]))30  31def predict_and_ctc_decode(input_file, model_name):32  processor, model = return_processor_and_model(model_name)33  speech = load_and_fix_data(input_file)34 35  input_values = processor(speech, return_tensors="pt", sampling_rate=16000).input_values36  logits = model(input_values).logits.cpu().detach().numpy()[0]37  38  vocab_list = list(processor.tokenizer.get_vocab().keys())  39  decoder = build_ctcdecoder(vocab_list)40  pred = decoder.decode(logits)41 42  transcribed_text = fix_transcription_casing(pred.lower())43 44  return transcribed_text45 46def predict_and_ctc_lm_decode(input_file, model_name):47  processor, model = return_processor_and_model(model_name)48  speech = load_and_fix_data(input_file)49 50  input_values = processor(speech, return_tensors="pt", sampling_rate=16000).input_values51  logits = model(input_values).logits.cpu().detach().numpy()[0]52  53  vocab_list = list(processor.tokenizer.get_vocab().keys())  54  vocab_dict = processor.tokenizer.get_vocab()55  sorted_dict = {k.lower(): v for k, v in sorted(vocab_dict.items(), key=lambda item: item[1])}56 57  decoder = build_ctcdecoder(58    list(sorted_dict.keys()),59    "4gram_small.arpa.gz",60    )61 62  pred = decoder.decode(logits)63 64  transcribed_text = fix_transcription_casing(pred.lower())65 66  return transcribed_text67 68def predict_and_greedy_decode(input_file, model_name):69  processor, model = return_processor_and_model(model_name)70  speech = load_and_fix_data(input_file)71 72  input_values = processor(speech, return_tensors="pt", sampling_rate=16000).input_values73  logits = model(input_values).logits74 75  predicted_ids = torch.argmax(logits, dim=-1)76  pred = processor.batch_decode(predicted_ids)77 78  transcribed_text = fix_transcription_casing(pred[0].lower())79 80  return transcribed_text81 82def return_all_predictions(input_file, model_name):83  return predict_and_ctc_decode(input_file, model_name), predict_and_ctc_lm_decode(input_file, model_name), predict_and_greedy_decode(input_file, model_name)84 85 86gr.Interface(return_all_predictions,87             inputs = [gr.inputs.Audio(source="microphone", type="filepath", label="Record/ Drop audio"), gr.inputs.Dropdown(["facebook/wav2vec2-base-960h", "facebook/hubert-large-ls960-ft"], label="Model Name")],88             outputs = [gr.outputs.Textbox(label="Beam CTC decoding"), gr.outputs.Textbox(label="Beam CTC decoding w/ LM"), gr.outputs.Textbox(label="Greedy decoding")],89             title="ASR using Wav2Vec2/ Hubert & pyctcdecode",90             description = "Comparing greedy decoder with beam search CTC decoder, record/ drop your audio!",91             layout = "horizontal",92             examples = [["test1.wav", "facebook/wav2vec2-base-960h"], ["test2.wav", "facebook/hubert-large-ls960-ft"]], 93             theme="huggingface",94             enable_queue=True).launch()