Update app.py

app.py

@@ -1,81 +1,27 @@
-import torch
+import streamlit
+import pandas as pd
+#import torch
+from transformers import pipeline
 import streamlit as st
-from transformers import BartTokenizer, BartForConditionalGeneration
-from transformers import T5Tokenizer, T5ForConditionalGeneration
 
-st.title('Text Summarization Demo')
-st.markdown('Using BART and T5 transformer model')
 
-model = st.selectbox('Select the model', ('BART', 'T5'))
+def app():
+    st.title("Text Summarization  🤓")
 
-if model == 'BART':
-    _num_beams = 4
-    _no_repeat_ngram_size = 3
-    _length_penalty = 1
-    _min_length = 12
-    _max_length = 128
-    _early_stopping = True
-else:
-    _num_beams = 4
-    _no_repeat_ngram_size = 3
-    _length_penalty = 2
-    _min_length = 30
-    _max_length = 200
-    _early_stopping = True
+    st.markdown("This is a Web application that Summarizes Text 😎")
+    text=st.text_area('Enter Text')
+    
+    summarize = st.button("Summarize") 
 
-col1, col2, col3 = st.beta_columns(3)
-_num_beams = col1.number_input("num_beams", value=_num_beams)
-_no_repeat_ngram_size = col2.number_input("no_repeat_ngram_size", value=_no_repeat_ngram_size)
-_length_penalty = col3.number_input("length_penalty", value=_length_penalty)
+    # Check to see if a file has been uploaded
+    if summarize:
+        st.success("Summarizing Text, Please wait...")
+        summarizer = pipeline("summarization", model="google/bigbird-pegasus-large-bigpatent"
+        # If it has then do the following:
+        out=summarizer(text,min_length=100, max_length=400)
+        st.json(out)
 
-col1, col2, col3 = st.beta_columns(3)
-_min_length = col1.number_input("min_length", value=_min_length)
-_max_length = col2.number_input("max_length", value=_max_length)
-_early_stopping = col3.number_input("early_stopping", value=_early_stopping)
 
-text = st.text_area('Text Input')
 
-
-def run_model(input_text):
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
-    if model == "BART":
-        bart_model = BartForConditionalGeneration.from_pretrained("facebook/bart-base")
-        bart_tokenizer = BartTokenizer.from_pretrained("facebook/bart-base")
-        input_text = str(input_text)
-        input_text = ' '.join(input_text.split())
-        input_tokenized = bart_tokenizer.encode(input_text, return_tensors='pt').to(device)
-        summary_ids = bart_model.generate(input_tokenized,
-                                          num_beams=_num_beams,
-                                          no_repeat_ngram_size=_no_repeat_ngram_size,
-                                          length_penalty=_length_penalty,
-                                          min_length=_min_length,
-                                          max_length=_max_length,
-                                          early_stopping=_early_stopping)
-
-        output = [bart_tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=False) for g insummary_ids]
-        st.write('Summary')
-        st.success(output[0])
-
-    else:
-        t5_model = T5ForConditionalGeneration.from_pretrained("t5-base")
-        t5_tokenizer = T5Tokenizer.from_pretrained("t5-base")
-        input_text = str(input_text).replace('\n', '')
-        input_text = ' '.join(input_text.split())
-        input_tokenized = t5_tokenizer.encode(input_text, return_tensors="pt").to(device)
-        summary_task = torch.tensor([[21603, 10]]).to(device)
-        input_tokenized = torch.cat([summary_task, input_tokenized], dim=-1).to(device)
-        summary_ids = t5_model.generate(input_tokenized,
-                                        num_beams=_num_beams,
-                                        no_repeat_ngram_size=_no_repeat_ngram_size,
-                                        length_penalty=_length_penalty,
-                                        min_length=_min_length,
-                                        max_length=_max_length,
-                                        early_stopping=_early_stopping)
-        output = [t5_tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=False) for g in summary_ids]
-        st.write('Summary')
-        st.success(output[0])
-
-
-if st.button('Submit'):
-    run_model(text)
+if __name__ == "__main__":
+    app()