Evaluating Evo2 with GFMBench-API (BioNeMo Recipe)¶
This notebook evaluates the Evo2 genomic foundation model using the BioNeMo Recipes (bionemo-evo2 + NeMo/Megatron) on the GFMBench-API benchmark suite.
GFMBench-API is a framework for evaluating genomic foundation models. It exposes a single, model-agnostic API that decouples model development from benchmark tasks and metrics:
you implement inference methods once, and the same interface drives zero-shot scoring, supervised evaluation, and reporting across all tasks. See the bioRxiv preprint for the full benchmark design and task definitions.
This recipe evaluates Evo2 on both task categories in GFMBench-API: zero-shot tasks and linear probing for supervised tasks, following the Evo2 paper protocol. For supervised tasks, because GFMBench-API decouples model development from task-specific training, you interact with task modules to load training data, extract frozen Evo2 embeddings, train lightweight classifiers, and evaluate via the same task API.
Note: In practice, each developer can choose any supervised training strategy (full fine-tuning, LoRA, adapter heads, etc.) by interacting with GFMBench-API task modules the same way: load training data through the task API, train your model, then evaluate with
eval_test_set().
Prerequisites¶
- BioNeMo Evo2 package installed (
bionemo-evo2)
Evaluation Strategy¶
GFMBench-API contains 20 tasks in two categories:
Zero-Shot Tasks (no training required)¶
These use Evo2's autoregressive log-likelihood and embeddings directly:
- Log-Likelihood Ratio (LLR):
log P(variant_seq) - log P(reference_seq)scores variant pathogenicity - Embedding Cosine Similarity / L2 Distance: Compare representative embeddings of variant vs reference sequences
- SNV-specific metrics: Position-specific embedding and probability comparisons
| Task | Description | Type |
|---|---|---|
| VepevalClinvar | ClinVar SNV pathogenicity | SNV |
| IndelClinvar | ClinVar indel pathogenicity | Indel |
| BendVEPExpression | SNV expression effects | SNV |
| BendVEPDisease | SNV disease effects | SNV |
| SonglabClinvar | ClinVar SNV (likelihood) | SNV |
| BRCA1 | BRCA1 functional impact | SNV |
| TraitGymComplex | Complex trait variants | SNV |
| TraitGymMendelian | Mendelian disease variants | SNV |
| LrbPathogenicOmim | OMIM pathogenic variants | SNV |
| LoleveCausalEqtl | Causal eQTL indels | Indel |
Supervised Tasks (linear probing)¶
For these tasks we:
- Extract Evo2 embeddings (frozen backbone) for each training sample
- Train a lightweight
LogisticRegressionclassifier on the embeddings - Wrap the trained head + Evo2 into a model that returns label probabilities
- Evaluate via GFMBench-API
| Task | Type | Description |
|---|---|---|
| GuePromoterAll | Single-seq | Promoter prediction (2-class, 300bp) |
| GueSpliceSite | Single-seq | Splice site prediction (3-class, 400bp) |
| GueTranscriptionFactor | Single-seq | TF binding prediction (2-class) |
| VariantBenchmarksCoding | Variant-pair | Coding variant pathogenicity |
| VariantBenchmarksNonCoding | Variant-pair | Non-coding variant pathogenicity |
| VariantBenchmarksExpression | Variant-pair | Expression-affecting variants |
| VariantBenchmarksCommonVsRare | Variant-pair | Common vs rare variants |
| VariantBenchmarksMEQTL | Variant-pair | Methylation eQTL variants |
| VariantBenchmarksSQTL | Variant-pair | Splicing eQTL variants |
| LRBCausalEqtl (supervised) | Variant-pair | Causal eQTL (with tissue metadata) |
1. Setup GFMBench-API¶
Linux commands (run from your working directory):
git clone https://github.com/NVIDIA/GFMBench-api
export GFMBENCH_PATH=./GFMBench-api
cd "$GFMBENCH_PATH"
pip install -r basic_requirements.txt
export PYTHONPATH="$GFMBENCH_PATH"
The following Python code will do the same:
In [1]:
Copied!
import os
import sys
from pathlib import Path
# Clone target is relative to the notebook working directory; resolve once for stable absolute paths.
GFMBENCH_PATH_RELATIVE = Path("GFMBench-api")
GFMBENCH_PATH = GFMBENCH_PATH_RELATIVE.resolve()
# Clone only if GFMBench-api is not already present (safe to re-run the cell)
if not GFMBENCH_PATH.exists():
os.system("git clone https://github.com/NVIDIA/GFMBench-api")
else:
print(f"Using existing clone: {GFMBENCH_PATH}")
REQUIREMENTS = GFMBENCH_PATH / "basic_requirements.txt"
if not REQUIREMENTS.exists():
raise FileNotFoundError(f"{REQUIREMENTS} not found — run `git pull` in {GFMBENCH_PATH}")
print(f"Installing GFMBench-API deps from {REQUIREMENTS}")
os.system(f"{sys.executable} -m pip install -r {REQUIREMENTS}")
os.environ["GFMBENCH_PATH"] = str(GFMBENCH_PATH)
os.environ["PYTHONPATH"] = str(GFMBENCH_PATH)
sys.path.insert(0, str(GFMBENCH_PATH))
print(f"GFMBENCH_PATH={GFMBENCH_PATH}")
print(f"cwd={Path.cwd()}")
import os
import sys
from pathlib import Path
# Clone target is relative to the notebook working directory; resolve once for stable absolute paths.
GFMBENCH_PATH_RELATIVE = Path("GFMBench-api")
GFMBENCH_PATH = GFMBENCH_PATH_RELATIVE.resolve()
# Clone only if GFMBench-api is not already present (safe to re-run the cell)
if not GFMBENCH_PATH.exists():
os.system("git clone https://github.com/NVIDIA/GFMBench-api")
else:
print(f"Using existing clone: {GFMBENCH_PATH}")
REQUIREMENTS = GFMBENCH_PATH / "basic_requirements.txt"
if not REQUIREMENTS.exists():
raise FileNotFoundError(f"{REQUIREMENTS} not found — run `git pull` in {GFMBENCH_PATH}")
print(f"Installing GFMBench-API deps from {REQUIREMENTS}")
os.system(f"{sys.executable} -m pip install -r {REQUIREMENTS}")
os.environ["GFMBENCH_PATH"] = str(GFMBENCH_PATH)
os.environ["PYTHONPATH"] = str(GFMBENCH_PATH)
sys.path.insert(0, str(GFMBENCH_PATH))
print(f"GFMBENCH_PATH={GFMBENCH_PATH}")
print(f"cwd={Path.cwd()}")
Cloning into 'GFMBench-api'...
Installing GFMBench-API deps from /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt Requirement already satisfied: numpy==1.26.4 in /usr/local/lib/python3.12/dist-packages (from -r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 2)) (1.26.4) Requirement already satisfied: pandas>=1.3.3 in /usr/local/lib/python3.12/dist-packages (from -r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 3)) (2.2.3) Requirement already satisfied: torch<3,>=2.0.0 in /usr/local/lib/python3.12/dist-packages (from -r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 4)) (2.8.0a0+5228986c39.nv25.6) Requirement already satisfied: tqdm==4.67.3 in /usr/local/lib/python3.12/dist-packages (from -r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 5)) (4.67.3) Requirement already satisfied: datasets>=2.0.0 in /usr/local/lib/python3.12/dist-packages (from -r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 8)) (4.3.0) Requirement already satisfied: huggingface_hub==0.36.2 in /usr/local/lib/python3.12/dist-packages (from -r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 9)) (0.36.2) Requirement already satisfied: pyfaidx==0.9.0.4 in /usr/local/lib/python3.12/dist-packages (from -r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 10)) (0.9.0.4) Requirement already satisfied: pyarrow>=6.0.0 in /usr/local/lib/python3.12/dist-packages (from -r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 11)) (22.0.0) Requirement already satisfied: openpyxl==3.1.5 in /usr/local/lib/python3.12/dist-packages (from -r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 12)) (3.1.5) Requirement already satisfied: biopython==1.87 in /usr/local/lib/python3.12/dist-packages (from -r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 15)) (1.87) Requirement already satisfied: scikit-learn==1.7.2 in /usr/local/lib/python3.12/dist-packages (from -r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 16)) (1.7.2) Requirement already satisfied: filelock in /usr/local/lib/python3.12/dist-packages (from huggingface_hub==0.36.2->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 9)) (3.18.0) Requirement already satisfied: fsspec>=2023.5.0 in /usr/local/lib/python3.12/dist-packages (from huggingface_hub==0.36.2->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 9)) (2024.12.0) Requirement already satisfied: hf-xet<2.0.0,>=1.1.3 in /usr/local/lib/python3.12/dist-packages (from huggingface_hub==0.36.2->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 9)) (1.2.0) Requirement already satisfied: packaging>=20.9 in /usr/local/lib/python3.12/dist-packages (from huggingface_hub==0.36.2->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 9)) (24.2) Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.12/dist-packages (from huggingface_hub==0.36.2->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 9)) (6.0.2) Requirement already satisfied: requests in /usr/local/lib/python3.12/dist-packages (from huggingface_hub==0.36.2->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 9)) (2.32.3) Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.12/dist-packages (from huggingface_hub==0.36.2->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 9)) (4.14.0) Requirement already satisfied: et-xmlfile in /usr/local/lib/python3.12/dist-packages (from openpyxl==3.1.5->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 12)) (2.0.0) Requirement already satisfied: scipy>=1.8.0 in /usr/local/lib/python3.12/dist-packages (from scikit-learn==1.7.2->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 16)) (1.15.3) Requirement already satisfied: joblib>=1.2.0 in /usr/local/lib/python3.12/dist-packages (from scikit-learn==1.7.2->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 16)) (1.5.1) Requirement already satisfied: threadpoolctl>=3.1.0 in /usr/local/lib/python3.12/dist-packages (from scikit-learn==1.7.2->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 16)) (3.6.0) Requirement already satisfied: sympy>=1.13.3 in /usr/local/lib/python3.12/dist-packages (from torch<3,>=2.0.0->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 4)) (1.14.0) Requirement already satisfied: networkx in /usr/local/lib/python3.12/dist-packages (from torch<3,>=2.0.0->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 4)) (3.5) Requirement already satisfied: jinja2 in /usr/local/lib/python3.12/dist-packages (from torch<3,>=2.0.0->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 4)) (3.1.6) Requirement already satisfied: setuptools in /usr/local/lib/python3.12/dist-packages (from torch<3,>=2.0.0->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 4)) (78.1.1) Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.12/dist-packages (from pandas>=1.3.3->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 3)) (2.9.0) Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.12/dist-packages (from pandas>=1.3.3->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 3)) (2023.4) Requirement already satisfied: tzdata>=2022.7 in /usr/local/lib/python3.12/dist-packages (from pandas>=1.3.3->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 3)) (2025.2) Requirement already satisfied: dill<0.4.1,>=0.3.0 in /usr/local/lib/python3.12/dist-packages (from datasets>=2.0.0->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 8)) (0.4.0) Requirement already satisfied: httpx<1.0.0 in /usr/local/lib/python3.12/dist-packages (from datasets>=2.0.0->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 8)) (0.28.1) Requirement already satisfied: xxhash in /usr/local/lib/python3.12/dist-packages (from datasets>=2.0.0->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 8)) (3.6.0) Requirement already satisfied: multiprocess<0.70.17 in /usr/local/lib/python3.12/dist-packages (from datasets>=2.0.0->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 8)) (0.70.16) Requirement already satisfied: aiohttp!=4.0.0a0,!=4.0.0a1 in /usr/local/lib/python3.12/dist-packages (from fsspec[http]<=2025.9.0,>=2023.1.0->datasets>=2.0.0->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 8)) (3.12.7) Requirement already satisfied: anyio in /usr/local/lib/python3.12/dist-packages (from httpx<1.0.0->datasets>=2.0.0->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 8)) (4.9.0) Requirement already satisfied: certifi in /usr/local/lib/python3.12/dist-packages (from httpx<1.0.0->datasets>=2.0.0->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 8)) (2025.4.26) Requirement already satisfied: httpcore==1.* in /usr/local/lib/python3.12/dist-packages (from httpx<1.0.0->datasets>=2.0.0->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 8)) (1.0.9) Requirement already satisfied: idna in /usr/local/lib/python3.12/dist-packages (from httpx<1.0.0->datasets>=2.0.0->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 8)) (3.10) Requirement already satisfied: h11>=0.16 in /usr/local/lib/python3.12/dist-packages (from httpcore==1.*->httpx<1.0.0->datasets>=2.0.0->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 8)) (0.16.0) Requirement already satisfied: aiohappyeyeballs>=2.5.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.9.0,>=2023.1.0->datasets>=2.0.0->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 8)) (2.6.1) Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.9.0,>=2023.1.0->datasets>=2.0.0->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 8)) (1.3.2) Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.9.0,>=2023.1.0->datasets>=2.0.0->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 8)) (25.3.0) Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.9.0,>=2023.1.0->datasets>=2.0.0->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 8)) (1.6.0) Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.9.0,>=2023.1.0->datasets>=2.0.0->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 8)) (6.4.4) Requirement already satisfied: propcache>=0.2.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.9.0,>=2023.1.0->datasets>=2.0.0->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 8)) (0.3.1) Requirement already satisfied: yarl<2.0,>=1.17.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.9.0,>=2023.1.0->datasets>=2.0.0->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 8)) (1.20.0) Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.12/dist-packages (from python-dateutil>=2.8.2->pandas>=1.3.3->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 3)) (1.16.0) Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.12/dist-packages (from requests->huggingface_hub==0.36.2->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 9)) (3.4.2) Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.12/dist-packages (from requests->huggingface_hub==0.36.2->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 9)) (1.26.20) Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.12/dist-packages (from sympy>=1.13.3->torch<3,>=2.0.0->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 4)) (1.3.0) Requirement already satisfied: sniffio>=1.1 in /usr/local/lib/python3.12/dist-packages (from anyio->httpx<1.0.0->datasets>=2.0.0->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 8)) (1.3.1) Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.12/dist-packages (from jinja2->torch<3,>=2.0.0->-r /data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api/basic_requirements.txt (line 4)) (3.0.2) GFMBENCH_PATH=/data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples/GFMBench-api cwd=/data/sense/alarey/code/bionemo-framework/bionemo-recipes/recipes/evo2_megatron/examples
WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager, possibly rendering your system unusable. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv. Use the --root-user-action option if you know what you are doing and want to suppress this warning.
Note: You can also use GFMBench-API's example script,
usage_examples/run_benchmark.pyon Linux (from$GFMBENCH_PATH):python ./usage_examples/run_benchmark.py \ --model Evo2 \ --linear_prob \ --csv_path ./evo2_gfmbench_results.csv \ --report_algo_name evo2_run \ --root_data_dir_path ./benchmark_data
2. Evo2 Evaluation Pipeline¶
This notebook walks through the Evo2-paper linear probing protocol step by step: interact with GFMBench-API task modules to load supervised train/test data, extract frozen Evo2 embeddings, train LogisticRegression heads, and evaluate via task.eval_test_set().
Use this approach when you need full control over embedding extraction, probe training, or per-task customization.
2.1 Setup & Imports¶
In [2]:
Copied!
import json
import logging
import os
import time
import traceback
import warnings
from typing import List, Optional, Tuple
import numpy as np
import torch
from tqdm.auto import tqdm
warnings.filterwarnings("ignore", category=FutureWarning)
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)
import json
import logging
import os
import time
import traceback
import warnings
from typing import List, Optional, Tuple
import numpy as np
import torch
from tqdm.auto import tqdm
warnings.filterwarnings("ignore", category=FutureWarning)
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)
/usr/local/lib/python3.12/dist-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html from .autonotebook import tqdm as notebook_tqdm
2.2 Configuration¶
Adjust these settings for your environment. Paths below are relative to the current working directory.
In [3]:
Copied!
# ---------------------------------------------------------------------------
# Model configuration
# ---------------------------------------------------------------------------
# BioNeMo checkpoint tag — automatically downloaded + cached from NGC.
# Available tags:
# evo2/1b-8k-bf16:1.0 — 1B params, 8K context, bf16
# evo2/1b-8k:1.0 — 1B params, 8K context, fp8
# evo2/7b-8k:1.0 — 7B params, 8K context
# evo2/7b-1m:1.0 — 7B params, 1M context
# Can also be a local path to an extracted NeMo2 checkpoint directory.
# ---------------------------------------------------------------------------
# Data & output paths
# ---------------------------------------------------------------------------
DATA_ROOT = "benchmark_data" # Downloaded benchmark datasets (relative to cwd)
RESULTS_CSV = "evo2_gfmbench_results.csv" # Results output path (relative to cwd)
CKPT_TAG = "evo2/1b-8k-bf16:1.0"
# ---------------------------------------------------------------------------
# Evaluation settings
# ---------------------------------------------------------------------------
MAX_SEQ_LENGTH = 8192 # Max tokens per sequence
BATCH_SIZE = 8
MAX_NUM_SAMPLES = None # limit the num of samples to lower number (e.g. 100) for fast runs
# ---------------------------------------------------------------------------
# Linear probing settings (for supervised tasks)
# ---------------------------------------------------------------------------
LINEAR_PROBE_MAX_ITER = 5000
LINEAR_PROBE_C = 1.0
os.makedirs(DATA_ROOT, exist_ok=True)
print(f"Checkpoint tag: {CKPT_TAG}")
print(f"Model max seq: {MAX_SEQ_LENGTH} | Batch size: {BATCH_SIZE}")
print(f"Data root: {DATA_ROOT}")
print(f"Results CSV: {RESULTS_CSV}")
# ---------------------------------------------------------------------------
# Model configuration
# ---------------------------------------------------------------------------
# BioNeMo checkpoint tag — automatically downloaded + cached from NGC.
# Available tags:
# evo2/1b-8k-bf16:1.0 — 1B params, 8K context, bf16
# evo2/1b-8k:1.0 — 1B params, 8K context, fp8
# evo2/7b-8k:1.0 — 7B params, 8K context
# evo2/7b-1m:1.0 — 7B params, 1M context
# Can also be a local path to an extracted NeMo2 checkpoint directory.
# ---------------------------------------------------------------------------
# Data & output paths
# ---------------------------------------------------------------------------
DATA_ROOT = "benchmark_data" # Downloaded benchmark datasets (relative to cwd)
RESULTS_CSV = "evo2_gfmbench_results.csv" # Results output path (relative to cwd)
CKPT_TAG = "evo2/1b-8k-bf16:1.0"
# ---------------------------------------------------------------------------
# Evaluation settings
# ---------------------------------------------------------------------------
MAX_SEQ_LENGTH = 8192 # Max tokens per sequence
BATCH_SIZE = 8
MAX_NUM_SAMPLES = None # limit the num of samples to lower number (e.g. 100) for fast runs
# ---------------------------------------------------------------------------
# Linear probing settings (for supervised tasks)
# ---------------------------------------------------------------------------
LINEAR_PROBE_MAX_ITER = 5000
LINEAR_PROBE_C = 1.0
os.makedirs(DATA_ROOT, exist_ok=True)
print(f"Checkpoint tag: {CKPT_TAG}")
print(f"Model max seq: {MAX_SEQ_LENGTH} | Batch size: {BATCH_SIZE}")
print(f"Data root: {DATA_ROOT}")
print(f"Results CSV: {RESULTS_CSV}")
Checkpoint tag: evo2/1b-8k-bf16:1.0 Model max seq: 8192 | Batch size: 8 Data root: benchmark_data Results CSV: evo2_gfmbench_results.csv
2.3 Load Evo2 Model via BioNeMo¶
We use the BioNeMo Recipes (bionemo-evo2) to load Evo2. The notebook-compatible approach:
- Initializes Megatron parallel state manually (single-GPU, no NeMo Trainer)
- Creates the model from NeMo's
HYENA_MODEL_OPTIONS/MAMBA_MODEL_OPTIONSconfigs - Loads checkpoint weights via Megatron
dist_checkpointing(torch_dist format) - Registers a forward hook on
output_layerto capture embeddings (hidden states before the lm_head)
Supported architectures and sizes:
| Config | Arch | Hidden | Layers | Seq Len |
|---|---|---|---|---|
| 1b | Hyena | 1920 | 25 | 8,192 |
| 7b | Hyena | 4096 | 32 | 8,192 |
| 7b_arc_longcontext | Hyena | 4096 | 32 | 1,048,576 |
| 40b | Hyena | 8192 | 50 | 8,192 |
| 1b_nv / 7b_nv / 40b_nv | Hyena | varies | varies | varies |
| hybrid_mamba_8b | Mamba | - | - | - |
In [4]:
Copied!
from usage_examples.sanity_models.evo2_model import Evo2BioNeMoModel
print("Evo2BioNeMoModel class defined.")
from usage_examples.sanity_models.evo2_model import Evo2BioNeMoModel
print("Evo2BioNeMoModel class defined.")
2026-06-14 22:14:59,305 [INFO] The multistorageclient package is available.
Evo2BioNeMoModel class defined.
Instantiate the model¶
In [5]:
Copied!
t0 = time.time()
evo2_model = Evo2BioNeMoModel(
ckpt_tag=CKPT_TAG,
max_length=MAX_SEQ_LENGTH,
)
elapsed = time.time() - t0
print(f"\nLoaded in {elapsed:.1f}s")
print(f" Hidden dim: {evo2_model.hidden_dim}")
print(f" Max length: {evo2_model.max_length}")
# Quick sanity check
test_seqs = ["ACGTACGTACGTACGT", "NNNNACGTACGTNNNN"]
probs, embeddings, representative = evo2_model.infer_sequence_to_sequence(test_seqs)
print("\nSanity check - infer_sequence_to_sequence:")
print(f" token_probs shape: {probs.shape}")
print(f" seq_embeddings shape: {embeddings.shape if embeddings is not None else 'None'}")
print(f" representative shape: {representative.shape if representative is not None else 'None'}")
print(f" token_probs range: [{probs.min():.4f}, {probs.max():.4f}]")
if representative is not None:
print(f" representative norm: {np.linalg.norm(representative, axis=1)}")
# Verify variant scoring
var_probs = evo2_model.infer_variant_ref_sequences_to_labels_probs(
variant_sequences=["ACGTACGTACGTACGT"],
ref_sequences=["ACGTACGAACGTACGT"],
)
print("\ninfer_variant_ref_sequences_to_labels_probs:")
print(f" output shape: {var_probs.shape}")
print(f" probs: {var_probs}")
print(f" probs sum: {var_probs.sum():.6f}")
t0 = time.time()
evo2_model = Evo2BioNeMoModel(
ckpt_tag=CKPT_TAG,
max_length=MAX_SEQ_LENGTH,
)
elapsed = time.time() - t0
print(f"\nLoaded in {elapsed:.1f}s")
print(f" Hidden dim: {evo2_model.hidden_dim}")
print(f" Max length: {evo2_model.max_length}")
# Quick sanity check
test_seqs = ["ACGTACGTACGTACGT", "NNNNACGTACGTNNNN"]
probs, embeddings, representative = evo2_model.infer_sequence_to_sequence(test_seqs)
print("\nSanity check - infer_sequence_to_sequence:")
print(f" token_probs shape: {probs.shape}")
print(f" seq_embeddings shape: {embeddings.shape if embeddings is not None else 'None'}")
print(f" representative shape: {representative.shape if representative is not None else 'None'}")
print(f" token_probs range: [{probs.min():.4f}, {probs.max():.4f}]")
if representative is not None:
print(f" representative norm: {np.linalg.norm(representative, axis=1)}")
# Verify variant scoring
var_probs = evo2_model.infer_variant_ref_sequences_to_labels_probs(
variant_sequences=["ACGTACGTACGTACGT"],
ref_sequences=["ACGTACGAACGTACGT"],
)
print("\ninfer_variant_ref_sequences_to_labels_probs:")
print(f" output shape: {var_probs.shape}")
print(f" probs: {var_probs}")
print(f" probs sum: {var_probs.sum():.6f}")
[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0 [Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0 [Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0
2026-06-14 22:15:07,431 [INFO] Downloading/caching checkpoint: evo2/1b-8k-bf16:1.0
Using cached path='/root/.cache/bionemo/ea4a3f5c9c26d5edc10bdc85165c090ad0ff23ac2670d4f61244f5f0d9d5e817-nemo2_evo2_1b_8k_bf16.tar.gz.untar' from checked=PosixPath('/root/.cache/bionemo/ea4a3f5c9c26d5edc10bdc85165c090ad0ff23ac2670d4f61244f5f0d9d5e817-nemo2_evo2_1b_8k_bf16.tar.gz.checked')
2026-06-14 22:15:07,446 [INFO] Checkpoint root: /root/.cache/bionemo/ea4a3f5c9c26d5edc10bdc85165c090ad0ff23ac2670d4f61244f5f0d9d5e817-nemo2_evo2_1b_8k_bf16.tar.gz.untar
2026-06-14 22:15:08,593 [INFO] Note: detected 256 virtual cores but NumExpr set to maximum of 64, check "NUMEXPR_MAX_THREADS" environment variable.
2026-06-14 22:15:08,594 [INFO] Note: NumExpr detected 256 cores but "NUMEXPR_MAX_THREADS" not set, so enforcing safe limit of 16.
2026-06-14 22:15:08,595 [INFO] NumExpr defaulting to 16 threads.
Import of quick_gelu from megatron.core.fusions.fused_bias_geglu failed with: Traceback (most recent call last):
File "/usr/local/lib/python3.12/dist-packages/nemo/utils/import_utils.py", line 319, in safe_import_from
return getattr(imported_module, symbol), True
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: module 'megatron.core.fusions.fused_bias_geglu' has no attribute 'quick_gelu'
2026-06-14 22:15:10,263 [INFO] Import of quick_gelu from megatron.core.fusions.fused_bias_geglu failed with: Traceback (most recent call last):
File "/usr/local/lib/python3.12/dist-packages/nemo/utils/import_utils.py", line 319, in safe_import_from
return getattr(imported_module, symbol), True
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: module 'megatron.core.fusions.fused_bias_geglu' has no attribute 'quick_gelu'
[NeMo I 2026-06-14 22:15:10 nemo_logging:393] Padded vocab_size: 512, original vocab_size: 512, dummy tokens: 0.
2026-06-14 22:15:11,201 [INFO] Using hybrid override pattern 2026-06-14 22:15:11,202 [INFO] Hybrid allocation (S is hyena_short_conv, D is hyena_medium_conv, H is hyena, * is attention, 2026-06-14 22:15:11,202 [INFO] SDH*SDHSDH*SDHSDH*SDHSDH* 2026-06-14 22:15:11,203 [INFO] 7 heyna_short_conv layers in 25 total layers. 2026-06-14 22:15:11,204 [INFO] 7 heyna_medium_conv layers in 25 total layers. 2026-06-14 22:15:11,205 [INFO] 7 heyna layers in 25 total layers. 2026-06-14 22:15:11,206 [INFO] 4 attention layers in 25 total layers.
[NeMo I 2026-06-14 22:15:11 nemo_logging:393] Using byte-level tokenization
2026-06-14 22:15:11,440 [INFO] Model created: 1,108,204,800 parameters, hidden=1920 2026-06-14 22:15:14,001 [INFO] Checkpoint loaded successfully
Loaded in 6.8s Hidden dim: 1920 Max length: 8192 Sanity check - infer_sequence_to_sequence: token_probs shape: (2, 16) seq_embeddings shape: (2, 16, 1920) representative shape: (2, 1920) token_probs range: [0.0000, 0.9951] representative norm: [22.061846 17.920654] infer_variant_ref_sequences_to_labels_probs: output shape: (1, 2) probs: [[0.00686961 0.9931304 ]] probs sum: 1.000000
In [6]:
Copied!
from gfmbench_api.benchmark_report.benchmark_report import BenchmarkReport
# Common task config (passed to every GFMBench-API task constructor)
TASK_CONFIG = {
# Extraction window: min(task_default, this value). Keep well below MAX_SEQ_LENGTH
# to avoid quadratic attention OOM at large batch sizes.
"max_sequence_length": MAX_SEQ_LENGTH,
"batch_size": BATCH_SIZE,
"max_num_samples": MAX_NUM_SAMPLES,
# Fail fast while validating this recipe. If this is False, GFMBench-API
# converts model-interface exceptions into None outputs and metrics become N/A.
"disable_safe_model_call": True,
}
MODEL_NAME = CKPT_TAG.replace("/", "_").replace(":", "_")
# Initialize the benchmark report (loads existing CSV if present)
report = BenchmarkReport(csv_path=RESULTS_CSV)
print(f"Benchmark report initialized: {RESULTS_CSV}")
print(f"Model name for reporting: {MODEL_NAME}")
print(f"TASK_CONFIG: {TASK_CONFIG}")
from gfmbench_api.benchmark_report.benchmark_report import BenchmarkReport
# Common task config (passed to every GFMBench-API task constructor)
TASK_CONFIG = {
# Extraction window: min(task_default, this value). Keep well below MAX_SEQ_LENGTH
# to avoid quadratic attention OOM at large batch sizes.
"max_sequence_length": MAX_SEQ_LENGTH,
"batch_size": BATCH_SIZE,
"max_num_samples": MAX_NUM_SAMPLES,
# Fail fast while validating this recipe. If this is False, GFMBench-API
# converts model-interface exceptions into None outputs and metrics become N/A.
"disable_safe_model_call": True,
}
MODEL_NAME = CKPT_TAG.replace("/", "_").replace(":", "_")
# Initialize the benchmark report (loads existing CSV if present)
report = BenchmarkReport(csv_path=RESULTS_CSV)
print(f"Benchmark report initialized: {RESULTS_CSV}")
print(f"Model name for reporting: {MODEL_NAME}")
print(f"TASK_CONFIG: {TASK_CONFIG}")
Benchmark report initialized: evo2_gfmbench_results.csv
Model name for reporting: evo2_1b-8k-bf16_1.0
TASK_CONFIG: {'max_sequence_length': 8192, 'batch_size': 8, 'max_num_samples': None, 'disable_safe_model_call': True}
Define Benchmark Tasks¶
Choose which GFMBench-API tasks to evaluate in the cell below. Uncomment the tasks you want to run in zero_shot_tasks and/or supervised_tasks.
Each enabled task is instantiated when you run that cell, which may download and cache datasets under DATA_ROOT on the first run. Later runs reuse the cache.
Note: First-time downloads (reference genomes, HuggingFace datasets, parquet files, etc.) can take several minutes to tens of minutes per task, depending on network speed and how many tasks are enabled.
In [ ]:
Copied!
# Uncomment the import for each task you enable below.
# isort: off
# --- Zero-shot tasks imports ---
# from gfmbench_api.tasks.concrete.clinvar_indel_task import IndelClinvarTask
# from gfmbench_api.tasks.concrete.loleve_causal_eqtl_task import LoleveCausalEqtlTask
from gfmbench_api.tasks.concrete.brca1_task import BRCA1Task
from gfmbench_api.tasks.concrete.bend_vep_disease_task import BendVEPDisease
# from gfmbench_api.tasks.concrete.bend_vep_expression_task import BendVEPExpression
# from gfmbench_api.tasks.concrete.songlab_clinvar_task import SonglabClinvarTask
# from gfmbench_api.tasks.concrete.traitgym_complex_task import TraitGymComplexTask
# from gfmbench_api.tasks.concrete.traitgym_mendelian_task import TraitGymMendelianTask
# from gfmbench_api.tasks.concrete.lrb_pathogenic_omim_task import LrbVariantEffectPathogenicOmimTask
# --- Supervised tasks imports ---
# from gfmbench_api.tasks.concrete.gue_tf_all_task import GueTranscriptionFactorTask
# from gfmbench_api.tasks.concrete.gue_promoter_all_task import GuePromoterAllTask
from gfmbench_api.tasks.concrete.gue_splice_site_task import GueSpliceSiteTask
from gfmbench_api.tasks.concrete.variant_benchmarks_coding_task import VariantBenchmarksCodingTask
# from gfmbench_api.tasks.concrete.variant_benchmarks_non_coding_task import VariantBenchmarksNonCodingTask
# from gfmbench_api.tasks.concrete.variant_benchmarks_expression_task import VariantBenchmarksExpressionTask
# from gfmbench_api.tasks.concrete.variant_benchmarks_common_vs_rare_task import VariantBenchmarksCommonVsRareTask
# from gfmbench_api.tasks.concrete.variant_benchmarks_meqtl_task import VariantBenchmarksMEQTLTask
# from gfmbench_api.tasks.concrete.variant_benchmarks_sqtl_task import VariantBenchmarksSQTLTask
# from gfmbench_api.tasks.concrete.lrb_causal_eqtl_task import LRBCausalEqtlTask
# isort: on
# ---------------------------------------------------------------------------
# Zero-shot tasks (no training — model runs inference directly)
# ---------------------------------------------------------------------------
# Uncomment the tasks you want to run.
zero_shot_tasks = [
# IndelClinvarTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # ClinVar indels
# LoleveCausalEqtlTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # Causal eQTL indels
BRCA1Task(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # BRCA1 functional impact (SNV)
BendVEPDisease(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # SNV disease effects
# BendVEPExpression(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # SNV expression effects
# SonglabClinvarTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # ClinVar SNV (likelihood)
# TraitGymComplexTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # Complex trait variants
# TraitGymMendelianTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # Mendelian disease variants
# LrbVariantEffectPathogenicOmimTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # OMIM pathogenic variants
]
# ---------------------------------------------------------------------------
# Supervised tasks (linear probing — extract embeddings, train a classifier head)
# ---------------------------------------------------------------------------
# Uncomment the tasks you want to run.
supervised_tasks = [
# GueTranscriptionFactorTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # TF binding (single-seq, binary)
# GuePromoterAllTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # Promoter prediction (single-seq, binary)
GueSpliceSiteTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # Splice site (single-seq, 3-class)
VariantBenchmarksCodingTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # Coding variant pathogenicity
# VariantBenchmarksNonCodingTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # Non-coding variant pathogenicity
# VariantBenchmarksExpressionTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # Expression-affecting variants
# VariantBenchmarksCommonVsRareTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # Common vs rare variants
# VariantBenchmarksMEQTLTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # Methylation eQTL variants
# VariantBenchmarksSQTLTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # Splicing eQTL variants
# LRBCausalEqtlTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # Causal eQTL (with tissue metadata)
]
print(f"Zero-shot tasks enabled: {len(zero_shot_tasks)}")
for task in zero_shot_tasks:
print(f" - {task.get_task_name()}")
print(f"Supervised tasks enabled: {len(supervised_tasks)}")
for task in supervised_tasks:
print(f" - {task.get_task_name()}")
# Uncomment the import for each task you enable below.
# isort: off
# --- Zero-shot tasks imports ---
# from gfmbench_api.tasks.concrete.clinvar_indel_task import IndelClinvarTask
# from gfmbench_api.tasks.concrete.loleve_causal_eqtl_task import LoleveCausalEqtlTask
from gfmbench_api.tasks.concrete.brca1_task import BRCA1Task
from gfmbench_api.tasks.concrete.bend_vep_disease_task import BendVEPDisease
# from gfmbench_api.tasks.concrete.bend_vep_expression_task import BendVEPExpression
# from gfmbench_api.tasks.concrete.songlab_clinvar_task import SonglabClinvarTask
# from gfmbench_api.tasks.concrete.traitgym_complex_task import TraitGymComplexTask
# from gfmbench_api.tasks.concrete.traitgym_mendelian_task import TraitGymMendelianTask
# from gfmbench_api.tasks.concrete.lrb_pathogenic_omim_task import LrbVariantEffectPathogenicOmimTask
# --- Supervised tasks imports ---
# from gfmbench_api.tasks.concrete.gue_tf_all_task import GueTranscriptionFactorTask
# from gfmbench_api.tasks.concrete.gue_promoter_all_task import GuePromoterAllTask
from gfmbench_api.tasks.concrete.gue_splice_site_task import GueSpliceSiteTask
from gfmbench_api.tasks.concrete.variant_benchmarks_coding_task import VariantBenchmarksCodingTask
# from gfmbench_api.tasks.concrete.variant_benchmarks_non_coding_task import VariantBenchmarksNonCodingTask
# from gfmbench_api.tasks.concrete.variant_benchmarks_expression_task import VariantBenchmarksExpressionTask
# from gfmbench_api.tasks.concrete.variant_benchmarks_common_vs_rare_task import VariantBenchmarksCommonVsRareTask
# from gfmbench_api.tasks.concrete.variant_benchmarks_meqtl_task import VariantBenchmarksMEQTLTask
# from gfmbench_api.tasks.concrete.variant_benchmarks_sqtl_task import VariantBenchmarksSQTLTask
# from gfmbench_api.tasks.concrete.lrb_causal_eqtl_task import LRBCausalEqtlTask
# isort: on
# ---------------------------------------------------------------------------
# Zero-shot tasks (no training — model runs inference directly)
# ---------------------------------------------------------------------------
# Uncomment the tasks you want to run.
zero_shot_tasks = [
# IndelClinvarTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # ClinVar indels
# LoleveCausalEqtlTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # Causal eQTL indels
BRCA1Task(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # BRCA1 functional impact (SNV)
BendVEPDisease(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # SNV disease effects
# BendVEPExpression(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # SNV expression effects
# SonglabClinvarTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # ClinVar SNV (likelihood)
# TraitGymComplexTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # Complex trait variants
# TraitGymMendelianTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # Mendelian disease variants
# LrbVariantEffectPathogenicOmimTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # OMIM pathogenic variants
]
# ---------------------------------------------------------------------------
# Supervised tasks (linear probing — extract embeddings, train a classifier head)
# ---------------------------------------------------------------------------
# Uncomment the tasks you want to run.
supervised_tasks = [
# GueTranscriptionFactorTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # TF binding (single-seq, binary)
# GuePromoterAllTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # Promoter prediction (single-seq, binary)
GueSpliceSiteTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # Splice site (single-seq, 3-class)
VariantBenchmarksCodingTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # Coding variant pathogenicity
# VariantBenchmarksNonCodingTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # Non-coding variant pathogenicity
# VariantBenchmarksExpressionTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # Expression-affecting variants
# VariantBenchmarksCommonVsRareTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # Common vs rare variants
# VariantBenchmarksMEQTLTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # Methylation eQTL variants
# VariantBenchmarksSQTLTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # Splicing eQTL variants
# LRBCausalEqtlTask(root_data_dir_path=DATA_ROOT, task_config=TASK_CONFIG), # Causal eQTL (with tissue metadata)
]
print(f"Zero-shot tasks enabled: {len(zero_shot_tasks)}")
for task in zero_shot_tasks:
print(f" - {task.get_task_name()}")
print(f"Supervised tasks enabled: {len(supervised_tasks)}")
for task in supervised_tasks:
print(f" - {task.get_task_name()}")
BRCA1 dataset or reference genome not found. Downloading... Downloading and processing BRCA1 dataset...
2026-06-14 22:15:21,588 [INFO] Downloading from: https://github.com/ArcInstitute/evo2/raw/refs/heads/main/notebooks/brca1/GRCh37.p13_chr17.fna.gz
chrom pos ref alt score label 0 17 41276135 T G -0.372611 1 1 17 41276135 T C -0.045313 1 2 17 41276135 T A -0.108254 1 3 17 41276134 T G -0.277963 1 4 17 41276134 T C -0.388414 1 Saved BRCA1 DMS dataset with 3893 variants to parquet. Downloading reference genome... Progress: 100.0% (23.2 MB) Unzipping and normalizing FASTA header... Reference genome ready: benchmark_data/brca1/GRCh37.p13_chr17.fna Successfully downloaded BRCA1 dataset and reference to benchmark_data/brca1 Loading reference genome: benchmark_data/brca1/GRCh37.p13_chr17.fna
2026-06-14 22:15:24,742 [INFO] Reference genome not found. Downloading hg38.fa... 2026-06-14 22:15:24,743 [INFO] Downloading reference genome hg38.fa (~3GB)... 2026-06-14 22:15:24,743 [INFO] Source: https://hgdownload.soe.ucsc.edu/goldenPath/hg38/bigZips/hg38.fa.gz 2026-06-14 22:15:24,744 [INFO] This may take several minutes... 2026-06-14 22:15:24,745 [INFO] Downloading from: https://hgdownload.soe.ucsc.edu/goldenPath/hg38/bigZips/hg38.fa.gz
Extracting sequences (window size: 8192bp)... Successfully extracted 3893 sequence pairs Progress: 100.0% (983.7 MB)
2026-06-14 22:16:37,649 [INFO] Extracting hg38.fa.gz...
2026-06-14 22:16:55,504 [INFO] Reference genome saved to: benchmark_data/reference_genome/hg38.fa 2026-06-14 22:16:55,505 [INFO] Loading reference genome: benchmark_data/reference_genome/hg38.fa 2026-06-14 22:17:18,046 [INFO] Downloading bend_variant_effects_disease from BEND repository... 2026-06-14 22:17:18,048 [INFO] Downloading from: https://sid.erda.dk/share_redirect/aNQa0Oz2lY/data/variant_effects/variant_effects_disease.bed
Progress: 100.0% (20.4 MB)
2026-06-14 22:17:19,675 [INFO] Data saved to: benchmark_data/bend_variant_effects_disease
2026-06-14 22:17:19,909 [INFO] Extracting sequences (window size: 512bp)... 2026-06-14 22:17:36,816 [INFO] Successfully extracted 295495 sequence pairs 2026-06-14 22:17:37,884 [INFO] Downloading gue_splice_site from HuggingFace... 2026-06-14 22:17:37,887 [INFO] Dataset not found locally, downloading Saving the dataset (1/1 shards): 100%|██████████| 36496/36496 [00:00<00:00, 1763317.08 examples/s] Saving the dataset (1/1 shards): 100%|██████████| 4562/4562 [00:00<00:00, 887866.68 examples/s] Saving the dataset (1/1 shards): 100%|██████████| 4562/4562 [00:00<00:00, 899850.21 examples/s] 2026-06-14 22:17:41,539 [INFO] Data saved to: benchmark_data/gue_splice_site 2026-06-14 22:17:41,770 [INFO] GUE: wrote benchmark_data/gue_splice_site/train.csv from HuggingFace split 'train' 2026-06-14 22:17:41,797 [INFO] GUE: wrote benchmark_data/gue_splice_site/dev.csv from HuggingFace split 'dev' 2026-06-14 22:17:41,825 [INFO] GUE: wrote benchmark_data/gue_splice_site/test.csv from HuggingFace split 'test' 2026-06-14 22:17:42,458 [INFO] Downloading var_bench_coding_pathogenicity from HuggingFace... 2026-06-14 22:17:42,459 [INFO] Dataset not found locally, downloading Saving the dataset (1/1 shards): 100%|██████████| 82872/82872 [00:00<00:00, 290897.37 examples/s] 2026-06-14 22:17:45,945 [INFO] Data saved to: benchmark_data/var_bench_coding_pathogenicity Filter: 100%|██████████| 82872/82872 [00:00<00:00, 127208.08 examples/s] Filter: 100%|██████████| 82872/82872 [00:00<00:00, 132380.18 examples/s]
Zero-shot tasks enabled: 2 - brca1 - bend_variant_effects_disease Supervised tasks enabled: 2 - gue_splice_site - var_bench_coding_pathogenicity
2.5 Zero-Shot Evaluation¶
For zero-shot tasks, we pass the Evo2BioNeMoModel directly to GFMBench-API. The benchmark calls:
infer_sequence_to_sequence(sequences)->(token_probs, seq_embeddings, representative)- Computes LLR, cosine similarity, and L2 distance metrics automatically
- For SNV tasks, also calls
sequence_pos_to_prob_pos
No training is needed -- the model's pre-trained representations are evaluated as-is.
In [8]:
Copied!
zero_shot_results = {}
for task in zero_shot_tasks:
try:
t0 = time.time()
attrs = task.get_task_attributes()
print(f" Task: {task.get_task_name()}")
print(f" SNV only: {attrs.get('is_snv_only_variants', 'N/A')}")
print(f" Test samples: {len(task.test_dataset)}")
# Evaluate directly -- Evo2BioNeMoModel implements the full GFMBench-API interface
results = task.eval_test_set(evo2_model)
elapsed = time.time() - t0
print(f"\n Results ({elapsed:.1f}s):")
for metric, score in results.items():
val_str = f"{score:.4f}" if score is not None else "N/A"
print(f" {metric}: {val_str}")
# Store and save incrementally
zero_shot_results[task.get_task_name()] = results
report.add_scores(task.get_task_name(), MODEL_NAME, results)
report.save_csv()
except Exception as e:
print(f" FAILED: {e}")
traceback.print_exc()
finally:
torch.cuda.empty_cache()
print(f"\nZero-shot evaluation complete. {len(zero_shot_results)}/{len(zero_shot_tasks)} tasks succeeded.")
zero_shot_results = {}
for task in zero_shot_tasks:
try:
t0 = time.time()
attrs = task.get_task_attributes()
print(f" Task: {task.get_task_name()}")
print(f" SNV only: {attrs.get('is_snv_only_variants', 'N/A')}")
print(f" Test samples: {len(task.test_dataset)}")
# Evaluate directly -- Evo2BioNeMoModel implements the full GFMBench-API interface
results = task.eval_test_set(evo2_model)
elapsed = time.time() - t0
print(f"\n Results ({elapsed:.1f}s):")
for metric, score in results.items():
val_str = f"{score:.4f}" if score is not None else "N/A"
print(f" {metric}: {val_str}")
# Store and save incrementally
zero_shot_results[task.get_task_name()] = results
report.add_scores(task.get_task_name(), MODEL_NAME, results)
report.save_csv()
except Exception as e:
print(f" FAILED: {e}")
traceback.print_exc()
finally:
torch.cuda.empty_cache()
print(f"\nZero-shot evaluation complete. {len(zero_shot_results)}/{len(zero_shot_tasks)} tasks succeeded.")
Task: brca1 SNV only: True Test samples: 3893
Evaluating (Zero-Shot): 0%| | 0/487 [00:00<?, ?it/s]2026-06-14 22:17:54,870 [WARNING] Model did not provide required inputs for 'snv_variant_effect_prediction_masked_llr_auroc' metric. This metric will return None. 2026-06-14 22:17:54,871 [WARNING] Model did not provide required inputs for 'snv_variant_effect_prediction_masked_llr_auprc' metric. This metric will return None. Evaluating (Zero-Shot): 100%|██████████| 487/487 [18:24<00:00, 2.27s/it]
Results (1109.0s):
sum_probs_llr_auroc: 0.7423
sum_probs_llr_auprc: 0.5332
sequence_embeddings_cosinesim_auroc: 0.7495
sequence_embeddings_cosinesim_auprc: 0.5426
sequence_embeddings_l2_auroc: 0.7781
sequence_embeddings_l2_auprc: 0.5773
snv_variant_effect_cosinesim_auroc: 0.5328
snv_variant_effect_cosinesim_auprc: 0.2221
snv_variant_effect_prediction_masked_llr_auroc: N/A
snv_variant_effect_prediction_masked_llr_auprc: N/A
Task: bend_variant_effects_disease
SNV only: True
Test samples: 295495
Evaluating (Zero-Shot): 0%| | 0/36937 [00:00<?, ?it/s]2026-06-14 22:36:21,823 [WARNING] Model did not provide required inputs for 'snv_variant_effect_prediction_masked_llr_auroc' metric. This metric will return None. 2026-06-14 22:36:21,823 [WARNING] Model did not provide required inputs for 'snv_variant_effect_prediction_masked_llr_auprc' metric. This metric will return None. Evaluating (Zero-Shot): 100%|██████████| 36937/36937 [1:26:49<00:00, 7.09it/s]
Results (5213.1s):
sum_probs_llr_auroc: 0.9195
sum_probs_llr_auprc: 0.7422
sequence_embeddings_cosinesim_auroc: 0.8723
sequence_embeddings_cosinesim_auprc: 0.4737
sequence_embeddings_l2_auroc: 0.9115
sequence_embeddings_l2_auprc: 0.5955
snv_variant_effect_cosinesim_auroc: 0.5755
snv_variant_effect_cosinesim_auprc: 0.0822
snv_variant_effect_prediction_masked_llr_auroc: N/A
snv_variant_effect_prediction_masked_llr_auprc: N/A
Zero-shot evaluation complete. 2/2 tasks succeeded.
3.6 Supervised Evaluation via Linear Probing¶
For supervised tasks, Evo2 has no built-in classification head. We perform linear probing:
- Extract embeddings (frozen Evo2 backbone) for every training sample
- Single-seq tasks: mean-pooled embedding per sequence
- Variant-pair tasks: concatenation of
[variant_repr, ref_repr, variant_repr - ref_repr]
- Train a
LogisticRegressionon the extracted embeddings - Wrap the trained head + Evo2 into a model adapter that GFMBench-API can call
- Evaluate using the standard GFMBench-API
eval_test_set()API
This approach avoids any gradient-based fine-tuning of the large backbone.
In [9]:
Copied!
from sklearn.linear_model import LogisticRegression
from torch.utils.data import DataLoader
def extract_embeddings_single_seq(
model,
dataset,
batch_size: int,
) -> Tuple[np.ndarray, np.ndarray]:
"""Extract mean-pooled embeddings for single-sequence dataset.
Dataset items: (sequence, label, conditional_input)
Returns: (embeddings [N, hidden_dim], labels [N])
"""
loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
all_embeds, all_labels = [], []
for batch in tqdm(loader, desc="Extracting embeddings (single-seq)"):
sequences, labels, _ = batch
# infer_sequence_to_sequence returns (probs, seq_embeddings, representative)
_, _, representative = model.infer_sequence_to_sequence(list(sequences))
all_embeds.append(representative)
all_labels.append(np.array(labels))
return np.concatenate(all_embeds, axis=0), np.concatenate(all_labels, axis=0)
def extract_embeddings_variant_pair(
model,
dataset,
batch_size: int,
) -> Tuple[np.ndarray, np.ndarray]:
"""Extract embeddings for variant-pair dataset.
Dataset items: (variant_seq, ref_seq, label, conditional_input)
Returns: (embeddings [N, 3*hidden_dim], labels [N])
Embedding = [variant_repr | ref_repr | variant_repr - ref_repr]
"""
loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
all_embeds, all_labels = [], []
for batch in tqdm(loader, desc="Extracting embeddings (variant-pair)"):
var_seqs, ref_seqs, labels, _ = batch
_, _, var_repr = model.infer_sequence_to_sequence(list(var_seqs))
_, _, ref_repr = model.infer_sequence_to_sequence(list(ref_seqs))
# Concatenate: [var, ref, var-ref] to give the linear head richer signal
combined = np.concatenate([var_repr, ref_repr, var_repr - ref_repr], axis=1)
all_embeds.append(combined)
all_labels.append(np.array(labels))
return np.concatenate(all_embeds, axis=0), np.concatenate(all_labels, axis=0)
def train_linear_probe(
x_train: np.ndarray,
y_train: np.ndarray,
num_labels: int,
max_iter: int = 5000,
c: float = 1.0,
) -> LogisticRegression:
"""Train a LogisticRegression on extracted embeddings."""
multi_class = "multinomial" if num_labels > 2 else "auto"
clf = LogisticRegression(
max_iter=max_iter,
C=c,
solver="lbfgs",
multi_class=multi_class,
n_jobs=-1,
)
clf.fit(x_train, y_train)
train_acc = clf.score(x_train, y_train)
print(f" Linear probe train accuracy: {train_acc:.4f}")
return clf
print("Embedding extraction and linear probing utilities defined.")
from sklearn.linear_model import LogisticRegression
from torch.utils.data import DataLoader
def extract_embeddings_single_seq(
model,
dataset,
batch_size: int,
) -> Tuple[np.ndarray, np.ndarray]:
"""Extract mean-pooled embeddings for single-sequence dataset.
Dataset items: (sequence, label, conditional_input)
Returns: (embeddings [N, hidden_dim], labels [N])
"""
loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
all_embeds, all_labels = [], []
for batch in tqdm(loader, desc="Extracting embeddings (single-seq)"):
sequences, labels, _ = batch
# infer_sequence_to_sequence returns (probs, seq_embeddings, representative)
_, _, representative = model.infer_sequence_to_sequence(list(sequences))
all_embeds.append(representative)
all_labels.append(np.array(labels))
return np.concatenate(all_embeds, axis=0), np.concatenate(all_labels, axis=0)
def extract_embeddings_variant_pair(
model,
dataset,
batch_size: int,
) -> Tuple[np.ndarray, np.ndarray]:
"""Extract embeddings for variant-pair dataset.
Dataset items: (variant_seq, ref_seq, label, conditional_input)
Returns: (embeddings [N, 3*hidden_dim], labels [N])
Embedding = [variant_repr | ref_repr | variant_repr - ref_repr]
"""
loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
all_embeds, all_labels = [], []
for batch in tqdm(loader, desc="Extracting embeddings (variant-pair)"):
var_seqs, ref_seqs, labels, _ = batch
_, _, var_repr = model.infer_sequence_to_sequence(list(var_seqs))
_, _, ref_repr = model.infer_sequence_to_sequence(list(ref_seqs))
# Concatenate: [var, ref, var-ref] to give the linear head richer signal
combined = np.concatenate([var_repr, ref_repr, var_repr - ref_repr], axis=1)
all_embeds.append(combined)
all_labels.append(np.array(labels))
return np.concatenate(all_embeds, axis=0), np.concatenate(all_labels, axis=0)
def train_linear_probe(
x_train: np.ndarray,
y_train: np.ndarray,
num_labels: int,
max_iter: int = 5000,
c: float = 1.0,
) -> LogisticRegression:
"""Train a LogisticRegression on extracted embeddings."""
multi_class = "multinomial" if num_labels > 2 else "auto"
clf = LogisticRegression(
max_iter=max_iter,
C=c,
solver="lbfgs",
multi_class=multi_class,
n_jobs=-1,
)
clf.fit(x_train, y_train)
train_acc = clf.score(x_train, y_train)
print(f" Linear probe train accuracy: {train_acc:.4f}")
return clf
print("Embedding extraction and linear probing utilities defined.")
Embedding extraction and linear probing utilities defined.
Model Adapter for Linear Probing¶
The GFMBench-API supervised evaluation calls infer_sequence_to_labels_probs (single-seq tasks) or infer_variant_ref_sequences_to_labels_probs (variant-pair tasks) and expects [batch, num_labels] probability arrays.
We wrap Evo2 + the trained LogisticRegression to satisfy this interface.
In [10]:
Copied!
class Evo2LinearProbeModel:
"""Wraps Evo2 backbone + a trained LogisticRegression to satisfy the GFM-Bench model API.
For single-seq tasks: calls infer_sequence_to_labels_probs(sequences)
For variant-pair tasks: calls infer_variant_ref_sequences_to_labels_probs(var_seqs, ref_seqs)
Also delegates infer_sequence_to_sequence / sequence_pos_to_prob_pos /
infer_masked_sequence_to_token_probs to the underlying Evo2 model so that
any zero-shot metrics that the benchmark may also compute still work.
"""
def __init__(self, backbone, clf: LogisticRegression, is_variant_pair: bool):
"""Initialize adapter with Evo2 backbone and trained linear probe."""
self.backbone = backbone
self.clf = clf
self.is_variant_pair = is_variant_pair
# --- Single-sequence classification ---
def infer_sequence_to_labels_probs(
self,
sequences: List[str],
conditional_input=None,
) -> Optional[np.ndarray]:
"""Return class probabilities for single-sequence inputs."""
_, _, representative = self.backbone.infer_sequence_to_sequence(sequences)
probs = self.clf.predict_proba(representative)
return probs.astype(np.float32)
# --- Variant-pair classification ---
def infer_variant_ref_sequences_to_labels_probs(
self,
variant_sequences: List[str],
ref_sequences: List[str],
conditional_input=None,
) -> Optional[np.ndarray]:
"""Return class probabilities for variant/reference sequence pairs."""
_, _, var_repr = self.backbone.infer_sequence_to_sequence(variant_sequences)
_, _, ref_repr = self.backbone.infer_sequence_to_sequence(ref_sequences)
combined = np.concatenate([var_repr, ref_repr, var_repr - ref_repr], axis=1)
probs = self.clf.predict_proba(combined)
return probs.astype(np.float32)
# --- Delegate zero-shot methods to backbone ---
def infer_sequence_to_sequence(self, sequences, conditional_input=None):
"""Delegate sequence inference to the Evo2 backbone."""
return self.backbone.infer_sequence_to_sequence(sequences, conditional_input)
def sequence_pos_to_prob_pos(self, sequences, pos):
"""Delegate position-specific probability lookup to the Evo2 backbone."""
return self.backbone.sequence_pos_to_prob_pos(sequences, pos)
def infer_masked_sequence_to_token_probs(
self, sequences, variant_pos, variant_letters, reference_letters, conditional_input=None
):
"""Delegate masked-token probability inference to the Evo2 backbone."""
return self.backbone.infer_masked_sequence_to_token_probs(
sequences, variant_pos, variant_letters, reference_letters, conditional_input
)
print("Evo2LinearProbeModel adapter defined.")
class Evo2LinearProbeModel:
"""Wraps Evo2 backbone + a trained LogisticRegression to satisfy the GFM-Bench model API.
For single-seq tasks: calls infer_sequence_to_labels_probs(sequences)
For variant-pair tasks: calls infer_variant_ref_sequences_to_labels_probs(var_seqs, ref_seqs)
Also delegates infer_sequence_to_sequence / sequence_pos_to_prob_pos /
infer_masked_sequence_to_token_probs to the underlying Evo2 model so that
any zero-shot metrics that the benchmark may also compute still work.
"""
def __init__(self, backbone, clf: LogisticRegression, is_variant_pair: bool):
"""Initialize adapter with Evo2 backbone and trained linear probe."""
self.backbone = backbone
self.clf = clf
self.is_variant_pair = is_variant_pair
# --- Single-sequence classification ---
def infer_sequence_to_labels_probs(
self,
sequences: List[str],
conditional_input=None,
) -> Optional[np.ndarray]:
"""Return class probabilities for single-sequence inputs."""
_, _, representative = self.backbone.infer_sequence_to_sequence(sequences)
probs = self.clf.predict_proba(representative)
return probs.astype(np.float32)
# --- Variant-pair classification ---
def infer_variant_ref_sequences_to_labels_probs(
self,
variant_sequences: List[str],
ref_sequences: List[str],
conditional_input=None,
) -> Optional[np.ndarray]:
"""Return class probabilities for variant/reference sequence pairs."""
_, _, var_repr = self.backbone.infer_sequence_to_sequence(variant_sequences)
_, _, ref_repr = self.backbone.infer_sequence_to_sequence(ref_sequences)
combined = np.concatenate([var_repr, ref_repr, var_repr - ref_repr], axis=1)
probs = self.clf.predict_proba(combined)
return probs.astype(np.float32)
# --- Delegate zero-shot methods to backbone ---
def infer_sequence_to_sequence(self, sequences, conditional_input=None):
"""Delegate sequence inference to the Evo2 backbone."""
return self.backbone.infer_sequence_to_sequence(sequences, conditional_input)
def sequence_pos_to_prob_pos(self, sequences, pos):
"""Delegate position-specific probability lookup to the Evo2 backbone."""
return self.backbone.sequence_pos_to_prob_pos(sequences, pos)
def infer_masked_sequence_to_token_probs(
self, sequences, variant_pos, variant_letters, reference_letters, conditional_input=None
):
"""Delegate masked-token probability inference to the Evo2 backbone."""
return self.backbone.infer_masked_sequence_to_token_probs(
sequences, variant_pos, variant_letters, reference_letters, conditional_input
)
print("Evo2LinearProbeModel adapter defined.")
Evo2LinearProbeModel adapter defined.
Run Supervised Benchmarks¶
For each supervised task:
- Initialize the task (downloads data if needed)
- Extract embeddings from the training set
- Train a linear probe (LogisticRegression)
- Wrap into
Evo2LinearProbeModel - Run
task.eval_test_set(model)which uses the trained head for predictions
In [11]:
Copied!
supervised_results = {}
for task in supervised_tasks:
try:
t0 = time.time()
attrs = task.get_task_attributes()
is_variant = attrs["is_variant_effect_prediction"]
num_labels = attrs["num_labels"]
train_dataset = task.get_finetune_dataset()
if train_dataset is None:
print(" SKIPPED: no training data available")
continue
print(f" Task: {task.get_task_name()}")
print(f" Type: {'variant-pair' if is_variant else 'single-seq'}")
print(f" Num labels: {num_labels}")
print(f" Train samples: {len(train_dataset)}")
print(f" Test samples: {len(task.test_dataset)}")
# Step 1: Extract embeddings from training set
print("\n Step 1: Extracting training embeddings...")
if is_variant:
x_train, y_train = extract_embeddings_variant_pair(evo2_model, train_dataset, BATCH_SIZE)
else:
x_train, y_train = extract_embeddings_single_seq(evo2_model, train_dataset, BATCH_SIZE)
print(f" Embeddings shape: {x_train.shape}, Labels shape: {y_train.shape}")
# Step 2: Train linear probe
print("\n Step 2: Training linear probe...")
clf = train_linear_probe(x_train, y_train, num_labels, LINEAR_PROBE_MAX_ITER, LINEAR_PROBE_C)
# Step 3: Wrap into GFMBench-API compatible model
probe_model = Evo2LinearProbeModel(
backbone=evo2_model,
clf=clf,
is_variant_pair=is_variant,
)
# Step 4: Evaluate via GFMBench-API
print("\n Step 3: Evaluating on test set...")
results = task.eval_test_set(probe_model)
elapsed = time.time() - t0
print(f"\n Results ({elapsed:.1f}s):")
for metric, score in results.items():
val_str = f"{score:.4f}" if score is not None else "N/A"
print(f" {metric}: {val_str}")
# Store and save incrementally
supervised_results[task.get_task_name()] = results
report.add_scores(task.get_task_name(), MODEL_NAME, results)
report.save_csv()
# Free training embeddings
del x_train, y_train, clf, probe_model
except Exception as e:
print(f" FAILED: {e}")
traceback.print_exc()
finally:
torch.cuda.empty_cache()
print(f"\nSupervised evaluation complete. {len(supervised_results)}/{len(supervised_tasks)} tasks succeeded.")
supervised_results = {}
for task in supervised_tasks:
try:
t0 = time.time()
attrs = task.get_task_attributes()
is_variant = attrs["is_variant_effect_prediction"]
num_labels = attrs["num_labels"]
train_dataset = task.get_finetune_dataset()
if train_dataset is None:
print(" SKIPPED: no training data available")
continue
print(f" Task: {task.get_task_name()}")
print(f" Type: {'variant-pair' if is_variant else 'single-seq'}")
print(f" Num labels: {num_labels}")
print(f" Train samples: {len(train_dataset)}")
print(f" Test samples: {len(task.test_dataset)}")
# Step 1: Extract embeddings from training set
print("\n Step 1: Extracting training embeddings...")
if is_variant:
x_train, y_train = extract_embeddings_variant_pair(evo2_model, train_dataset, BATCH_SIZE)
else:
x_train, y_train = extract_embeddings_single_seq(evo2_model, train_dataset, BATCH_SIZE)
print(f" Embeddings shape: {x_train.shape}, Labels shape: {y_train.shape}")
# Step 2: Train linear probe
print("\n Step 2: Training linear probe...")
clf = train_linear_probe(x_train, y_train, num_labels, LINEAR_PROBE_MAX_ITER, LINEAR_PROBE_C)
# Step 3: Wrap into GFMBench-API compatible model
probe_model = Evo2LinearProbeModel(
backbone=evo2_model,
clf=clf,
is_variant_pair=is_variant,
)
# Step 4: Evaluate via GFMBench-API
print("\n Step 3: Evaluating on test set...")
results = task.eval_test_set(probe_model)
elapsed = time.time() - t0
print(f"\n Results ({elapsed:.1f}s):")
for metric, score in results.items():
val_str = f"{score:.4f}" if score is not None else "N/A"
print(f" {metric}: {val_str}")
# Store and save incrementally
supervised_results[task.get_task_name()] = results
report.add_scores(task.get_task_name(), MODEL_NAME, results)
report.save_csv()
# Free training embeddings
del x_train, y_train, clf, probe_model
except Exception as e:
print(f" FAILED: {e}")
traceback.print_exc()
finally:
torch.cuda.empty_cache()
print(f"\nSupervised evaluation complete. {len(supervised_results)}/{len(supervised_tasks)} tasks succeeded.")
Task: gue_splice_site Type: single-seq Num labels: 3 Train samples: 36496 Test samples: 4562 Step 1: Extracting training embeddings...
Extracting embeddings (single-seq): 100%|██████████| 4562/4562 [02:48<00:00, 27.13it/s]
Embeddings shape: (36496, 1920), Labels shape: (36496,) Step 2: Training linear probe... Linear probe train accuracy: 0.5877 Step 3: Evaluating on test set...
Evaluating: 100%|██████████| 571/571 [00:29<00:00, 19.14it/s]
Results (301.6s):
classification_accuracy: 0.5848
classification_mcc: 0.1851
classification_auroc: 0.7120
classification_auprc: 0.5081
Task: var_bench_coding_pathogenicity
Type: variant-pair
Num labels: 2
Train samples: 74143
Test samples: 8729
Step 1: Extracting training embeddings...
Extracting embeddings (variant-pair): 100%|██████████| 9268/9268 [29:57<00:00, 5.16it/s]
Embeddings shape: (74143, 5760), Labels shape: (74143,) Step 2: Training linear probe... Linear probe train accuracy: 0.7424 Step 3: Evaluating on test set...
Evaluating: 100%|██████████| 1092/1092 [03:32<00:00, 5.13it/s]
Results (2300.7s):
classification_accuracy: 0.7043
classification_mcc: 0.4027
classification_auroc: 0.7627
classification_auprc: 0.7379
Supervised evaluation complete. 2/2 tasks succeeded.
3.7 Results Summary¶
In [12]:
Copied!
import pandas as pd
# Zero-shot and supervised tasks report different metric families; keep tables separate.
ZERO_SHOT_METRIC_PREFIXES = ("sum_probs_", "sequence_embeddings_", "snv_variant_effect_")
SUPERVISED_METRIC_PREFIXES = ("classification_",)
def _filter_metric_columns(columns, prefixes):
return [c for c in columns if any(c.startswith(p) for p in prefixes)]
def _results_dict_to_table(results_dict, title, metric_prefixes=None):
"""Build a wide results table from eval dicts; drop empty metric columns."""
print(f"{title}")
print(f"{'=' * 70}")
if not results_dict:
print(" (no results)\n")
return None
rows = [{"task": task, **metrics} for task, metrics in results_dict.items()]
table = pd.DataFrame(rows).set_index("task")
table = table.apply(pd.to_numeric, errors="coerce")
if metric_prefixes:
cols = _filter_metric_columns(table.columns, metric_prefixes)
table = table[cols]
table = table.dropna(axis=1, how="all")
print(f"Tasks: {len(table)} | Metrics: {len(table.columns)}")
print(table.round(4).to_string())
print()
return table
pd.set_option("display.max_columns", None)
pd.set_option("display.width", 200)
pd.set_option("display.float_format", "{:.4f}".format)
print(f"GFMBench-API Results for {MODEL_NAME}\n")
zero_shot_summary = _results_dict_to_table(
zero_shot_results,
"Zero-shot tasks",
metric_prefixes=ZERO_SHOT_METRIC_PREFIXES,
)
supervised_summary = _results_dict_to_table(
supervised_results,
"Supervised tasks (linear probing)",
metric_prefixes=SUPERVISED_METRIC_PREFIXES,
)
import pandas as pd
# Zero-shot and supervised tasks report different metric families; keep tables separate.
ZERO_SHOT_METRIC_PREFIXES = ("sum_probs_", "sequence_embeddings_", "snv_variant_effect_")
SUPERVISED_METRIC_PREFIXES = ("classification_",)
def _filter_metric_columns(columns, prefixes):
return [c for c in columns if any(c.startswith(p) for p in prefixes)]
def _results_dict_to_table(results_dict, title, metric_prefixes=None):
"""Build a wide results table from eval dicts; drop empty metric columns."""
print(f"{title}")
print(f"{'=' * 70}")
if not results_dict:
print(" (no results)\n")
return None
rows = [{"task": task, **metrics} for task, metrics in results_dict.items()]
table = pd.DataFrame(rows).set_index("task")
table = table.apply(pd.to_numeric, errors="coerce")
if metric_prefixes:
cols = _filter_metric_columns(table.columns, metric_prefixes)
table = table[cols]
table = table.dropna(axis=1, how="all")
print(f"Tasks: {len(table)} | Metrics: {len(table.columns)}")
print(table.round(4).to_string())
print()
return table
pd.set_option("display.max_columns", None)
pd.set_option("display.width", 200)
pd.set_option("display.float_format", "{:.4f}".format)
print(f"GFMBench-API Results for {MODEL_NAME}\n")
zero_shot_summary = _results_dict_to_table(
zero_shot_results,
"Zero-shot tasks",
metric_prefixes=ZERO_SHOT_METRIC_PREFIXES,
)
supervised_summary = _results_dict_to_table(
supervised_results,
"Supervised tasks (linear probing)",
metric_prefixes=SUPERVISED_METRIC_PREFIXES,
)
GFMBench-API Results for evo2_1b-8k-bf16_1.0
Zero-shot tasks
======================================================================
Tasks: 2 | Metrics: 8
sum_probs_llr_auroc sum_probs_llr_auprc sequence_embeddings_cosinesim_auroc sequence_embeddings_cosinesim_auprc sequence_embeddings_l2_auroc sequence_embeddings_l2_auprc snv_variant_effect_cosinesim_auroc snv_variant_effect_cosinesim_auprc
task
brca1 0.7423 0.5332 0.7495 0.5426 0.7781 0.5773 0.5328 0.2221
bend_variant_effects_disease 0.9195 0.7422 0.8723 0.4737 0.9115 0.5955 0.5755 0.0822
Supervised tasks (linear probing)
======================================================================
Tasks: 2 | Metrics: 4
classification_accuracy classification_mcc classification_auroc classification_auprc
task
gue_splice_site 0.5848 0.1851 0.7120 0.5081
var_bench_coding_pathogenicity 0.7043 0.4027 0.7627 0.7379
Key AUROC metrics at a glance (aggregated by task type)¶
In [13]:
Copied!
def _print_auroc_aggregates(table, title):
"""Print mean AUROC per metric (per-task values are in the cell above)."""
print(f"{title}")
print(f"{'=' * 70}")
if table is None or table.empty:
print(" (no results)\n")
return
auroc_cols = [c for c in table.columns if "auroc" in c.lower()]
if not auroc_cols:
print(" (no AUROC metrics)\n")
return
auroc_df = table[auroc_cols].dropna(axis=1, how="all")
rows = []
for col in auroc_df.columns:
valid = auroc_df[col].dropna()
if len(valid) > 0:
rows.append(
{
"metric": col,
"mean_auroc": valid.mean(),
"n_tasks": len(valid),
}
)
if not rows:
print(" (no valid AUROC scores)\n")
return
agg = pd.DataFrame(rows).set_index("metric")
print(agg.round(4).to_string())
print()
print(f"Aggregated AUROC for Evo2 ({CKPT_TAG})\n")
_print_auroc_aggregates(zero_shot_summary, "Zero-shot (mean AUROC)")
_print_auroc_aggregates(supervised_summary, "Supervised (mean AUROC)")
def _print_auroc_aggregates(table, title):
"""Print mean AUROC per metric (per-task values are in the cell above)."""
print(f"{title}")
print(f"{'=' * 70}")
if table is None or table.empty:
print(" (no results)\n")
return
auroc_cols = [c for c in table.columns if "auroc" in c.lower()]
if not auroc_cols:
print(" (no AUROC metrics)\n")
return
auroc_df = table[auroc_cols].dropna(axis=1, how="all")
rows = []
for col in auroc_df.columns:
valid = auroc_df[col].dropna()
if len(valid) > 0:
rows.append(
{
"metric": col,
"mean_auroc": valid.mean(),
"n_tasks": len(valid),
}
)
if not rows:
print(" (no valid AUROC scores)\n")
return
agg = pd.DataFrame(rows).set_index("metric")
print(agg.round(4).to_string())
print()
print(f"Aggregated AUROC for Evo2 ({CKPT_TAG})\n")
_print_auroc_aggregates(zero_shot_summary, "Zero-shot (mean AUROC)")
_print_auroc_aggregates(supervised_summary, "Supervised (mean AUROC)")
Aggregated AUROC for Evo2 (evo2/1b-8k-bf16:1.0)
Zero-shot (mean AUROC)
======================================================================
mean_auroc n_tasks
metric
sum_probs_llr_auroc 0.8309 2
sequence_embeddings_cosinesim_auroc 0.8109 2
sequence_embeddings_l2_auroc 0.8448 2
snv_variant_effect_cosinesim_auroc 0.5541 2
Supervised (mean AUROC)
======================================================================
mean_auroc n_tasks
metric
classification_auroc 0.7374 2
3.8 Export¶
In [14]:
Copied!
# Final save
report.save_csv()
print(f"Results saved to: {RESULTS_CSV}")
json_path = RESULTS_CSV.replace(".csv", ".json")
json_summary = {
"model": MODEL_NAME,
"ckpt_tag": CKPT_TAG,
"max_seq_length": MAX_SEQ_LENGTH,
"task_config": TASK_CONFIG,
"zero_shot_results": zero_shot_results,
"supervised_results": supervised_results,
}
with open(json_path, "w") as f:
json.dump(json_summary, f, indent=2, default=str)
print(f"JSON summary saved to: {json_path}")
n_tasks = len(zero_shot_results) + len(supervised_results)
print(
f"\nDone! Evaluated {n_tasks} tasks total "
f"({len(zero_shot_results)} zero-shot, {len(supervised_results)} supervised)."
)
# Final save
report.save_csv()
print(f"Results saved to: {RESULTS_CSV}")
json_path = RESULTS_CSV.replace(".csv", ".json")
json_summary = {
"model": MODEL_NAME,
"ckpt_tag": CKPT_TAG,
"max_seq_length": MAX_SEQ_LENGTH,
"task_config": TASK_CONFIG,
"zero_shot_results": zero_shot_results,
"supervised_results": supervised_results,
}
with open(json_path, "w") as f:
json.dump(json_summary, f, indent=2, default=str)
print(f"JSON summary saved to: {json_path}")
n_tasks = len(zero_shot_results) + len(supervised_results)
print(
f"\nDone! Evaluated {n_tasks} tasks total "
f"({len(zero_shot_results)} zero-shot, {len(supervised_results)} supervised)."
)
Results saved to: evo2_gfmbench_results.csv JSON summary saved to: evo2_gfmbench_results.json Done! Evaluated 4 tasks total (2 zero-shot, 2 supervised).
In [ ]:
Copied!