Augmenting Large Language Models via Vector Embeddings to Improve Domain-Specific Responsiveness

Summary

In this protocol, foundation large language model response quality is improved via augmentation with peer-reviewed, domain-specific scientific articles through a vector embedding mechanism. Additionally, code is provided to aid in performance comparison across large language models.

Abstract

Large language models (LLMs) have emerged as a popular resource for generating information relevant to a user query. Such models are created through a resource-intensive training process utilizing an extensive, static corpus of textual data. This static nature results in limitations for adoption in domains with rapidly changing knowledge, proprietary information, and sensitive data. In this work, methods are outlined for augmenting general-purpose LLMs, known as foundation models, with domain-specific information using an embeddings-based approach for incorporating up-to-date, peer-reviewed scientific manuscripts. This is achieved through open-source tools such as Llama-Index and publicly available models such as Llama-2 to maximize transparency, user privacy and control, and replicability. While scientific manuscripts are used as an example use case, this approach can be extended to any text data source. Additionally, methods for evaluating model performance following this enhancement are discussed. These methods enable the rapid development of LLM systems for highly specialized domains regardless of the comprehensiveness of information in the training corpus.

Introduction

Large language models (LLMs) such as OpenAI's ChatGPT or Meta AI's Llama have rapidly become a popular resource for generating text relevant to a user prompt. Originally functioning to predict the next lexical items in a sequence, these models have evolved to understand context, encode clinical information, and demonstrate high performance on a variety of tasks^1,2,3,4. Though language models predate such capabilities and their current level of popularity by decades⁵, recent advances in deep learning and computing capabilities have made pretrained, high-quality commercial LLMs broadly available to users via web-based technologies and application program interfaces (APIs)⁶. However, there are several notable limitations to consuming LLMs in this format.

Challenge 1: Static training corpus
LLMs are trained on an enormous (e.g., two trillion tokens in the case of Llama 2⁷) but static body of text data. This poses a challenge to producing accurate responses pertaining to fields undergoing rapid development or changing literature. In this static approach, LLMs would require frequent retraining to keep up with the latest data, which is neither practical nor scalable. Moreover, prompts that require responses based on information not present in the training data may prevent useful text generation or lead to hallucinations⁸. Instances of hallucinations or fact fabrication raise significant concerns about the reliability of LLMs, particularly in settings where the accuracy of information is critical⁹.

Challenge 2: Lack of domain specificity
Pretrained models are often created for general use, while users may require a model specifically optimized for performance in a particular domain. Additionally, the computational resources and data required for training a model de novo or performing significant fine-tuning are prohibitive to many users.

Challenge 3: Lack of privacy
Users pursuing applications involving sensitive data or proprietary information may be unwilling or unable to use certain LLM services as they lack information on how data may be stored or utilized.

Challenge 4: Lack of guaranteed stability
Services with proprietary LLMs may change available models or alter behavior at any time, making stability a concern for the implementation of applications relying on these services.

Retrieval-augmented generation (RAG) is a technique developed to improve LLM performance, particularly on queries related to information outside the model's training corpus^10,11. These systems augment LLMs by incorporating contextual information to be considered when generating a response to a user query. Various recent works have described applications of RAG systems and their potential advantages^12,13,14.

The goal of the method outlined in this work is to demonstrate the construction of such a system and provide a framework for researchers to rapidly experiment on domain-specific, augmented LLMs. This method is applicable for users seeking to augment an LLM with an external text-based data source. Specifically, an overarching aim of this protocol is to provide step-by-step code that is extensible to a variety of practical LLM and RAG experiments without the need for significant technical expertise in the language-modeling domain, though a working knowledge of Python is required to apply this approach without modification. To maximize user control, transparency, portability, and affordability of solutions, open-source, publicly available tools are utilized. The proposed system addresses the previously stated issues in the following ways:

