Fine-tune ChatGPT-3.5 and GPT-4 with Weights & Biases

If you use OpenAI's API to fine-tune ChatGPT-3.5, you can now use the WandbLogger integration to track experiments, models, and datasets in your central dashboard with just two lines of code:

from wandb.integration.openai.fine_tuning import WandbLogger

# Your fine-tuning logic

WandbLogger.sync(id=fine_tune_job_id)

See the OpenAI section in the Weights & Biases documentation for full details of the integration.

In this colab notebook, we will be finetuning GPT 3.5 model on the LegalBench dataset. The notebook will show how to prepare and validate the dataset, upload it to OpenAI and setup a fine-tune job. Finally, the notebook shows how to use the WandbLogger.

Imports and initial set-up

Initialize the OpenAI client

You can add the api key to your environment variable by doing os.environ['OPENAI_API_KEY'] = "sk-....".

Import the WandbLogger

Dataset Preparation

We download a dataset from LegalBench, a project to curate tasks for evaluating legal reasoning, specifically the Contract NLI Explicit Identification task.

This comprises of a total of 117 examples, from which we will create our own train and test datasets

Let's look at a few samples.

Format our Data for Chat Completion Models

We modify the base_prompt from the LegalBench task to make it a zero-shot prompt, as we are training the model instead of using few-shot prompting

We now split it into training/validation dataset, lets train on 30 samples and test on the remainder

Save the data to Weights & Biases

Save the data in a train and test file first

Run the OpenAI data validation script

Next, we validate that our training data is in the correct format using a script from the OpenAI fine-tuning documentation

Validate train data

Validate test data

Upload the training and validation data to OpenAI

We will first upload the data to OpenAI. This might take a few minutes depending on the size of your dataset.

Notice the unique ids for both training and validation data. OpenAI uses these ids to access the uploaded data to fine-tune GPT 3.5 on.

Train the model and log to Weights & Biases

Let's define our ChatGPT-3.5 fine-tuning hyper-parameters.

this takes around 5 minutes to train.

Start Weight & Biases Sync

Calling WandbLogger.sync will start polling OpenAI for the fine-tuning job results and log them when they are retrieved, see the docs for how to modify this behaviour

Thats it!

Now your model is training on OpenAI's machines.

Logging the fine-tuning job to W&B is straight forward. The integration will automatically log the following to W&B:

• training and validation metrics (if validation data is provided)
• log the training and validation data as W&B Tables for storage and versioning
• log the fine-tuned model's metadata.

The integration automatically creates the DAG lineage between the data and the model.

You can call the WandbLogger with the job id. The cell will keep running till the fine-tuning job is not complete. Once the job's status is succeeded, the WandbLogger will log metrics and data to W&B. This way you don't have to wait for the fine-tune job to be completed to call WandbLogger.sync.

Calling WandbLogger.sync without any id will log all un-synced fine-tuned jobs to W&B

See the OpenAI section in the Weights & Biases documentation for full details of the integration

The fine-tuning job is now successfully synced to Weights and Biases. Click on the URL above to open the W&B run page. The following will be logged to W&B:

Training and validation metrics

Training and validation data as W&B Tables

The data and model artifacts for version control (go to the overview tab)

The configuration and hyperparameters (go to the overview tab)

The data and model DAG

Run evalution and log the results

The best way to evaluate a generative model is to explore sample predictions from your evaluation set.

Let's generate a few inference samples and log them to W&B and see how the performance compares to a baseline ChatGPT-3.5 model

We will be evaluating using the validation dataset. In the overview tab of the run page, find the "validation_files" in the Artifact Inputs section. Clicking on it will take you to the artifacts page. Copy the artifact URI (full name) as shown in the image below.

The code snippet below, download the logged validation data and prepare a pandas dataframe from it.

We will need to package the data in the dataframe in the format acceptable by GPT 3.5. The format is:

{"messages": [{"role": "system", "content": "some system prompt"}, {"role": "user", "content": "some user prompt"}, {"role": "assistant", "content": "completion text"}]}

For evaluation we don't need to pack the {"role": "assistant", "content": "completition text"} in messages as this is meant to be generated by GPT 3.5.

Run evaluation on the Fine-Tuned Model

Next up we will get the fine-tuned model's id from the logged model_metadata. In the overview tab of the run page, find the "model" in the Artifact Outputs section. Clicking on it will take you to the artifacts page. Copy the artifact URI (full name) as shown in the image below.

Run evaluation and log results to W&B

Calculate the accuracy of the fine-tuned model and log to W&B

Run evaluation on a Baseline model for comparison

Lets compare our model to the baseline model, gpt-3.5-turbo

Calculate the accuracy of the fine-tuned model and log to W&B

And thats it! In this example we have prepared our data, logged it to Weights & Biases, fine-tuned an OpenAI model using that data, logged the results to Weights & Biases and then run evaluation on the fine-tuned model.

From here you can start to train on larger or more complex tasks, or else explore other ways to modify ChatGPT-3.5 such as giving it a different tone and style or response.

Resources

• OpenAI Fine-Tuning Guide
• W&B Integration with OpenAI API Documentation