Retrieval-augmented generation for Chatbot

I developed a chatbot-like agent that utilizes Retrieval-Augmented Generation, integrating a pre-trained sequence-to-sequence model with a document database. This innovative synergy facilitates the creation of a question-answering application, which is capable of delivering contextually relevant and accurate responses by querying specific information.

Chatbot and Data Base

In the advanced domain of Natural Language Processing (NLP), large pre-trained language models have become pivotal tools, showcasing superior performance in embedding factual knowledge and achieving state-of-the-art outcomes on various downstream tasks. However, these models exhibit constraints, particularly in their ability to accurately access and manipulate knowledge. Such limitations become evident in knowledge-intensive tasks, where their efficacy can be surpassed by specialized architectures.

In this context, I constructed a prototype to explore this technology. My goal was to integrate the capabilities of pre-trained parametric memory with the expansive nature of non-parametric one. Thus, I developed a chatbot with access to a knowledge base. Utilizing retrieval-augmented generation, this system incorporates a pre-trained seq2seq model as its parametric memory and a dense vector index of documents as its non-parametric counterpart.

Dataset Overview

The chosen dataset come from the VAST 2014 Challenge The Kronos Incident. The 2014 IEEE Visual Analytics Science and Technology (VAST) Challenge researchers with a fictional scenario: the disappearance of staff members from the GASTech oil and gas company on Kronos Island. The dataset comprises:

• Geospatial Data: Geographic coordinates, maps, event locations, trajectories of vehicles or individuals.
• Temporal Data: Time series, timestamps, chronological events.
• Network Data: Information on connections between entities, such as phone calls, social media interactions, data traffic flows.
• Textual Data: Information in text form, like emails, chat logs, incident reports, etc.
• Numerical Data: Numeric values that might be related to the incident, such as counts, measurements, quantities.

A total of 845 emails, 39 Word documents, 2 Excel sheets, and 1 CSV text file were processed to construct a vector database comprising 887 documents.

Solution’s Architecture

The project is challenging as it proposes the construction of a complex system, comprised of two textual data processing modules, each with its respective language models. The final system outcome is conditioned by the performance of both modules individually and their correct integration. Several experiments were conducted, adjusting the model composition and also switching the computing architecture from CPU to GPU. Two system versions were built, with a third variation for GPU processing.

Architecture Overview

The architecture consists of two major components:

1. Searcher

Its primary goal is to find fragments of text within the corpus that can be used to answer the user-generated question. For our searcher module, we utilized all-MiniLM-L6-v2, developed by Google AI. This language model, based on transformers, has 137 billion parameters and is an enhanced version of the MiniLM-L6 model. It’s smaller and lighter than the BERT model. It’s trained on a massive dataset that includes books, articles, code, and other forms of text.

2. Generator

Focused on text construction.The generation module is built using Llama 2, developed by Meta AI. This LLM (Language Learning Model) is a pre-trained transformer model that can be used for various tasks, including translation, summarization, crafting different types of creative content, and informatively answering questions. Due to the model’s size, we used a quantized version of Llama 2 for our application. Quantization is the technique that reduces the number of bits used to represent a number or value. In the context of LLMs, this means decreasing the precision of the model’s parameters by storing the weights in lower precision data types. Specifically, the original values, encoded in 16 bits, were converted to 8 bits.

Question-Answering Solution Process

The solution addresses the Question-Answering problem in two steps. It first operates on a vectorized database of our corpus C, where each 500-character document has its associated embedding. The search is performed based on similarity between the vectorized question p and the corpus documents [d1,...,dN] using their respective indices. For similarity calculation, we use the Euclidean distance between vectors with the L2 norm (vector subtraction, where the resulting vector’s L2 norm gives us the distance between them: the smaller the distance, the greater the similarity).

After this stage, we proceed to a second step which involves generating an answer r conditioned on the retrieved information. The generation takes the question p and the 2 most similar documents (smallest distance), combines them into a single question-context (prompt) document, and passes it to a language generation model to formulate the answer r. In this final step, we are using Llama2 with 7 billion parameters, optimized for Question-Answering (chat).

A primary challenge in implementing the described solution on a personal-type CPU (4 processors and 12GB of RAM) is the size of the models. For this reason, in Version 1 of our solution, we worked with the following specifications: for embedding operations (on C and p) we used the all-MiniLM-L6-v2 model that produces a dense 384-dimensional vector space. For the r answer generation operation, we worked with a quantized version of Llama 2, 7b, 8-bit (an exploratory attempt was made with the 13 billion parameter model, thwarted by lack of RAM). Similarity operations are performed using the FAISS library for efficient similarity search among dense vectors.

Evaluation

To evaluate our model, several Question-Answering experiments were conducted, taking the answers provided as a solution to the Vast-Challenge competition as the ground truth. That is, we compared the answers formulated by analysts who participated in the competition with the answers provided by our solution. For this comparison, we used cosine distance.

Experiments: CPU

