Update extract_abs.py

extract_abs.py

@@ -19,6 +19,7 @@ import json
 import codecs
 from unidecode import unidecode
 from collections import OrderedDict
+import streamlit as st
 from typing import (
     Dict,
     List,
@@ -275,6 +276,52 @@ def search_term_extraction(search_term:Union[int,str], maxResults:int, filtering
     
     print(len(results),'abstracts classified as epidemiological.')
     return results.sort_values('EPI_PROB', ascending=False)
+    
+#Returns a Pandas dataframe                                                                                                          
+def streamlit_extraction(search_term:Union[int,str], maxResults:int, filtering:str, #for abstract search
+                           NER_pipeline:Any, entity_classes:Union[Set[str],List[str]], #for biobert extraction 
+                           extract_diseases:bool, GARD_dict:Dict[str,str], max_length:int, #for disease extraction
+                           classify_model_vars:Tuple[Any,Any,Any,Any,Any]) -> Any: #for classification
+                                                                                                                                             
+    #Format of Output
+    ordered_labels = order_labels(entity_classes)
+    if extract_diseases:
+        columns = ['PMID', 'ABSTRACT','EPI_PROB','IsEpi','IDS','DIS']+ordered_labels
+    else:
+        columns = ['PMID', 'ABSTRACT','EPI_PROB','IsEpi']+ordered_labels
+    
+    results = pd.DataFrame(columns=columns)
+    
+    ##Check to see if search term maps to anything in the GARD dictionary, if so it pulls up all synonyms for the search
+    search_term_list = autosearch(search_term, GARD_dict)
+    if len(search_term_list)>1:
+        st.write("SEARCH TERM MATCHED TO GARD DICTIONARY. SEARCHING FOR: "+ str(search_term_list))
+    else:
+        st.write("SEARCHING FOR: "+ str(search_term_list))
+    
+    #Gather title+abstracts into a dictionary {pmid:abstract}
+    pmid_abs = classify_abs.search_getAbs(search_term_list, maxResults, filtering)
+    st.write("GATHERED " +str(len(pmid_abs))+" PubMed IDs.")
+    
+    i = 0
+    my_bar = st.progress(i)
+    percent_at_step = 100/len(pmid_abs)
+    for pmid, abstract in pmid_abs.items():
+        epi_prob, isEpi = classify_abs.getTextPredictions(abstract, classify_model_vars)
+        if isEpi:
+            #Preprocessing Functions for Extraction
+            sentences = str2sents(abstract)
+            model_outputs = [NER_pipeline(sent) for sent in sentences]
+            extraction = parse_info(sentences, model_outputs, entity_classes, extract_diseases, GARD_dict, max_length)
+            if extraction:
+                extraction.update({'PMID':pmid, 'ABSTRACT':abstract, 'EPI_PROB':epi_prob, 'IsEpi':isEpi})
+                #Slow dataframe update
+                results = results.append(extraction, ignore_index=True)
+        i+=1
+        my_bar.progress(i*percent_at_step)
+    
+    st.write(len(results),'abstracts classified as epidemiological.')
+    return results.sort_values('EPI_PROB', ascending=False)
 
 #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