Initial commit

tokenizer/special_tokens_map.json

@@ -0,0 +1 @@
+{"bos_token": {"content": "<s>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "eos_token": {"content": "</s>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "unk_token": {"content": "<unk>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "sep_token": {"content": "</s>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "pad_token": {"content": "<pad>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "cls_token": {"content": "<s>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "mask_token": {"content": "<mask>", "single_word": false, "lstrip": true, "rstrip": false, "normalized": true}}

tokenizer/vocab.json

@@ -0,0 +1 @@
+{"<unk>":0,"<s>":1,"</s>":2,"<pad>":3,"<mask>":4,"\n":5,"#":6,"+":7,"-":8,".":9,"/":10,"0":11,"1":12,"2":13,"3":14,"4":15,"5":16,"6":17,"7":18,"8":19,"9":20,"=":21,"@":22,"A":23,"B":24,"C":25,"F":26,"H":27,"I":28,"K":29,"L":30,"M":31,"N":32,"O":33,"P":34,"R":35,"S":36,"T":37,"Z":38,"\\":39,"a":40,"b":41,"c":42,"e":43,"g":44,"h":45,"i":46,"l":47,"n":48,"r":49,"s":50,"Br":51,"an":52,"ch":53,"Bran":54,"Branch":55,"Branch1":56,"=C":57,"Ri":58,"ng":59,"Ring":60,"Ring1":61,"=Branch1":62,"Branch2":63,"=O":64,"Ring2":65,"H1":66,"C@":67,"=N":68,"#Branch1":69,"C@@":70,"=Branch2":71,"C@H1":72,"C@@H1":73,"#Branch2":74,"#C":75,"Cl":76,"/C":77,"NH1":78,"=Ring1":79,"+1":80,"-1":81,"O-1":82,"N+1":83,"\\C":84,"#N":85,"/N":86,"=Ring2":87,"=S":88,"=N+1":89,"\\N":90,"Na":91,"Na+1":92,"/O":93,"\\O":94,"Br-1":95,"Branch3":96,"\\S":97,"S+1":98,"Cl-1":99,"I-1":100,"/C@@H1":101,"Si":102,"/C@H1":103,"/S":104,"=N-1":105,"Se":106,"=P":107,"N-1":108,"Ring3":109,"2H":110,"P+1":111,"K+1":112,"\\C@@H1":113,"\\C@H1":114,"/N+1":115,"@@":116,"C-1":117,"#N+1":118,"B-1":119,"+3":120,"Cl+3":121,"\\NH1":122,"Li":123,"Li+1":124,"PH1":125,"18":126,"18F":127,"@+1":128,"3H":129,"P@@":130,"H0":131,"OH0":132,"12":133,"P@":134,"+2":135,"@@+1":136,"S-1":137,"/Br":138,"-/":139,"\\Cl":140,"-/Ring2":141,"\\O-1":142,"11":143,"5I":144,"125I":145,"11C":146,"H3":147,"\\N+1":148,"-\\":149,"/C@@":150,"S@+1":151,"As":152,"/Cl":153,"11CH3":154,"=Se":155,"S@@+1":156,"N@+1":157,"14":158,"-\\Ring2":159,"14C":160,"\\F":161,"/C@":162,"Te":163,"H2":164,"H1-1":165,"=O+1":166,"N@@+1":167,"C+1":168,"=S+1":169,"Zn":170,"/P":171,"a+2":172,"/I":173,"OH1-1":174,"Ca+2":175,"\\Br":176,"Mg":177,"Zn+2":178,"Al":179,"/F":180,"Mg+2":181,"123":182,"123I":183,"13":184,"I+1":185,"/O-1":186,"-\\Ring1":187,"BH2":188,"BH2-1":189,"\\I":190,"/NH1":191,"O+1":192,"131":193,"131I":194,"=14C":195,"/S+1":196,"=Ring3":197,"\\C@@":198,"H2+1":199,"\\C@":200,"Ag":201,"=As":202,"=Se+1":203,"NH2+1":204,"SeH1":205,"-/Ring1":206,"=Te":207,"Al+3":208,"NaH1":209,"=Te+1":210,"NH1+1":211,"Ag+1":212,"H1+1":213,"NH1-1":214,"\\P":215,"14CH2":216,"13C":217,"14CH1":218,"=11C":219,"S@@":220,"=P@@":221,"SiH2":222,"H3-1":223,"14CH3":224,"BH3-1":225,"S@":226,"=14CH1":227,"=PH1":228,"=P@":229,"=NH1+1":230,"\\S+1":231,"124":232,"CH1-1":233,"Sr":234,"=Si":235,"124I":236,"Sr+2":237,"#C-1":238,"/C-1":239,"N@":240,"/N-1":241,"13CH1":242,"/B":243,"19":244,"Ba+2":245,"H4":246,"SH1+1":247,"Se+1":248,"19F":249,"/125I":250,"P@+1":251,"Rb":252,"Cl+1":253,"SiH4":254,"Rb+1":255,"=Branch3":256,"N@@":257,"As+1":258,"/Si":259,"BH1-1":260,"SH1":261,"/123I":262,"32":263,"=Mg":264,"H+1":265,"\\B":266,"SiH1":267,"P@@+1":268,"-2":269,"15":270,"17":271,"35":272,"=13CH1":273,"Cs":274,"=NH2+1":275,"=SH1":276,"MgH2":277,"32P":278,"17F":279,"35S":280,"Cs+1":281,"#11C":282,"/131I":283,"Bi":284,"\\125I":285,"=S@@":286,"\\S-1":287,"6Br":288,"7I":289,"76Br":290,"=B":291,"eH1":292,"\\N-1":293,"18O":294,"127I":295,"11CH2":296,"14C@@H1":297,"TeH2":298,"15NH1":299,"Bi+3":300,"/P+1":301,"/13C":302,"/13CH1":303,"0B":304,"10B":305,"=Al":306,"=18O":307,"BH0":308,"F-1":309,"NH3":310,"S-2":311,"Br+2":312,"Cl+2":313,"\\Si":314,"/S-1":315,"=PH2":316,"14C@H1":317,"NH3+1":318,"#14C":319,"#O+1":320,"-3":321,"22":322,"4H":323,"5Se":324,"5Sr+2":325,"75Se":326,"85Sr+2":327,"=B-1":328,"=13C":329,"@-1":330,"Be":331,"B@@":332,"B@-1":333,"Ca":334,"CH1":335,"I+3":336,"KH1":337,"OH1+1":338,"Ra+2":339,"SH1-1":340,"\\PH1":341,"\\123I":342,"=Ca":343,"\\CH1-1":344,"=S@":345,"\\SeH1":346,"/SeH1":347,"Se-1":348,"LiH1":349,"18F-1":350,"125IH1":351,"11CH1":352,"TeH1":353,"Zn+1":354,"Zn-2":355,"Al-3":356,"13CH3":357,"15N":358,"Be+2":359,"B@@-1":360,"#P":361,"#S":362,"-4":363,"/PH1":364,"/P@@":365,"/As":366,"/14C":367,"/14CH1":368,"2K+1":369,"2Rb+1":370,"3Se":371,"3Ra+2":372,"45":373,"47":374,"42K+1":375,"5I-1":376,"73Se":377,"89":378,"82Rb+1":379,"=32":380,"=32P":381,"CH0":382,"CH2":383,"I+2":384,"NH0":385,"NH4":386,"OH1":387,"PH2+1":388,"SH0":389,"SH2":390,"\\3H":391,"\\11CH3":392,"\\C-1":393,"\\Se":394,"Si@":395,"Si-1":396,"SiH1-1":397,"SiH3-1":398,"/Se":399,"Se-2":400,"\\NH1-1":401,"18FH1":402,"125I-1":403,"11C@@H1":404,"11C-1":405,"AsH1":406,"As-1":407,"14C@@":408,"Te-1":409,"Mg+1":410,"123I-1":411,"123Te":412,"123IH1":413,"135I":414,"131I-1":415,"Ag-4":416,"124I-1":417,"76BrH1":418,"18OH1":419,"22Na+1":420,"223Ra+2":421,"CaH2":422,"45Ca+2":423,"47Ca+2":424,"89Sr+2":425,"=32PH1":426,"NH4+1":427}

tokenizer/vocab.txt

@@ -0,0 +1,429 @@
+<unk>
+<s>
+</s>
+<pad>
+<mask>
+
+
+#
++
+-
+.
+/
+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+=
+@
+A
+B
+C
+F
+H
+I
+K
+L
+M
+N
+O
+P
+R
+S
+T
+Z
+\
+a
+b
+c
+e
+g
+h
+i
+l
+n
+r
+s
+Br
+an
+ch
+Bran
+Branch
+Branch1
+=C
+Ri
+ng
+Ring
+Ring1
+=Branch1
+Branch2
+=O
+Ring2
+H1
+C@
+=N
+#Branch1
+C@@
+=Branch2
+C@H1
+C@@H1
+#Branch2
+#C
+Cl
+/C
+NH1
+=Ring1
++1
+-1
+O-1
+N+1
+\C
+#N
+/N
+=Ring2
+=S
+=N+1
+\N
+Na
+Na+1
+/O
+\O
+Br-1
+Branch3
+\S
+S+1
+Cl-1
+I-1
+/C@@H1
+Si
+/C@H1
+/S
+=N-1
+Se
+=P
+N-1
+Ring3
+2H
+P+1
+K+1
+\C@@H1
+\C@H1
+/N+1
+@@
+C-1
+#N+1
+B-1
++3
+Cl+3
+\NH1
+Li
+Li+1
+PH1
+18
+18F
+@+1
+3H
+P@@
+H0
+OH0
+12
+P@
++2
+@@+1
+S-1
+/Br
+-/
+\Cl
+-/Ring2
+\O-1
+11
+5I
+125I
+11C
+H3
+\N+1
+-\
+/C@@
+S@+1
+As
+/Cl
+11CH3
+=Se
+S@@+1
+N@+1
+14
+-\Ring2
+14C
+\F
+/C@
+Te
+H2
+H1-1
+=O+1
+N@@+1
+C+1
+=S+1
+Zn
+/P
+a+2
+/I
+OH1-1
+Ca+2
+\Br
+Mg
+Zn+2
+Al
+/F
+Mg+2
+123
+123I
+13
+I+1
+/O-1
+-\Ring1
+BH2
+BH2-1
+\I
+/NH1
+O+1
+131
+131I
+=14C
+/S+1
+=Ring3
+\C@@
+H2+1
+\C@
+Ag
+=As
+=Se+1
+NH2+1
+SeH1
+-/Ring1
+=Te
+Al+3
+NaH1
+=Te+1
+NH1+1
+Ag+1
+H1+1
+NH1-1
+\P
+14CH2
+13C
+14CH1
+=11C
+S@@
+=P@@
+SiH2
+H3-1
+14CH3
+BH3-1
+S@
+=14CH1
+=PH1
+=P@
+=NH1+1
+\S+1
+124
+CH1-1
+Sr
+=Si
+124I
+Sr+2
+#C-1
+/C-1
+N@
+/N-1
+13CH1
+/B
+19
+Ba+2
+H4
+SH1+1
+Se+1
+19F
+/125I
+P@+1
+Rb
+Cl+1
+SiH4
+Rb+1
+=Branch3
+N@@
+As+1
+/Si
+BH1-1
+SH1
+/123I
+32
+=Mg
+H+1
+\B
+SiH1
+P@@+1
+-2
+15
+17
+35
+=13CH1
+Cs
+=NH2+1
+=SH1
+MgH2
+32P
+17F
+35S
+Cs+1
+#11C
+/131I
+Bi
+\125I
+=S@@
+\S-1
+6Br
+7I
+76Br
+=B
+eH1
+\N-1
+18O
+127I
+11CH2
+14C@@H1
+TeH2
+15NH1
+Bi+3
+/P+1
+/13C
+/13CH1
+0B
+10B
+=Al
+=18O
+BH0
+F-1
+NH3
+S-2
+Br+2
+Cl+2
+\Si
+/S-1
+=PH2
+14C@H1
+NH3+1
+#14C
+#O+1
+-3
+22
+4H
+5Se
+5Sr+2
+75Se
+85Sr+2
+=B-1
+=13C
+@-1
+Be
+B@@
+B@-1
+Ca
+CH1
+I+3
+KH1
+OH1+1
+Ra+2
+SH1-1
+\PH1
+\123I
+=Ca
+\CH1-1
+=S@
+\SeH1
+/SeH1
+Se-1
+LiH1
+18F-1
+125IH1
+11CH1
+TeH1
+Zn+1
+Zn-2
+Al-3
+13CH3
+15N
+Be+2
+B@@-1
+#P
+#S
+-4
+/PH1
+/P@@
+/As
+/14C
+/14CH1
+2K+1
+2Rb+1
+3Se
+3Ra+2
+45
+47
+42K+1
+5I-1
+73Se
+89
+82Rb+1
+=32
+=32P
+CH0
+CH2
+I+2
+NH0
+NH4
+OH1
+PH2+1
+SH0
+SH2
+\3H
+\11CH3
+\C-1
+\Se
+Si@
+Si-1
+SiH1-1
+SiH3-1
+/Se
+Se-2
+\NH1-1
+18FH1
+125I-1
+11C@@H1
+11C-1
+AsH1
+As-1
+14C@@
+Te-1
+Mg+1
+123I-1
+123Te
+123IH1
+135I
+131I-1
+Ag-4
+124I-1
+76BrH1
+18OH1
+22Na+1
+223Ra+2
+CaH2
+45Ca+2
+47Ca+2
+89Sr+2
+=32PH1
+NH4+1

utils/.ipynb_checkpoints/drug_tokenizer-checkpoint.py

@@ -0,0 +1,66 @@
+import json
+import re
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+class DrugTokenizer:
+    def __init__(self, vocab_path="tokenizer/vocab.json", special_tokens_path="tokenizer/special_tokens_map.json"):
+        self.vocab, self.special_tokens = self.load_vocab_and_special_tokens(vocab_path, special_tokens_path)
+        self.cls_token_id = self.vocab[self.special_tokens['cls_token']]
+        self.sep_token_id = self.vocab[self.special_tokens['sep_token']]
+        self.unk_token_id = self.vocab[self.special_tokens['unk_token']]
+        self.pad_token_id = self.vocab[self.special_tokens['pad_token']]
+        self.id_to_token = {v: k for k, v in self.vocab.items()}
+
+    def load_vocab_and_special_tokens(self, vocab_path, special_tokens_path):
+        with open(vocab_path, 'r', encoding='utf-8') as vocab_file:
+            vocab = json.load(vocab_file)
+        with open(special_tokens_path, 'r', encoding='utf-8') as special_tokens_file:
+            special_tokens_raw = json.load(special_tokens_file)
+
+        special_tokens = {key: value['content'] for key, value in special_tokens_raw.items()}
+        return vocab, special_tokens
+
+    def encode(self, sequence):
+        tokens = re.findall(r'\[([^\[\]]+)\]', sequence)
+        input_ids = [self.cls_token_id] + [self.vocab.get(token, self.unk_token_id) for token in tokens] + [self.sep_token_id]
+        attention_mask = [1] * len(input_ids)
+        return {
+            'input_ids': input_ids,
+            'attention_mask': attention_mask
+        }
+
+    def batch_encode_plus(self, sequences, max_length, padding, truncation, add_special_tokens, return_tensors):
+        input_ids_list = []
+        attention_mask_list = []
+
+        for sequence in sequences:
+            encoded = self.encode(sequence)
+            input_ids = encoded['input_ids']
+            attention_mask = encoded['attention_mask']
+
+            if len(input_ids) > max_length:
+                input_ids = input_ids[:max_length]
+                attention_mask = attention_mask[:max_length]
+            elif len(input_ids) < max_length:
+                pad_length = max_length - len(input_ids)
+                input_ids = input_ids + [self.vocab[self.special_tokens['pad_token']]] * pad_length
+                attention_mask = attention_mask + [0] * pad_length
+
+            input_ids_list.append(input_ids)
+            attention_mask_list.append(attention_mask)
+
+        return {
+            'input_ids': torch.tensor(input_ids_list, dtype=torch.long),
+            'attention_mask': torch.tensor(attention_mask_list, dtype=torch.long)
+        }
+
+    def decode(self, input_ids, skip_special_tokens=False):
+        tokens = []
+        for id in input_ids:
+            if skip_special_tokens and id in [self.cls_token_id, self.sep_token_id, self.pad_token_id]:
+                continue
+            tokens.append(self.id_to_token.get(id, self.special_tokens['unk_token']))
+        sequence = ''.join([f'[{token}]' for token in tokens])
+        return sequence

utils/drug_tokenizer.py

@@ -0,0 +1,66 @@
+import json
+import re
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+class DrugTokenizer:
+    def __init__(self, vocab_path="tokenizer/vocab.json", special_tokens_path="tokenizer/special_tokens_map.json"):
+        self.vocab, self.special_tokens = self.load_vocab_and_special_tokens(vocab_path, special_tokens_path)
+        self.cls_token_id = self.vocab[self.special_tokens['cls_token']]
+        self.sep_token_id = self.vocab[self.special_tokens['sep_token']]
+        self.unk_token_id = self.vocab[self.special_tokens['unk_token']]
+        self.pad_token_id = self.vocab[self.special_tokens['pad_token']]
+        self.id_to_token = {v: k for k, v in self.vocab.items()}
+
+    def load_vocab_and_special_tokens(self, vocab_path, special_tokens_path):
+        with open(vocab_path, 'r', encoding='utf-8') as vocab_file:
+            vocab = json.load(vocab_file)
+        with open(special_tokens_path, 'r', encoding='utf-8') as special_tokens_file:
+            special_tokens_raw = json.load(special_tokens_file)
+
+        special_tokens = {key: value['content'] for key, value in special_tokens_raw.items()}
+        return vocab, special_tokens
+
+    def encode(self, sequence):
+        tokens = re.findall(r'\[([^\[\]]+)\]', sequence)
+        input_ids = [self.cls_token_id] + [self.vocab.get(token, self.unk_token_id) for token in tokens] + [self.sep_token_id]
+        attention_mask = [1] * len(input_ids)
+        return {
+            'input_ids': input_ids,
+            'attention_mask': attention_mask
+        }
+
+    def batch_encode_plus(self, sequences, max_length, padding, truncation, add_special_tokens, return_tensors):
+        input_ids_list = []
+        attention_mask_list = []
+
+        for sequence in sequences:
+            encoded = self.encode(sequence)
+            input_ids = encoded['input_ids']
+            attention_mask = encoded['attention_mask']
+
+            if len(input_ids) > max_length:
+                input_ids = input_ids[:max_length]
+                attention_mask = attention_mask[:max_length]
+            elif len(input_ids) < max_length:
+                pad_length = max_length - len(input_ids)
+                input_ids = input_ids + [self.vocab[self.special_tokens['pad_token']]] * pad_length
+                attention_mask = attention_mask + [0] * pad_length
+
+            input_ids_list.append(input_ids)
+            attention_mask_list.append(attention_mask)
+
+        return {
+            'input_ids': torch.tensor(input_ids_list, dtype=torch.long),
+            'attention_mask': torch.tensor(attention_mask_list, dtype=torch.long)
+        }
+
+    def decode(self, input_ids, skip_special_tokens=False):
+        tokens = []
+        for id in input_ids:
+            if skip_special_tokens and id in [self.cls_token_id, self.sep_token_id, self.pad_token_id]:
+                continue
+            tokens.append(self.id_to_token.get(id, self.special_tokens['unk_token']))
+        sequence = ''.join([f'[{token}]' for token in tokens])
+        return sequence

utils/metric_learning_models_att_maps.py

@@ -0,0 +1,330 @@
+import logging
+import os
+import sys
+
+sys.path.append("../")
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from torch.cuda.amp import autocast
+from torch.nn import Module
+from tqdm import tqdm
+from torch.nn.utils.weight_norm import weight_norm
+from torch.utils.data import Dataset
+
+LOGGER = logging.getLogger(__name__)
+
+class FusionDTI(nn.Module):
+    def __init__(self, prot_out_dim, disease_out_dim, args):
+        super(FusionDTI, self).__init__()
+        self.fusion = args.fusion
+        self.drug_reg = nn.Linear(disease_out_dim, 512)
+        self.prot_reg = nn.Linear(prot_out_dim, 512)
+
+        if self.fusion == "CAN":
+            self.can_layer = CAN_Layer(hidden_dim=512, num_heads=8, args=args)
+            self.mlp_classifier = MlPdecoder_CAN(input_dim=1024)
+        elif self.fusion == "BAN":
+            self.ban_layer = weight_norm(BANLayer(512, 512, 256, 2), name='h_mat', dim=None)
+            self.mlp_classifier = MlPdecoder_CAN(input_dim=256)
+        elif self.fusion == "Nan":
+            self.mlp_classifier_nan = MlPdecoder_CAN(input_dim=1214)
+
+    def forward(self, prot_embed, drug_embed, prot_mask, drug_mask):
+        # print("drug_embed", drug_embed.shape)
+        if self.fusion == "Nan":
+            prot_embed = prot_embed.mean(1)  # query : [batch_size, hidden]
+            drug_embed = drug_embed.mean(1)  # query : [batch_size, hidden]
+            joint_embed = torch.cat([prot_embed, drug_embed], dim=1)
+            score = self.mlp_classifier_nan(joint_embed)
+        else:
+            prot_embed = self.prot_reg(prot_embed)
+            drug_embed = self.drug_reg(drug_embed)
+
+            if self.fusion == "CAN":
+                joint_embed, att = self.can_layer(prot_embed, drug_embed, prot_mask, drug_mask)
+            elif self.fusion == "BAN":
+                joint_embed, att = self.ban_layer(prot_embed, drug_embed)
+
+            score = self.mlp_classifier(joint_embed)
+
+        return score, att
+    
+class Pre_encoded(nn.Module):
+    def __init__(
+            self, prot_encoder, drug_encoder, args
+    ):
+        """Constructor for the model.
+
+        Args:
+            prot_encoder (_type_): Protein sturcture-aware sequence encoder.
+            drug_encoder (_type_): Drug SFLFIES encoder.
+            args (_type_): _description_
+        """
+        super(Pre_encoded, self).__init__()
+        self.prot_encoder = prot_encoder
+        self.drug_encoder = drug_encoder
+        
+    def encoding(self, prot_input_ids, prot_attention_mask, drug_input_ids, drug_attention_mask):
+        # Process inputs through encoders
+        prot_embed = self.prot_encoder(
+            input_ids=prot_input_ids, attention_mask=prot_attention_mask, return_dict=True
+        ).logits
+        # prot_embed = self.prot_reg(prot_embed)
+
+        drug_embed = self.drug_encoder(
+            input_ids=drug_input_ids, attention_mask=drug_attention_mask, return_dict=True
+        ).last_hidden_state  # .last_hidden_state
+        
+        # print("drug_embed", drug_embed.shape)
+        
+        return prot_embed, drug_embed
+    
+    
+class CAN_Layer(nn.Module):
+    def __init__(self, hidden_dim, num_heads, args):
+        super(CAN_Layer, self).__init__()
+        self.agg_mode = args.agg_mode
+        self.group_size = args.group_size  #  Control Fusion Scale
+        self.hidden_dim = hidden_dim
+        self.num_heads = num_heads
+        self.head_size = hidden_dim // num_heads
+
+        self.query_p = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.key_p = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.value_p = nn.Linear(hidden_dim, hidden_dim, bias=False)
+
+        self.query_d = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.key_d = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.value_d = nn.Linear(hidden_dim, hidden_dim, bias=False)
+
+    def alpha_logits(self, logits, mask_row, mask_col, inf=1e6):
+        N, L1, L2, H = logits.shape
+        mask_row = mask_row.view(N, L1, 1).repeat(1, 1, H)
+        mask_col = mask_col.view(N, L2, 1).repeat(1, 1, H)
+        mask_pair = torch.einsum('blh, bkh->blkh', mask_row, mask_col)
+
+        logits = torch.where(mask_pair, logits, logits - inf)
+        alpha = torch.softmax(logits, dim=2)
+        mask_row = mask_row.view(N, L1, 1, H).repeat(1, 1, L2, 1)
+        alpha = torch.where(mask_row, alpha, torch.zeros_like(alpha))
+        return alpha
+
+    def apply_heads(self, x, n_heads, n_ch):
+        s = list(x.size())[:-1] + [n_heads, n_ch]
+        return x.view(*s)
+
+    def group_embeddings(self, x, mask, group_size):
+        N, L, D = x.shape
+        groups = L // group_size
+        x_grouped = x.view(N, groups, group_size, D).mean(dim=2)
+        mask_grouped = mask.view(N, groups, group_size).any(dim=2)
+        return x_grouped, mask_grouped
+
+    def forward(self, protein, drug, mask_prot, mask_drug):
+        # Group embeddings before applying multi-head attention
+        protein_grouped, mask_prot_grouped = self.group_embeddings(protein, mask_prot, self.group_size)
+        drug_grouped, mask_drug_grouped = self.group_embeddings(drug, mask_drug, self.group_size)
+        
+        # print("protein_grouped:", protein_grouped.shape)
+        # print("mask_prot_grouped:", mask_prot_grouped.shape)
+
+        # Compute queries, keys, values for both protein and drug after grouping
+        query_prot = self.apply_heads(self.query_p(protein_grouped), self.num_heads, self.head_size)
+        key_prot = self.apply_heads(self.key_p(protein_grouped), self.num_heads, self.head_size)
+        value_prot = self.apply_heads(self.value_p(protein_grouped), self.num_heads, self.head_size)
+
+        query_drug = self.apply_heads(self.query_d(drug_grouped), self.num_heads, self.head_size)
+        key_drug = self.apply_heads(self.key_d(drug_grouped), self.num_heads, self.head_size)
+        value_drug = self.apply_heads(self.value_d(drug_grouped), self.num_heads, self.head_size)
+
+        # Compute attention scores
+        logits_pp = torch.einsum('blhd, bkhd->blkh', query_prot, key_prot)
+        logits_pd = torch.einsum('blhd, bkhd->blkh', query_prot, key_drug)
+        logits_dp = torch.einsum('blhd, bkhd->blkh', query_drug, key_prot)
+        logits_dd = torch.einsum('blhd, bkhd->blkh', query_drug, key_drug)
+        # print("logits_pp:", logits_pp.shape)
+
+        alpha_pp = self.alpha_logits(logits_pp, mask_prot_grouped, mask_prot_grouped)
+        alpha_pd = self.alpha_logits(logits_pd, mask_prot_grouped, mask_drug_grouped)
+        alpha_dp = self.alpha_logits(logits_dp, mask_drug_grouped, mask_prot_grouped)
+        alpha_dd = self.alpha_logits(logits_dd, mask_drug_grouped, mask_drug_grouped)
+
+        prot_embedding = (torch.einsum('blkh, bkhd->blhd', alpha_pp, value_prot).flatten(-2) +
+                   torch.einsum('blkh, bkhd->blhd', alpha_pd, value_drug).flatten(-2)) / 2
+        drug_embedding = (torch.einsum('blkh, bkhd->blhd', alpha_dp, value_prot).flatten(-2) +
+                   torch.einsum('blkh, bkhd->blhd', alpha_dd, value_drug).flatten(-2)) / 2
+        
+        # print("prot_embedding:", prot_embedding.shape)
+        
+        # Continue as usual with the aggregation mode
+        if self.agg_mode == "cls":
+            prot_embed = prot_embedding[:, 0]  # query : [batch_size, hidden]
+            drug_embed = drug_embedding[:, 0]  # query : [batch_size, hidden]
+        elif self.agg_mode == "mean_all_tok":
+            prot_embed = prot_embedding.mean(1)  # query : [batch_size, hidden]
+            drug_embed = drug_embedding.mean(1)  # query : [batch_size, hidden]
+        elif self.agg_mode == "mean":
+            prot_embed = (prot_embedding * mask_prot_grouped.unsqueeze(-1)).sum(1) / mask_prot_grouped.sum(-1).unsqueeze(-1)
+            drug_embed = (drug_embedding * mask_drug_grouped.unsqueeze(-1)).sum(1) / mask_drug_grouped.sum(-1).unsqueeze(-1)
+        else:
+            raise NotImplementedError()
+            
+        # print("prot_embed:", prot_embed.shape)
+
+        query_embed = torch.cat([prot_embed, drug_embed], dim=1)
+        
+
+        att = torch.zeros(1, 1, 1024, 1024)
+        att[:, :, :512, :512] = alpha_pp.mean(dim=-1)  # Protein to Protein
+        att[:, :, :512, 512:] = alpha_pd.mean(dim=-1)  # Protein to Drug
+        att[:, :, 512:, :512] = alpha_dp.mean(dim=-1)  # Drug to Protein
+        att[:, :, 512:, 512:] = alpha_dd.mean(dim=-1)  # Drug to Drug
+
+        # print("query_embed:", query_embed.shape)
+        return query_embed, att
+
+class MlPdecoder_CAN(nn.Module):
+    def __init__(self, input_dim):
+        super(MlPdecoder_CAN, self).__init__()
+        self.fc1 = nn.Linear(input_dim, input_dim)
+        self.bn1 = nn.BatchNorm1d(input_dim)
+        self.fc2 = nn.Linear(input_dim, input_dim // 2)
+        self.bn2 = nn.BatchNorm1d(input_dim // 2)
+        self.fc3 = nn.Linear(input_dim // 2, input_dim // 4)
+        self.bn3 = nn.BatchNorm1d(input_dim // 4)
+        self.output = nn.Linear(input_dim // 4, 1)
+
+    def forward(self, x):
+        x = self.bn1(torch.relu(self.fc1(x)))
+        x = self.bn2(torch.relu(self.fc2(x)))
+        x = self.bn3(torch.relu(self.fc3(x)))
+        x = torch.sigmoid(self.output(x))
+        return x
+    
+class MLPdecoder_BAN(nn.Module):
+    def __init__(self, in_dim, hidden_dim, out_dim, binary=1):
+        super(MLPdecoder_BAN, self).__init__()
+        self.fc1 = nn.Linear(in_dim, hidden_dim)
+        self.bn1 = nn.BatchNorm1d(hidden_dim)
+        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
+        self.bn2 = nn.BatchNorm1d(hidden_dim)
+        self.fc3 = nn.Linear(hidden_dim, out_dim)
+        self.bn3 = nn.BatchNorm1d(out_dim)
+        self.fc4 = nn.Linear(out_dim, binary)
+
+    def forward(self, x):
+        x = self.bn1(F.relu(self.fc1(x)))
+        x = self.bn2(F.relu(self.fc2(x)))
+        x = self.bn3(F.relu(self.fc3(x)))
+        # x = self.fc4(x)
+        x = torch.sigmoid(self.fc4(x)) 
+        return x
+
+class BANLayer(nn.Module):
+    """ Bilinear attention network
+    Modified from https://github.com/peizhenbai/DrugBAN/blob/main/ban.py
+    """
+    def __init__(self, v_dim, q_dim, h_dim, h_out, act='ReLU', dropout=0.2, k=3):
+        super(BANLayer, self).__init__()
+
+        self.c = 32
+        self.k = k
+        self.v_dim = v_dim
+        self.q_dim = q_dim
+        self.h_dim = h_dim
+        self.h_out = h_out
+
+        self.v_net = FCNet([v_dim, h_dim * self.k], act=act, dropout=dropout)
+        self.q_net = FCNet([q_dim, h_dim * self.k], act=act, dropout=dropout)
+        # self.dropout = nn.Dropout(dropout[1])
+        if 1 < k:
+            self.p_net = nn.AvgPool1d(self.k, stride=self.k)
+
+        if h_out <= self.c:
+            self.h_mat = nn.Parameter(torch.Tensor(1, h_out, 1, h_dim * self.k).normal_())
+            self.h_bias = nn.Parameter(torch.Tensor(1, h_out, 1, 1).normal_())
+        else:
+            self.h_net = weight_norm(nn.Linear(h_dim * self.k, h_out), dim=None)
+
+        self.bn = nn.BatchNorm1d(h_dim)
+
+    def attention_pooling(self, v, q, att_map):
+        fusion_logits = torch.einsum('bvk,bvq,bqk->bk', (v, att_map, q))
+        if 1 < self.k:
+            fusion_logits = fusion_logits.unsqueeze(1)  # b x 1 x d
+            fusion_logits = self.p_net(fusion_logits).squeeze(1) * self.k  # sum-pooling
+        return fusion_logits
+
+    def forward(self, v, q, softmax=False):
+        v_num = v.size(1)
+        q_num = q.size(1)
+        # print("v_num", v_num)
+        # print("v_num ", v_num)
+        if self.h_out <= self.c:
+            v_ = self.v_net(v)
+            q_ = self.q_net(q)
+            # print("v_", v_.shape)
+            # print("q_ ", q_.shape)
+            att_maps = torch.einsum('xhyk,bvk,bqk->bhvq', (self.h_mat, v_, q_)) + self.h_bias
+            # print("Attention map_1",att_maps.shape)
+        else:
+            v_ = self.v_net(v).transpose(1, 2).unsqueeze(3)
+            q_ = self.q_net(q).transpose(1, 2).unsqueeze(2)
+            d_ = torch.matmul(v_, q_)  # b x h_dim x v x q
+            att_maps = self.h_net(d_.transpose(1, 2).transpose(2, 3))  # b x v x q x h_out
+            att_maps = att_maps.transpose(2, 3).transpose(1, 2)  # b x h_out x v x q
+            # print("Attention map_2",att_maps.shape)
+        if softmax:
+            p = nn.functional.softmax(att_maps.view(-1, self.h_out, v_num * q_num), 2)
+            att_maps = p.view(-1, self.h_out, v_num, q_num)
+            # print("Attention map_softmax", att_maps.shape)
+        logits = self.attention_pooling(v_, q_, att_maps[:, 0, :, :])
+        for i in range(1, self.h_out):
+            logits_i = self.attention_pooling(v_, q_, att_maps[:, i, :, :])
+            logits += logits_i
+        logits = self.bn(logits)
+        return logits, att_maps
+
+    
+class FCNet(nn.Module):
+    """Simple class for non-linear fully connect network
+    Modified from https://github.com/jnhwkim/ban-vqa/blob/master/fc.py
+    """
+
+    def __init__(self, dims, act='ReLU', dropout=0):
+        super(FCNet, self).__init__()
+
+        layers = []
+        for i in range(len(dims) - 2):
+            in_dim = dims[i]
+            out_dim = dims[i + 1]
+            if 0 < dropout:
+                layers.append(nn.Dropout(dropout))
+            layers.append(weight_norm(nn.Linear(in_dim, out_dim), dim=None))
+            if '' != act:
+                layers.append(getattr(nn, act)())
+        if 0 < dropout:
+            layers.append(nn.Dropout(dropout))
+        layers.append(weight_norm(nn.Linear(dims[-2], dims[-1]), dim=None))
+        if '' != act:
+            layers.append(getattr(nn, act)())
+
+        self.main = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.main(x)
+    
+
+class BatchFileDataset_Case(Dataset):
+    def __init__(self, file_list):
+        self.file_list = file_list
+
+    def __len__(self):
+        return len(self.file_list)
+
+    def __getitem__(self, idx):
+        batch_file = self.file_list[idx]
+        data = torch.load(batch_file)
+        return data['prot'], data['drug'], data['prot_ids'], data['drug_ids'], data['prot_mask'], data['drug_mask'], data['y']