本节将带你实现一种典型的迭代推理能力的 RAG 工作流:IterRetGen。该流程支持基于模型输出反复更新查询,从而逐步收敛至更优的答案。
论文:https://arxiv.org/pdf/2305.15294

Step 1: 明确工作流结构

我们先来回顾论文提出的 IterRetGen 原始流程结构图: 在 UltraRAG 中,你可以基于已有模块快速实现该工作流,其整体架构可抽象为下图: 如图所示:
  • 模型每轮生成的答案会与原始问题拼接,作为下一轮检索的查询。
  • 该过程按照设定的最大循环次数执行 N 次
  • 除了 prompt 和 custom 两个模块需自定义实现,其他功能均可直接复用 UltraRAG 内置工具。

Step 2: 实现必要 Tool

IterRetGen 通过拼接查询 query 和当前轮次生成的答案作为本轮查询,为此在 servers/custom/src/custom.py 中添加如下代码:
servers/custom/src/custom.py
@app.tool(output="q_ls,ret_psg->nextq_ls")
def iterretgen_nextquery(
    q_ls: List[str],
    ans_ls: List[str | Any],
) -> Dict[str, List[str]]:
    ret = []
    for q, ans in zip(q_ls, ans_ls):
        next_query = f"{q} {ans}"
        ret.append(next_query)
    return {"nextq_ls": ret}

Step 3: 编写 Pipeline 配置文件

examples/ 目录下新建一个 YAML 文件,如 IterRetGen.yaml
/images/yaml.svgexamples/IterRetGen.yaml
# IterRetGen demo

# MCP Server
servers:
  benchmark: servers/benchmark
  retriever: servers/retriever
  prompt: servers/prompt
  generation: servers/generation
  evaluation: servers/evaluation
  custom: servers/custom

# MCP Client Pipeline
pipeline:
- benchmark.get_data
- retriever.retriever_deploy_search
- prompt.qa_rag_boxed
- generation.generate
- custom.output_extract_from_boxed
- loop:
    times: 3
    steps:
    - custom.iterretgen_nextquery:
        input:
          ans_ls: pred_ls
    - retriever.retriever_deploy_search:
        input:
          query_list: nextq_ls
    - prompt.qa_rag_boxed
    - generation.generate
    - custom.output_extract_from_boxed
- evaluation.evaluate

Step 4:配置 Pipeline参数

执行下列命令:
ultrarag build examples/IterRetGen.yaml
打开生成的 examples/parameter/IterRetGen_parameter.yaml,修改如下配置:
/images/yaml.svgexamples/parameter/IterRetGen_parameter.yaml
benchmark:
  benchmark:
    key_map:
      gt_ls: golden_answers
      q_ls: question
    limit: 2
    name: asqa
    path: data/sample_asqa_5.jsonl
custom: {}
evaluation:
  metrics:
  - acc
  - f1
  - em
  - coverem
  - stringem
  - rouge-1
  - rouge-2
  - rouge-l
  save_path: output/asqa.json
generation:
  base_url: http://localhost:8000/v1
  model_name: openbmb/MiniCPM4-8B
  sampling_params:
    extra_body:
      chat_template_kwargs:
        enable_thinking: false
      include_stop_str_in_output: true
      top_k: 20
    max_tokens: 2048
    temperature: 0.7
    top_p: 0.8
prompt:
  template: prompt/qa_boxed.jinja
retriever:
  query_instruction: 'Query: '
  retriever_url: http://localhost:8080
  top_k: 5

Step 5:运行你的推理流程!

一切准备就绪后,执行以下命令启动推理流程:
ultrarag run examples/IterRetGen.yaml