GraphPFN

https://github.com/yandex-research/graphpfn
Admin
GraphPFN is a graph foundation model for node-level tasks that enables zero-shot in-context learning...
Tokens:71,845
Snippets:427
Trust Score:9.2
License:Apache-2.0
Update:1 week ago
Show doc for...
Tokens:
Raw
### Install dependencies with uv sync

Source: https://github.com/yandex-research/graphpfn/blob/main/paper/README.md

Installs project dependencies using uv. Requires uv to be installed first.

```bash
uv sync
```

--------------------------------

### Install graphpfn from local clone

Source: https://github.com/yandex-research/graphpfn/blob/main/README.md

Install graphpfn from a local clone in editable mode.

```bash
git clone https://github.com/yandex-research/graphpfn
cd graphpfn
pip install -e .
```

--------------------------------

### Install graphpfn from PyPI

Source: https://github.com/yandex-research/graphpfn/blob/main/README.md

Install the graphpfn package from PyPI.

```bash
pip install graphpfn
```

--------------------------------

### Install Optuna for retrieval optimization

Source: https://github.com/yandex-research/graphpfn/blob/main/paper/vendor/limix/README.md

Installs Optuna, required for hyperparameter tuning of retrieval parameters.

```bash
pip install optuna
```

--------------------------------

### Work within the repo using uv

Source: https://github.com/yandex-research/graphpfn/blob/main/README.md

Set up the environment using uv: install uv, clone the repository, and run a script with uv run. The full environment including PyTorch and DGL is defined in uv.lock and set up automatically on first run.

```bash
# Install uv
curl -LsSf https://astral.sh/uv/install.sh | sh

# Clone the repo
git clone https://github.com/yandex-research/graphpfn
cd graphpfn

# Run a script
uv run examples/basic.py
```

--------------------------------

### Install Python dependencies

Source: https://github.com/yandex-research/graphpfn/blob/main/paper/vendor/limix/README.md

Installs Python dependencies for LimiX, including the downloaded flash_attn wheel and other required packages.

```bash
pip install python==3.12.7 torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1
pip install flash_attn-2.8.0.post2+cu12torch2.7cxx11abiTRUE-cp312-cp312-linux_x86_64.whl
pip install scikit-learn  einops  huggingface-hub matplotlib networkx numpy pandas  scipy tqdm typing_extensions xgboost kditransform hyperopt
```

--------------------------------

### Download prebuilt flash_attn wheel

Source: https://github.com/yandex-research/graphpfn/blob/main/paper/vendor/limix/README.md

Downloads the prebuilt flash_attn wheel file for CUDA 12 and PyTorch 2.7. This is required for manual installation.

```bash
wget -O flash_attn-2.8.0.post2+cu12torch2.7cxx11abiTRUE-cp312-cp312-linux_x86_64.whl https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.0.post2/flash_attn-2.8.0.post2+cu12torch2.7cxx11abiTRUE-cp312-cp312-linux_x86_64.whl
```

--------------------------------

### Install PyTorch and DGL with CUDA 11.8

Source: https://github.com/yandex-research/graphpfn/blob/main/README.md

Install PyTorch 2.4 and DGL with CUDA 11.8 support. DGL does not officially support more modern PyTorch versions, so graphpfn pins torch<2.5.

```bash
pip install "torch==2.4" --index-url https://download.pytorch.org/whl/cu118
pip install dgl -f https://data.dgl.ai/wheels/torch-2.4/cu118/repo.html
```

--------------------------------

### Run GraphPFN pretraining with distributed training

Source: https://github.com/yandex-research/graphpfn/blob/main/paper/README.md

Runs GraphPFN pretraining using PyTorch distributed launch with 8 processes per node. Sets the DGL backend to PyTorch.

```bash
DGLBACKEND=pytorch uv run -m torch.distributed.run --nproc-per-node 8 bin/go.py exp/graphpfn/pretrain/main/pretrain.toml
```

--------------------------------

### Clone and enter LimiX repository

Source: https://github.com/yandex-research/graphpfn/blob/main/paper/vendor/limix/README.md

Clones the LimiX repository and changes into the project directory.

```bash
git clone https://github.com/limix-ldm/LimiX.git
cd LimiX
```

--------------------------------

### Build LimiX Docker image

Source: https://github.com/yandex-research/graphpfn/blob/main/paper/vendor/limix/README.md

Builds the LimiX Docker image using the provided Dockerfile. The --network=host flag is used for network access during build, and the FROM_IMAGES build argument specifies the base CUDA image.

```bash
docker build --network=host -t limix/infe:v1 --build-arg FROM_IMAGES=nvidia/cuda:12.2.0-base-ubuntu22.04 -f Dockerfile .
```

