MemoRAG: Moving Towards Next-Gen RAG Via Memory-Inspired Knowledge Discovery

Empowering RAG with a memory-based data interface for all-purpose applications!

Quick-Start | Colab | Usage | Demo | Dataset | FAQs

Overview

MemoRAG is an innovative RAG framework built on top of a highly efficient, super-long memory model. Unlike standard RAG, which primarily handles queries with explicit information needs, MemoRAG leverages its memory model to achieve a global understanding of the entire database. By recalling query-specific clues from memory, MemoRAG enhances evidence retrieval, resulting in more accurate and contextually rich response generation.​

MemoRAG Demo

We will provide a toy demo to demonstrate MemoRAG, you can try with the following scripts:

streamlit run demo/demo.py

Afterwards, you can view the demo as bellow:

:sparkles: Features

• Global Memory: Handles up to 1 million tokens in a single context, providing comprehensive understanding across massive datasets.
• Optimizable & Flexible: Adapts to new tasks with ease, achieving optimized performance with just a few hours of additional training.
• Contextual Clues: Generates precise clues from global memory, bridging raw input to answers and unlocking hidden insights from complex data.
• Efficient Caching: Speeds up context pre-filling by up to 30x, with support for caching chunking, indexing, and encoding.
• Context Reuse: Encodes long contexts once and supports repeated usage, boosting efficiency in tasks that require recurring data access.

:mag_right: Roadmap

MemoRAG is currently under active development, with resources and prototypes continuously being published at this repository.

• Codes / Models / Dataset Release
• Support OpenAI/Azure models
• Technical Report Release
• Demo Codes Release
• Training Codes for Memory model Release
• Support More Languages (e.g., Chinese)
• Light-Weight Optimization
• Enrich the Memory Ability

Note: The recent goals of MemoRAG are to achieve light-weight optimization through engineering improvements and to enhance its memory capabilities, enabling it to adapt to a wider range of applications and support additional languages.

:page_with_curl: Changelog

[10/09/24] We release MemoRAG's Technical Report.

[09/09/24] You can try MemoRAG on Google Colab for free.

[05/09/24] A Qwen2-based memory model is available at TommyChien/memorag-qwen2-7b-inst.

[03/09/24] A Mistral-based memory model is available at TommyChien/memorag-mistral-7b-inst.

[01/09/24] The project launched!

:rocket: Quick Start

:free: You can directly try MemoRAG on Google Colab for free.

In this notebook, we run the complete MemoRAG pipeline (Memory Model + Retriever + Generation Model) on a single T4 GPU with 15GiB of memory provided by Google Colab. Despite the limited resources, MemoRAG can process half of the content from the example book (~68K tokens) and perform all of its functions.

Installation

To use Memorizer and MemoRAG, you need to have Python installed along with the required libraries. You can install the necessary dependencies using the following command:

Install Dependencies

pip install torch==2.3.1
conda install -c pytorch -c nvidia faiss-gpu=1.8.0

Install from source

# clone this repo first
cd MemoRAG
pip install -e .

Install via pip

pip install memorag

For Quick Start, We provide a notebook to illustrate all functions of MemoRAG here.

:notebook: Usage

MemoRAG is easy to use and can be initialized with HuggingFace models directly. By using the MemoRAG.memorize() method, the memory model builds a global memory over a long input context. Empirically, with default parameter settings, TommyChien/memorag-qwen2-7b-inst can handle contexts of up to 400K tokens, while TommyChien/memorag-mistral-7b-inst can manage contexts up to 128K tokens. By increasing the beacon_ratio parameter, the model’s capacity to handle longer contexts can be extended. For example, TommyChien/memorag-qwen2-7b-inst can process up to one million tokens with beacon_ratio=16.

Basic Usage of MemoRAG

from memorag import MemoRAG

# Initialize MemoRAG pipeline
pipe = MemoRAG(
    mem_model_name_or_path="TommyChien/memorag-mistral-7b-inst",
    ret_model_name_or_path="BAAI/bge-m3", 
    gen_model_name_or_path="mistralai/Mistral-7B-Instruct-v0.2", # Optional: if not specify, use memery model as the generator
    cache_dir="path_to_model_cache",  # Optional: specify local model cache directory
    access_token="hugging_face_access_token",  # Optional: Hugging Face access token
    beacon_ratio=4
)

context = open("examples/harry_potter.txt").read()
query = "How many times is the Chamber of Secrets opened in the book?"

# Memorize the context and save to cache
pipe.memorize(context, save_dir="cache/harry_potter/", print_stats=True)

# Generate response using the memorized context
res = pipe(context=context, query=query, task_type="memorag", max_new_tokens=256)
print(f"MemoRAG generated answer: \n{res}")

When running the above code, the encoded key-value (KV) cache, Faiss index, and chunked passages are stored in the specified save_dir. Afterward, if the same context is used again, the data can be quickly loaded from the disk:

pipe.load("cache/harry_potter/", print_stats=True)

