A LoRA-fine-tuned ESM2 model for multi-label Gene Ontology (GO) term prediction on the Biological Process Ontology (BPO), trained on the CAFA6 competition dataset. Given a raw protein amino-acid sequence, the model outputs confidence scores for 5,333 BPO GO terms.

Model Details

Model Description

This model adapts ESM2-8M (facebook/esm2_t6_8M_UR50D) with Low-Rank Adaptation (LoRA) for protein function prediction. The encoder is kept frozen except for the injected LoRA matrices; a two-layer MLP classification head maps mean-pooled residue embeddings to per-GO-term sigmoid scores.

Developed by: Afzal Hossan
Model type: Protein sequence encoder + multi-label classifier
Language(s): Protein sequences (amino-acid alphabet)
License: MIT
Finetuned from: facebook/esm2_t6_8M_UR50D
Ontology covered: Biological Process (BPO) — 5,333 GO terms

Model Sources

Repository: AfzalHosaan-2005021/cafa6-esm2-lora-bpo
Competition: CAFA6 on Kaggle

Uses

Direct Use

Load the model and run inference on any protein sequence to obtain BPO GO term confidence scores:

python

import torch
import sys
sys.path.insert(0, "src")

from models.esm2_classifier import ESM2Classifier
from transformers import AutoTokenizer

model = ESM2Classifier.from_pretrained("AfzalHosaan-2005021/cafa6-esm2-lora-bpo")
model.eval()

tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D")

sequence = "MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGDGTQDNLSGAEKAVQVKVKALPDAQFEVVHSLAKWKRQTLGQHDFSAGEGLYTHMKALRPDEDRLSPLHSVYVDQWDWERVMGDGERQFSTLKSTVEAIWAGIKATEAAVSEEFGLAPFLPDQIHFVHSQELLSRYPDLDAKGRERAIAKDLGAVFLVGIGGKLSDGHRHDVRAPDYDDWSTPSELGHAGLNGDILVWNPVLEDAFELSSMGIRVDADTLKHQLALTGDEDRLELEWHQALLRGEMPQTIGGGIGQSRLTMLLLQLPHIGQVQAGVWPAAVRESVPSLL"

inputs = tokenizer(sequence, return_tensors="pt")
with torch.no_grad():
    out = model(inputs["input_ids"], inputs["attention_mask"])
    scores = torch.sigmoid(out["logits"])[0]

go_terms = model.config.go_terms
threshold = 0.3
predicted = [(go_terms[i], scores[i].item()) for i in range(len(go_terms)) if scores[i] > threshold]
predicted.sort(key=lambda x: -x[1])
for term, score in predicted[:10]:
    print(f"{term}  {score:.3f}")