--------------------------------

### Run GraphPFN evaluation (finetuning with 10 ensemble members)

Source: https://github.com/yandex-research/graphpfn/blob/main/paper/README.md

Executes a minimal evaluation run for GraphPFN finetuning with 10 ensemble members on the tolokers-2 dataset. Requires datasets downloaded to the data directory.

```bash
uv run bin/go.py exp/graphpfn/eval/main/finetune/10/tolokers-2/tuning.toml --force
```

--------------------------------

### LimiXPredictor.__init__

Source: https://github.com/yandex-research/graphpfn/blob/main/paper/vendor/limix/README.md

Creates a LimiXPredictor instance by loading a model from a specified path and configuring inference parameters such as device, precision, task type, and outlier removal.

```APIDOC
## LimiXPredictor.__init__

### Description
Initializes a LimiXPredictor object with the given device, model path, and inference configuration. It sets up the model for classification, regression, or missing value imputation based on the task_type and mask_prediction flags.

### Method
Constructor

### Signature
```python
def __init__(self,
             device: torch.device,
             model_path: str,
             mix_precision: bool = True,
             inference_config: list | str,
             categorical_features_indices: List[int] | None = None,
             outlier_remove_std: float = 12,
             softmax_temperature: float = 0.9,
             task_type: Literal['Classification', 'Regression'] = 'Classification',
             mask_prediction: bool = False,
             inference_with_DDP: bool = False,
             seed: int = 0)
```

### Parameters
- **device** (torch.device) - Required - The hardware that loads the model.
- **model_path** (str) - Required - The path to the model that needs to be loaded.
- **mix_precision** (bool) - Optional - Whether to enable the mixed precision inference. Default: True.
- **inference_config** (list|str) - Required - Configuration file used for inference.
- **categorical_features_indices** (List[int]|None) - Optional - The indices of categorical columns in the tabular data. Default: None.
- **outlier_remove_std** (float) - Optional - The threshold is employed to remove outliers, defined as values that are multiples of the standard deviation. Default: 12.
- **softmax_temperature** (float) - Optional - The temperature used to control the behavior of softmax operator. Default: 0.9.
- **task_type** (str) - Optional - The task type which can be either "Classification" or "Regression". Default: 'Classification'.
- **mask_prediction** (bool) - Optional - Whether to enable missing value imputation. Default: False.
- **inference_with_DDP** (bool) - Optional - Whether to enable DDP during inference. Default: False.
- **seed** (int) - Optional - The seed to control random states. Default: 0.

### Returns
None
```

--------------------------------

### GraphPFN.from_pretrained

Source: https://github.com/yandex-research/graphpfn/blob/main/README.md

Creates a GraphPFN model instance from a pretrained checkpoint. This is the recommended way to interact with the model directly.

```APIDOC
## GraphPFN.from_pretrained

### Description
Creates a GraphPFN model instance from a pretrained checkpoint. This is the recommended way to interact with the model directly.

### Method
from_pretrained

### Endpoint
GraphPFN.from_pretrained

### Parameters
- **n_features** (int) - Required - Number of input features.
- **n_classes** (int or None) - Optional - Number of classes; None for regression.
- **device** (str) - Optional - Device to load the model on (e.g., "cuda:0").
- **checkpoint** (str) - Optional - Path or HF identifier to the checkpoint. Default: "hf://eremeev-d/graphpfn-1.3/graphpfn-adapters-1_3.pt".
- **ensemble_kwargs** (dict) - Optional - Keyword arguments for the ensemble, e.g., {"n_members": 10}. Default: {"n_members": 10}.

### Request Example
```python
from graphpfn import GraphPFN

model = GraphPFN.from_pretrained(
    n_features=64,
    n_classes=5,           # None for regression
    device="cuda:0",
    checkpoint="hf://eremeev-d/graphpfn-1.3/graphpfn-adapters-1_3.pt",  # default
    ensemble_kwargs={"n_members": 10},  # default
)
```

### Response
Returns a GraphPFN model instance.

```

--------------------------------

### Finetune with grid search using tune_and_predict_ft

Source: https://github.com/yandex-research/graphpfn/blob/main/README.md

Finetune the model on the target graph with a grid search over learning rates. Returns predictions and the best hyperparameters found. Requires a dataset object with a task_type attribute and a CUDA device.

```python
from functools import partial
from graphpfn import tune_and_predict_ft
from graphpfn.inference.util import compute_score, compute_loss

preds, best_hparams = tune_and_predict_ft(
    dataset=dataset,
    score_fn=partial(compute_score, task_type=dataset.task_type),
    loss_fn=partial(compute_loss, task_type=dataset.task_type),
    device="cuda:0",
)
print(best_hparams)  # {"lr": 0.0005}
```

