ChemFLM - Unified Chemical Fingerprint Language Model

Status: Phase 1 & 2 Complete ✅ | Phase 3 In Progress 🔄

ChemFLM unifies and extends two previous implementations (chemfpformer and chemfpformer-main) into a comprehensive, production-ready repository for molecular property prediction using transformers.

🎯 Project Overview

ChemFLM is a pretrained ModernBERT-based transformer model trained on molecular SMILES and DNA-encoded library (DEL) fingerprints. The model performs hit prediction on DEL fingerprints through fine-tuning, with a strong focus on out-of-distribution (OOD) generalization to molecules significantly different from the training set.

Key Features

✅ Large-scale pretraining across multiple molecular datasets (40M SMILES, CHEMBL, MOSES, DEL libraries)
✅ Multiple OOD evaluation strategies to rigorously test generalization
✅ Dynamic split generation for flexible cross-validation
✅ Advanced evaluation metrics (PR-AUC, ROC-AUC, Precision@k, Enrichment@k)
✅ Span masking for learning SMILES-fingerprint relationships
✅ Configuration-driven experiments using Hydra

📊 Implementation Status

✅ Phase 1: Core Infrastructure - COMPLETE

Repository structure and organization
HuggingFace-compatible tokenizer
Base model architecture (ModernBERT)
Configuration system (Hydra)
Utility functions (seed, logging)

✅ Phase 2: Data Pipeline - COMPLETE

Fingerprint Generation (8 types: ECFP4/6, FCFP4/6, MACCS, Avalon, AtomPair, TopTor)
On-demand Clustering with hierarchy rectification
Three OOD Split Strategies:
- Hierarchical clustering (3 levels: 0.85/50%, 0.75/25%, 0.65/10%)
- Library-based splits (3 levels: 30%, 20%, 10% exclusion)
- Building block overlap (k=0, 1, 2 test buckets)
Dynamic split architecture (splits generated on-demand during CV)
Comprehensive data preparation script with metadata tracking

🔄 Phase 3: Model Training - IN PROGRESS

Pretraining infrastructure (needs verification)
Fine-tuning with LoRA and focal loss
Early stopping and model selection

✅ Phase 4-6: Evaluation + Analysis (Refactored)

Cross-validation output analysis has been refactored into a reusable, config-driven framework under analysis/src/.
Metrics, curves, ranking metrics, statsig exports, and advanced analyses are generated by thin notebooks in analysis/notebooks/.

See:

docs/ANALYSIS_FRAMEWORK_USAGE.md
analysis/notebooks/README.md

See docs/IMPLEMENTATION_STATUS.md for detailed progress tracking.

🏗️ Architecture Overview

Legacy Implementations (Now Unified)

ChemFLM combines the best features from two previous implementations:

From chemfpformer:

✅ Large-scale pretraining (40M SMILES, CHEMBL, MOSES)
✅ Hierarchical clustering for OOD evaluation
- Now implemented with dynamic on-demand generation
- Hierarchy rectification for proper nesting
✅ Baseline comparisons (LGBM, Decision Trees)
✅ End-to-end workflows from pretraining to fine-tuning

From chemfpformer-main:

✅ Library-based splits for OOD testing
✅ Advanced metrics (Prec@k, Enrich@k, PR/ROC curves)
✅ Cross-validation framework for statistical significance
✅ Comprehensive baselines (LGBM, XGBoost)

New in ChemFLM:

🆕 Dynamic split generation - Splits created on-demand during CV
🆕 Building block overlap splits - Novel OOD strategy with k=0,1,2 evaluation
🆕 Unified vocabulary - Single tokenizer for pretrain and finetune
🆕 Span masking - Learn SMILES-fingerprint relationships (configurable)
🆕 On-demand clustering - Memory-efficient with caching
🆕 Enhanced configurations - Complete Hydra-based system

🚀 Quick Start

Installation

# Clone repository
git clone git@github.com:GenerativeDrugDiscovery/ChemFLM.git
cd ChemFLM

# Activate conda env (recommended)
conda activate chemfp

# Install dependencies
pip install -r requirements.txt

If you need a more robust activation (e.g., non-interactive shells), you can also:

source scripts/activate_chemfp.sh

Data Preparation

# Prepare data with default configuration (DEL only)
python scripts/prepare_data.py

# Prepare data with SMILES + fingerprints
python scripts/prepare_data.py --config-name data_prep/data_prep_smiles_fp

# Prepare data with large-scale pretraining (40M SMILES)
python scripts/prepare_data.py --config-name data_prep/data_prep_40m

Dynamic Split Generation (During CV)

from src.data.splits import create_hierarchy_split
from datasets import load_from_disk

# Load finetune dataset (saved whole, unsplit)
dataset = load_from_disk("processed_data/default/dataset/ft_DCAF7")

