Update models.py

models.py

@@ -169,7 +169,11 @@ class Generator2(nn.Module):
         #pos_enc = self.pos_enc(lap)
         #drug_n = drug_n + pos_enc
                 
-        nodes_logits,akt1_annot, edges_logits, akt1_adj = self.TransformerDecoder(nodes_logits,akt1_annot,edges_logits,akt1_adj)
+        if self.submodel == "Ligand" or self.submodel == "RL" :
+            nodes_logits,akt1_annot, edges_logits, akt1_adj = self.TransformerDecoder(akt1_annot,nodes_logits,akt1_adj,edges_logits)   
+
+        else: 
+            nodes_logits,akt1_annot, edges_logits, akt1_adj = self.TransformerDecoder(nodes_logits,akt1_annot,edges_logits,akt1_adj)
      
         edges_logits = self.edges_output_layer(edges_logits)
         nodes_logits = self.nodes_output_layer(nodes_logits)
@@ -203,190 +207,4 @@ class simple_disc(nn.Module):
         
         #prediction = F.softmax(prediction,dim=-1)
         
-        return prediction
-
-"""class Discriminator(nn.Module):
-  
-    def __init__(self,deg,agg,sca,pna_in_ch,pna_out_ch,edge_dim,towers,pre_lay,post_lay,pna_layer_num, graph_add):
-        super(Discriminator, self).__init__()
-        self.degree = deg
-        self.aggregators = agg
-        self.scalers = sca
-        self.pna_in_channels = pna_in_ch
-        self.pna_out_channels = pna_out_ch
-        self.edge_dimension = edge_dim
-        self.towers = towers
-        self.pre_layers_num = pre_lay
-        self.post_layers_num = post_lay
-        self.pna_layer_num = pna_layer_num
-        self.graph_add = graph_add
-        self.PNA_layer = PNA(deg=self.degree, agg =self.aggregators,sca = self.scalers,
-                             pna_in_ch= self.pna_in_channels, pna_out_ch = self.pna_out_channels, edge_dim = self.edge_dimension,
-                             towers = self.towers, pre_lay = self.pre_layers_num, post_lay = self.post_layers_num,
-                             pna_layer_num = self.pna_layer_num, graph_add = self.graph_add)
-
-    def forward(self, x, edge_index, edge_attr, batch, activation=None):
-
-        h = self.PNA_layer(x, edge_index, edge_attr, batch)
-
-        h = activation(h) if activation is not None else h
-        
-        return h"""
-
-"""class Discriminator2(nn.Module):
-
-    def __init__(self,deg,agg,sca,pna_in_ch,pna_out_ch,edge_dim,towers,pre_lay,post_lay,pna_layer_num, graph_add):
-        super(Discriminator2, self).__init__()
-        self.degree = deg
-        self.aggregators = agg
-        self.scalers = sca
-        self.pna_in_channels = pna_in_ch
-        self.pna_out_channels = pna_out_ch
-        self.edge_dimension = edge_dim
-        self.towers = towers
-        self.pre_layers_num = pre_lay
-        self.post_layers_num = post_lay
-        self.pna_layer_num = pna_layer_num
-        self.graph_add = graph_add
-        self.PNA_layer = PNA(deg=self.degree, agg =self.aggregators,sca = self.scalers,
-                             pna_in_ch= self.pna_in_channels, pna_out_ch = self.pna_out_channels, edge_dim = self.edge_dimension,
-                             towers = self.towers, pre_lay = self.pre_layers_num, post_lay = self.post_layers_num,
-                             pna_layer_num = self.pna_layer_num, graph_add = self.graph_add)
-
-    def forward(self, x, edge_index, edge_attr, batch, activation=None):
-
-        h = self.PNA_layer(x, edge_index, edge_attr, batch)
-
-        h = activation(h) if activation is not None else h
-        
-        return h"""
-
-
-"""class Discriminator_old(nn.Module):
-
-    def __init__(self, conv_dim, m_dim, b_dim, dropout, gcn_depth):
-        super(Discriminator_old, self).__init__()
-
-        graph_conv_dim, aux_dim, linear_dim = conv_dim
-        
-        # discriminator
-        self.gcn_layer = GraphConvolution(m_dim, graph_conv_dim, b_dim, dropout,gcn_depth)
-        self.agg_layer = GraphAggregation(graph_conv_dim[-1], aux_dim, m_dim, dropout)
-
-        # multi dense layer
-        layers = []
-        for c0, c1 in zip([aux_dim]+linear_dim[:-1], linear_dim):
-            layers.append(nn.Linear(c0,c1))
-            layers.append(nn.Dropout(dropout))
-        self.linear_layer = nn.Sequential(*layers)
-        
-        self.output_layer = nn.Linear(linear_dim[-1], 1)
-
-    def forward(self, adj, hidden, node, activation=None):
-        
-        adj = adj[:,:,:,1:].permute(0,3,1,2)
-
-        annotations = torch.cat((hidden, node), -1) if hidden is not None else node
-        
-        h = self.gcn_layer(annotations, adj)
-        annotations = torch.cat((h, hidden, node) if hidden is not None\
-                                 else (h, node), -1)
-        
-        h = self.agg_layer(annotations, torch.tanh)
-        h = self.linear_layer(h)
-        
-        # Need to implement batch discriminator #
-        #########################################
-
-        output = self.output_layer(h)
-        output = activation(output) if activation is not None else output
-        
-        return output, h"""
-    
-"""class Discriminator_old2(nn.Module):
-
-    def __init__(self, conv_dim, m_dim, b_dim, dropout, gcn_depth):
-        super(Discriminator_old2, self).__init__()
-
-        graph_conv_dim, aux_dim, linear_dim = conv_dim
-        
-        # discriminator
-        self.gcn_layer = GraphConvolution(m_dim, graph_conv_dim, b_dim, dropout, gcn_depth)
-        self.agg_layer = GraphAggregation(graph_conv_dim[-1], aux_dim, m_dim, dropout)
-
-        # multi dense layer
-        layers = []
-        for c0, c1 in zip([aux_dim]+linear_dim[:-1], linear_dim):
-            layers.append(nn.Linear(c0,c1))
-            layers.append(nn.Dropout(dropout))
-        self.linear_layer = nn.Sequential(*layers)
-        
-        self.output_layer = nn.Linear(linear_dim[-1], 1)
-
-    def forward(self, adj, hidden, node, activation=None):
-        
-        adj = adj[:,:,:,1:].permute(0,3,1,2)
-
-        annotations = torch.cat((hidden, node), -1) if hidden is not None else node
-        
-        h = self.gcn_layer(annotations, adj)
-        annotations = torch.cat((h, hidden, node) if hidden is not None\
-                                 else (h, node), -1)
-        
-        h = self.agg_layer(annotations, torch.tanh)
-        h = self.linear_layer(h)
-        
-        # Need to implement batch discriminator #
-        #########################################
-
-        output = self.output_layer(h)
-        output = activation(output) if activation is not None else output
-        
-        return output, h"""
-    
-"""class Discriminator3(nn.Module):
-    
-    def __init__(self,in_ch):
-        super(Discriminator3, self).__init__()
-        self.dim = in_ch
-        
-        
-        self.TraConv_layer = TransformerConv(in_channels = self.dim,out_channels =  self.dim//4,edge_dim = self.dim)  
-        self.mlp = torch.nn.Sequential(torch.nn.Tanh(), torch.nn.Linear(self.dim//4,1))
-    def forward(self, x, edge_index, edge_attr, batch, activation=None):
-
-        h = self.TraConv_layer(x, edge_index, edge_attr)
-        h = global_add_pool(h,batch)
-        h = self.mlp(h)
-        h = activation(h) if activation is not None else h
-        
-        return h"""
-    
-    
-"""class PNA_Net(nn.Module):
-    def __init__(self,deg):
-        super().__init__()
-
-      
-
-        self.convs = nn.ModuleList()
-        
-        self.lin = nn.Linear(5, 128)
-        for _ in range(1):
-            conv = DenseGCNConv(128, 128, improved=False, bias=True)
-            self.convs.append(conv)
-            
-        self.agg_layer = GraphAggregation(128, 128, 0, dropout=0.1)
-        self.mlp = nn.Sequential(nn.Linear(128, 64), nn.Tanh(), nn.Linear(64, 32), nn.Tanh(),
-                              nn.Linear(32, 1))
-
-    def forward(self, x, adj,mask=None):
-        x = self.lin(x)
-        
-        for conv in self.convs:
-            x = F.relu(conv(x, adj,mask=None))
-
-        x = self.agg_layer(x,torch.tanh)
-       
-        return self.mlp(x) """    
-    
+        return prediction