End-to-End Text-to-SQL with Dataset Selection: Leveraging LLMs for Adaptive Query Generation

2025-08-08

AI Summary

To address the overlooked problem of database intent identification in text-to-SQL tasks under multi-database settings, this paper proposes the first end-to-end joint modeling framework that explicitly incorporates a database selection stage into the conventional text-to-SQL pipeline, enabling collaborative optimization of db_id prediction and SQL generation. Methodologically, the approach integrates large language model (LLM)-based semantic understanding, a fine-tuned RoBERTa encoder, structured prompt engineering, and a critic-agent-driven SQL correction mechanism. Crucially, it innovatively leverages LLMs to extract implicit database clues from natural language queries, thereby enhancing the robustness of database intent recognition. Evaluated on standard multi-database benchmarks, the framework achieves significant improvements over existing state-of-the-art models in both database selection accuracy and SQL execution accuracy.

Technology Category

• Data Mining & Knowledge Management: Intelligent Query Processing
• Natural Language Processing: Sentence-level Semantics, Textual Inference, etc.
• Machine Learning: Large Multimodal Models (LMMs)

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Large language models for search
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

Abstract

Text-to-SQL bridges the gap between natural language and structured database language, thus allowing non-technical users to easily query databases. Traditional approaches model text-to-SQL as a direct translation task, where a given Natural Language Query (NLQ) is mapped to an SQL command. Recent advances in large language models (LLMs) have significantly improved translation accuracy, however, these methods all require that the target database is pre-specified. This becomes problematic in scenarios with multiple extensive databases, where identifying the correct database becomes a crucial yet overlooked step. In this paper, we propose a three-stage end-to-end text-to-SQL framework to identify the user's intended database before generating SQL queries. Our approach leverages LLMs and prompt engineering to extract implicit information from natural language queries (NLQs) in the form of a ruleset. We then train a large db_id prediction model, which includes a RoBERTa-based finetuned encoder, to predict the correct Database identifier (db_id) based on both the NLQ and the LLM-generated rules. Finally, we refine the generated SQL by using critic agents to correct errors. Experimental results demonstrate that our framework outperforms the current state-of-the-art models in both database intent prediction and SQL generation accuracy.

Problem

• Identify intended database from natural language queries
• Generate accurate SQL queries using LLMs
• Improve database intent prediction and SQL accuracy

Innovation

• Three-stage end-to-end text-to-SQL framework
• LLM and prompt engineering for ruleset extraction
• RoBERTa-based encoder for db_id prediction