PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Haoyang Li, Jing Zhang, Hanbing Liu, Ju Fan, Xiaokang Zhang, Jun Zhu, Renjie Wei, Hongyan Pan, Cuiping Li, Hong Chen

arXiv:2402.16347v1cs.CLcs.DB

TL;DR

Text-to-SQL SOTA systems often rely on closed-source LLMs with opaque details, privacy risks, and high inference overheads. CodeS introduces fully open-source 1B–15B models using SQL-centric incremental pre-training, strategic prompts, and bidirectional augmentation. The paper reports comparable or superior accuracy at much smaller sizes and new SOTA accuracy and robustness on nearly all challenging benchmarks.

  • Problem

    Most SOTA text-to-SQL approaches rely on closed-source LLMs, limiting transparency, risking data privacy through APIs, and creating substantial inference overheads.

  • Method

    CodeS incrementally pre-trains StarCoder on a curated SQL-centric corpus and uses strategic schema prompts plus bidirectional data augmentation for domain adaptation.

  • Results

    CodeS models range from 1B to 15B parameters and achieve comparable or superior accuracy to much larger SOTA LLMs, with new SOTA accuracy and robustness on nearly all challenging benchmarks.

  • Takeaways & Limitations

    Fully open-source CodeS models provide a smaller-scale alternative for text-to-SQL applications, with code, models, and data released publicly.

  • Takeaways & Limitations

    The best SFT CodeS-15B model requires at least 35GB of GPU memory for inference, exceeding the 12GB GPUs used by Spider’s submission platform.

Abstract

from arXiv · show

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.

1 INTRODUCTION

CodeS addresses the reliance of text-to-SQL on closed-source LLMs by providing smaller, open-source models designed for SQL generation. The paper combines incremental pre-training, prompt construction, and bidirectional augmentation to improve capability, schema linking, and domain adaptation.

  • Motivation: Closed-source SOTA text-to-SQL models obscure architectures, risk data privacy through API calls, and incur substantial inference overheads.These limitations motivate an open-source alternative for real-world applications.
  • Results: CodeS is 10x-100x smaller than existing SOTA LLMs while achieving comparable or superior accuracy on Spider and BIRD.Figure 1 compares CodeS with SOTA LLMs on these two challenging benchmarks.
  • CodeS: CodeS is a fully open-source language-model series built on StarCoder, spanning 1B to 15B parameters.The range lets users select a model according to available computational resources.
  • Challenges and Solutions: Strategic schema prompts filter relevant tables, columns, and values, while bidirectional data augmentation supports new-domain adaptation with limited annotation overhead.The schema strategy addresses numerous tables, wide tables, ambiguous names, and large tables; augmentation creates SQL-to-question and question-to-SQL data.
  • Results: Extensive evaluations report that CodeS surpasses notable open-source pre-trained models and achieves new SOTA accuracy and robustness on almost all challenging text-to-SQL benchmarks.The paper also states that its code, models, and data are open-sourced.

2 RELATED WORK

Prior text-to-SQL work includes supervised encoder-decoder modeling, schema- and grammar-aware architectures, database-related pre-training, prompting, and code language models. CodeS differs by directly targeting SQL generation through language-model pre-training.

  • Supervised Fine-Tuning-Based Text-to-SQL: Before LLMs, text-to-SQL commonly used supervised fine-tuning of encoder-decoder neural networks.Later sequence-to-sequence formulations flattened questions and database information into inputs paired with target SQL queries.
  • Supervised Fine-Tuning-Based Text-to-SQL: Graph-relational methods enhance encoder representations using structural relationships among query tokens, tables, and columns.These approaches incorporate database and query structure into the encoder.
  • Supervised Fine-Tuning-Based Text-to-SQL: Some methods inject SQL grammar into the decoder to constrain its output space and ensure syntactically correct SQL queries.The grammar is used as a decoding constraint rather than as a direct generation objective.
  • Pre-training: Database-related pre-training studies primarily improve encoders’ representations of questions and schemas, whereas CodeS directly enhances SQL generation capability.The related work passage explicitly contrasts these goals with CodeS.
  • Prompting-Based Text-to-SQL: Prompting-based methods achieve SOTA benchmark performance but face implementation challenges from API costs and potential data-privacy concerns.This motivates research into open-source alternatives for real-world deployment.

