---
license: mit
base_model: facebook/esm2_t33_650M_UR50D
tags:
- biology
- peptide
- multi-task-learning
- protein
- classification
---

# MTL Peptide Classifier (22 Tasks)

Multi-Task Learning peptide classifier covering 22 binary peptide-activity tasks. Built on a frozen ESM-2 (650M) backbone with a parallel Transformer + CNN feature extractor and per-task heads, following a PDeepPP-inspired design.

## Architecture

- **Shared encoder**: frozen ESM-2 (`facebook/esm2_t33_650M_UR50D`, 650M params) + learnable base embedding, mixed at `esm_ratio=0.9`
- **Feature extraction (parallel)**: 4-layer Transformer + CNN (kernel=7, padding=3) → concatenated to 2560-dim features
- **Heads**: 22 binary classifiers (`2560 → 256 → 128 → 2`) with masked average pooling
- **Loss**: TIM (Threshold-Independent Multi-task) loss + label smoothing 0.1

## Tasks

| # | Task | Source |
|---|---|---|
| 1 | ACE_inhibitory | UniDL4BioPep |
| 2 | DPPIV_inhibitory | UniDL4BioPep |
| 3 | Bitter | UniDL4BioPep |
| 4 | Umami | UniDL4BioPep |
| 5 | Antimicrobial | UniDL4BioPep |
| 6 | Antimalarial (main) | UniDL4BioPep |
| 7 | Antimalarial_alt | UniDL4BioPep |
| 8 | Quorum_sensing | UniDL4BioPep |
| 9 | Anticancer (main) | UniDL4BioPep |
| 10 | Anticancer_alt | UniDL4BioPep |
| 11 | AntiMRSA | UniDL4BioPep |
| 12 | TTCA | UniDL4BioPep |
| 13 | BBP | UniDL4BioPep |
| 14 | Anti_parasitic | UniDL4BioPep |
| 15 | NeuroPred | UniDL4BioPep |
| 16 | Antibacterial | UniDL4BioPep |
| 17 | Antifungal | UniDL4BioPep |
| 18 | Antiviral | UniDL4BioPep |
| 19 | Toxicity | UniDL4BioPep |
| 20 | Anti_inflammatory | local dataset |
| 21 | Signal_peptide | local dataset |
| 22 | Antioxidant | UniDL4BioPep (antioxidant_FRS) |

## Usage

```python
import os
from huggingface_hub import hf_hub_download
import torch
from transformers import EsmTokenizer

from mtl_peptide_classifier import MTLPeptideClassifier, get_all_peptide_tasks

REPO = "tuankg1028/MTL-Peptide-Classifier"
checkpoint_dir = "MTL-Peptide-Classifier"
os.makedirs(checkpoint_dir, exist_ok=True)

for fname in ["heads.pt", "shared_backbone.pt", "ablation_config.json"]:
    hf_hub_download(repo_id=REPO, filename=fname, local_dir=checkpoint_dir)

tokenizer = EsmTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
task_configs = get_all_peptide_tasks("datasets")  # needs local datasets/ dir for task names

model = MTLPeptideClassifier(
    task_configs=task_configs,
    hidden_dim=1280,
    esm_ratio=0.9,
    num_transformer_layers=4,
    dropout=0.3,
    use_transformer=True,
    use_cnn=True,
    unfreeze_esm=False,
)

device = "cuda" if torch.cuda.is_available() else "cpu"
backbone = torch.load(f"{checkpoint_dir}/shared_backbone.pt", map_location=device)
heads = torch.load(f"{checkpoint_dir}/heads.pt", map_location=device)

model.base_embed.load_state_dict(backbone["base_embed"])
if "transformer" in backbone:
    model.transformer.load_state_dict(backbone["transformer"])
if "cnn" in backbone:
    model.cnn.load_state_dict(backbone["cnn"])
    model.layer_norm.load_state_dict(backbone["layer_norm"])
for name, head in model.heads.items():
    if name in heads:
        head.load_state_dict(heads[name])

model = model.to(device).eval()

sequence = "MKWVTFISLLFLFSSAYSRGVFRR"
tokens = " ".join(list(sequence))
inputs = tokenizer(tokens, return_tensors="pt", max_length=128, padding="max_length", truncation=True)
with torch.no_grad():
    logits = model(inputs["input_ids"].to(device), inputs["attention_mask"].to(device), task_name="Antimicrobial")
    probs = torch.softmax(logits, dim=-1)
```

## Training

- Base model: `facebook/esm2_t33_650M_UR50D` (frozen)
- Batch size: 16, learning rate: 1e-4, 50 epochs, dropout: 0.3
- 3-way split per task: 80% train / 20% val (checkpoint selection) / held-out test CSV evaluated once
- Mixed precision, gradient clipping 1.0, cosine LR with 5 warmup epochs
- TIM loss + label smoothing 0.1

## Files

- `heads.pt` — per-task classifier heads
- `shared_backbone.pt` — base embedding, Transformer, CNN, LayerNorm
- `ablation_config.json` — architecture configuration for reproducibility
- `test_results.json` — held-out test metrics (per task + averages)
- `mtl_peptide_classifier.py` — model code.

## Requirements

```
torch>=2.0.0
transformers>=4.30.0
huggingface_hub
numpy
pandas
scikit-learn
```