hjxdl 0.0.1__py3-none-any.whl

hdl/models/fast_transformer.py

@@ -0,0 +1,234 @@
+import torch
+# import torch.nn.functional as F
+from torch import nn, einsum
+
+from einops import rearrange, reduce
+from rotary_embedding_torch import apply_rotary_emb, RotaryEmbedding
+
+# helper functions
+
+
+def exists(val):
+    return val is not None
+
+
+def default(val, d):
+    return val if exists(val) else d
+
+
+# helper classes
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.fn = fn
+
+    def forward(self, x, **kwargs):
+        x = self.norm(x)
+        return self.fn(x, **kwargs)
+
+# blocks
+
+
+def FeedForward(dim, mult=4):
+    return nn.Sequential(
+        nn.Linear(dim, dim * mult),
+        nn.GELU(),
+        nn.Linear(dim * mult, dim)
+    )
+
+
+class FastAttention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        *,
+        heads=8,
+        dim_head=64,
+        max_seq_len=None,
+        pos_emb=None
+    ):
+        super().__init__()
+        inner_dim = heads * dim_head
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
+
+        # rotary positional embedding
+
+        assert not (exists(pos_emb) and not exists(max_seq_len)), \
+            'max_seq_len must be passed in if to use rotary positional embeddings'
+
+        self.pos_emb = pos_emb
+        self.max_seq_len = max_seq_len
+
+        # if using relative positional encoding, make sure to reduce pairs of 
+        # consecutive feature dimension before doing projection to attention logits
+
+        kv_attn_proj_divisor = 1 if not exists(pos_emb) else 2
+
+        # for projecting queries to query attention logits
+        self.to_q_attn_logits = nn.Linear(dim_head, 1, bias=False) 
+        self.to_k_attn_logits = nn.Linear(
+            dim_head // kv_attn_proj_divisor,
+            1,
+            bias=False
+        )  # for projecting keys to key attention logits
+
+        # final transformation of values to "r" as in the paper
+
+        self.to_r = nn.Linear(dim_head // kv_attn_proj_divisor, dim_head)
+
+        self.to_out = nn.Linear(inner_dim, dim)
+
+    def forward(self, x, mask=None):
+        n, device, h, use_rotary_emb = x.shape[1], x.device, self.heads, exists(self.pos_emb)
+
+        qkv = self.to_qkv(x).chunk(3, dim=-1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), qkv)
+
+        mask_value = -torch.finfo(x.dtype).max
+        mask = rearrange(mask, 'b n -> b () n')
+
+        # if relative positional encoding is needed
+
+        if use_rotary_emb:
+            freqs = self.pos_emb(torch.arange(self.max_seq_len, device=device), cache_key=self.max_seq_len)
+            freqs = rearrange(freqs[:n], 'n d -> () () n d')
+            q_aggr, k_aggr, v_aggr = map(lambda t: apply_rotary_emb(freqs, t), (q, k, v))
+        else:
+            q_aggr, k_aggr, v_aggr = q, k, v
+
+        # calculate query attention logits
+
+        q_attn_logits = rearrange(self.to_q_attn_logits(q), 'b h n () -> b h n') * self.scale
+        q_attn_logits = q_attn_logits.masked_fill(~mask, mask_value)
+        q_attn = q_attn_logits.softmax(dim=-1)
+
+        # calculate global query token
+
+        global_q = einsum('b h n, b h n d -> b h d', q_attn, q_aggr)
+        global_q = rearrange(global_q, 'b h d -> b h () d')
+
+        # bias keys with global query token
+
+        k = k * global_q
+
+        # if using rotary embeddings, do an inner product between adjacent pairs in the feature dimension
+
+        if use_rotary_emb:
+            k = reduce(k, 'b h n (d r) -> b h n d', 'sum', r=2)
+
+        # now calculate key attention logits
+
+        k_attn_logits = rearrange(self.to_k_attn_logits(k), 'b h n () -> b h n') * self.scale
+        k_attn_logits = k_attn_logits.masked_fill(~mask, mask_value)
+        k_attn = k_attn_logits.softmax(dim=-1)
+
+        # calculate global key token
+
+        global_k = einsum('b h n, b h n d -> b h d', k_attn, k_aggr)
+        global_k = rearrange(global_k, 'b h d -> b h () d')
+
+        # bias the values
+
+        u = v_aggr * global_k
+
+        # if using rotary embeddings, do an inner product between adjacent pairs in the feature dimension
+
+        if use_rotary_emb:
+            u = reduce(u, 'b h n (d r) -> b h n d', 'sum', r=2)
+
+        # transformation step
+
+        r = self.to_r(u)
+
+        # paper then says to add the queries as a residual
+
+        r = r + q
+
+        # combine heads
+
+        r = rearrange(r, 'b h n d -> b n (h d)')
+        return self.to_out(r)
+
+
+# main class
+class FastTransformer(nn.Module):
+    def __init__(
+        self,
+        *,
+        num_tokens,
+        dim,
+        depth,
+        max_seq_len,
+        heads=8,
+        dim_head=64,
+        ff_mult=4,
+        absolute_pos_emb=False
+    ):
+        super().__init__()
+        self.token_emb = nn.Embedding(num_tokens, dim)
+
+        # positional embeddings
+
+        self.abs_pos_emb = nn.Embedding(max_seq_len, dim) if absolute_pos_emb else None
+
+        layer_pos_emb = None
+        if not absolute_pos_emb:
+            assert (dim_head % 4) == 0, 'dimension of the head must be divisible by 4 to use rotary embeddings'
+            layer_pos_emb = RotaryEmbedding(dim_head // 2)
+
+        # layers
+
+        self.layers = nn.ModuleList([])
+
+        for _ in range(depth):
+            attn = FastAttention(
+                dim,
+                dim_head=dim_head,
+                heads=heads,
+                pos_emb=layer_pos_emb,
+                max_seq_len=max_seq_len
+            )
+            ff = FeedForward(dim, mult=ff_mult)
+
+            self.layers.append(nn.ModuleList([
+                PreNorm(dim, attn),
+                PreNorm(dim, ff)
+            ]))
+
+        # weight tie projections across all layers
+
+        first_block, _ = self.layers[0]
+        for block, _ in self.layers[1:]:
+            block.fn.to_q_attn_logits = first_block.fn.to_q_attn_logits
+            block.fn.to_k_attn_logits = first_block.fn.to_k_attn_logits
+
+        # to logits
+
+        self.to_logits = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, num_tokens)
+        )
+
+    def forward(
+        self,
+        x,
+        mask=None
+    ):
+        n, device = x.shape[1], x.device
+        if mask is None:
+            mask = torch.ones_like(x).bool().to(device)
+        x = self.token_emb(x)
+ 
+        if exists(self.abs_pos_emb):
+            pos_emb = self.abs_pos_emb(torch.arange(n, device=device))
+            x = x + rearrange(pos_emb, 'n d -> () n d')
+
+        for attn, ff in self.layers:
+            x = attn(x, mask=mask) + x
+            x = ff(x) + x
+
+        return self.to_logits(x)

