Tool use: a GPT that calls a calculator

Intro

0:00 Hello and welcome back to episode twenty, where we're going to teach the model we pretrained on the web to use a tool, a calculator, which it calls when a question needs one. Let's begin with a quick look at where we're going to be by the end of the episode.

0:14 By the end of today, our model answers a question like this one, about a café selling three hundred and eighty-six cups of coffee a day for twenty-nine days, by calling a calculator. It writes the call, the calculator works out eleven thousand one hundred and ninety-four, and the model answers with it.

0:33 On the website of the series, you can find more information about the project. The code is open source on GitHub, tagged at each step of each episode, so that you can follow along, and the docs get a page on tool use today. And each episode has a page here as well, with its chapters and its code.

0:52 Now, before we get to today's work, let's quickly go over what we did last time. We pretrained a model of about a hundred and ten million parameters on educational web pages, for a night, on one GPU that we rented for under four dollars. On HellaSwag, it got about as many right as GPT-2 small. It writes fluent pages, but what it writes often isn't true, and arithmetic is beyond it: with three worked examples in front of it, it gets none of fifty two-digit additions right. So today, we don't teach our model arithmetic: we teach it to hand the arithmetic to a calculator.

1:28 Today, we start with how a model calls a tool: it writes the call as text, and the program running the model watches for it. We write a calculator that takes numbers and arithmetic, and refuses anything else, as the model's calls are code we run. We generate fifty thousand questions, with held-out ones worded in ways training doesn't use. We fine-tune for ten minutes on the laptop, with the questions and the calculator's results left out of the loss. And we measure the calls on the held-out questions, with the calculator and with it turned off. All right, let's get started with how a language model calls a tool.

Calling a tool

2:05 Toolformer: Language Models Can Teach Themselves to Use Tools, from 2023, is by Timo Schick and others at Meta. Its abstract starts from a paradox: language models solve new tasks from a few examples, and still struggle with arithmetic or looking up a fact, where much simpler programs do well. Toolformer is trained to decide which tool to call, when, with what arguments, and how to use the result, for a calculator, a question answering system, a search engine, a translator and a calendar. Figure one shows its calls, written in square brackets in the middle of the text: the name of the tool, its arguments, then an arrow and the result. The model writes the call, the program running it fills in the result, and the model carries on with it in its context. Toolformer made its training calls by sampling them from the model and keeping the ones that helped, on a model of six point seven billion parameters. Ours is far smaller, so we generate its training calls with a program.

3:00 Our calls are plain text too, in a format of our own. The model writes two angle brackets, calc, the arithmetic, and two closing brackets. The program generating the text checks each new token, and when the call closes, it gives the arithmetic to the calculator, and appends an equals sign and the result to the text, as if the model had written them. Then the model carries on writing, with the result in its context, so its answer can use it. Plain text keeps the vocabulary as it is, and the model loads in Hugging Face's libraries as before. Bigger models use tokens of their own for this, like Llama 3.1's python tag, and the idea is the one here: text that the program watches for.

