Lazy loaded image5.5 Enterprise-level RAG Practice (II)

In [5.4 Enterprise-level RAG Practice (I)], we modified the meta info file for our own needs, vectorized a few papers, designed a few questions, and successfully retrieved related information and answered questions. It’s practical, but it’s still hard to fully understand how all of these work behind the scenes by just running a few lines of code, let alone adapting it to more complex cases.
Let’s dive deeper into it to learn the exact steps this project went through to put all of these fancy stuff together: parsing, chunking, vectorization, rephrasing the question, searching, reranking, etc. After understanding all it has done, you will be able to modify it for your own needs.
In this tutorial, we use PDFs from this website for demonstration, where each chapter is saved in a PDF file, for example: LLM Components, LLM Pre-training, and so on.

5.5.1 Workflows

Below is the full workflow used in the original repo, where the term company is used, but you can just replace it with anything you like. In our case, it’s chapter of the blog. It’s important to know each step of it so you can customize more complex workflows.
Note there are rhombus-shaped steps. They mainly check if there is company being indentified or is the user question/query for single or multiple companies. Because in the Enterprise RAG Challenge, it involves comparing numbers between companies.
In this tutorial, let’s just keep these workflows before we are able to customize them!
notion image
Where we use four tools:
  • LangChain: Chunking the JSONs with RecursiveCharacterTextSplitter, which tries to preserve natural text structure (paragraphs, sentences). It may look like using GPT-4o, but in fact, it does not call or use the LLM itself. It just uses the model’s tokenizer rules, which are defined by the tiktoken library (developed by OpenAI).
    • Faiss: This tool is from Meta for efficient similarity search and clustering of dense vectors. We will use this for indexing the embeddings (faiss.IndexFlatIP) and similarity matching between embeddings (faiss.read_index.search).
    • ChatGPT: We use this a lot, such as text-embedding-3-large for vectorization, rephrasing the questions into subquestions using gpt-4o, vectorizing the user question with text-embedding-3-large, reranking the retrieved pages using gpt-4o-mini, generating the answer based on the retrieved context using o3-mini.
      • You are free to change all these embeddings models and LLMs into whatever you think would work the best for you
    <ins/>

    5.5.2 Codeflows

    In the previous tutorial, the code is extremely simple as shown below, because everything is encapsulated in the Pipeline class.
    Now we present the code running logic to have a deeper understanding of the implementation and codeflows. The parsing, merging, exporting to markdown, chunking, and vectorizing belong to the preprocessing steps and are relatively simple and straightforward.
    Basically, you don’t have to modify them; just use them as functions as long as you are also retrieving information from PDFs; or removing the parsing PDFs part and adding, for example, markdown to JSON format if you are using your markdown notes.
    So let’s more focus on the last but not least part: processing questions, which contains the most critical and complex logic of RAG. And also, this is where you want to modify the code for your own needs.
    Let’s start with clarifying the files used in the whole workflow and then tracking the functions invoking logic.
    • pipeline.py
    • pdf_parsing.py
    • parsed_reports_merging.py
    • text_splitter.py
    • ingestion.py
    • questions_processing.py
    • retrieval.py
    • reranking.py
    • api_requests.py
    • prompts.py
    We use different colors for different Python files for better visualization.
    🏗️
    1. Parsing:
      1. pipeline.parse_pdf_reports_sequential()
        pdf_parsing.parse_and_export()
    1. Merging pieces into pages:
      1. pipeline.merge_reports()
        parsed_reports_merging.process_reports()
    1. Exporting to Markdown:
      1. pipeline.export_reports_to_markdown()
        parsed_reports_merging.export_to_markdown()
    1. Chunking:
      1. pipeline.chunk_reports()
        text_splitter.split_all_reports()
    1. Vectorization:
      1. pipeline.create_vector_dbs()
        ingestion.process_reports()
    1. Processing questions:
      1. pipeline.process_questions()
        questions_processing.process_all_questions() → process_questions_list() → _process_single_question() process_question()
        • If single company: questions_processing.get_answer_for_company()
        • If multiple companies: questions_processing.process_comparative_question()
        In either case of above, the logic is simple:
        • retrieval.VectorRetriever.retrieve_by_company_name
        • reranking.LLMReranker.rerank_documents
        • api_requests.get_answer_from_rag_context
        In the above reranking and api_requests functions, the prompts defined in prompts.py are used.
    The brief demonstration above may seem a bit messy, but it should give you an initial sense of how the files are organized and how the functions are invoked. It’s better you check the source code yourself and use this as a reference if you find yourself lost in the code.
    To fit them for your needs, by maintaining the overall processing logic and workflows, the only thing you need to pay attention to is the prompts.py, where all utilized systems prompts, task-specific prompts, general instructional prompts, and response schema prompts are defined. They are designed for the companies' annual reports retrieval; to adapt to our case: knowledge answering based on technical notes/PDFs, we have to modify these prompts.
    <ins/>

    5.5.3 Customizing Prompts

    Let’s go through all code snippets in prompts.py that need to be modified and customized, following the original order of class definitions so that you can compare them with the original prompts. You don’t have to copy them because the full code is provided in the end. Just get familiar with them first.
    1. RephrasedQuestionsPrompt: Rephrasing a comparative question to subquestions
      1. In the above code, we mainly change the company to chapter because the knowledge is organized in chapter on this website. You can modify it for your specific needs, like paper, class, code file, etc. Also, we change the few-shot example to provide context about a technical blog.
        You may notice that the companies variable is kept in the user_prompt above. This is OK because it represents a variable and would be replaced with our chapter names in the code like this: companies=", ".join([f'"{company}"' for company in companies]).
        📌
        You may also notice that there are many variables named company in files other than prompts.py like below, which is from questions_processing.py. If you want to modify them, feel free to do so, but if you don’t, you are totally fine because they are just variable names. As long as the prompts are correctly customized, you are good to go.
        notion image
    1. AnswerWithRAGContextSharedPrompt: A class containing shared instruction and user prompt that are used in other classes defined later: AnswerWithRAGContextNamePrompt, AnswerWithRAGContextNumberPrompt, AnswerWithRAGContextBooleanPrompt, and AnswerWithRAGContextNamesPrompt.
      1. Similarly, we change the company-related stuff to chapter
    1. AnswerWithRAGContextNamePrompt: The answer template and prompts for response with a name
      1. In this class, we modify words like company and product, and customize the few-shot prompt from querying the CEO of a company to the terminology in a chapter.
    1. AnswerWithRAGContextNumberPrompt: The answer template and prompts for response with a number
      1. In this class, we adapted both examples to technical use cases. As you can observe, designing effective prompts to guide the LLM is a non-trivial task that requires careful attention. It often demands significant time and effort, and it plays a critical role in determining the final performance of your RAG system.
    1. AnswerWithRAGContextBooleanPrompt: The answer template and prompts for response with a boolean
      1. AnswerWithRAGContextNamesPrompt: The answer template and prompts for response with multiple names
        1. ComparativeAnswerPrompt: The answer template and prompts for comparing multiple chapters
          1. That’s pretty much it! The remaining prompt definitions, such as RerankingPrompt, are generic enough to work under different semantics, so we can leave them as they are.
        Coming to this point, we’ve obtained the minimal version that can be run. By copying all the modified prompts classes into the prompts.py file, you can expect a better performance of this RAG system than not customizing these prompts.
        We tested three questions and all were correctly answered:
        • What's the hidden embedding length of Transformers PE in LLM Components?
          • Let’s see the page indices of 72 and 73 (numbering from 0); the one shown below image is indexing from 1
            notion image
        • What's the first step of LLM Pre-training?
          • Honestly, I didn’t really think of the answer when I typed this question and I just let the model find it, then it found the page index 16 (page 17 if starting from 1) as shown below
            notion image
        • Does the LLM Applications contain rotary positional embedding?
          • Clearly, the RoPE is not mentioned in the LLM Applications chapter and the model knows that, so the references are just two random pages with the highest similarities.
        The Colab link is here. In the code, you will download all results in my previous run (changes to prompts.py, faiss files, designed questions, and final answers), so you may want to skip some steps such as parsing the PDFs to save some time.
        Note that we also modify the parsed_reports_merging.py by adding the code into the supported block_type, which originally did not support merging the code type of text and raised an error.
        <ins/>

        5.5.4 Change the Workflow

        The RAG system works fine so far, but the problem is that we can only ask questions that lead to responses in the format of a name, a number, or a boolean. We cannot ask general questions like How RoPE works? and expect an explanation answer yet.
        Such questions do not specify the chapter name and need the system to automatically find the most relevant chapters and further find the most related pages, then explain it to us with a few sentences. Also, we need to customize a prompt class for such general questions.
        Before we dive into the code, let’s outline what needs to be done to implement this feature:
        • Absolutely adding an AnswerWithRAGContextExplanationPrompt class
        • Although the user didn’t specify a chapter, the system should match the potentially best one by default, regardless of the types of questions
        That’s the plan. Let’s start building it step by step!
        1. Adding the AnswerWithRAGContextExplanationPrompt to the prompts.py
          1. Add an elif condition statement to match our customized prompt class in the _build_rag_context_prompts function from api_requests.py
            1. When no chapter is detected from the user’s question, we match the best one for it, which may be irrelevant though. Let’s add a function match_the_best_chapter to the retrieval.py to find the best chapter first.
              1. The if distances[0][0] > best_score may look counterintuitive at first glance, but note that we are using IndexFlatIP with cosine similarity as metric, so the returned distance actually indicates similarity: the higher the better.
            1. Let’s invoke above added function at a proper place: process_question function in the questions_processing.py by changing the
              1. to
                Let’s just leave the variable name company_name as is; it’s not really causing any issues (yeah, I get the readability thing).

            That’s it! Let's design some questions and see how this RAG system performs
            • What's the hidden embedding length of Transformers Positional Encoding?
              • We use the same question as before but without specifying the corresponding chapter to see if the system can correctly find it
                The answer 512 is correct and the first related page is also just right!
            • What's the first step of LLM pretraining?
              • You can see, even without specifying the chapter, now the RAG system can automatically match the right one.
                Both the answer and the references (chapter and pages) are also correct for this question.
            • How does RoPE work?
              • Let’s ask a general question, which requires an explanation of the concept
                You can observe the chapter is located correctly at LLM_roadmap_1 and the answer is given as
                RoPE, or Rotary Positional Embedding, operates by applying position-dependent rotations to token embeddings. Specifically, it splits each embedding into two parts, treats these as the real and imaginary components of a complex number, and multiplies them by a complex rotation factor derived from sine and cosine functions with frequencies that vary by dimension. This rotation integrates positional information directly into the embeddings so that when the dot product between queries and keys is computed during attention, the resulting scores naturally reflect the relative position of tokens.
                Looks good. Let’s check the first page (index 84, i.e., page 85) being referenced. Interestingly, to better explain the concept deep down, this RAG system chose to reference the source code implementation of RoPE instead of referencing the text explanation.
                notion image
            • What's the steps to train a nanoGPT from scratch?
              • The last question asks for the steps of training the nanoGPT, in which case, the RAG system may have to go through the whole content of [8.1 nanoGPT: Training a GPT from Scratch] to summarize it
                We can see that the chapter is also correctly located and the answer is given as
                Training nanoGPT from scratch involves several clearly defined steps. First, set up the environment by installing necessary libraries, using either Anaconda or Google Colab, and then download the dataset (e.g., tinyShakespeare). Next, tokenize the text into numerical representations and split the data into training and validation sets. Define the model architecture including token/positional embeddings, transformer blocks with multi-head self-attention and feed-forward networks, and layer normalization. Configure training hyperparameters and set up an optimizer (such as AdamW). Proceed with a training loop that performs forward passes, computes loss, backpropagates, and updates parameters, while periodically evaluating performance on both training and validation data. Finally, use the trained model to generate new text from a given context.
                Well, it actually looks great.
            <ins/>

            Cost: around 100,000 tokens for the 4 questions
            Feel free to play with it using more diverse questions. The full Colab code can be found here
            📌
            Comments
            Thanks to IlyaRice, who open-sourced this incredible project, which, in my opinion, is a great material to start with RAG and build a high-performing RAG system. With the practice in this tutorial, you can now customize it for your needs. You can also use it for commercial purposes because it’s MIT-licensed.
            The MIT License is a permissive, free software license that grants broad rights to users, including the ability to use, copy, modify, and distribute the software, both personally and commercially.
            Of course, it’s still not perfect yet. An obvious inconvenience is that if you are frequently updating existing content (like this AI blog being updated frequently), you have to rebuild the vectorized embeddings for the modified chapters. So maybe the best use case for this is for retrieving information from a huge but frozen database, like your class notes from previous semesters, all papers from CVPR 2024, or just 100 annual reports of the companies.
            Prev
            5.4 Enterprise-level RAG Practice (I)
            Next
            6.1 GPU Usage Analysis
            Loading...
            Article List
            LLM Learning Roadmap
            ✨ Awesome-Anything
            🖼️ Digital Image Processing
            🍃 LLM Components
            🌱 LLM Pre-training
            ☘️ LLM Post-Training
            🍀 LLM Popular Models
            🪴 LLM Applications
            🌿 LLM Optimization
            🌾 LLM Compression
            🌵 LLM Hands-on Practice
            🌴 LLM Must-read Papers
            🌳 LLM Q&A
            🐝 VLM Image Encoders
            📝 MISC.