Leveraging model distillation to fine-tune a model

OpenAI recently released Distillation which allows to leverage the outputs of a (large) model to fine-tune another (smaller) model. This can significantly reduce the price and the latency for specific tasks as you move to a smaller model. In this cookbook we'll look at a dataset, distill the output of gpt-4o to gpt-4o-mini and show how we can get significantly better results than on a generic, non-distilled, 4o-mini.

We'll also leverage Structured Outputs for a classification problem using a list of enum. We'll see how fine-tuned model can benefit from structured output and how it will impact the performance. We'll show that Structured Ouputs work with all of those models, including the distilled one.

We'll first analyze the dataset, get the output of both 4o and 4o mini, highlighting the difference in performance of both models, then proceed to the distillation and analyze the performance of this distilled model.

Prerequisites

Let's install and load dependencies. Make sure your OpenAI API key is defined in your environment as "OPENAI_API_KEY" and it'll be loaded by the client directly.

Loading and understanding the dataset

For this cookbook, we'll load the data from the following Kaggle challenge: https://www.kaggle.com/datasets/zynicide/wine-reviews.

This dataset has a large number of rows and you're free to run this cookbook on the whole data, but as a biaised french wine-lover, I'll narrow down the dataset to only French wine to focus on less rows and grape varieties.

We're looking at a classification problem where we'd like to guess the grape variety based on all other criterias available, including description, subregion and province that we'll include in the prompt. It gives a lot of information to the model, you're free to also remove some information that can help significantly the model such as the region in which it was produced to see if it does a good job at finding the grape.

Let's filter the grape varieties that have less than 5 occurences in reviews.

Let's proceed with a subset of 500 random rows from this dataset.

Let's retrieve all grape varieties to include them in the prompt and in our structured outputs enum list.

Generating the prompt

Let's build out a function to generate our prompt and try it for the first wine of our list.

To get a understanding of the cost before running the queries, you can leverage tiktoken to understand the number of tokens we'll send and the cost associated to run this. This will only give you an estimate for to run the completions, not the fine-tuning process (used later in this cookbook when running the distillation), which depends on other factors such as the number of epochs, training set etc.

Preparing functions to Store Completions

As we're looking at a limited list of response (enumerate list of grape varieties), let's leverage structured outputs so we make sure the model will answer from this list. This also allows us to compare the model's answer directly with the grape variety and have a deterministic answer (compared to a model that could answer "I think the grape is Pinot Noir" instead of just "Pinot noir"), on top of improving the performance to avoid grape varieties not in our dataset.

If you want to know more on Structured Outputs you can read this cookbook and this documentation guide.

To distill a model, you need to store all completions from a model, allowing you to give it as a reference to the smaller model to fine-tune it. We're therefore adding a store=True parameter to our client.chat.completions.create method so we can store those completions from gpt-4o.

We're going to store all completions (even 4o-mini and our future fine-tuned model) so we are able to run Evals from OpenAI platform directly.

When storing those completions, it's useful to store them with a metadata tag, that will allow filtering from the OpenAI platform to run distillation & evals on the specific set of completions you'd like to run those.

Parallel processing

As we'll run this on a large number of rows, let's make sure we run those completions in parallel and use concurrent futures for this. We'll iterate on our dataframe and output progress every 20 rows. We'll store the completion from the model we run the completion for in the same dataframe using the column name {model}-variety.

Let's try out our call model function before processing the whole dataframe and check the output.

Great! We confirmed we can get a grape variety as an output, let's now process the dataset with both gpt-4o and gpt-4o-mini and compare the results.

Comparing gpt-4o and gpt-4o-mini

Now that we've got all chat completions for those two models ; let's compare them against the expected grape variety and assess their accuracy at finding it. We'll do this directly in python here as we've got a simple string check to run, but if your task involves more complex evals you can leverage OpenAI Evals or our open-source eval framework.

We can see that gpt-4o is better a finding grape variety than 4o-mini (12.80% higher or almost 20% relatively to 4o-mini!). Now I'm wondering if we're making gpt-4o drink wine during training!

Distilling gpt-4o outputs to gpt-4o-mini

Let's assume we'd like to run this prediction often, we want completions to be faster and cheaper, but keep that level of accuracy. That'd be great to be able to distill 4o accuracy to 4o-mini, wouldn't it? Let's do it!

We'll now go to OpenAI Stored completions page: https://platform.openai.com/chat-completions.

Let's select the model gpt-4o (make sure to do this, you don't want to distill the outputs of 4o-mini that we ran). Let's also select the metadata distillation: wine-distillation to get only stored completions ran from this cookbook.

Once you've selected completions, you can click on "Distill" on the top right corner to fine-tune a model based on those completions. Once we've done that, a file to run the fine-tuning process will automatically be created. Let's then select gpt-4o-mini as the base model, keep the default parameters (but you're free to change them or iterate with it to improve performance).

Once the fine-tuning job is starting, you can retrieve the fine tuning job ID from the fine-tuning page, we'll use it to monitor status of the fine-tuned job as well as retrieving the fine-tuned model id once done.

Running completions for the distilled model

Now that we've got our model fine-tuned, we can use this model to run completions and compare accuracy with both gpt4o and gpt4o-mini. Let's grab a different subset of french wines (as we restricted the outputs to french grape varieties, without outliers, we'll need to focus our validation dataset to this too). Let's run this on 300 entries for each models.

Let's compare accuracy of models

That's almost a 22% relative improvement over the non-distilled gpt-4o-mini! 🎉

Our fine-tuned model performs way better than gpt-4o-mini, while having the same base model. We'll be able to use this model to run inferences at a lower cost and lower latency for future grape variety prediction.