Upload 3 files

app.py

@@ -0,0 +1,44 @@
+import gradio as gr
+import torch
+import librosa
+import json
+from transformers import pipeline
+from stitched_model import CombinedModel
+
+
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+model = CombinedModel("ak3ra/wav2vec2-sunbird-speech-lug", "Sunbird/sunbird-mul-en-mbart-merged", device="cpu")
+
+
+def transcribe(audio_file_mic=None, audio_file_upload=None):
+    if audio_file_mic:
+        audio_file = audio_file_mic
+    elif audio_file_upload:
+        audio_file = audio_file_upload
+    else:
+        return "Please upload an audio file or record one"
+
+    # Make sure audio is 16kHz
+    speech, sample_rate = librosa.load(audio_file)
+    if sample_rate != 16000:
+        speech = librosa.resample(speech, orig_sr=sample_rate, target_sr=16000)
+    speech = torch.tensor([speech])
+
+    with torch.no_grad():
+        transcription, translation = model({"audio":speech})
+    
+    return transcription, translation[0]
+
+description = '''Luganda to English Speech Translation'''
+
+iface = gr.Interface(fn=transcribe,
+                     inputs=[
+                         gr.Audio(source="microphone", type="filepath", label="Record Audio"),
+                         gr.Audio(source="upload", type="filepath", label="Upload Audio")],
+                     outputs=[gr.Textbox(label="Transcription"),
+                              gr.Textbox(label="Translation")
+                     ],
+                     description=description
+                     )
+iface.launch()

requirements.txt

@@ -0,0 +1,4 @@
+transformers[torch]
+librosa
+sentencepiece
+jiwer

stitched_model.py

@@ -0,0 +1,27 @@
+import torch
+from torch import nn
+from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC, AutoTokenizer, AutoModelForSeq2SeqLM
+
+class CombinedModel(nn.Module):
+    def __init__(self, stt_model_name, nmt_model_name,device = "cuda"):
+        super(CombinedModel, self).__init__()
+
+        self.stt_processor = Wav2Vec2Processor.from_pretrained(stt_model_name)
+        self.stt_model = Wav2Vec2ForCTC.from_pretrained(stt_model_name)
+        self.nmt_tokenizer = AutoTokenizer.from_pretrained(nmt_model_name)
+        self.nmt_model = AutoModelForSeq2SeqLM.from_pretrained(nmt_model_name)
+        self.device = device
+
+    def forward(self, batch, *args, **kwargs):
+        # Use stt_model to transcribe the audio to text
+        device = self.device
+        audio = torch.tensor(batch["audio"][0]).to(self.device)
+        input_features = self.stt_processor(audio,sampling_rate=16000, return_tensors="pt",max_length=110000, padding=True, truncation=True)
+        stt_output = self.stt_model(input_features.input_values.to(device), attention_mask= input_features.attention_mask.to(device) )
+        transcription = self.stt_processor.decode(torch.squeeze(stt_output.logits.argmax(axis=-1)).to(device))
+        input_nmt_tokens = self.nmt_tokenizer(transcription, return_tensors="pt", padding=True, truncation=True)
+        output_nmt_output = self.nmt_model.generate(input_ids = input_nmt_tokens.input_ids.to(device), attention_mask= input_nmt_tokens.attention_mask.to(device))
+        decoded_nmt_output = self.nmt_tokenizer.batch_decode(output_nmt_output, skip_special_tokens=True)
+
+
+        return transcription, decoded_nmt_output