Update app.py

app.py

@@ -1,9 +1,9 @@
-from transformers import MT5ForConditionalGeneration, AutoTokenizer
+from transformers import MT5ForConditionalGeneration, AutoTokenizer, Text2TextGenerationPipeline
 import gradio as gr
 
 trans_mdl = MT5ForConditionalGeneration.from_pretrained("K024/mt5-zh-ja-en-trimmed")
 trans_tokenizer = AutoTokenizer.from_pretrained("K024/mt5-zh-ja-en-trimmed")
-
+trans_pipe = Text2TextGenerationPipeline(model=trans_mdl, tokenizer=trans_tokenizer)
 
 def translation_job(data):
     job = data[0]
@@ -16,9 +16,8 @@ def translation_job(data):
     job_map = dict(zip(job_key, job_value))
 
     input = job_map[job] + text
-    enc = trans_tokenizer(input, return_tensor="pt")
-    tokens = trans_mdl.generate(**enc)
-    response = trans_tokenizer.batch_decode(tokens)
+    print(input)
+    response = trans_tokenizer.batch_decode(input, max_length=100, num_beams=4)
     return response