added parameter sliders, added ONXX optimization for T5

app.py

@@ -7,6 +7,7 @@ from keybert import KeyBERT
 from sentence_transformers import SentenceTransformer
 from keyphrase_vectorizers import KeyphraseCountVectorizer
 from transformers import T5ForConditionalGeneration,T5Tokenizer
+from fastT5 import export_and_get_onnx_model, set_auth_token
 
 import nltk
 from nltk.tokenize import sent_tokenize
@@ -17,26 +18,26 @@ import streamlit as st
 import traceback
 import logging
 
-
-nltk.download('stopwords')
-nltk.download('punkt')
-
 logger = logging.getLogger(__name__)
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 HfFolder.save_token(st.secrets["hf-auth-token"])
-
+set_auth_token(st.secrets["hf-auth-token"])
 
 @st.cache(allow_output_mutation=True)
 def load_model():
     try:
+        nltk.download('stopwords')
+        nltk.download('punkt')
         # Load KeyBert Model
         tmp_model = SentenceTransformer('valurank/MiniLM-L6-Keyword-Extraction', use_auth_token=True)
         kw_extractor = KeyBERT(tmp_model)
         
         # Load T5 for Paraphrasing
-        t5_model = T5ForConditionalGeneration.from_pretrained('valurank/t5-paraphraser', use_auth_token=True)
+
+        # t5_model = T5ForConditionalGeneration.from_pretrained('valurank/t5-paraphraser', use_auth_token=True)
+        t5_model = export_and_get_onnx_model('valurank/t5-paraphraser')
         t5_tokenizer = T5Tokenizer.from_pretrained('t5-base')
         t5_model = t5_model.to(device)
         return kw_extractor, t5_model, t5_tokenizer
@@ -90,7 +91,7 @@ def t5_paraphraser(text, number_of_results=5):
     
     
 #### Extract Sentences with Keywords -> Paraphrase multiple versions -> Extract Keywords again
-def extract_paraphrased_sentences(article):
+def extract_paraphrased_sentences(article, number_of_keywords, number_of_paraphrases):
     try: 
         start1 = time.time()
         with st.spinner('Extraction Keywords from Original Document...'):
@@ -106,8 +107,8 @@ def extract_paraphrased_sentences(article):
         
             for sent in target_sentences:
                 ### T5
-                t5_paraphrased = t5_paraphraser(sent)
-                t5_keywords = [get_keybert_results_with_vectorizer(i) for i in t5_paraphrased]
+                t5_paraphrased = t5_paraphraser(sent, number_of_results = number_of_paraphrases)
+                t5_keywords = [get_keybert_results_with_vectorizer(i, number_of_results = number_of_keywords) for i in t5_paraphrased]
                 t5_keywords = [(word[0], word[1]) for s in t5_keywords for word in s]
         
                 t5_paraphrasing_keywords.extend(t5_keywords)
@@ -126,13 +127,25 @@ def extract_paraphrased_sentences(article):
         st.error('Error running Extraction Pipeline. Please contact admin')
         logger.error(traceback.format_exc())    
 
-doc = st.text_area("Enter a custom document")
 
+st.title('Exhaustive Keyword Extraction with Paraphrasing')
+with st.sidebar:
+    st.header('Overview')
+    st.markdown('This demo allows users to input text article and generate synonym-aware keywords. The pipeline includes the use of T5 Model for paraphrasing target sentences, and Sentence-transformers based Keyword Extraction')
+    
+    st.header('Parameters')
+    number_of_keywords = st.slider('Number of Keywords to extract for each target sentence', min_value=5, max_value=50, step=5, value=20)
+    number_of_paraphrases = st.slider('Number of Paraphrased versions to generate for each target sentence', min_value=1, max_value=20, step=1, value=5)
+    
+    st.header('Specifications')
+    st.markdown('To generate context aware and OOV keywords, we first run KeyBert for keyword extraction on the original article. The sentences which had Keywords are then passed through T5 for generating multiple paraphrased versions. These paraphrased sentences are then run through Keyword Extraction again to generate the final results')
+
+doc = st.text_area("Enter a custom document")
 if doc:
-    t5_keywords_df, original_keywords_df, unique_keywords_df, total_end = extract_paraphrased_sentences(doc)
+    t5_keywords_df, original_keywords_df, unique_keywords_df, total_end = extract_paraphrased_sentences(doc, number_of_keywords, number_of_paraphrases)
     
     # extract_paraphrased_article(input_list[0])
-    st.text(f'T5 PARAPHRASING RUNTIME: {total_end}\n')
+    st.text(f'PIPELINE RUNTIME: {total_end}\n')
     
     st.subheader('\nOriginal Keywords Extracted:\n\n')
     st.dataframe(original_keywords_df)