# Copyright (c) 2023 Graphcore Ltd. All rights reserved. # # Includes derived work from https://github.com/KiddoZhu/NBFNet-PyG # Copyright (c) 2021 MilaGraph # Licensed under the MIT License import torch import os import copy from torch.utils.data import DataLoader, IterableDataset, default_collate from torch.nn import functional as F import numpy as np from typing import Optional from torch_geometric.data import Data, InMemoryDataset, download_url from torch_geometric.datasets import RelLinkPredDataset, WordNet18RR import nbfnet_utils class IndRelLinkPredDataset(InMemoryDataset): def __init__( self, root, name, version, add_inverse_train=True, add_inverse_test=True, transform=None, pre_transform=None ): self.name = name self.version = version self.add_inverse_train = add_inverse_train self.add_inverse_test = add_inverse_test assert name in ["FB15k-237"] assert version in ["v1", "v2", "v3", "v4"] self.inv_train_entity_vocab = {} self.inv_test_entity_vocab = {} self.urls = { "FB15k-237": [ f"https://raw.githubusercontent.com/kkteru/grail/master/data/fb237_{version}_ind/train.txt", f"https://raw.githubusercontent.com/kkteru/grail/master/data/fb237_{version}_ind/test.txt", f"https://raw.githubusercontent.com/kkteru/grail/master/data/fb237_{version}/train.txt", f"https://raw.githubusercontent.com/kkteru/grail/master/data/fb237_{version}/valid.txt", f"{RelLinkPredDataset.urls['FB15k-237']}/relations.dict", ], } super().__init__(root, transform, pre_transform) self.data, self.slices = torch.load(self.processed_paths[0]) @property def num_relations(self): return int(self.data.edge_type.max()) + 1 @property def raw_dir(self): return os.path.join(self.root, self.name, self.version, "raw") @property def processed_dir(self): annot = "" if self.add_inverse_train: annot += "_inv_train" if self.add_inverse_test: annot += "_inv_test" return os.path.join(self.root, self.name, self.version, "processed" + annot) @property def processed_file_names(self): return "data.pt" @property def raw_file_names(self): return ["train_ind.txt", "test_ind.txt", "train.txt", "valid.txt", "relations.dict"] def download(self): for url, path in zip(self.urls[self.name], self.raw_paths): download_path = download_url(url, self.raw_dir) os.rename(download_path, path) def process(self): test_files = self.raw_paths[:2] train_files = self.raw_paths[2:4] triplets = [] num_samples = [] with open(os.path.join(self.raw_dir, "relations.dict")) as file: lines = [row.split("\t") for row in file.read().split("\n")[:-1]] inv_relation_vocab = {key: int(value) for value, key in lines} for txt_file in train_files: with open(txt_file, "r") as fin: num_sample = 0 for line in fin: h_token, r_token, t_token = line.strip().split("\t") if h_token not in self.inv_train_entity_vocab: self.inv_train_entity_vocab[h_token] = len(self.inv_train_entity_vocab) h = self.inv_train_entity_vocab[h_token] assert r_token in inv_relation_vocab r = inv_relation_vocab[r_token] if t_token not in self.inv_train_entity_vocab: self.inv_train_entity_vocab[t_token] = len(self.inv_train_entity_vocab) t = self.inv_train_entity_vocab[t_token] triplets.append((h, t, r)) num_sample += 1 num_samples.append(num_sample) for txt_file in test_files: with open(txt_file, "r") as fin: num_sample = 0 for line in fin: h_token, r_token, t_token = line.strip().split("\t") if h_token not in self.inv_test_entity_vocab: self.inv_test_entity_vocab[h_token] = len(self.inv_test_entity_vocab) h = self.inv_test_entity_vocab[h_token] assert r_token in inv_relation_vocab r = inv_relation_vocab[r_token] if t_token not in self.inv_test_entity_vocab: self.inv_test_entity_vocab[t_token] = len(self.inv_test_entity_vocab) t = self.inv_test_entity_vocab[t_token] triplets.append((h, t, r)) num_sample += 1 num_samples.append(num_sample) triplets = torch.tensor(triplets) edge_index = triplets[:, :2].t() edge_type = triplets[:, 2] num_relations = len(inv_relation_vocab) train_fact_slice = slice(None, sum(num_samples[:1])) test_fact_slice = slice(sum(num_samples[:2]), sum(num_samples[:3])) train_fact_index = edge_index[:, train_fact_slice] train_fact_type = edge_type[train_fact_slice] test_fact_index = edge_index[:, test_fact_slice] test_fact_type = edge_type[test_fact_slice] # add flipped triplets for the fact graphs train_fact_index = torch.cat([train_fact_index, train_fact_index.flip(0)], dim=-1) train_fact_type = torch.cat([train_fact_type, train_fact_type + num_relations]) test_fact_index = torch.cat([test_fact_index, test_fact_index.flip(0)], dim=-1) test_fact_type = torch.cat([test_fact_type, test_fact_type + num_relations]) train_slice = slice(None, sum(num_samples[:1])) valid_slice = slice(sum(num_samples[:1]), sum(num_samples[:2])) test_slice = slice(sum(num_samples[:3]), sum(num_samples)) train_target_index = edge_index[:, train_slice] train_target_type = edge_type[train_slice] valid_target_index = edge_index[:, valid_slice] valid_target_type = edge_type[valid_slice] test_target_index = edge_index[:, test_slice] test_target_type = edge_type[test_slice] # add flipped training triples if self.add_inverse_train: train_target_index = torch.cat([train_target_index, train_target_index.flip(0)], dim=-1) train_target_type = torch.cat([train_target_type, train_target_type + num_relations]) # add flipped validation and test triples if self.add_inverse_test: valid_target_index = torch.cat([valid_target_index, valid_target_index.flip(0)], dim=-1) valid_target_type = torch.cat([valid_target_type, valid_target_type + num_relations]) test_target_index = torch.cat([test_target_index, test_target_index.flip(0)], dim=-1) test_target_type = torch.cat([test_target_type, test_target_type + num_relations]) train_data = Data( edge_index=train_fact_index, edge_type=train_fact_type, num_nodes=len(self.inv_train_entity_vocab), target_edge_index=train_target_index, target_edge_type=train_target_type, ) valid_data = Data( edge_index=train_fact_index, edge_type=train_fact_type, num_nodes=len(self.inv_train_entity_vocab), target_edge_index=valid_target_index, target_edge_type=valid_target_type, ) test_data = Data( edge_index=test_fact_index, edge_type=test_fact_type, num_nodes=len(self.inv_test_entity_vocab), target_edge_index=test_target_index, target_edge_type=test_target_type, ) if self.pre_transform is not None: train_data = self.pre_transform(train_data) valid_data = self.pre_transform(valid_data) test_data = self.pre_transform(test_data) torch.save((self.collate([train_data, valid_data, test_data])), self.processed_paths[0]) def __repr__(self): return "%s()" % self.name def build_dataset( name: str, path: str, version: str = None, add_inverse_train: bool = True, add_inverse_test: bool = True ): if name == "FB15k-237": dataset = RelLinkPredDataset(name=name, root=path) data = dataset.data train_target_index = data.train_edge_index train_target_type = data.train_edge_type valid_target_index = data.valid_edge_index valid_target_type = data.valid_edge_type test_target_index = data.test_edge_index test_target_type = data.test_edge_type # add flipped training triples if add_inverse_train: train_target_index = torch.cat([train_target_index, train_target_index.flip(0)], dim=-1) train_target_type = torch.cat([train_target_type, train_target_type + dataset.num_relations // 2]) # add flipped validation and test triples if add_inverse_test: valid_target_index = torch.cat([valid_target_index, valid_target_index.flip(0)], dim=-1) valid_target_type = torch.cat([valid_target_type, valid_target_type + dataset.num_relations // 2]) test_target_index = torch.cat([test_target_index, test_target_index.flip(0)], dim=-1) test_target_type = torch.cat([test_target_type, test_target_type + dataset.num_relations // 2]) train_data = Data( edge_index=data.edge_index, edge_type=data.edge_type, num_nodes=data.num_nodes, target_edge_index=train_target_index, target_edge_type=train_target_type, ) valid_data = Data( edge_index=data.edge_index, edge_type=data.edge_type, num_nodes=data.num_nodes, target_edge_index=valid_target_index, target_edge_type=valid_target_type, ) test_data = Data( edge_index=data.edge_index, edge_type=data.edge_type, num_nodes=data.num_nodes, target_edge_index=test_target_index, target_edge_type=test_target_type, ) dataset.data, dataset.slices = dataset.collate([train_data, valid_data, test_data]) elif name == "WN18RR": dataset = WordNet18RR(root=path) # convert wn18rr into the same format as fb15k-237 data = dataset.data num_nodes = int(data.edge_index.max()) + 1 num_relations = int(data.edge_type.max()) + 1 edge_index = data.edge_index[:, data.train_mask] edge_type = data.edge_type[data.train_mask] edge_index = torch.cat([edge_index, edge_index.flip(0)], dim=-1) edge_type = torch.cat([edge_type, edge_type + num_relations]) if add_inverse_train: raise NotImplementedError if add_inverse_test: raise NotImplementedError train_data = Data( edge_index=edge_index, edge_type=edge_type, num_nodes=num_nodes, target_edge_index=data.edge_index[:, data.train_mask], target_edge_type=data.edge_type[data.train_mask], ) valid_data = Data( edge_index=edge_index, edge_type=edge_type, num_nodes=num_nodes, target_edge_index=data.edge_index[:, data.val_mask], target_edge_type=data.edge_type[data.val_mask], ) test_data = Data( edge_index=edge_index, edge_type=edge_type, num_nodes=num_nodes, target_edge_index=data.edge_index[:, data.test_mask], target_edge_type=data.edge_type[data.test_mask], ) dataset.data, dataset.slices = dataset.collate([train_data, valid_data, test_data]) dataset.num_relations = num_relations * 2 elif name.startswith("Ind"): dataset = IndRelLinkPredDataset( name=name[3:], root=path, version=version, add_inverse_train=add_inverse_train, add_inverse_test=add_inverse_test, ) else: raise ValueError("Unknown dataset `%s`" % name) return dataset class DataWrapper(IterableDataset): def __init__(self, data): super(DataWrapper).__init__() self.data = data def __len__(self): return self.data.n_batches def __iter__(self): return self.data.batches() class NBFData: def __init__( self, data: Data, batch_size: int, is_training: bool, num_relations: Optional[int] = None, num_negatives: Optional[int] = None, check_negatives: Optional[bool] = True, edge_dropout: Optional[float] = 0.0, ): self.batch_size = batch_size self.is_training = is_training self.num_relations = num_relations self.num_negatives = num_negatives self.check_negatives = check_negatives self.edge_dropout = edge_dropout self.data = transpose_dataset(data) self.num_nodes = data.num_nodes self.num_edges = len(data.edge_type) self.graph = torch.cat([self.data.edge_index, self.data.edge_type.unsqueeze(1)], dim=1) self.dataloader = DataLoader( torch.cat([self.data.target_edge_index, self.data.target_edge_type.unsqueeze(1)], dim=1), batch_size, drop_last=False, shuffle=is_training, ) @staticmethod def pad(value, shape, pad_value): dims = list(zip(value.shape, shape)) assert all( actual <= target for actual, target in dims ), f"original shape {value.shape} larger than target {shape}" padding = [] for actual, target in reversed(dims): padding.extend([0, target - actual]) return F.pad(value, padding, value=pad_value) def batches(self): for n, batch in enumerate(self.dataloader): if self.num_negatives: head_id, tail_id, relation_id = nbfnet_utils.negative_sampling( batch=batch, graph=self.graph, num_nodes=self.num_nodes, num_negative=self.num_negatives, strict=self.check_negatives, ) num_negative = self.num_negatives else: head_id, tail_id, relation_id = nbfnet_utils.all_negative(batch, self.num_nodes) num_negative = self.num_nodes - 1 if self.is_training: graph = self.remove_easy_edges(self.graph, head_id, tail_id, relation_id, self.edge_dropout) else: graph = self.graph yield dict( graph=graph + 1, num_nodes=self.num_nodes + 1, head_id=self.pad(head_id + 1, [self.batch_size], 0), tail_id=self.pad(tail_id + 1, [self.batch_size, 1 + num_negative], 0), relation_id=self.pad(relation_id + 1, [self.batch_size], 0), ) def remove_easy_edges( self, graph: torch.Tensor, head_id: torch.Tensor, tail_id: torch.Tensor, relation_id: torch.Tensor, edge_dropout: Optional[float] = 0.0, ): """For a given batch remove direct edges between head and tail entities and their inverse from the graph :param graph: The full set of triples (head, tail, relation). Shape [num_triples, 3] :param head_id: Head of edges to be removed. Shape [batch_size] :param tail_id: Tail of edges to be removed. Shape [batch_size, 1 + num_negatives] :param relation_id: Relation of edges to be removed. Shape [batch_size] :param edge_dropout: Optional dropout rate for remaining edges. :return: graph with direct edges replaced by padding tokens (0, 0, 0) """ if not self.check_negatives: num_tails = tail_id.shape[1] head_id = head_id.repeat(num_tails) tail_id = tail_id.t().flatten() relation_id = relation_id.repeat(num_tails) else: tail_id = tail_id[:, 0] # Also remove inverse edges head_id_ext = torch.cat([head_id, tail_id], dim=-1) tail_id_ext = torch.cat([tail_id, head_id], dim=-1) relation_id_ext = torch.cat([relation_id, relation_id + self.num_relations // 2 % self.num_relations], dim=-1) to_remove = torch.stack([head_id_ext, tail_id_ext, relation_id_ext], -1) id_remove = nbfnet_utils.edge_match(graph, to_remove)[0] if edge_dropout: id_remove = torch.cat([id_remove, torch.randint(0, self.num_edges, [int(self.num_edges * edge_dropout)])]) mask_remove = torch.zeros(graph.shape[0], dtype=torch.bool) mask_remove[id_remove] = True modified_graph = copy.deepcopy(graph) modified_graph[mask_remove, :] = 0 # Padding token 0 for node and relation ids return modified_graph @property def n_items(self): return len(self.data.target_edge_type) @property def n_batches(self): return int(np.ceil(self.n_items / self.batch_size)) def transpose_dataset(data): return Data( edge_index=data.edge_index.t(), edge_type=data.edge_type, num_nodes=data.num_nodes, target_edge_index=data.target_edge_index.t(), target_edge_type=data.target_edge_type, ) def custom_collate(batches: list): out = dict() for key in batches[0].keys(): if key == "num_nodes": out[key] = [batch[key] for batch in batches] else: out[key] = default_collate([batch[key] for batch in batches]) return out