DAG在检索增强中的应用

General0805
437 阅读6分钟

DAG概述

在图论中，图是由顶点和连接这些顶点的边所构成，如果一个有向图从任意顶点出发无法经过若干条边回到该点，则这个图是一个有向无环图（英语：Directed Acyclic Graph，缩写：DAG）
每条边都带有从一个顶点指向另一个顶点的方向的图为有向图。有向图中的道路为一系列的边，系列中每条边的终点都是下一条边的起点。如果一条路径的起点是这条路径的终点，那么这条路径就是一个环。有向无环图即为没有环出现的有向图。

组成结构：

  • 顶点(vertices)
  • 边(edges)
  • 出度：从一个顶点出发的边的总数。
  • 入度：指向一个顶点的边的总数。

基本表示方法

简单一点来说，DAG实现状态流的核心在于描述单个任务（Task）他完成了什么样的动作（Action），以及他依赖什么样的任务（Task）。

TaskA:
  Actions:
  ○ 动作 A
  Dependencies:
TaskB:
  Actions:
  ○ 动作 B
  Dependencies:
  ○ TaskA
TaskC:
  Actions:
  ○ 动作 C
  Dependencies:
  ○ TaskB
TaskD:
  Actions:
  ○ 动作 D
  Dependencies:
  ○ TaskB
下图为依赖图
A -> B -> C
      |
       -> D

调度器设计

调度是指如何让节点一个一个的执行；一个节点执行完成，如何让其他节点感知并开始执行，常见的思路是，节点的分层调度，它的含义及特点如下：

  • 依赖分层算法（如广度优先遍历）提前计算好每一层需要执行的节点；节点一批一批的调度，无需任何通知和驱动机制。
  • 在同批次多节点时，由于各节点执行时间不同，容易出现长板效应。
  • 在多串行节点的图调度时，有较好的性能优势。

另一种常见的思路是，基于流水线思想的队列通知驱动模式：

  • 某节点执行完成后，立即发送信号给消息队列；消费侧在收到信号后，执行后续节点。如上图DAG中，B执行完成后，D/E收到通知开始执行，不需要关心C的状态。
  • 由于不关心兄弟节点的执行状态，不会出现分层调度的长板效应。
  • 在多并行节点的图调度时，有非常好的并行性能；但在多串行节点的图中，由于额外存在线程切换和队列通知开销，性能会稍差。

image.png

在大模型检索增强中的应用

检索增强是提升大模型生成内容准确性和时效性的成熟手段，在实践中，检索源往往来自多个，包括开放的搜索引擎以及自建库，汇总的各个库信息在输出给大模型进行推理之前需要经过粗排、精排、干预等一系列聚合逻辑。这种场景业务上有两个特点：

  1. 检索、排序等能力具有高度可复用性
  2. 绝大部分场景都是在1的基础能力之上进行自定义组合。

因此检索增强的业务逻辑可以很好的抽象为DAG，大大提升业务迭代过程的人效 检索流程可以大致抽象为干预->多路召回->[粗排]->精排-> 过滤 ->[答案抽取] 这样一个DAG模型

3.1 节点抽象

其中【过滤】就是平时一般说的【策略】，其逻辑是频繁变化的，对它的抽象是方案的核心。广义上上【粗排】、【精排】、【答案抽取】等也符合过滤操作的特征，可以统一抽象为输入输出均为召回列表的处理单元。

【干预】、【召回】可以抽象为输入是query，输出为召回列表的另一类处理单元

3.2 节点调度

  • DAG框架：负责流程配置解析和执行
  • processor：
  • 上下文：

3.3 流程

DAG分为静态和动态两种

静态DAG以文件形式维护在欧若拉代码库中，在线上运行；

动态DAG由用户通过接口将其配置传入，实时解析运行，使用场景是线下测试，为确保稳定性，线上不开放这个能力

静态DAG执行

动态DAG执行

4 DAG框架

框架在实现上有以下几个要求

  • 一个节点的所有子节点能够并行执行
  • 子节点需等待所有父节点执行完成后执行
  • 无环

4.1 描述

需要定义DSL来描述DAG，语法可以用yaml、json等，简单起见先用json来描述，以下示例表达了 【note库】->【排序】 这个简单的DAG

{
    "name": "base",
    "nodes": [
        {
            "name": "note",
            "processor": "ann_es",
            "processor_conf": {
                "index": ""
            },
        },
        {
            "name": "satisfy_rank",
            "processor": "rank"
            "processor_conf": {
                "model": "rank-v5"
            },
        },
    ]，
    "edges": [
        {
            "source": "A",
            "target": "B",
            "condition": "",
        }
    ]
}

4.2 DAG

4.2.1 DAG存储

使用邻接表作为DAG的存储方式，但维护指向子节点的指针的同时，还需要维护指向父节点的指针，原因是运行时子节点需要等待所有父节点完成后才能执行

type Node struct {
    Name string
    Processor
    Pre []*Node
    Next []*Node
    Conditions map[string]Condition
}

4.2.2 校验

  • 不能有环
  • 不能引用未知节点，因为所有处理节点都必须编码好的，目前不开放自定义节点的能力
  • 校验的同时完成拓扑排序并保存起来

4.2.3 调度

要求：

  • 子节点等待所有父节点执行完成后才能执行，并且merge所有父节点的hits
  • 一个节点的所有子节点能并行执行

实现要点

  • 使用无缓冲的channel作为父节点向子节点结果传递方式，同时能保证所有父节点执行完成后子节点才能执行
  • 使用waitgroup进行流程控制
func (d *Flow) Execute(c *ExecContext) []Hit {
    nodeRuntimeMap := make(map[string]*NodeRuntime)
    wg := sync.WaitGroup{}
    // 根据拓扑排序遍历DAG
    for _, name := range d.order {
        wg.Add(1)
        node := d.nodeMap[name]
        nodeRT := &NodeRuntime{Node: node, HitsChan: make(chan []Hit)}
        nodeRuntimeMap[name] = nodeRT
        go func() {
            defer wg.Done()
            defer close(nodeRT.HitsChan)
            
            hits := []Hit{}
            for _, pre := range node.Pre {
                // channel是无缓冲的，当父节点没有将结果写入时子节点将在此处阻塞
                hits = append(hits, (<- nodeRuntimeMap[pre.Name].HitsChan)...)
            }
            hits = node.Process(c, hits)
            for _, next := range node.Next {
                nodeRuntimeMap[next.Name].HitsChan <- hits
            }
        }()
    }
    wg.Wait()
    lastNodeName := d.order[len(d.order) - 1]
    hits := <- nodeRuntimeMap[lastNodeName].HitsChan
    return hits
}

4.2.4 Processor

processor是一个节点的执行器，包含具体的处理逻辑。

type Processor interface {
    func Process(c *ExecContext, hits []Hit) []Hit
}

降级策略： 所有节点都看作非关键节点，因此processor若出现失败情况则视为直接跳过

4.2.5 边规则

边规则是一类输出为布尔值的表达式，用来在运行时动态决策执行流程

规则引擎使用开源的expr，内置常用操作函数，支持用.访问属性值，因此规则可以写成len(Hits) == 0 或者 isSafe == 1的形式，开发和使用成本均比较低

type Condition struct {
    rule string
}
func (c *ExprCondition) IsSatisfied(env any) (bool, error) {}

avatar
文章
阅读
粉丝