In this article, I tune a base LLama2 LLM to output SQL code. I use Parameter Efficient Fine-Tuning techniques to optimise the process.
In a previous article, I began to make a case for why you would consider training your own LLM. I also provided a brief introduction to the hardware requirements, as well as methods for optimising the training and inference. In this article, I will cover exactly how to fine-tune an open-source LLM and provide code snippets for you to follow along and reproduce the results. We will tune a Llama2–7B model to provide us with SQL output based on natural language input — in other words, the model will convert a question we ask in natural language:
“How many customers decided to buy eggs in the month of November?”
To a SQL query that fetches the corresponding result:
SELECT COUNT(DISTINCT customer_id) AS num_customers
FROM purchases
WHERE product_name = 'eggs'
AND EXTRACT(MONTH FROM purchase_date) = 11;
In each case, the schema of the database (DB) will be provided as the context for the LLM to work with:
CREATE TABLE purchases (
purchase_id INT PRIMARY KEY,
customer_id INT,
product_name VARCHAR(255),
purchase_date DATE
);
We will be using this dataset in the tuning process. Whilst this article is focussed primarily on achieving the above task, the methodology will be provided in such a way that you can adapt the tuning process to suit your requirements.
In this article I will be using Google Colab to fine-tune the LLM. We will be using the know_sql dataset (OpenRAIL license) that I mentioned previously. We will also be using the axolotl framework to handle the fine-tuning process. They have some great documentation on their GitHub page. Rather than writing the ~100 lines of code to manually handle the fine-tuning process, axolotl allows us to simply edit a YAML config file for the respective model we are looking to fine-tune. I will be running through the exact process in this article but I would suggest reading through the axolotl documentation if anything is unclear.