How yahiyav/generative-ai-for-drug-discovery works
Purpose and entry points
This is an end-to-end molecular-generation workflow: train a character-level SMILES VAE, generate candidate molecules, predict their properties with a graph neural network, evaluate chemical quality, and expose generation through a Flask dashboard.
Primary entry points:
• Data acquisition/preprocessing: data/download_data.py
• VAE training: src/train_vae.py
• GNN training: src/train_gnn.py
• Batch generation/evaluation: src/evaluate.py
• Web/API runtime: server.py (with app.py also present as an application-level module)
1. Dataset and chemical feature boundary
The pipeline starts from ZINC250K drug-like molecules, downloads/cleans them, and computes molecular descriptors into the persisted training dataset at data/processed/molecules.csv.
Cheminformatics conversion is centralized in src/mol_utils.py: it bridges SMILES text, RDKit molecules, graph representations, fingerprints, and tokenization. This is the key boundary between chemistry-domain objects and ML-ready tensors.
Paths: data/download_data.py, data/processed/molecules.csv, src/mol_utils.py.
2. Shared experiment configuration
Central hyperparameters—including dataset size, VAE latent dimension/batching/epochs, and GNN hidden size/epochs—are defined in src/config.py. Training and evaluation stages should consume this module rather than independently defining model/data assumptions.
Path: src/config.py.
3. SMILES variational autoencoder
The generative model encodes tokenized SMILES into a continuous latent distribution, samples from that space, and decodes token sequences back into candidate SMILES. Its architecture resides in src/vae_model.py; training is orchestrated by src/train_vae.py.
The VAE output is textual molecular candidates, not directly graph tensors. Generated SMILES must therefore pass RDKit validation and graph conversion before downstream property prediction.
Paths: src/vae_model.py, src/train_vae.py, src/mol_utils.py.
4. Molecular graph property predictor
The predictor consumes molecular graphs derived from valid SMILES and produces multiple molecular property estimates: molecular weight, LogP, TPSA, and QED. Model definition and optimization are separated into src/gnn_model.py and src/train_gnn.py.
This stage depends on consistent atom/bond graph featurization from src/mol_utils.py, creating a contract between generation, preprocessing, and inference.
Paths: src/gnn_model.py, src/train_gnn.py, src/mol_utils.py.
5. Generation and evaluation pipeline
src/evaluate.py composes trained VAE generation, validity filtering, novelty checks against the processed dataset, structural similarity/fingerprint analysis, Lipinski drug-likeness assessment, and GNN property inference.
It materializes reproducible artifacts rather than only console output: aggregate metrics in results/evaluation_metrics.json, generated SMILES in results/generated_molecules.csv, and predicted properties in results/generated_properties.csv.
Paths: src/evaluate.py, results/evaluation_metrics.json, results/generated_properties.csv.
6. Interactive serving layer
The Flask server provides the interactive runtime on port 7860, serving the static dashboard and connecting browser actions to molecule generation/evaluation behavior. The browser client is a static HTML/CSS/JavaScript UI rather than a separate frontend build system.
Paths: server.py, static/index.html, static/script.js.
7. Runtime and infrastructure
The core runtime is Python 3.9+ with PyTorch for VAE training/inference, PyTorch Geometric for graph convolution, RDKit for chemistry parsing/descriptors/fingerprints, and Flask for HTTP serving. GPU acceleration is supported, while CPU training is supported but slower.
The repository is intended for Docker-based Hugging Face Spaces deployment, where the Flask process is exposed on port 7860. ZINC is the external source dataset; no database, message queue, or remote model-serving dependency is indicated.
Main entry points: evaluate.py Batch evaluation · molecules.csv Processed molecules · mol_utils.py Chemistry conversion · train_vae.py VAE training
Components 16
Data & chemistry
- ZINC250K · Source of drug-like molecules.
- Download & preprocess
data/download_data.py - Processed molecules · Cleaned SMILES and molecular descriptors.
data/processed/molecules.csv - Chemistry conversion · Converts SMILES, molecules, graphs, fingerprints, and tokens.
src/mol_utils.py
Model training
- Experiment config
src/config.py - VAE training
src/train_vae.py - SMILES VAE · Encodes tokenized SMILES to latent distributions and decodes candidates.
src/vae_model.py - GNN training
src/train_gnn.py - Property GNN · Predicts molecular weight, LogP, TPSA, and QED.
src/gnn_model.py
Generation & evaluation
- Batch evaluation · Generates, validates, assesses quality, and predicts properties.
src/evaluate.py - Valid candidate SMILES · VAE output after RDKit validation and novelty checks.
results/generated_molecules.csv - Generated properties
results/generated_properties.csv - Evaluation metrics · Validity, novelty, similarity, and Lipinski assessment results.
results/evaluation_metrics.json
Interactive serving
- Flask server · Serves generation behavior on port 7860.
server.py - Browser dashboard
static/index.html - Dashboard actions
static/script.js
Connections 21
- ZINC250K → Download & preprocess: download
- Download & preprocess → Processed molecules: clean and persist
- Download & preprocess → Chemistry conversion: compute descriptors
- Experiment config → VAE training: hyperparameters
- Experiment config → GNN training: hyperparameters
- Processed molecules → VAE training: training SMILES
- Chemistry conversion → VAE training: tokenization
- VAE training → SMILES VAE: optimize
- Processed molecules → GNN training: molecules
- Chemistry conversion → GNN training: graph features
- GNN training → Property GNN: optimize
- SMILES VAE → Batch evaluation: generate SMILES
- Processed molecules → Batch evaluation: novelty reference
- Chemistry conversion → Batch evaluation: validate, graph, fingerprints
- Batch evaluation → Property GNN: valid molecular graphs
- Batch evaluation → Valid candidate SMILES: write
- Batch evaluation → Generated properties: write predictions
- Batch evaluation → Evaluation metrics: write quality metrics
- Browser dashboard → Dashboard actions: loads
- Dashboard actions → Flask server: generation requests
- Flask server → Batch evaluation: invoke generation