[v0.2.3.beta.1] RAG

[liyin2015]

• Designing RAG
• Benchmark RAG
• Optimize RAG

(make any task pipeline auto-differentiable), generator used in retriever (rewrite, retrieve, read), retrieve + read (train these two RAGs on hotpot qa datasets)

need to improve the clarity on GradComponent, tensor + parameter (potentially differentiate trainable parameters with other intermediates outputs/inputs in the pipeline)

• Defined clear parameter type
• defined a forward func which is a wrapper on call function to handle parameter and dynamic computation graph
• Adapted default gradcomponnet's forward
• customized retriever's forward to save top_k as hyperparam and input as input type

Here is the computation graph using the following code for a vanila rag pipeline

    def foward(self, question: str, id: str = None) -> adal.Parameter:
        retriever_out = self.retriever.forward(input=question)
        successor_map_fn = lambda x: (
            "\n\n".join(x.data[0].documents)
            if x.data and x.data[0] and x.data[0].documents
            else ""
        )
        retriever_out.add_successor_map_fn(successor=self.llm, map_fn=successor_map_fn)
        generator_out = self.llm.forward(
            prompt_kwargs={"question": question, "context": retriever_out}, id=id
        )
        return generator_out

The computation graph

TODO:

• backward tests on classification + qa use case, potentially simplify previous code.
• tutorial + notebook for rag optimization