CodeS: Towards Building Open-source Language Models for Text-to-SQL

TL;DR

CodeS introduces open-source 1B–15B language models tailored for text-to-SQL via incremental pre-training on a SQL-centric corpus, achieving state-of-the-art results on Spider and BIRD benchmarks while offering strong robustness and domain adaptation. The approach combines a comprehensive database prompt construction strategy with a bi-directional data augmentation pipeline to enable rapid adaptation to new domains, and demonstrates effectiveness in both few-shot prompting and supervised fine-tuning settings. Extensive experiments across standard benchmarks and real-world domains show CodeS surpasses many baselines in SQL generation accuracy and robustness while delivering efficient, deployable inference. By releasing code, models, and data publicly, the work aims to accelerate research and practical adoption of open-source Text-to-SQL solutions.

Abstract

Language models have shown promising performance on the task of translating natural language questions into SQL queries (Text-to-SQL). However, most of the state-of-the-art (SOTA) approaches rely on powerful yet closed-source large language models (LLMs), such as ChatGPT and GPT-4, which may have the limitations of unclear model architectures, data privacy risks, and expensive inference overheads. To address the limitations, we introduce CodeS, a series of pre-trained language models with parameters ranging from 1B to 15B, specifically designed for the text-to-SQL task. CodeS is a fully open-source language model, which achieves superior accuracy with much smaller parameter sizes. This paper studies the research challenges in building CodeS. To enhance the SQL generation abilities of CodeS, we adopt an incremental pre-training approach using a specifically curated SQL-centric corpus. Based on this, we address the challenges of schema linking and rapid domain adaptation through strategic prompt construction and a bi-directional data augmentation technique. We conduct comprehensive evaluations on multiple datasets, including the widely used Spider benchmark, the newly released BIRD benchmark, robustness-diagnostic benchmarks such as Spider-DK, Spider-Syn, Spider-Realistic, and Dr.Spider, as well as two real-world datasets created for financial and academic applications. The experimental results show that our CodeS achieves new SOTA accuracy and robustness on nearly all challenging text-to-SQL benchmarks.

Figures (5)

• Figure 1: Comparisons between CodeS and SOTA LLMs on two challenging text-to-SQL benchmarks, Spider yu2018@spider and BIRD li2023@bird. While 10x-100x smaller than the existing SOTA LLMs, CodeS achieves comparable or even superior accuracy.
• Figure 2: An example of text-to-SQL in the finance domain.
• Figure 3: Illustration of the comprehensive framework which encompasses: (a) CodeS that is incrementally pre-trained on top of StarCoder using our specially curated SQL-focused dataset. (b) Our unique method for database prompt construction. (c) The proposed bi-directional data augmentation technique for adapting to new domains. CodeS can be employed in two distinct manners: (d) Inferring after a supervised fine-tuning of CodeS on a training dataset, sourced from text-to-SQL benchmarks along with our enriched (question, SQL) pairs. (e) Direct inference through few-shot in-context learning on CodeS.
• Figure 4: A text-to-SQL sample in Spider's training set, consisting of a triplet of <database prompt, question, SQL query>. The database prompt is crafted by our proposed method.
• Figure 5: Prompt formats used in the bi-directional data augmentation. DDL stands for data definition language.