--------------------------------

### Load GraphPFN and run inference

Source: https://github.com/yandex-research/graphpfn/blob/main/README.md

Instantiate the recommended GraphPFN model and run inference. Use n_classes=None for regression. The default checkpoint is hf://eremeev-d/graphpfn-1.3/graphpfn-adapters-1_3.pt and the default ensemble has n_members=10. The forward call returns a tensor of shape [n_nodes, n_classes].

```python
from graphpfn import GraphPFN

model = GraphPFN.from_pretrained(
    n_features=64,
    n_classes=5,           # None for regression
    device="cuda:0",
    checkpoint="hf://eremeev-d/graphpfn-1.3/graphpfn-adapters-1_3.pt",  # default
    ensemble_kwargs={"n_members": 10},  # default
)

preds = model(
    graph=graph,                  # DGL Graph
    features=features,            # Tensor [n_nodes, n_features]
    y_train=y_train,              # Tensor [n_train]
    train_mask=train_mask,        # BoolTensor [n_nodes]
    task_type="multiclass",
)  # → Tensor [n_nodes, n_classes]
```

--------------------------------

### Run regression inference with retrieval

Source: https://github.com/yandex-research/graphpfn/blob/main/paper/vendor/limix/README.md

Runs regression inference with sample retrieval. Requires a GPU with specifications higher than the NVIDIA RTX 4090.

```bash
python inference_regression.py --save_name your_save_name --inference_config_path path_to_retrieval_config --data_dir path_to_data
```

--------------------------------

### Run regression inference with multi-GPU distributed

Source: https://github.com/yandex-research/graphpfn/blob/main/paper/vendor/limix/README.md

Runs regression inference with multi-GPU distributed inference using torchrun and the --inference_with_DDP flag.

```bash
torchrun --nproc_per_node=8  inference_regression.py --save_name your_save_name --inference_config_path path_to_retrieval_config --data_dir path_to_data --inference_with_DDP
```

--------------------------------

### Regression with LimiXPredictor

Source: https://github.com/yandex-research/graphpfn/blob/main/paper/vendor/limix/README.md

Shows how to use LimiXPredictor for regression on the California housing dataset. Requires a CUDA device and downloads the LimiX-16M checkpoint from Hugging Face. Normalizes the target variable and outputs RMSE and R² scores.

```python
from functools import partial

from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split
from sklearn.metrics import r2_score
from huggingface_hub import hf_hub_download
try:
    from sklearn.metrics import root_mean_squared_error as mean_squared_error
except:
    from sklearn.metrics import mean_squared_error
    mean_squared_error = partial(mean_squared_error, squared=False)
import os, sys

os.environ["RANK"] = "0"
os.environ["WORLD_SIZE"] = "1"
os.environ["MASTER_ADDR"] = "127.0.0.1"
os.environ["MASTER_PORT"] = "29500"

ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
if ROOT_DIR not in sys.path:
    sys.path.insert(0, ROOT_DIR)
from inference.predictor import LimiXPredictor

house_data = fetch_california_housing()
X, y = house_data.data, house_data.target
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=42)

y_mean = y_train.mean()
y_std = y_train.std()
y_train_normalized = (y_train - y_mean) / y_std
y_test_normalized = (y_test - y_mean) / y_std

model_path = hf_hub_download(repo_id="stableai-org/LimiX-16M", filename="LimiX-16M.ckpt", local_dir="./cache")

model = LimiXPredictor(device=torch.device('cuda'), model_path=model_path, inference_config='config/reg_default_retrieval.json')
y_pred = model.predict(X_train, y_train_normalized, X_test)    

# Compute RMSE and R²
y_pred = y_pred.to('cpu').numpy()
rmse = mean_squared_error(y_test_normalized, y_pred)
r2 = r2_score(y_test_normalized, y_pred)

print(f'RMSE: {rmse}')
print(f'R2: {r2}')
```

--------------------------------

### Generate inference configuration file

Source: https://github.com/yandex-research/graphpfn/blob/main/paper/vendor/limix/README.md

Generates the inference configuration file for customizing data preprocessing.

```python
generate_inference_config()
```

--------------------------------

### Run classification inference with retrieval

Source: https://github.com/yandex-research/graphpfn/blob/main/paper/vendor/limix/README.md

Runs classification inference with sample retrieval. Requires a GPU with specifications higher than the NVIDIA RTX 4090.

