LEMON: Layered Extraction of Molecular Ordering from Nature

- 203M parameter protein sequence encoder with ZEST tokenizer
- Hierarchical contrastive learning for fold/superfamily similarity
- Bundled benchmarks: SCOPe, SCOP 1.75, CATH-S20
- Circular Permutation detection (CIRPIN SCOPe40)
- Test-Time Augmentation with trie-dropout
- MIT License

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Team-LEMON
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,192 @@
+---
+license: mit
+tags:
+  - protein
+  - biology
+  - sequence-encoder
+  - contrastive-learning
+  - lemon
+---
+
+# LEMON: Layered Extraction of Molecular Ordering from Nature
+
+LEMON is a protein sequence encoder trained with hierarchical contrastive
+learning for family and fold similarity search, using the ZEST tokenizer
+(Zoned Encoding of Sequence Traits).  Submitted anonymously for double-blind
+peer review.
+
+## Architecture
+
+| Component | Details |
+|---|---|
+| **Encoder** | 24-layer transformer, 768-d, 12 heads, SwiGLU FFN (ff_mult=4), RoPE with linear scaling |
+| **Pooling** | Learned-query multi-head attention aggregator → 768-d sequence vector |
+| **Projector** | Bottleneck MLP (768 → 768 → 384-d), L2-normalised output |
+| **Tokenizer** | ZEST 32K (Zoned Encoding of Sequence Traits) — greedy max-match trie over biochemically-substitutable amino-acid n-gram clusters |
+| **Context** | 1,024 tokens (linear RoPE scaling for longer sequences) |
+| **Dropout** | 0.04 |
+
+**Parameter breakdown (203.72M total):**
+
+| Module | Params |
+|---|---|
+| Core transformer | 194.52M |
+| Attention aggregator | 2.95M |
+| Profile expansion head | 2.41M |
+| Global position embedding | 2.36M |
+| Projector | 1.48M |
+| **Total** | **203.72M** |
+
+## Quickstart
+
+```python
+import torch
+from huggingface_hub import snapshot_download
+import sys, os
+
+path = snapshot_download("Team-LEMON/lemon")
+sys.path.insert(0, path)
+
+from modeling_lemon import LemonEncoder
+from tokenization_zest import ZESTTokenizer
+
+tok   = ZESTTokenizer.from_pretrained(path)
+model = LemonEncoder.from_pretrained(
+    os.path.join(path, "model.safetensors"),
+    os.path.join(path, "config.json"),
+)
+model.eval()
+
+seqs = ["MKTAYIAKQRQISFVKSHFSRQ", "ACDEFGHIKLMNPQRSTVWY"]
+enc  = tok.batch_encode_plus(seqs, max_length=512, padding=True)
+with torch.no_grad():
+    emb = model.embed(enc["input_ids"], enc["attention_mask"])  # [2, 384]
+print(emb.shape)   # torch.Size([2, 384])
+
+sim = model.similarity(emb[:1], emb[1:])
+print("cosine-like similarity:", sim.item())
+```
+
+## Reproducing Table 1
+
+The `eval_retrieval.py` script and all three benchmark datasets are bundled in this
+repo. No external downloads required.
+
+**Run all three datasets in one command:**
+
+```python
+from huggingface_hub import snapshot_download
+path = snapshot_download("Team-LEMON/lemon")
+```
+
+```bash
+cd /path/to/snapshot
+python eval_retrieval.py          # runs SCOPe + SCOP + CATH-S20
+python eval_retrieval.py --scope  # SCOPe only
+python eval_retrieval.py --cath   # CATH-S20 only
+python eval_retrieval.py --scop   # SCOP only
+```
+
+**Test-Time Augmentation (TTA) with Trie-Dropout:**
+
+TTA improves retrieval by averaging embeddings from multiple stochastic tokenizations.
+
+```bash
+python eval_retrieval.py --dropout 0.45 --tta 5   # 5 stochastic passes, averaged
+```
+
+**TTA Gain (SCOPe, seed=42):**
+
+| Level | Metric | Baseline | TTA (d=0.45, k=5) | Gain |
+|-------|--------|----------|-------------------|------|
+| fold | AUROC | 0.9025 | 0.9080 | +0.0055 |
+| fold | mAP | 0.3067 | 0.3197 | +0.0130 |
+| superfamily | AUROC | 0.9443 | 0.9519 | +0.0076 |
+| superfamily | mAP | 0.4700 | 0.4803 | +0.0103 |
+
+To reproduce:
+```bash
+# Baseline
+python eval_retrieval.py --scope --seed 42
+
+# With TTA
+python eval_retrieval.py --scope --seed 42 --dropout 0.45 --tta 5
+```
+
+**Or from a Jupyter notebook:**
+
+```python
+import sys
+from huggingface_hub import snapshot_download
+
+path = snapshot_download("Team-LEMON/lemon")
+sys.path.insert(0, path)
+
+from eval_retrieval import run_benchmark, display_results
+
+results = run_benchmark(repo=path, seed=42)   # deterministic with seed=42
+display_results(results)
+
+# With TTA:
+# results = run_benchmark(repo=path, seed=42, dropout=0.1, tta_passes=8)
+```
+
+**Expected output (seed=42, deterministic):**
+
+| Dataset  | Level        | AUROC  | mAP    |
+|----------|--------------|--------|--------|
+| SCOPe    | fold         | 0.8847 | 0.3149 |
+| CATH-S20 | architecture | 0.8129 | 0.3418 |
+| SCOP     | fold         | 0.9062 | 0.2919 |
+
+> Results are deterministic with `--seed 42` (default).
+> CATH uses Architecture/Topology levels; SCOP/SCOPe uses Fold/Superfamily.
+
+**Bundled dataset provenance:**
+
+| File | Sequences | Original source |
+|------|-----------|----------------|
+| `data/scope_10_2.08.fa` | 7 117 | SCOPe 2.08, 10% seq-id — [scop.berkeley.edu](https://scop.berkeley.edu/downloads/scopeseq-2.08/) |
+| `data/cath_s20.fa` | 15 043 | CATH v4.4.0 S20 — [cathdb.info](https://www.cathdb.info/wiki/doku/?id=data:index) |
+| `data/cath_s20_labels.tsv` | 15 043 | CATH domain list v4.4.0 (S20 subset) — [cathdb.info](https://www.cathdb.info/wiki/doku/?id=data:index) |
+| `data/scop175.fa` | 31 073 | SCOP 1.75 — [plm-zero-shot-remote-homology-evaluation](https://github.com/amoldwin/plm-zero-shot-remote-homology-evaluation) |
+
+## Circular Permutation Detection (CIRPIN SCOPe40)
+
+Zero-shot detection of circularly permuted protein pairs using cosine similarity of LEMON embeddings.
+Benchmark: CIRPIN SCOPe40 — 18,127 pairs (1,967 positive CP pairs) from ASTRAL SCOPe 2.08 at 40% identity.
+
+**Results (seed=42):**
+
+| Configuration | AUROC | AUPRC | Accuracy |
+|---------------|-------|-------|----------|
+| Baseline | 0.7413 | 0.3035 | 0.8990 |
+| TTA (d=0.45, k=5) | **0.7576** | **0.3066** | 0.8987 |
+| Gain | +0.0163 | +0.0031 | - |
+
+TTA improves CP detection by averaging embeddings over multiple stochastic tokenizations.
+
+To reproduce:
+```bash
+# Baseline
+python eval_circular_permutation.py --fasta data/cirpin/scope40.fa --pairs data/cirpin/pairs.tsv --seed 42
+
+# With TTA
+python eval_circular_permutation.py --fasta data/cirpin/scope40.fa --pairs data/cirpin/pairs.tsv --seed 42 --dropout 0.45 --tta 5
+```
+
+## Requirements
+
+```
+torch>=2.0
+safetensors
+huggingface_hub
+```
+
+## Notes
+
+- Input sequences should be standard single-letter amino-acid strings.
+- The tokenizer handles unknown characters via `<MASK>` token fallback.
+- `model.embed()` returns L2-normalised embeddings; use dot product for
+  cosine similarity.
+- `model.similarity()` applies a learned temperature scalar.

config.json

@@ -0,0 +1,21 @@
+{
+  "model_class": "LemonEncoder",
+  "architecture": {
+    "vocab_size": 32000,
+    "embed_dim": 768,
+    "num_layers": 24,
+    "nhead": 12,
+    "ff_mult": 4,
+    "dropout": 0.04,
+    "max_tokens": 1024,
+    "proj_dim": 384,
+    "padding_idx": 0,
+    "rope_scaling_type": "linear",
+    "num_proj_layers": 1,
+    "proj_ff_mult": 1
+  },
+  "epoch": 3,
+  "global_step": 40371,
+  "best_val_loss": 0.2449137058109045,
+  "best_epoch": 2
+}

eval_retrieval.py

@@ -0,0 +1,589 @@
+"""
+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()

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5e49be719f03463908e08660036292f6b727bd3dbef1a740ecd41e1264dea893
+size 814907772

modeling_lemon.py

@@ -0,0 +1,401 @@
+"""
+Protein sequence encoder: transformer with RoPE, SwiGLU FFN, attention pooling,
+and a bottleneck projector for contrastive sequence similarity.
+
+No external dependencies beyond PyTorch.
+"""
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint as grad_checkpoint
+
+
+# ── Positional Embeddings ─────────────────────────────────────────────────
+
+class RotaryPositionalEmbedding(nn.Module):
+    def __init__(self, dim: int, max_tokens: int = 512, base: int = 10000,
+                 scaling_type=None):
+        super().__init__()
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+        self.max_tokens = max_tokens
+        self.scaling_type = scaling_type
+
+    def _scaled_positions(self, seq_len, device, dtype):
+        positions = torch.arange(seq_len, device=device, dtype=dtype)
+        if seq_len <= self.max_tokens or not self.scaling_type:
+            return positions
+        if self.scaling_type == "linear":
+            return positions * (self.max_tokens / seq_len)
+        raise ValueError(f"Unknown scaling_type: {self.scaling_type}")
+
+    def forward(self, x):
+        t = self._scaled_positions(x.shape[1], x.device, self.inv_freq.dtype)
+        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+        emb = torch.cat([freqs, freqs], dim=-1)
+        return emb[None, :, :]
+
+
+def rotate_half(x):
+    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
+    return torch.cat([-x2, x1], dim=-1)
+
+def apply_rotary_pos_emb(q, k, pos_emb):
+    cos, sin = pos_emb.cos(), pos_emb.sin()
+    return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)
+
+
+class ScaledPositionalEmbedding(nn.Module):
+    def __init__(self, num_embeddings, embedding_dim, extend_strategy="alibi",
+                 padding_idx=None):
+        super().__init__()
+        self.embedding = nn.Embedding(num_embeddings, embedding_dim, padding_idx=padding_idx)
+        self.extend_strategy = extend_strategy
+        self.num_embeddings = num_embeddings
+        self.padding_idx = padding_idx
+        max_index = num_embeddings - 1
+        if padding_idx is not None:
+            content_max_index = max(0, min(max_index - 1, padding_idx - 1))
+        else:
+            content_max_index = max_index
+        self.register_buffer("max_index", torch.tensor(max_index), persistent=False)
+        self.register_buffer("content_max_index", torch.tensor(content_max_index), persistent=False)
+        if extend_strategy == "alibi":
+            slopes = torch.logspace(0, -3, steps=embedding_dim, base=2.0)
+            self.register_buffer("alibi_slopes", slopes, persistent=False)
+        else:
+            self.register_buffer("alibi_slopes", None, persistent=False)
+
+    def forward(self, positions):
+        positions = positions.long()
+        if self.padding_idx is not None:
+            pad_mask = positions == self.padding_idx
+            clamped = positions.clamp(0, int(self.content_max_index.item()))
+            clamped = torch.where(pad_mask, torch.full_like(clamped, self.padding_idx), clamped)
+        else:
+            pad_mask = torch.zeros_like(positions, dtype=torch.bool)
+            clamped = positions.clamp(0, int(self.max_index.item()))
+        embeddings = self.embedding(clamped)
+        if self.extend_strategy == "alibi":
+            excess = (positions - self.content_max_index).clamp_min(0)
+            if self.padding_idx is not None:
+                excess = torch.where(pad_mask, torch.zeros_like(excess), excess)
+            embeddings = embeddings + excess.unsqueeze(-1) * self.alibi_slopes
+        if self.padding_idx is not None:
+            embeddings = embeddings.masked_fill(pad_mask.unsqueeze(-1), 0.0)
+        return embeddings
+
+
+# ── Transformer Block ─────────────────────────────────────────────────────
+
+class AttentionBlock(nn.Module):
+    def __init__(self, embed_dim, nhead, ff_mult=4, dropout=0.1):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.nhead = nhead
+        self.head_dim = embed_dim // nhead
+        self.norm1 = nn.LayerNorm(embed_dim)
+        self.norm2 = nn.LayerNorm(embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.o_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        ff_dim = int(embed_dim * ff_mult * 2 / 3)
+        self.w1 = nn.Linear(embed_dim, ff_dim, bias=False)
+        self.w2 = nn.Linear(embed_dim, ff_dim, bias=False)
+        self.w3 = nn.Linear(ff_dim, embed_dim, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, mask=None, rope_emb=None):
+        B, L, _ = x.shape
+        if mask is not None:
+            mask = mask.to(torch.bool)
+        h = self.norm1(x)
+        if mask is not None:
+            h = h.masked_fill(~mask.unsqueeze(-1), 0.0)
+        q = self.q_proj(h).view(B, L, self.nhead, self.head_dim).transpose(1, 2)
+        k = self.k_proj(h).view(B, L, self.nhead, self.head_dim).transpose(1, 2)
+        v = self.v_proj(h).view(B, L, self.nhead, self.head_dim).transpose(1, 2)
+        if mask is not None:
+            m = mask.unsqueeze(1).unsqueeze(-1)
+            k = k.masked_fill(~m, 0.0)
+            v = v.masked_fill(~m, 0.0)
+        if rope_emb is not None:
+            q, k = apply_rotary_pos_emb(q, k, rope_emb.unsqueeze(1))
+        out = F.scaled_dot_product_attention(
+            q, k, v, attn_mask=None,
+            dropout_p=self.dropout.p if self.training else 0.0,
+            is_causal=False,
+        )
+        out = out.transpose(1, 2).contiguous().view(B, L, self.embed_dim)
+        out = self.o_proj(out)
+        if mask is not None:
+            out = out.masked_fill(~mask.unsqueeze(-1), 0.0)
+        x = x + self.dropout(out)
+        h = self.norm2(x)
+        if mask is not None:
+            h = h.masked_fill(~mask.unsqueeze(-1), 0.0)
+        ff = self.w3(F.silu(self.w1(h)) * self.w2(h))
+        if mask is not None:
+            ff = ff.masked_fill(~mask.unsqueeze(-1), 0.0)
+        x = x + self.dropout(ff)
+        return x
+
+
+# ── Encoder Stack ─────────────────────────────────────────────────────────
+
+class Encoder(nn.Module):
+    def __init__(self, vocab_size=32000, embed_dim=256, num_layers=4, nhead=8,
+                 ff_mult=4, dropout=0.1, max_tokens=1024, padding_idx=0,
+                 rope_scaling_type=None):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_layers = num_layers
+        self.gradient_checkpointing = False
+        self.token_emb = nn.Embedding(vocab_size, embed_dim, padding_idx=padding_idx)
+        self.emb_dropout = nn.Dropout(dropout)
+        self.rope = RotaryPositionalEmbedding(
+            embed_dim // nhead, max_tokens, scaling_type=rope_scaling_type)
+        self.layers = nn.ModuleList([
+            AttentionBlock(embed_dim, nhead, ff_mult, dropout) for _ in range(num_layers)
+        ])
+        self.final_norm = nn.LayerNorm(embed_dim)
+        self._init_weights()
+
+    def _init_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.Embedding):
+                nn.init.normal_(m.weight, std=0.02)
+                if m.padding_idx is not None:
+                    m.weight.data[m.padding_idx].zero_()
+
+    def forward(self, tokens, mask=None):
+        x = self.emb_dropout(self.token_emb(tokens))
+        padding_mask = mask.to(torch.bool) if mask is not None else None
+        rope_emb = self.rope(x)
+        for layer in self.layers:
+            if self.gradient_checkpointing and self.training:
+                x = grad_checkpoint(layer, x, padding_mask, rope_emb, use_reentrant=False)
+            else:
+                x = layer(x, padding_mask, rope_emb)
+        x = self.final_norm(x)
+        if padding_mask is not None:
+            x = x.masked_fill(~padding_mask.unsqueeze(-1), 0.0)
+        return x
+
+
+# ── Attention Pooling ─────────────────────────────────────────────────────
+
+class AttentionAggregator(nn.Module):
+    def __init__(self, embed_dim, num_heads=4, dropout=0.1):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        self.dropout = dropout
+        self.pool_query = nn.Parameter(torch.randn(1, 1, embed_dim) * 0.02)
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.o_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.norm = nn.LayerNorm(embed_dim)
+        self.align = nn.Linear(embed_dim, embed_dim, bias=False)
+        nn.init.eye_(self.align.weight)
+
+    def forward(self, token_embs, mask=None):
+        B, L, _ = token_embs.shape
+        H, D = self.num_heads, self.head_dim
+        query = self.pool_query.expand(B, -1, -1)
+        q = self.q_proj(query).view(B, 1, H, D).transpose(1, 2)
+        k = self.k_proj(token_embs).view(B, L, H, D).transpose(1, 2)
+        v = self.v_proj(token_embs).view(B, L, H, D).transpose(1, 2)
+        attn_mask = None
+        if mask is not None:
+            bool_mask = mask.to(torch.bool)
+            attn_mask = torch.zeros(B, 1, 1, L, device=token_embs.device, dtype=token_embs.dtype)
+            attn_mask = attn_mask.masked_fill(~bool_mask.unsqueeze(1).unsqueeze(2), float("-inf"))
+            m = bool_mask.unsqueeze(1).unsqueeze(-1)
+            k = k.masked_fill(~m, 0.0)
+            v = v.masked_fill(~m, 0.0)
+        out = F.scaled_dot_product_attention(
+            q, k, v, attn_mask=attn_mask,
+            dropout_p=self.dropout if self.training else 0.0,
+            is_causal=False,
+        )
+        out = out.transpose(1, 2).contiguous().view(B, 1, self.embed_dim)
+        pooled = self.norm(self.o_proj(out)).squeeze(1)
+        return self.align(pooled)
+
+
+# ── Residue Expansion Head ────────────────────────────────────────────────
+
+class ExpansionHead(nn.Module):
+    def __init__(self, embed_dim, max_residues, vocab_size=20, padding_index=31,
+                 local_extend_strategy="alibi"):
+        super().__init__()
+        self.max_residues = max_residues
+        self.PADDING_INDEX = padding_index
+        self.local_pos_emb = ScaledPositionalEmbedding(
+            num_embeddings=self.PADDING_INDEX + 1, embedding_dim=embed_dim,
+            extend_strategy=local_extend_strategy, padding_idx=self.PADDING_INDEX,
+        )
+        self.residue_expansion = self._bottleneck_mlp(embed_dim, embed_dim, 0.1, 2)
+        self.aa_logits_proj = nn.Linear(embed_dim, vocab_size)
+
+    def _bottleneck_mlp(self, input_dim, hidden, dropout, num_layers, output_dim=None):
+        def lng(i, h, d):
+            return [nn.Linear(i, h), nn.LayerNorm(h), nn.GELU(), nn.Dropout(d)]
+        layers = []
+        for _ in range(num_layers):
+            layers.extend(lng(input_dim, hidden, dropout) + lng(hidden, input_dim, dropout))
+        if output_dim is not None:
+            layers.append(nn.Linear(input_dim, output_dim))
+        return nn.Sequential(*layers)
+
+    def _expand_tokens(self, z, token_lengths, target_len):
+        B, T, D = z.shape
+        device = z.device
+        if not torch.is_tensor(token_lengths):
+            token_lengths = torch.tensor(token_lengths, device=device, dtype=torch.long)
+        token_lengths = torch.clamp(token_lengths.clone(), min=0)
+        cumsum = torch.cumsum(token_lengths, dim=1)
+        total_residues = cumsum[:, -1]
+        positions = torch.arange(target_len, device=device).unsqueeze(0).expand(B, -1)
+        token_indices = torch.clamp(torch.searchsorted(cumsum, positions + 1), 0, T - 1)
+        z_expanded = torch.gather(z, 1, token_indices.unsqueeze(-1).expand(-1, -1, D))
+        cumsum_shifted = F.pad(cumsum[:, :-1], (1, 0), value=0)
+        local_indices = positions - torch.gather(cumsum_shifted, 1, token_indices)
+        padding_mask = positions >= total_residues.unsqueeze(1)
+        z_expanded = z_expanded.masked_fill(padding_mask.unsqueeze(-1), 0.0)
+        local_indices = torch.where(
+            padding_mask, torch.full_like(local_indices, self.PADDING_INDEX),
+            local_indices.clamp(min=0))
+        return z_expanded, local_indices
+
+    def forward(self, z, token_lengths, global_indices, global_pos_emb):
+        target_len = global_indices.shape[1]
+        z_exp, local_idx = self._expand_tokens(z, token_lengths, target_len)
+        pos_global = global_pos_emb(torch.clamp(global_indices, min=0)).to(z_exp.dtype)
+        pos_local = self.local_pos_emb(local_idx).to(z_exp.dtype)
+        z_final = z_exp + pos_global + pos_local
+        return self.aa_logits_proj(self.residue_expansion(z_final))
+
+
+# ── Main Encoder Model ────────────────────────────────────────────────────
+
+class LemonEncoder(nn.Module):
+    """
+    Protein sequence encoder: trie-tokenised input → per-token embeddings →
+    attention-pooled sequence embedding → L2-normalised projector output.
+
+    Suitable for contrastive similarity search (family / fold retrieval).
+
+    Loading
+    -------
+    >>> from modeling_protein_encoder import LemonEncoder
+    >>> model = LemonEncoder.from_pretrained("model.safetensors",
+    ...                                                 config_path="config.json")
+    >>> model.eval()
+    """
+
+    def __init__(
+        self,
+        vocab_size: int = 32000,
+        embed_dim: int = 256,
+        num_layers: int = 8,
+        nhead: int = 8,
+        ff_mult: int = 4,
+        dropout: float = 0.1,
+        max_tokens: int = 1024,
+        proj_dim: int = 128,
+        padding_idx: int = 0,
+        rope_scaling_type=None,
+        num_proj_layers: int = 2,
+        proj_ff_mult: int = 2,
+    ):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.max_tokens = max_tokens
+        self.embed_dim = embed_dim
+        self.proj_dim = proj_dim or embed_dim
+        self.num_proj_layers = num_proj_layers
+        self.proj_ff_mult = proj_ff_mult
+
+        self.core = Encoder(
+            vocab_size=vocab_size, embed_dim=embed_dim, num_layers=num_layers,
+            nhead=nhead, ff_mult=ff_mult, dropout=dropout, max_tokens=max_tokens,
+            padding_idx=padding_idx, rope_scaling_type=rope_scaling_type,
+        )
+        self.log_temperature = nn.Parameter(torch.tensor(0.07).log())
+        self.aggregator = AttentionAggregator(embed_dim, dropout=dropout)
+        self.profile_expansion_head = ExpansionHead(embed_dim, max_residues=3 * max_tokens)
+        self.global_pos_emb = nn.Embedding(3 * max_tokens, embed_dim)
+
+        hidden = embed_dim * proj_ff_mult
+        self.projector = self._bottleneck_mlp(embed_dim, hidden, dropout,
+                                              num_proj_layers, output_dim=proj_dim)
+        self.config = dict(
+            vocab_size=vocab_size, embed_dim=embed_dim, num_layers=num_layers,
+            nhead=nhead, ff_mult=ff_mult, dropout=dropout, max_tokens=max_tokens,
+            proj_dim=proj_dim, padding_idx=padding_idx,
+            rope_scaling_type=rope_scaling_type,
+            num_proj_layers=num_proj_layers, proj_ff_mult=proj_ff_mult,
+        )
+
+    def _bottleneck_mlp(self, input_dim, hidden, dropout, num_layers, output_dim=None):
+        def lng(i, h, d):
+            return [nn.Linear(i, h), nn.LayerNorm(h), nn.GELU(), nn.Dropout(d)]
+        layers = []
+        for _ in range(num_layers):
+            layers.extend(lng(input_dim, hidden, dropout) + lng(hidden, input_dim, dropout))
+        if output_dim is not None:
+            layers.append(nn.Linear(input_dim, output_dim))
+        return nn.Sequential(*layers)
+
+    @property
+    def temperature(self):
+        return self.log_temperature.exp().clamp(min=0.01, max=1.0)
+
+    def encode_tokens(self, tokens, mask=None):
+        return self.core(tokens, mask)
+
+    def embed(self, tokens, mask=None, normalize=True):
+        """Encode tokens → L2-normalised sequence embedding [B, proj_dim]."""
+        token_embs = self.encode_tokens(tokens, mask)
+        agg = self.aggregator(token_embs, mask)
+        proj = self.projector(agg)
+        if normalize:
+            proj = F.normalize(proj, p=2, dim=-1)
+        return proj
+
+    def similarity(self, emb_a, emb_b):
+        """Temperature-scaled cosine similarity matrix [B_a, B_b]."""
+        return (emb_a @ emb_b.T) / self.temperature
+
+    def forward(self, tokens, mask=None):
+        token_embs = self.encode_tokens(tokens, mask)
+        return F.linear(token_embs, self.core.token_emb.weight)
+
+    @classmethod
+    def from_pretrained(cls, weights_path: str, config_path: str = None,
+                        device: str = "cpu"):
+        """Load from safetensors weights + JSON config."""
+        import json, os
+        from safetensors.torch import load_file
+
+        if config_path is None:
+            config_path = os.path.join(os.path.dirname(weights_path), "config.json")
+
+        with open(config_path) as f:
+            cfg = json.load(f)
+
+        arch = cfg.get("architecture", cfg)
+        model = cls(**arch)
+        state = load_file(weights_path, device=device)
+        model.load_state_dict(state)
+        return model

tokenization_zest.py

@@ -0,0 +1,217 @@
+"""
+Protein motif tokenizer: greedy max-match trie over amino-acid n-gram clusters.
+
+No external dependencies beyond the standard library.
+
+Usage
+-----
+>>> from tokenization_protein_encoder import ZESTTokenizer
+>>> tok = ZESTTokenizer.from_pretrained(".")
+>>> ids = tok.encode("MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGDGTQDNLSGAEKAVQVKVKALPDAQFEVVHSLAKWKRQTLGQHDFSAGEGLYTHMKALRPDEDRLSPLHSVYVDQWDWERVMGDGERQFSTLKSTVEAIWAGIKATEAAVSEEFGLAPFLPDQIHFVHSQELLSRYPDLDAKGRERAIAKDLGAVFLVGIGGKLSDGHRHDVRAPDYDDWSTPSELGHAGLNGDILVWNPSVSMEFQKIPIHRLATLKKMRHSSMCGQDKTAFGKELQDLQTELESMSGQGRFFLASTPYLRPQLNQLPGLKVNLNVIEQYVQKQNQWSTILTVYRQKGKLSAEPFQPTSHQLSAEKLNEGNDNLSLAAFVQLLNTSPTLAQATAVQVQNPIDKLPNLNQDSIQALQPEDLHQVLNLPKR")
+>>> print(ids[:10])
+"""
+import json
+import os
+import re
+import random
+from typing import Callable, Dict, List, Optional, Tuple
+
+
+class _TrieNode:
+    __slots__ = ["children", "token_id"]
+    def __init__(self):
+        self.children: Dict[str, "_TrieNode"] = {}
+        self.token_id: int = -1
+
+
+class ZESTTokenizer:
+    """
+    Greedy max-match tokenizer backed by a symbol-level trie.
+    Clusters biochemically substitutable amino-acid n-grams into shared tokens,
+    analogous to BPE but guided by substitutability rather than frequency.
+    """
+
+    SPECIAL_TOKENS = ["<PAD>", "<UNK>", "<CLS>", "<EOS>", "<MASK>"]
+    DEFAULT_ALPHABET = list("ACDEFGHIKLMNPQRSTVWY")
+
+    def __init__(self, vocab_path: str, alphabet=None, alphabet_path=None,
+                 verbose=False):
+        self.path = vocab_path
+        self.vocab: Dict[str, int] = {}
+        self.id_to_token: Dict[int, str] = {}
+        self._root = _TrieNode()
+        self._symbol_to_id: Dict[str, int] = {}
+        self._ws = re.compile(r"\s+")
+
+        for i, tok in enumerate(self.SPECIAL_TOKENS):
+            self.vocab[tok] = i
+            self.id_to_token[i] = tok
+
+        if alphabet is not None:
+            self._alphabet = list(alphabet)
+        elif alphabet_path is not None and os.path.exists(alphabet_path):
+            with open(alphabet_path) as f:
+                self._alphabet = json.load(f)
+        else:
+            auto = vocab_path.replace(".json", "_alphabet.json")
+            if os.path.exists(auto):
+                with open(auto) as f:
+                    self._alphabet = json.load(f)
+            else:
+                self._alphabet = list(self.DEFAULT_ALPHABET)
+
+        offset = len(self.SPECIAL_TOKENS)
+        for i, sym in enumerate(self._alphabet):
+            tid = offset + i
+            self.vocab[sym] = tid
+            self.id_to_token[tid] = sym
+            self._symbol_to_id[sym] = tid
+            self._trie_insert([sym], tid)
+
+        with open(vocab_path) as f:
+            clusters: Dict[str, List[str]] = json.load(f)
+
+        offset = len(self.vocab)
+        for i, (centroid, members) in enumerate(clusters.items()):
+            tid = offset + i
+            self.id_to_token[tid] = centroid
+            self.vocab[centroid] = tid
+            self._trie_insert(self._pattern_to_symbols(centroid), tid)
+            for member in (members or []):
+                if member == centroid:
+                    continue
+                self.vocab[member] = tid
+                self._trie_insert(self._pattern_to_symbols(member), tid)
+
+        self.pad_id   = self.vocab["<PAD>"]
+        self.unk_id   = self.vocab["<UNK>"]
+        self.cls_id   = self.vocab["<CLS>"]
+        self.eos_id   = self.vocab["<EOS>"]
+        self.mask_id  = self.vocab["<MASK>"]
+        self.vocab_size = len(self.id_to_token)
+
+        if verbose:
+            n_many = sum(1 for m in clusters.values() if m and len(m) > 1)
+            print(f"ZESTTokenizer — vocab size: {self.vocab_size}")
+            print(f"  Alphabet:  {len(self._alphabet)} symbols")
+            print(f"  Patterns:  {len(clusters)} ({n_many} many-to-one)")
+
+    def _pattern_to_symbols(self, pattern):
+        if self._alphabet and all(len(s) == 1 for s in self._alphabet):
+            return list(pattern)
+        return pattern.split()
+
+    def _trie_insert(self, symbols, token_id):
+        node = self._root
+        for sym in symbols:
+            if sym not in node.children:
+                node.children[sym] = _TrieNode()
+            node = node.children[sym]
+        node.token_id = token_id
+
+    def _trie_match(self, symbols, start):
+        matches = []
+        node = self._root
+        for i in range(start, len(symbols)):
+            sym = symbols[i]
+            if sym not in node.children:
+                break
+            node = node.children[sym]
+            if node.token_id != -1:
+                matches.append((i - start + 1, node.token_id))
+        return matches[::-1]
+
+    def _segment(self, symbols, dropout=0.0):
+        segments = []
+        i, n = 0, len(symbols)
+        while i < n:
+            matches = self._trie_match(symbols, i)
+            if not matches:
+                segments.append((1, self.mask_id))
+                i += 1
+                continue
+            length, tid = matches[0]
+            if dropout > 0 and random.random() < dropout and len(matches) > 1:
+                idx = random.randint(1, len(matches) - 1)
+                length, tid = matches[idx]
+            segments.append((length, tid))
+            i += length
+        return segments
+
+    def encode(self, text: str, dropout: float = 0.0,
+               add_special_tokens: bool = False) -> List[int]:
+        """Encode a raw amino-acid sequence string to token IDs."""
+        symbols = list(re.sub(r"\s+", "", text).upper())
+        ids = [tid for _, tid in self._segment(symbols, dropout)]
+        if add_special_tokens:
+            ids = [self.cls_id] + ids + [self.eos_id]
+        return ids
+
+    def decode(self, ids: List[int], skip_special: bool = True) -> str:
+        skip = {self.pad_id, self.cls_id, self.eos_id, self.mask_id}
+        parts = []
+        for i in ids:
+            if skip_special and i in skip:
+                continue
+            parts.append(self.id_to_token.get(i, ""))
+        if self._alphabet and len(self._alphabet[0]) == 1:
+            return "".join(parts)
+        return " | ".join(parts)
+
+    def batch_encode_plus(self, sequences: List[str], padding: bool = True,
+                          max_length: Optional[int] = None, truncation: bool = True,
+                          dropout: float = 0.0, return_tensors: str = "pt",
+                          add_special_tokens: bool = False):
+        """Batch encode and optionally pad/truncate."""
+        try:
+            import torch
+        except ImportError:
+            raise ImportError("PyTorch is required for return_tensors=\'pt\'")
+        batch = [self.encode(s, dropout=dropout) for s in sequences]
+        n_special = 2 if add_special_tokens else 0
+        processed = []
+        for ids in batch:
+            if max_length and truncation and len(ids) + n_special > max_length:
+                ids = ids[: max_length - n_special]
+            if add_special_tokens:
+                ids = [self.cls_id] + ids + [self.eos_id]
+            processed.append(ids)
+        if padding:
+            target = max_length or max(len(ids) for ids in processed)
+            padded, masks = [], []
+            for ids in processed:
+                pad_n = max(0, target - len(ids))
+                padded.append(ids + [self.pad_id] * pad_n)
+                masks.append([1] * len(ids) + [0] * pad_n)
+        else:
+            padded = processed
+            masks = [[1] * len(ids) for ids in processed]
+        if return_tensors == "pt":
+            return {
+                "input_ids": torch.tensor(padded, dtype=torch.long),
+                "attention_mask": torch.tensor(masks, dtype=torch.long),
+            }
+        return {"input_ids": padded, "attention_mask": masks}
+
+    def save_pretrained(self, directory: str):
+        os.makedirs(directory, exist_ok=True)
+        clusters: Dict[str, List[str]] = {}
+        offset = len(self.SPECIAL_TOKENS) + len(self._alphabet)
+        centroid_ids = {tid for tid in self.id_to_token if tid >= offset}
+        for tok, tid in self.vocab.items():
+            if tid not in centroid_ids or tok in self.id_to_token.values():
+                continue
+            centroid = self.id_to_token[tid]
+            clusters.setdefault(centroid, []).append(tok)
+        with open(os.path.join(directory, "vocab_map.json"), "w") as f:
+            json.dump(clusters, f, indent=2)
+        with open(os.path.join(directory, "vocab_map_alphabet.json"), "w") as f:
+            json.dump(self._alphabet, f, indent=2)
+
+    @classmethod
+    def from_pretrained(cls, directory: str, **kwargs) -> "ZESTTokenizer":
+        vocab_path = os.path.join(directory, "vocab_map.json")
+        alpha_path = os.path.join(directory, "vocab_map_alphabet.json")
+        if os.path.exists(alpha_path):
+            return cls(vocab_path, alphabet_path=alpha_path, **kwargs)
+        return cls(vocab_path, **kwargs)

tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "tokenizer_class": "ZESTTokenizer",
+  "vocab_size": 32000,
+  "pad_token": "<PAD>",
+  "unk_token": "<UNK>",
+  "cls_token": "<CLS>",
+  "eos_token": "<EOS>",
+  "mask_token": "<MASK>",
+  "model_max_length": 1024
+}

vocab_map_alphabet.json

@@ -0,0 +1,22 @@
+[
+  "A",
+  "C",
+  "D",
+  "E",
+  "F",
+  "G",
+  "H",
+  "I",
+  "K",
+  "L",
+  "M",
+  "N",
+  "P",
+  "Q",
+  "R",
+  "S",
+  "T",
+  "V",
+  "W",
+  "Y"
+]