Update extract_abs.py

extract_abs.py

@@ -290,9 +290,9 @@ def streamlit_extraction(search_term:Union[int,str], maxResults:int, filtering:s
     #Format of Output
     ordered_labels = order_labels(entity_classes)
     if extract_diseases:
-        columns = ['PMID', 'ABSTRACT','EPI_PROB','IsEpi','IDS','DIS']+ordered_labels
+        columns = ['PMID', 'ABSTRACT','PROB_OF_EPI','IsEpi','IDS','DIS']+ordered_labels
     else:
-        columns = ['PMID', 'ABSTRACT','EPI_PROB','IsEpi']+ordered_labels
+        columns = ['PMID', 'ABSTRACT','PROB_OF_EPI','IsEpi']+ordered_labels
     
     results = pd.DataFrame(columns=columns)
     
@@ -321,7 +321,7 @@ def streamlit_extraction(search_term:Union[int,str], maxResults:int, filtering:s
                 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, 'PROB_OF_EPI':epi_prob, 'IsEpi':isEpi})
                     #Slow dataframe update
                     results = results.append(extraction, ignore_index=True)
                     epidemiologic+=1
@@ -330,7 +330,7 @@ def streamlit_extraction(search_term:Union[int,str], maxResults:int, filtering:s
         
         sankey_data = (gathered, relevant,epidemiologic)
         st.write(len(results),'abstracts classified as epidemiological.')
-        return results.sort_values('EPI_PROB', ascending=False), sankey_data
+        return results.sort_values('PROB_OF_EPI', 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