You can deploy this model and start earning money today!

derrickzhu/
EllieSQL_Router_Checkpoints

safetensors

2503.22402

mit

Views

No views yet

🐘EllieSQL: Cost-Efficient Text-to-SQL with Complexity-Aware Routing

Checkpoints for the trained routers in the paper "EllieSQL: Cost-Efficient Text-to-SQL with Complexity-Aware Routing".

arXiv

GitHub

Overview

Text-to-SQL automatically translates natural language queries to SQL, allowing non-technical users to retrieve data from databases without specialized SQL knowledge. Despite the success of advanced LLM-based Text-to-SQL approaches on leaderboards, their unsustainable computational costs—often overlooked—stand as the "elephant in the room" in current leaderboard-driven research, limiting their economic practicability for real-world deployment and widespread adoption.

To tackle this, we exploratively propose EllieSQL, a complexity-aware routing framework that assigns queries to suitable SQL generation pipelines based on estimated complexity. We investigate multiple routers to direct simple queries to efficient approaches while reserving computationally intensive methods for complex cases. Drawing from economics, we introduce the Token Elasticity of Performance (TEP) metric, capturing cost-efficiency by quantifying the responsiveness of performance gains relative to token investment in SQL generation. Experiments show that compared to always using the most advanced methods in our study, EllieSQL with the Qwen2.5-0.5B-DPO router reduces token use by over 40% without compromising performance on Bird development set, achieving more than a 2x boost in TEP over non-routing approaches. This not only advances the pursuit of cost-efficient Text-to-SQL but also invites the community to weigh resource efficiency alongside performance, contributing to progress in sustainable Text-to-SQL.

Usage

This repository contains checkpoints for the EllieSQL routers, which are designed to classify the complexity of natural language queries for Text-to-SQL. These models are PEFT (LoRA) adapters, meaning they are fine-tuned on top of a base model. To use them, you will first need to load the corresponding base model and then apply the adapter.

First, ensure you have the necessary libraries installed:

pip install transformers peft torch

Here's an example of how to load and use one of the router checkpoints (e.g., a Qwen2.5-0.5B based router):

python

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from peft import PeftModel, PeftConfig

# The model_id refers to the path on the Hugging Face Hub where the PEFT adapter is located.
peft_model_id = "derrickzhu/EllieSQL_Router_Checkpoints"

# Load the PEFT configuration. This helps identify the base model if it's not explicitly named.
config = PeftConfig.from_pretrained(peft_model_id)

# The base model for these routers is typically Qwen2.5-0.5B, as indicated in the paper and adapter configs.
# Make sure to use the correct Hugging Face Hub ID for the base model.
base_model_name_or_path = "Qwen/Qwen2.5-0.5B" 

# Load the tokenizer for the base model.
tokenizer = AutoTokenizer.from_pretrained(base_model_name_or_path)

# The router performs sequence classification (e.g., simple, medium, complex queries).
# The `config.json` inside the repository indicates `num_labels=3` for this classification task.
model = AutoModelForSequenceClassification.from_pretrained(
    base_model_name_or_path,
    num_labels=3, 
    torch_dtype=torch.bfloat16, # Adjust dtype based on your hardware and desired precision
    low_cpu_mem_usage=True      # Optimize memory usage for large models
)

# Load the PEFT adapter weights onto the base model.
model = PeftModel.from_pretrained(model, peft_model_id)
model.eval() # Set the model to evaluation mode
model.to("cuda") # Move the model to GPU if available for faster inference

# Example natural language query to classify
text_query = "Find the names of all students who scored above 90 in Math and live in New York."

# Tokenize the input query.
inputs = tokenizer(text_query, return_tensors="pt").to(model.device)

# Perform inference to get the classification logits.
with torch.no_grad():
    outputs = model(**inputs)
    logits = outputs.logits
    # Get the predicted class ID (0, 1, or 2).
    predicted_class_id = torch.argmax(logits, dim=-1).item()

# Map the predicted ID back to a human-readable label.
# The `id2label` mapping is typically found in the model's `config.json`.
# For this model, labels are "LABEL_0", "LABEL_1", "LABEL_2".
# Refer to the EllieSQL paper or project documentation for the exact mapping of these labels
# to complexity categories (e.g., LABEL_0 -> "simple", LABEL_1 -> "medium", LABEL_2 -> "complex").
id_to_label = model.config.id2label 
predicted_label = id_to_label[predicted_class_id]

print(f"Text query: '{text_query}'")
print(f"Predicted raw label: {predicted_label}")
# If the exact mapping to complexity is known, e.g.:
# complexity_map = {"LABEL_0": "simple", "LABEL_1": "medium", "LABEL_2": "complex"}
# print(f"Predicted complexity: {complexity_map.get(predicted_label, 'unknown')}")

Citation

If you find our work useful or inspiring, please kindly cite:

bibtex

@misc{zhu2025elliesql,
      title={EllieSQL: Cost-Efficient Text-to-SQL with Complexity-Aware Routing}, 
      author={Yizhang Zhu and Runzhi Jiang and Boyan Li and Nan Tang and Yuyu Luo},
      year={2025},
      eprint={2503.22402},
      archivePrefix={arXiv},
      primaryClass={cs.DB},
      url={https://arxiv.org/abs/2503.22402}, 
}