reaction subpackage transformed to module.

dropped support of pytorch 1.x.
added llama mlp.

Author: stsouko

chytorch/nn/reaction/encoder.py chytorch/nn/reaction.pychytorch/nn/reaction/encoder.py renamed to chytorch/nn/reaction.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 #
-# Copyright 2021-2023 Ramil Nugmanov <nougmanoff@protonmail.com>
+# Copyright 2021-2024 Ramil Nugmanov <nougmanoff@protonmail.com>
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the “Software”), to deal
@@ -24,9 +24,9 @@
 from torch import zeros_like, float as t_float
 from torch.nn import Embedding, GELU, Module
 from torchtyping import TensorType
-from ..molecule import MoleculeEncoder
-from ..transformer import EncoderLayer
-from ...utils.data import ReactionEncoderDataBatch
+from .molecule import MoleculeEncoder
+from .transformer import EncoderLayer
+from ..utils.data import ReactionEncoderDataBatch
 
 
 class ReactionEncoder(Module):

chytorch/nn/reaction/__init__.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 #
-# Copyright 2023 Ramil Nugmanov <nougmanoff@protonmail.com>
+# Copyright 2023, 2024 Ramil Nugmanov <nougmanoff@protonmail.com>
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the “Software”), to deal
@@ -25,4 +25,6 @@
 
 
 __all__ = ['EncoderLayer',
+           'MLP',
+           'LLaMAMLP',
            'GraphormerAttention']

chytorch/nn/transformer/__init__.py

@@ -21,7 +21,7 @@
 # SOFTWARE.
 #
 from torch import Tensor, nn
-from torch.nn import Dropout, GELU, LayerNorm, Module
+from torch.nn import Dropout, GELU, LayerNorm, Module, SiLU
 from typing import Tuple, Optional, Type
 from warnings import warn
 from .attention import GraphormerAttention
@@ -54,6 +54,23 @@ def forward(self, x):
         return self.linear2(self.dropout(self.activation(self.linear1(x))))
 
 
