finishgit add app.py

app.py

@@ -14,7 +14,7 @@ en2ko_tokenizer = AutoTokenizer.from_pretrained(en2ko)
 ko2en_tokenizer = AutoTokenizer.from_pretrained(ko2en)
 style_tokenizer = AutoTokenizer.from_pretrained(style)
 
-def translate(source, target, text):
+def translation(source, target, text):
     formats = {"English":"eng_Latn", "Korean":"kor_Hang"}
     src = formats[source]
     tgt = formats[target]
@@ -23,7 +23,7 @@ def translate(source, target, text):
         translator = pipeline(
             'translation',
             model=en2ko_model,
-            tokenizer=ko2en_tokenizer,
+            tokenizer=en2ko_tokenizer,
             src_lang=src, 
             tgt_lang=tgt, 
         )
@@ -32,7 +32,7 @@ def translate(source, target, text):
         translator = pipeline(
             'translation',
             model=ko2en_model,
-            tokenizer=en2ko_tokenizer,
+            tokenizer=ko2en_tokenizer,
             src_lang=src, 
             tgt_lang=tgt
         )
@@ -42,6 +42,32 @@ def translate(source, target, text):
 
     return translated_text
 
+def augmentation(text):
+    ko2en_translator = pipeline(
+            'translation',
+            model=ko2en_model,
+            tokenizer=ko2en_tokenizer,
+            src_lang="kor_Hang", 
+            tgt_lang="eng_Latn"
+        )
+    
+    output = ko2en_translator(text)
+    ko2en_text = output[0]['translation_text']
+    
+    en2ko_translator = pipeline(
+            'translation',
+            model=en2ko_model,
+            tokenizer=en2ko_tokenizer,
+            src_lang="eng_Latn", 
+            tgt_lang="kor_Hang", 
+        )
+    
+    output = en2ko_translator(ko2en_text)
+    en2ko_text = output[0]['translation_text']
+
+    return en2ko_text
+    
+
 def conversion(source, text):
     formats = {
                 "formal":"문어체",
@@ -77,19 +103,28 @@ def conversion(source, text):
 
     return generated_text
 
-title = 'KoTAN Translator & Speech-style converter. If you want to download as pip package, please visit our github. (https://github.com/KoJLabs/KoTAN)'
+
+
 lang = ['English','Korean']
 style = ['formal', 'informal', 'android', 'azae', 'chat', 'choding', 'emoticon', 'enfp', \
          'gentle', 'halbae', 'halmae', 'joongding', 'king', 'naruto', 'seonbi', 'sosim', 'translator']
 
     
-translator_app = gr.Interface(
-    fn=translate,
+translation_app = gr.Interface(
+    fn=translation,
     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,
     enable_queue=True,
 )
+
+augmentation_app = gr.Interface(
+    fn=augmentation,
+    inputs=[gr.inputs.Textbox(lines=5, label='Korean Text to Augmentation')],
+    outputs=[gr.outputs.Textbox(label='Augmented Text')],
+    title=title,
+    enable_queue=True,
+)
     
 conversion_app = gr.Interface(
     fn=conversion,
@@ -99,6 +134,11 @@ conversion_app = gr.Interface(
     enable_queue=True,
 )
 
-demo = gr.TabbedInterface([translator_app, conversion_app], ["Translator", "Speech Style converter"])
+demo = gr.TabbedInterface([translation_app, augmentation_app, conversion_app], \
+                          ["Translation", "Augmentation", "Speech Style conversion"],\
+                          title = 'KoTAN Translator & Speech-style converter. \n \
+                            🔥If you want to download as pip package, \
+                            please visit our github. (https://github.com/KoJLabs/KoTAN) 🔥'
+                            )
 
 demo.launch()