Log In to continue learning
Tabular Foundation Models for Data Science Agents
Intermediate
1h32m
You've seen how TabPFN can make and explain predictions. Next, you'll give that capability to an AI agent. You'll use TabPFN MCP and custom skills and tools, so an AI agent, in this case Claude Code, can reason about a task and call TabPFN when it needs a specialized tabular prediction. Let's dive in. Lesson 1 with Frank, you had a brief look into the difference between large language models and tabular foundation models, and how while large language models are great with dealing with unstructured data, how tabular foundation models really focus on dealing with structured data like tables. In this lesson, we're going to also see how they work better together to create agents that are able to make predictions over your tabular data. Let's have a look at how we can make that reality. Large language models are great for reasoning, not necessarily prediction. But let's think of a scenario where we have all this information about customers, how long they've been your customer, how much they pay you, and all the information about past customers who've churned or not churned. If we were to ask a large language model, "Did customer C-1042 churn?" in this particular scenario, we'll get an answer, but if you pay attention, you'll also see that the answer is in fact in the data. Now, asking some questions about unseen data is a slightly different task. Instead, let's assume we have a new customer, and we don't yet know that they've churned or not, and let's ask the same question. will customer C-1046 churn? It will give you an answer, but it will be a plausible answer, and we won't know where this answer really is coming from statistically. We're going to look at how we can use TabPFN the same way we've used it before, but instead of a large language model answering the question on its own, it will use our historical data as training sets, or in this case, context sets, and our new row as the query. This is really a supervised learning problem. Large language models are trained to predict the next-best token. So really, it's only looking at text and generating what the next value of that text might be. For example, this customer will, and then what's the next best token. Whereas a tabular foundation model is working fundamentally differently. It's going to be giving us a distribution over the outcomes of what our question is. We're going to be predicting a value not based on what the large language model knows about the world, but specifically on what the tabular foundation model has learned about our data. What this would look like is the following. The pre-trained transformer TabPFN would receive the context, i.e., the history, as well as the query. We'll then forward that to TabPFN, and then TabPFN will return a prediction. This is different to how classical machine learning models work, and why we can use TabPFN relatively easily as a tool within an agent, because only one TabPFN model can act as a tool for all datasets and all domains. And this means we can avoid having to retrain the model every time we have a new domain. Let's look at the same scenario as we saw before. We had customer C-1046, who was a new customer and we don't yet know whether they've churned or not. our agent that's using a large language model should reason that this is a prediction task because it's data that we haven't seen before. So, it will decide to call our TabPFN tool with all of the relevant inputs for the TabPFN tool to run a prediction. TabPFN will then return a prediction, and then our large language model can do what it's best at, and produce a sentence giving us the answer. We're going to do all of this using the TabPFN MCP server and Claude Code. The MCP server will come with five tools that allow you to fit and predict and even get statuses on how your fit is going. Additionally, the TabPFN MCP server is using TabPFN via the API. It also has the capability to cache your fitted models, so once you've fitted the model once, you'll just use the cached versions, so you don't have to spend time fitting again. Let's see this in action. In this demo, we're going to use Claude Code to try to predict if a customer is likely to churn or not. We're going to use a historical data set of customers and whether they've churned or not, and then try to predict the same for current customers. Before we start, though, let's have a quick look at the structure of this repository. You'll notice that there's a .mcp.json file. This lists all of the MCP servers we want Claude to start off with. In this case, we have only one of our own, which is the official TabPFN MCP server by Prior Labs. We also have a skills directory with one skill for TabPFN handoff. It's a long skills file which is ultimately explaining to Claude when to hand over to TabPFN. In this case, to look out for a question from a user that might actually require a prediction over an existing data set. As you scroll down you'll also notice that there are instructions on how to fetch data, how to run baselines and evaluations, and towards the end you'll also notice a template for a markdown file which is generated when a user asks for a report. We also have some scripts. These are some helper scripts that are also mentioned in the skill, to run the baseline model, to evaluate TabPFN, and to fetch new data. We're going to fetch telco data, which includes customers and whether they've churned or not. But if you like, you can use german_credit, ames_housing, wine_quality, heart_disease, as well. Let's start with starting Claude and having a look at the MCP servers. You'll notice that I have TabPFN here. You may be asked for authorization to enable TabPFN, and once enabled, because we're using the official MCP server, you'll have to authenticate with Prior Labs. This may require you to create a sign up. If you already have an account, it will just ask for authentication. Now, because I don't have any data to start off with, and I want to imagine that I'm a company that has a bunch of data on our customers, I can ask Claude to fetch telco churn data. The fetch data script already splits up the dataset for us into historical, current, and test sets. The difference here is that the historical set will have the labels as to whether the customer has churned or not. The current set we're going to pretend are our current customers where we don't know if they will churn or not. The test set is the same exact one which we can use to evaluate the model. In this case, it has the same exact current set, but it does have the labels. Claude has been set up with the skill to make sure that it's never fitting on the test set, and it only uses historical sets to fit the model with. We can have a quick look inside the data sets. This is our current set. We have a bunch of information about each customer: their gender, whether they're a senior citizen or not, how they pay us, etc. But if we look at the historical set, we have a similar distribution, but we also have the Churn label. The test set is the exact same as current, and we can use it to just have a look at individual predictions and whether they were correct or not. Now to start off, we can do something where we ask for the prediction over a single customer. Let's get row two from our current set, which, notice, does not have the label, and let's ask Claude, "Given our historical data, what's the likelihood of this customer to churn?" And we can copy paste the entire information about the customer that we already have. You'll see that TabPFN has decided that this is a prediction skill, a prediction task, and has handed over to TabPFN. If you've already fit TabPFN with the data set, you'll notice that Claude will tell you that there is a cached model that exists and will not bother fitting again. In this case, I haven't, When you're first using Claude, it may ask for permission to run certain commands. And we see the result where Claude is telling us there's a 1.3% chance of this customer churning. We can go to the test set where the label does exist. It was the first customer in our dataset, and the Churn label is "No". So in this case, the prediction by TabPFN was correct. Instead of running it per customer, we can also ask Claude to predict the top 10 customers risking churn. "Who are my top 10 risking churn right now?" You'll notice that it'll decide that this, again, is a TabPFN task, but instead use our entire telco_churn_current dataset to run predictions over. And one more important thing to notice is that when we're predicting for a class label, the TabPFN MCP server will by default run the predict_proba method. So every time we're running a prediction task for classification, we're going to be predicting all the probabilities for all of the classes. And here we have it. We don't have the names of our customers, but we do have row numbers as IDs. And here we have the list of the probabilities for the top 10 customers risking the highest probability of churn. Now, once you're done trying all of this out, you can also ask for a report. As we go along, we already create some level of reports for all of the handovers we're making to Claude and TabPFN, the train-test splits, etc. However, we can ask for an entire report as well as an evaluation against a baseline. "Create a report for all of this, as well as an evaluation against a benchmark." The benchmark in this case is set to be XGBoost. When evaluation is running, our skill has made it clear that the data set that should be used is the test data set. And once it's ready, we'll have a report.md file. The report will highlight the questions we asked, the results of those questions, and then also the test results. The test results will be using the entire test set we have available under the data directory for the dataset we're using. If you try this agent out with any of the other datasets available through the fetch_data script, reports will be created for each of them under the name of the dataset. Try them out. Try some tasks that are regression tasks instead of classification. The housing data set is a good one for this because you'll be trying to predict the price of houses, and see how well your model does against the baseline.