```bash
python inference_classifier.py --save_name your_save_name --inference_config_path path_to_retrieval_config --data_dir path_to_data
```

--------------------------------

### Load GraphLand dataset with from_graphland

Source: https://github.com/yandex-research/graphpfn/blob/main/README.md

Loads a GraphLand dataset with pre-defined splits. Automatically separates numerical, fractional, and categorical features based on dataset metadata. split options: "RL" (default), "RH", "TH".

```python
dataset = GraphDataset.from_graphland(
    "tolokers-2",
    split="RL",          # "RL" (default) | "RH" | "TH"
    root_dir="data",     # path where datasets are stored
)
```

--------------------------------

### Run regression inference on single GPU or CPU

Source: https://github.com/yandex-research/graphpfn/blob/main/paper/vendor/limix/README.md

Runs regression inference on a single GPU or CPU.

```bash
python  inference_regression.py --save_name your_save_name --inference_config_path path_to_retrieval_config --data_dir path_to_data
```

--------------------------------

### Run classification inference with multi-GPU distributed

Source: https://github.com/yandex-research/graphpfn/blob/main/paper/vendor/limix/README.md

Runs classification inference with multi-GPU distributed inference using torchrun and the --inference_with_DDP flag.

```bash
torchrun --nproc_per_node=8  inference_classifier.py --save_name your_save_name --inference_config_path path_to_retrieval_config --data_dir path_to_data --inference_with_DDP
```

--------------------------------

### Search retrieval hyperparameters with Optuna

Source: https://github.com/yandex-research/graphpfn/blob/main/paper/vendor/limix/README.md

Launches an Optuna study to find the best combination of retrieval parameters for your dataset. Uses the RetrievalSearchHyperparameters class and the search method with n_trials=10 and metric="AUC".

```python
searchInference = RetrievalSearchHyperparameters(
           dict(device_id=0,model_path=model_path), X_train, y_train, X_test, y_test,
)
config, result = searchInference.search(n_trials=10, metric="AUC",
              inference_config='config/cls_default_retrieval.json',task_type="cls")
```

--------------------------------

### Classification with LimiXPredictor

Source: https://github.com/yandex-research/graphpfn/blob/main/paper/vendor/limix/README.md

Shows how to use LimiXPredictor for classification on the breast cancer dataset. Requires a CUDA device and downloads the LimiX-16M checkpoint from Hugging Face. Outputs ROC AUC and accuracy scores.

```python
from sklearn.datasets import load_breast_cancer
from sklearn.metrics import accuracy_score, roc_auc_score
from sklearn.model_selection import train_test_split
from huggingface_hub import hf_hub_download
import numpy as np
import os, sys

os.environ["RANK"] = "0"
os.environ["WORLD_SIZE"] = "1"
os.environ["MASTER_ADDR"] = "127.0.0.1"
os.environ["MASTER_PORT"] = "29500"

ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
if ROOT_DIR not in sys.path:
    sys.path.insert(0, ROOT_DIR)
from inference.predictor import LimiXPredictor

X, y = load_breast_cancer(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.5, random_state=42)

model_file = hf_hub_download(repo_id="stableai-org/LimiX-16M", filename="LimiX-16M.ckpt", local_dir="./cache")

clf = LimiXPredictor(device=torch.device('cuda'), model_path=model_file, inference_config='config/cls_default_retrieval.json')
prediction = clf.predict(X_train, y_train, X_test)

print("roc_auc_score:", roc_auc_score(y_test, prediction[:, 1]))
print("accuracy_score:", accuracy_score(y_test, np.argmax(prediction, axis=1)))
```

--------------------------------

### Create GraphDataset from PyG with from_pyg

Source: https://github.com/yandex-research/graphpfn/blob/main/README.md

Creates a GraphDataset from a PyTorch Geometric data object. PyG node features (pyg_data.x) are loaded as "other" features. For high-dimensional neural embeddings, consider setting d_pca=64 in preprocessing_kwargs.

```python
from graphpfn import GraphDataset

dataset = GraphDataset.from_pyg(
    pyg_data,               # torch_geometric.data.Data
    task_type="multiclass",
    name="my-dataset",
    split_index=0,          # which column of *_mask to use (default: 0)
)
```

--------------------------------

### Run classification inference on single GPU or CPU

Source: https://github.com/yandex-research/graphpfn/blob/main/paper/vendor/limix/README.md

Runs classification inference on a single GPU or CPU.

```bash
python  inference_classifier.py --save_name your_save_name --inference_config_path path_to_retrieval_config --data_dir path_to_data
```

--------------------------------

### Use compute_score, compute_loss, and compute_metrics from graphpfn.inference.util

