Upload app.py

app.py

@@ -0,0 +1,131 @@
+import gradio as gr
+from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
+import torch
+
+
+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_XX",
+            "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)
+        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")
+        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=False)
+
+        return translated_text[0]
+
+
+
+def translate_text(source_lang, target_lang, text):
+    translator = MBartTranslator()
+    return translator.translate(text, source_lang, target_lang)
+
+translation_interface = gr.Interface(fn=translate_text, 
+                                      inputs=[gr.inputs.Dropdown(choices=["en_XX", "es_XX", "fr_XX", "zh_XX", "hi_IN"], label="Source Language"), 
+                                              gr.inputs.Dropdown(choices=["en_XX", "es_XX", "fr_XX", "zh_XX", "hi_IN"], label="Target Language"),
+                                              gr.inputs.Textbox(label="Text to translate")], 
+                                      outputs=gr.outputs.Textbox(label="Translated text"))
+translation_interface.launch()