Skip to content

Biobridge Primekg

Class for loading BioBridgePrimeKG dataset and PrimeKG nodes and edges data set.

BioBridgePrimeKG

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 vpeleaderboard/data/utils/biobridge_primekg.py 
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
class BioBridgePrimeKG:
    """
    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, cfg: DictConfig):
        """
        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.cfg = cfg
        self.primekg_dir: str = cfg.data.primekg_dir
        self.local_dir: str = cfg.data.biobridge_dir
        self.random_seed: int = cfg.data.random_seed if "random_seed" in cfg.data else 0
        self.n_neg_samples: int = cfg.data.n_neg_samples if "n_neg_samples" in cfg.data else 5
        # 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.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
        os.makedirs(os.path.dirname(self.local_dir), exist_ok=True)

    def _load_primekg(self) -> PrimeKG:
        """
        Private method to load related files of PrimeKG dataset.

        Returns:
            The PrimeKG dataset.
        """
        primekg_data = PrimeKG(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= self.cfg.biobridge_urls.data_config,
            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 = self.cfg.biobridge_urls.embeddings_base
            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 = self.cfg.biobridge_urls.processed_base
            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 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()

    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_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__(cfg)

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.

Source code in vpeleaderboard/data/utils/biobridge_primekg.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
def __init__(self, cfg: DictConfig):
    """
    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.cfg = cfg
    self.primekg_dir: str = cfg.data.primekg_dir
    self.local_dir: str = cfg.data.biobridge_dir
    self.random_seed: int = cfg.data.random_seed if "random_seed" in cfg.data else 0
    self.n_neg_samples: int = cfg.data.n_neg_samples if "n_neg_samples" in cfg.data else 5
    # 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.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
    os.makedirs(os.path.dirname(self.local_dir), exist_ok=True)

_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 vpeleaderboard/data/utils/biobridge_primekg.py 
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
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 = self.cfg.biobridge_urls.processed_base
        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 vpeleaderboard/data/utils/biobridge_primekg.py 
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
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 = self.cfg.biobridge_urls.embeddings_base
        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

_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 vpeleaderboard/data/utils/biobridge_primekg.py 
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
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 vpeleaderboard/data/utils/biobridge_primekg.py 
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
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= self.cfg.biobridge_urls.data_config,
        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 vpeleaderboard/data/utils/biobridge_primekg.py 
63
64
65
66
67
68
69
70
71
72
73
def _load_primekg(self) -> PrimeKG:
    """
    Private method to load related files of PrimeKG dataset.

    Returns:
        The PrimeKG dataset.
    """
    primekg_data = PrimeKG(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 vpeleaderboard/data/utils/biobridge_primekg.py 
278
279
280
281
282
283
284
285
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 vpeleaderboard/data/utils/biobridge_primekg.py 
287
288
289
290
291
292
293
294
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 vpeleaderboard/data/utils/biobridge_primekg.py 
305
306
307
308
309
310
311
312
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 vpeleaderboard/data/utils/biobridge_primekg.py 
269
270
271
272
273
274
275
276
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 vpeleaderboard/data/utils/biobridge_primekg.py 
296
297
298
299
300
301
302
303
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

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 vpeleaderboard/data/utils/biobridge_primekg.py 
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
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()

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 vpeleaderboard/data/utils/biobridge_primekg.py 
260
261
262
263
264
265
266
267
def set_random_seed(self, seed: int):
    """
    Set the random seed for reproducibility.

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