Add STAMP hybrid AO-GPT 30M d=512/l=8 ep7 (GenMol 79.00%)

Checkpoint + hybrid motif+char vocab (2481 tokens) + self-contained HybridVocab class + config + model card.

README.md

@@ -0,0 +1,201 @@
+---
+license: apache-2.0
+tags:
+  - chemistry
+  - molecular-generation
+  - smiles
+  - stamp
+  - drug-discovery
+---
+
+# STAMP Hybrid AO-GPT (31M, d=512/l=8, epoch 7)
+
+Pretrained AO-GPT (any-order GPT) over **STAMP** molecular token sequences
+with a **hybrid motif + character vocabulary**. Trained on 30M unique
+filtered molecules. Achieves **79.00% GenMol quality** — matching the
+112M-parameter AR baseline (79.64%) with 28% of the parameters and 20% of
+the vocabulary size.
+
+## Highlights
+
+- **Small vocab (2481)**: 2387 high-frequency atomic motifs (freq ≥ 5000)
+  + 49 SMILES character tokens + ~45 STAMP structural tokens. Covers ~91%
+  of motif occurrences as atomic tokens; rare motifs expand to chars.
+- **Training-time char fallback** with log-interpolated probability
+  (~2% at the most frequent motif, ~15% at the cutoff). The model sees
+  each atomic motif in both atomic and char form, closing the
+  train/inference gap for OOV motifs.
+- **STAMP structural tokens** (`[J_*]`, `[B_*]`, `[S_*]`, `[END]`) act as
+  natural motif boundaries — no extra `[MS]`/`[ME]` markers needed.
+- **Drug-like outputs**: 100% validity, 100% uniqueness (at N=1024),
+  79.00% pass the GenMol filter (QED ≥ 0.6 AND SA ≤ 4.0).
+
+## Files
+
+| file | what it is |
+|---|---|
+| `model.pt` | torch checkpoint: `{model_state, cfg, epoch, representation, model_type}` |
+| `hybrid_vocab.json` | full vocab with atomic motif map, frequencies, and char expansions |
+| `motif_vocab.txt` | source motif-freq file (v3_cm_union format: `smiles\tn_heavy\tfreq`) |
+| `hybrid_vocab.py` | self-contained `HybridVocab` class for decoding |
+| `config.json` | architecture summary + default sampling + eval numbers |
+
+## Evaluation (N=1024 at T=0.95, top_p=0.85)
+
+| metric | value |
+|---|---:|
+| validity | 100.00% |
+| uniqueness (raw SMILES) | 100.00% |
+| quality over valid (QED ≥ 0.6 ∧ SA ≤ 4) | 79.16% |
+| **GenMol score** | **79.00%** |
+| QED mean | 0.727 |
+| SA mean | 2.92 |
+| diversity (1 − pairwise Tanimoto, 1024-bit Morgan r=2) | 0.860 |
+
+**Reference (AR baseline, old 12573-token vocab, d=768/l=12, 112M params): 79.64%.**
+The hybrid model matches within noise at 28% of the parameter count.
+
+## Usage
+
+### 1. Load vocab
+
+```python
+from hybrid_vocab import HybridVocab
+vocab = HybridVocab.load("hybrid_vocab.json")
+# vocab.itos          -> list of 2481 token strings
+# vocab.atomic_motifs -> {smiles: id} for the 2387 motifs
+# vocab.motif_freq    -> {smiles: freq}
+# vocab.motif_expansion -> {smiles: [char_id, ...]}
+```
+
+### 2. Load model
+
+```python
+import torch
+from dataclasses import dataclass, field
+from typing import Optional
+
+# Option A: clone https://github.com/... (STAMP repo) to get `stamp.benchmark.lm`
+from stamp.benchmark.lm import LMConfig, TinyDecoderLM
+
+ckpt = torch.load("model.pt", map_location="cpu", weights_only=False)
+cfg = LMConfig(**ckpt["cfg"])
+cfg.use_adaln = True  # AO-GPT arch
+model = TinyDecoderLM(vocab_size=len(vocab.itos), cfg=cfg, bidirectional=False)
+
+state = ckpt["model_state"]
+# Strip torch.compile prefix if present.
+if any(k.startswith("_orig_mod.") for k in state):
+    state = {k.replace("_orig_mod.", "", 1): v for k, v in state.items()}
+model.load_state_dict(state)
+model.eval().cuda()
+```
+
+### 3. Sample (AR, top-p)
+
+```python
+import torch
+from hybrid_vocab import is_stamp_structural
+
+BOS, EOS = vocab.bos_id, vocab.eos_id
+PAD, UNK, MASK = vocab.pad_id, vocab.unk_id, vocab.mask_id
+struct_ids = {vocab.stoi[t] for t in vocab.itos if is_stamp_structural(t)}
+suppress = {PAD, BOS, MASK, UNK}
+T, P = 0.95, 0.85
+n, max_new = 64, 64
+
+@torch.no_grad()
+def sample(n_samples=64):
+    x = torch.full((n_samples, 1), BOS, dtype=torch.long, device="cuda")
+    finished = torch.zeros(n_samples, dtype=torch.bool, device="cuda")
+    for step in range(max_new):
+        orders = torch.arange(x.size(1), device="cuda").unsqueeze(0).expand(x.size(0), -1)
+        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+            logits = model(x[:, -cfg.max_seq_len:], orders=orders)[:, -1, :].float()
+        for sid in suppress:
+            logits[:, sid] = float("-inf")
+        # top-p
+        sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1)
+        sorted_probs = torch.softmax(sorted_logits / T, dim=-1)
+        cum = torch.cumsum(sorted_probs, dim=-1)
+        remove = cum > P
+        remove[..., 1:] = remove[..., :-1].clone()
+        remove[..., 0] = False
+        sorted_logits = sorted_logits.masked_fill(remove, float("-inf"))
+        logits = torch.zeros_like(logits).scatter_(-1, sorted_idx, sorted_logits)
+        probs = torch.softmax(logits / T, dim=-1)
+        nxt = torch.multinomial(probs, 1).squeeze(-1)
+        nxt = torch.where(finished, torch.full_like(nxt, EOS), nxt)
+        x = torch.cat([x, nxt.unsqueeze(1)], dim=1)
+        finished = finished | (nxt == EOS)
+        if finished.all():
+            break
+    return x
+```
+
+### 4. Decode token stream → SMILES
+
+```python
+def decode_to_stamp_tokens(ids):
+    """Flush character runs to motif SMILES at structural token boundaries."""
+    special = {PAD, BOS, EOS, MASK, UNK}
+    out, buf = [], []
+    for i in ids:
+        if i in special: continue
+        tok = vocab.itos[i]
+        if i in struct_ids:
+            if buf: out.append("".join(buf)); buf = []
+            out.append(tok)
+        else:
+            buf.append(tok)
+    if buf: out.append("".join(buf))
+    return out
+
+# Then run through the STAMP codec in the stamp repo:
+#     from stamp.benchmark.representations import build_representation
+#     rep = build_representation("stamp")
+#     text = rep.detokenize(stamp_tokens)
+#     mol  = rep.codec.decode_stamp_to_mol(text)
+```
+
+## Sample outputs
+
+Ten representative draws from this checkpoint (all drug-like, QED ≥ 0.6 ∧ SA ≤ 4):
+
+```
+Cn1nc(CNCc2cc(Cl)ccc2Cl)n(C)c1=O                   QED=0.935 SA=2.46 MW=300
+CCn1ncc(NC[C@@H]2CCCC[C@@H]2C)c(Br)c1=O            QED=0.923 SA=3.31 MW=327
+Cc1cccc(Cl)c1NC(=O)CN1CCO[C@@H](C(F)F)CC1          QED=0.921 SA=2.86 MW=332
+CCN1CCN(CC(=O)Nc2cc(C(F)(F)F)ccc2Cl)CC1            QED=0.908 SA=1.94 MW=349
+COc1ccc(F)c(CNC(=O)C2=CCCCC2)c1                    QED=0.907 SA=2.16 MW=263
+CN1CC[C@@H]2[C@@H](CCCN2C(=O)NCc2ccc(OC(F)F)cc2)C1 QED=0.905 SA=3.03 MW=353
+CC[C@H](C(=O)NCc1c(F)cc(F)cc1F)N1CCCC1=O           QED=0.905 SA=2.93 MW=314
+Cc1ccc(C2CCN(C(=O)NCc3cccc(F)c3F)CC2)c(=O)n1C      QED=0.897 SA=2.45 MW=375
+CN(CC(=O)NCc1ccccc1)C(=O)C12CC3CC(CC(C3)C1)C2      QED=0.896 SA=3.37 MW=340
+NCC1CCN(c2cc3c(cc2F)c(=O)c(C(=O)O)cn3C2CC2)C1      QED=0.884 SA=2.96 MW=345
+```
+
+## Architecture notes
+
+- **AO-GPT**: decoder-only transformer with causal attention over a
+  shuffled token order per batch (random permutation of middle tokens,
+  BOS/EOS pinned at ends). Target position is conditioned via AdaLN so
+  the model learns "any-order" decoding.
+- **Hybrid vocab**: structural tokens + SMILES char tokens + atomic motif
+  tokens share a single id space. At training time, atomic motif tokens
+  may be expanded to their SMILES char form with a
+  log-frequency-weighted probability (`HybridVocab.fallback_prob`) so
+  the model is not brittle at char-level decoding.
+- **Decoder**: the STAMP structural tokens delimit motifs; consecutive
+  character tokens between structural tokens concatenate to a single
+  motif SMILES, which the STAMP codec parses to a molecule via a stack
+  machine with safety fallbacks.
+
+## License
+
+Apache-2.0.
+
+## Citation
+
+Cite the STAMP representation paper and this repository. (Placeholder —
+fill in with your actual citation info.)