3:43 A new file, tools dot py, holds the calculator and the loop that runs it. It imports ast, Python's module for reading Python into a syntax tree, operator, for the four operations as functions, re and torch, and our KV cache. The comment at the top says how a call works: plain text the model writes, with the program running it watching for the closing brackets, so the vocabulary doesn't change. Open and close are the markers, call is a regular expression for a whole call, and ops maps the syntax tree's four operations to the functions that do them. Calculate is the calculator, and it has a job to be careful with. The model writes its calls, and we run them, so we can't hand them to Python's eval, which would run any code at all. Instead, value walks the syntax tree, and takes a number, one of the four operations on two values, or a minus sign in front of one, and anything else raises an error. So powers, function calls and names are refused, and a call can't do more than arithmetic. An expression over a hundred characters is refused too. Ast parse reads the expression, and a syntax error, a refusal, or a division by zero all come back as the word error, which the model then reads. Show writes a result the way our examples do: whole numbers without a decimal point, others to four places, and result text puts an equals sign in front. Generate with tools writes after the prompt's ids, a token at a time, and returns the new ids, along with the calls. The prompt goes through the model once, filling the KV cache. Then, at each step, the next token is the most likely one, or a sampled one with a temperature, and end of text stops it. The text so far is decoded, and when it ends with the opening marker, opened remembers where the arithmetic starts. When it ends with the closing brackets, the call is closed, and what's between the markers goes to the calculator. Its result is encoded and added to the output, and to the tokens the model reads next, along with its own token. With use tool false, the call is noted, and the model goes on to write its own result, which is how we'll see what it does without the calculator. Each step runs the newest tokens through the model, with the cache holding the keys and values of the ones before. Main asks a model a question from the command line, in the format of our examples, and prints the answer and each call. The tests check the calculator's arithmetic, from a multiplication to a division with decimals, and what it refuses: an import, a power, a function, a name, a division by zero. Scripted stands in for a model, writing the tokens of a script one per step, so the loop can be tested without training anything. With the calculator, the result goes in after the call, and the model reads the result's tokens with the closing bracket, in one step, and without it, the model writes its own. And a small real model runs through the loop without an error.

6:32 All five tests pass, in about a second, without a trained model.

Fifty thousand questions

6:37 There's no dataset of questions for our calculator, so tool data dot py generates them. Train is fifty thousand examples, held out is two hundred per held-out set, and share is how the training examples divide up: half arithmetic, a quarter word problems, a tenth questions with numbers that need no call, and the rest web pages. Ask holds ten ways of asking for each operation, from what is a plus b, to find the sum, to a plus b equals question mark. Ask held out has two more for each, which training doesn't use, like what is the total of a and b, or how many times does b go into a. They tell us whether the model calls the calculator because it understood the question, or because it recognises the wording. Word holds the word problems: the question, the call's arithmetic, and the sentence that answers it, with r for the result. A box holds a pencils, how many pencils are in b boxes, is b times a. Some take two operations, like a pack of cards with some lost. The last six are held out, on things training's problems don't mention, like a baker, a ferry and a garden. Numbers holds questions that mention numbers but need no calculator, which is bigger, which comes first, and one more wording is held out. Number picks a number with a given number of digits. Question puts a question into our format: question, colon, the question, and answer on the next line. With call builds an example from a question, its arithmetic and its sentence, and computes the result with our calculator. Its comment explains the segments: the question and the result are text the model is given, and the call and the answer are what it learns to write. Arithmetic picks an operation and two numbers, of one to six digits each. Most divisions are made to come out whole, as most questions people ask do. Word problem fills a problem with numbers from two up to four digits, so no sentence says one cars, keeps a subtraction from going below zero, and makes divisions come out whole. Number question fills a question that needs no call. Web pages takes pages from the last FineWeb-Edu file, after the held-out pages, a file pretraining didn't read, each cut to its first sentences. They keep the model writing text, and teach it that text doesn't call for a calculator. Build makes it all with one seed, so each run gets those examples again, and shuffles the training set. Seen collects the training questions, so no held-out question repeats one. Then the five held-out sets: the training questions with new numbers, arithmetic asked in the held-out ways, the held-out word problems, questions that need no call, and the start of web pages, cut in the middle of a sentence, for the model to carry on. Each set is written as JSON, a line per example, with how many of each kind it has. Load reads a set back, and show prints an example with the text the model doesn't learn in square brackets.

9:28 Tool data dot py writes fifty thousand training examples, about twenty-five thousand of them arithmetic and thirteen thousand word problems, and two hundred in each held-out set, in a few seconds, and prints one example of each kind. The question and the calculator's result are in square brackets, as the model is given them, and the call and the sentence after it are what it learns to write. The question with numbers ends without a call, and the web page is a page.

9:56 Before fine-tuning, we measure the pretrained model, with three worked examples in front of each question, which is the most a pretrained model can go on. With the calculator available, it calls it on about thirty-four percent of the training-style questions, and none of those calls is right. It gets about three point five percent of the answers, and as many without the calculator, as it doesn't use it. On the questions that need no call, it calls anyway about seventeen percent of the time. The wrong calls show why: it writes calc and one of the numbers, or repeats a number from the question as its answer, or repeats the question. It has picked up the look of a call from the examples, but not what goes in one.