hdl/models/ginet.py

@@ -0,0 +1,189 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+# from torch_geometric.nn import MessagePassing
+# from torch_geometric.utils import add_self_loops
+from torch_geometric.nn import global_add_pool, global_mean_pool, global_max_pool
+
+from hdl.layers.graph.gin import GINEConv
+from hdl.layers.general.linear import (
+    # BNReLULinear,
+    BNReLULinearBlock,
+)
+from hdl.models.utils import load_model
+from hdl.ops.utils import get_activation
+
+
+__all__ = [
+    "GINet",
+    "GINMLPR",
+]
+
+
+num_atom_type = 119  # including the extra mask tokens
+num_chirality_tag = 3
+
+num_bond_type = 5  # including aromatic and self-loop edge
+num_bond_direction = 3 
+
+
+class GINet(nn.Module):
+    """
+    Args:
+        num_layer (int): the number of GNN layers
+        emb_dim (int): dimensionality of embeddings
+        max_pool_layer (int): the layer from which we use max pool rather than add pool for neighbor aggregation
+        drop_ratio (float): dropout rate
+        gnn_type: gin, gcn, graphsage, gat
+    Output:
+        node representations
+    """
+    def __init__(
+        self,
+        num_layer=5,
+        emb_dim=300,
+        feat_dim=512,
+        drop_ratio=0,
+        pool='mean'
+    ):
+        super(GINet, self).__init__()
+        self.init_args = {
+            'num_layer': num_layer,
+            'emb_dim': emb_dim,
+            'feat_dim': feat_dim,
+            'drop_ratio': drop_ratio,
+            'pool': pool
+        }
+        self.num_layer = num_layer
+        self.emb_dim = emb_dim
+        self.feat_dim = feat_dim
+        self.drop_ratio = drop_ratio
+
+        self.x_embedding1 = nn.Embedding(num_atom_type, emb_dim)
+        self.x_embedding2 = nn.Embedding(num_chirality_tag, emb_dim)
+        nn.init.xavier_uniform_(self.x_embedding1.weight.data)
+        nn.init.xavier_uniform_(self.x_embedding2.weight.data)
+
+        # List of MLPs
+        self.gnns = nn.ModuleList()
+        for layer in range(num_layer):
+            self.gnns.append(GINEConv(emb_dim))
+
+        # List of batchnorms
+        self.batch_norms = nn.ModuleList()
+        for layer in range(num_layer):
+            self.batch_norms.append(nn.BatchNorm1d(emb_dim))
+        
+        if pool == 'mean':
+            self.pool = global_mean_pool
+        elif pool == 'max':
+            self.pool = global_max_pool
+        elif pool == 'add':
+            self.pool = global_add_pool
+        
+        self.feat_lin = nn.Linear(
+            self.emb_dim,
+            self.feat_dim
+        )
+
+        self.out_lin = nn.Sequential(
+            nn.Linear(self.feat_dim, self.feat_dim), 
+            nn.ReLU(inplace=True),
+            nn.Linear(
+                self.feat_dim,
+                self.feat_dim // 2
+            )
+        )
+
+    def forward(self, data):
+        x = data.x
+        edge_index = data.edge_index
+        edge_attr = data.edge_attr
+
+        h = self.x_embedding1(x[:,0]) + self.x_embedding2(x[:,1])
+
+        for layer in range(self.num_layer):
+            h = self.gnns[layer](h, edge_index, edge_attr)
+            h = self.batch_norms[layer](h)
+            if layer == self.num_layer - 1:
+                h = F.dropout(h, self.drop_ratio, training=self.training)
+            else:
+                h = F.dropout(F.relu(h), self.drop_ratio, training=self.training)
+
+        h = self.pool(h, data.batch)
+        h = self.feat_lin(h)
+        out = self.out_lin(h)
+        
+        return h, out
+
+
+class GINMLPR(nn.Module):
+    def __init__(
+        self,
+        num_layer=5,
+        emb_dim=300,
+        feat_dim=512,
+        out_dim=1,
+        drop_ratio=0,
+        pool='mean',
+        ckpt_file: str = None,
+        num_smiles: int = 1,
+    ) -> None:
+        super().__init__()
+        self.init_args = {
+            "num_layer": num_layer,
+            "emb_dim": emb_dim,
+            "feat_dim": feat_dim,
+            "out_dim": out_dim,
+            "drop_ratio": drop_ratio,
+            "pool": pool,
+            "ckpt_file": ckpt_file,
+            "num_smiles": num_smiles
+        }
+        self.gins = nn.ModuleList([])
+        for _ in range(num_smiles):
+            self.gins.append(
+                GINet(
+                    num_layer=num_layer,
+                    emb_dim=emb_dim,
+                    feat_dim=feat_dim,
+                    drop_ratio=drop_ratio,
+                    pool=pool,
+                )
+            )
+        self.ckpt_file = ckpt_file
+        self.num_smiles = num_smiles
+
+        self.ffn = BNReLULinearBlock(
+            in_features=feat_dim // 2 * num_smiles,
+            out_features=out_dim,
+            num_layers=num_layer,
+            hidden_size=feat_dim // 2
+        )
+        self.out_act = get_activation('sigmoid')
+
+        if ckpt_file is not None:
+            self.load_ckpt()
+
+    def load_ckpt(self):
+        if self.ckpt_file is not None:
+            for i in range(self.num_smiles):
+                load_model(
+                    self.ckpt_file,
+                    model=self.gins[i]
+                )
+ 
+    def forward(
+        self,
+        data 
+    ):
+        out_list = []
+        for data_i, gin in zip(data, self.gins):
+            out_list.append(gin(data_i[0])[1])
+        out = torch.hstack(out_list)  # (batch_size, feat_dim//2 * num_smiles)
+        out = self.ffn(out)
+        out = self.out_act(out)
+
+        return out
+            

