QueryCraft (by SuperKnowa)

Welcome to the QueryCraft repository, your comprehensive solution for fine-tuning Large Language Models (LLMs) for the task of generating SQL queries from natural language (Text2SQL, Text2GraphQL, NL2Query). This framework is designed to create a simple higher abtraction framework so that AI engineerrs can quickly construct full GenAI pipeline without having to worry about inner complexittes of full fine-tuning.

Enterprise Text2SQL applications are complex and need lot more than just plugging in LLMs. The specific metadata make it even more challenging case of fine-tuning. 'QueryCraft' streamline the process of adapting LLMs for Text2SQL tasks, providing a robust framework and pipeline that accelerates the initiation and fine-tuning process. Whether you're a developer or a expert data scientist keen on harnessing the power of LLMs for database querying, QueryCraft offers the tools and infrastructure to make your journey smoother and more efficient.

Cover Blog- https://medium.com/@KunalSavvy/querycraft-the-simplest-and-most-comprehensive-open-framework-yet-to-build-nl2sql-at-scale-474b193e902e

---

What QueryCraft comes with

An easy frameworrk goverrned by relational databases using natural language questions instead of writing complex SQL queries

• A configurable TextToSQL pipeline on your Database
• A tool to evaluate the Text To SQL pipeline on your Database
• Script to correct the LLM-generated queries
• HIGHER abstraction framework to fine-tune an LLM for SQL query generation on your database
• Tools to speed up inferencing for SQL query generation
• Analytics Dashboard to analyze the quality of generated queries

You need the following to run the QueryCraft pipeline:

• Database to run SQL query on and to evaluate the generated queries

• Instruct set for fine tuning (you can curate instruct it with Annotation tool)

• 

---

Benchmark Results

We conducted various experiments on the Spider development dataset using different pre-trained and fine-tuned language models (LLM) in the QueryCraft pipeline. Below are the highest results achieved with their respective LLMs. The highest result, 70.57% execution accuracy, was obtained with the DB-Chat SQL model, followed by 69% execution accuracy achieved by the Granite 20B code instruction-based model.

Outcomes:

1. QueryCraft evaluation module applies post processing to the generated SQL queries and significantly improve the execution accuracy by correcting the SQL queries. This can be observed for all the models used in the QueryCraft pipeline as shown below:

   

2. Smaller fine-tuned models outperforms some larger ones pretrained models using the QueryCraft pipeline. It can be observed from the chart below QueryCraft pipeline can finetune a smaller model (7B) to outperform larger models of 34B or even 70B parameters for Text To SQL task.

   

Features/Modules in QueryCraft

The QueryCraft pipeline is built of 8 modules/components.

1. Instruct dataset curation: This is a UI tool to curate instruct dataset for finetuning or evaluation of TextToSQL task. Read details about this module on our blog on Golden dataset curation for text2sql.

   

2. Data ingestion: Data ingestion module can be used for ingesting your data to Db2 on IBM Cloud.

   

3. Context Retriever: Context retriever can be used to extract context from the instruct dataset given the instruct dataset. Read details about this module on our blog on Diving Deep into Context Retrieval Techniques

   

4. Finetune: Finetune service can be used for PEFT (Parameter Efficient Fine Tuning) for Text To SQL task. Read details about this module on our blog on Text-to-SQL: A Deep Dive into Fine-Tuning Strategies

   

5. Inference: Inference service can be used to generate SQL queries either using a fine-tuned model or a pre-trained model on your NL-SQL pairs dataset. Read details about this module on our blog on LLM inference text2sql

   

6. Query Correction: Inference service can be used to fix incorrect queries generated by LLM.

Read details about this module on our blog on Improving Text-To-SQL Accuracy with Query Correction Service

7. Evaluation: Evaluation service consists of two modules - Query correction and Execution evaluation. Read details about our Evaluation framework on our blog Text To SQL Evaluation Framework. Execution evaluation service can be used to calculate the execution accuracy on your database. Read details about our Evaluation framework on our blog Text To SQL Execution Accuracy

   

8. Query Analysis dashboard Query Analysis dashboard service can be used to analyze the results of generated queries and compare output of multiple LLMs. Read details about our Evaluation framework on our blog Exploring LLM performance in Text-To-SQL generation with Query Analysis Dashboard

   

Getting Started

To get started with SuperKnowa-QueryCraft, clone this repository to your local machine:

git clone https://github.com/ibm-ecosystem-engineering/SuperKnowa-QueryCraft.git
cd SuperKnowa-QueryCraft

Installation

Before diving into fine-tuning or inference, ensure that your environment is set up with all the necessary dependencies:

pip install -r requirements.txt

Configuration Files

Configure your environment and services by editing the simpleConfig.ini and expertConfig.ini files. You have two options of configuring the QueryCraft pipeline

