test

app.py

@@ -1,57 +1,104 @@
+# 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"
 
-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)
+checkpoint = 'facebook/nllb-200-distilled-1.3B'
+model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint)
+tokenizer = AutoTokenizer.from_pretrained(checkpoint)
 
 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
-        )
-
+    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)
+    
     output = translator(text)
     translated_text = output[0]['translation_text']
-
     return translated_text
 
-title = 'KoTAN Translator & Speech-style converter'
-lang = ['English','Korean']
-
+# 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']
 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',
+    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,
     enable_queue=True,
 )