# Generate hierarchical split dynamically
train_idx, val_idx, test_idx = create_hierarchy_split(
    dataset,
    dataset_path="processed_data/default",
    dataset_name="DCAF7",
    seed=42,
    config={'threshold': 0.85, 'ood_percentage': 0.5, 'val_fraction': 0.1}
)

# Or library split
from src.data.splits import create_library_split
train_idx, val_idx, test_idx = create_library_split(
    dataset,
    dataset_path="processed_data/default",
    dataset_name="DCAF7",
    seed=42,
    config={'exclusion_percentage': 0.30, 'min_library_size': 1000, 'val_fraction': 0.1}
)

# Or building block split (returns 3 test buckets)
from src.data.splits import create_building_block_split
train_idx, val_idx, test_dict = create_building_block_split(
    dataset,
    dataset_path="processed_data/default",
    dataset_name="DCAF7",
    seed=42,
    config={'holdout_fraction': 0.40, 'label_aware_sampling': True, 'val_fraction': 0.1}
)
# test_dict contains: 'test_0bb', 'test_1bb', 'test_2bb'

📁 Repository Structure

ChemFLM/
├── configs/                    # Hydra configurations
│   ├── config.yaml            # Main config
│   ├── data/                  # Data prep configs
│   │   ├── data_prep.yaml
│   │   ├── data_prep_del_only.yaml
│   │   └── data_prep_smiles_fp.yaml
│   ├── splits/                # Split strategy configs
│   │   ├── hierarchy.yaml
│   │   ├── library.yaml
│   │   └── building_block.yaml
│   ├── model/                 # Model configs
│   │   ├── modernbert_base.yaml
│   │   └── modernbert_small.yaml
│   └── training/              # Training configs
│       └── default.yaml
│
├── src/                       # Source code
│   ├── data/                  # Data processing
│   │   ├── fingerprints.py    # 8 fingerprint types
│   │   ├── clustering.py      # On-demand clustering
│   │   ├── splits.py          # 3 OOD split strategies
│   │   ├── tokenizer.py       # HuggingFace tokenizer
│   │   ├── transforms.py      # Molecular tokenization
│   │   └── dataloader.py      # PyTorch data loaders
│   ├── model/                 # Model components
│   │   ├── chemfpformer_model.py
│   │   ├── pretrain_model.py
│   │   ├── finetune_model.py
│   │   └── finetuning_strategies.py  # LoRA, etc.
│   └── utils/                 # Utilities
│       ├── seed.py
│       └── logging_utils.py
│
├── scripts/                   # Execution scripts
│   └── prepare_data.py        # Main data prep script
│
├── docs/                      # Documentation
│   ├── IMPLEMENTATION_STATUS.md
│   ├── IMPLEMENTATION_PLAN_OVERVIEW.md
│   ├── PHASE2_IMPLEMENTATION_PLAN_REVISED.md
│   ├── IMPLEMENTATION_SUMMARY.md
│   └── ANALYSIS_FRAMEWORK_USAGE.md
│
├── analysis/                   # Analysis framework + thin notebooks
│   ├── configs/                # Analysis configs (YAML)
│   ├── notebooks/              # Thin Jupyter notebooks (call analysis/src)
│   ├── reports/                # Generated outputs (figures/results)
│   └── src/                    # Reusable analysis modules
│
└── data/                      # Raw data (not in repo)
    ├── SMILES-40M/
    ├── CHEMBL/
    ├── MOSES/
    └── [DEL libraries]/

📈 Analysis (Refactored)

The analysis workflow is now config-driven and writes outputs into:

analysis/reports/<analysis.name>/

1) Activate environment

conda activate chemfp
# or (more robust in scripts/non-interactive shells)
source scripts/activate_chemfp.sh

2) Choose an analysis config

Start from the example:

analysis/configs/example_LRRK2_library_initial_baseline.yaml

3) Run thin notebooks

Open and run (in order as needed):

analysis/notebooks/1_cv_folds_threshold_comparison.ipynb (core tables + metric bar plots)
analysis/notebooks/2_curves.ipynb (mean ROC/PR curves)
analysis/notebooks/3_ranking_metrics.ipynb (precision@k / recall@k / enrichment@k / ndcg@k)
analysis/notebooks/4_summary_export.ipynb (Phase 7-style exports: statsig + recommendations + RESULTS_SUMMARY.md)
analysis/notebooks/5_advanced.ipynb (Phase 8 advanced: fold variability, radar, calibration, ensemble)

4) Output structure

Example:

analysis/reports/<analysis.name>/
  figures/
    metrics_comparison/
    roc_curves/
    pr_curves/
    ranking_metrics/
    advanced/
  results/
    metrics_long.csv
    metrics_summary.csv
    statistical_tests.csv
    ranking_metrics_complete_summary.csv
    ...
  RESULTS_SUMMARY.md