Source: https://github.com/yandex-research/graphpfn/blob/main/README.md

Use the provided score, loss, and metrics functions for finetuning. compute_score returns a scalar (higher is better) for early stopping, compute_loss returns a differentiable tensor loss, and compute_metrics returns a dict of all relevant metrics.

```python
from graphpfn.inference.util import compute_score, compute_loss, compute_metrics

# compute_score: single scalar, higher is always better (used for early stopping)
# Defaults: average_precision (binclass), accuracy (multiclass), R^2 (regression)
score = compute_score(y_true=..., y_pred=..., task_type="multiclass")

# compute_loss: differentiable Tensor loss (cross-entropy or MSE)
loss = compute_loss(y_true=..., y_pred=..., task_type="multiclass")

# compute_metrics: dict of all relevant metrics
metrics = compute_metrics(y_true=..., y_pred=..., task_type="binclass")
# {"average_precision": ..., "roc_auc": ...}
```

--------------------------------

### Run ICL inference with predict_icl

Source: https://github.com/yandex-research/graphpfn/blob/main/README.md

Runs ICL inference on an amazon-ratings dataset using pretrained GraphPFN weights downloaded automatically from HuggingFace. Requires a CUDA device (device="cuda:0") and prints accuracy metrics.

```python
from torch_geometric.datasets import HeterophilousGraphDataset
from graphpfn import GraphDataset, predict_icl
from graphpfn.inference.util import compute_metrics

pyg_data = HeterophilousGraphDataset("data", "amazon-ratings")[0]
dataset = GraphDataset.from_pyg(pyg_data, task_type="multiclass", name="amazon-ratings")

preds = predict_icl(dataset=dataset, device="cuda:0")

test_mask = dataset.masks["test"]
metrics = compute_metrics(
    y_true=dataset.targets[test_mask],
    y_pred=preds[test_mask],
    task_type=dataset.task_type,
)
print(metrics)  # {"accuracy": ...}
```

--------------------------------

### GraphPFNLight.forward

Source: https://github.com/yandex-research/graphpfn/blob/main/README.md

Performs a single forward pass without ensembling. The signature is identical to GraphPFN.forward. Intended for advanced use cases where you manage ensembling yourself.

```APIDOC
## GraphPFNLight.forward

### Description
Performs a single forward pass without ensembling. The signature is identical to GraphPFN.forward. Intended for advanced use cases where you manage ensembling yourself.

### Method
forward

### Endpoint
GraphPFNLight.forward

### Parameters
- **graph** (DGL Graph) - Required - The input graph.
- **features** (Tensor) - Required - Feature tensor of shape [n_nodes, n_features].
- **y_train** (Tensor) - Required - Training labels tensor of shape [n_train].
- **train_mask** (BoolTensor) - Required - Boolean mask of shape [n_nodes] indicating training nodes.
- **task_type** (str) - Required - Type of task, e.g., "multiclass".

### Response
Returns a Tensor of shape [n_nodes, n_classes] with predictions.

```

--------------------------------

### tune_and_predict_ft

Source: https://github.com/yandex-research/graphpfn/blob/main/README.md

Performs a grid search over learning rates and returns predictions from the best trial. It updates the model weights on training nodes of the target graph.

```APIDOC
## tune_and_predict_ft

### Description
Performs a grid search over learning rates and returns predictions from the best trial. It updates the model weights on training nodes of the target graph.

### Method
Function call

### Endpoint
`tune_and_predict_ft(dataset, score_fn, loss_fn, device, lr_grid, n_steps, steps_per_epoch, patience, train_batch_size, model_kwargs, preprocessing_kwargs, amp, verbose)`

### Parameters
- **dataset** (object) - Required - The target graph dataset.
- **score_fn** (callable) - Required - Scalar metric used for validation and early stopping (higher = better).
- **loss_fn** (callable) - Required - Differentiable loss used for gradient updates.
- **device** (str) - Required - Device to run on (e.g., "cuda:0").
- **lr_grid** (list[float]) - Optional - List of learning rates to try; the trial with the best validation score wins. Default: 10 log-spaced values 5e-6 to 5e-4.
- **n_steps** (int) - Optional - Total training steps. `-1` means run until early stopping triggers. Default: `-1`.
- **steps_per_epoch** (int) - Optional - Steps between validation evaluations. Default: `10`.
- **patience** (int) - Optional - Early stopping patience in epochs. Default: `4`.
- **train_batch_size** (int) - Optional - Number of training nodes predicted per gradient step. Default: `1024`.
- **model_kwargs** (dict) - Optional - Same as in `predict_icl`. Default: `{}`.
- **preprocessing_kwargs** (dict) - Optional - Same as in `predict_icl`. Default: `{}`.
- **amp** (bool) - Optional - Same as in `predict_icl`. Default: `True`.
- **verbose** (bool) - Optional - Print validation scores during training. Default: `True`.

### Request Example
```python
from functools import partial
from graphpfn import tune_and_predict_ft
from graphpfn.inference.util import compute_score, compute_loss

