griffith-bigdata/GRAST-SQL-8B-BIRD-Reranker

1---2license: apache-2.03pipeline_tag: text-ranking4library_name: transformers5tags:6- text-to-sql7- llm8- schema-filtering9- graph-reranker10- qwen311---12 13# GRAST-SQL: Scaling Text2SQL via LLM-efficient Schema Filtering with Functional Dependency Graph Rerankers14 15The **GRAST-SQL** model was introduced in the paper [Scaling Text2SQL via LLM-efficient Schema Filtering with Functional Dependency Graph Rerankers](https://huggingface.co/papers/2512.16083).16 17**Authors:** Thanh Dat Hoang, Thanh Tam Nguyen, Thanh Trung Huynh, Hongzhi Yin, Quoc Viet Hung Nguyen18 19GRAST-SQL is an open-source, LLM-efficient schema filtering framework designed to scale Text2SQL systems to real-world, large databases that often exceed LLM context limits. It compacts Text2SQL prompts by employing a multi-step approach:20(i) ranking columns with a query-aware LLM encoder enriched with values and metadata,21(ii) reranking inter-connected columns via a lightweight graph transformer over functional dependencies, and22(iii) selecting a connectivity-preserving sub-schema with a Steiner-tree heuristic.23 24This framework achieves near-perfect recall and higher precision than existing methods (CodeS, SchemaExP, Qwen rerankers, embedding retrievers), maintains sub-second median latency, scales to schemas with 23,000+ columns, and significantly reduces prompt tokens while often improving accuracy in end-to-end systems.25 26For more details on the project, including training and full evaluation scripts, visit the [GitHub repository](https://github.com/thanhdath/grast-sql).27 28## Sample Usage29 30GRAST-SQL models are often served via `vLLM` for efficient embedding generation, as suggested by the project's GitHub repository. The following example demonstrates how to set up a `vLLM` server and use the model to generate embeddings for text inputs. This is a crucial step for the ranking and filtering pipeline described in the paper.31 32First, ensure you have `vLLM` installed. You can typically install it via pip:33```bash34pip install vllm35```36 37### Step 1: Start the vLLM Server38 39Start the `vLLM` server for the GRAST-SQL model in a separate terminal or background process. This command specifies using `griffith-bigdata/GRAST-SQL-0.6B-BIRD-Reranker` as the model, enabling embedding generation.40 41```bash42CUDA_VISIBLE_DEVICES=0,1 vllm serve griffith-bigdata/GRAST-SQL-0.6B-BIRD-Reranker \43  --port 8000 \44  --max-model-len 8192 \45  --tensor-parallel-size 2 \46  --task embedding \47  --gpu-memory-utilization 0.848```49 50### Step 2: Generate Embeddings using Python51 52Once the `vLLM` server is running, you can connect to it and generate embeddings programmatically:53 54```python55from vllm import LLM, SamplingParams56 57# Ensure the vLLM server is running at the specified port.58# Replace with the actual path to your GRAST-SQL model checkpoint if different.59model_path = "griffith-bigdata/GRAST-SQL-0.6B-BIRD-Reranker"60llm = LLM(61    model=model_path,62    tensor_parallel_size=1, # Adjust based on your GPU setup63    dtype="auto",64    max_model_len=8192,65    enforce_eager=True,66    trust_remote_code=True,67    gpu_memory_utilization=0.8,68    task="embedding", # Essential for embedding models69)70 71# Example texts for which to generate embeddings72text_list = [73    "List all tables related to user activity.",74    "Find columns for product price and description."75]76 77# Generate embeddings78# The `llm.encode` method is used when the vLLM server is started with --task embedding.79embeddings = llm.encode(texts=text_list)80 81for i, text in enumerate(text_list):82    print(f"Text: '{text}'")83    print(f"Embedding shape: {embeddings[i].shape}")84    print(f"First 5 embedding dimensions: {embeddings[i][:5]}85")86 87# These embeddings can then be utilized by the GRAST-SQL framework88# for tasks like column ranking and schema filtering.89```90 91## Datasets92 93The GRAST-SQL framework was evaluated on the following datasets:94-   **Spider**: [Spider Evaluation Dataset](https://huggingface.co/datasets/griffith-bigdata/GRAST-SQL-Spider)95-   **BIRD**: [BIRD Training/Evaluation Dataset](https://huggingface.co/datasets/griffith-bigdata/GRAST-SQL-BIRD)96-   **Spider-2.0-lite**: [Spider 2.0-lite Eval Dataset](https://huggingface.co/datasets/griffith-bigdata/GRAST-SQL-Spider2.0-lite)97 98## Models99 100Other GRAST-SQL models available on the Hugging Face Hub:101-   **GRAST-SQL 0.6B**: [GRAST-SQL 0.6B BIRD](https://huggingface.co/griffith-bigdata/GRAST-SQL-0.6B-BIRD-Reranker)102-   **GRAST-SQL 4B**: [GRAST-SQL 4B BIRD](https://huggingface.co/griffith-bigdata/GRAST-SQL-4B-BIRD-Reranker)103-   **GRAST-SQL 8B**: [GRAST-SQL 8B BIRD](https://huggingface.co/griffith-bigdata/GRAST-SQL-8B-BIRD-Reranker)104 105More models can be found in the [Huggingface collection](https://huggingface.co/collections/griffith-bigdata/grast-sql).106 107## System Flow108 109![GRAST-SQL main flow](https://raw.githubusercontent.com/thanhdath/grast-sql/main/figures/main-flow.png)110 111## Citation112 113If you find this work useful for your research, please cite the paper:114 115```bibtex116@misc{hoang2025scalingtext2sqlllmefficientschema,117      title={Scaling Text2SQL via LLM-efficient Schema Filtering with Functional Dependency Graph Rerankers}, 118      author={Thanh Dat Hoang and Thanh Tam Nguyen and Thanh Trung Huynh and Hongzhi Yin and Quoc Viet Hung Nguyen},119      year={2025},120      eprint={2512.16083},121      archivePrefix={arXiv},122      primaryClass={cs.DB},123      url={https://arxiv.org/abs/2512.16083}, 124}125```