parthsarthi03 / raptor

The official implementation of RAPTOR: Recursive Abstractive Processing for Tree-Organized Retrieval
https://arxiv.org/abs/2401.18059
MIT License
687 stars 98 forks source link

Return Citation #22

Open junmagic-ai opened 3 months ago

junmagic-ai commented 3 months ago

Hi team, really awesome library. Is there a way to return the source text to be used in the retrieval as source citation? Similar to llama_index's CitationQueryEngine.

tskippervold commented 3 months ago

You can just adjust your prompts to the LLM of your choosing.

Here is an example from Anthropic: https://docs.anthropic.com/claude/page/cite-your-sources

My point is, llamaIndex and all these libraries very often don't do anything special, but just have prompt templates they use agains the LLM you choose.

You can inspect the source code yourself here: https://github.com/run-llama/llama_index/blob/cfb2d7a58a8b8fe070ece26322b7df39e8d2b804/llama-index-core/llama_index/core/query_engine/citation_query_engine.py#L32

parthsarthi03 commented 3 months ago

Adding on to @tskippervold answer, you can adjust define your own custom prompt by defining your own QAModel.

For example, to use your GPT-3 Turbo as a QA model with a prompt asking for citations, you can do the following:

from raptor import RetrievalAugmentation, RetrievalAugmentationConfig, BaseQAModel
from openai import OpenAI
from tenacity import retry, stop_after_attempt, wait_random_exponential

class GPT3TurboCitationQAModel(BaseQAModel):

    def __init__(self, model="gpt-3.5-turbo"):
        self.model = model
        self.client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])

    @retry(wait=wait_random_exponential(min=1, max=20), stop=stop_after_attempt(6))
    def _attempt_answer_question(
        self, context, question, max_tokens=150, stop_sequence=None
    ):
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": "You are an expert research assistant."},
                {
                    "role": "user",
                    "content": f"Given the following context:\n\n{context}\n\nAnswer the question: {question}\n\nPlease provide the most relevant and concise answer, and cite the specific parts of the context that support your answer.",
                },
            ],
            temperature=0,
        )
        return response.choices[0].message.content.strip()

    @retry(wait=wait_random_exponential(min=1, max=20), stop=stop_after_attempt(6))
    def answer_question(self, context, question, max_tokens=150, stop_sequence=None):
        try:
            return self._attempt_answer_question(
                context, question, max_tokens=max_tokens, stop_sequence=stop_sequence
            )
        except Exception as e:
            print(e)
            return e

RAC = RetrievalAugmentationConfig(qa_model=GPT3TurboCitationQAModel())
RA = RetrievalAugmentation(config=RAC)
RA.add_documents(text)

You can try experimenting with different prompts, or even try stronger models. If you have a sample dataset of questions and answers with citations, you could also use DSPy (https://github.com/stanfordnlp/dspy) to automatically optimize the prompt for your specific use case.