--- a
+++ b/aiagents4pharma/talk2knowledgegraphs/datasets/biobridge_primekg.py
@@ -0,0 +1,562 @@
+"""
+Class for loading BioBridgePrimeKG dataset.
+"""
+
+import os
+import pickle
+import json
+import requests
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+from .dataset import Dataset
+from .primekg import PrimeKG
+
+class BioBridgePrimeKG(Dataset):
+    """
+    Class for loading BioBridgePrimeKG dataset.
+    It downloads the data from the BioBridge repo and stores it in the local directory.
+    The data is then loaded into pandas DataFrame of nodes and edges.
+    This class was adapted from the BioBridge repo:
+    https://github.com/RyanWangZf/BioBridge
+    """
+
+    def __init__(self,
+                 primekg_dir: str = "../../../data/primekg/",
+                 local_dir: str = "../../../data/biobridge_primekg/",
+                 random_seed: int=0,
+                 n_neg_samples: int=5):
+        """
+        Constructor for BioBridgePrimeKG class.
+
+        Args:
+            primekg_dir (str): The directory of PrimeKG dataset.
+            local_dir (str): The directory to store the downloaded data.
+            random_seed (int): The random seed value.
+        """
+        self.name: str = "biobridge_primekg"
+        self.primekg_dir: str = primekg_dir
+        self.local_dir: str = local_dir
+        self.random_seed = random_seed
+        self.n_neg_samples = n_neg_samples
+        # Preselected node types:
+        # protein, molecular function, cellular component, biological process, drug, disease
+        self.preselected_node_types = ["protein", "mf", "cc", "bp", "drug", "disease"]
+        self.node_type_map = {
+            "protein": "gene/protein",
+            "mf": "molecular_function",
+            "cc": "cellular_component",
+            "bp": "biological_process",
+            "drug": "drug",
+            "disease": "disease",
+        }
+
+        # Attributes to store the data
+        self.primekg = None
+        self.primekg_triplets = None
+        self.primekg_triplets_negative = None
+        self.data_config = None
+        self.emb_dict = None
+        self.df_train = None
+        self.df_node_train = None
+        self.df_test = None
+        self.df_node_test = None
+        self.node_info_dict = None
+
+        # Set up the dataset
+        self.setup()
+
+    def setup(self):
+        """
+        A method to set up the dataset.
+        """
+        # Make the directories if it doesn't exist
+        os.makedirs(os.path.dirname(self.primekg_dir), exist_ok=True)
+        os.makedirs(os.path.dirname(self.local_dir), exist_ok=True)
+
+        # Set the random seed
+        self.set_random_seed(self.random_seed)
+
+        # Set SettingWithCopyWarning  warnings to none
+        pd.options.mode.chained_assignment = None
+
+    def _load_primekg(self) -> PrimeKG:
+        """
+        Private method to load related files of PrimeKG dataset.
+
+        Returns:
+            The PrimeKG dataset.
+        """
+        primekg_data = PrimeKG(local_dir=self.primekg_dir)
+        primekg_data.load_data()
+
+        return primekg_data
+
+    def _download_file(self,
+                       remote_url:str,
+                       local_dir: str,
+                       local_filename: str):
+        """
+        A helper function to download a file from remote URL to the local directory.
+
+        Args:
+            remote_url (str): The remote URL of the file to be downloaded.
+            local_dir (str): The local directory to store the downloaded file.
+            local_filename (str): The local filename to store the downloaded file.
+        """
+        # Make the local directory if it does not exist
+        if not os.path.exists(local_dir):
+            os.makedirs(local_dir)
+        # Download the file from remote URL to local directory
+        local_path = os.path.join(local_dir, local_filename)
+        if os.path.exists(local_path):
+            print(f"File {local_filename} already exists in {local_dir}.")
+        else:
+            print(f"Downloading {local_filename} from {remote_url} to {local_dir}...")
+            response = requests.get(remote_url, stream=True, timeout=300)
+            response.raise_for_status()
+            progress_bar = tqdm(
+                total=int(response.headers.get("content-length", 0)),
+                unit="iB",
+                unit_scale=True,
+            )
+            with open(os.path.join(local_dir, local_filename), "wb") as file:
+                for data in response.iter_content(1024):
+                    progress_bar.update(len(data))
+                    file.write(data)
+            progress_bar.close()
+
+    def _load_data_config(self) -> dict:
+        """
+        Load the data config file of BioBridgePrimeKG dataset.
+
+        Returns:
+            The data config file of BioBridgePrimeKG dataset.
+        """
+        # Download the data config file of BioBridgePrimeKG
+        self._download_file(
+            remote_url= ('https://raw.githubusercontent.com/RyanWangZf/BioBridge/'
+                         'refs/heads/main/data/BindData/data_config.json'),
+            local_dir=self.local_dir,
+            local_filename='data_config.json')
+
+        # Load the downloaded data config file
+        with open(os.path.join(self.local_dir, 'data_config.json'), 'r', encoding='utf-8') as f:
+            data_config = json.load(f)
+
+        return data_config
+
+    def _build_node_embeddings(self) -> dict:
+        """
+        Build the node embeddings for BioBridgePrimeKG dataset.
+
+        Returns:
+            The dictionary of node embeddings.
+        """
+        processed_file_path = os.path.join(self.local_dir, "embeddings", "embedding_dict.pkl")
+        if os.path.exists(processed_file_path):
+            # Load the embeddings from the local directory
+            with open(processed_file_path, "rb") as f:
+                emb_dict_all = pickle.load(f)
+        else:
+            # Download the embeddings from the BioBridge repo and further process them
+            # List of embedding source files
+            url = ('https://media.githubusercontent.com/media/RyanWangZf/BioBridge/'
+                   'refs/heads/main/data/embeddings/esm2b_unimo_pubmedbert/')
+            file_list = [f"{n}.pkl" for n in self.preselected_node_types]
+
+            # Download the embeddings
+            for file in file_list:
+                self._download_file(remote_url=os.path.join(url, file),
+                                    local_dir=os.path.join(self.local_dir, "embeddings"),
+                                    local_filename=file)
+
+            # Unified embeddings
+            emb_dict_all = {}
+            for file in file_list:
+                with open(os.path.join(self.local_dir, "embeddings", file), "rb") as f:
+                    emb = pickle.load(f)
+                emb_ar = emb["embedding"]
+                if not isinstance(emb_ar, list):
+                    emb_ar = emb_ar.tolist()
+                emb_dict_all.update(dict(zip(emb["node_index"], emb_ar)))
+
+            # Store embeddings
+            with open(processed_file_path, "wb") as f:
+                pickle.dump(emb_dict_all, f)
+
+        return emb_dict_all
+
+    def _build_full_triplets(self) -> tuple[pd.DataFrame, dict]:
+        """
+        Build the full triplets for BioBridgePrimeKG dataset.
+
+        Returns:
+            The full triplets for BioBridgePrimeKG dataset.
+            The dictionary of node information.
+        """
+        processed_file_path = os.path.join(self.local_dir, "processed", "triplet_full.tsv.gz")
+        if os.path.exists(processed_file_path):
+            # Load the file from the local directory
+            with open(processed_file_path, "rb") as f:
+                primekg_triplets = pd.read_csv(f, sep="\t", compression="gzip", low_memory=False)
+
+            # Load each dataframe in the local directory
+            node_info_dict = {}
+            for i, node_type in enumerate(self.preselected_node_types):
+                with open(os.path.join(self.local_dir, "processed",
+                                       f"{node_type}.csv"), "rb") as f:
+                    df_node = pd.read_csv(f)
+                node_info_dict[self.node_type_map[node_type]] = df_node
+        else:
+            # Download the related files from the BioBridge repo and further process them
+            # List of processed files
+            url = ('https://media.githubusercontent.com/media/RyanWangZf/BioBridge/'
+                   'refs/heads/main/data/Processed/')
+            file_list = ["protein", "molecular", "cellular", "biological", "drug", "disease"]
+
+            # Download the processed files
+            for i, file in enumerate(file_list):
+                self._download_file(remote_url=os.path.join(url, f"{file}.csv"),
+                                    local_dir=os.path.join(self.local_dir, "processed"),
+                                    local_filename=f"{self.preselected_node_types[i]}.csv")
+
+            # Build the node index list
+            node_info_dict = {}
+            node_index_list = []
+            for i, file in enumerate(file_list):
+                df_node = pd.read_csv(os.path.join(self.local_dir, "processed",
+                                                   f"{self.preselected_node_types[i]}.csv"))
+                node_info_dict[self.node_type_map[self.preselected_node_types[i]]] = df_node
+                node_index_list.extend(df_node["node_index"].tolist())
+
+            # Filter the PrimeKG dataset to take into account only the selected node types
+            primekg_triplets = self.primekg.get_edges().copy()
+            primekg_triplets = primekg_triplets[
+                primekg_triplets["head_index"].isin(node_index_list) &\
+                primekg_triplets["tail_index"].isin(node_index_list)
+            ]
+            primekg_triplets = primekg_triplets.reset_index(drop=True)
+
+            # Perform mapping of node types
+            primekg_triplets["head_type"] = primekg_triplets["head_type"].apply(
+                lambda x: self.data_config["node_type"][x]
+            )
+            primekg_triplets["tail_type"] = primekg_triplets["tail_type"].apply(
+                lambda x: self.data_config["node_type"][x]
+            )
+
+            # Perform mapping of relation types
+            primekg_triplets["display_relation"] = primekg_triplets["display_relation"].apply(
+                lambda x: self.data_config["relation_type"][x]
+            )
+
+            # Store the processed triplets
+            primekg_triplets.to_csv(processed_file_path, sep="\t", compression="gzip", index=False)
+
+        return primekg_triplets, node_info_dict
+
+    def _build_train_test_split(self) -> tuple[pd.DataFrame, pd.DataFrame,
+                                               pd.DataFrame, pd.DataFrame, pd.DataFrame]:
+        """
+        Build the train-test split for BioBridgePrimeKG dataset.
+
+        Returns:
+            The train triplets for BioBridgePrimeKG dataset.
+            The train nodes for BioBridgePrimeKG dataset.
+            The test triplets for BioBridgePrimeKG dataset.
+            The test nodes for BioBridgePrimeKG dataset.
+            The full triplets for BioBridgePrimeKG dataset.
+        """
+        if os.path.exists(os.path.join(self.local_dir, "processed",
+                                       "triplet_full_altered.tsv.gz")):
+            # Load each dataframe in the local directory
+            with open(os.path.join(self.local_dir, "processed",
+                                   "triplet_train.tsv.gz"), "rb") as f:
+                df_train = pd.read_csv(f, sep="\t", compression="gzip", low_memory=False)
+
+            with open(os.path.join(self.local_dir, "processed",
+                                   "node_train.tsv.gz"), "rb") as f:
+                df_node_train = pd.read_csv(f, sep="\t", compression="gzip", low_memory=False)
+
+            with open(os.path.join(self.local_dir, "processed",
+                                   "triplet_test.tsv.gz"), "rb") as f:
+                df_test = pd.read_csv(f, sep="\t", compression="gzip", low_memory=False)
+
+            with open(os.path.join(self.local_dir, "processed",
+                                   "node_test.tsv.gz"), "rb") as f:
+                df_node_test = pd.read_csv(f, sep="\t", compression="gzip", low_memory=False)
+
+            with open(os.path.join(self.local_dir, "processed",
+                                   "triplet_full_altered.tsv.gz"), "rb") as f:
+                triplets = pd.read_csv(f, sep="\t", compression="gzip", low_memory=False)
+        else:
+            # Filtering out some nodes in the embedding dictionary
+            triplets = self.primekg_triplets.copy()
+            triplets = triplets[
+                triplets["head_index"].isin(list(self.emb_dict.keys())) &\
+                triplets["tail_index"].isin(list(self.emb_dict.keys()))
+            ].reset_index(drop=True)
+
+            # Perform splitting of the triplets
+            list_split = {
+                "train": [],
+                "test": [],
+            }
+            node_split = {
+                "train": {
+                    "node_index": [],
+                    "node_type": [],
+                },
+                "test": {
+                    "node_index": [],
+                    "node_type": [],
+                }
+            }
+            # Loop over the node types
+            for node_type in triplets["head_type"].unique():
+                df_sub = triplets[triplets["head_type"] == node_type]
+                all_x_indexes = df_sub["head_index"].unique()
+                # By default, we use 90% of the nodes for training and 10% for testing
+                te_x_indexes = np.random.choice(
+                    all_x_indexes, size=int(0.1*len(all_x_indexes)), replace=False
+                )
+                df_subs = {}
+                df_subs["test"] = df_sub[df_sub["head_index"].isin(te_x_indexes)]
+                df_subs["train"] = df_sub[~df_sub["head_index"].isin(te_x_indexes)]
+                list_split["train"].append(df_subs["train"])
+                list_split["test"].append(df_subs["test"])
+
+                # record the split
+                node_index = {}
+                node_index["train"] = df_subs["train"]["head_index"].unique()
+                node_split["train"]["node_index"].extend(node_index["train"].tolist())
+                node_split["train"]["node_type"].extend([node_type]*len(node_index["train"]))
+                node_index["test"] = df_subs["test"]["head_index"].unique()
+                node_split["test"]["node_index"].extend(node_index["test"].tolist())
+                node_split["test"]["node_type"].extend([node_type]*len(node_index["test"]))
+
+                print(f"Number of {node_type} nodes in train: {len(node_index['train'])}")
+                print(f"Number of {node_type} nodes in test: {len(node_index['test'])}")
+
+            # Prepare train and test DataFrames
+            df_train = pd.concat(list_split["train"])
+            df_node_train = pd.DataFrame(node_split["train"])
+            df_test = pd.concat(list_split["test"])
+            df_node_test = pd.DataFrame(node_split["test"])
+
+            # Store each dataframe in the local directory
+            df_train.to_csv(os.path.join(self.local_dir, "processed", "triplet_train.tsv.gz"),
+                            sep="\t", compression="gzip", index=False)
+            df_node_train.to_csv(os.path.join(self.local_dir, "processed", "node_train.tsv.gz"),
+                                sep="\t", compression="gzip", index=False)
+            df_test.to_csv(os.path.join(self.local_dir, "processed", "triplet_test.tsv.gz"),
+                           sep="\t", compression="gzip", index=False)
+            df_node_test.to_csv(os.path.join(self.local_dir, "processed", "node_test.tsv.gz"),
+                                sep="\t", compression="gzip", index=False)
+            # Store altered full triplets as well
+            triplets.to_csv(os.path.join(self.local_dir, "processed",
+                                         "triplet_full_altered.tsv.gz"),
+                            sep="\t", compression="gzip", index=False)
+
+        return df_train, df_node_train, df_test, df_node_test, triplets
+
+    # def _negative_sampling(self,
+    #                        batch_df: pd.DataFrame,
+    #                        process_index: int,
+    #                        index_map: dict,
+    #                        node_train_dict: dict) -> pd.DataFrame:
+    #     """
+    #     A helper function to perform negative sampling for a batch of triplets.
+    #     """
+    #     negative_y_index_list = []
+    #     for _, row in tqdm(batch_df.iterrows(),
+    #                        total=batch_df.shape[0],
+    #                        desc=f"Process {process_index}"):
+    #         x_index = row['head_index']
+    #         # y_index = row['y_index']
+    #         y_index_type = row['tail_type']
+    #         paired_y_index_list = index_map[x_index]
+
+    #         # sample a list of negative y_index
+    #         node_train_sub = node_train_dict[y_index_type]
+    #         negative_y_index = node_train_sub[
+    #             ~node_train_sub['node_index'].isin(paired_y_index_list)
+    #         ]['node_index'].sample(self.n_neg_samples).tolist()
+    #         negative_y_index_list.append(negative_y_index)
+
+    #     batch_df.loc[:, 'negative_tail_index'] = negative_y_index_list
+    #     return batch_df
+
+    # def _build_negative_triplets(self,
+    #                              chunk_size: int=100000,
+    #                              n_neg_samples: int=10):
+    #     """
+    #     Build the negative triplets for BioBridgePrimeKG dataset.
+    #     """
+    #     processed_file_path = os.path.join(self.local_dir,
+    #                                        "processed",
+    #                                        "triplet_train_negative.tsv.gz")
+    #     if os.path.exists(processed_file_path):
+    #         # Load the negative triplets from the local directory
+    #         with open(processed_file_path, "rb") as f:
+    #             triplets_negative = pd.read_csv(f, sep="\t", compression="gzip", low_memory=False)
+    #     else:
+    #         # Set the number samples for negative sampling
+    #         self.n_neg_samples = n_neg_samples
+
+    #         # Split node list by type
+    #         node_train_dict = {}
+    #         type_list = self.df_node_train['node_type'].unique()
+    #         for node_type in type_list:
+    #             node_train_dict[node_type] = self.df_node_train[
+    #                 self.df_node_train['node_type'] == node_type
+    #             ].reset_index(drop=True)
+
+    #         # create an index mapping from x_index to y_index
+    #         index_map = self.df_train[
+    #             ['head_index', 'tail_index']
+    #         ].drop_duplicates().groupby('head_index').agg(list).to_dict()['tail_index']
+
+    #         # Negative sampling
+    #         batch_df_list = []
+    #         for i in tqdm(range(0, self.df_train.shape[0], chunk_size)):
+    #             batch_df_list.append(self.df_train.iloc[i:i+chunk_size])
+    #         # Process negative sampling
+    #         results = [
+    #             self._negative_sampling(batch_df,
+    #                                     num_piece,
+    #                                     index_map,
+    #                                     node_train_dict)
+    #                                     for num_piece, batch_df in enumerate(batch_df_list)
+    #         ]
+
+    #         # Store the negative triplets
+    #         triplets_negative = pd.concat(results, axis=0)
+    #         triplets_negative.to_csv(processed_file_path,
+    #                                  sep="\t", compression="gzip", index=False)
+
+    #     # Set attribute
+    #     self.primekg_triplets_negative = triplets_negative
+
+    #     return triplets_negative
+
+    # def load_data(self,
+    #               build_neg_triplest: bool= False,
+    #               chunk_size: int=100000,
+    #               n_neg_samples: int=10):
+
+    def load_data(self):
+        """
+        Load the BioBridgePrimeKG dataset into pandas DataFrame of nodes and edges.
+
+        Args:
+            build_neg_triplest (bool): Whether to build negative triplets.
+            chunk_size (int): The chunk size for negative sampling.
+            n_neg_samples (int): The number of negative samples for negative sampling.
+        """
+        # Load PrimeKG dataset
+        print("Loading PrimeKG dataset...")
+        self.primekg = self._load_primekg()
+
+        # Load data config file of BioBridgePrimeKG
+        print("Loading data config file of BioBridgePrimeKG...")
+        self.data_config = self._load_data_config()
+
+        # Build node embeddings
+        print("Building node embeddings...")
+        self.emb_dict = self._build_node_embeddings()
+
+        # Build full triplets
+        print("Building full triplets...")
+        self.primekg_triplets, self.node_info_dict = self._build_full_triplets()
+
+        # Build train-test split
+        print("Building train-test split...")
+        self.df_train, self.df_node_train, self.df_test, self.df_node_test, self.primekg_triplets =\
+        self._build_train_test_split()
+
+        # if build_neg_triplest:
+        #     # Build negative triplets
+        #     print("Building negative triplets...")
+        #     self.primekg_triplets_negative = self._build_negative_triplets(
+        #         chunk_size=chunk_size,
+        #         n_neg_samples=n_neg_samples
+        #     )
+
+    def set_random_seed(self, seed: int):
+        """
+        Set the random seed for reproducibility.
+
+        Args:
+            seed (int): The random seed value.
+        """
+        np.random.seed(seed)
+
+    def get_primekg(self) -> PrimeKG:
+        """
+        Get the PrimeKG dataset.
+
+        Returns:
+            The PrimeKG dataset.
+        """
+        return self.primekg
+
+    def get_data_config(self) -> dict:
+        """
+        Get the data config file of BioBridgePrimeKG dataset.
+
+        Returns:
+            The data config file of BioBridgePrimeKG dataset.
+        """
+        return self.data_config
+
+    def get_node_embeddings(self) -> dict:
+        """
+        Get the node embeddings for BioBridgePrimeKG dataset.
+
+        Returns:
+            The dictionary of node embeddings.
+        """
+        return self.emb_dict
+
+    def get_primekg_triplets(self) -> pd.DataFrame:
+        """
+        Get the full triplets for BioBridgePrimeKG dataset.
+
+        Returns:
+            The full triplets for BioBridgePrimeKG dataset.
+        """
+        return self.primekg_triplets
+
+    # def get_primekg_triplets_negative(self) -> pd.DataFrame:
+    #     """
+    #     Get the negative triplets for BioBridgePrimeKG dataset.
+
+    #     Returns:
+    #         The negative triplets for BioBridgePrimeKG dataset.
+    #     """
+    #     return self.primekg_triplets_negative
+
+    def get_train_test_split(self) -> dict:
+        """
+        Get the train-test split for BioBridgePrimeKG dataset.
+
+        Returns:
+            The train-test split for BioBridgePrimeKG dataset.
+        """
+        return {
+            "train": self.df_train,
+            "node_train": self.df_node_train,
+            "test": self.df_test,
+            "node_test": self.df_node_test
+        }
+
+    def get_node_info_dict(self) -> dict:
+        """
+        Get the node information dictionary for BioBridgePrimeKG dataset.
+
+        Returns:
+            The node information dictionary for BioBridgePrimeKG dataset.
+        """
+        return self.node_info_dict