UNFRAMED

UNFRAMED is a script-based molecular graph generation and optimization project. Given one or more seed molecules, the Graph Alteration Position Model (GAPM) selects positions to alter, append, or delete, and the Graph Fragment Prediction Model (GFPM) proposes new atoms or fragments. The generated molecules are filtered by similarity, synthetic accessibility, and optional molecular-weight constraints, then ranked by the selected optimization objective.

Run the entry scripts directly from the project directory; installing UNFRAMED itself as a Python package is not required.

Project layout

• batch_evaluate.py: batch molecule-generation entry script.
• train_gfpm.py: GFPM fragment-prediction training entry script.
• train_gapm.py: GAPM edit-position training entry script.
• trajectory_monitor.py: trajectory inspection script.
• unframed/: chemistry, model, and optimization implementation.
• unframed/training/: shared training datasets, graph construction, and loops.
• oracles/: QED and penalized-logP optimization objectives.
• checkpoints/GFPM.ckpt: default fragment/cloze model checkpoint.
• checkpoints/GAPM.ckpt: default position model checkpoint.
• data/: vocabularies, functional-group resources, and example inputs.
• requirements.txt: third-party Python dependencies.

Environment

Python 3.10 is required. Molecule generation and model training can run on either CUDA or CPU. CUDA is recommended for performance, while CPU mode is available for systems without a compatible GPU.

The provided requirements.txt pins CUDA 12.4 builds of PyTorch and DGL. Use it for the CUDA environment:

pip install -r requirements.txt

For a CPU-only environment, install CPU-compatible builds of PyTorch and DGL for the target platform instead of the CUDA-pinned builds, together with the remaining dependencies.

TensorBoard is installed separately and is required for scalar logging during GFPM or GAPM training:

pip install tensorboard

If TensorBoard logging is not needed, omit this installation and run either training script with --disable-tensorboard.

Molecule generation

Enter the project directory and inspect all generation options:

cd UNFRAMED
python batch_evaluate.py --help

The optimizer supports exactly two objectives:

• -oracle qed: maximize QED improvement relative to the seed molecule.
• -oracle plogp: maximize penalized-logP improvement relative to the seed molecule.

A small example is:

python batch_evaluate.py \
  -dir result/example \
  -smi data/test_set/plogp/test.txt \
  -oracle plogp \
  --device cuda \
  -gen 3 \
  -pop 100

To run the same generation workflow without CUDA, use --device cpu. CPU inference uses the same checkpoints but is usually slower:

python batch_evaluate.py \
  -dir result/example_cpu \
  -smi data/test_set/plogp/test.txt \
  -oracle plogp \
  --device cpu \
  -gen 3 \
  -pop 100

The program appends oracle and sampling settings to the -dir prefix when it creates the final result directory. The default checkpoints, vocabulary, and functional-group file are resolved relative to batch_evaluate.py, so no PYTHONPATH setup is needed.

Each line in the -smi input file must contain exactly one non-empty SMILES string. For every seed, one generation performs the following steps:

1. GAPM samples alter, append, and delete positions.
2. GFPM proposes replacement or appended vocabulary items for those positions.
3. Invalid molecules and candidates outside the similarity, synthetic-accessibility, or molecular-weight constraints are removed.
4. The oracle ranks the remaining candidates and retains at most -pop molecules for the next generation.

Optimization stops after -gen generations or earlier if no valid candidates remain.

Generation parameter reference

Options accept the long name and, where shown, a short alias. Required parameters have no default.

