add toxicity mitigation to m4tv2

app.py

@@ -60,11 +60,13 @@ if torch.cuda.is_available():
 else:
     device = torch.device("cpu")
     dtype = torch.float32
+
 translator = Translator(
     model_name_or_card="seamlessM4T_v2_large",
     vocoder_name_or_card="vocoder_v2",
     device=device,
     dtype=dtype,
+    apply_mintox=True,
 )
 
 
@@ -78,12 +80,16 @@ def preprocess_audio(input_audio: str) -> None:
     torchaudio.save(input_audio, new_arr, sample_rate=int(AUDIO_SAMPLE_RATE))
 
 
-def run_s2st(input_audio: str, target_language: str) -> tuple[tuple[int, np.ndarray] | None, str]:
+def run_s2st(
+    input_audio: str, source_language: str, target_language: str
+) -> tuple[tuple[int, np.ndarray] | None, str]:
     preprocess_audio(input_audio)
+    source_language_code = LANGUAGE_NAME_TO_CODE[source_language]
     target_language_code = LANGUAGE_NAME_TO_CODE[target_language]
     out_texts, out_audios = translator.predict(
         input=input_audio,
         task_str="S2ST",
+        src_lang=source_language_code,
         tgt_lang=target_language_code,
     )
     out_text = str(out_texts[0])
@@ -91,13 +97,15 @@ def run_s2st(input_audio: str, target_language: str) -> tuple[tuple[int, np.ndar
     return (int(AUDIO_SAMPLE_RATE), out_wav), out_text
 
 
-def run_s2tt(input_audio: str, target_language: str) -> str:
+def run_s2tt(input_audio: str, source_language: str, target_language: str) -> str:
     preprocess_audio(input_audio)
+    source_language_code = LANGUAGE_NAME_TO_CODE[source_language]
     target_language_code = LANGUAGE_NAME_TO_CODE[target_language]
     out_texts, _ = translator.predict(
         input=input_audio,
         task_str="S2TT",
         tgt_lang=target_language_code,
+        src_lang=source_language_code,
     )
     return str(out_texts[0])
 
@@ -144,6 +152,11 @@ with gr.Blocks() as demo_s2st:
         with gr.Column():
             with gr.Group():
                 input_audio = gr.Audio(label="Input speech", type="filepath")
+                source_language = gr.Dropdown(
+                    label="Source language",
+                    choices=ASR_TARGET_LANGUAGE_NAMES,
+                    value="English",
+                )
                 target_language = gr.Dropdown(
                     label="Target language",
                     choices=S2ST_TARGET_LANGUAGE_NAMES,
@@ -162,12 +175,12 @@ with gr.Blocks() as demo_s2st:
 
     gr.Examples(
         examples=[
-            ["assets/sample_input.mp3", "French"],
-            ["assets/sample_input.mp3", "Mandarin Chinese"],
-            ["assets/sample_input_2.mp3", "Hindi"],
-            ["assets/sample_input_2.mp3", "Spanish"],
+            ["assets/sample_input.mp3", "English", "French"],
+            ["assets/sample_input.mp3", "English", "Mandarin Chinese"],
+            ["assets/sample_input_2.mp3", "English", "Hindi"],
+            ["assets/sample_input_2.mp3", "English", "Spanish"],
         ],
-        inputs=[input_audio, target_language],
+        inputs=[input_audio, source_language, target_language],
         outputs=[output_audio, output_text],
         fn=run_s2st,
         cache_examples=CACHE_EXAMPLES,
@@ -176,7 +189,7 @@ with gr.Blocks() as demo_s2st:
 
     btn.click(
         fn=run_s2st,
-        inputs=[input_audio, target_language],
+        inputs=[input_audio, source_language, target_language],
         outputs=[output_audio, output_text],
         api_name="s2st",
     )
@@ -186,6 +199,11 @@ with gr.Blocks() as demo_s2tt:
         with gr.Column():
             with gr.Group():
                 input_audio = gr.Audio(label="Input speech", type="filepath")
+                source_language = gr.Dropdown(
+                    label="Source language",
+                    choices=ASR_TARGET_LANGUAGE_NAMES,
+                    value="English",
+                )
                 target_language = gr.Dropdown(
                     label="Target language",
                     choices=S2TT_TARGET_LANGUAGE_NAMES,
@@ -197,12 +215,12 @@ with gr.Blocks() as demo_s2tt:
 
     gr.Examples(
         examples=[
-            ["assets/sample_input.mp3", "French"],
-            ["assets/sample_input.mp3", "Mandarin Chinese"],
-            ["assets/sample_input_2.mp3", "Hindi"],
-            ["assets/sample_input_2.mp3", "Spanish"],
+            ["assets/sample_input.mp3", "English", "French"],
+            ["assets/sample_input.mp3", "English", "Mandarin Chinese"],
+            ["assets/sample_input_2.mp3", "English", "Hindi"],
+            ["assets/sample_input_2.mp3", "English", "Spanish"],
         ],
-        inputs=[input_audio, target_language],
+        inputs=[input_audio, source_language, target_language],
         outputs=output_text,
         fn=run_s2tt,
         cache_examples=CACHE_EXAMPLES,
@@ -211,7 +229,7 @@ with gr.Blocks() as demo_s2tt:
 
     btn.click(
         fn=run_s2tt,
-        inputs=[input_audio, target_language],
+        inputs=[input_audio, source_language, target_language],
         outputs=output_text,
         api_name="s2tt",
     )