Update app.py

app.py

@@ -164,7 +164,23 @@ def run_short_extraction(article, number_of_paraphrases):
         st.error('Error running Extraction Pipeline. Please contact admin')
         logger.error(traceback.format_exc())     
         
-        
+
+
+
+def check_document_length(article, number_of_paraphrases):
+    total_tokens = len(t5_tokenizer.encode(article))
+    st.info(f'Token Counts for Encoded Document: {total_tokens}')
+
+    if total_tokens >= 512:
+        st.info('Running Extraction for Long Document')
+        t5_keywords_df, original_keywords_df, unique_keywords_df, total_end = run_long_extraction(article, number_of_paraphrases)
+    else:
+        st.info('Running Extraction for Short Document')
+        t5_keywords_df, original_keywords_df, unique_keywords_df, total_end = run_short_extraction(article, number_of_paraphrases)
+    
+    return t5_keywords_df, original_keywords_df, unique_keywords_df, total_end
+    
+    
 
 kw_extractor, t5_model, t5_tokenizer = load_base_model()