preds, best_hparams = tune_and_predict_ft(
    dataset=dataset,
    score_fn=partial(compute_score, task_type=dataset.task_type),
    loss_fn=partial(compute_loss, task_type=dataset.task_type),
    device="cuda:0",
)
```

### Response
#### Success Response
Returns a tuple `(predictions, best_hyperparams)`. Predictions have the same shape as in `predict_icl`.

#### Response Example
```python
(preds, {"lr": 0.0005})
```
```

--------------------------------

### Run ICL inference with predict_icl

Source: https://github.com/yandex-research/graphpfn/blob/main/README.md

Run in-context learning inference with the pretrained GraphPFN model. No weight updates are performed; training nodes are passed as context and all other nodes are predicted in a single forward pass. Returns a NumPy array of predictions for all nodes: shape [n_nodes] for binary classification and regression, or [n_nodes, n_classes] for multiclass. The device parameter is required; amp enables bfloat16 autocast (default True).

```python
from graphpfn import predict_icl

preds = predict_icl(
    dataset=dataset,
    preprocessing_kwargs={"d_pca": 64},  # recommended for neural-embedding features
    amp=True,       # bfloat16 autocast (default: True)
    device="cuda:0",
)
```

--------------------------------

### Build GraphDataset from custom DGL graph

Source: https://github.com/yandex-research/graphpfn/blob/main/README.md

Builds a GraphDataset directly from a DGL graph and NumPy arrays. Requires src_nodes, dst_nodes, N, and the feature/target/mask arrays to be defined beforehand. targets must be float32 with NaN for unlabeled nodes.

```python
import dgl, numpy as np, torch
from graphpfn import GraphDataset

graph = dgl.graph((src_nodes, dst_nodes), num_nodes=N)
features = {"num": numerical_array, "cat": categorical_array}
targets = np.array(...)   # shape [N], float32; NaN for unlabeled nodes
masks = {"train": train_bool_array, "val": val_bool_array, "test": test_bool_array}

dataset = GraphDataset(
    name="my-dataset",
    graph=graph,
    features=features,
    targets=targets,
    masks=masks,
    task_type="binclass",
)
```

--------------------------------

### GraphPFN.train_step_forward

Source: https://github.com/yandex-research/graphpfn/blob/main/README.md

Samples a single random ensemble member and returns query predictions and targets in the transformed target space, ready for a loss and backward pass. Use this for custom finetuning loops instead of forward.

```APIDOC
## GraphPFN.train_step_forward

### Description
Samples a single random ensemble member and returns query predictions and targets in the transformed target space, ready for a loss and backward pass. Use this for custom finetuning loops instead of forward.

### Method
train_step_forward

### Endpoint
GraphPFN.train_step_forward

### Parameters
- **graph** (DGL Graph) - Required - The input graph.
- **features** (Tensor) - Required - Feature tensor of shape [n_nodes, n_features].
- **y_train** (Tensor) - Required - Training labels tensor of shape [n_train].
- **train_mask** (BoolTensor) - Required - Boolean mask of shape [n_nodes] indicating training nodes.
- **task_type** (str) - Required - Type of task, e.g., "multiclass".

### Response
Returns a tuple (query_preds, query_targets) in the transformed target space.

```

--------------------------------

### Replace dataset masks with custom splits

Source: https://github.com/yandex-research/graphpfn/blob/main/README.md

Replaces the masks of an existing GraphDataset with custom splits loaded from a .npz file. Useful for using custom train/val/test splits.

```python
custom_splits = dict(np.load("split.npz"))
dataset.masks = custom_splits
```

--------------------------------

### GraphPFN.forward

Source: https://github.com/yandex-research/graphpfn/blob/main/README.md

Performs a forward pass through the model, running an ensemble of feature/target permutations and aggregating the outputs.