config.json

@@ -0,0 +1,50 @@
+{
+  "model_type": "ao_gpt_hybrid",
+  "architecture": "TinyDecoderLM",
+  "vocab_size": 2481,
+  "atomic_motifs": 2387,
+  "freq_cutoff": 5000,
+  "d_model": 512,
+  "n_heads": 8,
+  "n_layers": 8,
+  "d_ff": 2048,
+  "max_seq_len": 64,
+  "dropout": 0.1,
+  "use_adaln": true,
+  "bidirectional": false,
+  "dtype": "bfloat16",
+  "epoch": 7,
+  "n_params_total": 31099825,
+  "training": {
+    "dataset": "30M STAMP molecules (train split, all_pass=True)",
+    "train_rows": 19148578,
+    "valid_rows_sampled": 20000,
+    "optimizer": "AdamW (fused, bf16)",
+    "lr": 5e-4,
+    "weight_decay": 0.01,
+    "micro_batch_size": 6144,
+    "global_batch_size": 6144,
+    "grad_accum_steps": 1,
+    "random_ratio": 0.9,
+    "torch_compile": true,
+    "fallback_p_low": 0.02,
+    "fallback_p_high": 0.15
+  },
+  "default_sampling": {
+    "temperature": 0.95,
+    "top_p": 0.85,
+    "top_k": 0,
+    "max_new_tokens": 64
+  },
+  "eval": {
+    "N": 1024,
+    "validity_pct": 100.0,
+    "uniqueness_pct": 100.0,
+    "quality_over_valid_pct": 79.16,
+    "genmol_pct": 79.00,
+    "qed_mean": 0.727,
+    "sa_mean": 2.92,
+    "diversity": 0.860,
+    "reference_ar_baseline_genmol_pct": 79.64
+  }
+}

