Views
No views yet
1git clone https://github.com/IBM/otter-knowledge.git
2cd otter-knowledge1import torch
2from torch import nn
3
4
5class BindingAffinity(nn.Module):
6
7 def __init__(self, gnn, drug_modality):
8 super(BindingAffinity, self).__init__()
9 self.drug_modality = drug_modality
10 self.protein_modality = 'protein-sequence-mean'
11 self.drug_entity_name = 'Drug'
12 self.protein_entity_name = 'Protein'
13 self.drug_rel_id = 1
14 self.protein_rel_id = 2
15 self.protein_drug_rel_id = 0
16 self.gnn = gnn
17 self.device = 'cpu'
18 hd1 = 512
19 num_input = 2
20 self.combine = torch.nn.ModuleList([nn.Linear(num_input * hd1, hd1), nn.ReLU(),
21 nn.Linear(hd1, hd1), nn.ReLU(),
22 nn.Linear(hd1, 1)])
23 self.to(self.device)
24
25 def forward(self, drug_embedding, protein_embedding):
26 nodes = {
27 self.drug_modality: {
28 'embeddings': drug_embedding.unsqueeze(0).to(self.device),
29 'node_indices': torch.tensor([1]).to(self.device)
30 },
31 self.drug_entity_name: {
32 'embeddings': [None],
33 'node_indices': torch.tensor([0]).to(self.device)
34 },
35 self.protein_modality: {
36 'embeddings': protein_embedding.unsqueeze(0).to(self.device),
37 'node_indices': torch.tensor([3]).to(self.device)
38 },
39 self.protein_entity_name: {
40 'embeddings': [None],
41 'node_indices': torch.tensor([2]).to(self.device)
42 }
43 }
44 triples = torch.tensor([[1, 3],
45 [3, 4],
46 [0, 2]]).to(self.device)
47 gnn_embeddings = self.gnn.encoder(nodes, triples)
48 node_gnn_embeddings = []
49 all_indices = [0, 2]
50
51 for indices in all_indices:
52 node_gnn_embedding = torch.index_select(gnn_embeddings, dim=0, index=torch.tensor(indices).to(self.device))
53 node_gnn_embeddings.append(node_gnn_embedding)
54
55 c = torch.cat(node_gnn_embeddings, dim=-1)
56 for m in self.combine:
57 c = m(c)
58
59 return c```
60
61- Run the inference with the initial embeddings (embeddings obtained after using the handlers (Chemberta, ESM1b) over the SMILES and the protein sequence):
62
63```python
64p = net(drug_embedding=drug_embedding, protein_embedding=protein_embedding)
65print(p)```