deepFRI2

deepFRI2 is an upgraded version of the well-established deepFRI (Deep Functional Residue Identification) framework for predicting protein function using Gene Ontology (GO) terms and Enzyme Commission (EC) numbers.

Like its predecessor, deepFRI2 operates in two complementary modes: sequence-based and sequence–structure-based. This dual approach enables robust functional inference in metagenomic settings — where protein structures are often unavailable — as well as structure-informed functional annotation when structural information is available.

For training, deepFRI2 leverages FRIData, a scalable and efficient library for generating large, non-redundant protein datasets.

While maintaining similar input and output formats, the model architecture has been completely redesigned to incorporate recent advances in the field, particularly the use of protein language models as powerful representations of protein sequences. It consists of two submodules: sequence analyzer (utilizing ESM embeddings and lightweight attention pooling) and structural prober (processing distograms with shallow convolutional network). Signals from both models are merged using an ESM-conditioned gating mechanism. deepFRI2 outputs sequence-, structure-, and fusion-based predictions for each ontology (MF, CC, BP). The architecture is intentionally simple yet robust, enabling accurate functional annotation while maintaining high scalability and interpretability.

Installation

Prerequisites

• conda or mamba
• For GPU inference: an NVIDIA GPU with a driver supporting CUDA 12.6
• For CPU inference: no additional requirements

Install deepFRI2

# Clone the repository
git clone https://github.com/Tomasz-Lab/deepFRI2.git
cd deepFRI2

# Create the environment (choose ONE)
conda env create -f environment-gpu.yml    # Recommended (GPU)
# conda env create -f environment-cpu.yml  # CPU-only

# Activate the environment
conda activate deepfri2

# Download the ESM-2 language model (~2.5 GB)
python src/deepFRI2/download_esm.py

Notes

• GPU environment is named deepfri2, whereas CPU environment is named deepfri2_cpu (name can be changed in the corressponding .yml file).
• The deepFRI2 model checkpoints are already included under params/<ontology>/. Only the ESM-2 weights (downloaded in the last step) need to be fetched.
• Once the ESM-2 weights are downloaded, all inference runs entirely offline.

Usage

Predict GO terms for a folder of protein structures (.cif / .pdb):

python src/deepFRI2/deepfri2.py --input path/to/structures

Or run the sequence model on a FASTA file of sequences (no structures needed):

python src/deepFRI2/deepfri2.py --input sequences.fasta

Options (run python src/deepFRI2/deepfri2.py --help for the full list):

Flag	Description
-i, --input	Either a folder with .cif / .pdb structures or a single FASTA file of sequences. A FASTA input runs the sequence model only: --model and --ids_file are ignored, and --threshold uses only its first value. (required)
-o, --output_dir	Folder for results (default: <repo>/results).
-f, --ids_file	Text file listing structures to run, one per line. Each entry is an id, optionally with a .cif / .pdb extension and/or a relative subfolder path (abCD, abCD.cif, sub/abCD, sub1/sub2/abCD.cif), resolved under the input folder. An id without an extension resolves to .cif if present, else .pdb. Default: all top-level files in the input folder.
-a, --aspect	Comma-separated GO aspects (ontologies) to run: any of MF, CC, BP (case-insensitive). Default: mf,cc,bp.
-m, --model	Which model to run: sequence (embeddings only), structure (distograms only) or fusion (both). Only the needed inputs are computed and outputs carry only that model's columns. Default: fusion.
-b, --batch_size	Proteins per inference batch (default: 32).
-t, --threshold	Keep a GO term in the summary if any model (sequence, structure, fusion) scores ≥ threshold. Either one float applied to all models (0.1) or two comma-separated floats applied to fusion/sequence and structure respectively (0.1,0.2). The structural prober is trained with a different loss and outputs higher probabilities on average, hence the higher default for it. 0 (or 0,0) keeps everything; 1 (or 1,1) keeps nothing (default: 0.1,0.2).
-k, --top_k	Maximum GO terms per protein in the summary (default: all selected).
-p, --prop	Propagate scores up the GO hierarchy. Adds the preds_propagated/ folder and the propagated columns to the summary (default: off).
-s, --summary	Write only prediction_summary.csv, skipping the preds/ (and preds_propagated/) folders (default: off). With --prop, the summary still includes the propagated columns.
-v, --verbose	Enable debug logging (default: off).

