flexrag 0.1.1

A RAG Framework for Information Retrieval and Generation.

FlexRAG is an open-source framework designed for Retrieval-Augmented Generation (RAG), combining state-of-the-art information retrieval techniques with large language models (LLMs) for enhanced, context-aware responses.

:book: Table of Contents

• :book: Table of Contents
• :sparkles: Key Features
• :rocket: Getting Started
  • Step 0. Installation
    • Install from pip
    • Install from source
  • Step 1. Prepare the Retriever
    • Download the Corpus
    • Prepare the Index
  • Step 2. Run your RAG Application
    • Run the FlexRAG Example RAG Application with GUI
    • Run the FlexRAG Example Assistants for Knowledge Intensive Tasks
    • Build your own RAG Application
• :bar_chart: Benchmarks
• :label: License
• :pen: Citation
• :heart: Acknowledgements

:sparkles: Key Features

• Multiple Retriever Types: Supports a wide range of retrievers, including sparse, dense, network-based, and multimodal retrievers.
• Diverse Information Sources: Allows integration of multiple data types, such as pure text, images, documents, web snapshots, and more.
• Unified Configuration Management: Leveraging OmegaConf and dataclasses, FlexRAG simplifies configuration and management across your entire project.
• Out-of-the-Box Performance: Comes with carefully optimized default configurations for retrievers, delivering solid performance without the need for extensive parameter tuning.
• High Performance: Built with persistent cache system and asynchronous methods to significantly improve speed and reduce latency in RAG workflows.
• Research & Development Friendly: Supports multiple development modes and includes a companion repository, flexrag_examples, to help you reproduce various RAG algorithms with ease.
• Lightweight: Designed with minimal overhead, FlexRAG is efficient and easy to integrate into your project.

:rocket: Getting Started

Step 0. Installation

Install from pip

To install FlexRAG via pip:

pip install flexrag

Install from source

Alternatively, to install from the source:

pip install pybind11

git clone https://github.com/ZhuochengZhang98/flexrag.git
cd flexrag
pip install ./

You can also install the FlexRAG in editable mode with the -e flag.

Step 1. Prepare the Retriever

Download the Corpus

Before starting you RAG application, you need to download the corpus. In this example, we will use the wikipedia corpus provided by DPR as the corpus. You can download the corpus by running the following command:

# Download the corpus
wget https://dl.fbaipublicfiles.com/dpr/wikipedia_split/psgs_w100.tsv.gz
# Unzip the corpus
gzip -d psgs_w100.tsv.gz

Prepare the Index

After downloading the corpus, you need to build the index for the retriever. If you want to employ the dense retriever, you can simply run the following command to build the index:

CORPUS_PATH=psgs_w100.tsv.gz
CORPUS_FIELDS='[title,text]'
DB_PATH=<path_to_database>

python -m flexrag.entrypoints.prepare_index \
    corpus_path=$CORPUS_PATH \
    saving_fields=$CORPUS_FIELDS \
    retriever_type=dense \
    dense_config.database_path=$DB_PATH \
    dense_config.encode_fields='[text]' \
    dense_config.passage_encoder_config.encoder_type=hf \
    dense_config.passage_encoder_config.hf_config.model_path='facebook/contriever' \
    dense_config.passage_encoder_config.hf_config.device_id=[0,1,2,3] \
    dense_config.index_type=faiss \
    dense_config.faiss_config.batch_size=4096 \
    dense_config.faiss_config.log_interval=100000 \
    dense_config.batch_size=4096 \
    dense_config.log_interval=100000 \
    reinit=True

If you want to employ the sparse retriever, you can run the following command to build the index:

CORPUS_PATH=psgs_w100.tsv.gz
CORPUS_FIELDS='[title,text]'
DB_PATH=<path_to_database>

python -m flexrag.entrypoints.prepare_index \
    corpus_path=$CORPUS_PATH \
    saving_fields=$CORPUS_FIELDS \
    retriever_type=bm25s \
    bm25s_config.database_path=$DB_PATH \
    bm25s_config.indexed_fields='[title,text]' \
    bm25s_config.method=lucene \
    bm25s_config.batch_size=512 \
    bm25s_config.log_interval=100000 \
    reinit=True

Step 2. Run your RAG Application

When the index is ready, you can run your RAG application. Here is an example of how to run a RAG application.

Run the FlexRAG Example RAG Application with GUI

python -m flexrag.entrypoints.run_interactive \
    assistant_type=modular \
    modular_config.used_fields=[title,text] \
    modular_config.retriever_type=dense \
    modular_config.dense_config.top_k=5 \
    modular_config.dense_config.database_path=${DB_PATH} \
    modular_config.dense_config.query_encoder_config.encoder_type=hf \
    modular_config.dense_config.query_encoder_config.hf_config.model_path='facebook/contriever' \
    modular_config.dense_config.query_encoder_config.hf_config.device_id=[0] \
    modular_config.response_type=short \
    modular_config.generator_type=openai \
    modular_config.openai_config.model_name='gpt-4o-mini' \
    modular_config.openai_config.api_key=$OPENAI_KEY \
    modular_config.do_sample=False

