.. _models:

Models
======

`unimol_tools.models <https://github.com/deepmodeling/Uni-Mol/tree/main/unimol_tools/unimol_tools/models>`_ contains the models of Uni-Mol.


Uni-Mol
-------

`unimol_tools.models.unimol.py <https://github.com/deepmodeling/Uni-Mol/tree/main/unimol_tools/unimol_tools/models/unimol.py>`_ contains the :class:`~unimol_tools.models.UniMolModel`, which is the backbone of Uni-Mol model.

.. automodule:: unimol_tools.models.unimol
   :members:

Model
-----

`unimol_tools.models.nnmodel.py <https://github.com/deepmodeling/Uni-Mol/tree/main/unimol_tools/unimol_tools/models/unimol.py>`_ contains the :class:`~unimol_tools.models.NNModel`, which is responsible for initializing the model.

.. automodule:: unimol_tools.models.nnmodel
   :members:

Loss
-----

`unimol_tools.models.loss.py <https://github.com/deepmodeling/Uni-Mol/tree/main/unimol_tools/unimol_tools/models/loss.py>`_ contains different loss functions.

.. automodule:: unimol_tools.models.loss
   :members:

Transformers
------------

`unimol_tools.models.transformers.py <https://github.com/deepmodeling/Uni-Mol/tree/main/unimol_tools/unimol_tools/models/transformers.py>`_ contains a custom Transformer Encoder module that extends `PyTorch's nn.Module <https://pytorch.org/docs/stable/generated/torch.nn.Module.html>`_.

.. automodule:: unimol_tools.models.transformers
   :members: