Update app.py

app.py

@@ -1,4 +1,4 @@
-from transformers import T5Tokenizer, T5Model, T5ForConditionalGeneration, BertTokenizer, EncoderDecoderModel
+from transformers import T5Tokenizer, T5Model, T5ForConditionalGeneration, BertTokenizer, EncoderDecoderModel, AutoTokenizer, AutoModelForSeq2SeqLM
 
 import gradio as gr
 from gradio.mix import Parallel
@@ -11,6 +11,10 @@ tokenizer_bert.bos_token = tokenizer_bert.cls_token
 tokenizer_bert.eos_token = tokenizer_bert.sep_token
 model_bert = EncoderDecoderModel.from_pretrained("cahya/bert2bert-indonesian-summarization")
 
+t5_para_tokenizer = AutoTokenizer.from_pretrained("Wikidepia/IndoT5-base-paraphrase")  
+t5_para_model = AutoModelForSeq2SeqLM.from_pretrained("Wikidepia/IndoT5-base-paraphrase")
+
+
 def summ_t5(text):
     input_ids = tokenizer_t5.encode(text, return_tensors='pt')
     summary_ids = model_t5.generate(input_ids,
@@ -25,8 +29,21 @@ def summ_t5(text):
     return summary_text
     
 def summ_bert(text):
-    input_ids = tokenizer_bert.encode(text, return_tensors='pt')
-    summary_ids = model_bert.generate(input_ids,
+    encoding = tokenizer(text, padding='longest', return_tensors="pt")
+    outputs = model.generate(
+                input_ids=encoding["input_ids"], attention_mask=encoding["attention_mask"],
+                max_length=512,
+                do_sample=True,
+                top_k=200,
+                top_p=0.95,
+                early_stopping=True,
+                num_return_sequences=5)
+    summary_text = tokenizer_bert.decode(summary_ids[0], skip_special_tokens=True)
+    return summary_text
+    
+ def para_t5(text):
+    input_ids = t5_para_tokenizer.encode(text, return_tensors='pt')
+    outputs = t5_para_model .generate(input_ids,
                 min_length=20,
                 max_length=100,
                 num_beams=10,
@@ -39,8 +56,12 @@ def summ_bert(text):
                 temperature = 0.8,
                 top_k = 50,
                 top_p = 0.95)
-    summary_text = tokenizer_bert.decode(summary_ids[0], skip_special_tokens=True)
-    return summary_text
+    return [
+        t5_para_tokenizer.decode(
+            output, skip_special_tokens=True, clean_up_tokenization_spaces=True
+        )
+        for output in outputs
+    ]
 
 t5_demo = gr.Interface(
     fn=summ_t5,
@@ -52,8 +73,13 @@ bert_demo = gr.Interface(
     inputs="text",
     outputs=gr.Textbox(lines=10, label="Bert2Bert Base Output")
     )
+para_demo = gr.Interface(
+    fn=para_t5,
+    inputs="text",
+    outputs=gr.Textbox(lines=10, label="T5 Paraphrase Output")
+    )
 
 if __name__ == "__main__":
-    Parallel(t5_demo, bert_demo,
+    Parallel(t5_demo, bert_demo, para_demo,
     inputs=gr.Textbox(lines=10, label="Input Text", placeholder="Enter article here..."),
     title="Summary of Summarizer - Indonesia").launch()