Fine-tuning gpt-3.5-turbo to learn to play "Connections"

I started playing the NYTimes word game "Connections" recently, by the recommendation of a few friends. It has the type of freshness that Wordle lost for me a long time ago. After playing Connections for a few days, I wondered if an OpenAI language model could solve the game (the objective is to group the 16 words into 4 categories of 4 words). I tried with gpt-4-32k and gpt-4-1106-preview, tweaking prompts for a few hours and wasn't able to make much progress. It's certainly possible prompt engineering alone could solve this problem, but it wasn't easy for me for find a path forward. I imagine it will involve a bit of creativity. I decided this was as good a time as any to try and fine tune a model to do a thing I couldn't easily get it to do with prompts.

Getting the dataset

I remembered seeing at some point that Wordle had an API to return the day's word and it does at https://www.nytimes.com/svc/wordle/v2/yyyy-mm-dd.json. I figured Connections might have a similar JSON API and it does. I wrote a script to get all the solutions to past game and put the in the connections_data folder with the naming convention yyyy-mm-dd.json.

It turns out the first game was published on 2023-06-12. Next, I wrote some not pretty code to create a jsonl file to upload to OpenAI. They describe the structure of the file to fine-tune the gpt-3.5-turbo model here.

Estimating Price

I estimated the price of the job by counting the number of tokens in my fine-tune file. The OpenAI pricing page lists the price of the fine-tune at $0.0080 / 1K tokens. Additionally, the fine-tuning guide notes

To estimate the costs for a specific fine-tuning job, use the following formula:

For a training file with 100,000 tokens trained over 3 epochs, the expected cost would be ~$2.40 USD.

Using tiktoken to count tokens

which output

Running the fine-tune job

I took 15 of the rows (number chosen pretty randomly) from the result and separated them out into a training file, so that OpenAI can used it to calculate loss for the fine-tuning (I think). I called these files connections_prompts_train.jsonl and connections_prompts_test.jsonl. I went to OpenAI and created a new fine-tuning job. I selected the gpt-3.5-turbo-1106 model then uploaded my training and test data, and clicked "create" and off it went. As the job ran, I could see this dashboard. Here's what it looked like when it finished.

!Fine-tuning progress. A loss graph with two lines, slowly decreasing starting around 3 and down to around 0.2

And it cost \$0.90. According to the fine-tune metrics, 3 epochs (given my estimated cost of \$0.37) as mentioned fine-tune training guide though I don't recall setting that or seeing it anywhere else.

!Fine-tuning cost screenshot

Testing the fine-tune

With the fine-tuning done, I went to the playground to try it out. I fetched the words for the next day's puzzle, which was mercifully already available. I forgot to make a validation set. Shame on me. It would have take a bit longer to fine-tune again, and I had a lot of anticipation to see the results, it being my first fine-tune and not really understanding the loss graph or knowing if it worked.

I copied the same system prompt from the training

then added the user message with the words

and clicked submit. The model outputted

A definite bummer. I did a quick search in my project directory for "words: level members", and found it everywhere. Oops! I built the dataset improperly.

Fixing a dataset bug

After a bit of a reset, I found some issues in the prompt generation python script and fixed them here

I checked my data and validated the correct words for each category were in the output jsonl file. This time around, I remembered to create an external validation set so I had a few cases to run once the fine-tune was ready. I split the data up 60% to train, 20% to validate for the fine-tune and 20% to use myself to validate after the fine-tune.

I put up my feet and waited for the fine-tune to run.

The first run of the fine-tune.

!fine-tune run model output

I'm very glad I had more validation data because I couldn't believe it when I saw it work. I ran several more and each time, it got the word groups correct. Though occasionally it seemed to struggle to get the category right with x ___ or ___ x categories. But honestly, I was a little stunned. Before this fine-tuning, the model could not come all that close to solving Connections puzzles and now it could generally get the 4 word groups right every time and the categories right most of the time. It seemed to good to be true. It was! Can you spot the problem? The input words

and the output categories

are in the same order.

When I change the input order of the words for this fine-tune

it all falls apart.

The categories stop making sense. The model hallucinates words. Some categories don't even have 4 words in them. This fine-tune is toast.

Fixing another dataset bug

Back to the code. I modified my data set generation code, then split my data into three sets again. I used a stable random seed, so that the generation would be the same each time.

I ran a few tests

and

The results are ok, but not incredible. It does get some groupings correct. There are hallucinations, adding words that don't exist. Also, most categories and grouping are not correct.

Evaluating the results

To figure out if this fine-tune was worthwhile, I needed benchmark against the performance of gpt-3.5-turbo-1106.

I wrote some quick and dirty code to read the validation.jsonl, make an OpenAI call, parse the response and compare it to the known, correct answer. I decided to measure percentage of puzzles correct (all four categories) and percentage of categories correct (4 words correctly grouped). I also augmented the user prompt a bit to get the un-fine-tuned model to output its response the same way the fine-tuned model would for consistent parsing.

I started by running the code against gpt-3.5-turbo-1106.

Next, I ran it against my fine-tuned model.

Not a huge difference. I did a few more runs and none of the results we're too out of the ordinary. Other than consistent formatting, it's not clear the model got all that much better at the game after fine-tuning.

Wrap up

This experience was an interesting introduction to model fine-tuning. The results weren't that amazing, but I learned a lot about easy pitfalls and mistakes one can make and had some fun.

Future work

To see if there might be more to explore for this project, I ran the validation set through gpt-4 with the following results

This improvement is notable over gpt-3.5's ~20-25%. I requested access to fine-tune gpt-4, so we'll see if that comes through and I can try it out.