```APIDOC
## GraphPFN.forward

### Description
Performs a forward pass through the model, running an ensemble of feature/target permutations and aggregating the outputs.

### Method
forward

### Endpoint
GraphPFN.forward

### Parameters
- **graph** (DGL Graph) - Required - The input graph.
- **features** (Tensor) - Required - Feature tensor of shape [n_nodes, n_features].
- **y_train** (Tensor) - Required - Training labels tensor of shape [n_train].
- **train_mask** (BoolTensor) - Required - Boolean mask of shape [n_nodes] indicating training nodes.
- **task_type** (str) - Required - Type of task, e.g., "multiclass".

### Request Example
```python
preds = model(
    graph=graph,                  # DGL Graph
    features=features,            # Tensor [n_nodes, n_features]
    y_train=y_train,              # Tensor [n_train]
    train_mask=train_mask,        # BoolTensor [n_nodes]
    task_type="multiclass",
)  # → Tensor [n_nodes, n_classes]
```

### Response
Returns a Tensor of shape [n_nodes, n_classes] with aggregated predictions.

```

--------------------------------

### predict_icl

Source: https://github.com/yandex-research/graphpfn/blob/main/README.md

Runs in-context learning (ICL) inference with the pretrained GraphPFN model. It takes a dataset as context and predicts all other nodes in a single forward pass without any weight updates.

```APIDOC
## [predict_icl] graphpfn.predict_icl

### Description
Runs in-context learning (ICL) inference with the pretrained GraphPFN model. It takes a dataset as context and predicts all other nodes in a single forward pass without any weight updates.

### Method
predict_icl

### Endpoint
graphpfn.predict_icl

### Parameters
#### Parameters
- **dataset** (object) - Required - The dataset object containing features and labels for training and test nodes.
- **preprocessing_kwargs** (dict) - Optional - Feature preprocessing options. See the preprocessing table below.
- **amp** (bool) - Optional - Enables bfloat16 autocast for faster inference on supported GPUs. Default: `True`.
- **device** (str) - Required - Target CUDA device, e.g. `"cuda:0"`.
- **model_kwargs** (dict) - Optional - Passed to `GraphPFN.from_pretrained`. Use `"checkpoint"` key to specify a local path or HuggingFace URI. Default: `{}`.

#### Preprocessing Options (via `preprocessing_kwargs`)
- **num_transform** (str) - Optional - Transform for `"num"` features. One of `"standard"`, `"quantile-normal"`, `"quantile-uniform"`, or `None`. Default: `"quantile-normal"`.
- **frac_transform** (str) - Optional - Transform for `"frac"` features. Same options as `num_transform`. Default: `"quantile-uniform"`.
- **cat_transform** (str) - Optional - Transform for `"cat"` features. One of `"ordinal"`, `"one-hot"`, or `None`. Default: `"ordinal"`.
- **other_transform** (str) - Optional - Transform for `"other"` features. Same options as `num_transform`. Default: `"standard"`.
- **d_pca** (int) - Optional - If set, all features are concatenated and reduced to `d_pca` dimensions with PCA. Recommended (e.g. `64`) when features are high-dimensional neural embeddings. Default: `None`.

### Request Example
```python
from graphpfn import predict_icl

preds = predict_icl(
    dataset=dataset,
    preprocessing_kwargs={"d_pca": 64},  # recommended for neural-embedding features
    amp=True,       # bfloat16 autocast (default: True)
    device="cuda:0",
)
```

### Response
#### Success Response
- **predictions** (numpy.ndarray) - A NumPy array of predictions for all nodes. Shape `[n_nodes]` for binary classification and regression, or `[n_nodes, n_classes]` for multiclass. Index with a mask to get per-split results.

#### Response Example
```python
# Example: binary classification
preds = np.array([0.1, 0.9, 0.4, ...])  # shape [n_nodes]

# Example: multiclass
preds = np.array([[0.1, 0.8, 0.1], ...])  # shape [n_nodes, n_classes]
```
```

--------------------------------

### compute_metrics

Source: https://github.com/yandex-research/graphpfn/blob/main/README.md

Computes a dict of all relevant metrics for the given task type.

```APIDOC
## compute_metrics

### Description
Computes a dict of all relevant metrics for the given task type.

### Method
Function call

### Endpoint
`compute_metrics(y_true, y_pred, task_type)`

### Parameters
- **y_true** (array-like) - Required - Ground truth labels.
- **y_pred** (array-like) - Required - Predicted values.
- **task_type** (str) - Required - Type of task: "binclass", "multiclass", or "regression".

### Response
#### Success Response
Returns a dict of metrics, e.g., `{"average_precision": ..., "roc_auc": ...}` for binclass.
```

--------------------------------

### GraphDataset dataclass definition

Source: https://github.com/yandex-research/graphpfn/blob/main/README.md

Shows the GraphDataset dataclass definition. Features are split by type (num, frac, cat, other) so the preprocessing pipeline can apply the right transform to each. Only keys present in the dict are processed.

