eval_retrieval.py

"""
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()