Update epi_pipeline.py

epi_pipeline.py

@@ -1,30 +1,21 @@
 from typing import List, Dict, Union, Optional, Set, Tuple
 
 # coding=utf-8
-# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
-# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ALSO See NCATS LICENSE
+##             PUBLIC DOMAIN NOTICE
+##     National Center for Advancing Translational Sciences
 
-# Written by William Kariampuzha at NIH/NCATS. Adapted from code written by Jennifer John, et al. above
+## This software/database is a "United States Government Work" under the terms of the United States Copyright Act. It was written as part of the author's official duties as United States Government employee and thus cannot be copyrighted. This software is freely available to the public for use. The National Center for Advancing Translational Science (NCATS) and the U.S. Government have not placed any restriction on its use or reproduction.  Although all reasonable efforts have been taken to ensure the accuracy and reliability of the software and data, the NCATS and the U.S. Government do not and cannot warrant the performance or results that may be obtained by using this software or data. The NCATS and the U.S.  Government disclaim all warranties, express or implied, including warranties of performance, merchantability or fitness for any particular purpose.  Please cite the authors in any work or product based on this material.
 
+# Written by William Kariampuzha @ NIH/NCATS. Adapted from code written by Jennifer John, et al. 
+# The transformer-based pipeline code has its own copyright notice under the Apache License. 
+# The code was compiled into a single python file to make adding additional features and importing into other modules easy.
 # Each section has its own import statements to facilitate clean code reuse, except for typing which applies to all.
-# the `Any` type is used in place of the specific class variable, not necessarily to mean that any object type can go there...
 
 ## Section: GATHER ABSTRACTS FROM APIs
 import requests
 import xml.etree.ElementTree as ET
+import nltk
+nltk.data.path.extend(["/home/user/app/nltk_data","./nltk_data"])
 from nltk.corpus import stopwords
 STOPWORDS = set(stopwords.words('english'))
 from nltk import tokenize as nltk_tokenize
@@ -464,6 +455,22 @@ import pandas as pd
 from more_itertools import pairwise
 
 # Subsection: Processing the abstracts into the correct data format
