trans options

app.py

@@ -1,15 +1,58 @@
+import os
 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
-api = gr.load("huggingface/EleutherAI/gpt-j-6B")
+en2ko = "KoJLabs/nllb-finetuned-en2ko"
+ko2en = "KoJLabs/nllb-finetuned-ko2en"
+style = "KoJLabs/bart-speech-style-converter"
 
-def complete_with_gpt(text):
-    # Use the last 50 characters of the text as context
-    return text[:-50] + api(text[-50:])
+en2ko_model = AutoModelForSeq2SeqLM.from_pretrained(en2ko)
+ko2en_model = AutoModelForSeq2SeqLM.from_pretrained(ko2en)
+style_model = AutoModelForSeq2SeqLM.from_pretrained(style)
 
-with gr.Blocks() as demo:
-    textbox = gr.Textbox(placeholder="Type here and press enter...", lines=4)
-    btn = gr.Button("Generate")
-    
-    btn.click(complete_with_gpt, textbox, textbox)
-    
-demo.launch()
+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()