translation

app.py

@@ -1,104 +1,57 @@
-# import os
-# import gradio as gr
-# from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
-
-# en2ko = "KoJLabs/nllb-finetuned-en2ko"
-# ko2en = "KoJLabs/nllb-finetuned-ko2en"
-# style = "KoJLabs/bart-speech-style-converter"
-
-# en2ko_model = AutoModelForSeq2SeqLM.from_pretrained(en2ko)
-# ko2en_model = AutoModelForSeq2SeqLM.from_pretrained(ko2en)
-# style_model = AutoModelForSeq2SeqLM.from_pretrained(style)
-
-# en2ko_tokenizer = AutoTokenizer.from_pretrained(en2ko)
-# ko2en_tokenizer = AutoTokenizer.from_pretrained(ko2en)
-# style_tokenizer = AutoTokenizer.from_pretrained(style)
-
-# def translate(source, target, text):
-#     formats = {"English":"eng_Latn", "Korean":"kor_Hang"}
-#     src = formats[source]
-#     tgt = formats[target]
-
-#     if src == "eng_Latn":
-#         translator = pipeline(
-#             'translation',
-#             model=en2ko_model,
-#             tokenizer=ko2en_tokenizer,
-#             src_lang=src, 
-#             tgt_lang=tgt, 
-#         )
-
-#     if src == "kor_Hang":
-#         translator = pipeline(
-#             'translation',
-#             model=ko2en_model,
-#             tokenizer=en2ko_tokenizer,
-#             src_lang=src, 
-#             tgt_lang=tgt
-#         )
-
-#     output = translator(text)
-#     translated_text = output[0]['translation_text']
-
-#     return translated_text
-
-# title = 'KoTAN Translator & Speech-style converter'
-# lang = ['English','Korean']
-
-# translator_app = gr.Interface(
-#     fn=translate,
-#     inputs=[gr.inputs.Dropdown(choices=lang, label='Source Language'), gr.inputs.Dropdown(choices=lang, label='Target Language'), gr.inputs.Textbox(lines=5, label='Text to Translate')],
-#     outputs=[gr.outputs.Textbox(label='Translated Text')],
-#     title=title,
-#     description = 'KoTAN: Korean Translation and Augmentation with fine-tuned NLLB. If you want to download as pip package, please visit our github. (https://github.com/KoJLabs/KoTAN)',
-#     article='Jisu, Kim. Juhwan, Lee',
-#     enable_queue=True,
-# )
-
-# translator_app.launch()
-
 import os
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
+en2ko = "KoJLabs/nllb-finetuned-en2ko"
+ko2en = "KoJLabs/nllb-finetuned-ko2en"
+style = "KoJLabs/bart-speech-style-converter"
 
-checkpoint = 'facebook/nllb-200-distilled-1.3B'
-model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint)
-tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+en2ko_model = AutoModelForSeq2SeqLM.from_pretrained(en2ko)
+ko2en_model = AutoModelForSeq2SeqLM.from_pretrained(ko2en)
+style_model = AutoModelForSeq2SeqLM.from_pretrained(style)
+
+en2ko_tokenizer = AutoTokenizer.from_pretrained(en2ko)
+ko2en_tokenizer = AutoTokenizer.from_pretrained(ko2en)
+style_tokenizer = AutoTokenizer.from_pretrained(style)
 
 def translate(source, target, text):
-    formats = {'English':'eng_Latn', 'Asante': 'aka_Latn', 'Akuapem': 'twi_Latn', 'Ewe':'ewe_Latn', 'Hausa': 'hau_Latn'}
-    source_fmt = formats[source]
-    target_fmt = formats[target]
-    translator = pipeline('translation', 
-                      model=model, 
-                      tokenizer=tokenizer, 
-                      src_lang=source_fmt, 
-                      tgt_lang=target_fmt, 
-                      max_length = 400)
-    
+    formats = {"English":"eng_Latn", "Korean":"kor_Hang"}
+    src = formats[source]
+    tgt = formats[target]
+
+    if src == "eng_Latn":
+        translator = pipeline(
+            'translation',
+            model=en2ko_model,
+            tokenizer=ko2en_tokenizer,
+            src_lang=src, 
+            tgt_lang=tgt, 
+        )
+
+    if src == "kor_Hang":
+        translator = pipeline(
+            'translation',
+            model=ko2en_model,
+            tokenizer=en2ko_tokenizer,
+            src_lang=src, 
+            tgt_lang=tgt
+        )
+
     output = translator(text)
     translated_text = output[0]['translation_text']
+
     return translated_text
 
-# hf_token = os.getenv('HF_TOKEN')
-title = 'Ananse AI | Ghanaian Language Translator'
-# hf_writer = gr.HuggingFaceDatasetSaver(hf_token, 'crowdsourced-GLT')
-lang = ['Akuapem','Asante', 'English', 'Ewe','Hausa']
+title = 'KoTAN Translator & Speech-style converter'
+lang = ['English','Korean']
+
 translator_app = gr.Interface(
     fn=translate,
     inputs=[gr.inputs.Dropdown(choices=lang, label='Source Language'), gr.inputs.Dropdown(choices=lang, label='Target Language'), gr.inputs.Textbox(lines=5, label='Text to Translate')],
     outputs=[gr.outputs.Textbox(label='Translated Text')],
     title=title,
-    description = 'The Ghanaian Languages considered now are the Asante, Akuapem, Ewe, and Hausa. Kindly use the Flag button to rate (👍🏼, 👎🏼) your translation to help us improve. Thanks',
-    article='Ananse AI | hnmensah',
-    examples = [['English','Asante','Kwame went to Kaneshie to buy tomates.'],
-                ['English','Ewe','The event should be hosted at the Accra Mall.'],
-                ['English','Akuapem','The trader is suffering from Malaria so she did not go to work.'],
-                ['English','Hausa','The last person to get to the class will be sacked.']],
-    #allow_flagging='manual',
-    #flagging_options=['👍🏼','👎🏼'],
-    #flagging_callback=hf_writer,
+    description = 'KoTAN: Korean Translation and Augmentation with fine-tuned NLLB. If you want to download as pip package, please visit our github. (https://github.com/KoJLabs/KoTAN)',
+    article='Jisu, Kim. Juhwan, Lee',
     enable_queue=True,
 )
 

requirements.txt

@@ -0,0 +1,3 @@
+transformers
+torch
+sentencepiece