Typically, loading cached weights is highly efficient. For example, encoding, chunking, and indexing a 200K-token context takes approximately 35 seconds using TommyChien/memorag-qwen2-7b-inst as the memory model, but only 1.5 seconds when loading from cached files.

Summarization Task

To perform summarization tasks, use the following script:

res = pipe(context=context, task_type="summarize", max_new_tokens=512)
print(f"MemoRAG summary of the full book:\n {res}")

Using APIs as Generators

If you want to use APIs as a generator, refer to the script below:

from memorag import Agent, MemoRAG

# API configuration
api_dict = {
    "endpoint": "",
    "api_version": "2024-02-15-preview",
    "api_key": ""
}
model = "gpt-35-turbo-16k"
source = "azure"

# Initialize Agent with the API
agent = Agent(model, source, api_dict)
print(agent.generate("hi!"))  # Test the API

# Initialize MemoRAG pipeline with a customized generator model
pipe = MemoRAG(
    mem_model_name_or_path="TommyChien/memorag-qwen2-7b-inst",
    ret_model_name_or_path="BAAI/bge-m3",
    cache_dir="path_to_model_cache",  # Optional: specify local model cache directory
    customized_gen_model=agent,
)

# Load previously cached context
pipe.load("cache/harry_potter_qwen/", print_stats=True)

# Use the loaded context for question answering
query = "How are the mutual relationships between the main characters?"
context = open("harry_potter.txt").read()

res = pipe(context=context, query=query, task_type="memorag", max_new_tokens=256)
print(f"MemoRAG with GPT-3.5 generated answer: \n{res}")

Supported APIs for Generators

The built-in Agent object supports models from both openai and deepseek. Below are the configurations for initializing these models:

# Using deepseek models
model = ""
source = "deepseek"
api_dict = {
    "base_url": "",
    "api_key": ""
}

# Using openai models
model = ""
source = "openai"
api_dict = {
    "api_key": ""
}

Usage for Memory Model

The Memory model can be used independently to store, recall, and interact with the context. Here’s an example:

from memorag import Memory

# Initialize the Memory model
memo_model = Memory(
    "TommyChien/memorag-qwen2-7b-inst",
    cache_dir="path_to_model_cache",  # Optional: specify local model cache directory
    beacon_ratio=4  # Adjust beacon ratio for handling longer contexts
)

# Load and memorize the context
context = open("harry_potter.txt").read()
memo_model.memorize(context)

# Save the memorized context to disk
memo_model.save("cache/harry_potter/memory.bin")

# Query the model for answers
query = "How are the mutual relationships between the main characters?"

res = memo_model.answer(query)
print("Using memory to answer the query:\n", res)

# Recall text clues for evidence retrieval
res = memo_model.recall(query)
print("Using memory to recall text clues to support evidence retrieval:\n", res)

# Rewrite the query into more specific surrogate queries
res = memo_model.rewrite(query)
print("Using memory to rewrite the input query into more specific surrogate queries:\n", res)

Usage for Memory-Augmented Retrieval

In addition to the standalone Memory Model, MemoRAG provides memory-augmented retrieval functionality. This allows for improved evidence retrieval based on recalled clues from memory.

from memorag import MemoRAG

# Initialize MemoRAG pipeline
pipe = MemoRAG(
    mem_model_name_or_path="TommyChien/memorag-qwen2-7b-inst",
    ret_model_name_or_path="BAAI/bge-m3",
    cache_dir="path_to_model_cache",  # Optional: specify local model cache directory
    access_token="hugging_face_access_token"  # Optional: Hugging Face access token
)

# Load and memorize the context
test_txt = open("harry_potter.txt").read()
pipe.memorize(test_txt, save_dir="cache/harry_potter/", print_stats=True)

# Define the query
query = "How are the mutual relationships between the main characters?"

# Recall clues from memory
clues = pipe.mem_model.recall(query).split("\n")
clues = [q for q in clues if len(q.split()) > 3]  # Filter out short or irrelevant clues
print("Clues generated from memory:\n", clues)

# Retrieve relevant passages based on the recalled clues
retrieved_passages = pipe._retrieve(clues)
print("\n======\n".join(retrieved_passages[:3]))

Evaluation on Benchmark

Below are experiments results for the memory model, incorporating with three generation models.

Evaluation

To evaluate MemoRAG, use the following script:

cd examples
bash longbench/eval.sh

We will update other evaluation scripts soon.

Dataset

UltraDomain Benchmark: this repo.

Other Evaluation Data: this repo.

:raised_hands: FAQs

:bookmark: License

MemoRAG is licensed under the MIT License.

Citation

If you use MemoRAG in your research, please cite our paper:

@misc{qian2024memorag,
      title={MemoRAG: Moving towards Next-Gen RAG Via Memory-Inspired Knowledge Discovery}, 
      author={Hongjin Qian and Peitian Zhang and Zheng Liu and Kelong Mao and Zhicheng Dou},
      year={2024},
      eprint={2409.05591},
      url={https://arxiv.org/abs/2409.05591}, 
}