3 PRELIMINARIES

Text-to-SQL generates an executable SQL query from a natural-language question and a database. The database input includes schema and metadata used by the parser to produce the query.

  • Text-to-SQL Task: The text-to-SQL objective is to generate a SQL query S from a natural-language question Q and database D that can execute to answer the question.The formulation treats executable query generation as the task objective.
  • Text-to-SQL Task: The parser interprets Q using D and produces S.This identifies the input-output role of the parser in the task formulation.
  • Text-to-SQL Task: D contains database schema and metadata, including tables, columns, types, comments, values, primary keys, and foreign-key relations.These elements provide the database information available to the parser.

4 OVERVIEW

CodeS combines incremental pre-training, database prompt construction, and bi-directional data augmentation into a compact open-source text-to-SQL framework. It supports both supervised fine-tuning and few-shot in-context inference, with retrieval improving online prompt processing.

  • CodeS integrates SQL-focused incremental pre-training, database prompt construction, and bi-directional data augmentation for compact text-to-SQL modeling.The framework also supports supervised fine-tuned inference and direct few-shot in-context learning.
  • The framework supports supervised fine-tuning when abundant training data is available and few-shot in-context learning when data is limited.
  • A demonstration retriever selects examples using both question similarity and question-pattern similarity for in-context inference.
  • BM25 indexing accelerates the value retriever’s online processing, while the schema filter uses a compact neural network for fast inference.

5 INCREMENTAL PRE-TRAINING

CodeS is incrementally pre-trained from StarCoder on a mixed corpus spanning SQL, natural language, and natural-language-to-code data. The training objective maximizes the likelihood of each token in the input sequence.

  • 5.1 Pre-training Corpus: The pre-training corpus contains 11GB of SQL-related data, 4.5GB of natural-language data, and 6GB of natural-language-to-code data.
  • 5.1 Pre-training Corpus: The NL-to-code component includes automatically collected pairs, instruction-following data, structured notebooks, and the authors’ 458K-pair NL-SQL dataset.
  • 5.2 Pre-Training Details: CodeS incrementally pre-trains StarCoder for two epochs on SQL-related data and one epoch each on NL-related and NL-to-code data.
  • 5.2 Pre-Training Details: The objective maximizes the likelihood of the entire token sequence by multiplying conditional probabilities for successive tokens.
  • 5.2 Pre-Training Details: Training uses AdamW with β1 = 0.9, β2 = 0.95, ϵ = 10^-8, a 5e-5 learning rate, cosine decay, and 4M-token batches.

6 DATABASE PROMPT CONSTRUCTION

CodeS constructs database prompts by filtering relevant schema elements, retrieving useful values, and adding metadata such as types, comments, representative values, and key relations. These choices reduce prompt length while preserving information for schema linking and SQL generation.

  • 6 DATABASE PROMPT CONSTRUCTION: The prompt-construction pipeline combines a schema filter, a value retriever, and database metadata.Algorithm 1 specifies the overall construction process.
  • 6.1 Schema Filter: The schema filter minimizes long prompts by retaining relevant tables and columns when databases exceed the model’s context length.
  • 6.1 Schema Filter: The resulting schema filter reduces database-prompt length and alleviates the model’s schema-linking burden.
  • 6.2 Value Retriever: The value retriever can identify question-matched values such as district.a2 = ’Jesenik’ and insert them into the database prompt.
  • 6.3 Database Metadata: Prompt metadata includes column data types and comments to clarify permissible operations and ambiguous schema terminology.Comments are incorporated into schema-item classification and database prompts.
  • 6.3 Database Metadata: Representative values expose formats and categorical options, helping generate predicates such as client.gender = ’F’ and year-based date conditions.
  • 6.3 Database Metadata: Primary and foreign keys guide join-path inference and JOIN ON clause generation.The prompt represents key relationships explicitly for the language model.
  • 6.3 Database Metadata: Figure 4 presents a database-prompt, question, and SQL triplet, while Figure 5 shows prompt formats using DDL for bi-directional augmentation.

