README.md

---
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.)