Generate SQL from plain English with Code Llama

Turn a question written in plain English into a SQL query. You deploy Code Llama Instruct on vLLM and send it your question together with the table definitions. Anyone on your team can then ask for data in plain English.

Before you begin

• Permission to create model endpoints in the project. A Member needs a policy that grants create permission for the Inference service. See Create an endpoint.

• A Hugging Face access token with Read scope, saved as a Hugging Face integration. See Create a Hugging Face integration.

• Python 3 on your computer or in an instance, and the OpenAI client:

  pip install openai

Create the endpoint

Create a vLLM endpoint that downloads the model from Hugging Face, with these values. Leave every other field as the form fills it. For what each field does, see Deploy a model from Hugging Face.

Field	Value
Framework on Choose Framework	vLLM
Endpoint Name	codellama-sql, or a name of your own
Download Type	Download from Huggingface
HuggingFace Model	codellama/CodeLlama-7b-Instruct-hf
HF Token Integration	Your Hugging Face integration
Model Cache Storage	Local drive, 100 GB, which holds the 7B model; choose a larger size for the 34B model
Select Computing Resource	GPU tab, a card with more than 14 GB of GPU memory for the 7B model, and a size under Configuration
Active Worker and Max Workers	1 and 1

For better SQL, enter codellama/CodeLlama-34b-Instruct-hf instead, with a plan that has about 70 GB of GPU memory or more. Select Deploy Endpoint, and wait until the state changes from Deploying to Running, which can take 10 to 15 minutes.

Ask for a query

On the endpoint’s details page, select API Request. In Sample API Request, choose a token under Auth token, or select Generate Token to make one, and copy the Root Endpoint. Set both as variables in your terminal, where <root-endpoint> is the address you copied and <auth-token> is the token:

export ROOT_ENDPOINT="<root-endpoint>"
export AUTH_TOKEN="<auth-token>"

The model writes better SQL when the question comes with the tables. Put the table definitions in the first message and the question in the second:

import os
from openai import OpenAI

client = OpenAI(
    base_url=os.environ["ROOT_ENDPOINT"] + "/v1/",
    api_key=os.environ["AUTH_TOKEN"],
)

SCHEMA = """
Tables, MySQL dialect:
employees(id, name, state, hire_date, salary, department_id)
departments(id, name)
"""

def to_sql(question):
    reply = client.chat.completions.create(
        model="codellama/CodeLlama-7b-Instruct-hf",
        messages=[
            {"role": "system", "content": "Write one SQL query for the question. Use only these tables." + SCHEMA},
            {"role": "user", "content": question},
        ],
        max_tokens=300,
        temperature=0.1,
    )
    return reply.choices[0].message.content

print(to_sql("Find the names of all employees who are from California."))

If you deployed the 34B model, set model to codellama/CodeLlama-34b-Instruct-hf. The reply is a SQL query, often with a short explanation after it.

Try more questions

Describe the tables each time, and name the columns and the time range you mean. Two questions to try with the same tables:

• “Average salary per department, for departments with more than 10 employees.”
• “Employees hired after 2023-01-01, with their department names.”

A reply for the first question looks like this:

SELECT d.name AS department_name,
       AVG(e.salary) AS average_salary,
       COUNT(e.id) AS total_employees,
       MAX(e.salary) AS highest_salary
FROM departments d
JOIN employees e ON d.id = e.department_id
GROUP BY d.name
HAVING COUNT(e.id) > 10;

The exact text varies from run to run.

Another question, “Retrieve customers whose orders total more than 100 in the last month.”, with the tables customers and orders, can give:

SELECT name, email
FROM customers
WHERE id IN (
  SELECT customer_id
  FROM orders
  WHERE total_amount > 100
  AND created_at > DATE_SUB(NOW(), INTERVAL 1 MONTH)
);

Get better results

• Give the full context: the tables, the filters, and the time range.
• Use SQL-style references, such as hire_date > ‘2023-01-01’.
• Define vague words such as “recent”, “large sales”, and “active users”. Say what each one means.
• Name the things you mean, such as departments or customers.

Problem	What to do
The query uses another dialect	Name your database in the first message, for example MySQL or PostgreSQL
A query with many joins is wrong	Split the question into smaller steps and refine the result
The model invents columns	List the table and column names in the first message
The SQL has a syntax error	Check the query in your application before you run it
The data is private	Run the query with a database user that has restricted rights

Read each generated query before you run it against production data.

Check that the endpoint answers

The script prints a query that uses your table and column names. If the call fails, check that the endpoint is Running, and that the model name in the script matches the endpoint’s Hugging Face ID exactly. For other ways to call the endpoint, see Call an endpoint with a token.

A team that asks for data in plain English usually calls the model only during working hours. Set the endpoint to scale to zero when idle so that you pay only while questions arrive.