7 NEW DOMAIN ADAPTION

CodeS adapts to new domains by generating authentic question–SQL pairs from a small set of manually annotated examples. Its bi-directional augmentation combines question-to-SQL synthesis with schema-based SQL-to-question generation and question refinement.

  • The method targets new databases where labeled training data is scarce by using minimal annotation to build domain-specific training pairs.
  • In question-to-SQL augmentation, GPT-3.5 first synthesizes questions from real examples and then generates corresponding SQL queries.
  • Shuffling demonstrations and using high-temperature generation promotes diversity while preserving user-intent patterns.
  • SQL-to-question augmentation extracts benchmark SQL templates, fills them with a new schema, and uses GPT-3.5 to rephrase artificial questions naturally.

8 USAGE OF CODES

CodeS supports supervised fine-tuning and few-shot in-context learning for generating SQL from database prompts and natural-language questions. Its demonstration retriever emphasizes question structure rather than surface entities.

  • Supervised Fine-Tuning: CodeS converts databases into database prompts and represents training samples as database-prompt, question, and SQL triplets.Test samples become a database prompt and question pair before SQL generation.
  • Supervised Fine-Tuning: Supervised fine-tuning combines the database prompt with the question and optimizes CodeS to predict the desired SQL query.After fine-tuning, the refined model generates SQL from the combined test inputs.
  • Few-Shot In-Context Learning: Without fine-tuning, few-shot learning uses retrieved demonstrations to leverage CodeS’s built-in text-to-SQL capabilities.The approach is intended for cases where fine-tuning is impractical.
  • Few-Shot In-Context Learning: The question-pattern-aware demonstration retriever removes entities and uses SimCSE to compare question structures before selecting examples.This avoids retrieving demonstrations primarily because they share entities with the test question.

9 EXPERIMENTS

The experiments evaluate CodeS through in-context learning, supervised fine-tuning, and robustness-oriented benchmarks, using execution- and test-suite-based metrics. Results show strong performance after incremental pre-training and fine-tuning across Spider and BIRD.

  • Experimental Setup: CodeS is evaluated on Spider, BIRD, four robustness-diagnostic benchmarks, and manually created financial and academic databases.Spider supplies the main training set for evaluations on its robustness variants.
  • Evaluation Metrics: Execution accuracy checks matching database results, whereas test-suite accuracy evaluates consistency across augmented database instances to reduce false positives.The latter addresses cases where incorrect SQL happens to produce the correct output on one database instance.
  • Implementation Details: CodeS experiments use SQLite, schema-item classifiers evaluated by AUC, beam search with four SQL candidates, and eight NVIDIA A800 80GB GPUs.The classifier results report consistently higher AUC scores for Spider than BIRD and BIRD with external knowledge.
  • Few-Shot Evaluation: Incremental pre-training substantially improves StarCoder’s SQL generation in few-shot evaluations using consistent prompts and 1-, 3-, and 5-shot settings.The comparison is conducted on Spider and BIRD development sets.
  • Supervised Fine-Tuning: SFT CodeS-3B outperforms the leading GPT-4-based method on Spider development, while SFT CodeS-7B and 15B achieve new SOTA performance.CodeS-7B has a marginal advantage over CodeS-15B, suggesting possible overfitting of the larger model to Spider training data.

9.4 Evaluation on Robustness Benchmarks

