Update epi_pipeline.py

epi_pipeline.py

@@ -735,7 +735,7 @@ class NER_Pipeline:
                 output = list(OrderedDict.fromkeys(output)) 
                 output_dict[entity] = output
 
-        if output_dict['EPI'] and (output_dict['STAT'] or output_dict['LOC'] or output_dict['DATE']):
+        if output_dict['EPI'] and output_dict['STAT']:
             return output_dict
     
     def align_predictions(self, predictions: np.ndarray, label_ids: np.ndarray) -> Tuple[List[int], List[int]]:
@@ -827,9 +827,9 @@ def streamlit_extraction(search_term:Union[int,str], maxResults:int, filtering:s
     #Format of Output
     ordered_labels = order_labels(epi_ner.labels)
     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)