alphagrammar 0.1.2

AlphaGrammar: Grammar-based molecular representation learning

AlphaGrammar

AlphaGrammar is a grammar-based molecular representation learning framework. It learns a hyperedge-replacement grammar (HRG) over molecular graphs via Monte Carlo Tree Search (MCMC) with a learned agent.

Installation

pip install -e /path/to/AlphaGrammar_pkg

Or from PyPI (once published):

pip install alphagrammar

Pretrained models

After installation, copy the pretrained model files to the package data directory:

DATA_DIR=$(python -c "import alphagrammar, os; print(os.path.join(os.path.dirname(alphagrammar.__file__), 'data'))")
cp /path/to/AlphaGrammar/ckpts/vocab_epoch5.pkl $DATA_DIR/
cp /path/to/AlphaGrammar/ckpts/best_agent_epoch0_R0.0000.pkl $DATA_DIR/
cp /path/to/AlphaGrammar/GCN/supervised_contextpred.pth $DATA_DIR/

Usage

Command-line interface

# Parse a single SMILES string (rollout mode)
alphagrammar parse "CCO"

# Parse from a file (one SMILES per line)
alphagrammar parse molecules.smi

# Parse with Bolinas parser (10-second timeout per molecule)
alphagrammar parse "CCO" --timeout 10

# Use only top-100 rules from vocab
alphagrammar parse "CCO" --vocab_size 100

Python API

from alphagrammar import MoleculeDataset, _collate_mol_batch, RuleStats
from alphagrammar.grammar_generation import MCMC_sampling
from alphagrammar.agent import Agent
import torch, pickle

# Load pretrained models
with open("data/vocab_epoch5.pkl", "rb") as f:
    rule_stats = pickle.load(f)

agent = Agent(feat_dim=300, hidden_size=256)
agent.load_state_dict(torch.load("data/best_agent_epoch0_R0.0000.pkl"))
agent.eval()

# Build input dataset
dataset = MoleculeDataset(["CCO", "c1ccccc1"], GNN_model_path="data/supervised_contextpred.pth")
batch = [dataset[i] for i in range(len(dataset))]
input_graphs_dict = _collate_mol_batch(batch)

# Run MCMC sampling
results, rules_per_mol, sequential_steps = MCMC_sampling(
    "output_dir", agent, input_graphs_dict, MCMC_size=1, debug=True
)

Package structure

src/alphagrammar/
├── __init__.py          # Public API
├── cli.py               # argparse CLI entry point
├── core.py              # Core functions (bolinas_evaluate, MoleculeDataset, RuleStats, ...)
├── grammar_generation.py# MCMC sampling and grammar generation
├── agent.py             # Neural agent (policy network)
├── hrg_td_parser_undirected.py  # Bolinas HRG parser
├── private/             # Internal hypergraph / grammar data structures
├── fuseprop/            # Molecular fragmentation utilities
├── GCN/                 # Graph neural network feature extraction
├── data/                # Pretrained model files (copy here after install)
└── commands/
    └── parse.py         # 'alphagrammar parse' subcommand