This is a dataset for training machine learning models to convert natural English language text into Postgres dialect SQL queries.

This dataset comprises 200,000 DPO pairs curated to support the rapid development of text-to-SQL generation models. The uniqueness of this dataset lies in its optimization process. The "chosen" field within each data pair contains SQL queries that have been canonicalized, optimized, and which are chosen from the candidate set which minimizes syntactic cyclomatic and asymptotic complexity against the given schema.

Direct Preference Optimization (see Rafailov et al, 2023) is a novel approach to refinement learning from positive and negative samples to modify the behavior of large-scale unsupervised language models to align with human preferences This method simplifies the fine-tuning process, making it more stable and computationally efficient without the need for extensive hyperparameter tuning or LM sampling, and has been shown to effectively control model outputs, matching or surpassing existing methods.

The source data is cleaned and filtered based on the following criteria:

• Remove queries which are not in English.
• Remove queries which are not valid SQL queries.
• Remove queries which are not executable against the given schema.
• Remove queries which are executed against tables with non-Latin characters.
• Remove queries which use features not supported by the given database.
• Remove long queries which contain domain-specific knowledge which cause model confusion.
• Remove queries which do not fit within a 4096 token context window.

To load the dataset using the HuggingFace datasets library:

from datasets import load_dataset

dataset = load_dataset("zerolink/zsql-postgres-dpo")

To use in model fine-tuning, apply the following chat tokenizer:

tokenizer = AutoTokenizer.from_pretrained(model)

def tokenize(element):
    schema = element["schema"]
    question = element["question"]
    answer = element["chosen"]

    prompt = f"""
    Using the schema:
    {schema}
    Generate SQL for the following question:
    {question}
    """

    system = "Translate English to Postgres SQL."
    message = [
        {"role": "system", "content": system},
        {"role": "user", "content": prompt},
        {"role": "assistant", "content": answer},
    ]
    output = tokenizer.apply_chat_template(
        message, add_generation_prompt=False, tokenize=True
    )
    return {"text": output}

The fields in this dataset are as follows:

This dataset is derived from the following sources:

• datetime - Use of Postgres date and time functions.
• json - Use of Postgres JSON functions.
• math - Use of Postgres math functions.
• postgis - Use of Postgres GIS functions.
• re - Use of Postgres regular expression functions.
• rollup - Use of Postgres rollup functions.
• set - Use of Postgres set functions.
• string - Use of Postgres string functions.
• vector - Use of PGVector functions.
• window - Use of Postgres window functions.

Composition:

This dataset is provided for academic and research purposes. Please adhere to the specified license terms and conditions for usage and distribution.