Autoeval Overview

The biotrainer autoeval module allows automatically evaluating a protein language model on downstream tasks. This can give a better impression of the model performance, and if it is actually creating useful embeddings.

Getting Started

Please refer to the two tutorial notebooks to get started with supervised and zero-shot evaluation.

Downstream Tasks

All relevant files for the downstream tasks are downloaded before the evaluation to /{user_cache_dir}/biotrainer/autoeval.

PBC (Supervised)

The autoeval module provides direct support of the PBC framework, namely for the supervised tasks.

Quick description of each task, you can find more information in the PBC Readme:

• conservation: Predict residue conservation scores for proteins.
• disorder: Predict disorder scores for proteins.
• scl: Predict protein subcellular localization for mixed human and non-human proteins.
• secondary_structure: Predict secondary structure of proteins (includes multiple test sets from CASP and the ProtT5 paper).

FLIP (Supervised)

WARNING: FLIP datasets are currently still supported but not actively maintained.

The autoeval module provides a curated subset of datasets in FLIP. The tasks have been chosen to be as hard as possible for a prediction model. We provide a flip_conversion_script.py in the module that shows, how the datasets have been preprocessed and curated.

Quick description of each task, you can find more information in the split descriptions:

• aav: Predict adeno-associated virus fitness.
  • low_vs_high: Training on low fitness values, testing on high fitness values
  • two_vs_many: Training on two mutations max, testing on many mutations
• bind: Predict binding sites of ligands for protein residues.
  • from_publication: Dataset split used in this publication
• gb1: Predict a binding score for each variant of the GB1 protein.
  • low_vs_high: Training on low fitness values, testing on high fitness values
  • two_vs_many: Training on two mutations max, testing on others
• meltome: Predict meltome temperature of proteins.
  • mixed_split: Mixed human and non-human proteins
• scl: Predict protein subcellular localization.
  • mixed_hard: Mixed human and non-human proteins
• secondary_structure: Predict secondary structure of proteins
  • sampled: Only available FLIP split

PGYM (Zero-Shot)

The autoeval module provides support for the ProteinGym DMS Supervised datasets that include various deep mutational scanning (DMS) fitness scores for protein mutations. For a given model, a zero-shot method must be selected. Then, it is evaluated on all datasets. The ProteinGym datasets were split into three tasks (using the reference file:

• virus: Consists of all datasets for which the taxon in the reference file equals virus. PLMs are often trained without or with only limited viral sequences, so this separate evaluation is interesting to see the effect of data distribution shifts after training.
• non-virus: All other datasets.
• total: virus and non-virus combined.

Please use a different output directory for different zero-shot methods, as the report is currently only able to capture the results of a single method at once.

How to use

CLI

The autoeval pipeline can be run directly through the cli:

biotrainer autoeval --embedder-name embedder_name --framework framework_name [--min-seq-length min_length] [--max-seq-length max_length] [--use-half-precision]

Script - Supervised

You can also integrate autoeval into your scripts or training pipelines:

from typing import Iterator

import torch
import numpy as np
from tqdm import tqdm

from biotrainer.autoeval import autoeval_pipeline
from biotrainer.utilities import seed_all


class ExampleRandomEmbedder:
    def __init__(self, embedding_dim: int = 21):
        self.embedding_dim = embedding_dim
        # Pre-generate random state for faster random number generation
        self.rng = np.random.default_rng()

    def embed_per_residue(self, sequences: Iterator[str]):
        for sequence in tqdm(sequences):
            # Generate all embeddings at once using rng
            embedding = self.rng.random((len(sequence), self.embedding_dim), dtype=np.float32)
            yield sequence, torch.tensor(embedding)

    def embed_per_sequence(self, sequences: Iterator[str]):
        for sequence in tqdm(sequences):
            # Generate single embedding using rng
            embedding = self.rng.random(self.embedding_dim, dtype=np.float32)
            yield sequence, torch.tensor(embedding)


seed_all(42)

embedder = ExampleRandomEmbedder()

for progress in autoeval_pipeline(embedder_name="example_random_embedder",
                                  framework="pbc",
                                  custom_embedding_function_per_residue=lambda seqs: embedder.embed_per_residue(
                                      seqs),
                                  custom_embedding_function_per_sequence=lambda
                                          seqs: embedder.embed_per_sequence(seqs),
                                  ):
    print(progress)

Some deeper explanations:

• The autoeval_pipeline function is a generator function that yields AutoEvalProgress objects to track progress. Therefore, you must "do something" with the pipeline return values, in order to execute the pipeline. Simply calling autoeval_pipeline(...) will not run the pipeline, see python generators.
• The custom embedding functions take an iterable of sequences as strings and need to yield a sequence and the according embedding. This is to make sure that the correct embedding is assigned to the respective sequence.
• All embeddings are calculated at once at the beginning of the training, to avoid duplicated embedding.

Script - Zero-Shot

Here is a minimal example on how to use the autoeval_pipeline for zero-shot evaluation:

from biotrainer.autoeval import autoeval_pipeline
from biotrainer.bioengineer import ZeroShotMethod

print("Starting AutoEval pipeline...")

current_progress = None
for progress in autoeval_pipeline(embedder_name="facebook/esm2_t6_8M_UR50D",
                                  framework="pgym",
                                  zero_shot_method=ZeroShotMethod.MASKED_MARGINALS,
                                  device="cuda",
                                  ):
    print(f"AutoEvalProgress: " + str(progress))
    current_progress = progress

print("**FINISHED**")

current_progress.final_report.summary()

Report

After all tasks have been successfully finished, a report is created in the output directory. All metadata and model results are tracked there. See the autoeval_report class for more details.

The report can be summarized:

final_report = progress.final_report
final_report.summary()

and compared to other reports:

final_report.compare([other_reports], plot=True)

Visualization and Leaderboard [BETA]

You can visualize and compare autoeval reports using the autoeval dashboard.

Features:

• Load one or more autoeval_report_*.json files by providing directories (scanned recursively) or uploading files.
• Leaderboard: simple average rank across tasks per framework (PBC, PGYM).
• Detailed view: per-task metrics and training/validation loss curves (if present in report results).
• Compare: side-by-side comparison across multiple reports for supervised and zero-shot metrics.

API Interaction

To directly upload your report from a script to the dashboard, you can use the following function:

final_report.compare_with_public_leaderboard()

If successful, this will return a link to the autoeval dashboard that includes your report. Note that your report is deleted from the autoeval_service after 24h.

If you want to make your results public, you can also publish them on the leaderboard. Your results will then be permanently and publicly available.

To do so, you can use the following method:

final_report.publish(name="Your Name", email="your.mail@example.com", citation="https://doi.org/VALID-DOI_URL")

Note that your report is internally checked for validity and compliance with the autoeval standards. Only reports for models publicly available on huggingface are accepted at the moment. We will use your name and email to credit you for your contribution and to reach out to you in case of any questions.

Local execution

To run the dashboard locally:

pip install streamlit  # if not already installed
streamlit run run_autoeval-frontend.py