Update app.py

app.py

@@ -23,22 +23,19 @@ 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():
+def load_base_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 = export_and_get_onnx_model('valurank/t5-paraphraser')
+        # Load T5 for Paraphrasing
         t5_tokenizer = T5Tokenizer.from_pretrained('t5-base')
+        t5_model = T5ForConditionalGeneration.from_pretrained('valurank/t5-paraphraser', use_auth_token=True)
         t5_model = t5_model.to(device)
         return kw_extractor, t5_model, t5_tokenizer
     except Exception:
@@ -49,11 +46,10 @@ def load_model():
 kw_extractor, t5_model, t5_tokenizer = load_model()
 
 
-@st.cache()
 def get_keybert_results_with_vectorizer(text, number_of_results=20):
-
     try:
         keywords = kw_extractor.extract_keywords(text, vectorizer=KeyphraseCountVectorizer(), stop_words=None, top_n=number_of_results)
+        keywords = [i for i in keywords if i[1] >= 0.25]
         return keywords
     except Exception:
         st.error('Error running Keybert. Please contact admin')
@@ -61,24 +57,22 @@ def get_keybert_results_with_vectorizer(text, number_of_results=20):
 
 
 
-@st.cache()
 def t5_paraphraser(text, number_of_results=5):
     try:
-        text =  "paraphrase: " + text + " </s>"
-        max_len = 2048
+        text =  "paraphrase: " + text
         encoding = t5_tokenizer.encode_plus(text, pad_to_max_length=True, return_tensors="pt")
         input_ids, attention_masks = encoding["input_ids"].to(device), encoding["attention_mask"].to(device)
-    
+
         beam_outputs = t5_model.generate(
             input_ids=input_ids, attention_mask=attention_masks,
             do_sample=True,
-            max_length=2048,
+            max_length=1024,
             top_k=50,
             top_p=0.95,
             early_stopping=True,
             num_return_sequences=number_of_results
         )
-        
+
         final_outputs =[]
         for beam_output in beam_outputs:
             sent = t5_tokenizer.decode(beam_output, skip_special_tokens=True, clean_up_tokenization_spaces=True)
@@ -90,69 +84,119 @@ def t5_paraphraser(text, number_of_results=5):
     
     
     
-#### Extract Sentences with Keywords -> Paraphrase multiple versions -> Extract Keywords again
-def extract_paraphrased_sentences(article, number_of_keywords, number_of_paraphrases):
-    try: 
+def run_long_extraction(article, number_of_paraphrases):
+    try:
         start1 = time.time()
         with st.spinner('Extraction Keywords from Original Document...'):
-            original_keywords = get_keybert_results_with_vectorizer(article)
-        
+            original_keywords = get_keybert_results_with_vectorizer(article, number_of_results=30)
             article_sentences = sent_tokenize(article)
-            target_sentences = [sent for sent in article_sentences if any(kw[0] in sent for kw in original_keywords)]        
+            target_sentences = [sent for sent in article_sentences if any(kw[0] in sent for kw in original_keywords)]
+
         st.success('Keyword Extraction from Original Document finished in {}'.format(time.time() - start1))
-   
+        st.info(f'Total Sentences in Article : {len(article_sentences)}')
+        st.info(f'Total Target Sentences Selected : {len(target_sentences)}')
+
         start2 = time.time()
         with st.spinner('Extracting Keywords from Paraphrased Target Sentences...'):
             t5_paraphrasing_keywords = []
-        
             for sent in target_sentences:
                 ### T5
                 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 = [get_keybert_results_with_vectorizer(i) 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)
         st.success('Keyword Extraction from Paraphrased Target Sentences finished in {}'.format(time.time() - start2))
-    
+
         original_keywords_df = pd.DataFrame(original_keywords, columns=['Keyword', 'Score'])
