Added translation controls and code

app.py

@@ -1,47 +1,46 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import gradio as gr
-import torch
-
-
-title = "Translation Chatbot"
-#description = "A State-of-the-Art Large-scale Pretrained Response generation model (DialoGPT)"
-examples = [["How are you?"]]
-
-
-tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-large")
-model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-large")
-
-
-def predict(input, history=[]):
-    # tokenize the new input sentence
-    new_user_input_ids = tokenizer.encode(
-        input + tokenizer.eos_token, return_tensors="pt"
-    )
-
-    # append the new user input tokens to the chat history
-    bot_input_ids = torch.cat([torch.LongTensor(history), new_user_input_ids], dim=-1)
-
-    # generate a response
-    history = model.generate(
-        bot_input_ids, max_length=4000, pad_token_id=tokenizer.eos_token_id
-    ).tolist()
-
-    # convert the tokens to text, and then split the responses into lines
-    response = tokenizer.decode(history[0]).split("<|endoftext|>")
-    # print('decoded_response-->>'+str(response))
-    response = [
-        (response[i], response[i + 1]) for i in range(0, len(response) - 1, 2)
-    ]  # convert to tuples of list
-    # print('response-->>'+str(response))
-    return response, history
-
-
-gr.Interface(
-    fn=predict,
-    title=title,
-    #description=description,
-    examples=examples,
-    inputs=["text", "state"],
-    outputs=["chatbot", "state"],
-    theme="earneleh/paris",
-).launch()
+import gradio as gr
+import torch
+
+from translate import Translator
+
+# https://medium.com/analytics-vidhya/make-a-translate-app-with-huggingface-transformers-ce9203f84c79
+# https://huggingface.co/docs/transformers/en/model_doc/mbart
+title = "Translation Chatbot"
+description = "A simple implementation of translating one language to another"
+examples = [["UN Chief Says There Is No Military Solution in Syria","en_XX","ja_XX"]]
+
+translator_obj = Translator()
+
+def translate_sentence(sentence):
+    return pipe(f'<-ja2zh-> {sentence}')[0]['translation_text']
+
+def predict(input, 
+    history=[],
+    original_language="en_XX", 
+    translated_language="ro_RO"):
+    response = translator_obj.translate(input, original_language, translated_language)
+    history.append((input, response))
+    return history, history
+
+if __name__ == "__main__":
+    gr.Interface(
+        fn=predict,
+        title=title,
+        description=description,
+        examples=examples,
+        inputs=[
+        gr.Textbox(),
+        "state",
+        gr.Dropdown(
+            [("English","en_XX"), ("French","fr_XX"), ("German","de_DE"), ("Japanese","ja_XX"), ("Russian","ru_RU")], value="en_XX", label="Input Language", info="Choose the language the input text is in."
+        ),
+        gr.Dropdown(
+            [("French","fr_XX"), ("German","de_DE"), ("Japanese","ja_XX"), ("Russian","ru_RU"), ("English","en_XX")], value="fr_XX", label="Language to translate to", info="Choose the language to convert the text to."
+        )
+        ],
+        outputs=[
+        gr.Chatbot(),
+        "state"
+        ],
+        theme='earneleh/paris',
+    ).launch()

requirements.txt

@@ -1,2 +1,5 @@
-transformers
-torch
+transformers
+torch
+gradio
+sentencepiece
+protobuf

translate.py

@@ -0,0 +1,47 @@
+from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
+
+
+class Translator:
+
+    '''
+        Install Requirements -
+        pip install pickle5 transformers==4.12.2 sentencepiece
+        MBart Documentation
+        https://huggingface.co/transformers/model_doc/mbart.html
+        Get the supported lang codes
+        https://huggingface.co/facebook/mbart-large-50-one-to-many-mmt
+        Class - Translator
+        Initializes MBart Seq2Seq Model and Tokenizer
+        Helper func to translate input language to desired target language
+        Supported Languages: English, Gujarati, Hindi, Bengali, Malayalam, Marathi, Tamil, Telugu
+        
+    '''
+
+    def __init__(self):
+        
+        self.model = MBartForConditionalGeneration.from_pretrained('facebook/mbart-large-50-many-to-many-mmt')
+        self.tokenizer = MBart50TokenizerFast.from_pretrained('facebook/mbart-large-50-many-to-many-mmt')
+        # self.model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50")
+        # self.tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50")
+        # , src_lang="en_XX", tgt_lang="ro_RO"
+        # https://dl-translate.readthedocs.io/en/latest/available_languages/
+        self.supported_langs = ['en_XX', 'fr_XX', 'de_DE', 'ru_RU', 'ja_XX']
+
+
+
+    def translate(self, input_text, src_lang, tgt_lang):
+
+        if src_lang not in self.supported_langs:
+            raise RuntimeError('Unsupported source language.')
+        if tgt_lang not in self.supported_langs:
+            raise RuntimeError('Unsupported target language.')
+
+        self.tokenizer.src_lang = src_lang
+        encoded_text = self.tokenizer(input_text, return_tensors='pt')
+        generated_tokens = self.model.generate(**encoded_text, forced_bos_token_id=self.tokenizer.lang_code_to_id[tgt_lang])
+        output_text_arr = self.tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
+
+        if len(output_text_arr) > 0:
+            return output_text_arr[0]
+        else:
+            raise RuntimeError('Failed to generate output. Output Text Array is empty.')