Version 1 provides unsatisfactory results not only due to the answer content but also due to system latency, caused by the prolonged generation process. For example, we have this question-answer:

We see that the model fails to identify the girl, offering an answer with a similarity score of 0.59. This difficulty is associated with the retrieval module, particularly the document vectorization model. This is because, although the documents identified as similar to the question contained relevant information, they had a marginal treatment of the question’s subject. Therefore, the failure in the generator model is strongly conditioned by the context that served as input. In addition to the above, the delay in formulating the answer was 272 seconds, which is unfeasible for a solution like ours.

Given these difficulties, we implemented a Version 2 of the solution with the following transformations. We replaced the all-MiniLM-L6-v2 model with the gte-base, which scored higher on the Overall MTEB English leaderboard. With this, we not only sought to improve the representation of the documents to favor the search but also the transformation of the question for similarity analysis. Several adjustments were made to the solution architecture to improve execution time, considering that 0.62% of the delay is in the generation module. With these transformations, we obtained the following results:

The improvement in the answer content was accompanied by a reduction in system latency, which lasted 163 seconds (+- 20). A series of experiments were conducted with questions that implied extractive and generative answers, with promising results in both cases. Therefore, we took a further step, reconfiguring the entire architecture to run in a GPU environment.

Experiments: GPU

Version 3 of the solution was readapted to run in a higher capacity environment on Google Colab, which implied significant changes. Although we kept the language models selected in Version 2, all processing performed in both modules was adjusted to run on GPU. Even the similarity calculation between the question and the database runs on GPU.

To evaluate this third version, we conducted 10 question-answer experiments. Regarding the question presented above, we obtained - as expected - the same answer obtained in Version 2 but with a latency reduction of -40 seconds. This is a positive result.

In another satisfactorily resolved question, we have these results:

More complex questions generally yielded unsatisfactory results. For example, in the following case, we formulated a question that assumes a significant level of synthesis and abstraction of the corpus documents that the model fails to resolve.

The model captures a general idea of the topic but does not articulate an appropriate answer. It fails to identify motives, circumstances, and people. It suffers from a systematic phenomenon in generative models related to ‘hallucination’.

However, the reason for the failure is not attributable to the results of the modules considered individually but to the solution in general. In effect, if we analyze the elements of this third experiment, we will find that the answer has a pertinent relationship with the available information, i.e., ‘question + context -> answer’ are acceptable, although the final result does not relate to the truth as formulated by analysts. This is because: a) the retrieval module only takes 2 documents as relevant context information due to a model limit, leaving with partial information, b) this partial information is strongly conditioned by the bias that the information has, which translates into a text with categorical and polarizing terms, and c) the decontextualization that needs to be performed for the generation instance lacks a weighting or rebalancing strategy of the corpus’s structuring semantic contents and therefore reproduces a partial and decontextualized idea of its information.

A significant number of conducted experiments show critical solution failures. They particularly point to poor performance of the retrieval module. In these cases, the model responds that it ‘does not have information to answer the question (‘I don’t know…’). That is, the modules conflict, arriving at an unacceptable result.

Finally, other experiments failed because they arrived at answers that deviate from the truth. In these cases, a failure in the retrieval module can also be observed.

Python, LLama2 and HuggingFace

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57

import box
import timeit
import yaml
import argparse
from dotenv import find_dotenv, load_dotenv
from src.utils import setup_dbqa
import torch
from langchain.embeddings import HuggingFaceEmbeddings

# Load environment variables from .env file
load_dotenv(find_dotenv())


gpu = torch.cuda.is_available()

# Import config vars
with open('config/config.yml', 'r', encoding='utf8') as ymlfile:
    cfg = box.Box(yaml.safe_load(ymlfile))

if gpu: 
    embeddings = HuggingFaceEmbeddings(model_name="thenlper/gte-base",
                                           model_kwargs={'device': 'cuda'})
else:
    embeddings = HuggingFaceEmbeddings(model_name="thenlper/gte-base",
                                       model_kwargs={'device': 'cpu'})


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('input',
                        type=str,
                        default='What were the motivations behind the kidnapping carried out by the more violent wing of the POK?',
                        help='Enter the query to pass into the LLM')
    args = parser.parse_args()

    # Setup DBQA
    start = timeit.default_timer()
    dbqa = setup_dbqa(embeddings=embeddings)
    
    response = dbqa({'query': args.input})
    end = timeit.default_timer()

    print(f'\nAnswer: {response["result"]}')
    print('='*50)

    # Process source documents
    source_docs = response['source_documents']
    for i, doc in enumerate(source_docs):
        print(f'\nSource Document {i+1}\n')
        print(f'Source Text: {doc.page_content}')
        print(f'Document Name: {doc.metadata["source"]}')
        #print(f'Page Number: {doc.metadata["page"]}\n')
        print('='* 60)

        #

    print(f"Time to retrieve response: {end - start}")

Pic by @Luke Tanis