Update app.py

app.py

@@ -1,15 +1,20 @@
 import gradio as gr
-from parrot import Parrot
 import torch
-import warnings
-warnings.filterwarnings("ignore")
+from transformers import PegasusForConditionalGeneration, PegasusTokenizer
+from sentence_splitter import SentenceSplitter, split_text_into_sentences
 
-parrot = Parrot(model_tag="prithivida/parrot_paraphraser_on_T5", use_gpu=False)
+model_name = 'tuner007/pegasus_paraphrase'
+torch_device = 'cuda' if torch.cuda.is_available() else 'cpu'
+tokenizer = PegasusTokenizer.from_pretrained(model_name)
+model = PegasusForConditionalGeneration.from_pretrained(model_name).to(torch_device)
 
-def paraphrase_text(input_text):
-  para_phrases = parrot.augment(input_phrase=input_text, 
-                               max_return_phrases = 3)
-  return para_phrases[0][0], para_phrases[1][0] if len(para_phrases) > 1 else '' , para_phrases[2][0] if len(para_phrases) > 2 else ''
+def paraphrase_text(input_text, num_return_sequences = 3):
+    batch = tokenizer.prepare_seq2seq_batch([input_text], truncation=True, padding='longest', max_length=60,
+                                            return_tensors="pt").to(torch_device)
+    translated = model.generate(**batch, max_length=60, num_beams=10, num_return_sequences=num_return_sequences,
+                                temperature=1.5)
+    paraphrased_text = tokenizer.batch_decode(translated, skip_special_tokens=True)
+    return paraphrased_text[0], paraphrased_text[1], paraphrased_text[2]
 
 examples = [["Uploading a video to YouTube can help exposure for your business.", "45"], ["Niagara Falls is viewed by thousands of tourists every year.", "30"]]