hdl/models/linear.py

@@ -0,0 +1,137 @@
+import typing as t
+
+import torch
+from torch import nn
+from torch.nn import functional as nnfunc
+import numpy as np
+
+from hdl.layers.general.linear import (
+    BNReLULinearBlock,
+    BNReLULinear
+)
+# from hdl.ops.utils import get_activation
+
+
+class MMIterLinear(nn.Module):
+    _NAME = 'mumc_linear'
+
+    def __init__(
+        self,
+        num_fp_bits: int,
+        num_in_feats: int,
+        nums_classes: t.List[int] = [3, 3],
+        target_names: t.List[str] = None,
+        hidden_size: int = 128,
+        num_hidden_layers: int = 10,
+        activation: str = 'elu',
+        out_act: str = 'softmax',
+        hard_select: bool = False,
+        iterative: bool = True,
+        **kwargs,
+    ):
+        super().__init__()
+
+        if target_names is None:
+            self.target_names = list(range(len(nums_classes)))
+        else:
+            self.target_names = target_names
+        
+        self.init_args = {
+            'num_fp_bits': num_fp_bits,
+            'num_in_feats': num_in_feats,
+            'nums_classes': nums_classes,
+            'target_names': target_names,
+            'hidden_size': hidden_size,
+            'num_hidden_layers': num_hidden_layers,
+            'activation': activation,
+            'out_act': out_act,
+            'hard_select': hard_select,
+            'iterative': iterative,
+            **kwargs
+        }
+        self.hard_select = hard_select
+        self.iterative = iterative
+        self._freeze_classifier = [True] * len(target_names)
+        
+        # self.w1 = BNReLULinear(num_fp_bits, num_in_feats, activation)
+        self.w1 = nn.Linear(num_fp_bits, num_in_feats)
+        # self.w2 = BNReLULinear(num_fp_bits, num_in_feats, activation)
+        self.w2 = nn.Linear(num_fp_bits, num_in_feats)
+        # self.w3 = BNReLULinear(num_fp_bits, num_in_feats, activation)
+        self.w3 = nn.Linear(num_fp_bits, num_in_feats)
+
+        nums_in_feats = [num_in_feats]
+        if iterative:
+            nums_in_feats.extend(nums_classes)
+            nums_in_feats = np.cumsum(np.array(nums_in_feats, dtype=np.int))[:-1]
+        else:
+            nums_in_feats = nums_in_feats * len(nums_classes)
+        
+        if isinstance(out_act, str):
+            self.out_acts = [out_act] * len(nums_classes)
+        else:
+            self.out_acts = out_act
+
+        self.classifiers = nn.ModuleList([
+            nn.Sequential(
+                BNReLULinearBlock(
+                    num_in,
+                    hidden_size,
+                    num_hidden_layers,
+                    hidden_size,
+                    activation,
+                    **kwargs
+                ),
+                BNReLULinear(
+                    hidden_size,
+                    num_out,
+                    out_act,
+                    **kwargs
+                )
+            )
+            for num_in, num_out, out_act in zip(
+                nums_in_feats, nums_classes, self.out_acts
+            )
+        ])
+
+    @property
+    def freeze_classifier(self):
+        return self._freeze_classifier
+
+    @freeze_classifier.setter
+    def freeze_classifier(self, freeze: t.List = []):
+        self._freeze_classifier = freeze
+        self.change_classifier_grad([not f for f in freeze])
+
+    def change_classifier_grad(self, requires_grads: t.List = []):
+        for requires_grad, classifier in zip(requires_grads, self.classifiers):
+            for param in classifier.parameters():
+                param.requires_grad = requires_grad
+    
+    def forward(self, fps, target_tensors=None, teach=True):
+        result_dict = {}
+        fp1, fp2, fp3 = fps
+        fp1 = self.w1(fp1)
+        fp2 = self.w2(fp2)
+        fp3 = self.w3(fp3)
+        X = fp3 - (fp1 + fp2)
+        if target_tensors is None:
+            target_tensors = [None] * len(self.target_names)
+        for classifier, target_name, target_tensor in zip(
+            self.classifiers, self.target_names, target_tensors
+        ):
+            result = classifier(X)
+            result_dict[target_name] = result
+            if self.iterative:
+                if teach:
+                    assert target_tensors is not None
+                    X = torch.cat((X, target_tensor), -1) 
+                else: 
+                    if not self.hard_select:
+                        X = torch.cat((X, result), -1)
+                    else:
+                        X = torch.cat(
+                            (X, nnfunc.gumbel_softmax(result, tau=1, hard=True)),
+                            -1
+                        )
+        return result_dict

