LLM Knowledge-Injection of Document Chunks for Table Retrieval

Overview • Setup • Workflow

Overview

It is highly recommended to develop in an x86-64 Linux-based environment for ease of use. This project allows for knowledge-injection of document chunks into various objects. It is currently set up to enrich serialized tables and natural language queries for table retrieval, but can be easily modified. Both embedding and text serach are used in conjunction for the retrieval algorithm.

Setup

The vast majority of needed packages can be installed through pip3 using requirements.txt in a virtual environment. It is likely that the dependency resolution will get angry on newer Python3 versions, so it is recommended to use a tool like pyenv to downgrade to the more stable 3.12.3 before creating the virtual environment.

# create virtual environment
python3 -m venv .venv
# activate virtual environment
source ./.venv/bin/activate (if you're using bash)

# install packages
pip3 install -r requirements.txt

You will additionally need to install xformers in a separate step after installing other requirements. There may be further packages that need to be installed when prompted, such as NVIDIA libraries.

You will also have to add other external packages on your system as needed, such as Python3, cmake, etc. The system which the code is running on must have a GPU that is accessible with at least 8GB VRAM.

Workflow

Chunking and Embedding

First, collect documents in text form and dump them into a directory named texts/{doc_num} as .txt files, where doc_num is the number of .txt files in the directory you will be including. There should be many existing documents database sizes in the shared Dropbox available to the team. To chunk them, chunk.py will take each document and split it into chunks of specified length and overlap. These can be configured in this file or set later on. This will create a new chunk directory of the name chunks/{doc_num}_{chunk_size}, so you can have multiple chunking models available for use. This pattern will carry on in later parts as well.

Code in embed.py will use Snowflake/snowflake-arctic-embed-m-v2.0 to output the embeddings of the chunks, natural language, and tables to the embeddings directory. Right now, the file filters only natural language and tables used in the dw data warehouse, but this can be configured in this file.

Both steps can be run conveniently by using scripts/rechunk.sh, which takes in the chunk size as the singular argument and runs both the chunk generation as well as the embedding. The *_jsons directories contain many auxiliary files needed to run all steps of the pipeline.

Knowledge-Injection

Once all embeddings have been generated, similarity.py works as a retrieval library to add document chunks to needed objects using both embeddings and text search. Curently, it is set up to do this for both table serializations and natural language questions. This file is important, as it contains the core of the knowledge-injection, which is the dual embedding and text search algorithms.

There are two important arguments that can be passed in when running this file.

• The first is --ecat, which gives the embedding category to use. This will be in the form {doc_num}_{chunk_size}, that is, if you have 1 doucument and generated chunks with size 100, the respective embeddings would have been generated into embeddings/1_100/, and thus that directory should be passed into the code. By default, 6143_40 is used, as the default large doucment database has 6143 entries with a default chunk size of 40.
• The second is --k, which tells how many chunks should enrich each object. The default value is 5

These will then generate their respective mapping of each object to its most similar chunks in the rank_jsons direcotry, with the ecat and k postpended with underscores. For example, the enriched tables might be stored in rank_jsons/question_chunk_6143_40_5.json. The most similar values are stored as numbers, i.e. if a mapping contains 5, then the fifth document chunk should be substituted in its place if there is to be actual use with it.

Table Retrieval

Applications are essentially the above steps repeated, except now, instead of adding document chunks to enrich objects, we find the similarity between tables and natural language. The example code for this is in table_benchmark.py, which takes in a single argument k, the number of table to retrieve for a given natural language question.

Reranking

Before doing the previous step, we attempted to do LLM Reranking, basically asking the LLM to reorder the top matches from before. To do this, we would use an artificially high k=100, and then ask the LLM to filter it down to the true k=5 or any other desired value. This code can be run by itself by using the rerank.py file.