Fine-tuning

10:38 As in episode three's fine-tuning, the loss leaves out the question, which the model is given. It learns to write the call, which is the part that's its job. The result is left out too, as the calculator writes it, and a model trained to predict it would be learning arithmetic, the thing we're handing over. And it learns the sentence after, which uses the result.

11:00 Tool SFT dot py follows sft dot py from episode three. The batch is thirty-two examples, and the learning rate is ten to the minus four. Its comment gives the reasons: half of sft dot py's, for a model eight times the size, and a narrow task that shouldn't push out the rest of what it knows. Demo is a held-out word problem the model answers at each evaluation. Encode turns an example into ids, segment by segment, and for each id, whether the model learns to write it. Each segment is encoded on its own, which gives the tokens that encoding the whole text would, as each segment starts at a space or a punctuation mark. Learned repeats the segment's flag once for each of its tokens, so ids and learned stay side by side, and both are cut to the context window, plus one, as a window of inputs needs one more token for its last target. Our examples are far shorter than a thousand and twenty-four tokens, so the cut is there for safety. Make batch pads the examples to the longest one in the batch. Each position predicts the next token, which counts if it's one the model writes, and the others keep minus one, which the loss skips. The inputs are the ids without the last one, and targets are the ids from the second on. NumPy's where keeps a target where learned is true, and puts minus one everywhere else, in a single line, for the whole row. So a question's last position predicts the call's first token, which counts, and the call's last position, the closing brackets, predicts the calculator's equals sign, which doesn't. Estimate loss measures all the held-out examples, in order, so each evaluation measures those ones again. Main loads the model we pretrained on the web, encodes the training examples and the seen held-out set, and counts the tokens there are to learn. The optimiser is AdamW, with pretraining's settings for weight decay. Each evaluation measures the held-out loss, and answers the demo with the calculator. The loop draws random batches of thirty-two, as sft dot py does: the forward pass gives the loss over the tokens that count, then zero grad, backward, the gradients clipped to a norm of one, and the optimiser's step. Each fifty steps, it prints the batch's loss and the minutes so far, and each two hundred and fifty, it evaluates, with the model put back in training mode after. At the end, a last evaluation, the model is saved to checkpoints, tools, and it prints the steps, the examples it saw, and the minutes.

13:30 Then we fine-tune for ten minutes on the laptop. There are one million five hundred thousand tokens to learn, as most of each example is given. At step zero, the pretrained model treats the word problem as the start of a page, and carries on with it. Within two hundred and fifty steps, it calls the calculator with a hundred and twenty-eight times thirty-seven, and answers with the result, on a word problem training didn't include. The batch loss jumps around, as some batches hold more web pages, which are harder to predict than a call. At the end, after five hundred steps, it has seen sixteen thousand examples, about a third of the training set, and its last answer uses the right result, though it says trips where it means people.

Measuring the calls

14:13 Tool eval dot py measures the calls the model makes. Number finds numbers in text, whole or with decimals, and call sets are the three held-out sets whose questions need a call. After call takes the answer the model writes in words, after its call and the result, as the result is the calculator's, and with no call, it's all of the text. Score checks one answer, by the kind of question. On a web page, all that counts is whether the model called. On a question that needs no call, it's right if there's no call and the first number is the answer. On the others, the call is right if our calculator gets the answer from its arithmetic, and the answer is right if the sentence after it has the answer in it. Run answers a set's questions, with the calculator or without it, and with examples in front, for a pretrained model. A pretrained model carries on with another question after its answer, so the answer ends before the next question. Main takes a checkpoint, and dash dash shots, which puts three worked examples from the training set in front of each question, for the pretrained model: the first arithmetic question, the first word problem and the first question that needs no call, each written out in full, with its call and result, and a blank line after it. A pretrained model has no idea what our format is, and examples are how we show it. For each set, it answers with the calculator, and for the call sets, again without it. By the length of the longest number in the question, it keeps the right answers both ways. Rate is the share of rows where a key is true. By digits is a default dictionary of two lists for each length, the answers with the calculator and without it, and each length becomes a row with both rates and how many examples it has. The table prints a line per set: the calls, the right calls, the right answers, and without the tool, while the sets that need no call print what applies to them. Each set keeps three of its answers as samples, and three of its wrong calls, as the question and what the model wrote up to the closing brackets. After the table, it prints two wrong calls from each call set, and saves it all to results, as tools, then the model's name, dot JSON.