hdl/models/model_dict.py

@@ -0,0 +1,18 @@
+from hdl.layers.general.linear import (
+    MultiTaskMultiClassBlock,
+    MuMcHardBlock
+)
+from .linear import MMIterLinear
+from .chiral_gnn import GNN
+from .ginet import GINet
+from .ginet import GINMLPR
+
+
+MODEL_DICT = {
+    'rxn_trans': MultiTaskMultiClassBlock,
+    'rxn_trans_hard': MuMcHardBlock,
+    'mmiter_linear': MMIterLinear,
+    'chiral_gnn': GNN,
+    'ginet': GINet,
+    'ginmlpr': GINMLPR
+}

hdl/models/norm_flows.py

@@ -0,0 +1,33 @@
+import torch
+import torch.nn as nn
+
+
+class NormalizingFlowModel(nn.Module):
+
+    def __init__(self, prior, flows):
+        super().__init__()
+        self.prior = prior
+        self.flows = nn.ModuleList(flows)
+
+    def forward(self, x):
+        m, _ = x.shape
+        log_det = torch.zeros(m)
+        for flow in self.flows:
+            x, ld = flow.forward(x)
+            log_det += ld
+        z, prior_logprob = x, self.prior.log_prob(x)
+        return z, prior_logprob, log_det
+
+    def inverse(self, z):
+        m, _ = z.shape
+        log_det = torch.zeros(m)
+        for flow in self.flows[::-1]:
+            z, ld = flow.inverse(z)
+            log_det += ld
+        x = z
+        return x, log_det
+
+    def sample(self, n_samples):
+        z = self.prior.sample((n_samples,))
+        x, _ = self.inverse(z)
+        return x