Parameter	Default	Meaning
--dir_root, -dir	required	Prefix for the result directory. UNFRAMED appends the objective, sampling, constraint, population, and generation settings to this value.
--smiles_file, -smi	required	Plain-text input file containing exactly one non-empty seed SMILES per line. Blank lines are not skipped.
--group_num, -ng	8	Number of worker processes/groups used for parallel oracle screening.
--device	cuda	GFPM/GAPM inference device. Choose cuda or cpu. CUDA mode requires torch.cuda.is_available() to be true; CPU mode does not require a GPU but is generally slower.
--oracle_name, -oracle	plogp	Optimization objective. Choose qed for QED improvement or plogp for normalized penalized-logP improvement. In both cases the score is the candidate value minus the original seed value.
--generation, -gen	15	Maximum number of optimization generations for each seed.
--population_size, -pop	1000	Maximum number of top-ranked molecules retained as parents for the next generation.
--sim_cons, -sim	0.6	Minimum Tanimoto similarity to the original seed. Similarity uses radius-2, 2048-bit feature Morgan fingerprints; candidates must have similarity greater than or equal to this value.
--sa_cons_low, -sa_low	2.0	Exclusive lower bound for the RDKit synthetic-accessibility (SA) score. A candidate must have SA > sa_low.
--sa_cons_high, -sa_high	2.5	Inclusive upper bound for the SA score. A candidate must have SA <= sa_high; lower SA generally indicates easier synthesis.
--GFPM_path, -GFPM	checkpoints/GFPM.ckpt	GFPM checkpoint used to predict replacement and appended atoms/fragments.
--GAPM_path, -GAPM	checkpoints/GAPM.ckpt	GAPM checkpoint used to predict edit actions and positions.
--vocabulary_path, -vocab	data/train_set/chembl/vocabulary_new.txt	Base atom/fragment vocabulary shared by GFPM and GAPM. It must be compatible with both checkpoints.
--func_group_path, -fung	data/fgs/functionalgroups_cleaned_1R_combine_unique_std_rm2+.txt	Functional-group definitions appended to the base vocabulary. This file must also match the vocabulary used to train the checkpoints.
--num_alter_position_sample, -num_al_pos	10	Maximum number of existing graph nodes selected by GAPM as replacement positions per parent molecule.
--num_append_position_sample, -num_ap_pos	10	Maximum number of existing graph nodes selected by GAPM as attachment positions per parent molecule.
--num_delete_position_sample, -num_del_pos	3	Number of leaf-node deletion positions sampled per parent molecule. Only leaf nodes are eligible.
--num_alter_sample, -num_al	15	Number of top GFPM atom/fragment predictions tried at each selected replacement position.
--num_append_sample, -num_ap	15	Number of top GFPM atom/fragment predictions tried at each selected attachment position.
--enable_funcgroup, -enable_fg	True	Whether GFPM proposals may use entries from the functional-group vocabulary in addition to the base vocabulary.
--optim_step, -opt_step	-1	Optional generation-dependent oracle-score transformation. -1 disables it. Otherwise, scores at least optim_step * (generation_index + 1) are negated before ranking; leave disabled unless this behavior is explicitly required.
--trajectory_dir, -traj	<result_dir>/trajectories	Directory for per-seed JSONL trajectory files. The computed result directory is used when this option is omitted.
--disable_trajectory	off	Disable JSONL trajectory recording. This overrides --trajectory_dir.
--min_mw	none	Optional inclusive minimum RDKit molecular weight. Candidates with lower molecular weight are discarded.
--max_mw	none	Optional inclusive maximum RDKit molecular weight. Candidates with higher molecular weight are discarded.

--min_mw cannot be greater than --max_mw. Increasing any of the five position/proposal sampling parameters expands the search but also increases runtime and memory use.

The script uses these checkpoint paths by default:

checkpoints/GFPM.ckpt
checkpoints/GAPM.ckpt

Both checkpoints use the following vocabulary by default:

data/train_set/chembl/vocabulary_new.txt

Override them when needed with -GFPM <path> and -GAPM <path>.

Generation outputs

For input line index <i>, the result directory contains:

• <i>.pkl: all unique accepted molecules accumulated for that seed, including the seed itself unless removed by a molecular-weight constraint.
• <i>_gen_<g>.pkl: valid candidates produced before oracle selection in generation <g>.
• trajectories/<i>.jsonl: seed, candidate, selected, and per-generation summary events, unless trajectory recording is disabled.
• logfile.txt: seed processing and existing-result molecular-weight filtering messages.

Model training

The source-project ChEMBL split is included at:

data/train_set/chembl/train.csv
data/train_set/chembl/test.csv

Preparing a custom training dataset

GFPM and GAPM use the same molecule-only dataset. Property values and manual class labels are not required: both scripts construct their supervision from each input molecule during training.

Training and validation CSV files

Prepare separate training and validation files as UTF-8 CSV. By default, each file must have a case-sensitive column named smiles:

smiles
CCO
c1ccccc1
CC(=O)O
CCN(CC)CC

Additional columns are allowed and ignored. If the SMILES column has another name, select it with --smiles-column:

