BioBridge-PrimeKG

Class for loading BioBridgePrimeKG dataset.

BioBridgePrimeKG

Bases: 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

Source code in aiagents4pharma/talk2knowledgegraphs/datasets/biobridge_primekg.py

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

__init__(primekg_dir='../../../data/primekg/', local_dir='../../../data/biobridge_primekg/', random_seed=0, n_neg_samples=5)

Constructor for BioBridgePrimeKG class.

Parameters:

Name	Type	Description	Default
primekg_dir	str	The directory of PrimeKG dataset.	'../../../data/primekg/'
local_dir	str	The directory to store the downloaded data.	'../../../data/biobridge_primekg/'
random_seed	int	The random seed value.	0

Source code in aiagents4pharma/talk2knowledgegraphs/datasets/biobridge_primekg.py

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()

_build_full_triplets()

Build the full triplets for BioBridgePrimeKG dataset.

Returns:

Type	Description
DataFrame	The full triplets for BioBridgePrimeKG dataset.
dict	The dictionary of node information.

Source code in aiagents4pharma/talk2knowledgegraphs/datasets/biobridge_primekg.py

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

_build_node_embeddings()

Build the node embeddings for BioBridgePrimeKG dataset.

Returns:

Type	Description
dict	The dictionary of node embeddings.

Source code in aiagents4pharma/talk2knowledgegraphs/datasets/biobridge_primekg.py

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

_build_train_test_split()

Build the train-test split for BioBridgePrimeKG dataset.

Returns:

Type	Description
DataFrame	The train triplets for BioBridgePrimeKG dataset.
DataFrame	The train nodes for BioBridgePrimeKG dataset.
DataFrame	The test triplets for BioBridgePrimeKG dataset.
DataFrame	The test nodes for BioBridgePrimeKG dataset.
DataFrame	The full triplets for BioBridgePrimeKG dataset.

Source code in aiagents4pharma/talk2knowledgegraphs/datasets/biobridge_primekg.py

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

_download_file(remote_url, local_dir, local_filename)

A helper function to download a file from remote URL to the local directory.

Parameters:

Name	Type	Description	Default
remote_url	str	The remote URL of the file to be downloaded.	required
local_dir	str	The local directory to store the downloaded file.	required
local_filename	str	The local filename to store the downloaded file.	required

Source code in aiagents4pharma/talk2knowledgegraphs/datasets/biobridge_primekg.py

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()

_load_data_config()

Load the data config file of BioBridgePrimeKG dataset.

Returns:

Type	Description
dict	The data config file of BioBridgePrimeKG dataset.

Source code in aiagents4pharma/talk2knowledgegraphs/datasets/biobridge_primekg.py

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

_load_primekg()

Private method to load related files of PrimeKG dataset.

Returns:

Type	Description
PrimeKG	The PrimeKG dataset.

Source code in aiagents4pharma/talk2knowledgegraphs/datasets/biobridge_primekg.py

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

get_data_config()

Get the data config file of BioBridgePrimeKG dataset.

Returns:

Type	Description
dict	The data config file of BioBridgePrimeKG dataset.

Source code in aiagents4pharma/talk2knowledgegraphs/datasets/biobridge_primekg.py

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

get_node_embeddings()

Get the node embeddings for BioBridgePrimeKG dataset.

Returns:

Type	Description
dict	The dictionary of node embeddings.

Source code in aiagents4pharma/talk2knowledgegraphs/datasets/biobridge_primekg.py

def get_node_embeddings(self) -> dict:
    """
    Get the node embeddings for BioBridgePrimeKG dataset.

    Returns:
        The dictionary of node embeddings.
    """
    return self.emb_dict

get_node_info_dict()

Get the node information dictionary for BioBridgePrimeKG dataset.

Returns:

Type	Description
dict	The node information dictionary for BioBridgePrimeKG dataset.

Source code in aiagents4pharma/talk2knowledgegraphs/datasets/biobridge_primekg.py

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

get_primekg()

Get the PrimeKG dataset.

Returns:

Type	Description
PrimeKG	The PrimeKG dataset.

Source code in aiagents4pharma/talk2knowledgegraphs/datasets/biobridge_primekg.py

def get_primekg(self) -> PrimeKG:
    """
    Get the PrimeKG dataset.

    Returns:
        The PrimeKG dataset.
    """
    return self.primekg

get_primekg_triplets()

Get the full triplets for BioBridgePrimeKG dataset.

Returns:

Type	Description
DataFrame	The full triplets for BioBridgePrimeKG dataset.

Source code in aiagents4pharma/talk2knowledgegraphs/datasets/biobridge_primekg.py

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

get_train_test_split()

Get the train-test split for BioBridgePrimeKG dataset.

Returns:

Type	Description
dict	The train-test split for BioBridgePrimeKG dataset.

Source code in aiagents4pharma/talk2knowledgegraphs/datasets/biobridge_primekg.py

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
    }

load_data()

Load the BioBridgePrimeKG dataset into pandas DataFrame of nodes and edges.

Parameters:

Name	Type	Description	Default
build_neg_triplest	bool	Whether to build negative triplets.	required
chunk_size	int	The chunk size for negative sampling.	required
n_neg_samples	int	The number of negative samples for negative sampling.	required

Source code in aiagents4pharma/talk2knowledgegraphs/datasets/biobridge_primekg.py

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()

set_random_seed(seed)

Set the random seed for reproducibility.

Parameters:

Name	Type	Description	Default
seed	int	The random seed value.	required

Source code in aiagents4pharma/talk2knowledgegraphs/datasets/biobridge_primekg.py

def set_random_seed(self, seed: int):
    """
    Set the random seed for reproducibility.

    Args:
        seed (int): The random seed value.
    """
    np.random.seed(seed)

setup()

A method to set up the dataset.

Source code in aiagents4pharma/talk2knowledgegraphs/datasets/biobridge_primekg.py

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