16:24 Then we measure the model we fine-tuned, in the way we measured the other. On the training's kinds of questions, with numbers it hasn't seen, it calls each time, and about ninety-nine point five percent of the answers are right. With the calculator off, about two percent are: the model hasn't learned arithmetic, it has learned to hand it over. Asked in the held-out ways, it still calls each time, and about fifty-four point five percent of the calls are right, and on the held-out word problems, about forty percent. On questions that need no call, it calls about five point five percent of the time, and gets about seventy-two percent right, and on web pages, it carries on writing, with a call in about one in two hundred. The wrong calls it prints show where it fails. Its one miss on the training's questions is a nine-digit number, copied with a digit missing. The others pick the wrong operation: shared equally among thirty-seven becomes a subtraction, and lots of, or marbles, go into the call as words, which the calculator refuses. It copies the numbers well, and it knows when to call, but it picks the operation from the wording it learned, and not from what the question means.

17:31 Tools dot py asks it a question of our own: forty-eight thousand two hundred and thirteen times nine hundred and seven. It calls, and answers forty-three million, seven hundred and twenty-nine thousand, one hundred and ninety-one, which is right. With dash dash no tool, the call goes unanswered. The model writes the call a second time, then says the answer is nine hundred and seven, a number from the question. Without the result, it doesn't know what to say. And one of the held-out word problems, about a baker: it calls with a hundred and forty-four times six, and answers eight hundred and sixty-four loaves.

18:09 Side by side, ten minutes of fine-tuning took the pretrained model from about three point five percent of these answers to about ninety-nine point five on questions worded the way training words them. On wordings it hasn't seen, it gets about half, and on new word problems, about two in five. The model learned to call, and learned our fifty or so wordings well, and that's as far as ten minutes and fifty thousand generated questions go.

18:35 By the length of the numbers, on the training's kinds of questions: with the calculator, the answers stay right from two digits to six, as copying a number is one job at any length. Without it, the model gets a few small ones right, and none of the five-digit ones.

18:50 Last, the docs get a page on tool use, with the format and the calculator, the generated questions and their held-out sets, the mask, and the results, with the calculator and without it. It goes at the end of the navigation, after pretraining on the web.

Wrap-up

19:04 Now, let's go over what's real here, and what's scaled down. The loop is real: a model writes a call, the program running it runs the tool and puts the result in the context, and the model carries on, which is how assistants call tools. So is leaving the tool's output out of the loss. The scale of it isn't, not by a long way. We have one tool, and questions generated from about fifty wordings, while an assistant chooses between many tools, described to it in its prompt, after fine-tuning on a far wider range of requests. And our markers are plain text, where those models use tokens of their own and arguments as JSON. Our model calls the calculator well when asked the way it learned, and half as well when asked another way.

19:48 So, today we taught our model to call a tool. It writes a call as plain text, and the program running it watches for the closing brackets. Our calculator reads the call as a syntax tree and takes arithmetic alone, so a call can't run other code. We generated fifty thousand questions, with held-out sets worded in new ways. Ten minutes of fine-tuning on the laptop, with the questions and the results left out of the loss, took the answers from about three point five percent to about ninety-nine point five on the training's kinds of questions, about fifty-five on new wordings and forty on new word problems. With the calculator off, it gets almost none right: it learned to hand the arithmetic over.

20:30 If you'd give a small model a tool of its own, let me know in the comments which one, and what it should learn to ask it. And if you'd like to follow along, subscribe to the channel. I hope you find this useful, and thanks for watching.