Dropped useless datasets, NNs and utils.

Author: stsouko

README.md

@@ -21,15 +21,15 @@ Installed models can be imported as `from chytorch.zoo.<model_name> import Model
 Usage
 -----
 
-`chytorch.nn.MoleculeEncoder` and `chytorch.nn.ReactionEncoder` - core graphormer layers for molecules and reactions.
+`chytorch.nn.MoleculeEncoder` - core graphormer layer for molecules encoding.
 API is combination of `torch.nn.TransformerEncoderLayer` with `torch.nn.TransformerEncoder`. 
 
 **Batch preparation:**
 
-`chytorch.utils.data.MoleculeDataset` and `chytorch.utils.data.ReactionDataset` - Map-like on-the-fly dataset generators for molecules and reactions.
-Supported `chython.MoleculeContainer` and `chython.ReactionContainer` objects, and bytes-packed structures.
+`chytorch.utils.data.MoleculeDataset` - Map-like on-the-fly dataset generators for molecules.
+Supported `chython.MoleculeContainer` objects, and PaCh structures.
 
-`chytorch.utils.data.collate_molecules` and `chytorch.utils.data.collate_reactions` - collate functions for `torch.utils.data.DataLoader`.
+`chytorch.utils.data.collate_molecules` - collate function for `torch.utils.data.DataLoader`.
 
 Note: torch DataLoader automatically do proper collation since 1.13 release.
 
@@ -52,10 +52,6 @@ Molecules coded as tensors of:
 * topological distances' matrix shifted by 2 with upper limit.
   0 - reserved for padding, 1 - reserved for not-connected graph components coding, 2 - self-loop, 3 - connected atoms.
 
-Reactions coded in similar way. Molecules atoms and neighbors matrices just stacked. Distance matrices stacked on diagonal.
-Reactions include additional tensor with reaction role codes for each token.
-0 - padding, 1 - reaction CLS, 2 - reactants, 3 - products.
-
     from chytorch.nn import MoleculeEncoder
 
     encoder = MoleculeEncoder()
@@ -73,11 +69,6 @@ Reactions include additional tensor with reaction role codes for each token.
     dl = DataLoader(TensorDataset(MoleculeDataset(molecules_list), properties_tensor),
         collate_fn=chained_collate(collate_molecules, stack))
 
-
-**Scheduler:**
-
-`chytorch.optim.lr_scheduler.WarmUpCosine` - Linear warmup followed with cosine-function for 0-pi range rescaled to lr_rate - decrease_coef * lr_rate interval.
-
 **Voting NN with single hidden layer:**
 
 `chytorch.nn.VotingClassifier`, `chytorch.nn.BinaryVotingClassifier` and `chytorch.nn.VotingRegressor` - speed optimized multiple heads for ensemble predictions.
@@ -94,8 +85,6 @@ All wrappers have `torch.utils.data.Dataset` interface.
 * `SizedList` - list wrapper with `size()` method. Useful with `torch.utils.data.TensorDataset`. 
 * `SMILESDataset` - on-the-fly smiles to `chython.MoleculeContainer` or `chython.ReactionContainer` parser.
 * `LMDBMapper` - LMDB KV storage to dataset mapper.
-* `PostgresMapper` - Postgres DB table to dataset mapper.
-* `SMILESTokenizerDataset` - on-the-fly generator of tokenized SMILES.
 * `TensorUnpack`, `StructUnpack`, `PickleUnpack` - bytes to tensor/object unpackers
 
 

chytorch/nn/__init__.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
@@ -20,8 +20,6 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 #
-from .activation import *
-from .converters import *
 from .losses import *
 from .molecule import *
 from .reaction import *
@@ -33,9 +31,7 @@
            'ReactionEncoder',
            'Slicer',
            'VotingClassifier', 'VotingRegressor', 'BinaryVotingClassifier',
-           'PulingHardtanh',
            'MultiTaskLoss',
            'CensoredLoss',
            'MaskedNaNLoss',
-           'MSLELoss',
-           'Converters', 'MultiColumnConverters']
+           'MSLELoss']

chytorch/nn/converters.py

@@ -20,9 +20,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 #
-from .embedding import *
 from .encoder import *
 
 
-__all__ = ['MoleculeEncoder',
-           'EmbeddingBag']
+__all__ = ['MoleculeEncoder']

chytorch/nn/functional/__init__.py

@@ -20,10 +20,12 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 #
+from itertools import repeat
 from torch.nn import GELU, Module, ModuleList, LayerNorm
 from torchtyping import TensorType
+from typing import Tuple, Optional, List
 from warnings import warn
-from .embedding import EmbeddingBag
+from ._embedding import EmbeddingBag
 from ..lora import Embedding
 from ..transformer import EncoderLayer
 from ...utils.data import MoleculeDataBatch
@@ -49,6 +51,7 @@ def __init__(self, max_neighbors: int = 14, max_distance: int = 10, d_model: int
                  shared_attention_bias: bool = True, dropout: float = 0.1, activation=GELU,
                  layer_norm_eps: float = 1e-5, norm_first: bool = False, post_norm: bool = False,
                  zero_bias: bool = False, perturbation: float = 0., max_tokens: int = 121,
+                 projection_bias: bool = True, ff_bias: bool = True,
                  lora_r: int = 0, lora_alpha: float = 1., lora_dropout: float = 0.):
         """
         Molecule Graphormer from https://doi.org/10.1021/acs.jcim.2c00344.
@@ -103,16 +106,21 @@ def __init__(self, max_neighbors: int = 14, max_distance: int = 10, d_model: int
         self.shared_weights = shared_weights
         if shared_weights:
             self.layer = EncoderLayer(d_model, nhead, dim_feedforward, dropout, activation, layer_norm_eps, norm_first,
+                                      projection_bias=projection_bias, ff_bias=ff_bias,
                                       lora_r=lora_r, lora_alpha=lora_alpha, lora_dropout=lora_dropout)
             self.layers = [self.layer] * num_layers
         else:
             # layers sharing scheme can be manually changed. e.g. pairs of shared encoders
             self.layers = ModuleList(EncoderLayer(d_model, nhead, dim_feedforward, dropout, activation,
                                                   layer_norm_eps, norm_first, lora_r=lora_r, lora_alpha=lora_alpha,
+                                                  projection_bias=projection_bias, ff_bias=ff_bias,
                                                   lora_dropout=lora_dropout) for _ in range(num_layers))
         self._register_load_state_dict_pre_hook(_update)
 
-    def forward(self, batch: MoleculeDataBatch) -> TensorType['batch', 'atoms', 'embedding']:
+    def forward(self, batch: MoleculeDataBatch, /, *,
+                cache: Optional[List[Tuple[TensorType['batch', 'atoms+conditions', 'embedding'],
+                                           TensorType['batch', 'atoms+conditions', 'embedding']]]] = None) -> \
+            TensorType['batch', 'atoms', 'embedding']:
         """
         Use 0 for padding.
         Atoms should be coded by atomic numbers + 2.
@@ -122,15 +130,16 @@ def forward(self, batch: MoleculeDataBatch) -> TensorType['batch', 'atoms', 'emb
         Distances should be coded from 2 (means self-loop) to max_distance + 2.
         Non-reachable atoms should be coded by 1.
         """
+        cache = repeat(None) if cache is None else iter(cache)
         atoms, neighbors, distances = batch
 
         x = self.embedding(atoms, neighbors)
 
-        for lr, d in zip(self.layers, self.distance_encoders):
+        for lr, d, c in zip(self.layers, self.distance_encoders, cache):
             if d is not None:
                 d_mask = d(distances).permute(0, 3, 1, 2)  # BxNxNxH > BxHxNxN
             # else: reuse previously calculated mask
-            x, _ = lr(x, d_mask)  # noqa
+            x, _ = lr(x, d_mask, cache=c)  # noqa
 
         if self.post_norm:
             return self.norm(x)