Add application file

app.py

@@ -0,0 +1,194 @@
+#!/usr/bin/env python
+# Copyright (c) Meta Platforms, Inc. and affiliates
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# MIT_LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+import os
+import pathlib
+import getpass
+from typing import Any, Dict
+
+import gradio as gr
+import numpy as np
+import torch
+import torchaudio
+from fairseq2.assets import InProcAssetMetadataProvider, asset_store
+from huggingface_hub import snapshot_download
+from seamless_communication.inference import Translator
+
+from langlist_slt import (
+    LANGUAGE_NAME_TO_CODE,
+    S2TT_TARGET_LANGUAGE_NAMES,
+    ASR_TARGET_LANGUAGE_NAMES,
+)
+os.environ["GRADIO_TEMP_DIR"] = "/data/bhavana/app/tmp"
+os.makedirs(os.environ["GRADIO_TEMP_DIR"], exist_ok=True)
+user = getpass.getuser() # this is not portable on windows
+CHECKPOINTS_PATH = pathlib.Path(os.getenv("CHECKPOINTS_PATH", f"/data/bhavana/app/models"))
+if not CHECKPOINTS_PATH.exists():
+    snapshot_download(repo_id="facebook/hf-seamless-m4t-medium", repo_type="model", local_dir=CHECKPOINTS_PATH)
+asset_store.env_resolvers.clear()
+asset_store.env_resolvers.append(lambda: "demo")
+demo_metadata = [
+    {
+        "name": "seamlessM4T_v2_medium@demo",
+        "checkpoint": f"file://{CHECKPOINTS_PATH}/seamlessM4T_medium.pt",
+        "char_tokenizer": f"file://{CHECKPOINTS_PATH}/spm_char_lang38_tc.model",
+    },
+    {
+        "name": "vocoder_v2@demo",
+        "checkpoint": f"file://{CHECKPOINTS_PATH}/vocoder_v2.pt",
+    },
+]
+asset_store.metadata_providers.append(InProcAssetMetadataProvider(demo_metadata))
+
+DESCRIPTION = """\
+#"IIITH-SLT: End-to-End Speech Translation demo for low-resource Indian languages using weakly labeled data. Supports ST models for Bengali-Hindi, 
+# Malayalam-Hindi, Odia-Hindi, and Telugu-Hindi, trained on the Shrutilipi-anuvaad dataset."
+[Paper](https://arxiv.org/pdf/2506.16251)
+"""
+
+CACHE_EXAMPLES = os.getenv("CACHE_EXAMPLES") == "1" and torch.cuda.is_available()
+
+AUDIO_SAMPLE_RATE = 16000.0
+MAX_INPUT_AUDIO_LENGTH = 60  # in seconds
+DEFAULT_TARGET_LANGUAGE = "Hindi"
+
+if torch.cuda.is_available():
+    device = torch.device("cuda:0")
+    dtype = torch.float16
+else:
+    device = torch.device("cpu")
+    dtype = torch.float32
+
+FINETUNED_MODEL_MAP = {
+    "Telugu": "/data/aishwarya/seamless_communication/src/seamless_communication/cli/m4t_original/finetune_original/icassp/checkpoint_te_hi_v5.pt",
+    "Malayalam": "/data/aishwarya/seamless_communication/src/seamless_communication/cli/m4t_original/finetune_original/icassp/checkpoint_ml_hi_v5.pt",
+    "Odiya": "/data/aishwarya/seamless_communication/src/seamless_communication/cli/m4t_original/finetune_original/icassp/checkpoint_od_hi_v5.pt",
+    "Bengali": "/data/aishwarya/seamless_communication/src/seamless_communication/cli/m4t_original/finetune_original/icassp/checkpoint_bn_hi_v5.pt ",
+    # Add more
+}
+
+
+from functools import lru_cache
+
+@lru_cache(maxsize=None)
+def load_translator_for_language(language: str) -> Translator:
+    # Load base model (Meta's SeamlessM4T)
+    translator = Translator(
+        model_name_or_card="seamlessM4T_v2_large",
+        vocoder_name_or_card=None,
+        device=device,
+        dtype=dtype,
+        apply_mintox=False,
+    )
+
+    # Apply language-specific fine-tuned weights
+    ckpt_path = FINETUNED_MODEL_MAP.get(language)
+    if ckpt_path and os.path.exists(ckpt_path):
+        print(f"Loading fine-tuned checkpoint for {language} from {ckpt_path}")
+        saved_model = torch.load(ckpt_path, map_location=device)["model"]
+        saved_model = {k.replace("module.", ""): v for k, v in saved_model.items()}
+
+        def _select_keys(state_dict: Dict[str, Any], prefix: str) -> Dict[str, Any]:
+            return {key.replace(prefix, ""): value for key, value in state_dict.items() if key.startswith(prefix)}
+
+        translator.model.speech_encoder_frontend.load_state_dict(_select_keys(saved_model, "model.speech_encoder_frontend."))
+        translator.model.speech_encoder.load_state_dict(_select_keys(saved_model, "model.speech_encoder."))
+
+        assert translator.model.text_decoder_frontend is not None
+        translator.model.text_decoder_frontend.load_state_dict(_select_keys(saved_model, "model.text_decoder_frontend."))
+
+        assert translator.model.text_decoder is not None
+        translator.model.text_decoder.load_state_dict(_select_keys(saved_model, "model.text_decoder."))
+
+        assert translator.model.final_proj is not None
+        translator.model.final_proj.load_state_dict(_select_keys(saved_model, "model.final_proj."))
+
+    return translator
+
+
+def preprocess_audio(input_audio: str) -> None:
+    arr, org_sr = torchaudio.load(input_audio)
+    new_arr = torchaudio.functional.resample(arr, orig_freq=org_sr, new_freq=AUDIO_SAMPLE_RATE)
+    max_length = int(MAX_INPUT_AUDIO_LENGTH * AUDIO_SAMPLE_RATE)
+    if new_arr.shape[1] > max_length:
+        new_arr = new_arr[:, :max_length]
+        gr.Warning(f"Input audio is too long. Only the first {MAX_INPUT_AUDIO_LENGTH} seconds is used.")
+    torchaudio.save(input_audio, new_arr, sample_rate=int(AUDIO_SAMPLE_RATE))
+
+
+def run_s2tt(input_audio: str, source_language: str, target_language: str):
+    preprocess_audio(input_audio)
+
+    translator = load_translator_for_language(source_language)
+
+    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="S2ST",
+        src_lang=source_language_code,
+        tgt_lang=target_language_code,
+    )
+    return str(out_texts[0])
+
+
+
+with gr.Blocks() as demo_s2tt:
+    with gr.Row():
+        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="Telugu",
+                )
+                target_language = gr.Dropdown(
+                    label="Target language",
+                    choices=S2TT_TARGET_LANGUAGE_NAMES,
+                    value=DEFAULT_TARGET_LANGUAGE,
+                )
+            btn = gr.Button("Translate")
+        with gr.Column():
+            output_text = gr.Textbox(label="Translated text")
+
+    gr.Examples(
+        examples=[],
+        inputs=[input_audio, source_language, target_language],
+        outputs=output_text,
+        fn=run_s2tt,
+        cache_examples=CACHE_EXAMPLES,
+        api_name=False,
+    )
+
+    btn.click(
+        fn=run_s2tt,
+        inputs=[input_audio, source_language, target_language],
+        outputs=output_text,
+        api_name="s2tt",
+    )
+
+
+
+
+
+with gr.Blocks(css="style.css") as demo:
+    gr.Markdown(DESCRIPTION)
+    gr.DuplicateButton(
+        value="Duplicate Space for private use",
+        elem_id="duplicate-button",
+        visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
+    )
+
+    with gr.Tabs():
+        demo_s2tt.render()
+
+
+if __name__ == "__main__":
+    demo.queue(max_size=50).launch(share=True)

langlist_slt.py

@@ -0,0 +1,26 @@
+# Language dict
+language_code_to_name = {
+    "ben": "Bengali",
+    "hin": "Hindi",
+    "mal": "Malayalam",
+    "tel": "Telugu",
+    "odi": "Odiya",
+}
+LANGUAGE_NAME_TO_CODE = {v: k for k, v in language_code_to_name.items()}
+
+# Source langs: S2ST / S2TT / ASR don't need source lang
+# T2TT / T2ST use this
+text_source_language_codes = [
+    "tel",
+    "odi",
+    "mal",
+    "ben",
+    "hin",
+]
+TEXT_SOURCE_LANGUAGE_NAMES = sorted([language_code_to_name[code] for code in text_source_language_codes])
+
+
+
+# S2TT / T2TT / ASR
+S2TT_TARGET_LANGUAGE_NAMES = TEXT_SOURCE_LANGUAGE_NAMES
+ASR_TARGET_LANGUAGE_NAMES = TEXT_SOURCE_LANGUAGE_NAMES

requirements.txt

@@ -0,0 +1,5 @@
+gradio~=4.5.0
+omegaconf~=2.3.0
+torch~=2.1.0
+torchaudio~=2.1.0
+fairseq2~=0.2.0