Run the FlexRAG Example Assistants for Knowledge Intensive Tasks

You can evaluate your RAG application on several knowledge intensive datasets with great ease. The following command let you evaluate the modular RAG assistant with dense retriever on the Natural Questions (NQ) dataset:

OUTPUT_PATH=<path_to_output>
DB_PATH=<path_to_database>
OPENAI_KEY=<your_openai_key>

python -m flexrag.entrypoints.run_assistant \
    data_path=flash_rag/nq/test.jsonl \
    output_path=${OUTPUT_PATH} \
    assistant_type=modular \
    modular_config.used_fields=[title,text] \
    modular_config.retriever_type=dense \
    modular_config.dense_config.top_k=10 \
    modular_config.dense_config.database_path=${DB_PATH} \
    modular_config.dense_config.query_encoder_config.encoder_type=hf \
    modular_config.dense_config.query_encoder_config.hf_config.model_path='facebook/contriever' \
    modular_config.dense_config.query_encoder_config.hf_config.device_id=[0] \
    modular_config.response_type=short \
    modular_config.generator_type=openai \
    modular_config.openai_config.model_name='gpt-4o-mini' \
    modular_config.openai_config.api_key=$OPENAI_KEY \
    modular_config.do_sample=False \
    eval_config.metrics_type=[retrieval_success_rate,generation_f1,generation_em] \
    eval_config.retrieval_success_rate_config.context_preprocess.processor_type=[simplify_answer] \
    eval_config.retrieval_success_rate_config.eval_field=text \
    eval_config.response_preprocess.processor_type=[simplify_answer] \
    log_interval=10

Similarly, you can evaluate the modular RAG assistant with sparse retriever on the Natural Questions dataset:

OUTPUT_PATH=<path_to_output>
DB_PATH=<path_to_database>
OPENAI_KEY=<your_openai_key>

python -m flexrag.entrypoints.run_assistant \
    data_path=flash_rag/nq/test.jsonl \
    output_path=${OUTPUT_PATH} \
    assistant_type=modular \
    modular_config.used_fields=[title,text] \
    modular_config.retriever_type=bm25s \
    modular_config.bm25s_config.top_k=10 \
    modular_config.bm25s_config.database_path=${DB_PATH} \
    modular_config.response_type=short \
    modular_config.generator_type=openai \
    modular_config.openai_config.model_name='gpt-4o-mini' \
    modular_config.openai_config.api_key=$OPENAI_KEY \
    modular_config.do_sample=False \
    eval_config.metrics_type=[retrieval_success_rate,generation_f1,generation_em] \
    eval_config.retrieval_success_rate_config.context_preprocess.processor_type=[simplify_answer] \
    eval_config.retrieval_success_rate_config.eval_field=text \
    eval_config.response_preprocess.processor_type=[simplify_answer] \
    log_interval=10

You can also evaluate your own assistant by adding the user_module=<your_module_path> argument to the command.

Build your own RAG Application

To build your own RAG application, you can create a new Python file and import the necessary FlexRAG modules. Here is an example of how to build a RAG application:

from flexrag.retrievers import DenseRetriever, DenseRetrieverConfig
from flexrag.models import OpenAIGenerator, OpenAIGeneratorConfig
from flexrag.prompt import ChatPrompt, ChatTurn


def main():
    # Initialize the retriever
    retriever = DenseRetriever(
        DenseRetrieverConfig(
            database_path='path_to_database',
            encode_fields=['text'],
            top_k=1,
            passage_encoder_config=HfEncoderConfig(
                encoder_type='hf',
                hf_config=HfConfig(
                    model_path='facebook/contriever',
                    device_id=[0],
                )
            ),
        )
    )

    # Initialize the generator
    generator = OpenAIGenerator(
        OpenAIGeneratorConfig(
            model_name='gpt-4o-mini',
            api_key='your_openai_key',
            do_sample=False
        )
    )

    # Run your RAG application
    Prompt = ChatPrompt()
    while True:
        query = input('Please input your query: ')
        context = retriever.retrieve(query)[0]
        prompt_str = f"Question: {query}\nContext: {context.data["text"]}"
        prompt.update(ChatTurn(role="user", content=prompt_str))
        response = generator.chat(prompt)
        prompt.update(ChatTurn(role="assistant", content=response))
        print(response)
    return

if __name__ == "__main__":
    main()

We also provide several detailed examples of how to build a RAG application in the flexrag examples repository.

:bar_chart: Benchmarks

We have conducted extensive benchmarks using the FlexRAG framework. For more details, please refer to the benchmarks page.

:label: License

This repository is licensed under the MIT License. See the LICENSE file for details.

:pen: Citation

If you find this project helpful, please consider citing it:

@software{FlexRAG,
  author = {Zhang Zhuocheng},
  doi = {10.5281/zenodo.14306984},
  month = {12},
  title = {{FlexRAG}},
  url = {https://github.com/ZhuochengZhang98/flexrag},
  version = {0.1.0},
  year = {2024}
}

:heart: Acknowledgements

This project benefits from the following open-source projects:

• Faiss
• FlashRAG
• LanceDB
• ANN Benchmarks
• Chonkie
• rerankers