🎨 Key Features Explained

Dynamic Split Architecture

Unlike traditional approaches that pre-generate all splits, ChemFLM uses on-demand split generation:

Data Preparation - Saves finetune datasets whole (unsplit)
Cross-Validation - Generates splits dynamically for each fold
Benefits:
- Simpler file structure
- Easy parameter experimentation
- Memory efficient
- Better reproducibility tracking

Three OOD Strategies

1. Hierarchical Clustering ✅

Based on molecular similarity (Tanimoto distance)
Three levels: threshold (0.85/0.75/0.65) × OOD% (50%/25%/10%)
Rectified hierarchy ensures proper nesting
On-demand cluster generation with caching

2. Library-Based Splits ✅

Groups molecules by library prefix
Three exclusion levels: 30%, 20%, 10%
Small libraries merged to "OTHER" (always in training)
No library appears in both train and test

3. Building Block Overlap ✅

Based on combinatorial chemistry BBs
40% of building blocks held out
Three test buckets by overlap:
- k=0: All 3 BBs novel (highest OOD)
- k=1: 1 BB shared, 2 novel
- k=2: 2 BBs shared, 1 novel
Label-aware sampling for enrichment

Fingerprint Types Supported

ChemFLM supports 8 molecular fingerprint types:

ECFP4, ECFP6 - Extended Connectivity (binary)
FCFP4, FCFP6 - Functional Connectivity (binary)
MACCS - 167-bit MACCS keys
Avalon - Avalon fingerprints
AtomPair - Atom pair fingerprints
TopTor - Topological torsion fingerprints

All with robust error handling and batch processing.

📖 Documentation

docs/IMPLEMENTATION_STATUS.md - Comprehensive progress tracking
docs/IMPLEMENTATION_PLAN_OVERVIEW.md - Project overview and phases
docs/PHASE2_IMPLEMENTATION_PLAN_REVISED.md - Dynamic splits architecture
docs/IMPLEMENTATION_SUMMARY.md - Finalized design decisions
docs/IMPLEMENTATION_PLAN_DATA_PIPELINE.md - Data pipeline details
docs/ANALYSIS_FRAMEWORK_USAGE.md - How to run the refactored analysis pipeline

🔧 Analysis TODOs (Known Gaps)

The refactor intentionally focused on preserving existing analytics while standardizing the workflow. Remaining TODOs:

Multi-run comparisons (pretraining vs no-pretraining)
- Current modules support run_label but the notebooks generally analyze cfg.runs[0].
- TODO: add comparative plots/tables across multiple analysis.runs in a single report.
Generalized “level comparison” module
- The legacy analysis/2. threshold_comparison.ipynb logic is not yet ported into analysis/src/plots/.
- TODO: implement a threshold_comparison.py (generalized to hierarchy/library/building_block “levels”).
Smarter filesystem discovery
- Discovery currently depends on analysis.levels being enumerated in config.
- TODO: optionally auto-discover existing levels/testsets from output folders/files (useful for incomplete runs).
Ensemble alignment robustness
- Ensemble currently assumes transformer and LGBM prediction rows align by order.
- TODO: align on a stable sample key if available (e.g., id / smiles / compound_id).

🧪 Testing & Validation

Data Quality Checks

SMILES parsing success rate
Fingerprint consistency
No data leakage between splits
Label distribution per split
Cluster hierarchy validation

Reproducibility

Fixed seeds throughout
Fold assignments saved to YAML
Metadata tracking for all experiments
Configuration snapshots

🤝 Contributing

ChemFLM is under active development. Phase 3 (model training) and Phase 4 (cross-validation) are the current focus areas.

📄 License

[Add license information]

📚 Citation

[Add citation information]

🔗 References

ModernBERT: [Link to paper/repo]
Original chemfpformer: /home/bhux/workplace/chemfpformer
Original chemfpformer-main: /home/bhux/workplace/chemfpformer-main

Name		Name	Last commit message	Last commit date
analysis		analysis
configs		configs
data		data
scripts		scripts
src		src
.gitignore		.gitignore
README.md		README.md
__init__.py		__init__.py
config.yaml		config.yaml
requirements.txt		requirements.txt
run_finetune.sh		run_finetune.sh
run_finetune_all_baseline.sh		run_finetune_all_baseline.sh
run_prepare_data.sh		run_prepare_data.sh
run_prepare_embeddings.sh		run_prepare_embeddings.sh
run_prepare_fingerprints.sh		run_prepare_fingerprints.sh
run_prepare_hub.sh		run_prepare_hub.sh
run_prepare_pubchem.sh		run_prepare_pubchem.sh
run_train.sh		run_train.sh

Latest commit

History

Folders and files

Repository files navigation