agedi.models.head
=================

.. py:module:: agedi.models.head


Classes
-------

.. autoapisummary::

   agedi.models.head.Head


Module Contents
---------------

.. py:class:: Head(score_clip: Optional[float] = None, **kwargs)

   Bases: :py:obj:`abc.ABC`, :py:obj:`torch.nn.Module`


   Abstract base class for any score model heads.

   The head is responsible for taking the translated batch with precalculated
   representation and returning a score tensor.

   The score tensor should have the same shape as the original tensor for the
   key of the head.

   :rtype: Head


   .. py:attribute:: _key
      :type:  str


   .. py:attribute:: _score_clip
      :value: None



   .. py:property:: key
      :type: str


      The key of the attribute to be noised and denoised.


   .. py:method:: get_hparams() -> Dict

      Return hyperparameters sufficient to reconstruct this head.

      Returns a dictionary with a ``_target_`` key (the fully-qualified class
      name) plus ``score_clip`` from the base class.  Subclasses should call
      ``super().get_hparams()`` and merge in their own constructor parameters.

      :returns: Hyperparameter dictionary.
      :rtype: dict



   .. py:method:: forward(translated_batch: Any) -> torch.Tensor

      Forward pass of the head using a translated batch

      The output shape must match the either the positions (pos), types (x) or
      cell (cell) of the original batch.

      :param translated_batch: The translated batch to be used in the forward pass
      :type translated_batch: Any

      :returns: The output of the forward pass. The shape of the tensor depends on the key of the head.
      :rtype: torch.Tensor



   .. py:method:: _score(translated_batch: Any) -> torch.Tensor
      :abstractmethod:


      Abstract method for the forward pass of the head.

      Must be implemented by the subclass.

      :param translated_batch: The translated batch to be used in the forward pass
      :type translated_batch: Any

      :returns: The output of the forward pass. The shape of the tensor depends on the key of the head.
      :rtype: torch.Tensor