rishivk144/enzyme-hybrid-sp

Leak CURBER enzyme classification: DIAMOND kNN over training + Swiss-Prot, with an ESM-2 neural fallback

Predicts one fully specified EC number per enzyme sequence. It uses train.parquet, the external dataset UniProtKB/Swiss-Prot (release 2026_03) under a leakage cap described below, and the pretrained ESM-2 8M protein language model (facebook/esm2_t6_8M_UR50D) as a frozen feature extractor. It runs on CPU with numpy; PyTorch is needed only for training.

External data (disclosure)

Dataset	UniProtKB/Swiss-Prot, reviewed entries
Release	2026_03 (2 September 2026), from the X-UniProt-Release response header
Licence	CC BY 4.0 (UniProt Consortium)
Download	curl -sS -D swissprot_headers.txt -o swissprot.tsv.gz "https://rest.uniprot.org/uniprotkb/stream?compressed=true&format=tsv&fields=accession%2Cec%2Corganism_name%2Csequence&query=reviewed%3Atrue"
Processing	python prepare_external.py --train ../data/train.parquet --swissprot swissprot.tsv.gz
Accessions used	weights/ext_accessions.json (one list of UniProt accessions per indexed sequence)

Filtering performed by prepare_external.py (counts also in weights/ext_meta.json):

Step	Remaining
Entries downloaded	575,748
With at least one fully specified EC in ec_vocabulary.json (partial ECs such as 1.1.1.- dropped)	238,288
Unique sequences (entries with identical sequences merged, union of ECs)	202,065
Not identical to any training sequence (170,080 identical ones dropped; the training label is used)	31,985 indexed

Leakage control

Swiss-Prot is the source the benchmark itself was built from: 9,663 of the 11,367 unique validation inputs (85%) appear in it verbatim, and the test inputs very likely do too. Using it naively would copy the answers.

The benchmark separates its training set from its hidden splits by sequence similarity: validation inputs are mostly 30–55% identical to their closest training sequence. predict.py applies the same rule to the external data at prediction time. Any Swiss-Prot hit that is ≥ 50% identical to the input is discarded before voting. This applies identically to validation and test inputs, without needing to know the test set. Exact copies (100%) and close relatives are never used. External data can only contribute relatives at most as similar as the training set is allowed to be, including the 1,387 ECs that appear in Swiss-Prot but not in training.

The external search is run separately from the training search, keeping up to 100 hits before the cap. In a shared search, the many discarded near-copies of a validation input (86% of validation inputs have one) would otherwise crowd usable hits out of the top 25.

Pipeline

1. Quick alignment pass. DIAMOND default mode aligns every unique input against the 186,488 training sequences.
2. Upgrade pass. DIAMOND very-sensitive mode in batches of 3,000, replacing quick-pass hits while time remains.
3. External pass. DIAMOND very-sensitive mode against the 31,985 Swiss-Prot sequences (default mode if time is short). Hits ≥ 50% identical to the input are dropped, and the rest are merged with the training hits.
4. Nearest-neighbour vote. The 25 strongest hits (training and external together) vote for their recorded ECs with weight (bitscore / best bitscore) ** 16, split 1/n over a hit's n ECs.
5. Rescue pass. DIAMOND ultra-sensitive mode against training, for inputs still without hits. It runs only if it can finish in the time left.
6. Neural pass, for inputs that no search could match (about 0.4%):
   • ESM-2 8M embedding. Run in numpy, checked against PyTorch to a maximum difference of 2e-6. The output is the mean of the last layer over residues.
   • MLP classifier. 320 → 1024 → 1024 → 4,614 ECs, with GELU and dropout 0.3, trained with soft targets (1/n per recorded EC) on the embeddings of all training sequences. It is trained on training data only.
7. Fallback. The most frequent training EC (2.7.7.48), used only if every stage fails.

When time is short, the priority order is: upgrade, then external, then rescue, then neural.

Local validation

