record realtime

app.py

@@ -120,7 +120,6 @@ def speech2text_en(input_file):
     return transcription
 
 
-
 """Machine translation"""
 vien_model_checkpoint = "datnth1709/finetuned_HelsinkiNLP-opus-mt-vi-en_PhoMT"
 envi_model_checkpoint = "datnth1709/finetuned_HelsinkiNLP-opus-mt-en-vi_PhoMT"
@@ -140,13 +139,47 @@ def translate_en2vi(English):
 def inference_vien(audio):
     vi_text = speech2text_vi(audio)
     en_text = translate_vi2en(vi_text)
-    return en_text
+    return vi_text, en_text
 
 def inference_envi(audio):
     en_text = speech2text_en(audio)
     vi_text = translate_en2vi(en_text)
-    return vi_text
+    return en_text, vi_text
 
+def transcribe_vi(audio, state_vi="", state_en=""):
+    ds = speech_file_to_array_fn(audio.name)
+    # infer model
+    input_values = processor(
+          ds["speech"],
+          sampling_rate=ds["sampling_rate"], 
+          return_tensors="pt"
+    ).input_values
+    # decode ctc output
+    logits = vi_model(input_values).logits[0]
+    pred_ids = torch.argmax(logits, dim=-1)
+    greedy_search_output = processor.decode(pred_ids)
+    beam_search_output = ngram_lm_model.decode(logits.cpu().detach().numpy(), beam_width=500)
+    state_vi += beam_search_output + " "
+    en_text = translate_vi2en(beam_search_output)
+    state_en += en_text + " "
+    return state_vi, state_en, state_vi, state_en
+
+def transcribe_en(audio, state_en="", state_vi=""):
+    speech = load_data(audio)
+    # Tokenize
+    input_values = eng_tokenizer(speech, return_tensors="pt").input_values
+    # Take logits
+    logits = eng_model(input_values).logits
+    # Take argmax
+    predicted_ids = torch.argmax(logits, dim=-1)
+    # Get the words from predicted word ids
+    transcription = eng_tokenizer.decode(predicted_ids[0])
+    # Output is all upper case
+    transcription = correct_casing(transcription.lower())
+    state_en += transcription + " "
+    vi_text = translate_en2vi(transcription)
+    state_vi += vi_text + " "
+    return state_en, state_vi, state_en, state_vi
 
 """Gradio demo"""
 
@@ -177,14 +210,25 @@ with gr.Blocks() as demo:
         with gr.TabItem("Speech2text and Vi-En Translation"):
             with gr.Row():
                 with gr.Column():
-                    vi_audio = gr.Audio(source="microphone", label="Input Vietnamese Audio", type="file", streaming=True)
+                    vi_audio = gr.Audio(source="microphone", label="Input Vietnamese Audio", type="file", streaming=False)
                     translate_button_vien_2 = gr.Button(value="Translate To English")
                 with gr.Column():
+                    speech2text_vi1 = gr.Textbox(label="Vietnamese Text")
                     english_out_2 = gr.Textbox(label="English Text")
                 
-            translate_button_vien_2.click(lambda vi_voice: inference_vien(vi_voice), inputs=vi_audio, outputs=english_out_2)
+            translate_button_vien_2.click(lambda vi_voice: inference_vien(vi_voice), inputs=vi_audio, outputs=[speech2text_vi1, english_out_2])
             gr.Examples(examples=vi_example_voice,
                         inputs=[vi_audio])
+        with gr.TabItem("Vi-En Realtime Translation"):
+            with gr.Row():
+                with gr.Column():
+                    vi_audio = gr.Audio(source="microphone", label="Input Vietnamese Audio", type="file", streaming=True)
+                    translate_button_vien_2 = gr.Button(value="Translate To English")
+                with gr.Column():
+                    speech2text_vi2 = gr.Textbox(label="Vietnamese Text")
+                    english_out_3 = gr.Textbox(label="English Text")
+            vi_audio.change(transcribe_vi, [vi_audio, "state_vi", "state_en"], [speech2text_vi2, english_out_3, "state_vi", "state_en"])
+
 
     with gr.Tabs():
         with gr.TabItem("Translation: English to Vietnamese"):
@@ -197,17 +241,28 @@ with gr.Blocks() as demo:
             translate_button_envi_1.click(lambda text: translate_en2vi(text), inputs=english_text, outputs=vietnamese_out_1)
             gr.Examples(examples=en_example_text,
                         inputs=[english_text])
+        
         with gr.TabItem("Speech2text and En-Vi Translation"):
             with gr.Row():
                 with gr.Column():
-                    en_audio = gr.Audio(source="microphone", label="Input English Audio", type="filepath", streaming=True)
+                    en_audio_1 = gr.Audio(source="microphone", label="Input English Audio", type="filepath", streaming=False)
                     translate_button_envi_2 = gr.Button(value="Translate To Vietnamese")
                 with gr.Column():
+                    speech2text_en1 = gr.Textbox(label="English Text")
                     vietnamese_out_2 = gr.Textbox(label="Vietnamese Text")
-
-            translate_button_envi_2.click(lambda en_voice: inference_envi(en_voice), inputs=en_audio, outputs=vietnamese_out_2)
+            translate_button_envi_2.click(lambda en_voice: inference_envi(en_voice), inputs=en_audio_1, outputs=[speech2text_en1, vietnamese_out_2])
             gr.Examples(examples=en_example_voice,
-                        inputs=[en_audio])
+                        inputs=[en_audio_1])
+        
+        with gr.TabItem("En-Vi Realtime Translation"):
+            with gr.Row():
+                with gr.Column():
+                    en_audio_2 = gr.Audio(source="microphone", label="Input English Audio", type="filepath", streaming=True)
+                    # translate_button_envi_2 = gr.Button(value="Translate To Vietnamese")
+                with gr.Column():
+                    speech2text_en2 = gr.Textbox(label="English Text")
+                    vietnamese_out_3 = gr.Textbox(label="Vietnamese Text")
+            en_audio_2.change(transcribe_en, [en_audio_2, "state_en", "state_vi"], [speech2text_en2, vietnamese_out_3, "state_en", "state_vi"])
 
 if __name__ == "__main__":
     demo.launch()