graph LR
Model_Base["Model Base"]
MolE_Model["MolE Model"]
Encoder_Base["Encoder Base"]
Atom_Environment_Embeddings["Atom Environment Embeddings"]
BERT_Embeddings["BERT Embeddings"]
BERT_Layer["BERT Layer"]
BERT_Encoder["BERT Encoder"]
Disentangled_Self_Attention["Disentangled Self-Attention"]
Task_Prediction_Head["Task Prediction Head"]
Metrics_Dictionary["Metrics Dictionary"]
MolE_Model -- "inherits from" --> Model_Base
Model_Base -- "uses" --> Metrics_Dictionary
MolE_Model -- "inherits from" --> Model_Base
MolE_Model -- "uses" --> Atom_Environment_Embeddings
MolE_Model -- "uses" --> Task_Prediction_Head
Atom_Environment_Embeddings -- "inherits from" --> Encoder_Base
Task_Prediction_Head -- "inherits from" --> Encoder_Base
Atom_Environment_Embeddings -- "inherits from" --> Encoder_Base
Atom_Environment_Embeddings -- "uses" --> BERT_Embeddings
Atom_Environment_Embeddings -- "uses" --> BERT_Encoder
Atom_Environment_Embeddings -- "uses" --> BERT_Embeddings
BERT_Encoder -- "uses" --> BERT_Layer
Atom_Environment_Embeddings -- "uses" --> BERT_Encoder
BERT_Encoder -- "uses" --> BERT_Layer
Task_Prediction_Head -- "inherits from" --> Encoder_Base
MolE_Model -- "uses" --> Task_Prediction_Head
Model_Base -- "uses" --> Metrics_Dictionary
MolE_Model -- "uses" --> Metrics_Dictionary
The Deep Learning Models subsystem is the core of the mole project, encapsulating the fundamental neural network architectures and components essential for molecular property prediction. It adheres to a modular design, separating concerns into abstract base models, specific model implementations, and reusable neural network building blocks. This structure facilitates extensibility, maintainability, and the integration of various deep learning techniques.
This is the foundational abstract class (mole.training.models.base.Model) for all trainable deep learning models. It establishes a standardized interface and common lifecycle methods (e.g., training_step, validation_step, test_step), ensuring consistency across different model implementations. It also integrates with the Metrics Dictionary for performance tracking.
Related Classes/Methods:
The primary deep learning model (mole.training.models.mole.MolE) implemented in this project, specifically designed for molecular tasks. It extends Model Base and orchestrates the integration of various neural network components, such as Atom Environment Embeddings and Task Prediction Head, to perform end-to-end molecular property prediction. It encapsulates the full forward pass, loss calculation, and metric logging for training and validation steps.
Related Classes/Methods:
An abstract base class (mole.training.models.encoder.Encoder) providing a common interface or base functionality for various encoder components within the deep learning models. This promotes reusability and a consistent structure for modules responsible for generating representations.
Related Classes/Methods:
A specialized component (mole.training.models.mole.AtomEnvEmbeddings) responsible for generating contextualized embeddings for atomic environments within molecules. It leverages BERT Embeddings and BERT Encoder to capture rich, local chemical information, forming crucial input representations for downstream tasks.
Related Classes/Methods:
Implements the initial embedding layer (mole.training.nn.bert.BertEmbeddings) for BERT-like models. It converts input tokens (e.g., atom features or indices) into dense vector representations, which serve as the starting point for the transformer encoder.
Related Classes/Methods:
A fundamental building block (mole.training.nn.bert.BertLayer) that encapsulates a single transformer layer. It typically includes self-attention mechanisms (like BERT Attention) and feed-forward networks, processing input sequences to generate higher-level representations.
Related Classes/Methods:
The core transformer encoder block (mole.training.nn.bert.BertEncoder), similar to the architecture found in BERT models. It consists of multiple BERT Layer instances and processes the input embeddings through self-attention mechanisms and feed-forward networks to generate higher-level, context-aware representations.
Related Classes/Methods:
A specialized attention mechanism (mole.training.nn.disentangled_attention.DisentangledSelfAttention) used within the BERT Attention component of the BERT Encoder. It aims to improve representation learning by separating different aspects of attention (e.g., content-based vs. position-based), leading to more robust and interpretable models.
Related Classes/Methods:
A neural network module (mole.training.nn.bert.TaskPredictionHead) responsible for taking the learned representations from the BERT Encoder (or other encoders) and transforming them into final predictions for specific downstream molecular tasks (e.g., property regression, classification).
Related Classes/Methods:
A utility class (mole.training.utils.metrics.MetricsDict) for managing and updating various performance metrics during model training and validation. It aggregates and organizes metrics for easy access, logging, and evaluation.
Related Classes/Methods: