Add evaluation code

JohannesHog · JohannesHog · commit af6d3bf94599 · 2025-09-08T16:07:30.000+02:00
diff --git a/nanotabpfn/evaluation.py b/nanotabpfn/evaluation.py
@@ -0,0 +1,164 @@
+import argparse
+
+import numpy as np
+import openml
+import torch
+from openml.config import set_root_cache_directory
+from openml.tasks import TaskType
+from sklearn.preprocessing import LabelEncoder
+
+from nanotabpfn.interface import NanoTabPFNRegressor, NanoTabPFNClassifier
+
+TOY_TASKS_REGRESSION = [
+362443, # diabetes
+]
+
+TOY_TASKS_CLASSIFICATION = [
+    59, # iris
+    2382, # wine
+    9946, # breast_cancer
+]
+
+@torch.no_grad()
+def get_openml_predictions(
+        *,
+        model: NanoTabPFNRegressor | NanoTabPFNClassifier,
+        tasks: list[int] | str = "tabarena-v0.1",
+        max_n_features=500,
+        max_n_instances=10_000,
+        classification: bool | None = None,
+        cache_directory: str | None = None,
+):
+    """
+    Evaluates a model on a set of OpenML tasks and returns predictions.
+
+    Retrieves datasets from OpenML, applies preprocessing, and evaluates the given model on each task.
+    Returns true targets, predicted labels, and predicted probabilities for each dataset.
+
+    Args:
+        model (NanoTabPFNRegressor | NanoTabPFNClassifier): A scikit-learn compatible model or classifier to be evaluated.
+        tasks (list[int] | str, optional): A list of OpenML task IDs or the name of a benchmark suite.
+        max_n_features (int, optional): Maximum number of features allowed for a task. Tasks exceeding this limit are skipped.
+        max_n_instances (int, optional): Maximum number of instances allowed for a task. Tasks exceeding this limit are skipped.
+        classification (bool | None, optional): Whether the model is a classifier (True) or regressor (False). If None, it is inferred from the model type.
+        cache_directory (str | None, optional): Directory to save OpenML data. If None, default cache path is used.
+    Returns:
+        dict: A dictionary where keys are dataset names and values are tuples of (true targets, predicted labels, predicted probabilities).
+    """
+    if classification is None:
+        classification = isinstance(model, NanoTabPFNClassifier)
+
+    if cache_directory is not None:
+        set_root_cache_directory(cache_directory)
+
+    if isinstance(tasks, str):
+        benchmark_suite = openml.study.get_suite(tasks)
+        task_ids = benchmark_suite.tasks
+    else:
+        task_ids = tasks
+
+    dataset_predictions = {}
+
+    for task_id in task_ids:
+        task = openml.tasks.get_task(task_id, download_splits=False)
+
+        if classification and task.task_type_id != TaskType.SUPERVISED_CLASSIFICATION:
+            continue # skip task, only classification
+        if not classification and task.task_type_id != TaskType.SUPERVISED_REGRESSION:
+            continue # skip task, only regression
+
+        dataset = task.get_dataset(download_data=False)
+
+        n_features = dataset.qualities["NumberOfFeatures"]
+        n_instances = dataset.qualities["NumberOfInstances"]
+        if n_features > max_n_features or n_instances > max_n_instances:
+            continue  # skip task, too big
+
+        _, folds, _ = task.get_split_dimensions()
+        tabarena_light = True
+        if tabarena_light:
+            folds = 1 # code supports multiple folds but tabarena_light only has one
+        repeat = 0 # code only supports one repeat
+        targets = []
+        predictions = []
+        probabilities = []
+        for fold in range(folds):
+            X, y, categorical_indicator, attribute_names = dataset.get_data(
+                target=task.target_name, dataset_format="dataframe"
+            )
+            train_indices, test_indices = task.get_train_test_split_indices(
+                fold=fold, repeat=repeat
+            )
+            X_train = X.iloc[train_indices].to_numpy()
+            y_train = y.iloc[train_indices].to_numpy()
+            X_test = X.iloc[test_indices].to_numpy()
+            y_test = y.iloc[test_indices].to_numpy()
+
+            if classification:
+                label_encoder = LabelEncoder()
+                y_train = label_encoder.fit_transform(y_train)
+                y_test = label_encoder.transform(y_test)
+            targets.append(y_test)
+
+            model.fit(X_train, y_train)
+            y_pred = model.predict(X_test)
+            predictions.append(y_pred)
+            if classification:
+                y_proba = model.predict_proba(X_test)
+                if y_proba.shape[1] == 2:  # binary classification
+                    y_proba = y_proba[:, 1]
+                probabilities.append(y_proba)
+
+        y_pred = np.concatenate(predictions, axis=0)
+        targets = np.concatenate(targets, axis=0)
+        probabilities = np.concatenate(probabilities, axis=0) if len(probabilities) > 0 else None
+        dataset_predictions[str(dataset.name)] = (targets, y_pred, probabilities)
+    return dataset_predictions
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-model_type", type=str, choices=["regression", "classification"], required=True,
+                        help="Whether to use the regressor or classifier model")
+    parser.add_argument("-checkpoint", type=str, default=None,
+                        help="Path to load the model weights from. If None, default weights are used.")
+    parser.add_argument("-dist_path", type=str, default=None,
+                        help="Path to load the bucket edges for the support bar distribution from. Only needed for regression.")
+    parser.add_argument("-tasks", type=str, default="tabarena-v0.1",
+                        choices=["tabarena-v0.1", "toy_tasks"], help="Which OpenML tasks to evaluate on.")
+    parser.add_argument("-cache_directory", type=str, default=None,
+                        help="Directory to save OpenML data. If None, default cache path is used.")
+    parser.add_argument("-max_n_features", type=int, default=500,
+                        help="Maximum number of features allowed for a task. Tasks exceeding this limit are skipped.")
+    parser.add_argument("-max_n_instances", type=int, default=10_000,
+                        help="Maximum number of instances allowed for a task. Tasks exceeding this limit are skipped.")
+    args = parser.parse_args()
+
+    if args.model_type == "classification":
+        model = NanoTabPFNClassifier(model=args.checkpoint)
+    else:
+        model = NanoTabPFNRegressor(model=args.checkpoint, dist=args.dist_path)
+    model.model.eval()
+
+    if args.tasks == "toy_tasks" and args.model_type == "regression":
+        tasks = TOY_TASKS_REGRESSION
+    elif args.tasks == "toy_tasks" and args.model_type == "classification":
+        tasks = TOY_TASKS_CLASSIFICATION
+    else:
+        tasks = args.tasks
+
+    predictions = get_openml_predictions(
+        model=model, tasks=tasks, max_n_features=args.max_n_features, max_n_instances=args.max_n_instances,
+        classification=(args.model_type=="classification"), cache_directory=args.cache_directory
+    )
+
+    for dataset_name, (y_true, y_pred, y_proba) in predictions.items():
+        if args.model_type == "classification":
+            from sklearn.metrics import roc_auc_score, balanced_accuracy_score
+            acc = balanced_accuracy_score(y_true, y_pred)
+            auc = roc_auc_score(y_true, y_proba, multi_class='ovr')
+            print(f"Dataset: {dataset_name} | ROC AUC: {auc:.4f} | Balanced Accuracy: {acc:.4f}")
+        else:
+            from sklearn.metrics import r2_score
+            r2 = r2_score(y_true, y_pred)
+            print(f"Dataset: {dataset_name} | R2: {r2:.4f}")
diff --git a/nanotabpfn/interface.py b/nanotabpfn/interface.py
@@ -1,13 +1,20 @@
 import os