+
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 @dataclass
 class NERInput:
     """
@@ -546,7 +553,10 @@ class NerDataset(Dataset):
     def abstract2NERinputs(self, abstract:str) -> List[NERInput]:
         guid_index = 0
         sentences = self.str2sents(abstract)
-        ner_inputs = [NERInput(str(guid), nltk_tokenize.word_tokenize(sent), ["O" for i in range(len(nltk_tokenize.word_tokenize(sent)))]) for guid, sent in enumerate(sentences)]
+        ner_inputs = [NERInput(str(guid), 
+                      nltk_tokenize.word_tokenize(sent), 
+                      ["O" for i in range(len(nltk_tokenize.word_tokenize(sent)))]) 
+                          for guid, sent in enumerate(sentences)]
         return ner_inputs
     
     def convert_NERinputs_to_features(self, 
@@ -662,6 +672,7 @@ class NerDataset(Dataset):
 class NER_Pipeline:
     def __init__(self, name_or_path_to_model_folder:str = "ncats/EpiExtract4GARD-v2"):
         self.bert_tokenizer = BertTokenizer.from_pretrained(name_or_path_to_model_folder)
+        #no need for model variable because trainer wraps model and has more functions
         #model = AutoModelForTokenClassification.from_pretrained(name_or_path_to_model_folder)
         self.config = BertConfig.from_pretrained(name_or_path_to_model_folder)
         self.labels = {re.sub(".-","",label) for label in self.config.label2id.keys() if label != "O"}
@@ -670,6 +681,7 @@ class NER_Pipeline:
     def __str__(self):
         return "Instantiation: pipe = NER_Pipeline(name_or_path_to_model_folder)"+"\n Calling: output_dict = pipe(text)"
     
+    #Custom pipeline by WKariampuzha @NCATS (not Huggingface/Google/NVIDIA copyright)
     def __call__(self, text:str, rd_identify:Union[GARD_Search,None] = None):
         output_dict = {label:[] for label in self.labels}
         
@@ -767,6 +779,8 @@ class NER_Pipeline:
         return bi, tag
 
     
+# Unattached function -- not a method
+# move this to the NER_pipeline as a method??
 #This ensures that there is a standardized ordering of df columns while ensuring dynamics with multiple models. This is used by search_term_extraction.
 def order_labels(entity_classes:Union[Set[str],List[str]]) -> List[str]:
     ordered_labels = []
@@ -778,9 +792,72 @@ def order_labels(entity_classes:Union[Set[str],List[str]]) -> List[str]:
             ordered_labels.append(entity)
     return ordered_labels
 
+## SECTION: PIPELINES
+## This section combines all of the previous code into pipelines so that usage of these models and search functions are easy to implement in apps.
+
+# Given a search term and max results to return, this will acquire PubMed IDs and Title+Abstracts and Classify them as epidemiological.
+# results = search_term_extraction(search_term, maxResults, filering, GARD_dict, classify_model_vars)
+#Returns a Pandas dataframe   
+def search_term_classification(search_term:Union[int,str], maxResults:int, 
+                               filtering:str, rd_identify:GARD_Search, #for abstract search & filtering 
+                               epi_classify:Classify_Pipeline) -> pd.DataFrame: #for classification
+    
+    results = pd.DataFrame(columns=['PMID', 'ABSTRACT','EPI_PROB','IsEpi'])
+    
+    ##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 = rd_identify.autosearch(search_term)
+    
+    #Gather title+abstracts into a dictionary {pmid:abstract}
+    pmid_abs = search_getAbs(search_term_list, maxResults, filtering)
+    
+    for pmid, abstract in pmid_abs.items():
+        epi_prob, isEpi = epi_classify(abstract)
+        result = {'PMID':pmid, 'ABSTRACT':abstract, 'EPI_PROB':epi_prob, 'IsEpi':isEpi}
+        #Slow dataframe update
+        results = results.append(result, ignore_index=True)
+    
+    return results.sort_values('EPI_PROB', ascending=False)
+
+#Identical to search_term_classification, except it returns a JSON-compatible dictionary instead of a df
+def API_search_classification(search_term:Union[int,str], maxResults:int, 
+                              filtering:str, GARD_Search:GARD_Search, #for abstract search & filtering
+                              epi_classify:Classify_Pipeline) -> Dict[str,str]: #for classification
+
+    #Format of Output
+    results = {'entries':[]} 
+    
+    ##Check to see if search term maps to anything in the GARD dictionary, if so it pulls up all synonyms for the search
+    print('Inside `API_search_classification`. this is `search_term`:',search_term,type(search_term))
+    search_term_list = GARD_Search.autosearch(search_term)
+    
+    #Gather title+abstracts into a dictionary {pmid:abstract}
+    pmid_abs = search_getAbs(search_term_list, maxResults, filtering)
+    
+    for pmid, abstract in pmid_abs.items():
+        epi_prob, isEpi = epi_classify(abstract)
+        result = {'PMID':pmid, 'ABSTRACT':abstract, 'EPI_PROB':epi_prob, 'IsEpi':isEpi}
+        results['entries'].append(result)
+    
+    #sort 
+    results['entries'].sort(reverse=True, key=lambda x:x['EPI_PROB'])
+    
+    # float is not JSON serializable, so must convert all epi_probs to str
+    # This returns a map object, which is not JSON serializable
+    # results['entries'] = map(lambda entry:str(entry['EPI_PROB']),results['entries'])
+    # so must convert floats to str the boring and slow way
+    
+    for entry in results['entries']:
+        entry['EPI_PROB'] = str(entry['EPI_PROB'])
+        
+    return results
+
+def API_text_classification(text:str,epi_classify:Classify_Pipeline) ->  Dict[str,str]: 
+    epi_prob, isEpi = epi_classify(text)
+    return {'ABSTRACT':text, 'EPI_PROB':str(epi_prob), 'IsEpi':isEpi}
+
 # Given a search term and max results to return, this will acquire PubMed IDs and Title+Abstracts and Classify them as epidemiological.
 # It then extracts Epidemiologic Information[Disease GARD ID, Disease Name, Location, Epidemiologic Identifier, Epidemiologic Statistic] for each abstract
-# results = search_term_extraction(search_term, maxResults, filering, NER_pipeline, labels, extract_diseases, GARD_dict, max_length, classify_model_vars)
+# results = search_term_extraction(search_term, maxResults, filering, NER_pipeline, extract_diseases, GARD_Search, Classify_Pipeline)
 #Returns a Pandas dataframe                                                                                                          
 def search_term_extraction(search_term:Union[int,str], maxResults:int, filtering:str, #for abstract search
                            epi_ner:NER_Pipeline, #for biobert extraction 
@@ -882,11 +959,11 @@ def streamlit_extraction(search_term:Union[int,str], maxResults:int, filtering:s
         
         return results.sort_values('PROB_OF_EPI', ascending=False), sankey_data, disease_gardID
 
-#Identical to search_term_extraction, except it returns a JSON object instead of a df
+#Identical to search_term_extraction, except it returns a JSON-compatible dictionary instead of a df
 def API_search_extraction(search_term:Union[int,str], maxResults:int, filtering:str, #for abstract search
                    epi_ner:NER_Pipeline, #for biobert extraction 
                    GARD_Search:GARD_Search, extract_diseases:bool, #for disease extraction
-                   epi_classify:Classify_Pipeline) -> json: #for classification
+                   epi_classify:Classify_Pipeline) ->  Dict[str,str]: #for classification
                                                                                                                                              
     #Format of Output
     ordered_labels = order_labels(epi_ner.labels)
@@ -926,20 +1003,19 @@ def API_search_extraction(search_term:Union[int,str], maxResults:int, filtering:
         entry['EPI_PROB'] = str(entry['EPI_PROB'])
         
     return results
-    #return json.dumps(results)
 
-#Identical to search_term_extraction, except it returns a JSON object instead of a df
+#Identical to search_term_extraction, except it returns a JSON-compatible dictionary instead of a df
 def API_text_extraction(text:str, #Text to be extracted
                    epi_ner:NER_Pipeline, #for biobert extraction 
                    GARD_Search:GARD_Search, extract_diseases:bool, #for disease extraction
-                   ) -> json: #for classification
-                                                                                                                                             
+                   ) ->  Dict[str,str]:
+                                                                                                           
     #Format of Output
     ordered_labels = order_labels(epi_ner.labels)
     if extract_diseases:
         json_output = ['ABSTRACT','IDS','DIS']+ordered_labels
     else:
-        json_output = ['ABSTRACT',]+ordered_labels
+        json_output = ['ABSTRACT']+ordered_labels
     
     results = {'entries':[]} 
     #Do the extraction
@@ -954,7 +1030,34 @@ def API_text_extraction(text:str, #Text to be extracted
         results['entries'].append(extraction)
         
     return results
-    #return json.dumps(results)
+
+def API_text_classification_extraction(text:str, #Text to be extracted
+                           epi_ner:NER_Pipeline, #for biobert extraction 
+                           GARD_Search:GARD_Search, extract_diseases:bool, #for disease extraction
+                           epi_classify:Classify_Pipeline) ->  Dict[str,str]:
+                                                                                                                  
+    #Format of Output
+    ordered_labels = order_labels(epi_ner.labels)
+    if extract_diseases:
+        json_output = ['ABSTRACT','IsEpi','EPI_PROB','IDS','DIS']+ordered_labels
+    else:
+        json_output = ['ABSTRACT','IsEpi','EPI_PROB']+ordered_labels
+    
+    #Do the extraction
+    if extract_diseases:
+        extraction = epi_ner(text, GARD_Search)
+    else:
+        extraction = epi_ner(text)
+    
+    if extraction:
+        #Add the epidemiology probability and result
+        #Does not matter which order these are done in but doing classification after may save some time if there is no valid extraction
+        epi_prob, isEpi = epi_classify(text)
+        extraction.update({'EPI_PROB':str(epi_prob),'IsEpi':isEpi})
+        
+        #Re-order the dictionary into desired JSON output
+        output = OrderedDict([(term, extraction[term]) for term in json_output if term in extraction.keys()])
+        return output
 
 ## Section: Deprecated Functions
 import requests