molecule_id,canonical_smiles,source
mol_0001,CCO,example
mol_0002,c1ccccc1,example

python train_gfpm.py \
  --train-file data/train_set/custom/train.csv \
  --valid-file data/train_set/custom/valid.csv \
  --smiles-column canonical_smiles \
  --experiment custom

The CSV loader applies the following rules:

• The first row must be a header, and the selected SMILES column must exist in both files.
• Leading and trailing whitespace is removed from each SMILES value; empty values are skipped.
• Extra columns are ignored, and duplicate molecules are not removed automatically.
• Training and validation splitting is not automatic. Supply two separate files with no unintended overlap.
• SMILES must be parseable by RDKit. Single-component, sanitized molecules using aromatic, single, double, or triple bonds are recommended.
• Every atom token needed after functional-group detection must occur in the base vocabulary. Out-of-vocabulary chemistry can cause samples to be rejected or GFPM graph construction to fail.

Malformed SMILES may be skipped during batch collation, but the dataset should be validated in advance rather than relying on this behavior. Each file must contain at least one non-empty usable molecule.

Base vocabulary file

--vocabulary-path points to a plain UTF-8 text file with one token per line and no header. Tokens may be atom symbols or canonical ring-fragment SMILES:

C
O
N
F
Cl
Br
c1ccccc1
c1ccncc1
C1CCCCC1

The graph builder represents recognized rings as single fragment nodes. Ring entries must use the same canonical spelling produced by RDKit. A ring that is absent from the vocabulary is decomposed into its atoms, so all of those atom symbols must then be present as individual tokens.

Do not include blank or duplicate vocabulary lines. Line order determines token IDs and the GFPM output classes, so a trained checkpoint is tied to both the vocabulary contents and their exact order.

Functional-group file

--functional-group-path points to another plain UTF-8 text file with one valid RDKit SMARTS pattern per line and no header, for example:

[C]#[N&D1]
[C]([O&D1])=[O&D1]
[S]([N&D1])(=[O&D1])=[O&D1]
[N;+](=[O;D1])[O;D1;-]

At startup, these SMARTS entries are appended to the base vocabulary. Their order therefore also determines token IDs. Pattern order can additionally affect which group is selected when multiple functional-group matches overlap; place higher-priority patterns first and keep the file unchanged between training and inference.

The same base vocabulary and functional-group files, including their exact line order, must be used to train GFPM and GAPM and to run batch_evaluate.py. Changing either file changes the model vocabulary size or token mapping and makes existing checkpoints incompatible.

Automatically generated training targets

• GFPM converts each molecule into a fragment graph, randomly masks one node, and learns to recover that node's atom, ring, or functional-group token.
• GAPM synthetically assigns molecules to four graph-state types: unchanged, one altered node, one extra appended node, and one removed leaf node. Node-level edit-position labels are generated automatically from these corruptions.

Consequently, both models need only valid molecular structures in the CSV; oracle scores such as QED or penalized logP are not training columns.

Custom-data training example

One possible directory layout is:

data/train_set/custom/
├── train.csv
├── valid.csv
└── vocabulary.txt
data/fgs/custom_functional_groups.txt

Train both models with the same data resources:

python train_gfpm.py \
  --train-file data/train_set/custom/train.csv \
  --valid-file data/train_set/custom/valid.csv \
  --vocabulary-path data/train_set/custom/vocabulary.txt \
  --functional-group-path data/fgs/custom_functional_groups.txt \
  --experiment custom

python train_gapm.py \
  --train-file data/train_set/custom/train.csv \
  --valid-file data/train_set/custom/valid.csv \
  --vocabulary-path data/train_set/custom/vocabulary.txt \
  --functional-group-path data/fgs/custom_functional_groups.txt \
  --experiment custom

Before a full run, --max-train-samples and --max-valid-samples can be used for a short end-to-end smoke test. If the custom molecules are fully covered by the bundled vocabulary and functional-group definitions, those two path options may be omitted to keep the defaults.

Both training scripts use vocabulary_new.txt and the default functional-group file, resolved relative to the script location. Train GFPM directly with:

python train_gfpm.py --device cuda --experiment chembl

Train GAPM with:

python train_gapm.py --device cuda --experiment chembl

The two training scripts share the same parameters, but several model-specific defaults differ. By default, checkpoints are written without overwriting the supplied inference models:

