Document-level Clinical Entity and Relation Extraction via Knowledge Base-Guided Generation

UMLS Concept Mapping
We first map UMLS concepts from clinical text using MetaMap (Aronson, 2001). Given clinical text $X=\{x_{1},x_{2},\ldots,x_{n}\}$ where $x_{i}$ represents the $i$th clinical text, we map $C_{i}=\{c_{i1},c_{i2},\ldots,c_{in}\}$, where $C_{i}$ denote the set of concepts identified by MetaMap from $x_{i}$. $n$ denotes the number of concepts identified from $x_{i}$. These concepts are extracted leveraging MetaMap’s lexical parsing, syntactic analysis, semantic mapping, and concept mapping techniques.
Next, we filter the mapped concepts to include only those categorized as ’organic chemical’, ’antibiotic’, or ’pharmacologic substance’ within the UMLS concept hierarchy as these groups contains the medications. For this work, we only target and filter medication-related concepts for augmentation and for further analysis. Let’s denote the filtered set of concepts for the $i$th input clinical text $x_{i}$ as $C_{{filtered}}$, such that $C_{filtered,i}=\{c_{ij}\in C_{i}|c_{ij}\in\text{filtered groups}\}$.

Prompt Strategy and Large Language Model Implementation
Next, we prompt the GPT model to extract entity-relation pairs from the text, leveraging the mapped UMLS concepts from MetaMap, and employing a few-shot prompt strategy. Let $P_{i}$ represent the prompt generated for each input text $x_{i}$, incorporating the relevant UMLS concepts $C_{filtered,i}$. The final prompt $P_{i}$ is constructed as the concatenation of the initial prompt and the set of UMLS concepts, i.e., $P_{i}=\text{Concat}(P,C_{filtered,i})$. We use OpenAI’s GPT-4-32k (Version 0613) and GPT-3.5-turbo (Version 0301) via HIPPA-compliant Microsoft Azure’s OpenAI REST API¹¹1https://learn.microsoft.com/en-us/azure/ai-services/openai/quickstart?tabs=command-line%2Cpython-new&pivots=programming-language-pythonendpoint. A sample prompt and hyperparameters used by the models for this task are available in Figure 2 and A.2 respectively. As our goal for the project was not to explore different prompting strategies, we tested a few prompts and selected the prompt that generated more specific result. We used the same format for all relation pairs replacing only the entity type for every run.

Retrieval Augmented Generation
We also explored another approach-RAG to leverage UMLS in a language model, which is a more conventional method involving the use of external data. RAG was chosen for its potential to enhance the generation process by incorporating domain-specific knowledge from sources like the UMLS knowledge base. Appendix A.4 includes details on our RAG implementation.

n2c2 Dataset
We used a curated National NLP Clinical Challenges (n2c2) dataset (Henry et al., 2019) consisting of 303 deidentified discharge summaries obtained from the MIMIC-III (Medical Information Mart for Intensive Care-III) critical care database (Table 1A) (Johnson et al., 2016). The data also contained annotations of medication-related entities and their relationship to other entities. Annotations conducted by 3 subject matter experts served as a gold standard to evaluate model performance.

ADE Dataset
The Adverse Drug Events (ADE) dataset annotated by 5 individuals consists of MEDLINE²²2https://www.nlm.nih.gov/medline/medline_home.html case reports with information on medications, dosages and adverse effects associated with the medications (Gurulingappa et al., 2012) (Table 1B). It also contains relationships between medications, dosages, and adverse effects. For our experiments, we used the second version of the dataset downloaded from Huggingface³³3https://huggingface.co/datasets/ade_corpus_v2.

Results on n2c2 and ADE Dataset
Our results suggests that integrating prior knowledge from UMLS in the prompts have significant performance improvement as demonstrated by the higher average F-1 scores across both n2c2 and ADE datasets (Table 4). The reported results are average across all entity-entity relation pairs across models and for 2 datasets. GPT-4-32k model with UMLS show 4% improvement of F-1 score on the n2c2 dataset, and 12% improvement on the ADE dataset from the F-1 score of GPT-4-32k model without knowledge integration. For every entity-entity relation pairs, there was a performance improvement by a few percentages for both models and across both datasets. Additional detailed results for each entity-entity relation pairs can be found in Appendices A.6 through A.11.
Upon a closer look at the results, we identified that prompts with UMLS resulted in additional concepts verifying that UMLS is able to identify medications from the text that GPT may not identify independently(Appendix A.5).

Comparison with Retrieval Augmented Generation
RAG model and GPT-3.5-turbo had low F-1 score and it improved with UMLS for both models, but it did not have higher score compared to the GPT-4-32k+UMLS.
We observed performance variations across entity-entity relation pairs with retrieval augmented generation (A.8, A.11). While some entity pairs showed performance enhancements, others did not show significant improvements. This discrepancy might arise from the limitations of RAG, particularly its inability to utilize entire UMLS thesaurarus in the generation process. Since UMLS data is partitioned into chunks for indexing and embedding and embedding models can only take 8192 tokens per index, some concepts may not be in the top-k extracted documents used for generation, potentially limiting the scope of augmentation and its impact on final relation pairs. Further experiments are required to confirm this hypothesis.