```python
@dataclass
class GraphDataset:
    name: str
    graph: dgl.DGLGraph
    features: FeatureArrays   # TypedDict with keys: "num", "frac", "cat", "other"
    targets: np.ndarray       # shape [n_nodes]; NaN for unlabeled nodes
    masks: MaskArrays         # TypedDict with keys: "train", "val", "test"
    task_type: TaskType       # "binclass" | "multiclass" | "regression"
```

--------------------------------

### predict_ft

Source: https://github.com/yandex-research/graphpfn/blob/main/README.md

Updates the model weights on training nodes of the target graph using a single learning rate. This typically improves performance over ICL at the cost of more compute.

```APIDOC
## predict_ft

### Description
Updates the model weights on training nodes of the target graph using a single learning rate. This typically improves performance over ICL at the cost of more compute.

### Method
Function call

### Endpoint
`predict_ft(dataset, score_fn, loss_fn, lr, device, n_steps, steps_per_epoch, patience, train_batch_size, model_kwargs, preprocessing_kwargs, amp, verbose)`

### Parameters
- **dataset** (object) - Required - The target graph dataset.
- **score_fn** (callable) - Required - Scalar metric used for validation and early stopping (higher = better).
- **loss_fn** (callable) - Required - Differentiable loss used for gradient updates.
- **lr** (float) - Required - Learning rate for the Adam optimizer.
- **device** (str) - Required - Device to run on (e.g., "cuda:0").
- **n_steps** (int) - Optional - Total training steps. `-1` means run until early stopping triggers. Default: `-1`.
- **steps_per_epoch** (int) - Optional - Steps between validation evaluations. Default: `10`.
- **patience** (int) - Optional - Early stopping patience in epochs. Default: `4`.
- **train_batch_size** (int) - Optional - Number of training nodes predicted per gradient step. Default: `1024`.
- **model_kwargs** (dict) - Optional - Same as in `predict_icl`. Default: `{}`.
- **preprocessing_kwargs** (dict) - Optional - Same as in `predict_icl`. Default: `{}`.
- **amp** (bool) - Optional - Same as in `predict_icl`. Default: `True`.
- **verbose** (bool) - Optional - Print validation scores during training. Default: `True`.

### Response
#### Success Response
Returns predictions with the same shape as in `predict_icl`.
```

--------------------------------

### LimiXPredictor.predict

Source: https://github.com/yandex-research/graphpfn/blob/main/paper/vendor/limix/README.md

Performs inference on test data using the trained model, given training features and targets. Returns predictions as a numpy array.

```APIDOC
## LimiXPredictor.predict

### Description
Predicts target values for the test set based on the provided training data. The method uses the model's learned patterns to generate predictions for classification or regression tasks.

### Method
Function

### Signature
```python
def predict(self, x_train: np.ndarray, y_train: np.ndarray, x_test: np.ndarray) -> np.ndarray:
```

### Parameters
- **x_train** (np.ndarray) - Required - The input features of the training set.
- **y_train** (np.ndarray) - Required - The target variable of the training set.
- **x_test** (np.ndarray) - Required - The input features of the test set.

### Returns
- **np.ndarray** - The predicted target values for the test set.
```

--------------------------------

### compute_score

Source: https://github.com/yandex-research/graphpfn/blob/main/README.md

Computes a single scalar metric for validation and early stopping. Higher is always better. Defaults: average_precision (binclass), accuracy (multiclass), R^2 (regression).

```APIDOC
## compute_score

### Description
Computes a single scalar metric for validation and early stopping. Higher is always better. Defaults: average_precision (binclass), accuracy (multiclass), R^2 (regression).

### Method
Function call

### Endpoint
`compute_score(y_true, y_pred, task_type)`

### Parameters
- **y_true** (array-like) - Required - Ground truth labels.
- **y_pred** (array-like) - Required - Predicted values.
- **task_type** (str) - Required - Type of task: "binclass", "multiclass", or "regression".

### Response
#### Success Response
Returns a scalar score.
```

--------------------------------

### compute_loss

Source: https://github.com/yandex-research/graphpfn/blob/main/README.md

Computes a differentiable Tensor loss (cross-entropy or MSE) for gradient updates.

```APIDOC
## compute_loss

### Description
Computes a differentiable Tensor loss (cross-entropy or MSE) for gradient updates.

### Method
Function call

### Endpoint
`compute_loss(y_true, y_pred, task_type)`

### Parameters
- **y_true** (array-like) - Required - Ground truth labels.
- **y_pred** (array-like) - Required - Predicted values.
- **task_type** (str) - Required - Type of task: "binclass", "multiclass", or "regression".

### Response
#### Success Response
Returns a differentiable Tensor loss.
```
Tokens:
Raw