Update app.py

app.py

@@ -1,4 +1,6 @@
-from transformers import AutoModelWithLMHead, AutoTokenizer
+# from transformers import AutoModelWithLMHead, AutoTokenizer
+# Translate
+from transformers import T5ForConditionalGeneration, T5Tokenizer
 import gradio as grad
 
 # make a question
@@ -6,8 +8,12 @@ import gradio as grad
 # mdl = AutoModelWithLMHead.from_pretrained('mrm8488/t5-base-finetuned-question-generation-ap')
 
 # summarize
-text2text_tkn = AutoTokenizer.from_pretrained('deep-learning-analytics/wikihow-t5-small')
-mdl = AutoModelWithLMHead.from_pretrained('deep-learning-analytics/wikihow-t5-small')
+# text2text_tkn = AutoTokenizer.from_pretrained('deep-learning-analytics/wikihow-t5-small')
+# mdl = AutoModelWithLMHead.from_pretrained('deep-learning-analytics/wikihow-t5-small')
+
+# translate
+text2text_tkn = T5Tokenizer.from_pretrained('KETI-AIR/ke-t5-small')
+mdl = T5ForConditionalGeneration.from_pretrained('KETI-AIR/ke-t5-small')
 
 def text2text(context, answer):
     input_text = "answer: %s context: %s </s>" % (answer, context)
@@ -36,17 +42,32 @@ def text2text_summary(para):
 
     return response
 
+def text2text_translate(text):
+    inp = "translate English to Korean:: " + text
+    enc = text2text_tkn(inpu, return_tensors = 'pt')
+    tokens = mdl.generate(**enc)
+    response = text2text_tkn.batch_decode(tokens)
+
+    return response
+
+# question
 # context = grad.Textbox(lines = 10, label = 'English', placeholder = 'Context')
 # ans = grad.Textbox(lines = 1, label = 'Answer')
 # out = grad.Textbox(lines = 1, label = 'Generated Question')
 
-para = grad.Textbox(lines = 10, label = 'Paragraph', placeholder = 'Copy paragraph')
-out = grad.Textbox(lines = 1, label = 'Summary')
+# summary
+# para = grad.Textbox(lines = 10, label = 'Paragraph', placeholder = 'Copy paragraph')
+# out = grad.Textbox(lines = 1, label = 'Summary')
+
+# tranlate
+para = grad.Textbox(lines = 1, label = 'English Text', placeholder = 'Text in English')
+out = grad.Textbox(lines = 1, label = 'Korean Translation')
 
 grad.Interface(
    # text2text,
     # inputs = [context, ans],
-    text2text_summary,
+    # text2text_summary,
+    text2text_translate,
     inputs = para,
     outputs = out
 ).launch()