 b/unimol/tasks/drugclip.py
+# Copyright (c) DP Technology.
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from IPython import embed as debug_embedded
+import logging
+import os
+from collections.abc import Iterable
+from sklearn.metrics import roc_auc_score
+from xmlrpc.client import Boolean
+import numpy as np
+import torch
+import pickle
+from tqdm import tqdm
+from unicore import checkpoint_utils
+import unicore
+from unicore.data import (AppendTokenDataset, Dictionary, EpochShuffleDataset,
+                          FromNumpyDataset, NestedDictionaryDataset,
+                          PrependTokenDataset, RawArrayDataset,LMDBDataset, RawLabelDataset,
+                          RightPadDataset, RightPadDataset2D, TokenizeDataset,SortDataset,data_utils)
+from unicore.tasks import UnicoreTask, register_task
+from unimol.data import (AffinityDataset, CroppingPocketDataset,
+                         CrossDistanceDataset, DistanceDataset,
+                         EdgeTypeDataset, KeyDataset, LengthDataset,
+                         NormalizeDataset, NormalizeDockingPoseDataset,
+                         PrependAndAppend2DDataset, RemoveHydrogenDataset,
+                         RemoveHydrogenPocketDataset, RightPadDatasetCoord,
+                         RightPadDatasetCross2D, TTADockingPoseDataset, AffinityTestDataset, AffinityValidDataset, AffinityMolDataset, AffinityPocketDataset, ResamplingDataset)
+#from skchem.metrics import bedroc_score
+from rdkit.ML.Scoring.Scoring import CalcBEDROC, CalcAUC, CalcEnrichment
+from sklearn.metrics import roc_curve
+logger = logging.getLogger(__name__)
+def re_new(y_true, y_score, ratio):
+    fp = 0
+    tp = 0
+    p = sum(y_true)
+    n = len(y_true) - p
+    num = ratio*n
+    sort_index = np.argsort(y_score)[::-1]
+    for i in range(len(sort_index)):
+        index = sort_index[i]
+        if y_true[index] == 1:
+            tp += 1
+        else:
+            fp += 1
+            if fp>= num:
+                break
+    return (tp*n)/(p*fp)
+def calc_re(y_true, y_score, ratio_list):
+    fpr, tpr, thresholds = roc_curve(y_true, y_score, pos_label=1)
+    #print(fpr, tpr)
+    res = {}
+    res2 = {}
+    total_active_compounds = sum(y_true)
+    total_compounds = len(y_true)
+    # for ratio in ratio_list:
+    #     for i, t in enumerate(fpr):
+    #         if t > ratio:
+    #             #print(fpr[i], tpr[i])
+    #             if fpr[i-1]==0:
+    #                 res[str(ratio)]=tpr[i]/fpr[i]
+    #             else:
+    #                 res[str(ratio)]=tpr[i-1]/fpr[i-1]
+    #             break
+    for ratio in ratio_list:
+        res2[str(ratio)] = re_new(y_true, y_score, ratio)
+    #print(res)
+    #print(res2)
+    return res2
+def cal_metrics(y_true, y_score, alpha):
+    """
+    Calculate BEDROC score.
+    Parameters:
+    - y_true: true binary labels (0 or 1)
+    - y_score: predicted scores or probabilities
+    - alpha: parameter controlling the degree of early retrieval emphasis
+    Returns:
+    - BEDROC score
+    """
+        # concate res_single and labels
+    scores = np.expand_dims(y_score, axis=1)
+    y_true = np.expand_dims(y_true, axis=1)
+    scores = np.concatenate((scores, y_true), axis=1)
+    # inverse sort scores based on first column
+    scores = scores[scores[:,0].argsort()[::-1]]
+    bedroc = CalcBEDROC(scores, 1, 80.5)
+    count = 0
+    # sort y_score, return index
+    index  = np.argsort(y_score)[::-1]
+    for i in range(int(len(index)*0.005)):
+        if y_true[index[i]] == 1:
+            count += 1
+    auc = CalcAUC(scores, 1)
+    ef_list = CalcEnrichment(scores, 1, [0.005, 0.01, 0.02, 0.05])
+    ef = {
+        "0.005": ef_list[0],
+        "0.01": ef_list[1],
+        "0.02": ef_list[2],
+        "0.05": ef_list[3]
+    }
+    re_list = calc_re(y_true, y_score, [0.005, 0.01, 0.02, 0.05])
+    return auc, bedroc, ef, re_list
+@register_task("drugclip")
+class DrugCLIP(UnicoreTask):
+    """Task for training transformer auto-encoder models."""
+    @staticmethod
+    def add_args(parser):
+        """Add task-specific arguments to the parser."""
+        parser.add_argument(
+            "data",
+            help="downstream data path",
+        )
+        parser.add_argument(
+            "--finetune-mol-model",
+            default=None,
+            type=str,
+            help="pretrained molecular model path",
+        )
+        parser.add_argument(
+            "--finetune-pocket-model",
+            default=None,
+            type=str,
+            help="pretrained pocket model path",
+        )
+        parser.add_argument(
+            "--dist-threshold",
+            type=float,
+            default=6.0,
+            help="threshold for the distance between the molecule and the pocket",
+        )
+        parser.add_argument(
+            "--max-pocket-atoms",
+            type=int,
+            default=256,
+            help="selected maximum number of atoms in a pocket",
+        )
+        parser.add_argument(
+            "--test-model",
+            default=False,
+            type=Boolean,
+            help="whether test model",
+        )
+        parser.add_argument("--reg", action="store_true", help="regression task")
+    def __init__(self, args, dictionary, pocket_dictionary):
+        super().__init__(args)
+        self.dictionary = dictionary
+        self.pocket_dictionary = pocket_dictionary
+        self.seed = args.seed
+        # add mask token
+        self.mask_idx = dictionary.add_symbol("[MASK]", is_special=True)
+        self.pocket_mask_idx = pocket_dictionary.add_symbol("[MASK]", is_special=True)
+        self.mol_reps = None
+        self.keys = None
+    @classmethod
+    def setup_task(cls, args, **kwargs):
+        mol_dictionary = Dictionary.load(os.path.join(args.data, "dict_mol.txt"))
+        pocket_dictionary = Dictionary.load(os.path.join(args.data, "dict_pkt.txt"))
+        logger.info("ligand dictionary: {} types".format(len(mol_dictionary)))
+        logger.info("pocket dictionary: {} types".format(len(pocket_dictionary)))
+        return cls(args, mol_dictionary, pocket_dictionary)
+    def load_dataset(self, split, **kwargs):
+        """Load a given dataset split.
+        'smi','pocket','atoms','coordinates','pocket_atoms','pocket_coordinates'
+        Args:
+            split (str): name of the data scoure (e.g., bppp)
+        """
+        data_path = os.path.join(self.args.data, split + ".lmdb")
+        dataset = LMDBDataset(data_path)
+        if split.startswith("train"):
+            smi_dataset = KeyDataset(dataset, "smi")
+            poc_dataset = KeyDataset(dataset, "pocket")
+            dataset = AffinityDataset(
+                dataset,
+                self.args.seed,
+                "atoms",
+                "coordinates",
+                "pocket_atoms",
+                "pocket_coordinates",
+                "label",
+                True,
+            )
+            tgt_dataset = KeyDataset(dataset, "affinity")
+        else:
+            dataset = AffinityDataset(
+                dataset,
+                self.args.seed,
+                "atoms",
+                "coordinates",
+                "pocket_atoms",
+                "pocket_coordinates",
+                "label",
+            )
+            tgt_dataset = KeyDataset(dataset, "affinity")
+            smi_dataset = KeyDataset(dataset, "smi")
+            poc_dataset = KeyDataset(dataset, "pocket")
+        def PrependAndAppend(dataset, pre_token, app_token):
+            dataset = PrependTokenDataset(dataset, pre_token)
+            return AppendTokenDataset(dataset, app_token)
+        dataset = RemoveHydrogenPocketDataset(
+            dataset,
+            "pocket_atoms",
+            "pocket_coordinates",
+            True,
+            True,
+        )
+        dataset = CroppingPocketDataset(
+            dataset,
+            self.seed,
+            "pocket_atoms",
+            "pocket_coordinates",
+            self.args.max_pocket_atoms,
+        )
+        dataset = RemoveHydrogenDataset(dataset, "atoms", "coordinates", True, True)
+        apo_dataset = NormalizeDataset(dataset, "coordinates")
+        apo_dataset = NormalizeDataset(apo_dataset, "pocket_coordinates")
+        src_dataset = KeyDataset(apo_dataset, "atoms")
+        mol_len_dataset = LengthDataset(src_dataset)
+        src_dataset = TokenizeDataset(
+            src_dataset, self.dictionary, max_seq_len=self.args.max_seq_len
+        )
+        coord_dataset = KeyDataset(apo_dataset, "coordinates")
+        src_dataset = PrependAndAppend(
+            src_dataset, self.dictionary.bos(), self.dictionary.eos()
+        )
+        edge_type = EdgeTypeDataset(src_dataset, len(self.dictionary))
+        coord_dataset = FromNumpyDataset(coord_dataset)
+        distance_dataset = DistanceDataset(coord_dataset)
+        coord_dataset = PrependAndAppend(coord_dataset, 0.0, 0.0)
+        distance_dataset = PrependAndAppend2DDataset(distance_dataset, 0.0)
+        src_pocket_dataset = KeyDataset(apo_dataset, "pocket_atoms")
+        pocket_len_dataset = LengthDataset(src_pocket_dataset)
+        src_pocket_dataset = TokenizeDataset(
+            src_pocket_dataset,
+            self.pocket_dictionary,
+            max_seq_len=self.args.max_seq_len,
+        )
+        coord_pocket_dataset = KeyDataset(apo_dataset, "pocket_coordinates")
+        src_pocket_dataset = PrependAndAppend(
+            src_pocket_dataset,
+            self.pocket_dictionary.bos(),
+            self.pocket_dictionary.eos(),
+        )
+        pocket_edge_type = EdgeTypeDataset(
+            src_pocket_dataset, len(self.pocket_dictionary)
+        )
+        coord_pocket_dataset = FromNumpyDataset(coord_pocket_dataset)
+        distance_pocket_dataset = DistanceDataset(coord_pocket_dataset)
+        coord_pocket_dataset = PrependAndAppend(coord_pocket_dataset, 0.0, 0.0)
+        distance_pocket_dataset = PrependAndAppend2DDataset(
+            distance_pocket_dataset, 0.0
+        )
+        nest_dataset = NestedDictionaryDataset(
+            {
+                "net_input": {
+                    "mol_src_tokens": RightPadDataset(
+                        src_dataset,
+                        pad_idx=self.dictionary.pad(),
+                    ),
+                    "mol_src_distance": RightPadDataset2D(
+                        distance_dataset,
+                        pad_idx=0,
+                    ),
+                    "mol_src_edge_type": RightPadDataset2D(
+                        edge_type,
+                        pad_idx=0,
+                    ),
+                    "pocket_src_tokens": RightPadDataset(
+                        src_pocket_dataset,
+                        pad_idx=self.pocket_dictionary.pad(),
+                    ),
+                    "pocket_src_distance": RightPadDataset2D(
+                        distance_pocket_dataset,
+                        pad_idx=0,
+                    ),
+                    "pocket_src_edge_type": RightPadDataset2D(
+                        pocket_edge_type,
+                        pad_idx=0,
+                    ),
+                    "pocket_src_coord": RightPadDatasetCoord(
+                        coord_pocket_dataset,
+                        pad_idx=0,
+                    ),
+                    "mol_len": RawArrayDataset(mol_len_dataset),
+                    "pocket_len": RawArrayDataset(pocket_len_dataset)
+                },
+                "target": {
+                    "finetune_target": RawLabelDataset(tgt_dataset),
+                },
+                "smi_name": RawArrayDataset(smi_dataset),
+                "pocket_name": RawArrayDataset(poc_dataset),
+            },
+        )
+        if split == "train":
+            with data_utils.numpy_seed(self.args.seed):
+                shuffle = np.random.permutation(len(src_dataset))
+            self.datasets[split] = SortDataset(
+                nest_dataset,
+                sort_order=[shuffle],
+            )
+            self.datasets[split] = ResamplingDataset(
+                self.datasets[split]
+            )
+        else:
+            self.datasets[split] = nest_dataset
+    def load_mols_dataset(self, data_path,atoms,coords, **kwargs):
+        dataset = LMDBDataset(data_path)
+        label_dataset = KeyDataset(dataset, "label")
+        dataset = AffinityMolDataset(
+            dataset,
+            self.args.seed,
+            atoms,
+            coords,
+            False,
+        )
+        smi_dataset = KeyDataset(dataset, "smi")
+        def PrependAndAppend(dataset, pre_token, app_token):
+            dataset = PrependTokenDataset(dataset, pre_token)
+            return AppendTokenDataset(dataset, app_token)
+        dataset = RemoveHydrogenDataset(dataset, "atoms", "coordinates", True, True)
+        apo_dataset = NormalizeDataset(dataset, "coordinates")
+        src_dataset = KeyDataset(apo_dataset, "atoms")
+        len_dataset = LengthDataset(src_dataset)
+        src_dataset = TokenizeDataset(
+            src_dataset, self.dictionary, max_seq_len=self.args.max_seq_len
+        )
+        coord_dataset = KeyDataset(apo_dataset, "coordinates")
+        src_dataset = PrependAndAppend(
+            src_dataset, self.dictionary.bos(), self.dictionary.eos()
+        )
+        edge_type = EdgeTypeDataset(src_dataset, len(self.dictionary))
+        coord_dataset = FromNumpyDataset(coord_dataset)
+        distance_dataset = DistanceDataset(coord_dataset)
+        coord_dataset = PrependAndAppend(coord_dataset, 0.0, 0.0)
+        distance_dataset = PrependAndAppend2DDataset(distance_dataset, 0.0)
+        nest_dataset = NestedDictionaryDataset(
+            {
+                "net_input": {
+                    "mol_src_tokens": RightPadDataset(
+                        src_dataset,
+                        pad_idx=self.dictionary.pad(),
+                    ),
+                    "mol_src_distance": RightPadDataset2D(
+                        distance_dataset,
+                        pad_idx=0,
+                    ),
+                    "mol_src_edge_type": RightPadDataset2D(
+                        edge_type,
+                        pad_idx=0,
+                    ),
+                },
+                "smi_name": RawArrayDataset(smi_dataset),
+                "target":  RawArrayDataset(label_dataset),
+                "mol_len": RawArrayDataset(len_dataset),
+            },
+        )
+        return nest_dataset
+    def load_retrieval_mols_dataset(self, data_path,atoms,coords, **kwargs):
+        dataset = LMDBDataset(data_path)
+        dataset = AffinityMolDataset(
+            dataset,
+            self.args.seed,
+            atoms,
+            coords,
+            False,
+        )
+        smi_dataset = KeyDataset(dataset, "smi")
+        def PrependAndAppend(dataset, pre_token, app_token):
+            dataset = PrependTokenDataset(dataset, pre_token)
+            return AppendTokenDataset(dataset, app_token)
+        dataset = RemoveHydrogenDataset(dataset, "atoms", "coordinates", True, True)
+        apo_dataset = NormalizeDataset(dataset, "coordinates")
+        src_dataset = KeyDataset(apo_dataset, "atoms")
+        len_dataset = LengthDataset(src_dataset)
+        src_dataset = TokenizeDataset(
+            src_dataset, self.dictionary, max_seq_len=self.args.max_seq_len
+        )
+        coord_dataset = KeyDataset(apo_dataset, "coordinates")
+        src_dataset = PrependAndAppend(
+            src_dataset, self.dictionary.bos(), self.dictionary.eos()
+        )
+        edge_type = EdgeTypeDataset(src_dataset, len(self.dictionary))
+        coord_dataset = FromNumpyDataset(coord_dataset)
+        distance_dataset = DistanceDataset(coord_dataset)
+        coord_dataset = PrependAndAppend(coord_dataset, 0.0, 0.0)
+        distance_dataset = PrependAndAppend2DDataset(distance_dataset, 0.0)
+        nest_dataset = NestedDictionaryDataset(
+            {
+                "net_input": {
+                    "mol_src_tokens": RightPadDataset(
+                        src_dataset,
+                        pad_idx=self.dictionary.pad(),
+                    ),
+                    "mol_src_distance": RightPadDataset2D(
+                        distance_dataset,
+                        pad_idx=0,
+                    ),
+                    "mol_src_edge_type": RightPadDataset2D(
+                        edge_type,
+                        pad_idx=0,
+                    ),
+                },
+                "smi_name": RawArrayDataset(smi_dataset),
+                "mol_len": RawArrayDataset(len_dataset),
+            },
+        )
+        return nest_dataset
+    def load_pockets_dataset(self, data_path, **kwargs):
+        dataset = LMDBDataset(data_path)
+        dataset = AffinityPocketDataset(
+            dataset,
+            self.args.seed,
+            "pocket_atoms",
+            "pocket_coordinates",
+            False,
+            "pocket"
+        )
+        poc_dataset = KeyDataset(dataset, "pocket")
+        def PrependAndAppend(dataset, pre_token, app_token):
+            dataset = PrependTokenDataset(dataset, pre_token)
+            return AppendTokenDataset(dataset, app_token)
+        dataset = RemoveHydrogenPocketDataset(
+            dataset,
+            "pocket_atoms",
+            "pocket_coordinates",
+            True,
+            True,
+        )
+        dataset = CroppingPocketDataset(
+            dataset,
+            self.seed,
+            "pocket_atoms",
+            "pocket_coordinates",
+            self.args.max_pocket_atoms,
+        )
+        apo_dataset = NormalizeDataset(dataset, "pocket_coordinates")
+        src_pocket_dataset = KeyDataset(apo_dataset, "pocket_atoms")
+        len_dataset = LengthDataset(src_pocket_dataset)
+        src_pocket_dataset = TokenizeDataset(
+            src_pocket_dataset,
+            self.pocket_dictionary,
+            max_seq_len=self.args.max_seq_len,
+        )
+        coord_pocket_dataset = KeyDataset(apo_dataset, "pocket_coordinates")
+        src_pocket_dataset = PrependAndAppend(
+            src_pocket_dataset,
+            self.pocket_dictionary.bos(),
+            self.pocket_dictionary.eos(),
+        )
+        pocket_edge_type = EdgeTypeDataset(
+            src_pocket_dataset, len(self.pocket_dictionary)
+        )
+        coord_pocket_dataset = FromNumpyDataset(coord_pocket_dataset)
+        distance_pocket_dataset = DistanceDataset(coord_pocket_dataset)
+        coord_pocket_dataset = PrependAndAppend(coord_pocket_dataset, 0.0, 0.0)
+        distance_pocket_dataset = PrependAndAppend2DDataset(
+            distance_pocket_dataset, 0.0
+        )
+        nest_dataset = NestedDictionaryDataset(
+            {
+                "net_input": {
+                    "pocket_src_tokens": RightPadDataset(
+                        src_pocket_dataset,
+                        pad_idx=self.pocket_dictionary.pad(),
+                    ),
+                    "pocket_src_distance": RightPadDataset2D(
+                        distance_pocket_dataset,
+                        pad_idx=0,
+                    ),
+                    "pocket_src_edge_type": RightPadDataset2D(
+                        pocket_edge_type,
+                        pad_idx=0,
+                    ),
+                    "pocket_src_coord": RightPadDatasetCoord(
+                        coord_pocket_dataset,
+                        pad_idx=0,
+                    ),
+                },
+                "pocket_name": RawArrayDataset(poc_dataset),
+                "pocket_len": RawArrayDataset(len_dataset),
+            },
+        )
+        return nest_dataset
+    def build_model(self, args):
+        from unicore import models
+        model = models.build_model(args, self)
+        if args.finetune_mol_model is not None:
+            print("load pretrain model weight from...", args.finetune_mol_model)
+            state = checkpoint_utils.load_checkpoint_to_cpu(
+                args.finetune_mol_model,
+            )
+            model.mol_model.load_state_dict(state["model"], strict=False)
+        if args.finetune_pocket_model is not None:
+            print("load pretrain model weight from...", args.finetune_pocket_model)
+            state = checkpoint_utils.load_checkpoint_to_cpu(
+                args.finetune_pocket_model,
+            )
+            model.pocket_model.load_state_dict(state["model"], strict=False)
+        return model
+    def train_step(
+        self, sample, model, loss, optimizer, update_num, ignore_grad=False
+    ):
+        """
+        Do forward and backward, and return the loss as computed by *loss*
+        for the given *model* and *sample*.
+        Args:
+            sample (dict): the mini-batch. The format is defined by the
+                :class:`~unicore.data.UnicoreDataset`.
+            model (~unicore.models.BaseUnicoreModel): the model
+            loss (~unicore.losses.UnicoreLoss): the loss
+            optimizer (~unicore.optim.UnicoreOptimizer): the optimizer
+            update_num (int): the current update
+            ignore_grad (bool): multiply loss by 0 if this is set to True
+        Returns:
+            tuple:
+                - the loss
+                - the sample size, which is used as the denominator for the
+                  gradient
+                - logging outputs to display while training
+        """
+        model.train()
+        model.set_num_updates(update_num)
+        with torch.autograd.profiler.record_function("forward"):
+            loss, sample_size, logging_output = loss(model, sample)
+        if ignore_grad:
+            loss *= 0
+        with torch.autograd.profiler.record_function("backward"):
+            optimizer.backward(loss)
+        return loss, sample_size, logging_output
+    def valid_step(self, sample, model, loss, test=False):
+        model.eval()
+        with torch.no_grad():
+            loss, sample_size, logging_output = loss(model, sample)
+        return loss, sample_size, logging_output
+    def test_pcba_target(self, name, model, **kwargs):
+        """Encode a dataset with the molecule encoder."""
+        #names = "PPARG"
+        data_path = "./data/lit_pcba/" + name + "/mols.lmdb"
+        mol_dataset = self.load_mols_dataset(data_path, "atoms", "coordinates")
+        num_data = len(mol_dataset)
+        bsz=64
+        #print(num_data//bsz)
+        mol_reps = []
+        mol_names = []
+        labels = []
+        # generate mol data
+        mol_data = torch.utils.data.DataLoader(mol_dataset, batch_size=bsz, collate_fn=mol_dataset.collater)
+        for _, sample in enumerate(tqdm(mol_data)):
+            sample = unicore.utils.move_to_cuda(sample)
+            dist = sample["net_input"]["mol_src_distance"]
+            et = sample["net_input"]["mol_src_edge_type"]
+            st = sample["net_input"]["mol_src_tokens"]
+            mol_padding_mask = st.eq(model.mol_model.padding_idx)
+            mol_x = model.mol_model.embed_tokens(st)
+            n_node = dist.size(-1)
+            gbf_feature = model.mol_model.gbf(dist, et)
+            gbf_result = model.mol_model.gbf_proj(gbf_feature)
+            graph_attn_bias = gbf_result
+            graph_attn_bias = graph_attn_bias.permute(0, 3, 1, 2).contiguous()
+            graph_attn_bias = graph_attn_bias.view(-1, n_node, n_node)
+            mol_outputs = model.mol_model.encoder(
+                mol_x, padding_mask=mol_padding_mask, attn_mask=graph_attn_bias
+            )
+            mol_encoder_rep = mol_outputs[0][:,0,:]
+            mol_emb = model.mol_project(mol_encoder_rep)
+            mol_emb = mol_emb / mol_emb.norm(dim=1, keepdim=True)
+            mol_emb = mol_emb.detach().cpu().numpy()
+            mol_reps.append(mol_emb)
+            mol_names.extend(sample["smi_name"])
+            labels.extend(sample["target"].detach().cpu().numpy())
+        mol_reps = np.concatenate(mol_reps, axis=0)
+        labels = np.array(labels, dtype=np.int32)
+        # generate pocket data
+        data_path = "./data/lit_pcba/" + name + "/pockets.lmdb"
+        pocket_dataset = self.load_pockets_dataset(data_path)
+        pocket_data = torch.utils.data.DataLoader(pocket_dataset, batch_size=bsz, collate_fn=pocket_dataset.collater)
+        pocket_reps = []
+        for _, sample in enumerate(tqdm(pocket_data)):
+            sample = unicore.utils.move_to_cuda(sample)
+            dist = sample["net_input"]["pocket_src_distance"]
+            et = sample["net_input"]["pocket_src_edge_type"]
+            st = sample["net_input"]["pocket_src_tokens"]
+            pocket_padding_mask = st.eq(model.pocket_model.padding_idx)
+            pocket_x = model.pocket_model.embed_tokens(st)
+            n_node = dist.size(-1)
+            gbf_feature = model.pocket_model.gbf(dist, et)
+            gbf_result = model.pocket_model.gbf_proj(gbf_feature)
+            graph_attn_bias = gbf_result
+            graph_attn_bias = graph_attn_bias.permute(0, 3, 1, 2).contiguous()
+            graph_attn_bias = graph_attn_bias.view(-1, n_node, n_node)
+            pocket_outputs = model.pocket_model.encoder(
+                pocket_x, padding_mask=pocket_padding_mask, attn_mask=graph_attn_bias
+            )
+            pocket_encoder_rep = pocket_outputs[0][:,0,:]
+            pocket_emb = model.pocket_project(pocket_encoder_rep)
+            pocket_emb = pocket_emb / pocket_emb.norm(dim=1, keepdim=True)
+            pocket_emb = pocket_emb.detach().cpu().numpy()
+            pocket_names = sample["pocket_name"]
+            pocket_reps.append(pocket_emb)
+        pocket_reps = np.concatenate(pocket_reps, axis=0)
+        res = pocket_reps @ mol_reps.T
+        res_single = res.max(axis=0)
+        auc, bedroc, ef_list, re_list = cal_metrics(labels, res_single, 80.5)
+        return auc, bedroc, ef_list, re_list
+    def test_pcba(self, model, **kwargs):
+        #ckpt_date = self.args.finetune_from_model.split("/")[-2]
+        #save_name = "/home/gaobowen/DrugClip/test_results/pcba/" + ckpt_date + ".txt"
+        save_name = ""
+        targets = os.listdir("./data/lit_pcba/")
+        #print(targets)
+        auc_list = []
+        ef_list = []
+        bedroc_list = []
+        re_list = {
+            "0.005": [],
+            "0.01": [],
+            "0.02": [],
+            "0.05": []
+        }
+        ef_list = {
+            "0.005": [],
+            "0.01": [],
+            "0.02": [],
+            "0.05": []
+        }
+        for target in targets:
+            auc, bedroc, ef, re = self.test_pcba_target(target, model)
+            auc_list.append(auc)
+            bedroc_list.append(bedroc)
+            for key in ef:
+                ef_list[key].append(ef[key])
+            # print("re", re)
+            # print("ef", ef)
+            for key in re:
+                re_list[key].append(re[key])
+        print(auc_list)
+        print(ef_list)
+        print("auc 25%", np.percentile(auc_list, 25))
+        print("auc 50%", np.percentile(auc_list, 50))
+        print("auc 75%", np.percentile(auc_list, 75))
+        print("auc mean", np.mean(auc_list))
+        print("bedroc 25%", np.percentile(bedroc_list, 25))
+        print("bedroc 50%", np.percentile(bedroc_list, 50))
+        print("bedroc 75%", np.percentile(bedroc_list, 75))
+        print("bedroc mean", np.mean(bedroc_list))
+        #print(np.median(auc_list))
+        #print(np.median(ef_list))
+        for key in ef_list:
+            print("ef", key, "25%", np.percentile(ef_list[key], 25))
+            print("ef",key, "50%", np.percentile(ef_list[key], 50))
+            print("ef",key, "75%", np.percentile(ef_list[key], 75))
+            print("ef",key, "mean", np.mean(ef_list[key]))
+        for key in re_list:
+            print("re",key, "25%", np.percentile(re_list[key], 25))
+            print("re",key, "50%", np.percentile(re_list[key], 50))
+            print("re",key, "75%", np.percentile(re_list[key], 75))
+            print("re",key, "mean", np.mean(re_list[key]))
+        return
+    def test_dude_target(self, target, model, **kwargs):
+        data_path = "./data/DUD-E/raw/all/" + target + "/mols.lmdb"
+        mol_dataset = self.load_mols_dataset(data_path, "atoms", "coordinates")
+        num_data = len(mol_dataset)
+        bsz=64
+        print(num_data//bsz)
+        mol_reps = []
+        mol_names = []
+        labels = []
+        # generate mol data
+        mol_data = torch.utils.data.DataLoader(mol_dataset, batch_size=bsz, collate_fn=mol_dataset.collater)
+        for _, sample in enumerate(tqdm(mol_data)):
+            sample = unicore.utils.move_to_cuda(sample)
+            dist = sample["net_input"]["mol_src_distance"]
+            et = sample["net_input"]["mol_src_edge_type"]
+            st = sample["net_input"]["mol_src_tokens"]
+            mol_padding_mask = st.eq(model.mol_model.padding_idx)
+            mol_x = model.mol_model.embed_tokens(st)
+            n_node = dist.size(-1)
+            gbf_feature = model.mol_model.gbf(dist, et)
+            gbf_result = model.mol_model.gbf_proj(gbf_feature)
+            graph_attn_bias = gbf_result
+            graph_attn_bias = graph_attn_bias.permute(0, 3, 1, 2).contiguous()
+            graph_attn_bias = graph_attn_bias.view(-1, n_node, n_node)
+            mol_outputs = model.mol_model.encoder(
+                mol_x, padding_mask=mol_padding_mask, attn_mask=graph_attn_bias
+            )
+            mol_encoder_rep = mol_outputs[0][:,0,:]
+            mol_emb = mol_encoder_rep
+            mol_emb = model.mol_project(mol_encoder_rep)
+            mol_emb = mol_emb / mol_emb.norm(dim=-1, keepdim=True)
+            #print(mol_emb.dtype)
+            mol_emb = mol_emb.detach().cpu().numpy()
+            #print(mol_emb.dtype)
+            mol_reps.append(mol_emb)
+            mol_names.extend(sample["smi_name"])
+            labels.extend(sample["target"].detach().cpu().numpy())
+        mol_reps = np.concatenate(mol_reps, axis=0)
+        labels = np.array(labels, dtype=np.int32)
+        # generate pocket data
+        data_path = "./data/DUD-E/raw/all/" + target + "/pocket.lmdb"
+        pocket_dataset = self.load_pockets_dataset(data_path)
+        pocket_data = torch.utils.data.DataLoader(pocket_dataset, batch_size=bsz, collate_fn=pocket_dataset.collater)
+        pocket_reps = []
+        for _, sample in enumerate(tqdm(pocket_data)):
+            sample = unicore.utils.move_to_cuda(sample)
+            dist = sample["net_input"]["pocket_src_distance"]
+            et = sample["net_input"]["pocket_src_edge_type"]
+            st = sample["net_input"]["pocket_src_tokens"]
+            pocket_padding_mask = st.eq(model.pocket_model.padding_idx)
+            pocket_x = model.pocket_model.embed_tokens(st)
+            n_node = dist.size(-1)
+            gbf_feature = model.pocket_model.gbf(dist, et)
+            gbf_result = model.pocket_model.gbf_proj(gbf_feature)
+            graph_attn_bias = gbf_result
+            graph_attn_bias = graph_attn_bias.permute(0, 3, 1, 2).contiguous()
+            graph_attn_bias = graph_attn_bias.view(-1, n_node, n_node)
+            pocket_outputs = model.pocket_model.encoder(
+                pocket_x, padding_mask=pocket_padding_mask, attn_mask=graph_attn_bias
+            )
+            pocket_encoder_rep = pocket_outputs[0][:,0,:]
+            #pocket_emb = pocket_encoder_rep
+            pocket_emb = model.pocket_project(pocket_encoder_rep)
+            pocket_emb = pocket_emb / pocket_emb.norm(dim=-1, keepdim=True)
+            pocket_emb = pocket_emb.detach().cpu().numpy()
+            pocket_reps.append(pocket_emb)
+        pocket_reps = np.concatenate(pocket_reps, axis=0)
+        print(pocket_reps.shape)
+        res = pocket_reps @ mol_reps.T
+        res_single = res.max(axis=0)
+        auc, bedroc, ef_list, re_list = cal_metrics(labels, res_single, 80.5)
+        print(target)
+        print(np.sum(labels), len(labels)-np.sum(labels))
+        return auc, bedroc, ef_list, re_list, res_single, labels
+    def test_dude(self, model, **kwargs):
+        targets = os.listdir("./data/DUD-E/raw/all/")
+        auc_list = []
+        bedroc_list = []
+        ef_list = []
+        res_list= []
+        labels_list = []
+        re_list = {
+            "0.005": [],
+            "0.01": [],
+            "0.02": [],
+            "0.05": [],
+        }
+        ef_list = {
+            "0.005": [],
+            "0.01": [],
+            "0.02": [],
+            "0.05": [],
+        }
+        for i,target in enumerate(targets):
+            auc, bedroc, ef, re, res_single, labels = self.test_dude_target(target, model)
+            auc_list.append(auc)
+            bedroc_list.append(bedroc)
+            for key in ef:
+                ef_list[key].append(ef[key])
+            for key in re_list:
+                re_list[key].append(re[key])
+            res_list.append(res_single)
+            labels_list.append(labels)
+        res = np.concatenate(res_list, axis=0)
+        labels = np.concatenate(labels_list, axis=0)
+        print("auc mean", np.mean(auc_list))
+        print("bedroc mean", np.mean(bedroc_list))
+        for key in ef_list:
+            print("ef", key, "mean", np.mean(ef_list[key]))
+        for key in re_list:
+            print("re", key, "mean",  np.mean(re_list[key]))
+        # save printed results
+        return
+    def encode_mols_once(self, model, data_path, emb_dir, atoms, coords, **kwargs):
+        # cache path is embdir/data_path.pkl
+        cache_path = os.path.join(emb_dir, data_path.split("/")[-1] + ".pkl")
+        if os.path.exists(cache_path):
+            with open(cache_path, "rb") as f:
+                mol_reps, mol_names = pickle.load(f)
+            return mol_reps, mol_names
+        mol_dataset = self.load_retrieval_mols_dataset(data_path,atoms,coords)
+        mol_reps = []
+        mol_names = []
+        bsz=32
+        mol_data = torch.utils.data.DataLoader(mol_dataset, batch_size=bsz, collate_fn=mol_dataset.collater)
+        for _, sample in enumerate(tqdm(mol_data)):
+            sample = unicore.utils.move_to_cuda(sample)
+            dist = sample["net_input"]["mol_src_distance"]
+            et = sample["net_input"]["mol_src_edge_type"]
+            st = sample["net_input"]["mol_src_tokens"]
+            mol_padding_mask = st.eq(model.mol_model.padding_idx)
+            mol_x = model.mol_model.embed_tokens(st)
+            n_node = dist.size(-1)
+            gbf_feature = model.mol_model.gbf(dist, et)
+            gbf_result = model.mol_model.gbf_proj(gbf_feature)
+            graph_attn_bias = gbf_result
+            graph_attn_bias = graph_attn_bias.permute(0, 3, 1, 2).contiguous()
+            graph_attn_bias = graph_attn_bias.view(-1, n_node, n_node)
+            mol_outputs = model.mol_model.encoder(
+                mol_x, padding_mask=mol_padding_mask, attn_mask=graph_attn_bias
+            )
+            mol_encoder_rep = mol_outputs[0][:,0,:]
+            mol_emb = model.mol_project(mol_encoder_rep)
+            mol_emb = mol_emb / mol_emb.norm(dim=-1, keepdim=True)
+            mol_emb = mol_emb.detach().cpu().numpy()
+            mol_reps.append(mol_emb)
+            mol_names.extend(sample["smi_name"])
+        mol_reps = np.concatenate(mol_reps, axis=0)
+        # save the results
+        with open(cache_path, "wb") as f:
+            pickle.dump([mol_reps, mol_names], f)
+        return mol_reps, mol_names
+    def retrieve_mols(self, model, mol_path, pocket_path, emb_dir, k, **kwargs):
+        os.makedirs(emb_dir, exist_ok=True)
+        mol_reps, mol_names = self.encode_mols_once(model, mol_path, emb_dir,  "atoms", "coordinates")
+        pocket_dataset = self.load_pockets_dataset(pocket_path)
+        pocket_data = torch.utils.data.DataLoader(pocket_dataset, batch_size=16, collate_fn=pocket_dataset.collater)
+        pocket_reps = []
+        pocket_names = []
+        for _, sample in enumerate(tqdm(pocket_data)):
+            sample = unicore.utils.move_to_cuda(sample)
+            dist = sample["net_input"]["pocket_src_distance"]
+            et = sample["net_input"]["pocket_src_edge_type"]
+            st = sample["net_input"]["pocket_src_tokens"]
+            pocket_padding_mask = st.eq(model.pocket_model.padding_idx)
+            pocket_x = model.pocket_model.embed_tokens(st)
+            n_node = dist.size(-1)
+            gbf_feature = model.pocket_model.gbf(dist, et)
+            gbf_result = model.pocket_model.gbf_proj(gbf_feature)
+            graph_attn_bias = gbf_result
+            graph_attn_bias = graph_attn_bias.permute(0, 3, 1, 2).contiguous()
+            graph_attn_bias = graph_attn_bias.view(-1, n_node, n_node)
+            pocket_outputs = model.pocket_model.encoder(
+                pocket_x, padding_mask=pocket_padding_mask, attn_mask=graph_attn_bias
+            )
+            pocket_encoder_rep = pocket_outputs[0][:,0,:]
+            pocket_emb = model.pocket_project(pocket_encoder_rep)
+            pocket_emb = pocket_emb / pocket_emb.norm(dim=-1, keepdim=True)
+            pocket_emb = pocket_emb.detach().cpu().numpy()
+            pocket_reps.append(pocket_emb)
+            pocket_names.extend(sample["pocket_name"])
+        pocket_reps = np.concatenate(pocket_reps, axis=0)
+        res = pocket_reps @ mol_reps.T
+        res = res.max(axis=0)
+        # get top k results
+        top_k = np.argsort(res)[::-1][:k]
+        # return names and scores
+        return [mol_names[i] for i in top_k], res[top_k]