-        
+
         t5_keywords_df = pd.DataFrame(t5_paraphrasing_keywords, columns=['Keyword', 'Score']).sort_values(by='Score', ascending=False).drop_duplicates(subset=['Keyword'], keep='first').reset_index(drop=True)
-    
+
         unique_keywords_df = pd.DataFrame([i for i in t5_paraphrasing_keywords if not original_keywords_df['Keyword'].str.contains(i[0], regex=False, case=False).any()], columns=['Keyword', 'Score']).sort_values(by='Score', ascending=False).drop_duplicates(subset=['Keyword'], keep='first').reset_index(drop=True)
-    
+
         total_end = time.time()-start1
-    
+
+        return t5_keywords_df, original_keywords_df, unique_keywords_df, total_end
+    except Exception:
+        st.error('Error running Extraction Pipeline. Please contact admin')
+        logger.error(traceback.format_exc())     
+          
+          
+          
+def run_short_extraction(article, number_of_paraphrases):
+    try:
+        start1 = time.time()
+        original_keywords = get_keybert_results_with_vectorizer(article)
+        article_sentences = sent_tokenize(article)
+        st.info(f'Total Sentences in Article : {len(article_sentences)}')
+
+        target_sentences = []
+        tmp = []
+        token_count = 0
+        for i in article_sentences:
+            enc = t5_tokenizer.encode(i)
+            if token_count + len(enc) <= 96:
+                tmp.append(i)
+                token_count += len(enc)
+            else:
+                target_sentences.append(' '.join(tmp))
+                token_count = len(enc)
+                tmp = [i]          
+     
+        start2 = time.time()
+        with st.spinner('Extracting Keywords from Paraphrased Sentences Groups...'):
+            t5_paraphrasing_keywords = []
+            for sent in target_sentences:
+                ### T5
+                t5_paraphrased = t5_paraphraser(sent, number_of_results = number_of_paraphrases)
+                t5_keywords = [get_keybert_results_with_vectorizer(i) 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)
+        st.success('Keyword Extraction from Paraphrased Grouped Sentences finished in {}'.format(time.time() - start2))
+
+        original_keywords_df = pd.DataFrame(original_keywords, columns=['Keyword', 'Score'])
+
+        t5_keywords_df = pd.DataFrame(t5_paraphrasing_keywords, columns=['Keyword', 'Score']).sort_values(by='Score', ascending=False).drop_duplicates(subset=['Keyword'], keep='first').reset_index(drop=True)
+
+        unique_keywords_df = pd.DataFrame([i for i in t5_paraphrasing_keywords if not original_keywords_df['Keyword'].str.contains(i[0], regex=False, case=False).any()], columns=['Keyword', 'Score']).sort_values(by='Score', ascending=False).drop_duplicates(subset=['Keyword'], keep='first').reset_index(drop=True)
+
+        total_end = time.time()-start1
+
         return t5_keywords_df, original_keywords_df, unique_keywords_df, total_end
     except Exception:
         st.error('Error running Extraction Pipeline. Please contact admin')
-        logger.error(traceback.format_exc())    
+        logger.error(traceback.format_exc())     
+        
+        
+
+kw_extractor, t5_model, t5_tokenizer = load_base_model()
 
 
 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_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')
+    # st.markdown('To generate context aware and OOV keywords for long, 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, number_of_keywords, number_of_paraphrases)
-    
+    t5_keywords_df, original_keywords_df, unique_keywords_df, total_end = check_document_length(doc, number_of_paraphrases)
+
     # extract_paraphrased_article(input_list[0])
     st.text(f'PIPELINE RUNTIME: {total_end}\n')
-    
+
     st.subheader('\nOriginal Keywords Extracted:\n\n')
     st.dataframe(original_keywords_df)
-    
+
     st.subheader('\nT5 Unique New Keywords Extracted:\n\n')
     st.dataframe(unique_keywords_df)
-       
+
     st.subheader('\nT5 Keywords Extracted:\n\n')
     st.dataframe(t5_keywords_df)
-