Defog SQLCoder

repository·main·Indexed 26 days ago

https://github.com/defog-ai/sqlcoder

A family of large language models designed to convert natural language questions into accurate SQL queries. The project includes a CLI for launching a web interface, support for various hardware configurations (NVIDIA GPU, Apple Silicon, and CPU via llama-cpp), and integration with the transformers library for inference. It features specific prompt structures for database schemas and provides models like SQLCoder-7b-2 and SQLCoder-34B.

Tokens
2.3K
Snippets
10
Records
14
Agent score
87%

What's inside SQLCoder

  1. Understand the SQLCoder prompt structure

    main

    SQLCoder uses a specific prompt template to transform natural language questions into SQL queries. The prompt consists of four main sections: a Task definition, Instructions for handling unknown questions, the Database Schema (provided as a metadata string), and an Answer section that guides the model to output the SQL within [SQL] tags.

    To use this effectively, you must provide the {user_question} and the {table_metadata_string} representing your database schema.

  2. Install SQLCoder

    main

    Installation commands for SQLCoder depend on your hardware and operating system:

    • NVIDIA GPU (>16GB VRAM): Best performance.
    • Apple Silicon: Uses llama-cpp with Metal support (note: performance may be lower due to quantization and lack of beam search).
    • Non-Apple Silicon (No GPU access): Use llama-cpp with OpenBLAS on Linux or Intel Mac.
    • Windows (No GPU access): Use llama-cpp with OpenBLAS via PowerShell.
  3. Load SQLCoder-7b-2 with optimized precision

    main

    Load the defog/sqlcoder-7b-2 model using AutoModelForCausalLM. The loading strategy depends on your available GPU VRAM:

    • > 15GB VRAM: Load in float16 using torch_dtype=torch.float16.
    • 8GB - 15GB VRAM: Load in 8-bit using load_in_8bit=True.
    • 5GB - 8GB VRAM: Load in 4-bit (requires additional configuration not shown in this snippet).

    Always set trust_remote_code=True and device_map="auto".

    model_name = "defog/sqlcoder-7b-2"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    if available_memory > 15e9:
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            trust_remote_code=True,
            torch_dtype=torch.float16,
            device_map="auto",
            use_cache=True,
        )
    else:
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            trust_remote_code=True,
            load_in_8bit=True,
            device_map="auto",
            use_cache=True,
        )
  4. Format the SQLCoder prompt with database schema

    main

    SQLCoder requires a specific prompt structure to generate accurate SQL. The prompt must include:

    1. Task: A description of the goal.
    2. Instructions: Specific rules (e.g., how to calculate revenue).
    3. Database Schema: The DDL (e.g., CREATE TABLE statements) representing your database.
    4. Answer Section: A concluding instruction that wraps the question in [QUESTION]{question}[/QUESTION] tags and ends with a [SQL] marker.
    prompt = """### Task
    Generate a SQL query to answer [QUESTION]{question}[/QUESTION]
    
    ### Instructions
    - If you cannot answer the question with the available database schema, return 'I do not know'
    - Remember that revenue is price multiplied by quantity
    - Remember that cost is supply_price multiplied by quantity
    
    ### Database Schema
    This query will run on a database whose schema is represented in this string:
    CREATE TABLE products (...);
    
    ### Answer
    Given the database schema, here is the SQL query that answers [QUESTION]{question}[/QUESTION]
    [SQL]
    """
  5. Launch SQLCoder with automatic model downloading

    main

    Running sqlcoder launch automates the setup process based on your hardware:

    1. Non-GPU Machines: It checks for an NVIDIA GPU using lspci. If none is found, it downloads the sqlcoder-7b-q5_k_m.gguf file (~5GB) from the defog/sqlcoder-7b-2 repository to ~/.defog/.
    2. GPU Machines: It downloads the full defog/sqlcoder-7b-2 model snapshot (~14GB) from Hugging Face.
    3. Execution: Once the model is available, it launches both the static asset server and the webserver processes simultaneously.

    To exit the application, press Ctrl+C.

    sqlcoder launch
  6. Use SQLCoder for Inference

    main

    You can perform inference using the transformers library by downloading the model weights from Hugging Face. To run inference on a sample database using the provided inference.py script, use the -q flag to pass your natural language question.

    python inference.py -q "Question about the sample database goes here"
  7. SQLCoder License Information

    main

    SQLCoder has two distinct licenses:

    • Code: Apache-2 license.
    • Model Weights: CC BY-SA 4.0 license.

    Commercial Use: You may use and modify the model for any purpose, including commercial use. However, if you modify the weights (e.g., via fine-tuning), you must open-source your modified weights under the same CC BY-SA 4.0 license terms.

  8. SQLCoder Hardware Requirements

    main

    Hardware needs vary by model size and quantization:

    • SQLCoder-34B (float16): Tested on a 4xA10 GPU.
    • Quantized versions (8-bit/4-bit): Can run on consumer GPUs with 20GB or more of VRAM (e.g., RTX 4090, RTX 3090) or Apple Silicon (M2 Pro, M2 Max, M2 Ultra) with 20GB or more of memory.
  9. Generate SQL queries using `generate_query`

    main

    Use the following pattern to generate formatted SQL from a natural language question. The function tokenizes the prompt, generates the response using the model, clears the CUDA cache to prevent memory crashes (crucial for Colab), and uses sqlparse to reindent the resulting SQL.

    import sqlparse
    
    def generate_query(question):
        updated_prompt = prompt.format(question=question)
        inputs = tokenizer(updated_prompt, return_tensors="pt").to("cuda")
        generated_ids = model.generate(
            **inputs,
            num_return_sequences=1,
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.eos_token_id,
            max_new_tokens=400,
            do_sample=False,
            num_beams=1,
        )
        outputs = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
    
        torch.cuda.empty_cache()
        torch.cuda.synchronize()
        return sqlparse.format(outputs[0].split("[SQL]")[-1], reindent=True)
    
    question = "What was our revenue by product in the New York region last month?"
    generated_sql = generate_query(question)
    print(generated_sql)
  10. Use the SQLCoder CLI

    main

    The sqlcoder CLI provides commands to launch the full application, serve the backend webserver, or serve the static frontend assets.

    Usage:

    sqlcoder <command>

    Available Commands:

    • sqlcoder launch: Downloads the necessary model files (GGUF for non-GPU machines or full model for GPU machines) and starts both the static asset server and the webserver.
    • sqlcoder serve-webserver: Starts the backend webserver using Uvicorn on localhost:1235.
    • sqlcoder serve-static: Starts a local HTTP server to serve static assets on port 8002.
    sqlcoder launch
    sqlcoder serve-webserver
    sqlcoder serve-static