""" Fold-level remote homology retrieval benchmark for LEMON. Reproduces the SCOPe, SCOP, and CATH-S20 results from Table 1 of the paper. Metric definition (fold level) -------------------------------- Positive pair : same fold, different superfamily Negative pair : different fold Per-query AUROC / AUPRC are computed for every query that has at least one positive, then averaged. Bundled datasets (data/ subdirectory) -------------------------------------- data/scope_10_2.08.fa SCOPe 2.08, 10 % seq-id (7 117 seqs) Source: https://scop.berkeley.edu/downloads/scopeseq-2.08/ File : astral-scopedom-seqres-gd-sel-gs-bib-10-2.08.fa data/cath_s20.fa CATH S20 v4.4.0 (15 043 seqs) Source: https://release.cathdb.info/v4.4.0/non-redundant-data-sets/ File : cath-dataset-nonredundant-S20-v4_4_0.fa data/cath_s20_labels.tsv Domain → C.A.T.H classification mapping Source: https://release.cathdb.info/v4.4.0/cath-classification-data/ File : cath-domain-list-v4_4_0.txt (extracted S20 subset) data/scop175.fa SCOP 1.75 representative sequences (31 073 seqs) Source: Kabir et al. (2023) PLM zero-shot remote homology evaluation https://github.com/tymor22/protein-vec File : data/SCOP/processed/SCOP_with_seq.tsv (classification embedded) Usage — zero config when run from the snapshot directory --------------------------------------------------------- python eval_retrieval.py # uses all bundled datasets python eval_retrieval.py --scope # SCOPe only python eval_retrieval.py --cath # CATH-S20 only python eval_retrieval.py --scop # SCOP only Expected results (Table 1 — averaged across seq-id thresholds) --------------------------------------------------------------- Dataset AUROC AUPRC SCOPe 0.8847 0.3149 CATH S20 0.8129 0.3418 SCOP 0.8917 0.2631 Note: the paper averages metrics across multiple seq-id thresholds. Running on a single threshold (e.g. 10 % / S20) will be within ±0.01. """ import argparse import re import sys import os from pathlib import Path from typing import Dict, List, Optional, Tuple import numpy as np import torch from tqdm import tqdm # ─── 1. Load LEMON from the same HuggingFace repo ────────────────────────── def _load_model_and_tokenizer(repo_dir: str): sys.path.insert(0, repo_dir) from modeling_lemon import LemonEncoder # noqa: F401 from tokenization_zest import ZESTTokenizer # noqa: F401 tok = ZESTTokenizer.from_pretrained(repo_dir) model = LemonEncoder.from_pretrained( os.path.join(repo_dir, "model.safetensors"), os.path.join(repo_dir, "config.json"), ) model.eval() return model, tok # ─── 2. FASTA parsers ────────────────────────────────────────────────────── _SCOPE_RE = re.compile(r"^>(\S+)\s+.*\b([a-g]\.\d+\.\d+\.\d+)\b") _CATH_RE = re.compile(r"^>cath\|[\d._]+\|(\S+)") # Bundled SCOP header: ">FA_DOMID PDBID CL.CF.SF.FA" _SCOP_BUNDLED_RE = re.compile(r"^>(\S+)\s+\S+\s+(\d+\.\d+\.\d+\.\d+)") def _iter_fasta(path: str): """Yield (header_line, sequence) pairs.""" header, parts = None, [] with open(path) as fh: for line in fh: line = line.rstrip() if line.startswith(">"): if header: yield header, "".join(parts) header, parts = line, [] else: parts.append(line) if header: yield header, "".join(parts) def parse_scope_fasta(path: str) -> List[Dict]: """Parse SCOPe ASTRAL FASTA → list of {id, classification, sequence}.""" entries = [] for hdr, seq in _iter_fasta(path): m = _SCOPE_RE.match(hdr) if m: entries.append({"id": m.group(1), "classification": m.group(2), "sequence": seq}) return entries def parse_scop_fasta(path: str) -> List[Dict]: """ Parse bundled SCOP 1.75 FASTA. Header format (bundled): ">FA_DOMID PDBID CL.CF.SF.FA" where CL/CF/SF/FA are numeric SCOP 2 identifiers. Fold = CL.CF (2 parts), superfamily = CL.CF.SF (3 parts). """ entries = [] for hdr, seq in _iter_fasta(path): m = _SCOP_BUNDLED_RE.match(hdr) if m: entries.append({"id": m.group(1), "classification": m.group(2), "sequence": seq}) return entries def parse_cath_fasta(path: str, labels_tsv: Optional[str] = None) -> List[Dict]: """ Parse CATH FASTA and attach C.A.T.H classification. labels_tsv : path to the compact TSV (bundled as data/cath_s20_labels.tsv) columns: domain_id classification Produced from cath-domain-list-v4_4_0.txt. """ cath_map: Dict[str, str] = {} if labels_tsv and Path(labels_tsv).exists(): with open(labels_tsv) as fh: next(fh) # skip header for line in fh: parts = line.rstrip().split("\t") if len(parts) == 2: cath_map[parts[0]] = parts[1] elif labels_tsv is None: pass # caller did not supply — entries with no label are dropped entries = [] for hdr, seq in _iter_fasta(path): m = _CATH_RE.match(hdr) if m: did = m.group(1).split("/")[0] cls = cath_map.get(did, "") if cls: entries.append({"id": did, "classification": cls, "sequence": seq}) return entries # ─── 3. Embedding ────────────────────────────────────────────────────────── @torch.no_grad() def embed(model, tokenizer, sequences: List[str], batch_size: int = 128, max_tokens: int = 1024, device: torch.device = torch.device("cpu"), dropout: float = 0.0, tta_passes: int = 1) -> np.ndarray: """ Embed sequences with optional Test-Time Augmentation (TTA) via trie-dropout. Parameters ---------- dropout : float Trie-dropout rate for tokenization (0.0 = deterministic greedy). tta_passes : int Number of stochastic tokenization passes to average (TTA). Only used when dropout > 0. """ model = model.to(device) pad_id = tokenizer.pad_id def _embed_once(seqs, use_dropout): all_embs = [] for start in tqdm(range(0, len(seqs), batch_size), desc="Embedding", leave=False): batch = seqs[start : start + batch_size] enc = [tokenizer.encode(s, dropout=use_dropout)[:max_tokens] for s in batch] L = max(len(e) for e in enc) ids = torch.full((len(enc), L), pad_id, dtype=torch.long) mask = torch.zeros(len(enc), L, dtype=torch.long) for i, e in enumerate(enc): ids[i, :len(e)] = torch.tensor(e, dtype=torch.long) mask[i, :len(e)] = 1 ids, mask = ids.to(device), mask.to(device) with torch.amp.autocast("cuda", dtype=torch.float16, enabled=device.type == "cuda"): emb = model.embed(ids, mask) # L2-normalised [B, D] all_embs.append(emb.float().cpu().numpy()) return np.vstack(all_embs) if dropout > 0 and tta_passes > 1: # Test-Time Augmentation: average K stochastic encodings print(f" TTA: {tta_passes} passes with dropout={dropout}") emb_sum = None for k in range(tta_passes): emb_k = _embed_once(sequences, dropout) emb_sum = emb_k if emb_sum is None else emb_sum + emb_k emb_avg = emb_sum / tta_passes # Re-normalize after averaging norms = np.linalg.norm(emb_avg, axis=1, keepdims=True) return emb_avg / np.clip(norms, 1e-8, None) else: return _embed_once(sequences, dropout) # ─── 4. Hierarchy parsing ────────────────────────────────────────────────── def _scope_levels(classifications: List[str]) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: """ Return (fold_arr, sf_arr, family_arr). Works for both: SCOPe a.b.c.d (letter class + 3 ints) SCOP CL.CF.SF.FA (4 numeric IDs) fold = parts[:2], superfamily = parts[:3], family = parts[:4]. """ fold_arr = np.array([".".join(c.split(".")[:2]) for c in classifications]) sf_arr = np.array([".".join(c.split(".")[:3]) for c in classifications]) fam_arr = np.array([".".join(c.split(".")[:4]) for c in classifications]) return fold_arr, sf_arr, fam_arr def _cath_levels(classifications: List[str]) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: """Return (fold_arr, sf_arr, family_arr) for CATH C.A.T.H.S strings.""" fold_arr = np.array([".".join(c.split(".")[:3]) for c in classifications]) sf_arr = np.array([".".join(c.split(".")[:4]) for c in classifications]) fam_arr = np.array([".".join(c.split(".")[:5]) for c in classifications]) return fold_arr, sf_arr, fam_arr # ─── 5. Metrics ──────────────────────────────────────────────────────────── def _per_query_metrics(sim: np.ndarray, pos_arr: np.ndarray, excl_arr: np.ndarray) -> Dict: """ Generic per-query AUROC / AUPRC. Positive = same pos_arr label, different excl_arr label Negative = different pos_arr label Ignored = same excl_arr (trivially easy — excluded from scoring) Self = always excluded Fold-level : pos_arr = fold_arr, excl_arr = sf_arr Superfamily : pos_arr = sf_arr, excl_arr = fam_arr """ from sklearn.metrics import roc_auc_score, average_precision_score N = sim.shape[0] same_pos = pos_arr[:, None] == pos_arr[None, :] # [N, N] same_excl = excl_arr[:, None] == excl_arr[None, :] # [N, N] np.fill_diagonal(same_pos, False) np.fill_diagonal(same_excl, False) positive = same_pos & ~same_excl ignore = same_excl aurocs, auprcs = [], [] for qi in tqdm(range(N), desc="Computing metrics", leave=False): pos_row = positive[qi] ign_row = ignore[qi] if not pos_row.any(): continue keep = ~ign_row keep[qi] = False y_true = pos_row[keep].astype(int) y_pred = sim[qi][keep] if y_true.sum() == 0 or (1 - y_true).sum() == 0: continue try: aurocs.append(roc_auc_score(y_true, y_pred)) auprcs.append(average_precision_score(y_true, y_pred)) except ValueError: pass mAP = float(np.mean(auprcs)) if auprcs else 0.0 return { "n_sequences" : N, "n_queries" : len(aurocs), "auroc" : float(np.mean(aurocs)) if aurocs else 0.0, "auprc" : mAP, "mAP" : mAP, } # ─── 6. Dataset runner ───────────────────────────────────────────────────── def run_dataset(model, tokenizer, entries: List[Dict], ds_name: str, level_fn, batch_size: int, device: torch.device, dropout: float = 0.0, tta_passes: int = 1) -> List[Dict]: """ Returns two result dicts: one for fold/architecture-level and one for superfamily/topology-level. """ sequences = [e["sequence"] for e in entries] classifications = [e["classification"] for e in entries] fold_arr, sf_arr, fam_arr = level_fn(classifications) print(f"\n {ds_name}: {len(sequences)} sequences", flush=True) emb = embed(model, tokenizer, sequences, batch_size=batch_size, device=device, dropout=dropout, tta_passes=tta_passes) sim = (emb @ emb.T).astype(np.float32) np.fill_diagonal(sim, -2.0) # CATH uses Architecture/Topology; SCOP/SCOPe uses Fold/Superfamily is_cath = ds_name.startswith("CATH") level1_name = "architecture" if is_cath else "fold" level2_name = "topology" if is_cath else "superfamily" fold_m = _per_query_metrics(sim, fold_arr, sf_arr) fold_m["dataset"] = ds_name fold_m["level"] = level1_name sf_m = _per_query_metrics(sim, sf_arr, fam_arr) sf_m["dataset"] = ds_name sf_m["level"] = level2_name return [fold_m, sf_m] # ─── 7. Public API (notebook + script) ───────────────────────────────────── _EXPECTED = { ("SCOPe", "fold"): {"auroc": 0.8847, "auprc": 0.3149}, ("CATH-S20", "architecture"): {"auroc": 0.8129, "auprc": 0.3418}, ("SCOP", "fold"): {"auroc": 0.9062, "auprc": 0.2919}, } def run_benchmark( repo: str = ".", scope=True, scop: bool = True, cath: bool = True, cath_labels: Optional[str] = None, batch_size: int = 128, device: Optional[str] = None, seed: Optional[int] = 42, dropout: float = 0.0, tta_passes: int = 1, ) -> List[Dict]: """ Run the fold-level remote homology benchmark for LEMON. This function is the primary entry point for both scripts and Jupyter notebooks. It returns a list of result dicts that can be inspected directly or converted to a pandas DataFrame. Parameters ---------- repo : str Path to the HuggingFace snapshot directory (the root that contains ``model.safetensors``, ``config.json``, ``data/``, etc.). Defaults to ``"."`` — i.e. run from inside the snapshot dir. scope : bool or str ``True`` → use bundled ``data/scope_10_2.08.fa`` ``False`` → skip SCOPe ``str`` → path to a custom SCOPe ASTRAL FASTA scop : bool or str Same semantics for SCOP 1.75 (``data/scop175.fa``). cath : bool or str Same semantics for CATH-S20 (``data/cath_s20.fa``). cath_labels : str or None Path to a CATH domain-to-class TSV. ``None`` uses the bundled ``data/cath_s20_labels.tsv``. batch_size : int Sequences per forward pass. Reduce if you run out of memory. device : str or None ``"cuda"``, ``"cpu"``, or ``None`` for auto-detect. seed : int or None Random seed for reproducibility. ``None`` disables seeding. dropout : float Trie-dropout rate for tokenization (0.0 = deterministic greedy). tta_passes : int Number of stochastic tokenization passes to average (TTA). Only used when dropout > 0. Returns ------- list of dict One dict per dataset with keys: ``dataset``, ``n_sequences``, ``n_queries``, ``auroc``, ``auprc``, ``mAP``. Notebook quick-start -------------------- >>> import sys >>> sys.path.insert(0, "/path/to/snapshot") # or os.chdir there >>> from eval_retrieval import run_benchmark, display_results >>> >>> results = run_benchmark() # all three datasets >>> display_results(results) # rich table in notebook >>> >>> # As a DataFrame: >>> import pandas as pd >>> df = pd.DataFrame(results)[["dataset", "auroc", "auprc"]] >>> print(df) """ import random # Seed for reproducibility if seed is not None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) repo_path = Path(repo).resolve() data_dir = repo_path / "data" def _resolve_dataset(flag, bundled_name: str) -> Optional[Path]: if flag is False or flag is None: return None if flag is True: p = data_dir / bundled_name if not p.exists(): raise FileNotFoundError( f"Bundled dataset not found: {p}\n" f"Re-run: snapshot_download('Team-LEMON/lemon')" ) return p return Path(flag) # custom path string _device = torch.device( device if device else ("cuda" if torch.cuda.is_available() else "cpu") ) print(f"Device: {_device}") print("Loading LEMON …") model, tokenizer = _load_model_and_tokenizer(str(repo_path)) results: List[Dict] = [] scope_path = _resolve_dataset(scope, "scope_10_2.08.fa") if scope_path: entries = parse_scope_fasta(str(scope_path)) entries = [e for e in entries if len(e["classification"].split(".")) >= 4] results.extend(run_dataset(model, tokenizer, entries, "SCOPe", _scope_levels, batch_size, _device, dropout=dropout, tta_passes=tta_passes)) scop_path = _resolve_dataset(scop, "scop175.fa") if scop_path: entries = parse_scop_fasta(str(scop_path)) entries = [e for e in entries if len(e["classification"].split(".")) >= 4] results.extend(run_dataset(model, tokenizer, entries, "SCOP", _scope_levels, batch_size, _device, dropout=dropout, tta_passes=tta_passes)) cath_path = _resolve_dataset(cath, "cath_s20.fa") if cath_path: lbl = cath_labels or str(data_dir / "cath_s20_labels.tsv") entries = parse_cath_fasta(str(cath_path), lbl) entries = [e for e in entries if len(e["classification"].split(".")) >= 5] results.extend(run_dataset(model, tokenizer, entries, "CATH-S20", _cath_levels, batch_size, _device, dropout=dropout, tta_passes=tta_passes)) return results def display_results(results: List[Dict]) -> None: """ Pretty-print benchmark results. In a Jupyter notebook this also renders a styled pandas DataFrame if pandas is available. Falls back to a plain text table otherwise. Parameters ---------- results : list of dict Return value of :func:`run_benchmark`. """ # ── plain-text table (always printed) ────────────────────────────────── W = 76 print("\n" + "=" * W) print(f" {'Dataset':<12} {'Level':<12} {'N':>6} {'Queries':>7} {'AUROC':>7} {'AUPRC':>7} {'mAP':>7}") print("-" * W) for r in results: print( f" {r['dataset']:<12} {r['level']:<12} {r['n_sequences']:>6}" f" {r['n_queries']:>7} {r['auroc']:>7.4f} {r['auprc']:>7.4f} {r['mAP']:>7.4f}" ) print("=" * W) print("\nReference — LEMON (Table 1, averaged across seq-id thresholds):") for (ds, lvl), vals in _EXPECTED.items(): print(f" {ds:<12} {lvl:<12} AUROC={vals['auroc']:.4f} AUPRC={vals['auprc']:.4f}") # ── rich DataFrame display (Jupyter only) ────────────────────────────── try: import pandas as pd from IPython.display import display as ipy_display rows = [] for r in results: rows.append({ "Dataset" : r["dataset"], "Level" : r["level"], "N" : r["n_sequences"], "Queries" : r["n_queries"], "AUROC" : round(r["auroc"], 4), "AUPRC" : round(r["auprc"], 4), "mAP" : round(r["mAP"], 4), }) df = pd.DataFrame(rows).set_index(["Dataset", "Level"]) ref_rows = [ {"Dataset": ds, "Level": lvl, "AUROC": vals["auroc"], "AUPRC": vals["auprc"]} for (ds, lvl), vals in _EXPECTED.items() ] ref_df = pd.DataFrame(ref_rows).set_index(["Dataset", "Level"]) ipy_display( df.style .format("{:.4f}", subset=["AUROC", "AUPRC", "mAP"]) .set_caption("LEMON — fold-level and superfamily-level remote homology retrieval") ) print("\nReference (Table 1):") ipy_display( ref_df.style .format("{:.4f}", subset=["AUROC", "AUPRC"]) .set_caption("Expected values (paper, averaged over thresholds)") ) except ImportError: pass # pandas / IPython not available — plain text is sufficient # ─── 8. CLI entry point ───────────────────────────────────────────────────── def main(): p = argparse.ArgumentParser( description="Fold-level remote homology benchmark for LEMON", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=( "Run with no dataset flags to evaluate all three bundled datasets.\n" "Examples:\n" " python eval_retrieval.py\n" " python eval_retrieval.py --scope --cath\n" " python eval_retrieval.py --scope /my/scope.fa\n" " python eval_retrieval.py --dropout 0.1 --tta 8 # TTA with 8 passes\n" ), ) p.add_argument("--repo", default=".", help="Path to HF snapshot dir (default: current dir)") p.add_argument("--scope", nargs="?", const=True, default=None, metavar="FASTA", help="SCOPe FASTA (omit path → bundled data/scope_10_2.08.fa)") p.add_argument("--scop", nargs="?", const=True, default=None, metavar="FASTA", help="SCOP 1.75 FASTA (omit path → bundled data/scop175.fa)") p.add_argument("--cath", nargs="?", const=True, default=None, metavar="FASTA", help="CATH-S20 FASTA (omit path → bundled data/cath_s20.fa)") p.add_argument("--cath-labels", default=None, metavar="TSV", help="CATH domain→class TSV (default: bundled data/cath_s20_labels.tsv)") p.add_argument("--batch-size", type=int, default=128) p.add_argument("--device", default=None) p.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility (default: 42)") p.add_argument("--dropout", type=float, default=0.0, help="Trie-dropout rate for tokenization (default: 0.0 = deterministic)") p.add_argument("--tta", type=int, default=1, dest="tta_passes", help="Number of TTA passes (default: 1, only used if dropout > 0)") args = p.parse_args() # No flags → run everything run_all = not any([args.scope, args.scop, args.cath]) results = run_benchmark( repo = args.repo, scope = True if run_all else (args.scope or False), scop = True if run_all else (args.scop or False), cath = True if run_all else (args.cath or False), cath_labels = args.cath_labels, batch_size = args.batch_size, device = args.device, seed = args.seed, dropout = args.dropout, tta_passes = args.tta_passes, ) display_results(results) if __name__ == "__main__": main()