Held-out sequence families: MMseqs2 clusters of the training set at 50% identity and 80% coverage; a holdout of 1,092 whole clusters (11,804 sequences). Scores are reweighted so each top-hit identity band counts as often as in the real validation inputs, using the official 1/n credit rule. Swiss-Prot sequences identical to any training sequence (including the holdout) are excluded from the external index, and the 50% cap applies to holdout inputs exactly as to validation inputs.

Model	L4	L3	L2	L1
Most frequent EC (starter baseline)	0.011	0.055	0.143	0.381
ESM-2 8M + MLP alone	0.681	0.796	0.814	0.855
DIAMOND kNN on training only	0.817	0.928	0.940	0.970
+ Swiss-Prot hits below 30% identity	0.819	0.928	0.941	0.971
+ Swiss-Prot hits below 40% identity	0.835	0.950	0.961	0.977
+ Swiss-Prot hits below 50% identity (this model)	0.837	0.951	0.962	0.978
+ Swiss-Prot unrestricted, for comparison only (leaky)	0.839	0.952	0.963	0.979

The cap costs only 0.2 L4 points on the holdout, so the gain comes from distant relatives and new ECs rather than near-copies. An end-to-end holdout run of predict.py scored higher (0.845), but only because the packaged ESM-2 classifier was trained on all training sequences including the holdout. The 0.837 figure is the fair estimate.

Leaderboard (validation) results of the earlier models: DIAMOND only 0.78097; this pipeline without Swiss-Prot 0.78123.

Error analysis (training-only model, holdout)

Lost L4 credit	Points
True EC never appears in training	5.3
Several true ECs, but one prediction scores at most 1/n	6.5
EC in training, but no DIAMOND hit carries it	3.0
A hit carries it, but the vote picked another EC	3.5

What did not help

• Using ESM-2 more widely. A perfect per-input switch between DIAMOND and ESM-2 would reach 0.825. Every confidence-based switching rule tried scored at or below DIAMOND alone.
• Replacing the rescue pass with ESM-2. This looked better locally on a handful of inputs, but scored lower on the leaderboard (0.78045).
• A learned listwise re-ranker over DIAMOND's candidate ECs (an MLP on 21 features per candidate, trained on cross-family searches of 4 training folds). It scored 0.8172 against 0.8167 for the hand-tuned vote.

Runtime and sandbox

• Grader limits. The grader allows 256 open files per process. DIAMOND therefore uses 32 temp bins (about 37 open files, largest temp file about 6 MiB), retrying with 16 or 8 bins on memory or open-file errors.
• Time. On the grader's 2 CPUs: about 40 s quick, about 240 s upgrade, about 60 s external, about 100 s rescue, and about 10 s neural, around 420 s in total. All stages stop 70 s before the 10-minute limit.
• Logging. Sandbox limits, per-stage timings, errors, and how many inputs' votes came from training only, Swiss-Prot only, or both.

Files

• predict.py: inference.
• prepare_external.py: builds the Swiss-Prot reference from the UniProt download.
• esm_numpy.py: ESM-2 forward pass in numpy.
• train.py: rebuilds the training DIAMOND index and labels.
• embed.py and train_neural.py: rebuild the ESM-2 embeddings and the MLP classifier (requires PyTorch and transformers).
• weights/train.dmnd, labels.json, meta.json: the training index and recorded ECs.
• weights/ext.dmnd, ext_labels.json, ext_accessions.json, ext_meta.json: the Swiss-Prot index, its ECs, accessions and processing record.
• weights/esm2_8M.npz: ESM-2 8M encoder weights (MIT licence, Meta AI).
• weights/ec_mlp.npz, ec_mlp_classes.json: the trained classifier.
• bin/diamond: DIAMOND v2.1.11 Linux binary (GPL-3.0; source at github.com/bbuchfink/diamond).
• inputs.py, task.json, ec_vocabulary.json: unchanged from the starter.

EC vocabulary: ExPASy ENZYME release 02-Sep-2026, SIB, CC BY 4.0.