Change generation interface

app.py

@@ -1,33 +1,18 @@
 import gradio as gr
 import torch
-from transformers import T5Tokenizer, AutoModelForCausalLM
+from transformers import T5Tokenizer, AutoModelForCausalLM, pipeline
 from utils import translate_from_jp_to_en
 
 tokenizer = T5Tokenizer.from_pretrained("rinna/japanese-gpt-1b")
 model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt-1b")
 
+generator = pipeline("text-generation", tokenizer=tokenizer, model=model)
 
 def generate(text, max_length=512):
+    out = generator(text, do_sample=True, max_length=max_length, num_return_sequences=1)
+    text = out[0]['generated_text']
 
-    token_ids = tokenizer.encode(
-        text, add_special_tokens=False, return_tensors="pt")
-
-    with torch.no_grad():
-        output_ids = model.generate(
-            token_ids,
-            max_length=max_length,
-            do_sample=True,
-            top_k=500,
-            top_p=0.95,
-            #pad_token_id=tokenizer.pad_token_id,
-            #bos_token_id=tokenizer.bos_token_id,
-            #eos_token_id=tokenizer.eos_token_id,
-            #bad_word_ids=[[tokenizer.unk_token_id]]
-            early_stopping=False,
-        )
-
-    output = tokenizer.decode(output_ids.tolist()[0], skip_specual_tokens=True)
-    return output, translate_from_jp_to_en(output)
+    return text, translate_from_jp_to_en(text)
 
 
 title = "JP GPT Demo"