-import requests
+
 import numpy as np
+import pandas as pd
+import requests
 import torch
 import torch.nn.functional as F
-
+from numpy import ndarray
 from pfns.bar_distribution import FullSupportBarDistribution
+from sklearn.compose import ColumnTransformer
+from sklearn.impute import SimpleImputer
+from sklearn.pipeline import Pipeline
+from sklearn.preprocessing import OrdinalEncoder, FunctionTransformer
 
-from nanotabpfn.utils import get_default_device
 from nanotabpfn.model import NanoTabPFNModel
+from nanotabpfn.utils import get_default_device
+
 
 def init_model_from_state_dict_file(file_path):
     """
@@ -29,6 +36,40 @@ def init_model_from_state_dict_file(file_path):
     model.load_state_dict(torch.load(file_path, map_location='cpu'))
     return model
 
+def get_feature_preprocessor(X: ndarray | pd.DataFrame) -> ColumnTransformer:
+    """
+    fits a preprocessor that replaces NaNs with the mean of the respective column
+    and scales each column to mean 0 and variance 1
+    """
+    X = pd.DataFrame(X)
+    num_mask = []
+    for col in X:
+        non_nan_entries = X[col].notna().sum()
+        numeric_entries = pd.to_numeric(X[col], errors='coerce').notna().sum() # in case numeric columns are stored as strings
+        num_mask.append(non_nan_entries == numeric_entries)
+        # num_mask.append(is_numeric_dtype(X[col]))  # Assumes pandas dtype is correct
+
+    num_mask = np.array(num_mask)
+
+    num_transformer = Pipeline([
+        ("to_pandas", FunctionTransformer(lambda x: pd.DataFrame(x) if not isinstance(x, pd.DataFrame) else x)), # to apply pd.to_numeric of pandas
+        ("to_numeric", FunctionTransformer(lambda x: x.apply(pd.to_numeric, errors='coerce').to_numpy())), # in case numeric columns are stored as strings
+        ('imputer', SimpleImputer(strategy='mean')) # median might be better because of outliers
+    ])
+    cat_transformer = Pipeline([
+        ('encoder', OrdinalEncoder(handle_unknown='use_encoded_value', unknown_value=np.nan)),
+        ('imputer', SimpleImputer(strategy='most_frequent')),
+    ])
+
+    preprocessor = ColumnTransformer(
+        transformers=[
+            ('num', num_transformer, num_mask),
+            ('cat', cat_transformer, ~num_mask)
+        ]
+    )
+    return preprocessor
+
+
 class NanoTabPFNClassifier():
     """ scikit-learn like interface """
     def __init__(self, model: NanoTabPFNModel|str|None = None, device=get_default_device()):
@@ -46,7 +87,8 @@ def __init__(self, model: NanoTabPFNModel|str|None = None, device=get_default_de
 
     def fit(self, X_train: np.array, y_train: np.array):
         """ stores X_train and y_train for later use, also computes the highest class number occuring in num_classes """
-        self.X_train = X_train
+        self.feature_preprocessor = get_feature_preprocessor(X_train)
+        self.X_train = self.feature_preprocessor.fit_transform(X_train)
         self.y_train = y_train
         self.num_classes = max(set(y_train))+1
 
@@ -60,7 +102,7 @@ def predict_proba(self, X_test: np.array) -> np.array:
         creates (x,y), runs it through our PyTorch Model, cuts off the classes that didn't appear in the training data
         and applies softmax to get the probabilities
         """
-        x = np.concatenate((self.X_train, X_test))
+        x = np.concatenate((self.X_train, self.feature_preprocessor.transform(X_test)))
         y = self.y_train
         with torch.no_grad():
             x = torch.from_numpy(x).unsqueeze(0).to(torch.float).to(self.device)  # introduce batch size 1
@@ -76,7 +118,7 @@ def predict_proba(self, X_test: np.array) -> np.array:
 class NanoTabPFNRegressor():
     """ scikit-learn like interface """
     def __init__(self, model: NanoTabPFNModel|str|None = None, dist: FullSupportBarDistribution|str|None = None, device=get_default_device()):
-        if model == None:
+        if model is None:
             model = 'nanotabpfn_regressor.pth'
             dist = 'nanotabpfn_regressor_buckets.pth'
             if not os.path.isfile(model):
@@ -105,7 +147,8 @@ def fit(self, X_train: np.array, y_train: np.array):
         """
         Stores X_train and y_train for later use. Computes target normalization. Builds normalized bar distribution from existing self.dist.
         """
-        self.X_train = X_train
+        self.feature_preprocessor = get_feature_preprocessor(X_train)
+        self.X_train = self.feature_preprocessor.fit_transform(X_train)
         self.y_train = y_train
 
         self.y_train_mean = np.mean(self.y_train)
@@ -121,7 +164,7 @@ def predict(self, X_test: np.array) -> np.array:
         """
         Performs in-context learning using X_train and y_train. Predicts the means of the output distributions for X_test.
         """
-        X = np.concatenate((self.X_train, X_test))
+        X = np.concatenate((self.X_train, self.feature_preprocessor.transform(X_test)))
         y = self.y_train_n
 
         with torch.no_grad():
diff --git a/pretrain_classification.py b/pretrain_classification.py
@@ -1,20 +1,16 @@
 import argparse
-import torch
-import numpy as np
 
+import torch
+from sklearn.metrics import accuracy_score
 from torch import nn
-from functools import partial
 
 from nanotabpfn.callbacks import ConsoleLoggerCallback
-from nanotabpfn.priors import PriorDumpDataLoader
+from nanotabpfn.evaluation import get_openml_predictions, TOY_TASKS_CLASSIFICATION
+from nanotabpfn.interface import NanoTabPFNClassifier
 from nanotabpfn.model import NanoTabPFNModel
+from nanotabpfn.priors import PriorDumpDataLoader
 from nanotabpfn.train import train
 from nanotabpfn.utils import get_default_device, set_randomness_seed
