Update app.py

app.py

@@ -1,26 +1,151 @@
+import re
 import os
+import sys
+
+import torch
 import gradio as gr
-from scipy.io.wavfile import write
-
-
-def inference(audio):
-  os.makedirs("out", exist_ok=True)
-  write('test.wav', audio[0], audio[1])
-  os.system("python3 -m demucs.separate -n mdx_extra_q test.wav -o out")
-  return "./out/mdx_extra_q/test/vocals.wav","./out/mdx_extra_q/test/bass.wav",\
-"./out/mdx_extra_q/test/drums.wav","./out/mdx_extra_q/test/other.wav"
-  
-title = "Demucs"
-description = "<p style='text-align: center'>基于波形域的音乐源分离模型，使用它非常简单，只需上传您的音频文件，或者点击其中一个示例以加载它们。</p>"
-article = "<p style='text-align: center'><a href='https://arxiv.org/abs/1911.13254' target='_blank'>Music Source Separation in the Waveform Domain</a> | <a href='https://github.com/facebookresearch/demucs' target='_blank'>Github Repo</a></p>"
-
-examples=[['test.mp3']]
-gr.Interface(
-    inference, 
-    gr.inputs.Audio(type="numpy", label="上传"), 
-    [gr.outputs.Audio(type="filepath", label="人声"),gr.outputs.Audio(type="filepath", label="低音部分"),gr.outputs.Audio(type="filepath", label="打击乐器"),gr.outputs.Audio(type="filepath", label="其他")],
-    title=title,
-    description=description,
-    article=article,
-    examples=examples
-    ).launch(enable_queue=True)
+from transformers import MBart50TokenizerFast, MBartForConditionalGeneration
+
+
+class MBartTranslator:
+    """MBartTranslator class provides a simple interface for translating text using the MBart language model.
+
+    The class can translate between 50 languages and is based on the "facebook/mbart-large-50-many-to-many-mmt"
+    pre-trained MBart model. However, it is possible to use a different MBart model by specifying its name.
+
+    Attributes:
+        model (MBartForConditionalGeneration): The MBart language model.
+        tokenizer (MBart50TokenizerFast): The MBart tokenizer.
+    """
+
+    def __init__(self, model_name="facebook/mbart-large-50-many-to-many-mmt", src_lang=None, tgt_lang=None):
+
+        self.supported_languages = [
+            "ar_AR",
+            "de_DE",
+            "en_XX",
+            "es_XX",
+            "fr_XX",
+            "hi_IN",
+            "it_IT",
+            "ja_XX",
+            "ko_XX",
+            "pt_XX",
+            "ru_RU",
+            "zh_XX",
+            "af_ZA",
+            "bn_BD",
+            "bs_XX",
+            "ca_XX",
+            "cs_CZ",
+            "da_XX",
+            "el_GR",
+            "et_EE",
+            "fa_IR",
+            "fi_FI",
+            "gu_IN",
+            "he_IL",
+            "hi_XX",
+            "hr_HR",
+            "hu_HU",
+            "id_ID",
+            "is_IS",
+            "ja_XX",
+            "jv_XX",
+            "ka_GE",
+            "kk_XX",
+            "km_KH",
+            "kn_IN",
+            "ko_KR",
+            "lo_LA",
+            "lt_LT",
+            "lv_LV",
+            "mk_MK",
+            "ml_IN",
+            "mr_IN",
+            "ms_MY",
+            "ne_NP",
+            "nl_XX",
+            "no_XX",
+            "pl_XX",
+            "ro_RO",
+            "si_LK",
+            "sk_SK",
+            "sl_SI",
+            "sq_AL",
+            "sr_XX",
+            "sv_XX",
+            "sw_TZ",
+            "ta_IN",
+            "te_IN",
+            "th_TH",
+            "tl_PH",
+            "tr_TR",
+            "uk_UA",
+            "ur_PK",
+            "vi_VN",
+            "war_PH",
+            "yue_XX",
+            "zh_CN",
+            "zh_TW",
+        ]
+        print("Building translator")
+        print("Loading generator (this may take few minutes the first time as I need to download the model)")
+        self.model = MBartForConditionalGeneration.from_pretrained(model_name).to(device)
+        print("Loading tokenizer")
+        self.tokenizer = MBart50TokenizerFast.from_pretrained(model_name, src_lang=src_lang, tgt_lang=tgt_lang)
+        print("Translator is ready")
+
+    def translate(self, text: str, input_language: str, output_language: str) -> str:
+        """Translate the given text from the input language to the output language.
+
+        Args:
+            text (str): The text to translate.
+            input_language (str): The input language code (e.g. "hi_IN" for Hindi).
+            output_language (str): The output language code (e.g. "en_US" for English).
+
+        Returns:
+            str: The translated text.
+        """
+        if input_language not in self.supported_languages:
+            raise ValueError(f"Input language not supported. Supported languages: {self.supported_languages}")
+        if output_language not in self.supported_languages:
+            raise ValueError(f"Output language not supported. Supported languages: {self.supported_languages}")
+
+        self.tokenizer.src_lang = input_language
+        encoded_input = self.tokenizer(text, return_tensors="pt").to(device)
+        generated_tokens = self.model.generate(
+            **encoded_input, forced_bos_token_id=self.tokenizer.lang_code_to_id[output_language]
+        )
+        translated_text = self.tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
+
+        return translated_text[0]
+
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+translator = MBartTranslator()
+
+
+def translate(src, dst, content):
+    outputText = translator.translate(content, "zh_CN", "en_XX")
+    return outputText
+
+
+demo = gr.Interface(
+    fn=translate,
+    inputs=[
+        gr.Dropdown(
+            ["简体中文", "繁体中文", "英文", "泰文"], label="源语言", value="简体中文", show_label=True
+        ),
+        gr.Dropdown(
+            ["简体中文", "繁体中文", "英文", "泰文"], label="目标语言", value="英文", show_label=True
+        ),
+        gr.Text(label='内容')
+    ],
+    outputs=[
+        gr.Text(label='结果')
+    ]
+)
+
+
+demo.launch(enable_queue=True)