Update app.py

app.py

@@ -1,65 +1,32 @@
+import librosa
+import gradio as gr
+import numpy as np
+from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
 import soundfile as sf
 import torch
-from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor,Wav2Vec2ProcessorWithLM
-import gradio as gr
-import sox
-import subprocess
-
 
-def read_file_and_process(wav_file):
-    filename = wav_file.split('.')[0]
-    filename_16k = filename + "16k.wav"
-    resampler(wav_file, filename_16k)
-    speech, _ = sf.read(filename_16k)
-    inputs = processor(speech, sampling_rate=16_000, return_tensors="pt", padding=True)
-    
-    return inputs
+# load model and tokenizer
+processor = Wav2Vec2Processor.from_pretrained("aditii09/facebook_english_asr")
+model = Wav2Vec2ForCTC.from_pretrained("aditii09/facebook_english_asr")
 
+def speech2text(audio):
+    sr, data = audio
 
-def resampler(input_file_path, output_file_path):
-    command = (
-        f"ffmpeg -hide_banner -loglevel panic -i {input_file_path} -ar 16000 -ac 1 -bits_per_raw_sample 16 -vn "
-        f"{output_file_path}"
-    )
-    subprocess.call(command, shell=True)
+    # resample to 16hz
+    data_16hz = librosa.resample(data[:,0].astype(np.float32),sr,16000)
 
+    # tokenize
+    input_values = processor([data_16hz], return_tensors="pt", padding="longest").input_values  # Batch size 1
 
-def parse_transcription_with_lm(logits):
-    result = processor_with_LM.batch_decode(logits.cpu().numpy())
-    text = result.text
-    transcription = text[0]
-    return transcription
+    # retrieve logits
+    logits = model(input_values).logits
 
-def parse_transcription(logits):
+    # take argmax and decode
     predicted_ids = torch.argmax(logits, dim=-1)
-    transcription = processor.decode(predicted_ids[0], skip_special_tokens=True)
-    return transcription
-
-def parse(wav_file, applyLM):
-    input_values = read_file_and_process(wav_file)
-    with torch.no_grad():
-        logits = model(**input_values).logits
-   
-    if applyLM:
-        return parse_transcription_with_lm(logits)
-    else:
-        return parse_transcription(logits)
-
-model_id = "aditii09/facebook_english_asr"
-
-processor = Wav2Vec2Processor.from_pretrained(model_id)
-processor_with_LM = Wav2Vec2ProcessorWithLM.from_pretrained(model_id)
-model = Wav2Vec2ForCTC.from_pretrained(model_id)
+    transcription = processor.batch_decode(predicted_ids)
 
-    
-input_ = gr.Audio(source="microphone", type="filepath") 
-txtbox = gr.Textbox(
-            label="Output from model will appear here:",
-            lines=5
-        )
-chkbox = gr.Checkbox(label="Apply LM", value=False)
+    return transcription[0].lower()  # batch size 1
 
+iface = gr.Interface(speech2text, "microphone", "text")
 
-gr.Interface(parse, inputs = [input_, chkbox],  outputs=txtbox,
-             streaming=True, interactive=True,
-             analytics_enabled=False, show_tips=False, enable_queue=True).launch(inline=False);
+iface.launch()