checkpoints/training/GFPM/chembl/GFPM_best.ckpt
checkpoints/training/GAPM/chembl/GAPM_best.ckpt

The saved files are plain model state_dict checkpoints and can be passed directly to evaluation:

python batch_evaluate.py \
  -GFPM checkpoints/training/GFPM/chembl/GFPM_best.ckpt \
  -GAPM checkpoints/training/GAPM/chembl/GAPM_best.ckpt \
  ...

Training parameter reference

Parameter	GFPM default	GAPM default	Meaning
--epochs, -ep	500	200	Number of complete passes through the training dataset. Must be positive.
--batch-size, -bs	128	128	Number of input molecules requested per mini-batch. Invalid graph samples may be skipped during collation. Must be positive.
--num-workers, -nw	6	5	Number of PyTorch DataLoader worker processes. Use 0 to load data in the main process; the value cannot be negative.
--num-edge-types, -ne	5	5	Number of edge/bond relation types represented by the model. Changing it changes checkpoint architecture.
--feature-dim, -feat	128	64	Node and edge embedding width used by the attention layers. It must be compatible with --num-heads; changing it changes checkpoint architecture.
--num-heads, -nh	4	4	Number of relation-aware attention heads in each multi-head layer. Changing it changes checkpoint architecture.
--num-multi-layers, -nmul	3	3	Number of stacked relation-aware multi-head attention layers. Changing it changes checkpoint architecture.
--learning-rate, -lr	0.003	0.001	Initial optimizer learning rate. GFPM uses Adam; GAPM uses SGD. Both scripts apply a step learning-rate scheduler.
--experiment, -exp	chembl	chembl	Experiment name used in the default checkpoint and TensorBoard directory paths.
--train-file	data/train_set/chembl/train.csv	same	Training CSV file. It must contain the column selected by --smiles-column.
--valid-file	data/train_set/chembl/test.csv	same	Validation CSV file used for loss reporting and best-checkpoint selection.
--vocabulary-path	data/train_set/chembl/vocabulary_new.txt	same	Base vocabulary used to build graph nodes and GFPM labels. Use the same vocabulary for training and inference.
--functional-group-path	default file under data/fgs/	same	Functional-group vocabulary appended to the base vocabulary. It must remain consistent with inference and saved checkpoints.
--smiles-column	smiles	smiles	Name of the CSV column containing SMILES strings. Empty values are skipped.
--max-train-samples	all	all	Optional positive limit on the number of non-empty training SMILES read from the beginning of the CSV. Useful for smoke tests.
--max-valid-samples	all	all	Optional positive limit on the number of non-empty validation SMILES read from the beginning of the CSV.
--output-dir	checkpoints/training/GFPM/<experiment>	checkpoints/training/GAPM/<experiment>	Checkpoint output directory. Overrides the experiment-derived default.
--log-dir	runs/GFPM/<experiment>	runs/GAPM/<experiment>	TensorBoard scalar-log directory. Overrides the experiment-derived default.
--device	auto	auto	Training device: auto, cpu, cuda, or cuda:N. auto selects CUDA when available and otherwise uses CPU.
--seed	42	42	Seed for Python, NumPy, PyTorch, CUDA, data shuffling, and worker initialization.
--save-every	1	1	Save an epoch-numbered checkpoint every N epochs. The latest and best checkpoints are still updated independently. Must be positive.
--disable-tensorboard	off	off	Disable TensorBoard scalar logging. Checkpoint saving and console loss reporting remain enabled.

The model-architecture parameters must match when a checkpoint is loaded for inference. With the default inference architecture, GFPM uses --feature-dim 128 and GAPM uses --feature-dim 64; both use 5 edge types, 4 heads, and 3 multi-head layers.

Trajectory monitoring

Trajectories are written by default to trajectories/*.jsonl inside the result directory. Use -traj <path> to choose another location or --disable_trajectory to turn recording off.

Monitor a trajectory file or directory directly with:

python trajectory_monitor.py result/.../trajectories --follow

Parameter	Default	Meaning
path	required	Positional path to one trajectory .jsonl file or a directory containing such files.
--follow, -f	off	Keep running and print refreshed summaries instead of exiting after one summary.
--interval	2.0	Seconds between summaries when --follow is enabled.

All entry scripts also support -h or --help to print their command-line usage.