TableRAG:一个专为表格理解设计的检索增强生成框架
TableRAG是一个用于解决异构文档问答中结构信息丢失和缺乏全局视图问题的框架。它结合了SQL执行和文本检索,能够统一理解表格数据并进行复杂操作。这个框架包括离线数据库构建和四步在线迭代推理过程。在新基准HeteQA和其他公共基准测试中,TableRAG表现优异,超越了现有的RAG和程序化方法,成为处理复杂异构文档推理的先进解决方案。
核心原理
查询扩展与检索:TableRAG通过查询扩展、模式检索和单元格检索来定位表格中的关键信息。它会生成针对模式和单元格值的独立查询,避免单一查询带来的冗余信息。
模式检索:ai语言模型可以通过列名识别关键列及其数据类型,无需对整个列进行编码。例如,对于问题“钱包的平均价格是多少?”语言模型会生成列名查询,如“产品”和“价格”。
单元格检索:基于语言模型生成的查询,检索与问题相关的列名和单元格。例如,检索与“钱包”相关的单元格值。
频率感知截断:引入单元格编码预算B,如果不同值的数量超过B,则只编码出现频率最高的前B个值。
TableRAG主要模块
online_inference/main.py:TableRAG的核心推理代码,包含TableRAG类,实现在线推理的主要逻辑。
offline_data_ingestion_and_query_interface/src/interface.py:使用Flask框架提供SQL查询接口。
online_inference/tools/sql_tool.py:包含SQL生成和执行的工具函数,支持重试机制。
offline_data_ingestion_and_query_interface/src/sql_alchemy_helper.py:提供SQLAlchemy数据库操作的辅助类,支持数据插入和查询。
offline_data_ingestion_and_query_interface/src/service.py:处理TableRAG请求的核心逻辑,包括从LLM获取SQL语句并执行。
online_inference/chat_utils.py:包含获取LLM生成结果的工具函数,支持不同API后端。
offline_data_ingestion_and_query_interface/src/data_persistent.py:用于清洗DataFrame列名的工具函数。
offline_data_ingestion_and_query_interface/src/handle_requests.py:调用大模型推理接口,包含异常捕获和指数退避重试机制。
online_inference/tools/retriever.py:实现语义检索和混合文档检索的功能,包括召回和重排序。
优势
高效数据编码:通过模式检索和单元格检索,TableRAG减少了语言模型的输入长度,使其与表格大小无关,提高了编码效率,降低了计算成本。
减少信息丢失:仅将检索到的模式和单元格提供给语言模型,避免了处理整个表格导致的关键信息丢失。
提升推理能力:TableRAG通过提供关键信息,增强了语言模型的表格推理能力。
TableRAG应用场景
数据分析和报告生成:适用于需要处理大规模表格数据的场景,如数据分析、报告生成和财务报表处理。
跨表问答:能够处理基于表格的事实验证、单跳和多跳表格问答等任务。
实验验证
研究者开发了基于Arcade和BIRD-SQL数据集的百万级标记基准测试,评估TableRAG在大规模表格理解任务中的表现。实验结果表明,TableRAG在检索质量和大规模表格理解性能方面优于现有方法。
项目链接
Github:https://github.com/yxh-y/TableRAG