-from nanotabpfn.interface import NanoTabPFNClassifier
-
-from sklearn.datasets import *
-from sklearn.model_selection import train_test_split
-from sklearn.metrics import accuracy_score, roc_auc_score
 
 parser = argparse.ArgumentParser()
 parser.add_argument("-priordump", type=str, default="/50x3_3_100k_classification.h5", help="path to the prior dump")
@@ -55,29 +51,22 @@
 if ckpt:
     model.load_state_dict(ckpt['model'])
 
-datasets = []
-datasets.append(train_test_split(*load_iris(return_X_y=True), test_size=0.5, random_state=42))
-datasets.append(train_test_split(*load_wine(return_X_y=True), test_size=0.5, random_state=42))
-datasets.append(train_test_split(*load_breast_cancer(return_X_y=True), test_size=0.5, random_state=42))
-
-
 class EvaluationLoggerCallback(ConsoleLoggerCallback):
-    def __init__(self, datasets):
-        self.datasets = datasets
+    def __init__(self, tasks):
+        self.tasks = tasks
 
     def on_epoch_end(self, epoch: int, epoch_time: float, loss: float, model, **kwargs):
         classifier = NanoTabPFNClassifier(model, device)
+        predictions = get_openml_predictions(model=classifier, tasks=self.tasks)
         scores = []
-        for X_train, X_test, y_train, y_test in self.datasets:
-            classifier.fit(X_train, y_train)
-            pred = classifier.predict(X_test)
-            scores.append(accuracy_score(y_test, pred))
+        for dataset_name, (y_true, y_pred, y_proba) in predictions.items():
+            scores.append(accuracy_score(y_true, y_pred))
         avg_score = sum(scores) / len(scores)
         print(f'epoch {epoch:5d} | time {epoch_time:5.2f}s | mean loss {loss:5.2f} | avg accuracy {avg_score:.3f}',
               flush=True)
 
 
-callbacks = [EvaluationLoggerCallback(datasets)]
+callbacks = [EvaluationLoggerCallback(TOY_TASKS_CLASSIFICATION)]
 
 trained_model, loss = train(
     model=model,
diff --git a/pretrain_regression.py b/pretrain_regression.py
@@ -1,18 +1,16 @@
 import argparse
+
 import torch
+from pfns.bar_distribution import FullSupportBarDistribution
+from sklearn.metrics import r2_score
 
 from nanotabpfn.callbacks import ConsoleLoggerCallback
-from nanotabpfn.priors import PriorDumpDataLoader
+from nanotabpfn.evaluation import get_openml_predictions, TOY_TASKS_REGRESSION
+from nanotabpfn.interface import NanoTabPFNRegressor
 from nanotabpfn.model import NanoTabPFNModel
+from nanotabpfn.priors import PriorDumpDataLoader
 from nanotabpfn.train import train
 from nanotabpfn.utils import get_default_device, set_randomness_seed, make_global_bucket_edges
-from nanotabpfn.interface import NanoTabPFNRegressor
-
-from pfns.bar_distribution import FullSupportBarDistribution
-
-from sklearn.datasets import load_diabetes
-from sklearn.model_selection import train_test_split
-from sklearn.metrics import r2_score
 
 parser = argparse.ArgumentParser()
 
@@ -66,27 +64,22 @@
 
 dist = FullSupportBarDistribution(bucket_edges)
 
-datasets = []
-datasets.append(train_test_split(*load_diabetes(return_X_y=True), test_size=0.5, random_state=42))
-
-
 class EvaluationLoggerCallback(ConsoleLoggerCallback):
-    def __init__(self, datasets):
-        self.datasets = datasets
+    def __init__(self, tasks):
+        self.tasks = tasks
 
     def on_epoch_end(self, epoch: int, epoch_time: float, loss: float, model, **kwargs):
         regressor = NanoTabPFNRegressor(model, dist, device)
+        predictions = get_openml_predictions(model=regressor, tasks=self.tasks)
         scores = []
-        for X_train, X_test, y_train, y_test in datasets:
-            regressor.fit(X_train, y_train)
-            pred = regressor.predict(X_test)
-            scores.append(r2_score(y_test, pred))
+        for dataset_name, (y_true, y_pred, _) in predictions.items():
+            scores.append(r2_score(y_true, y_pred))
         avg_score = sum(scores) / len(scores)
         print(f'epoch {epoch:5d} | time {epoch_time:5.2f}s | mean loss {loss:5.2f} | avg r2 score {avg_score:.3f}',
               flush=True)
 
 
-callbacks = [EvaluationLoggerCallback(datasets)]
+callbacks = [EvaluationLoggerCallback(TOY_TASKS_REGRESSION)]
 
 trained_model, loss = train(
     model=model,
