Update extract_abs.py

extract_abs.py

@@ -301,12 +301,13 @@ def streamlit_extraction(search_term:Union[int,str], maxResults:int, filtering:s
         st.write("SEARCHING FOR: "+ str(search_term_list))
     
     #Gather title+abstracts into a dictionary {pmid:abstract}
-    pmid_abs = classify_abs.streamlit_getAbs(search_term_list, maxResults, filtering)
+    pmid_abs, sankey_initial = classify_abs.streamlit_getAbs(search_term_list, maxResults, filtering)
     if len(pmid_abs)==0:
         st.error('No results were gathered. Enter a new search term.')
     else:
         st.write("Gathered " +str(len(pmid_abs))+" PubMed IDs. Classifying and extracting epidemiology information...")
-        
+        gathered, relevant = sankey_initial
+        epidemiologic = 0
         i = 0
         my_bar = st.progress(i)
         percent_at_step = 100/len(pmid_abs)
@@ -321,11 +322,13 @@ def streamlit_extraction(search_term:Union[int,str], maxResults:int, filtering:s
                     extraction.update({'PMID':pmid, 'ABSTRACT':abstract, 'EPI_PROB':epi_prob, 'IsEpi':isEpi})
                     #Slow dataframe update
                     results = results.append(extraction, ignore_index=True)
+                    epidemiologic+=1
             i+=1
             my_bar.progress(round(i*percent_at_step/100,1))
         
+        sankey_data = (gathered, relevant,epidemiologic)
         st.write(len(results),'abstracts classified as epidemiological.')
-        return results.sort_values('EPI_PROB', ascending=False)
+        return results.sort_values('EPI_PROB', ascending=False), sankey_data
 
 #Identical to search_term_extraction, except it returns a JSON object instead of a df
 def API_extraction(search_term:Union[int,str], maxResults:int, filtering:str, #for abstract search