【Text2SQL 论文】PET-SQL:用 Cross-Consistency 的 prompt 增强的两阶段 Text2SQL 框架

发布于:2024-06-12 ⋅ 阅读:(19) ⋅ 点赞:(0)

论文:PET-SQL: A Prompt-enhanced Two-stage Text-to-SQL Framework with Cross-consistency

⭐⭐⭐

arXiv:2403.09732,商汤 & 北大

Code:GitHub

一、论文速读

论文一开始提出了以往 prompt-based 的 Text2SQL 方法的一些缺点:

  • 缺少对 table cell value 的先验关注。比如查询 sex 时,WHERE sex='Male'WHERE sex='M' 都是有可能的,LLM 也拿不准用哪个
  • 通过 prompt 让 LLM 做 schema linking 时,让 LLM 根据 question 和 schema 列出相关 schema 的效果并不好,因为像对于 CodeLlama 这样的模型来说,基于指令生成非代码文本并不是他们的强项
  • 目前方法的 post-refinements 收益也不明显:
    • 如果一个由强大的 LLM 生成的 SQL 具有语义模糊而不是语法错误的话,那 LLM 是无法 self-debugging 的
    • 使用 self-consistency 时,LLM 的投资和产出是不成比例,多次调用一个 LLM 产生数倍的成本,并没有从根本上让 SQL 的输出更加多样化

由此,论文提出了 Prompt-Enhanced Two-stage text2SQL framework with cross-consistencyPET-SQL)—— 通过两个 stage 来完成任务。

二、PET-SQL

2.1 Stage 1:few-shots prompt 生成 PreSQL

这个阶段构造出一个 few-shots 的 prompt 用来输入给 LLM。

2.1.1 prompt 的表示风格

有多种指导 LLM 生成 SQL 的 prompt 表示风格。以往研究发现,Code Representation(CRp)和 OpenAI Demonstration(ODp)这两种是很不错的选择,示例如下图:

在这里插入图片描述

论文基于 ODp 进一步丰富了 prompt 的信息,提出了 Reference-Enchanced representation(REp)的 prompt,改动如下:

  • 在 instruction 最后加了一句 You must minimize SQL execution time while ensuring correctness 来指示 LLM 不仅要关注执行正确性,还要注意 SQL 语句的效率,避免冗余的字符和操作符。
  • prompt 中增加 table cell value 示例,也就是随机选择数据库表的 3 rows
  • prompt 中增加 schema 的外键关系

提出的 REp 示例如下:

在这里插入图片描述

2.1.2 增加 few-shots examples

这里从 dataset 中检索出与 user question 相似的样本来作为 ICL 的 demonstrations。

注意,这里要对被检索的 NL-SQL pairs 做去语义化:也就是把 question 中与 table schema 相关的 tokens 用 <mask> 去掉,从而得到一个只呈现问题意图的 question skeleton。然后再对这些 question skeleton 使用 embedding 做语义嵌入,从而用于检索。

检索到的 top-K 被用作 demonstrations 用于 few-shots ICL。

这样,demonstrations 和前面的 REp prompt 结构,构成了用于输入给 LLM 的 prompt:

在这里插入图片描述

由此,可以得到一个 LLM 生成的 preliminary SQL(PreSQL)。

2.2 Stage 2:Schema Linking 与 FinSQL 生成

这里的实现基于一个简单的原则:PreSQL 与 question 高度相关。

因此,解析 PreSQL 以识别其中提到的 db tables 和 columns,并用于数据库的 schema linking。

然后 schema linking 的结果被用来简化 prompt 中的 schema 信息,去除掉无关的 table 和 column 信息,如下所示:

在这里插入图片描述

使用这个简化后的 prompt 用来生成 FinSQL。

2.3 Cross Consistency

以往做 Text2SQL 时往往会使用 self-consistency 来修正 SQL 结果,但这有一些问题:使用 self-consistency 往往是将 LLM 的问题调高使其生成多样化,然后多次生成后投票选出最终结果。但是研究也发现,高温下的 LLM 会增加幻觉、降低性能,且对于确定性任务(如生成 SQL),单个 LLM 的多样性也不足。

所以本文提出:在较低的温度下指导多个 LLMs 生成 SQL,然后在这些 SQL 的执行结果之间投票。这也就是 Cross Consistency

关于投票的思路,这里提出了两种:

  • Naive voting:让每个 LLMs 都生成一个 SQL,执行后进行投票,将大多数结果作为最终答案
  • Fine-grained voting:由于不同的 LLM 的能力不同,所能处理的问题的复杂度不同,因此,根据 PreSQL 的语法解析结果将 question 分为四个难度,然后不同的难度的问题由不同的候选 LLMs 来解决,并进行投票。这样可以最大限度发挥 LLM 的潜力,并显著减轻投票偏见。

三、实验

在 Spider 数据集上做了测试,使用 EX 作为评估,使用的 LLM 包括 CodeLlama、SQLCode、IternLM、SenseChat、GPT-4 等。

在 Spider 上的表现,在所有非基于学习的方法中实现了最高的 EX,比 DAIL-SQL 高出 1%。

具体的实现可以参考原论文。

四、总结

PET-SQL 的两阶段思路还是挺好理解上,prompt 设计上主要多了 db content 样例,整个过程中,使用 question de-semanticization 来提取问题骨架表示问题意图在实现上应该是整个流程里最复杂的。