Example

Run a subset of structures with a stricter (global) thresholds:

python src/deepFRI2/deepfri2.py -i structures/ -o results/run1 -f ids.txt -t 0.3

Notes

• GPU environment is the default and recommended one.
• The model runs in float32 on both CPU and GPU, so CPU and GPU results agree closely — differences are at the float32-rounding level (~1e-4 on structure/fusion probabilities, ~1e-6 on sequence), driven only by the BLAS/cuDNN kernels. GO ranking and thresholded term sets are effectively identical.
• In the current setup, the model processes up to 1,020 aa. Longer proteins are truncated during inference (no functional signal beyond 1,020 aa), but ESM embeddings are generated for the full protein, which may be time-consuming. Therefore, it is recommended to run the model on functionally relevant domains. There is no lower limit; however, the structural prober is not sensitive to proteins shorter than 60 aa (in which case predictions equal the mean across the training data).
• The current version was trained on gapless structures, so fully resolved inputs (no missing residues) are recommended. For structures with gaps, a missing residue zeroes out its whole row/column in the residue–residue similarity map, breaking the backbone-adjacency band and pushing the structure model out of distribution. As a rough safeguard we fill only the immediate -1/+1 off-diagonals at gap positions with 1 (a non-zero, "these consecutive residues are neighbours" signal).

Output

Predictions are written under the output folder for ontologies selected by the --aspect argument (MF, CC, BP by default):

• prediction_summary.csv — top predicted GO terms per protein, with raw scores for the fusion, structure, and sequence models (and GO-hierarchy-propagated scores when --prop is set).
• preds/<protein>__<ontology>.csv — full per-term probabilities (fusion / structure / sequence + gate). Omitted when --summary is set.
• preds_propagated/<protein>__<ontology>.csv — full per-term probabilities after GO-hierarchy propagation (only when --prop is set; omitted when --summary is set).
• log.txt — the run log.

For a quick overview of predicted functions, please take a look at the pred_prob (raw probabilities) column in prediction_summary.csv — or, when you run with --prop, the pred_prop_prob column (consistent probabilities i.e., the more general the term, the higher its probability). In some cases, it is also useful to check purely structure- and sequence-based outputs (see struct_prob, seq_prob etc.). For a downstream analysis, you may wish to check the full output in preds (and, with --prop, preds_propagated) folders.

Runtime

End-to-end runtime (excluding model loading at startup, which usually takes a couple of seconds per run) depends primarily on the available compute resources and the protein length. Initial benchmarks with the default settings (batch size: 32) yielded the following throughput:

• 0.2–0.4 s/protein — GPU (NVIDIA A100)
• 0.7–1.5 s/protein — CPU (48-core server)

These measurements were obtained on protein datasets with median sequence lengths of 150–440 amino acids. Additional benchmarking is underway, and the results will be shared in future updates.

For large-scale inference, we recommend a GPU or a multi-core CPU cluster. On CPU, ESM embeddings are computed one sequence at a time, each forward using all available core.

Running the model on a personal computer (e.g., a laptop) is also possible. Initial tests on an Apple M3 Pro (11 CPU cores, 18 GB RAM) with a small set of proteins (median length ~150 aa; batch size: 32) took ~1.8 s/protein for embedding generation and ~32 s/protein for model inference — the structure model is markedly slower here because Apple-Silicon PyTorch ships a generic (non-MKL) CPU build. Local CPU inference is therefore best suited to small runs (e.g., select a subset of .cif/.pdb files with --ids_file) or with sequence-only mode (in such case, inference time drastically decreases to 2 s/protein).

Future releases

The model is still under development. We will soon add (among other things):

• interpretability module
• architectural details
• detailed benchmarks

In the nearest future we also plan to share the whole training pipeline in a fully reproducible manner.

Troubleshooting

If you run into installation problems, find a bug, or would like to propose an improvement, please raise an issue or write directly to p.szczerbiak[at]sanoscience.org.