sasudo2/cafa5_models

CAFA 5 Protein Function Prediction Models

This repository contains trained models, configurations, and inference scripts developed for multi-label Gene Ontology (GO) protein function prediction, aligned with the CAFA 5 tournament guidelines.

---

📋 Table of Contents

1. Overview
2. Model Artifacts & Architectures
   • esm2_ffn (ESM-2 + Feed-Forward Network Head)
   • blast_knn (BLAST & K-Nearest Neighbors Baseline)
   • lr3mer (Logistic Regression 3-mer Baseline)
3. Input & Output Schemas
4. Quick Usage Example

---

1. Overview

Predicting protein functions across Gene Ontology (GO) domains (Biological Process, Molecular Function, and Cellular Component) is crucial for functional annotation pipelines. This repository hosts multiple modeling approaches—ranging from deep learning embeddings (esm2_ffn) to sequence homology and baseline approaches (blast_knn, lr3mer)—to generate standardized functional predictions (EntryID, term, score).

---

2. Model Artifacts & Architectures

esm2_ffn

A hybrid model combining a frozen protein language model encoder with a custom feed-forward classification head.

• Base Model Encoder: facebook/esm2_t33_650M_UR50D (1,280 feature dimensions).
• Training Hardware: Tesla T4 GPU.
• Training Stats: Evaluated across 141,995 training proteins and 141,864 test proteins, generating over 200 million prediction rows.
• Head Configuration:
  • Hidden Dimension: 512
  • Dropout: 0.2
  • Learning Rate: 0.001
  • Weight Decay: 0.0001
  • Positive Weight Cap: 50.0

blast_knn

A sequence-similarity and neighborhood-based baseline utilizing sequence alignment scores and k-nearest neighbors over annotation spaces.

• Configuration: Configured via blast_knn/config.json and tracked with blast_knn/manifest.json.

lr3mer

A lightweight sequence-based baseline utilizing 3-mer frequency representations coupled with multi-label logistic regression classifiers.

• Configuration: Backed by sequence k-mer tracking matrices and lr3mer/go_terms.json.

---

3. Input & Output Schemas

• Required Input Fields:
  • EntryID: Unique identifier for the protein sequence.
  • sequence: Amino acid sequence string.
• Prediction Output Schema:
  • EntryID: Protein identifier.
  • term: Gene Ontology term ID (e.g., GO:0009987).
  • score: Confidence score ranging between $0.0$ and $1.0$.

---

4. Quick Usage Example

You can inspect, download, or execute inference for the esm2_ffn model bundle directly using the provided inference.py script:

import torch
from pathlib import Path
from esm2_ffn.inference import load_model, resolve_model_dir

# Load model bundle from the Hub
model_dir = resolve_model_dir("sasudo2/cafa5_models", subfolder="esm2_ffn")
print(f"Model bundle resolved at: {model_dir}")

# Example inference script execution
# python esm2_ffn/inference.py --model-dir {model_dir} --input proteins.fasta