AID-SQL: Adaptive In-Context Learning of Text-to-SQL with Difficulty-Aware Instruction and Retrieval-Augmented Generation

Recent research in Text-to-SQL translation has primarily adopted in-context learning methods leveraging large language models (LLMs), achieving significant progress. However, these methods face challenges in adapting to natural language questions of varying difficulty and the relevance of the few-shot examples provided. In this paper, we propose an adaptive in-context learning approach with difficulty-aware instruction and retrieval-augmented generation to enhance the performance of Text-to-SQL translation (AID-SQL). First, we introduce adaptive instructions for LLMs, which employ precise difficulty classification to apply difficulty-adaptive generative guidelines and chain of thought (CoT) templates for varying difficulty levels. We automatically incorporate few-shot examples retrieved through the knowledge base into the CoT template to construct CoT-enhanced examples, which improves the capability of LLMs with retrieval-augmented generation (RAG). Furthermore, considering that current RAG methods struggle to effectively measure the contribution of retrieved examples in solving the specific task of Text-to-SQL translation, we train a ranking model that can better bridge the semantic and structural gap between NL questions and SQL queries. This approach can better understand semantic information and allows for retrieving examples that are more beneficial to the final problem-solving. We evaluate our method on five benchmarks. Our method achieves competitive performance compared with existing methods.