+class LLaMAMLP(Module):
+    def __init__(self, d_model, dim_feedforward, dropout=0.1, activation=SiLU, bias: bool = False):
+        super().__init__()
+        self.linear1 = Linear(d_model, dim_feedforward, bias=bias)
+        self.linear2 = Linear(d_model, dim_feedforward, bias=bias)
+        self.linear3 = Linear(dim_feedforward, d_model, bias=bias)
+        self.dropout = Dropout(dropout)
+
+        # ad-hoc for resolving class from name
+        if isinstance(activation, str):
+            activation = getattr(nn, activation)
+        self.activation = activation()
+
+    def forward(self, x):
+        return self.linear3(self.dropout(self.activation(self.linear1(x))) * self.linear2(x))
+
+
 class EncoderLayer(Module):
     r"""EncoderLayer based on torch.nn.TransformerEncoderLayer, but batch always first and returns also attention.
 
@@ -96,4 +113,4 @@ def forward(self, x: Tensor, attn_mask: Optional[Tensor], *,
         return None, a
 
 
-__all__ = ['EncoderLayer', 'MLP']
+__all__ = ['EncoderLayer', 'MLP', 'LLaMAMLP']

chytorch/nn/transformer/encoder.py

@@ -30,7 +30,7 @@
 from ._utils import *
 
 
-__all__ = ['MoleculeDataset', 'collate_molecules',
+__all__ = ['MoleculeDataset', 'collate_molecules', 'left_padded_collate_molecules',
            'ConformerDataset', 'collate_conformers',
            'ReactionEncoderDataset', 'collate_encoded_reactions',
            'RDKitConformerDataset',

chytorch/utils/data/__init__.py

@@ -26,7 +26,8 @@
 from .rdkit import *
 
 
-__all__ = ['MoleculeDataset', 'MoleculeDataPoint', 'MoleculeDataBatch', 'collate_molecules',
+__all__ = ['MoleculeDataset', 'MoleculeDataPoint', 'MoleculeDataBatch',
+           'collate_molecules', 'left_padded_collate_molecules',
            'ConformerDataset', 'ConformerDataPoint', 'ConformerDataBatch', 'collate_conformers',
            'RDKitConformerDataset',
            'thiacalix_n_arene_dataset']

chytorch/utils/data/_abc.py

@@ -79,14 +79,12 @@ def unpack(const unsigned char[::1] data not None, unsigned short add_cls, unsig
 
     cdef cnp.ndarray[DTYPE_t, ndim=1] atoms, neighbors
     cdef cnp.ndarray[DTYPE_t, ndim=2] distance
-    cdef DTYPE_t d, attention
+    cdef DTYPE_t d
 
     # read header
     if data[0] != 2:
         raise ValueError('invalid pack version')
 
-    attention = 1 if components_attention else 0
-
     a, b, c = data[1], data[2], data[3]
     atoms_count = (a << 4| b >> 4) + add_cls
     cis_trans_count = (b & 0x0f) << 8 | c
@@ -170,7 +168,7 @@ def unpack(const unsigned char[::1] data not None, unsigned short add_cls, unsig
             d = distance[i, j]
             if d == 9999:
                 # set attention between subgraphs
-                distance[i, j] = distance[j, i] = attention
+                distance[i, j] = distance[j, i] = components_attention
             elif d > max_distance:
                 distance[i, j] = distance[j, i] = max_distance + 2
             else:

chytorch/utils/data/molecule/__init__.py

@@ -27,9 +27,9 @@
 from torch import IntTensor, Size, zeros, ones as t_ones, int32 as t_int32, eye
 from torch.nn.utils.rnn import pad_sequence
 from torch.utils.data import Dataset
+from torch.utils.data._utils.collate import default_collate_fn_map
 from torchtyping import TensorType
 from typing import Sequence, Tuple, Union, NamedTuple
-from .._abc import default_collate_fn_map
 
 
 class ConformerDataPoint(NamedTuple):
@@ -53,7 +53,7 @@ def cuda(self, *args, **kwargs):
         return ConformerDataBatch(*(x.cuda(*args, **kwargs) for x in self))
 
 
-def collate_conformers(batch, *, padding_left: bool = False, collate_fn_map=None) -> ConformerDataBatch:
+def collate_conformers(batch, *, collate_fn_map=None) -> ConformerDataBatch:
     """
     Prepares batches of conformers.
 
@@ -62,25 +62,16 @@ def collate_conformers(batch, *, padding_left: bool = False, collate_fn_map=None
     atoms, hydrogens, distances = [], [], []
 
     for a, h, d in batch:
-        if padding_left:
-            atoms.append(a.flipud())
-            hydrogens.append(h.flipud())
-        else:
-            atoms.append(a)
-            hydrogens.append(h)
+        atoms.append(a)
+        hydrogens.append(h)
         distances.append(d)
 
     pa = pad_sequence(atoms, True)
     b, s = pa.shape
     tmp = eye(s, dtype=t_int32).repeat(b, 1, 1)  # prevent nan in MHA softmax on padding
     for i, d in enumerate(distances):
         s = d.size(0)
-        if padding_left:
-            tmp[i, -s:, -s:] = d
-        else:
-            tmp[i, :s, :s] = d
-    if padding_left:
-        return ConformerDataBatch(pa.fliplr(), pad_sequence(hydrogens, True).fliplr(), tmp)
+        tmp[i, :s, :s] = d
     return ConformerDataBatch(pa, pad_sequence(hydrogens, True), tmp)
 
 
@@ -144,10 +135,9 @@ def __getitem__(self, item: int) -> ConformerDataPoint:
                 atoms = IntTensor(len(mol))
                 hydrogens = IntTensor(len(mol))
 
-            hgs = mol._hydrogens  # noqa
             for i, (n, a) in enumerate(mol.atoms(), self.add_cls):
                 atoms[i] = a.atomic_number + 2
-                hydrogens[i] = (hgs[n] or 0) + 2
+                hydrogens[i] = (a.implicit_hydrogens or 0) + 2
 
             xyz = empty((len(mol), 3))
             conformer = mol._conformers[0]  # noqa