Solutions 1 and 2: Static training corpus and lack of domain specificity
The provided methodology leverages a RAG approach, utilizing embeddings to supply domain-specific information not included in the original training data. At a high level, embedding models transform text or other data into a representation as a vector or single-dimension array of numbers. This technique is beneficial as it converts semantic information contained in text to a dense, numeric form. By projecting a user query into the same embedding space, various algorithms can be used to calculate the distance¹⁵, and therefore, approximate semantic similarity, between the user query and sections of text documents. Thus, creating a database of such vectors from documents broken into discrete sections can facilitate searching over a significant number of documents for text most relevant to a user query (Figure 1). This approach is extensible to any text document. While other approaches, such as online search capabilities, are beginning to be implemented to augment LLMs, this approach allows users to choose sources considered of sufficiently high quality for their use case.

Solution 2: Lack of privacy
In this implementation, a secure cloud environment was used for hosting, with no user prompts, generated responses, or other data leaving this ecosystem. All code is written in a platform-agnostic manner, however, to ensure that another cloud provider or local hardware may be substituted.

Solution 3: Lack of guaranteed stability
This approach utilizes open-source libraries and focuses on augmenting LLMs with publicly available weights, allowing a higher degree of transparency, stability, and versioning, if required.

A full schematic of our proposed system is shown in Figure 2, and detailed instructions on replicating this or a similar system is outlined in the protocol section. An additional consideration when altering model behavior through fine-tuning or augmentation is the evaluation of performance. In language-generating models, this presents a unique challenge as many traditional machine learning metrics are not applicable. Though a variety of techniques exist¹⁶, in this study, expert-written multiple-choice questions (MCQs) were used to assess accuracy and compare performance pre- and post-augmentation as well as against popular alternative LLMs.

Protocol

In the use case demonstrated in this paper, the vector store was generated using published guidelines from the Chicago Consensus Working Group¹⁷. This expert group was established to develop guidelines for the management of peritoneal cancers. The subject area was chosen as it is within the investigators' area of clinical expertise. The set of papers was accessed from online journal repositories including Cancer and the Annals of Surgical Oncology. A compact (33.4M parameters) embedding model created by the Beijing Academy for Artificial Intelligence (BAAI, https://www.baai.ac.cn/english.html), bge-small-en, was used to generate embeddings from source documents. The resulting database was then used to augment Llama 2 and Open-AI foundation models⁷. For the reader's convenience, the code is made available through GitHub (https://github.com/AnaiLab/AugmentedLLM). To ensure replicability, it is recommended to use the same versions of libraries used in the provided requirements list as well as the same version of Python. Additional details on installation or documentation regarding tools used in the following methods can be located at the official websites of the providers for Python (https://www.python.org), git (https://git-scm.com), Llama-Index (https://llamaindex.ai), and Chroma (https://trychroma.com).

1. Prerequisites: Review code and install required libraries

1. Verify that git, python, and pip are installed.
   1. In a terminal, run the following commands to verify the installation:
      git --version
      python3 --version
      pip3 --version
2. Check the code and install requirements.
   1. In a terminal, run the following commands:
      git clone https://github.com/AnaiLab/AugmentedLLM.git
      cd ./AugmentedLLM/
      pip3 install -r requirements.txt

2. Creation of a vector database with Llama-Index

1. Convert RTF file formats (only required if files are in an RTF format).
   1. Edit the file titled config.py, replacing the file paths in the following code with the location of the RTF articles to be converted and the location to which the plain text files are to be written by typing the following commands. Save the file.
      rtf_file_dir = './articles_rtf/'
      converted_file_dir = './articles_converted/'
   2. In a terminal, in the same directory, execute the code to generate plain text versions of RTF files by running the following command:
      python3 ./convert_rtf.py
2. Create and save the vector database.
   1. Edit the config.py file, replacing the value of the following variable with the file path of the folder containing documents with which the LLM is to be augmented; save the file.
      article_dir = './articles/'
   2. In a terminal, in the same directory, execute the code with the following command to create and persist the database. Verify the database is now saved in the vector_db folder.
      python3 ./build_index.py

3. Augmentation of a Llama model with vector database generated in section 2

1. Instantiate custom LLM Locally (optional)
   NOTE: Performing this step is only required if you would like to use a model other than the default Llama-2-7B. If you wish to use the default model, proceed to step 3.2.
   1. (Using a custom LLM) Specify an LLM to augment by editing run_augmented_llm.py and passing a llama-index LLM object in the constructor as the llm parameter in the following lines of code rather than None.
      augmentedLLM = AugmentedLLM(vector_store, llm=None)
2. Query augmented LLM
   1. Run the following command in the terminal:
      python3 ./run_augmented_llm.py
   2. Run user queries to get a response augmented by data in the set of manuscripts (Figure 3 and Figure 4). Press CTRL + C to exit when finished.

4. Programmatic comparison of alternative LLMs

1. Create MCQs.
   1. Edit the file questions.py, taking note of format of the examples. Add questions following a similar format. Save the file.
2. Connect to GPT-3.5, GPT-4, OpenChat, or other comparator models via API.
   1. Edit the config.py file, adding the API key for OpenAI or Huggingface if the objective is to benchmark against models from either provider. Save the file.
      huggingface_key = ''
      openai_key = ''
   2. Edit the compare_llms.py file, and choose the set of models to test against by uncommenting (deleting the '# ' characters at the begging of that line) models to compare against.
      NOTE: Some comparators require an API key as set in step 4.2.1. Additionally, edit the output_dir parameter to change where LLM output is stored if desired; otherwise, a default will be used.
      output_dir = './llm_responses/'
   3. In a terminal, execute the code with the following command. After execution, view the model responses in the folder specified in step 4.2.2 for grading or other review.
      python3 ./compare_llms.py
3. (Optional) Experiment with automated grading of MCQ responses. The example code uses the LMQL library to constrain LLM output into an expected format.
   1. Open the file automated_comparison.py and similarly to step 4.2.2, uncomment models to be included, edit the output_dir variable, or otherwise customize. Note that the model output will still be saved in a similar fashion. Save the file.
   2. Run the code from step 4.3.1 by running the following command in a terminal:
      python3 ./automated_comparison.py

Results

A set of 22 publications from the Chicago Consensus Working Group management guidelines were used to augment the base Llama-7b model¹⁷. The documents were converted into a vector index using the tool Llama-Index to generate Llama-2-7b-CCWG-Embed. Popular OpenAI models such as GPT-3.5 and GPT-4 were also augmented in a similar fashion to produce GPT-XX-CCWG-Embed models. A total of 20 multiple choice questions (MCQ) were developed to assess knowledge related to the management of a variety of peritoneal surface malignancies. The MCQs were created by a board-certified surgical oncologist to test at the level of knowledge expected of a surgical oncology fellow. Llama-2-7b-CCWG-Embed performed significantly better than the base model (see Table 1), as did augmented OpenAI models. This is considered a positive result for this method as model performance was improved with augmentation compared to baseline.

Figure 1: Schematic of generating vector database from academic papers. Please click here to view a larger version of this figure.

Figure 2: Overall system diagram for augmented LLM. Please click here to view a larger version of this figure.

Figure 3: Augmented Llama-2 model running in Linux terminal. Please click here to view a larger version of this figure.

Figure 4: Query result from augmented Llama-2 model. Please click here to view a larger version of this figure.

LLM	Score (% correct)
Llama-2-7b-chat-hf	65%
Llama-2-7b-CCWG-Embed	75%
Openchat-3.5-010618	65%
Mistral-7B-v0.119	70%
GPT-3.5	75%
GPT-3.5-CCWG-Embed	85%
GPT-4	85%
GPT-4-CCWG-Embed	90%

Table 1: Scores (percent correct) of various LLMs on a set of 20 questions pertaining to peritoneal malignancy management.

Discussion

The methods provided here aim to facilitate the research of domain-specific applications of LLMs without the need for de novo training or extensive fine-tuning. As LLMs are becoming an area of significant research interest, approaches for augmenting knowledge bases and improving the accuracy of responses will become increasingly important^18,19,20,21. As demonstrated in the provided results, the protocol outlined offers improvement in performance on domain-specific questions compared to the same LLM without augmentation and approaches the performance of more complex models such as OpenAI GPT models. As LLMs can require enormous computational resources to generate responses^22,23, augmenting a simpler model may provide an avenue to implementation in resource-constrained use cases. Furthermore, the RAG system also produced benefits when augmenting high-complexity models such as those provided by OpenAI. While the presented results are particular to peritoneal cancer management, recent studies on RAG systems with a variety of both domain and scale of data sources have been produced, indicating the broad applicability of such systems^21,24,25,26. As the quality of responses is closely related to the text data used for augmentation, however, it is critical users carefully consider which sources are optimal for their particular domain and desired application. This consideration is of particular importance in domains such as healthcare, which has been a particular focus of LLM-related studies. LLM use for diagnostics^27,28, clinical trial matching^29,30, and numerous other applications are being explored, and require validated, trustworthy text resources rather than relying on unknown training corpora.

Performance was measured by the percentage correct of MCQs written by an expert in peritoneal cancers and their clinical management. Responses were classified as correct or incorrect by human review; however, to minimize repetitive tasks for researchers, basic code to begin approaching more automated comparison of performance across LLMs is provided.

A key benefit of this protocol is the provided code to programmatically access a variety of LLMs. Previously, to assess performance on MCQs users without the requisite technical capabilities may have relied on repetitive manual testing through a web-based interface. The provided code gives basic classes to interface with several popular LLMs along with examples of how to execute the code. The goal of providing these tools is to allow more researchers to move towards automating workflows that assess responses to a prompt across a variety of LLMs, enhancing the ability to carry out comparative studies.

Furthermore, the outlined methods are designed to emphasize flexibility and extensibility. Though code can be used as-is, it is written with the intention of further tailoring to specific use cases as needed, such as the use of different LLMs and embedding models for augmentation or comparison. As the LLM landscape is rapidly evolving, this extensibility will become increasingly important to meet the varied needs of users in different domains. Additionally, the Llama-Index library provides a significant number of features not utilized in the demonstrated methods which can give users additional options when creating similar systems. These can be found in the Llama-Index official documentation.

Readers should also carefully consider evaluation methods that are optimal for their use case. While MCQs have gained popularity for LLM comparisons due to their easily quantifiable nature^31,32, caution should be taken in considering them a comprehensive performance metric for generative systems. Recent literature has implemented a diverse range of evaluations, including qualitative approaches, human review, Stanford Holistic Evaluation of Language Models (HELM)³³, and standard natural language processing metrics such as Bilingual Evaluation Understudy (BLEU), General Language Understanding Evaluation (GLUE), ROUGE, perplexity, and F1 score^{34,35,36,37,38,39}.

There are limitations to these methods as provided here. For example, while interfaces were implemented for several leading LLM providers, given the rapidly evolving nature of this field and how varied use cases may be, this implementation may not be comprehensive. The code, however, was designed with this in mind as previously discussed. Additionally, while basic code for moving toward automated grading was provided, there is room for expansion into using alternative tools for grading responses such as HELM or BLEU. Further, additional usage of tools to constrain output according to predefined grammars, such as Language model query language (LMQL), may be used to expand on this work and increase automation of comparative LLM studies. Additionally, given the wide variety of potential use cases, further studies into the effect of RAG systems on general, out-of-domain performance are required to characterize any performance trade-offs. Finally, additional methods for improving the reliability and evaluation of LLM responses must continue to be investigated.

Note that for technical reasons python 3.10 or greater is required. Shell prompts are written for bash, zsh, or other UNIX-like terminals. Commands can be executed (referred to as 'running the command' in the protocol) by simply typing the text followed by the return key. For each step in the protocol, a user may wish to examine the code in the file being executed to have a more comprehensive understanding of the mechanics behind the workflow.

Materials

List of materials used in this article

Name	Company	Catalog Number	Comments
pip3 version 22.0.2
Python version 3.10.12