Added file conversion to uploaded files. Added chunking to allow transcribing long audio files.

.gitignore

@@ -0,0 +1,2 @@
+.DS_Store
+experiments.ipynb

app.py

@@ -1,25 +1,28 @@
 import streamlit as st
-from transformers import Wav2Vec2ForCTC, Wav2Vec2ProcessorWithLM
-import torch
+from transformers import pipeline
+from torch import cuda
 import torchaudio
 import torchaudio.functional as F
+from pydub import AudioSegment
+import logging
+import io
 
 
 class ASR:
     def __init__(self):
         self.model_name = "viktor-enzell/wav2vec2-large-voxrex-swedish-4gram"
-        self.device = torch.device(
-            "cuda" if torch.cuda.is_available() else "cpu")
+        self.device = cuda.current_device() if cuda.is_available() else -1
         self.model = None
-        self.processor = None
 
     def load_model(self):
-        self.model = Wav2Vec2ForCTC.from_pretrained(
-            self.model_name).to(self.device)
-        self.processor = Wav2Vec2ProcessorWithLM.from_pretrained(
-            self.model_name)
+        self.model = pipeline(model=self.model_name, device=self.device)
 
     def run_inference(self, file):
+        audio = self.load_16khz_audio(file)
+        return self.model(audio, chunk_length_s=10)["text"].lower()
+
+    @staticmethod
+    def load_16khz_audio(file):
         waveform, sample_rate = torchaudio.load(file)
 
         if sample_rate == 16_000:
@@ -27,17 +30,7 @@ class ASR:
         else:
             waveform = F.resample(waveform, sample_rate, 16_000)[0]
 
-        inputs = self.processor(
-            waveform,
-            sampling_rate=16_000,
-            return_tensors="pt",
-            padding=True
-        ).to(self.device)
-
-        with torch.no_grad():
-            logits = self.model(**inputs).logits
-
-        return self.processor.batch_decode(logits.cpu().numpy()).text[0].lower()
+        return waveform.numpy()
 
 
 @st.cache(allow_output_mutation=True, show_spinner=False)
@@ -52,6 +45,26 @@ def run_inference(asr, file):
     return asr.run_inference(file)
 
 
+def convert_uploaded_file_to_wav(file):
+    try:
+        media_type = file.type.split("/")[0]
+        file_extension = file.name.split(".")[-1]
+
+        if media_type != "audio" and media_type != "video":
+            return None
+
+        if file_extension == "wav":
+            return file
+
+        audio = AudioSegment.from_file(file, file_extension)
+        in_memory_buffer = io.BytesIO()
+        return audio.export(in_memory_buffer, format="wav")
+
+    except Exception as e:
+        logging.exception(e)
+        return None
+
+
 if __name__ == "__main__":
     st.set_page_config(
         page_title="Swedish Speech-to-Text",
@@ -64,25 +77,25 @@ if __name__ == "__main__":
     st.markdown("""
     # Swedish Speech-to-text
 
-    Generate and download high-quality Swedish transcripts for your audio files. The speech-to-text model is KBLab's wav2vec 2.0 large VoxRex Swedish (C) with a 4-gram language model, which you can access [here](https://huggingface.co/viktor-enzell/wav2vec2-large-voxrex-swedish-4gram).
+    Generate and download high-quality Swedish transcripts for your audio and video files. The speech-to-text model is KBLab's wav2vec 2.0 large VoxRex Swedish (C) with a 4-gram language model, which you can access [here](https://huggingface.co/viktor-enzell/wav2vec2-large-voxrex-swedish-4gram).
     """)
 
     with st.spinner(text="Loading model..."):
         asr = load_model()
 
-    uploaded_file = st.file_uploader("Choose a file", type=[".wav"])
+    uploaded_file = st.file_uploader("Choose a file")
+
     if uploaded_file is not None:
-        if uploaded_file.type != "audio/wav":
-            pass
-            # TODO: convert to wav
-            # bytes = uploaded_file.getvalue()
-            # audio_input = ffmpeg.input(bytes).audio
-            # audio_output = ffmpeg.output(audio_input, "tmp.wav", format="wav")
-            # ffmpeg.run(audio_output)
-
-        with st.spinner(text="Transcribing..."):
-            transcript = run_inference(asr, uploaded_file)
-            st.download_button("Download transcript", transcript, "transcript.txt")
-
-        with st.expander("Transcript", expanded=True):
-            st.write(transcript)
+        file = convert_uploaded_file_to_wav(uploaded_file)
+
+        if file is None:
+            st.error(
+                "There was a problem handling the uploaded file. Try again using an audio or video file.")
+        else:
+            with st.spinner(text="Transcribing..."):
+                transcript = run_inference(asr, file)
+                st.download_button("Download transcript",
+                                   transcript, "transcript.txt")
+
+            with st.expander("Transcript", expanded=True):
+                st.write(transcript)

requirements.txt

@@ -3,3 +3,4 @@ torchaudio==0.10.1
 transformers==4.19.2
 pyctcdecode==0.3.0
 https://github.com/kpu/kenlm/archive/master.zip
+pydub==0.25.1