Models

unimol_tools.models contains the models of Uni-Mol.

Uni-Mol

unimol_tools.models.unimol.py contains the UniMolModel, which is the backbone of Uni-Mol model.

Model

unimol_tools.models.nnmodel.py contains the NNModel, which is responsible for initializing the model.

Loss

unimol_tools.models.loss.py contains different loss functions.

Transformers

unimol_tools.models.transformers.py contains a custom Transformer Encoder module that extends PyTorch’s nn.Module.