hybrid_vocab.py

@@ -0,0 +1,268 @@
+"""Hybrid motif+char STAMP vocabulary.
+
+Strategy:
+  - High-frequency motifs (freq >= cutoff) stay as atomic tokens — one vocab id
+    per motif. Rare motifs are always expanded to SMILES char tokens, letting
+    the model learn them via atom-level shared structure.
+  - During training, even atomic motifs are expanded to chars with a small
+    probability `p(freq)` (log-interpolated: higher near the cutoff, lower for
+    the very frequent motifs). This closes the train/inference gap for the
+    char form, which is the only form in which OOV motifs can appear.
+  - No special motif-boundary markers are added. The STAMP structural tokens
+    `[J_*]`, `[B_*]`, `[S_*]`, `[END]` already delimit motifs.
+
+Vocab layout (stoi / itos):
+  0..4:   <pad> <bos> <eos> <unk> <mask>
+  5..K-1: STAMP structural tokens ([J_*], [B_*], [S_*], [END], [BACK], [Ring])
+  K..K+C: SMILES char tokens (regex-level SMILES tokens, 49 distinct on our
+          v3 motif corpus)
+  K+C..:  atomic motif SMILES (freq >= cutoff). Single-token-SMILES motifs
+          like ``C`` or ``[CH2]`` are NOT duplicated here; they share the char
+          id that was allocated above.
+
+Tokens stream out of the STAMP canonical form (space-separated motif tokens
++ structural tokens) via ``encode_stream``, which expands rare motifs and
+probabilistically expands atomic motifs.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import random
+import re
+from collections import Counter
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Dict, Iterable, List, Optional, Sequence, Tuple
+
+
+# Covers all 12395 motif SMILES in v3_cm_union.txt (kekulized + isomeric).
+# Multi-char atoms `Br`, `Cl` and bracketed groups `[...]` come first so the
+# regex prefers them over their single-letter prefixes.
+SMILES_RE = re.compile(r"(\[[^\]]+\]|Br|Cl|[BCFINOPS]|[0-9]|\(|\)|=|#|-|\+|/|\\|@)")
+
+
+def smiles_char_tokens(smiles: str) -> List[str]:
+    """Tokenize a SMILES string into regex-level atomic tokens.
+
+    Raises ValueError if the regex does not cover the whole string.
+    """
+    toks = SMILES_RE.findall(smiles)
+    if "".join(toks) != smiles:
+        raise ValueError(f"SMILES char tokenizer did not cover input: {smiles!r} -> {toks!r}")
+    return toks
+
+
+STAMP_STRUCT_RE = re.compile(r"^\[(?:J_\d+|B_\d+|S_(?:R|S|E|Z)|Ring)\]$")
+
+
+def is_stamp_structural(tok: str) -> bool:
+    return bool(STAMP_STRUCT_RE.match(tok)) or tok in {"[END]", "[BACK]"}
+
+
+PAD = "<pad>"
+BOS = "<bos>"
+EOS = "<eos>"
+UNK = "<unk>"
+MASK = "<mask>"
+_SPECIAL = (PAD, BOS, EOS, UNK, MASK)
+
+
+@dataclass
+class HybridVocab:
+    """Vocabulary + atomic-motif expansion table.
+
+    ``itos`` is the master token list (indices = token ids). ``atomic_motifs``
+    maps each atomic-motif string to its id (may overlap char-token ids for
+    length-1 SMILES motifs). ``motif_freq`` stores the raw motif frequency
+    (for fallback-probability scheduling). ``motif_expansion`` caches the
+    SMILES char-token id list for every atomic motif — needed at train time
+    when we roll a fallback expansion.
+    """
+
+    itos: List[str]
+    atomic_motifs: Dict[str, int] = field(default_factory=dict)
+    motif_freq: Dict[str, int] = field(default_factory=dict)
+    motif_expansion: Dict[str, List[int]] = field(default_factory=dict)
+    freq_cutoff: int = 5000
+
+    def __post_init__(self) -> None:
+        self.stoi: Dict[str, int] = {tok: i for i, tok in enumerate(self.itos)}
+        self.pad_id = self.stoi[PAD]
+        self.bos_id = self.stoi[BOS]
+        self.eos_id = self.stoi[EOS]
+        self.unk_id = self.stoi[UNK]
+        self.mask_id = self.stoi[MASK]
+        # Freq bounds for log-interpolated fallback prob
+        freqs = list(self.motif_freq.values()) or [self.freq_cutoff]
+        self._log_freq_max = math.log(max(freqs))
+        self._log_freq_min = math.log(self.freq_cutoff)
+
+    # ---- Building ----
+
+    @classmethod
+    def build(
+        cls,
+        motif_vocab_path: str | Path,
+        freq_cutoff: int = 5000,
+        extra_stamp_tokens: Optional[Sequence[str]] = None,
+    ) -> "HybridVocab":
+        """Read a v3_cm_union-style motif vocab file and build a HybridVocab.
+
+        The file is tab-separated with three columns (smiles, num_heavy_atoms,
+        freq) and a leading JSON-config line.
+        """
+        path = Path(motif_vocab_path)
+        motifs: List[Tuple[str, int]] = []
+        with path.open("r", encoding="utf-8") as f:
+            first = f.readline()
+            try:
+                json.loads(first)
+            except json.JSONDecodeError:
+                # Not a JSON header; rewind by re-reading as a data line.
+                parts = first.rstrip("\n").split("\t")
+                if len(parts) == 3:
+                    motifs.append((parts[0], int(parts[2])))
+            for line in f:
+                parts = line.rstrip("\n").split("\t")
+                if len(parts) != 3:
+                    continue
+                motifs.append((parts[0], int(parts[2])))
+
+        atomic = [(smi, freq) for smi, freq in motifs if freq >= freq_cutoff]
+
+        # Collect SMILES char tokens across ALL motifs (so rare motifs can be
+        # expanded) plus the chars appearing inside atomic motifs (for the
+        # training-time fallback expansion).
+        char_set: set[str] = set()
+        for smi, _ in motifs:
+            for t in smiles_char_tokens(smi):
+                char_set.add(t)
+
+        itos: List[str] = list(_SPECIAL)
+
+        # STAMP structural tokens. Default includes the commonly used joint
+        # ([J_0]..[J_29]), branch ([B_1]..[B_12]), stereo and control tokens;
+        # callers may pass `extra_stamp_tokens` for anything custom.
+        stamp_struct = [f"[J_{i}]" for i in range(30)]
+        stamp_struct += [f"[B_{i}]" for i in range(1, 13)]
+        stamp_struct += ["[S_R]", "[S_S]", "[S_E]", "[S_Z]", "[Ring]", "[END]", "[BACK]"]
+        for t in stamp_struct:
+            if t not in itos:
+                itos.append(t)
+        if extra_stamp_tokens:
+            for t in extra_stamp_tokens:
+                if t not in itos:
+                    itos.append(t)
+
+        # SMILES char tokens (sorted for determinism).
+        for t in sorted(char_set):
+            if t not in itos:
+                itos.append(t)
+
+        # Atomic motifs (freq >= cutoff). Single-token SMILES motifs (e.g., 'C',
+        # '[CH2]') share ids with the SMILES-char entries; multi-char motifs
+        # get fresh ids.
+        # Sorted by descending freq for determinism + for convenient head/tail.
+        for smi, _ in sorted(atomic, key=lambda x: (-x[1], x[0])):
+            if smi not in itos:
+                itos.append(smi)
+
+        # Finalize.
+        vocab = cls(itos=itos, freq_cutoff=freq_cutoff)
+
+        atomic_ids: Dict[str, int] = {}
+        freq_map: Dict[str, int] = {}
+        expansion: Dict[str, List[int]] = {}
+        for smi, freq in atomic:
+            atomic_ids[smi] = vocab.stoi[smi]
+            freq_map[smi] = int(freq)
+            expansion[smi] = [vocab.stoi[t] for t in smiles_char_tokens(smi)]
+
+        vocab.atomic_motifs = atomic_ids
+        vocab.motif_freq = freq_map
+        vocab.motif_expansion = expansion
+        vocab.__post_init__()  # refresh freq bounds
+        return vocab
+
+    # ---- Fallback probability schedule ----
+
+    def fallback_prob(self, freq: int, p_low: float = 0.02, p_high: float = 0.15) -> float:
+        """Log-interpolate fallback prob. Most frequent motif → p_low; at the
+        freq cutoff → p_high. Monotone decreasing in freq."""
+        if freq <= self.freq_cutoff:
+            return p_high
+        f = math.log(max(1, freq))
+        denom = self._log_freq_max - self._log_freq_min
+        t = 0.0 if denom <= 1e-9 else (f - self._log_freq_min) / denom
+        t = max(0.0, min(1.0, t))
+        return p_high * (1.0 - t) + p_low * t
+
+    # ---- Encoding ----
+
+    def _expand_motif_to_chars(self, smi: str) -> List[int]:
+        if smi in self.motif_expansion:
+            return list(self.motif_expansion[smi])
+        # Encode on the fly for rare motifs (always expanded).
+        return [self.stoi.get(t, self.unk_id) for t in smiles_char_tokens(smi)]
+
+    def encode_stream(
+        self,
+        tokens: Sequence[str],
+        training: bool = False,
+        rng: Optional[random.Random] = None,
+        p_low: float = 0.02,
+        p_high: float = 0.15,
+    ) -> List[int]:
+        """Encode a STAMP token stream (list of structural tokens + motif
+        SMILES, as stored in `stamp_smiles`) into token ids.
+
+        During training, atomic motifs may be expanded to SMILES chars with
+        probability `fallback_prob(freq)`. Rare motifs (below cutoff) always
+        expand.
+        """
+        rnd = rng if rng is not None else random
+        out: List[int] = []
+        for tok in tokens:
+            if is_stamp_structural(tok):
+                out.append(self.stoi.get(tok, self.unk_id))
+                continue
+            if tok in self.atomic_motifs:
+                if training:
+                    p = self.fallback_prob(self.motif_freq[tok], p_low=p_low, p_high=p_high)
+                    if rnd.random() < p:
+                        out.extend(self._expand_motif_to_chars(tok))
+                        continue
+                out.append(self.atomic_motifs[tok])
+            else:
+                out.extend(self._expand_motif_to_chars(tok))
+        return out
+
+    # ---- Persistence ----
+
+    def to_dict(self) -> Dict:
+        return {
+            "itos": self.itos,
+            "atomic_motifs": self.atomic_motifs,
+            "motif_freq": self.motif_freq,
+            "motif_expansion": self.motif_expansion,
+            "freq_cutoff": self.freq_cutoff,
+        }
+
+    @classmethod
+    def from_dict(cls, d: Dict) -> "HybridVocab":
+        return cls(
+            itos=list(d["itos"]),
+            atomic_motifs=dict(d.get("atomic_motifs", {})),
+            motif_freq=dict(d.get("motif_freq", {})),
+            motif_expansion={k: list(v) for k, v in d.get("motif_expansion", {}).items()},
+            freq_cutoff=int(d.get("freq_cutoff", 5000)),
+        )
+
+    def save(self, path: str | Path) -> None:
+        Path(path).write_text(json.dumps(self.to_dict(), ensure_ascii=False))
+
+    @classmethod
+    def load(cls, path: str | Path) -> "HybridVocab":
+        return cls.from_dict(json.loads(Path(path).read_text()))

model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:611298f400ec12cfa68e1042e9d5c84c542add640b53db950e3e2115e5c7443e
+size 124436358