PET-SQL:基于大模型的两阶段Text2SQL方法

简介

PET-SQL出自论文《PET-SQL: A Prompt-enhanced Two-stage Text-to-SQL Framework with Cross-consistency》，将基于大模型的Text2SQL分为两个阶段进行，在第一阶段使用数据表schema信息、数据表采样数据、相似问答问答对生成初步的SQL(PreSQL)，在第二阶段，经过schema linking即过滤掉一些跟问题不相关的schema后生成最终的SQL(FinSQL)。接下来会使用Cross-consistency即用多个不同的LLM来投票决定生成的SQL是否正确。

第一阶段

pet-sql第一阶段的zero-shot prompt结构如下图(作者将其prompt模板称为reference-enhanced representation)：

• 指令里向LLM强调SQL生成时在保证正确性的前提下要考虑执行时间：“minimize SQL execution time while ensuring correctness”（图中第1行）
• 除了提供数据表的schema信息外，每个数据表采样了3条数据（图中第7-11行）
• 提供了数据表的外键信息（图中第12-15行）

除了zero-shot prompt外，作者考虑将与问题类似意图的question-SQL样例作为few-shot（如下图，将检索到的question-SQL放在zero-shot prompt前面）。few-shot样例采用了论文参考文献9：A Case-Based Reasoning Framework for Adaptive Prompting in Cross-Domain Text-to-SQL提到的question de-semanticization方法。具体来说，将问题里与特定数据表相关的token如表名、列名、数据取值等使用特殊token <mask>替换得到问题骨架（问题骨架被认为是问题意图的表示）。将训练集中所有问题骨架和question-SQL存储起来作为检索池，对查询问题骨架化处理后使用向量化检索方式得到与查询问题最相似的top k question-SQL作为few shot。

注：PET-SQL 引用的参考文献9：A Case-Based Reasoning Framework for Adaptive Prompting in Cross-Domain Text-to-SQL 与 Prompting GPT-3.5 for Text-to-SQL with De-semanticization and Skeleton Retrieval 的全部作者名单是相同的，思路也基本一致

第二阶段

根据第一阶段生成的PreSQL 内容，去掉与PreSQL中不相关的schema信息和外键信息，这个过程被称为schema linking (SL)（注：vanna 是通过向量检索方式来召回与问题最相关的schema，方法不一样，目的都是为了减少干扰信息）

schema linking之后简化的zero-shot prompt如下图所示。

cross consistency

作者认为self-consistency方法不是足够好的后处理方法，采取的是通过多个LLM来生成SQL的cross consistency策略，包括以下两种方法：

• Naive voting across several LLMs：使用更强的LLM如GPT4生成PreSQL，第二阶段使用不同的低温度设置的LLM生成FinSQL。将所有FinSQL和PreSQL在数据库引擎上执行并获取结果，将生成多数结果的SQL作为最后的答案。
• Fine-grained voting according to PreSQL complexity：在Naive voting 基础上考虑PreSQL的难易程度，按照难易程度进行分组投票。

总结

PET-SQL两阶段Text2SQL的思路还是比较好理解，prompt设计上主要多了数据样例，感觉使用question de-semanticization来提取问题骨架表示问题意图在实现上是整个流程里最复杂的。