DOCSEARCHbyAlgolia
Submit Docs

RAPTOR (Recursive Abstractive Processing for Tree-Organized Retrieval)

repository·master·Indexed 23 days ago

https://github.com/parthsarthi03/raptor

A system that builds a hierarchical tree of text summaries to enable high-quality retrieval and question answering across long documents. It integrates information at multiple levels of abstraction using the RetrievalAugmentation class to embed, cluster, and summarize text chunks. Supports OpenAI by default and allows extension with custom open-source models via BaseSummarizationModel, BaseQAModel, and BaseEmbeddingModel.

Tokens
1.2K
Snippets
5
Records
5
Agent score
33%

What's inside RAPTOR

  1. Build a tree using RetrievalAugmentation

    master

    To build a RAPTOR tree, instantiate RetrievalAugmentation and use the add_documents method. RAPTOR will recursively embed, cluster, and summarize text chunks to create a tree structure with varying levels of abstraction.

    from raptor import RetrievalAugmentation

RA = RetrievalAugmentation()
# construct the tree
RA.add_documents(text)

  2. Extend RAPTOR with custom Open Source models

    master

    RAPTOR allows you to use models like Llama, Mistral, or Gemma by extending the base classes for summarization, QA, and embeddings.

    To use custom models, follow these steps:

    1. Subclass BaseSummarizationModel and implement summarize(self, context, max_tokens=150).
    2. Subclass BaseQAModel and implement answer_question(self, context, question).
    3. Subclass BaseEmbeddingModel and implement create_embedding(self, text).
    4. Wrap these models in a RetrievalAugmentationConfig object.
    5. Pass the config to RetrievalAugmentation(config=...).
    from raptor import BaseSummarizationModel, BaseQAModel, BaseEmbeddingModel, RetrievalAugmentationConfig

# Example: Custom Summarization Model
class GEMMASummarizationModel(BaseSummarizationModel):
    def __init__(self, model_name="google/gemma-2b-it"):
        # ... initialization logic ...
        pass

    def summarize(self, context, max_tokens=150):
        # ... implementation ...
        return summary

# Example: Custom QA Model
class GEMMAQAModel(BaseQAModel):
    def __init__(self, model_name="google/gemma-2b-it"):
        # ... initialization logic ...
        pass

    def answer_question(self, context, question):
        # ... implementation ...
        return answer

# Example: Custom Embedding Model
from sentence_transformers import SentenceTransformer
class SBertEmbeddingModel(BaseEmbeddingModel):
    def __init__(self, model_name="sentence-transformers/multi-qa-mpnet-base-cos-v1"):
        self.model = SentenceTransformer(model_name)

    def create_embedding(self, text):
        return self.model.encode(text)

# Integration
RAC = RetrievalAugmentationConfig(
    summarization_model=GEMMASummarizationModel(), 
    qa_model=GEMMAQAModel(), 
    embedding_model=SBertEmbeddingModel()
)
RA = RetrievalAugmentation(config=RAC)

  3. Initialize RAPTOR with OpenAI

    master

    By default, RAPTOR requires an OpenAI API key for application initialization. If you are not using OpenAI models, you must still provide a placeholder string (e.g., "not_used") to the environment variable to prevent initialization errors.

    import os
os.environ["OPENAI_API_KEY"] = "your-openai-key"

  4. Save and load RAPTOR trees

    master

    You can persist a constructed tree to disk using RA.save(path) and reload it later by passing the saved path to the RetrievalAugmentation constructor via the tree parameter.

    # Save the tree
SAVE_PATH = "demo/cinderella"
RA.save(SAVE_PATH)

# Load back the tree
RA = RetrievalAugmentation(tree=SAVE_PATH)
answer = RA.answer_question(question=question)

  5. Query the tree with answer_question

    master

    Once the tree is built, you can perform queries across different abstraction levels using the answer_question method on a RetrievalAugmentation instance.

    question = "How did Cinderella reach her happy ending ?"
answer = RA.answer_question(question=question)
print("Answer: ", answer)