paraphrase

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 from transformers import pipeline, PegasusForConditionalGeneration, PegasusTokenizer
 from fill_in_summary import FillInSummary
+from paraphrase import PegasusParaphraser
 
 def paraphrase(text):
     return text
@@ -16,36 +17,43 @@ if select == "Summarization":
         # left_column.selectbox('Type',  ['Question Generator', 'Paraphrasing'])
         #st.selectbox('Model',  ['T5', 'GPT Neo-X'])
 
-        input = st.text_area("Input Text")
+        text_input = st.text_area("Input Text")
 
         submitted = st.form_submit_button("Generate")
 
         if submitted:
-            st.write(FillInSummary().summarize(input))
+            with st.spinner('Wait for it...'):
+                result = FillInSummary().summarize(text_input)
+            st.write(text_input)
 
 
 if select == "Fill in the blank":
     with st.form("summarization"):
-        input = st.text_area("Input Text")
+        text_input = st.text_area("Input Text")
 
         submitted = st.form_submit_button("Generate")
 
         if submitted:
-            fill = FillInSummary()
-            summarized = fill.summarize(input)
-            st.write(fill.blank_ne_out(summarized))
+            with st.spinner('Wait for it...'):
+                fill = FillInSummary()
+                result = fill.summarize(text_input)
+                result = fill.blank_ne_out(result)
+            st.write(result)
 
 
 if select == "Paraphrasing":
     with st.form("paraphrasing"):
-        st.selectbox('Model',  ['T5', 'GPT Neo-X'])
+        # st.selectbox('Model',  ['T5', 'GPT Neo-X'])
 
-        input = st.text_area("Input Text")
+        text_input = st.text_area("Input Text")
 
         submitted = st.form_submit_button("Generate")
 
         if submitted:
-            st.write(paraphrase(input))
+            with st.spinner('Wait for it...'):
+                paraphrase_model = PegasusParaphraser()
+                result = paraphrase_model.paraphrase(text_input)
+            st.write(result)
 
 
 

paraphrase.py

@@ -0,0 +1,45 @@
+from transformers import PegasusForConditionalGeneration, PegasusTokenizer
+
+class PegasusParaphraser:
+    """ Pegasus Model for Paraphrase"""
+
+    def __init__(self, num_return_sequences=3, num_beams=10, max_length=60,temperature=1.5, device="cpu"):
+        self.model_name = "tuner007/pegasus_paraphrase"
+        self.device = device
+        self.model = self.load_model()
+        self.tokenizer = PegasusTokenizer.from_pretrained(self.model_name)
+        self.num_return_sequences = num_return_sequences
+        self.num_beams = num_beams
+        self.max_length=max_length
+        self.temperature=temperature
+
+
+    def load_model(self):
+        model = PegasusForConditionalGeneration.from_pretrained(self.model_name).to(self.device)
+        return model
+
+
+    def paraphrase(self,input_text ):
+
+        batch = self.tokenizer(
+            [input_text],
+            truncation=True,
+            padding="longest",
+            max_length=self.max_length,
+            return_tensors="pt",
+        ).to(self.device)
+        translated = self.model.generate(
+            **batch,
+            max_length=self.max_length,
+            num_beams=self.num_beams,
+            num_return_sequences=self.num_return_sequences,
+            temperature=self.temperature
+        )
+        tgt_text = self.tokenizer.batch_decode(translated, skip_special_tokens=True)
+        return tgt_text
+
+
+
+
+
+