快速入门:构建您自己的 RAG 应用

快速入门:构建您自己的 RAG 应用#

除去使用 RAG 助手这一概念,您也可以将 FlexRAG 作为库来开发您的 RAG 应用。FlexRAG 提供了一个灵活且模块化的 API 来帮助您构建 RAG 应用。下面的代码向您展示了如何构建一个简单的命令行检索式 QA 系统:

from flexrag.models import OpenAIGenerator, OpenAIGeneratorConfig
from flexrag.retriever import LocalRetriever


def main():
    # load the retriever
    retriever = LocalRetriever.load_from_hub("FlexRAG/enwiki_2021_atlas")

    # load the generator
    generator = OpenAIGenerator(
        OpenAIGeneratorConfig(
            model_name="Qwen2-7B-Instruct",
            base_url="http://10.28.0.148:8000/v1",
        )
    )

    # build a QA loop
    while True:
        query = input("Please input your question (type /bye to quit): ")
        if query == "/bye":
            break
        # retrieve the contexts
        contexts = retriever.search(query, top_k=3)[0]
        # construct the prompt
        user_prompt = (
            "Please answer the following question based on the given contexts.\n"
            f"Question: {query}\n"
        )
        for i, ctx in enumerate(contexts):
            user_prompt += f"Context {i+1}: {ctx.data['text']}\n"
        # generate the response
        response = generator.chat([{"role": "user", "content": user_prompt}])[0][0]
        print(response)

    return


if __name__ == "__main__":
    main()