CodeS is tested on Spider variants and Dr.Spider to assess robustness under distribution shifts and perturbations. It generally outperforms prior systems, while retaining some weaknesses on database-content and SQL perturbations.

  • Spider Variants: 70.0% on Spider-Syn, 77.2% on Spider-Realistic, and 72.0% on Spider-DK exceed the best baselines for SFT CodeS-7B.The corresponding gains are 2.6%, 4.0%, and 4.5%; SFT CodeS-3B also exceeds previous SOTA methods across all three datasets.
  • Spider Variants: Training on Spider while testing on its variants demonstrates generalization under challenging distribution shifts.The reported robustness gains occur without training directly on the evaluated variants.
  • Dr.Spider: 74.3% and 75.2% on Dr.Spider’s natural-language perturbations surpass ChatGPT + ZeroNL2SQL’s 73.2%.CodeS slightly lags on database perturbations and SQL perturbations, while its global average slightly surpasses the prior best.
  • Ablation Studies: Removing pattern-aware retrieval, schema filtering, or value retrieval reduces performance, with value retrieval especially important on BIRD.Schema filtering also speeds generation by shortening input sequences.
  • Ablation Studies: Comments, representative database values, and primary/foreign keys materially support performance, whereas column data types have only minor impact.Comments matter particularly on BIRD’s ambiguous schemas, while keys help generate accurate JOIN ON clauses.

9.6 Evaluation on Real-World Scenarios

CodeS is evaluated on financial and academic databases and presented as a deployment-oriented system with pathways for zero, limited, or augmented annotations. The 7B model balances performance and efficiency for these scenarios.

  • Real-World Datasets: Bank-Financials challenges schema linking through many columns and ambiguous names, while Aminer-Simplified challenges complex table-join relationships.These datasets represent financial and academic real-world applications.
  • Evaluation Setup: CodeS-7B uses a BIRD-trained schema-item classifier to filter schemas on new databases and is compared with GPT-3.5 prompting baselines.Human evaluation supplements execution accuracy because exact-match execution can reject semantically valid SQL.
  • Deployment Pathways: Fine-tuned checkpoints can transfer to new databases without annotations, while three-shot CodeS-7B adapts without parameter tuning when annotations or resources are limited.The paper also evaluates augmented-data fine-tuning and a unified model trained on merged datasets.
  • Evaluation Caveat: Execution and human-evaluation scores differ substantially because benchmark and new-dataset annotation habits are different.This gap affects interpretation of real-world SQL accuracy.
  • Domain Adaptation: Merged training data prevents performance drops and improves results especially on Aminer-Simplified, while separate per-database fine-tuning has substantial overheads.This supports a unified text-to-SQL model for multiple domains.
  • Efficiency and Deployment: 0.6, 0.9, 1.1, and 1.5 seconds are the inference times for CodeS-1B, 3B, 7B, and 15B, versus approximately 60 seconds for DIN-SQL+GPT-4.Float16 GPU memory requirements are 10GB, 13GB, 20GB, and 35GB, respectively, enabling local deployment at suitable sizes.

10 CONLUSION

CodeS provides specialized open-source language models and supporting data and methods for text-to-SQL development. Extensive evaluations report new state-of-the-art accuracy and robustness across challenging benchmarks.

  • Automatic and human evaluations cover the Bank-Financials and Aminer-Simplified datasets.
  • CodeS gives developers a range of specialized pre-trained language models for building text-to-SQL applications.
  • The authors open-source their collected SQL-focused corpus, code, models, and data to support further research and innovation in SQL generation.
  • The approach combines database prompt construction with bi-directional data augmentation to support versatility and adaptation across domains.
  • CodeS is reported as a new state-of-the-art pre-trained language model for SQL generation, while SFT CodeS models achieve new state-of-the-art accuracy and robustness across many benchmarks.
Loading 2402.16347v1…