Update extract_abs.py

extract_abs.py

@@ -279,9 +279,9 @@ def search_term_extraction(search_term:Union[int,str], maxResults:int, filtering
     
 #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
+                         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)
@@ -301,7 +301,7 @@ def streamlit_extraction(search_term:Union[int,str], maxResults:int, filtering:s
     
     #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.")
+    st.write("Gathered " +str(len(pmid_abs))+" PubMed IDs. Classifying and extracting epidemiology information...")
     
     i = 0
     my_bar = st.progress(i)
@@ -352,7 +352,7 @@ def API_extraction(search_term:Union[int,str], maxResults:int, filtering:str, #f
             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})
+                extraction.update({'PMID':pmid, 'ABSTRACT':abstract, 'EPI_PROB':epi_prob})
                 extraction = OrderedDict([(term, extraction[term]) for term in json_output])
                 results['entries'].append(extraction)