generate.py

#!/usr/bin/env python
"""Generate molecules from a trained ChemLM checkpoint.

Examples
--------
Beam search, ranked by perplexity, valid molecules only::

    python generate.py --checkpoint checkpoints/lstm_pretrained.pt \\
        --vocab checkpoints/vocab.json --num_samples 50 --method beam \\
        --beam_size 10 --filter_valid --output generated_molecules.smi

Bioactivity-conditioned sampling from a fine-tuned model::

    python generate.py --checkpoint checkpoints/lstm_finetuned.pt \\
        --vocab checkpoints/vocab.json --condition active --num_samples 20
"""

from __future__ import annotations

import argparse
from pathlib import Path
from typing import List, Optional, Tuple

import torch
import torch.nn.functional as F

from data import SMILESTokenizer
from models import build_model

try:
    from rdkit import Chem

    RDKIT_AVAILABLE = True
except ImportError:
    RDKIT_AVAILABLE = False


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    p.add_argument("--checkpoint", type=str, required=True)
    p.add_argument("--vocab", type=str, required=True)
    p.add_argument("--num_samples", type=int, default=50)
    p.add_argument("--method", type=str, default="beam", choices=["beam", "sample"])
    p.add_argument("--beam_size", type=int, default=10)
    p.add_argument("--max_len", type=int, default=128)
    p.add_argument("--temperature", type=float, default=1.0)
    p.add_argument("--top_k", type=int, default=0, help="0 disables top-k filtering.")
    p.add_argument(
        "--condition",
        type=str,
        default=None,
        choices=["active", "inactive"],
        help="Bioactivity prefix token for conditioned generation (fine-tuned models only).",
    )
    p.add_argument(
        "--filter_valid",
        action="store_true",
        help="Keep only RDKit-valid, canonicalized molecules.",
    )
    p.add_argument(
        "--output",
        type=str,
        default=None,
        help="Path to write generated SMILES (one per line).",
    )
    p.add_argument("--seed", type=int, default=42)
    p.add_argument(
        "--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu"
    )
    return p.parse_args()


def load_model_and_tokenizer(
    checkpoint_path: str, vocab_path: str, device: torch.device
):
    tokenizer = SMILESTokenizer.load(vocab_path)
    state = torch.load(checkpoint_path, map_location=device)
    model = build_model(
        state["arch"], vocab_size=state["vocab_size"], pad_id=tokenizer.pad_id
    )
    model.load_state_dict(state["model_state_dict"])
    model.to(device).eval()
    return model, tokenizer


def _prefix_ids(tokenizer: SMILESTokenizer, condition: Optional[str]) -> List[int]:
    ids = [tokenizer.bos_id]
    if condition is not None:
        ids.append(tokenizer.condition_id(condition))
    return ids


@torch.no_grad()
def beam_search_generate(
    model,
    tokenizer: SMILESTokenizer,
    beam_size: int,
    max_len: int,
    device: torch.device,
    condition: Optional[str] = None,
) -> Tuple[List[int], float]:
    """Single-sequence beam search. Returns (token_ids, mean_log_prob)."""
    start = _prefix_ids(tokenizer, condition)
    # Each beam: (token_id_list, cumulative_log_prob, finished)
    beams: List[Tuple[List[int], float, bool]] = [(start, 0.0, False)]

    for _ in range(max_len - len(start)):
        candidates: List[Tuple[List[int], float, bool]] = []
        all_finished = all(b[2] for b in beams)
        if all_finished:
            break
        for seq, score, finished in beams:
            if finished:
                candidates.append((seq, score, finished))
                continue
            input_ids = torch.tensor([seq], dtype=torch.long, device=device)
            logits = model(input_ids=input_ids).logits[0, -1, :]
            log_probs = F.log_softmax(logits, dim=-1)
            topk_log_probs, topk_ids = log_probs.topk(beam_size)
            for lp, tid in zip(topk_log_probs.tolist(), topk_ids.tolist()):
                new_seq = seq + [tid]
                new_finished = tid == tokenizer.eos_id
                candidates.append((new_seq, score + lp, new_finished))
        candidates.sort(key=lambda x: x[1] / max(len(x[0]), 1), reverse=True)
        beams = candidates[:beam_size]

    best_seq, best_score, _ = max(beams, key=lambda x: x[1] / max(len(x[0]), 1))
    mean_log_prob = best_score / max(len(best_seq), 1)
    return best_seq, mean_log_prob


@torch.no_grad()
def sample_generate(
    model,
    tokenizer: SMILESTokenizer,
    max_len: int,
    device: torch.device,
    temperature: float = 1.0,
    top_k: int = 0,
    condition: Optional[str] = None,
) -> Tuple[List[int], float]:
    """Stochastic temperature/top-k sampling. Returns (token_ids, mean_log_prob)."""
    seq = _prefix_ids(tokenizer, condition)
    log_prob_sum = 0.0

    for _ in range(max_len - len(seq)):
        input_ids = torch.tensor([seq], dtype=torch.long, device=device)
        logits = model(input_ids=input_ids).logits[0, -1, :] / max(temperature, 1e-5)
        if top_k > 0:
            top_vals, top_idx = logits.topk(min(top_k, logits.size(-1)))
            mask = torch.full_like(logits, float("-inf"))
            mask[top_idx] = top_vals
            logits = mask
        probs = F.softmax(logits, dim=-1)
        next_id = torch.multinomial(probs, num_samples=1).item()
        log_prob_sum += F.log_softmax(logits, dim=-1)[next_id].item()
        seq.append(next_id)
        if next_id == tokenizer.eos_id:
            break

    return seq, log_prob_sum / max(len(seq), 1)


def perplexity_from_mean_log_prob(mean_log_prob: float) -> float:
    return float(torch.exp(torch.tensor(-mean_log_prob)))


def canonicalize(smiles: str) -> Optional[str]:
    if not RDKIT_AVAILABLE:
        return smiles
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return None
    return Chem.MolToSmiles(mol)


def main() -> None:
    args = parse_args()
    torch.manual_seed(args.seed)
    device = torch.device(args.device)
    model, tokenizer = load_model_and_tokenizer(args.checkpoint, args.vocab, device)

    if not RDKIT_AVAILABLE and args.filter_valid:
        print(
            "Warning: RDKit is not installed; --filter_valid cannot verify validity and will be skipped."
        )

    results = []
    for i in range(args.num_samples):
        if args.method == "beam":
            ids, mean_lp = beam_search_generate(
                model,
                tokenizer,
                args.beam_size,
                args.max_len,
                device,
                condition=args.condition,
            )
        else:
            ids, mean_lp = sample_generate(
                model,
                tokenizer,
                args.max_len,
                device,
                temperature=args.temperature,
                top_k=args.top_k,
                condition=args.condition,
            )
        smiles = tokenizer.decode(ids)
        ppl = perplexity_from_mean_log_prob(mean_lp)

        if args.filter_valid and RDKIT_AVAILABLE:
            canon = canonicalize(smiles)
            if canon is None:
                continue
            smiles = canon

        results.append((smiles, ppl))
        print(f"[{i + 1}/{args.num_samples}] ppl={ppl:8.3f}  {smiles}")

    results.sort(
        key=lambda x: x[1]
    )  # rank by ascending perplexity (most confident first)

    if args.output:
        out_path = Path(args.output)
        out_path.write_text("\n".join(smi for smi, _ in results) + "\n")
        print(f"Wrote {len(results)} molecules to {out_path}")


if __name__ == "__main__":
    main()