DOCSEARCHbyAlgolia
Submit Docs

MemoRAG

repository·main·Indexed 25 days ago

https://github.com/qhjqhj00/memorag

A next-generation Retrieval-Augmented Generation (RAG) framework that uses a memory-inspired model to provide global understanding of large datasets. MemoRAG utilizes a memory-based data interface to recall query-specific clues, enhancing evidence retrieval and response generation. It supports long-context LLMs, external APIs (OpenAI, DeepSeek, Azure) as generators, and provides a simplified 'MemoRAG Lite' version for quick experimentation.

Tokens
17.4K
Snippets
31
Records
75
Agent score
78%

What's inside MemoRAG

  1. Overview of MemoRAG

    main
    MemoRAG is a RAG (Retrieval-Augmented Generation) framework that utilizes a memory-based data interface to achieve a global understanding of entire databases. Unlike standard RAG which focuses on explicit information needs in queries, MemoRAG uses a memory model to recall query-specific clues, enhancing evidence retrieval and generating more contextually rich responses. It is designed to handle massive datasets by bridging raw input to answers through contextual clues generated from global memory.

  6. Use APIs (OpenAI, DeepSeek, Azure) as generators

    main

    You can use external APIs as the generation model by passing an Agent object to the customized_gen_model parameter of the MemoRAG constructor.

    Supported Sources:

    • openai
    • deepseek
    • azure (via custom api_dict configuration)

    Configuration Formats:

    For deepseek:

    model = "..."
source = "deepseek"
api_dict = {
    "base_url": "",
    "api_key": ""
}

    For openai:

    model = "..."
source = "openai"
api_dict = {
    "api_key": ""
}
    from memorag import Agent, MemoRAG

# API configuration for Azure
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)

# 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",
    customized_gen_model=agent,
)

# Use the loaded context
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)

  7. Use MemoRAG Lite for quick experimentation

    main

    MemoRAG Lite is a simplified version of the pipeline designed for a quick, user-friendly experience. It supports English or Chinese contexts up to millions of tokens. While it works with other languages, performance may degrade due to English-default prompts.

    Hardware Note: While a 24GiB GPU is recommended, a 16GiB GPU can typically handle the pipeline under default settings.

    from memorag import MemoRAGLite
pipe = MemoRAGLite()
context = open("examples/harry_potter.txt").read()
pipe.memorize(context, save_dir="harry_potter", print_stats=True)

query = "What's the book's main theme?"
print(pipe(query))

  8. Use long-context LLMs as memory models

    main

    MemoRAG supports leveraging long-context LLMs (like Llama 3.1) as memory models by utilizing MInference to optimize context prefilling. This allows the system to benefit from the native large context windows of these models.

    from memorag import MemoRAG
model = MemoRAG(
    mem_model_name_or_path="shenzhi-wang/Llama3.1-8B-Chinese-Chat",    # For Chinese
    # mem_model_name_or_path="meta-llama/Meta-Llama-3.1-8B-Instruct",  # For English
    ret_model_name_or_path="BAAI/bge-m3",
    # cache_dir="path_to_model_cache",
    # access_token="hugging_face_access_token"
    )

  9. Initialize and use the MemoRAG pipeline

    main

    The standard MemoRAG class allows for full control over the memory, retrieval, and generation models. You can initialize it directly with HuggingFace model paths.

    Context Capacity:

    • TommyChien/memorag-qwen2-7b-inst: ~400K tokens (can reach 1M tokens with beacon_ratio=16).
    • TommyChien/memorag-mistral-7b-inst: ~128K tokens.

    Caching: When you call .memorize(), the encoded KV cache, Faiss index, and chunked passages are stored in the save_dir. You can reload this data later using .load() to avoid re-processing the context.

    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
    cache_dir="path_to_model_cache",  # Optional
    access_token="hugging_face_access_token",  # Optional
    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}")

# To reload from cache later:
pipe.load("cache/harry_potter/", print_stats=True)

  10. Quick Start with MemoRAG via Google Colab

    main

    You can try MemoRAG for free using a Google Colab notebook. This setup allows you to run the complete MemoRAG pipeline (Memory Model + Retriever + Generation Model) on a single T4 GPU (15GiB memory). It is capable of processing approximately 68K tokens (e.g., half of the provided Harry Potter example book) and performing all core functions.

    https://colab.research.google.com/drive/1fPMXKyi4AwWSBkC7Xr5vBdpPpx9gDeFX?usp=sharing

  11. Understand the Beacon forward pass mechanism

    main

    The _beacon_forward method implements a specialized execution loop that interacts with a Memory object. Unlike a standard transformer forward pass, it operates in steps:

    1. Prepare: self.memory.prepare(...) is called with input_ids, attention_mask, and labels.
    2. Step Loop: While self.memory.finish is false, the model calls self.memory.step() to retrieve the next chunk of input_ids, attention_mask, position_ids, past_key_values, and labels.
    3. Native Forward: The model executes _native_forward(...) on the retrieved chunk.
    4. Update: The model updates the memory state using self.memory.update_memory(outputs.past_key_values) and self.memory.update_loss(...).
    5. Output: Once the loop finishes, self.memory.output(outputs) provides the final model outputs.