• simpleConfig: Use this to specify the paths for datasets, models, and other services as per your setup.
• expertConfig: Use this to configure the parameters of your experiment including fine-tuning.

SuperKnowa-QueryCraft provides the capability to run the whole pipeline (Context Retriever -> Fine-tuning -> Inference -> Query Correction -> Evaluation -> Query Analysis dashboard) together and also you can run each component individually.

Step by Step instructions

You need a GPU environment for running and fine-tuning LLM using QueryCraft framework. Read Setting up environment for more details. Please read the detailed documentation for step by step instructions

Step 0. Instruct dataset

There are three options for using your dataset to finetune/evaluate the Text to SQL (QueryCraft) pipeline:

1. Bring your dataset with golden queries in the following format: question, query, and db_id. Instruction for ingesting the dataset is provided in the next Step 1.
2. Curate the golden query dataset using our annotation tool: https://annotator.superknowa.tsglwatson.buildlab.cloud/
3. Use the example datasets provided below for testing: Spider and KaggleDBQA

Please read the Step 0. Golden Query Dataset Annotation for step by step instructions for the Instruct dataset.

Step 1. Data Ingestion

You have 3 options for Data Ingestion.

1. Bring Your Own Data

• If you have both databases and instruct set (golden query)
• If you only have database and not instruct set then use above annotation tool

2. Use the example set

• This comes with both source dataset and Instruct Db

Please read the Step 1. Data Ingestion for step by step instructions for the Instruct dataset.

Run the Data Ingestion module of the QueryCraft pipeline using the runQueryCraft.sh, file with the dataIngestion option after setting the simpleConfig.ini file to insert salary.csv into the a table in db2.

sh runQueryCraft.sh

Provide the option:

dataIngestion

Step 2. Context Retriever

Please read the Step 2. Context Retriever for step by step instructions for the Instruct dataset.

Execute the context retriever using the following command.

sh runQueryCraft.sh

Provide the option:

contextRetriever

Step 3. Fine-Tuning

Please read the Step 3. Finetuning for step by step instructions for the Instruct dataset.

To start fine-tuning your LLM for the Text to SQL task, run the below command.

sh runQueryCraft.sh

Provide the option:

finetune

Follow the prompts to specify your dataset and model configuration.

Step 4. Inference

Please read the Step 4. Inference for step by step instructions for the Instruct dataset.

To generate SQL queries using your fine-tuned or pre-trained model, execute:

sh runQueryCraft.sh

Provide the option:

inference

Step 5. Query Correction

Please read the Step 5. Query Correction for step by step instructions for the Instruct dataset.

sh runQueryCraft.sh

Provide the option:

querycorrection

Step 6. Evaluation

Please read the Step 6. Evaluation for step by step instructions for the Instruct dataset.

Evaluate the performance of your model against the SQLite database or DB2 by running the below command:

sh runQueryCraft.sh

Provide the option:

evaluation

Step 7. Query Analysis Dashboard

Please read the Step 7. Query Analysis for step by step instructions for the Instruct dataset.

For a visual analysis of your fine-tuning experiments and generated SQL queries, launch the streamlit dashboard:

sh runQueryCraft.sh

Provide the option:

queryanalysisDashboard

 cd code
 streamlit run streamlit_query_analysis_dashboard.py --server.port 8502 --server.fileWatcherType none

Step 8. Run pipeline (all)

Please read the Step 8. Run Full Pipeline for step by step instructions for the Instruct dataset.

To run all components together, you can change the required parameters in ssimpleConfig.ini. You must set the default path as shown in the designated section below.

[Finetune]
train_dataset =${Default:home_dir}input/datasets/${Default:exp_name}_contextRetriever.csv

[Inference]
input_dataset = input/datasets/${Default:exp_name}_validSet.csv


[QueryCorrection]
input_dataset = ${Default:home_dir}output/inference/${Default:exp_name}_inference.csv

[EXEvaluator]
input_dataset = ${Default:home_dir}output/inference/${Default:exp_name}_inference.csv

Run the below command:

sh runQueryCraft.sh

Provide the option:

all

Build Lab Team

Created & Architected By

• Kunal Sawarkar, Distinguished Engineer & Chief Data Scientist

Builders

• Shivam Solanki
• Kevin Huang
• Abhilasha Mangal
• Amit Khandelwal
• Sahil Desai
• Himadri Talukder
• Tyler Burns

Contact

• This framework is built and tested on top of IBM watsonx platform. In case of any questions or help, you can reach out to Kunal Sawarkar (kunal@ibm.com). For any issues, please create a GH issue.

License

This project is licensed under the Apache-2.0 license - see the LICENSE file for details. Copyright @ 2024 IBM Corporation.