hdl/models/optim_dict.py

@@ -0,0 +1,16 @@
+from torch.optim import (
+    Adadelta,
+    Adam,
+    SGD,
+    RMSprop,
+)
+from hdl.optims.nadam import Nadam
+
+
+OPTIM_DICT = {
+    'adam': Adam,
+    'adadelta': Adadelta,
+    'sgd': SGD,
+    'rmsprop': RMSprop,
+    'nadam': Nadam,
+}

hdl/models/rxn.py

@@ -0,0 +1,63 @@
+import pkg_resources
+from transformers import BertModel
+import torch
+# from torch import nn
+
+from hdl.layers.general.linear import (
+    MultiTaskMultiClassBlock,
+    MuMcHardBlock
+)
+# from hdl.data.seq.rxn import rxn_model
+
+
+def get_rxn_model(
+    model_path: str = None
+):
+    if model_path is None:
+        model_path = pkg_resources.resource_filename(
+            "rxnfp",
+            "models/transformers/bert_ft"
+        )
+        model = BertModel.from_pretrained(model_path)
+        model = model.eval().cpu()
+
+    return model
+
+
+def build_rxn_mu(
+    nums_classes,
+    hard=False,
+    hidden_size=128,
+    nums_hidden_layers=10,
+    encoder=get_rxn_model(),
+    # freeze_encoder=True,
+    device_id: int = 0,
+    **kwargs
+):
+    if not hard:
+        model = MultiTaskMultiClassBlock(
+            encoder=encoder,
+            nums_classes=nums_classes,
+            hidden_size=hidden_size,
+            num_hidden_layers=nums_hidden_layers,
+            # freeze_encoder=freeze_encoder,
+            **kwargs
+        )
+    else:
+        model = MuMcHardBlock(
+            encoder=encoder,
+            nums_classes=nums_classes,
+            hidden_size=hidden_size,
+            num_hidden_layers=nums_hidden_layers,
+            # freeze_encoder=freeze_encoder,
+            **kwargs
+        )
+    device = torch.device(f'cuda:{device_id}') \
+        if torch.cuda.is_available() \
+        else torch.device('cpu')
+
+    model = model.to(device)
+    
+    # if torch.cuda.device_count() > 1:
+    #     model = nn.DataParallel(model)
+    return model, device