$ show apple/ml-simplefold
SimpleFold is a research repository focused on protein structure prediction using flow-matching generative models. While it includes application-like features such as a CLI and inference scripts, its core architecture is designed for ML research, specifically training and evaluating large-scale diffusion transformers (FoldingDiT) on structural biology data. The codebase is hybrid in nature, combining rigorous data engineering pipelines (parsing MMCIF, filtering, tokenization) with standard deep learning frameworks (PyTorch Lightning, Hydra). It supports dual backends (PyTorch and MLX) for inference, indicating a focus on accessibility across hardware platforms. The repository's primary value lies in its data processing pipeline, which is heavily adapted from the Boltz-1 project. It handles complex structural data (PDB, SwissProt, AFESM) through a modular pipeline including parsing, filtering, and tokenization. The model architecture implements flow-matching with ESM-2 embeddings for sequence context. Evaluation metrics include pLDDT confidence, TM-score, and RMSD, typical for structural biology research. However, the codebase exhibits signs of research-to-production transition challenges. Hardcoded paths (e.g., `/mnt/task_runtime/...`) and reliance on external binaries (TMscore, OpenStructure) reduce portability. The data pipeline is a fork of an external project (`boltz`), introducing maintenance overhead for divergence. Despite these issues, the use of Hydra for configuration management and PyTorch Lightning for training orchestration demonstrates solid engineering practices for an ML research project.
The LLM identified 3 files worth highlighting for reuse:
A comprehensive data processing pipeline adapted from Boltz-1. It handles the full lifecycle of structural data: parsing raw MMCIF files, filtering based on static/dynamic criteria (resolution, size, date), tokenizing sequences, and writing processed data. The modular structure (parse, filter, tokenize, write) allows extraction for other structural biology projects. It demonstrates non-obvious skill in handling symmetry, ligand conformations, and alternate locations in protein structures.
The core LightningModule implementing the SimpleFold architecture. It encapsulates the training steps (flow matching, pLDDT prediction) and inference logic. The code demonstrates clean separation of concerns between model definition, training logic, and EMA (Exponential Moving Average) handling. It is reusable as a base for other flow-matching models, provided the data pipeline is adapted.
A hierarchical Hydra configuration system defining data, model, trainer, and experiment parameters. It uses `_target_` and `_partial_` for dynamic instantiation, enabling reproducible experiment management. The structure allows easy swapping of datasets, architectures, and training strategies without code changes. This is a high-value component for any ML project requiring rigorous experiment tracking and configuration management.