import os.path
import dgl
import math
import re
from copy import deepcopy
import numpy as np
import torch as th
import itertools
import random
from random import shuffle, choice
from collections import Counter
from os.path import join as joinpath
from os.path import isfile
from dgl.data.knowledge_graph import load_data
from . import BaseDataset, register_dataset
from . import AcademicDataset, HGBDataset, OHGBDataset, NBF_Dataset
from ..utils import add_reverse_edges
from collections import defaultdict
import os
from scipy.sparse import csr_matrix
__all__ = ['LinkPredictionDataset', 'HGB_LinkPrediction']
[docs]
@register_dataset('link_prediction')
class LinkPredictionDataset(BaseDataset):
"""
metric: Accuracy, multi-label f1 or multi-class f1. Default: `accuracy`
"""
def __init__(self, *args, **kwargs):
super(LinkPredictionDataset, self).__init__(*args, **kwargs)
self.target_link = None
self.target_link_r = None
def get_split(self, val_ratio=0.1, test_ratio=0.2):
"""
Get subgraphs for train, valid and test.
Generally, the original will have train_mask and test_mask in edata, or we will split it automatically.
If the original graph do not has the train_mask in edata, we default that there is no valid_mask and test_mask.
So we will split the edges of the original graph into train/valid/test 0.7/0.1/0.2.
The dataset has not validation_mask, so we split train edges randomly.
Parameters
----------
val_ratio : int
The ratio of validation. Default: 0.1
test_ratio : int
The ratio of test. Default: 0.2
Returns
-------
train_hg
"""
val_edge_dict = {}
test_edge_dict = {}
out_ntypes = []
train_graph = self.g
for i, etype in enumerate(self.target_link):
num_edges = self.g.num_edges(etype)
if 'train_mask' not in self.g.edges[etype].data:
"""
split edges into train/valid/test.
"""
random_int = th.randperm(num_edges)
val_index = random_int[:int(num_edges * val_ratio)]
val_edge = self.g.find_edges(val_index, etype)
test_index = random_int[int(num_edges * val_ratio):int(num_edges * (test_ratio + val_ratio))]
test_edge = self.g.find_edges(test_index, etype)
val_edge_dict[etype] = val_edge
test_edge_dict[etype] = test_edge
out_ntypes.append(etype[0])
out_ntypes.append(etype[2])
train_graph = dgl.remove_edges(train_graph, th.cat((val_index, test_index)), etype)
# train_graph = dgl.remove_edges(train_graph, val_index, etype)
if self.target_link_r is None:
pass
else:
reverse_edge = self.target_link_r[i]
train_graph = dgl.remove_edges(train_graph, th.arange(train_graph.num_edges(reverse_edge)),
reverse_edge)
edges = train_graph.edges(etype=etype)
train_graph = dgl.add_edges(train_graph, edges[1], edges[0], etype=reverse_edge)
else:
if 'valid_mask' not in self.g.edges[etype].data:
train_idx = self.g.edges[etype].data['train_mask']
random_int = th.randperm(int(train_idx.sum( )))
val_index = random_int[:int(train_idx.sum( ) * val_ratio)]
val_edge = self.g.find_edges(val_index, etype)
else:
val_mask = self.g.edges[etype].data['valid_mask'].squeeze( )
val_index = th.nonzero(val_mask).squeeze( )
val_edge = self.g.find_edges(val_index, etype)
test_mask = self.g.edges[etype].data['test_mask'].squeeze( )
test_index = th.nonzero(test_mask).squeeze( )
test_edge = self.g.find_edges(test_index, etype)
val_edge_dict[etype] = val_edge
test_edge_dict[etype] = test_edge
out_ntypes.append(etype[0])
out_ntypes.append(etype[2])
# self.val_label = train_graph.edges[etype[1]].data['label'][val_index]
self.test_label = train_graph.edges[etype[1]].data['label'][test_index]
train_graph = dgl.remove_edges(train_graph, th.cat((val_index, test_index)), etype)
# train_graph = dgl.remove_edges(train_graph, th.cat((val_index, test_index)), 'item-user')
self.out_ntypes = set(out_ntypes)
val_graph = dgl.heterograph(val_edge_dict,
{ntype: self.g.number_of_nodes(ntype) for ntype in set(out_ntypes)})
test_graph = dgl.heterograph(test_edge_dict,
{ntype: self.g.number_of_nodes(ntype) for ntype in set(out_ntypes)})
# todo: val/test negative graphs should be created before training rather than
# create them dynamically in every evaluation.
return train_graph, val_graph, test_graph, None, None
@register_dataset('demo_link_prediction')
class Test_LinkPrediction(LinkPredictionDataset):
def __init__(self, dataset_name):
super(Test_LinkPrediction, self).__init__( )
self.g = self.load_HIN('./openhgnn/debug/data.bin')
self.target_link = 'user-item'
self.has_feature = False
self.meta_paths_dict = None
self.preprocess( )
# self.generate_negative()
def preprocess(self):
test_mask = self.g.edges[self.target_link].data['test_mask']
index = th.nonzero(test_mask).squeeze( )
self.test_edge = self.g.find_edges(index, self.target_link)
self.pos_test_graph = dgl.heterograph({('user', 'user-item', 'item'): self.test_edge},
{ntype: self.g.number_of_nodes(ntype) for ntype in ['user', 'item']})
self.g.remove_edges(index, self.target_link)
self.g.remove_edges(index, 'item-user')
self.neg_test_graph, _ = dgl.load_graphs('./openhgnn/debug/neg.bin')
self.neg_test_graph = self.neg_test_graph[0]
return
def generate_negative(self):
k = 99
e = self.pos_test_graph.edges( )
neg_src = []
neg_dst = []
for i in range(self.pos_test_graph.number_of_edges( )):
src = e[0][i]
exp = self.pos_test_graph.successors(src)
dst = th.randint(high=self.g.number_of_nodes('item'), size=(k,))
for d in range(len(dst)):
while dst[d] in exp:
dst[d] = th.randint(high=self.g.number_of_nodes('item'), size=(1,))
src = src.repeat_interleave(k)
neg_src.append(src)
neg_dst.append(dst)
neg_edge = (th.cat(neg_src), th.cat(neg_dst))
neg_graph = dgl.heterograph({('user', 'user-item', 'item'): neg_edge},
{ntype: self.g.number_of_nodes(ntype) for ntype in ['user', 'item']})
dgl.save_graphs('./openhgnn/debug/neg.bin', neg_graph)
@register_dataset('hin_link_prediction')
class HIN_LinkPrediction(LinkPredictionDataset):
def __init__(self, dataset_name, *args, **kwargs):
super(HIN_LinkPrediction, self).__init__(*args, **kwargs)
self.g = self.load_HIN(dataset_name)
def load_link_pred(self, path):
u_list = []
v_list = []
label_list = []
with open(path) as f:
for i in f.readlines( ):
u, v, label = i.strip( ).split(', ')
u_list.append(int(u))
v_list.append(int(v))
label_list.append(int(label))
return u_list, v_list, label_list
def load_HIN(self, dataset_name):
self.dataset_name = dataset_name
if dataset_name == 'academic4HetGNN':
# which is used in HetGNN
dataset = AcademicDataset(name='academic4HetGNN', raw_dir='')
g = dataset[0].long( )
self.train_batch = self.load_link_pred('./openhgnn/dataset/' + dataset_name + '/a_a_list_train.txt')
self.test_batch = self.load_link_pred('./openhgnn/dataset/' + dataset_name + '/a_a_list_test.txt')
self.category = 'author'
elif dataset_name == 'Book-Crossing':
g, _ = dgl.load_graphs('./openhgnn/dataset/book_graph.bin')
g = g[0]
self.target_link = [('user', 'user-item', 'item')]
self.node_type = ['user', 'item']
elif dataset_name == 'amazon4SLICE':
dataset = AcademicDataset(name='amazon4SLICE', raw_dir='')
g = dataset[0].long( )
# self.target_link = [('product', 'product-1-product', 'product'),
# ('product', 'product-2-product', 'product')]
self.target_link = [('product', 'product-1-product', 'product')]
elif dataset_name == 'MTWM':
dataset = AcademicDataset(name='MTWM', raw_dir='')
g = dataset[0].long( )
g = add_reverse_edges(g)
self.target_link = [('user', 'user-buy-spu', 'spu')]
self.target_link_r = [('spu', 'user-buy-spu-rev', 'user')]
self.meta_paths_dict = {
'UPU1': [('user', 'user-buy-poi', 'poi'), ('poi', 'user-buy-poi-rev', 'user')],
'UPU2': [('user', 'user-click-poi', 'poi'), ('poi', 'user-click-poi-rev', 'user')],
'USU': [('user', 'user-buy-spu', 'spu'), ('spu', 'user-buy-spu-rev', 'user')],
'UPSPU1': [('user', 'user-buy-poi', 'poi'), ('poi', 'poi-contain-spu', 'spu'),
('spu', 'poi-contain-spu-rev', 'poi'), ('poi', 'user-buy-poi-rev', 'user')
],
'UPSPU2': [
('user', 'user-click-poi', 'poi'), ('poi', 'poi-contain-spu', 'spu'),
('spu', 'poi-contain-spu-rev', 'poi'), ('poi', 'user-click-poi-rev', 'user')
]
}
self.node_type = ['user', 'spu']
elif dataset_name == 'HGBl-ACM':
dataset = HGBDataset(name='HGBn-ACM', raw_dir='')
g = dataset[0].long( )
self.has_feature = True
self.target_link = [('paper', 'paper-ref-paper', 'paper')]
self.node_type = ['author', 'paper', 'subject', 'term']
self.target_link_r = [('paper', 'paper-cite-paper', 'paper')]
self.meta_paths_dict = {'PAP': [('paper', 'paper-author', 'author'), ('author', 'author-paper', 'paper')],
'PSP': [('paper', 'paper-subject', 'subject'),
('subject', 'subject-paper', 'paper')],
'PcPAP': [('paper', 'paper-cite-paper', 'paper'),
('paper', 'paper-author', 'author'),
('author', 'author-paper', 'paper')],
'PcPSP': [('paper', 'paper-cite-paper', 'paper'),
('paper', 'paper-subject', 'subject'),
('subject', 'subject-paper', 'paper')],
'PrPAP': [('paper', 'paper-ref-paper', 'paper'),
('paper', 'paper-author', 'author'),
('author', 'author-paper', 'paper')],
'PrPSP': [('paper', 'paper-ref-paper', 'paper'),
('paper', 'paper-subject', 'subject'),
('subject', 'subject-paper', 'paper')]
}
elif dataset_name == 'HGBl-DBLP':
dataset = HGBDataset(name='HGBn-DBLP', raw_dir='')
g = dataset[0].long( )
self.has_feature = True
self.target_link = [('author', 'author-paper', 'paper')]
self.node_type = ['author', 'paper', 'venue', 'term']
self.target_link_r = [('paper', 'paper-author', 'author')]
self.meta_paths_dict = {'APA': [('author', 'author-paper', 'paper'), ('paper', 'paper-author', 'author')],
'APTPA': [('author', 'author-paper', 'paper'), ('paper', 'paper-term', 'term'),
('term', 'term-paper', 'paper'), ('paper', 'paper-author', 'author')],
'APVPA': [('author', 'author-paper', 'paper'), ('paper', 'paper-venue', 'venue'),
('venue', 'venue-paper', 'paper'), ('paper', 'paper-author', 'author')],
'PAP': [('paper', 'paper-author', 'author'), ('author', 'author-paper', 'paper')],
'PTP': [('paper', 'paper-term', 'term'), ('term', 'term-paper', 'paper')],
'PVP': [('paper', 'paper-venue', 'venue'), ('venue', 'venue-paper', 'paper')],
}
elif dataset_name == 'HGBl-IMDB':
dataset = HGBDataset(name='HGBn-IMDB', raw_dir='')
g = dataset[0].long( )
self.has_feature = True
# self.target_link = [('author', 'author-paper', 'paper')]
# self.node_type = ['author', 'paper', 'subject', 'term']
# self.target_link_r = [('paper', 'paper-author', 'author')]
self.target_link = [('actor', 'actor->movie', 'movie')]
self.node_type = ['actor', 'director', 'keyword', 'movie']
self.target_link_r = [('movie', 'movie->actor', 'actor')]
self.meta_paths_dict = {
'MAM': [('movie', 'movie->actor', 'actor'), ('actor', 'actor->movie', 'movie')],
'MDM': [('movie', 'movie->director', 'director'), ('director', 'director->movie', 'movie')],
'MKM': [('movie', 'movie->keyword', 'keyword'), ('keyword', 'keyword->movie', 'movie')],
# 'DMD': [('director', 'director->movie', 'movie'), ('movie', 'movie->director', 'director')],
# 'DMAMD': [('director', 'director->movie', 'movie'), ('movie', 'movie->actor', 'actor'),
# ('actor', 'actor->movie', 'movie'), ('movie', 'movie->director', 'director')],
'AMA': [('actor', 'actor->movie', 'movie'), ('movie', 'movie->actor', 'actor')],
'AMDMA': [('actor', 'actor->movie', 'movie'), ('movie', 'movie->director', 'director'),
('director', 'director->movie', 'movie'), ('movie', 'movie->actor', 'actor')]
}
return g
def get_split(self, val_ratio=0.1, test_ratio=0.2):
if self.dataset_name == 'academic4HetGNN':
return None, None, None, None, None
else:
return super(HIN_LinkPrediction, self).get_split(val_ratio, test_ratio)
@register_dataset('HGBl_link_prediction')
class HGB_LinkPrediction(LinkPredictionDataset):
r"""
The HGB dataset will be used in task *link prediction*.
Dataset Name :
HGBn-amazon/HGBn-LastFM/HGBn-PubMed
So if you want to get more information, refer to
`HGB datasets <https://github.com/THUDM/HGB>`_
Attributes
-----------
has_feature : bool
Whether the dataset has feature. Except HGBl-LastFM, others have features.
target_link : list of tuple[canonical_etypes]
The etypes of test link. HGBl-amazon has two etypes of test link. other has only one.
"""
def __init__(self, dataset_name, *args, **kwargs):
super(HGB_LinkPrediction, self).__init__(*args, **kwargs)
self.dataset_name = dataset_name
self.target_link_r = None
if dataset_name == 'HGBl-amazon':
dataset = HGBDataset(name=dataset_name, raw_dir='')
g = dataset[0].long( )
self.has_feature = False
self.target_link = [('product', 'product-product-0', 'product'),
('product', 'product-product-1', 'product')]
self.target_link_r = None
self.link = [0, 1]
self.node_type = ["product"]
self.test_edge_type = {'product-product-0': 0, 'product-product-1': 1}
self.meta_paths_dict = {
'P0P': [('product', 'product-product-0', 'product'), ('product', 'product-product-1', 'product')],
'P1P': [('product', 'product-product-1', 'product'), ('product', 'product-product-0', 'product')]
}
elif dataset_name == 'HGBl-LastFM':
dataset = HGBDataset(name=dataset_name, raw_dir='')
g = dataset[0].long( )
self.has_feature = False
self.target_link = [('user', 'user-artist', 'artist')]
self.node_type = ['user', 'artist', 'tag']
self.test_edge_type = {'user-artist': 0}
g = add_reverse_edges(g)
self.target_link_r = [('artist', 'user-artist-rev', 'user')]
self.meta_paths_dict = {'UU': [('user', 'user-user', 'user')],
'UAU': [('user', 'user-artist', 'artist'), ('artist', 'user-artist-rev', 'user')],
'UATAU': [('user', 'user-artist', 'artist'), ('artist', 'artist-tag', 'tag'),
('tag', 'artist-tag-rev', 'artist'),
('artist', 'user-artist-rev', 'user')],
'AUA': [('artist', 'user-artist-rev', 'user'), ('user', 'user-artist', 'artist')],
'ATA': [('artist', 'artist-tag', 'tag'), ('tag', 'artist-tag-rev', 'artist')]
}
elif dataset_name == 'HGBl-PubMed':
dataset = HGBDataset(name=dataset_name, raw_dir='')
g = dataset[0].long( )
self.has_feature = True
self.target_link = [('1', '1_to_1', '1')]
self.node_type = ['0', '1', '2', '3']
self.test_edge_type = {'1_to_1': 2}
g = add_reverse_edges(g)
self.target_link_r = [('1', '1_to_1-rev', '1')]
self.meta_paths_dict = {'101': [('1', '0_to_1-rev', '0'), ('0', '0_to_1', '1')],
'111': [('1', '1_to_1', '1'), ('1', '1_to_1-rev', '1')],
'121': [('1', '2_to_1-rev', '2'), ('2', '2_to_1', '1')],
'131': [('1', '3_to_1-rev', '3'), ('3', '3_to_1', '1')]
}
self.g = g
self.shift_dict = self.calculate_node_shift( )
def load_link_pred(self, path):
return
def calculate_node_shift(self):
node_shift_dict = {}
count = 0
for type in self.node_type:
node_shift_dict[type] = count
count += self.g.num_nodes(type)
return node_shift_dict
def get_split(self):
r"""
Get graphs for train, valid or test.
The dataset has not validation_mask, so we split train edges randomly.
"""
val_edge_dict = {}
test_edge_dict = {}
out_ntypes = []
train_graph = self.g
val_ratio = 0.1
for i, etype in enumerate(self.target_link):
train_mask = self.g.edges[etype].data['train_mask'].squeeze( )
train_index = th.nonzero(train_mask).squeeze( )
random_int = th.randperm(len(train_index))[:int(len(train_index) * val_ratio)]
val_index = train_index[random_int]
val_edge = self.g.find_edges(val_index, etype)
test_mask = self.g.edges[etype].data['test_mask'].squeeze( )
test_index = th.nonzero(test_mask).squeeze( )
test_edge = self.g.find_edges(test_index, etype)
val_edge_dict[etype] = val_edge
test_edge_dict[etype] = test_edge
out_ntypes.append(etype[0])
out_ntypes.append(etype[2])
train_graph = dgl.remove_edges(train_graph, th.cat((val_index, test_index)), etype)
if self.target_link_r is None:
pass
else:
train_graph = dgl.remove_edges(train_graph, th.cat((val_index, test_index)), self.target_link_r[i])
self.out_ntypes = set(out_ntypes)
val_graph = dgl.heterograph(val_edge_dict,
{ntype: self.g.number_of_nodes(ntype) for ntype in set(out_ntypes)})
test_graph = dgl.heterograph(test_edge_dict,
{ntype: self.g.number_of_nodes(ntype) for ntype in set(out_ntypes)})
return train_graph, val_graph, test_graph, None, None
def save_results(self, hg, score, file_path):
with hg.local_scope( ):
src_list = []
dst_list = []
edge_type_list = []
for etype in hg.canonical_etypes:
edges = hg.edges(etype=etype)
src_id = edges[0] + self.shift_dict[etype[0]]
dst_id = edges[1] + self.shift_dict[etype[2]]
src_list.append(src_id)
dst_list.append(dst_id)
edge_type_list.append(th.full((src_id.shape[0],), self.test_edge_type[etype[1]]))
src_list = th.cat(src_list)
dst_list = th.cat(dst_list)
edge_type_list = th.cat(edge_type_list)
with open(file_path, "w") as f:
for l, r, edge_type, c in zip(src_list, dst_list, edge_type_list, score):
f.write(f"{l}\t{r}\t{edge_type}\t{round(float(c), 4)}\n")
@register_dataset('ohgb_link_prediction')
class OHGB_LinkPrediction(LinkPredictionDataset):
def __init__(self, dataset_name, *args, **kwargs):
super(OHGB_LinkPrediction, self).__init__(*args, **kwargs)
self.dataset_name = dataset_name
self.has_feature = True
if dataset_name == 'ohgbl-MTWM':
dataset = OHGBDataset(name=dataset_name, raw_dir='')
g = dataset[0].long( )
self.target_link = [('user', 'user-buy-spu', 'spu')]
self.target_link_r = [('spu', 'user-buy-spu-rev', 'user')]
self.node_type = ['user', 'spu']
elif dataset_name == 'ohgbl-yelp1':
dataset = OHGBDataset(name=dataset_name, raw_dir='')
g = dataset[0].long( )
self.target_link = [('user', 'user-buy-business', 'business')]
self.target_link_r = [('business', 'user-buy-business-rev', 'user')]
elif dataset_name == 'ohgbl-yelp2':
dataset = OHGBDataset(name=dataset_name, raw_dir='')
g = dataset[0].long( )
self.target_link = [('business', 'described-with', 'phrase')]
self.target_link_r = [('business', 'described-with-rev', 'phrase')]
elif dataset_name == 'ohgbl-Freebase':
dataset = OHGBDataset(name=dataset_name, raw_dir='')
g = dataset[0].long( )
self.target_link = [('BOOK', 'BOOK-and-BOOK', 'BOOK')]
self.target_link_r = [('BOOK', 'BOOK-and-BOOK-rev', 'BOOK')]
self.g = g
def build_graph_from_triplets(num_nodes, num_rels, triplets):
""" Create a DGL graph. The graph is bidirectional because RGCN authors
use reversed relations.
This function also generates edge type and normalization factor
(reciprocal of node incoming degree)
"""
g = dgl.graph(([], []))
g.add_nodes(num_nodes)
src, rel, dst = triplets
src, dst = np.concatenate((src, dst)), np.concatenate((dst, src))
rel = np.concatenate((rel, rel + num_rels))
edges = sorted(zip(dst, src, rel))
dst, src, rel = np.array(edges).transpose( )
g.add_edges(src, dst)
norm = comp_deg_norm(g)
print("# nodes: {}, # edges: {}".format(num_nodes, len(src)))
return g, rel.astype('int64'), norm.astype('int64')
def comp_deg_norm(g):
g = g.local_var( )
in_deg = g.in_degrees(range(g.number_of_nodes( ))).float( ).numpy( )
norm = 1.0 / in_deg
norm[np.isinf(norm)] = 0
return norm
@register_dataset('kg_sub_link_prediction')
class KG_RedDataset(LinkPredictionDataset):
def __init__(self, dataset_name, *args, **kwargs):
super(KG_RedDataset, self).__init__(*args, **kwargs)
self.trans_dir = os.path.join('openhgnn/dataset/data', dataset_name)
self.ind_dir = self.trans_dir + '_ind'
folder = os.path.exists(self.trans_dir)
if not folder:
os.makedirs(self.trans_dir)
url = "https://s3.cn-north-1.amazonaws.com.cn/dgl-data/dataset/openhgnn/fb237_v1.zip"
response = requests.get(url)
with zipfile.ZipFile(io.BytesIO(response.content)) as myzip:
myzip.extractall(self.trans_dir)
print("--- download data ---")
else:
print("--- There is data! ---")
folder = os.path.exists(self.ind_dir)
if not folder:
os.makedirs(self.ind_dir)
# 下载数据
url = "https://s3.cn-north-1.amazonaws.com.cn/dgl-data/dataset/openhgnn/fb237_v1_ind.zip"
response = requests.get(url)
with zipfile.ZipFile(io.BytesIO(response.content)) as myzip:
myzip.extractall(self.ind_dir)
print("--- download data ---")
else:
print("--- There is data! ---")
with open(os.path.join(self.trans_dir, 'entities.txt')) as f:
self.entity2id = dict()
for line in f:
entity, eid = line.strip().split()
self.entity2id[entity] = int(eid)
with open(os.path.join(self.trans_dir, 'relations.txt')) as f:
self.relation2id = dict()
id2relation = []
for line in f:
relation, rid = line.strip().split()
self.relation2id[relation] = int(rid)
id2relation.append(relation)
with open(os.path.join(self.ind_dir, 'entities.txt')) as f:
self.entity2id_ind = dict()
for line in f:
entity, eid = line.strip().split()
self.entity2id_ind[entity] = int(eid)
for i in range(len(self.relation2id)):
id2relation.append(id2relation[i] + '_inv')
id2relation.append('idd')
self.id2relation = id2relation
self.n_ent = len(self.entity2id)
self.n_rel = len(self.relation2id)
self.n_ent_ind = len(self.entity2id_ind)
self.tra_train = self.read_triples(self.trans_dir, 'train.txt')
self.tra_valid = self.read_triples(self.trans_dir, 'valid.txt')
self.tra_test = self.read_triples(self.trans_dir, 'test.txt')
self.ind_train = self.read_triples(self.ind_dir, 'train.txt', 'inductive')
self.ind_valid = self.read_triples(self.ind_dir, 'valid.txt', 'inductive')
self.ind_test = self.read_triples(self.ind_dir, 'test.txt', 'inductive')
self.val_filters = self.get_filter('valid')
self.tst_filters = self.get_filter('test')
for filt in self.val_filters:
self.val_filters[filt] = list(self.val_filters[filt])
for filt in self.tst_filters:
self.tst_filters[filt] = list(self.tst_filters[filt])
self.tra_KG, self.tra_sub = self.load_graph(self.tra_train)
self.ind_KG, self.ind_sub = self.load_graph(self.ind_train, 'inductive')
self.tra_train = np.array(self.tra_valid)
self.tra_val_qry, self.tra_val_ans = self.load_query(self.tra_test)
self.ind_val_qry, self.ind_val_ans = self.load_query(self.ind_valid)
self.ind_tst_qry, self.ind_tst_ans = self.load_query(self.ind_test)
self.valid_q, self.valid_a = self.tra_val_qry, self.tra_val_ans
self.test_q, self.test_a = self.ind_val_qry + self.ind_tst_qry, self.ind_val_ans + self.ind_tst_ans
self.n_train = len(self.tra_train)
self.n_valid = len(self.valid_q)
self.n_test = len(self.test_q)
def read_triples(self, directory, filename, mode='transductive'):
triples = []
with open(os.path.join(directory, filename)) as f:
for line in f:
h, r, t = line.strip().split()
if mode == 'transductive':
h, r, t = self.entity2id[h], self.relation2id[r], self.entity2id[t]
else:
h, r, t = self.entity2id_ind[h], self.relation2id[r], self.entity2id_ind[t]
triples.append([h, r, t])
triples.append([t, r + self.n_rel, h])
return triples
def load_graph(self, triples, mode='transductive'):
n_ent = self.n_ent if mode == 'transductive' else self.n_ent_ind
KG = np.array(triples)
idd = np.concatenate([np.expand_dims(np.arange(n_ent), 1), 2 * self.n_rel * np.ones((n_ent, 1)),
np.expand_dims(np.arange(n_ent), 1)], 1)
KG = np.concatenate([KG, idd], 0)
n_fact = KG.shape[0]
M_sub = csr_matrix((np.ones((n_fact,)), (np.arange(n_fact), KG[:, 0])), shape=(n_fact, n_ent))
return KG, M_sub
def load_query(self, triples):
triples.sort(key=lambda x: (x[0], x[1]))
trip_hr = defaultdict(lambda: list())
for trip in triples:
h, r, t = trip
trip_hr[(h, r)].append(t)
queries = []
answers = []
for key in trip_hr:
queries.append(key)
answers.append(np.array(trip_hr[key]))
return queries, answers
def get_neighbors(self, nodes, mode='transductive'):
# nodes: n_node x 2 with (batch_idx, node_idx)
if mode == 'transductive':
KG = self.tra_KG
M_sub = self.tra_sub
n_ent = self.n_ent
else:
KG = self.ind_KG
M_sub = self.ind_sub
n_ent = self.n_ent_ind
node_1hot = csr_matrix((np.ones(len(nodes)), (nodes[:, 1], nodes[:, 0])), shape=(n_ent, nodes.shape[0]))
edge_1hot = M_sub.dot(node_1hot)
edges = np.nonzero(edge_1hot)
sampled_edges = np.concatenate([np.expand_dims(edges[1], 1), KG[edges[0]]],
axis=1) # (batch_idx, head, rela, tail)
sampled_edges = th.LongTensor(sampled_edges)
# index to nodes
head_nodes, head_index = th.unique(sampled_edges[:, [0, 1]], dim=0, sorted=True, return_inverse=True)
tail_nodes, tail_index = th.unique(sampled_edges[:, [0, 3]], dim=0, sorted=True, return_inverse=True)
mask = sampled_edges[:, 2] == (self.n_rel * 2)
_, old_idx = head_index[mask].sort()
old_nodes_new_idx = tail_index[mask][old_idx]
sampled_edges = th.cat([sampled_edges, head_index.unsqueeze(1), tail_index.unsqueeze(1)], 1)
return tail_nodes, sampled_edges, old_nodes_new_idx
def get_batch(self, batch_idx, steps=2, data='train'):
if data == 'train':
return self.tra_train[batch_idx]
if data == 'valid':
# print(self.)
query, answer = np.array(self.valid_q), self.valid_a # np.array(self.valid_a)
n_ent = self.n_ent
if data == 'test':
query, answer = np.array(self.test_q), self.test_a # np.array(self.test_a)
n_ent = self.n_ent_ind
subs = []
rels = []
objs = []
subs = query[batch_idx, 0]
rels = query[batch_idx, 1]
objs = np.zeros((len(batch_idx), n_ent))
for i in range(len(batch_idx)):
objs[i][answer[batch_idx[i]]] = 1
return subs, rels, objs
def shuffle_train(self, ):
rand_idx = np.random.permutation(self.n_train)
self.tra_train = self.tra_train[rand_idx]
def get_filter(self, data='valid'):
filters = defaultdict(lambda: set())
if data == 'valid':
for triple in self.tra_train:
h, r, t = triple
filters[(h, r)].add(t)
for triple in self.tra_valid:
h, r, t = triple
filters[(h, r)].add(t)
for triple in self.tra_test:
h, r, t = triple
filters[(h, r)].add(t)
else:
for triple in self.ind_train:
h, r, t = triple
filters[(h, r)].add(t)
for triple in self.ind_valid:
h, r, t = triple
filters[(h, r)].add(t)
for triple in self.ind_test:
h, r, t = triple
filters[(h, r)].add(t)
return filters
@register_dataset('kg_subT_link_prediction')
class KG_RedTDataset(LinkPredictionDataset):
def __init__(self, dataset_name, *args, **kwargs):
super(KG_RedTDataset, self).__init__(*args, **kwargs)
self.task_dir = os.path.join('openhgnn/dataset/data', dataset_name)
task_dir = self.task_dir
folder = os.path.exists(self.task_dir)
if not folder:
os.makedirs(self.task_dir)
url = "https://s3.cn-north-1.amazonaws.com.cn/dgl-data/dataset/openhgnn/family.zip"
response = requests.get(url)
with zipfile.ZipFile(io.BytesIO(response.content)) as myzip:
myzip.extractall(self.task_dir)
print("--- download data ---")
else:
print("--- There is data! ---")
with open(os.path.join(task_dir, 'entities.txt')) as f:
self.entity2id = dict()
n_ent = 0
for line in f:
entity = line.strip()
self.entity2id[entity] = n_ent
n_ent += 1
with open(os.path.join(task_dir, 'relations.txt')) as f:
self.relation2id = dict()
n_rel = 0
for line in f:
relation = line.strip()
self.relation2id[relation] = n_rel
n_rel += 1
self.n_ent = n_ent
self.n_rel = n_rel
self.filters = defaultdict(lambda: set())
self.fact_triple = self.read_triples('facts.txt')
self.train_triple = self.read_triples('train.txt')
self.valid_triple = self.read_triples('valid.txt')
self.test_triple = self.read_triples('test.txt')
self.fact_data = self.double_triple(self.fact_triple)
self.train_data = np.array(self.double_triple(self.train_triple))
self.valid_data = self.double_triple(self.valid_triple)
self.test_data = self.double_triple(self.test_triple)
self.load_graph(self.fact_data)
self.load_test_graph(self.double_triple(self.fact_triple) + self.double_triple(self.train_triple))
self.valid_q, self.valid_a = self.load_query(self.valid_data)
self.test_q, self.test_a = self.load_query(self.test_data)
self.n_train = len(self.train_data)
self.n_valid = len(self.valid_q)
self.n_test = len(self.test_q)
for filt in self.filters:
self.filters[filt] = list(self.filters[filt])
print('n_train:', self.n_train, 'n_valid:', self.n_valid, 'n_test:', self.n_test)
def read_triples(self, filename):
triples = []
with open(os.path.join(self.task_dir, filename)) as f:
for line in f:
h, r, t = line.strip().split()
h, r, t = self.entity2id[h], self.relation2id[r], self.entity2id[t]
triples.append([h, r, t])
self.filters[(h, r)].add(t)
self.filters[(t, r + self.n_rel)].add(h)
return triples
def double_triple(self, triples):
new_triples = []
for triple in triples:
h, r, t = triple
new_triples.append([t, r + self.n_rel, h])
return triples + new_triples
def load_graph(self, triples):
idd = np.concatenate([np.expand_dims(np.arange(self.n_ent), 1), 2 * self.n_rel * np.ones((self.n_ent, 1)),
np.expand_dims(np.arange(self.n_ent), 1)], 1)
self.KG = np.concatenate([np.array(triples), idd], 0)
self.n_fact = len(self.KG)
self.M_sub = csr_matrix((np.ones((self.n_fact,)), (np.arange(self.n_fact), self.KG[:, 0])),
shape=(self.n_fact, self.n_ent))
def load_test_graph(self, triples):
idd = np.concatenate([np.expand_dims(np.arange(self.n_ent), 1), 2 * self.n_rel * np.ones((self.n_ent, 1)),
np.expand_dims(np.arange(self.n_ent), 1)], 1)
self.tKG = np.concatenate([np.array(triples), idd], 0)
self.tn_fact = len(self.tKG)
self.tM_sub = csr_matrix((np.ones((self.tn_fact,)), (np.arange(self.tn_fact), self.tKG[:, 0])),
shape=(self.tn_fact, self.n_ent))
def load_query(self, triples):
triples.sort(key=lambda x: (x[0], x[1]))
trip_hr = defaultdict(lambda: list())
for trip in triples:
h, r, t = trip
trip_hr[(h, r)].append(t)
queries = []
answers = []
for key in trip_hr:
queries.append(key)
answers.append(np.array(trip_hr[key]))
return queries, answers
def get_neighbors(self, nodes, mode='train'):
if mode == 'train':
KG = self.KG
M_sub = self.M_sub
else:
KG = self.tKG
M_sub = self.tM_sub
# nodes: n_node x 2 with (batch_idx, node_idx)
node_1hot = csr_matrix((np.ones(len(nodes)), (nodes[:, 1], nodes[:, 0])), shape=(self.n_ent, nodes.shape[0]))
edge_1hot = M_sub.dot(node_1hot)
edges = np.nonzero(edge_1hot)
sampled_edges = np.concatenate([np.expand_dims(edges[1], 1), KG[edges[0]]],
axis=1) # (batch_idx, head, rela, tail)
sampled_edges = torch.LongTensor(sampled_edges).cuda()
# index to nodes
head_nodes, head_index = torch.unique(sampled_edges[:, [0, 1]], dim=0, sorted=True, return_inverse=True)
tail_nodes, tail_index = torch.unique(sampled_edges[:, [0, 3]], dim=0, sorted=True, return_inverse=True)
sampled_edges = torch.cat([sampled_edges, head_index.unsqueeze(1), tail_index.unsqueeze(1)], 1)
mask = sampled_edges[:, 2] == (self.n_rel * 2)
_, old_idx = head_index[mask].sort()
old_nodes_new_idx = tail_index[mask][old_idx]
return tail_nodes, sampled_edges, old_nodes_new_idx
def get_batch(self, batch_idx, steps=2, data='train'):
if data == 'train':
return np.array(self.train_data)[batch_idx]
if data == 'valid':
query, answer = np.array(self.valid_q), self.valid_a
if data == 'test':
query, answer = np.array(self.test_q), self.test_a
subs = []
rels = []
objs = []
subs = query[batch_idx, 0]
rels = query[batch_idx, 1]
objs = np.zeros((len(batch_idx), self.n_ent))
for i in range(len(batch_idx)):
objs[i][answer[batch_idx[i]]] = 1
return subs, rels, objs
def shuffle_train(self, ):
fact_triple = np.array(self.fact_triple)
train_triple = np.array(self.train_triple)
all_triple = np.concatenate([fact_triple, train_triple], axis=0)
n_all = len(all_triple)
rand_idx = np.random.permutation(n_all)
all_triple = all_triple[rand_idx]
# increase the ratio of fact_data, e.g., 3/4->4/5, can increase the performance
self.fact_data = self.double_triple(all_triple[:n_all * 3 // 4].tolist())
self.train_data = np.array(self.double_triple(all_triple[n_all * 3 // 4:].tolist()))
self.n_train = len(self.train_data)
self.load_graph(self.fact_data)
@register_dataset('kg_link_prediction')
class KG_LinkPrediction(LinkPredictionDataset):
"""
From `RGCN <https://arxiv.org/abs/1703.06103>`_, WN18 & FB15k face a data leakage.
"""
def __init__(self, dataset_name, *args, **kwargs):
super(KG_LinkPrediction, self).__init__(*args, **kwargs)
if dataset_name in ['wn18', 'FB15k', 'FB15k-237']:
dataset = load_data(dataset_name)
g = dataset[0]
self.num_rels = dataset.num_rels
self.num_nodes = dataset.num_nodes
self.train_hg, self.train_triplets = self._build_hg(g, 'train')
self.valid_hg, self.valid_triplets = self._build_hg(g, 'valid')
self.test_hg, self.test_triplets = self._build_hg(g, 'test')
self.g = self.train_hg
self.category = '_N'
self.target_link = self.test_hg.canonical_etypes
def _build_hg(self, g, mode):
sub_g = dgl.edge_subgraph(g, g.edata[mode + '_edge_mask'], relabel_nodes=False)
src, dst = sub_g.edges( )
etype = sub_g.edata['etype']
edge_dict = {}
for i in range(self.num_rels):
mask = (etype == i)
edge_name = ('_N', str(i), '_N')
edge_dict[edge_name] = (src[mask], dst[mask])
hg = dgl.heterograph(edge_dict, {'_N': self.num_nodes})
return hg, th.stack((src, etype, dst)).T
def modify_size(self, eval_percent, dataset_type):
if dataset_type == 'valid':
self.valid_triplets = th.tensor(
random.sample(self.valid_triplets.tolist( ), math.ceil(self.valid_triplets.shape[0] * eval_percent)))
elif dataset_type == 'test':
self.test_triplets = th.tensor(
random.sample(self.test_triplets.tolist( ), math.ceil(self.test_triplets.shape[0] * eval_percent)))
def get_graph_directed_from_triples(self, triples, format='graph'):
s = th.LongTensor(triples[:, 0])
r = th.LongTensor(triples[:, 1])
o = th.LongTensor(triples[:, 2])
if format == 'graph':
edge_dict = {}
for i in range(self.num_rels):
mask = (r == i)
edge_name = (self.category, str(i), self.category)
edge_dict[edge_name] = (s[mask], o[mask])
return dgl.heterograph(edge_dict, {self.category: self.num_nodes})
def get_triples(self, g, mask_mode):
'''
:param g:
:param mask_mode: should be one of 'train_mask', 'val_mask', 'test_mask
:return:
'''
edges = g.edges( )
etype = g.edata['etype']
mask = g.edata.pop(mask_mode)
return th.stack((edges[0][mask], etype[mask], edges[1][mask]))
def get_all_triplets(self, dataset):
train_data = th.LongTensor(dataset.train)
valid_data = th.LongTensor(dataset.valid)
test_data = th.LongTensor(dataset.test)
return train_data, valid_data, test_data
def get_split(self):
return self.train_hg, self.valid_hg, self.test_hg, None, None
def split_graph(self, g, mode='train'):
"""
Parameters
----------
g: DGLGraph
a homogeneous graph fomat
mode: str
split the subgraph according to the mode
Returns
-------
hg: DGLHeterograph
"""
edges = g.edges( )
etype = g.edata['etype']
if mode == 'train':
mask = g.edata['train_mask']
elif mode == 'valid':
mask = g.edata['valid_edge_mask']
elif mode == 'test':
mask = g.edata['test_edge_mask']
hg = self.build_graph((edges[0][mask], edges[1][mask]), etype[mask])
return hg
def build_graph(self, edges, etype):
edge_dict = {}
for i in range(self.num_rels):
mask = (etype == i)
edge_name = (self.category, str(i), self.category)
edge_dict[edge_name] = (edges[0][mask], edges[1][mask])
hg = dgl.heterograph(edge_dict, {self.category: self.num_nodes})
return hg
def build_g(self, train):
s = train[:, 0]
r = train[:, 1]
o = train[:, 2]
edge_dict = {}
for i in range(self.num_rels):
mask = (r == i)
edge_name = (self.category, str(i), self.category)
edge_dict[edge_name] = (th.LongTensor(s[mask]), th.LongTensor(o[mask]))
hg = dgl.heterograph(edge_dict, {self.category: self.num_nodes})
return hg
import torch
import struct
import os
import json
import logging
from scipy.sparse import csc_matrix
from scipy.special import softmax
from tqdm import tqdm
import pickle
import scipy.sparse as ssp
import lmdb
import requests
import zipfile
import io
from torch.utils.data import Dataset
import networkx as nx
from ..utils.Grail_utils import *
class SubGraphDataset(Dataset):
def __init__(self, db_path, db_name_pos, db_name_neg, raw_data_paths, included_relations=None,
add_traspose_rels=False, num_neg_samples_per_link=1, use_kge_embeddings=False, dataset='',
kge_model='', file_name=''):
self.main_env = lmdb.open(db_path, readonly= True, max_dbs=3, lock=False)
self.db_pos = self.main_env.open_db(db_name_pos.encode())
self.db_neg = self.main_env.open_db(db_name_neg.encode())
self.node_features, self.kge_entity2id = get_kge_embeddings(dataset, kge_model) if use_kge_embeddings else (None, None)
self.num_neg_samples_per_link = num_neg_samples_per_link
self.file_name = file_name
self.add_traspose_rels = add_traspose_rels
ssp_graph, __, __, __, id2entity, id2relation = process_files(raw_data_paths, included_relations)
self.num_rels = len(ssp_graph)
# Add transpose matrices to handle both directions of relations.
if add_traspose_rels:
ssp_graph_t = [adj.T for adj in ssp_graph]
ssp_graph += ssp_graph_t
# the effective number of relations after adding symmetric adjacency matrices and/or self connections
self.aug_num_rels = len(ssp_graph)
self.graph = ssp_multigraph_to_dgl(ssp_graph)
self.ssp_graph = ssp_graph
self.id2entity = id2entity
self.id2relation = id2relation
self.max_n_label = np.array([0, 0])
with self.main_env.begin() as txn:
#a = txn.get('max_n_label_sub'.encode())
#print(a)
self.max_n_label[0] = int.from_bytes(txn.get('max_n_label_sub'.encode()), byteorder='little')
self.max_n_label[1] = int.from_bytes(txn.get('max_n_label_obj'.encode()), byteorder='little')
self.avg_subgraph_size = struct.unpack('f', txn.get('avg_subgraph_size'.encode()))
self.min_subgraph_size = struct.unpack('f', txn.get('min_subgraph_size'.encode()))
self.max_subgraph_size = struct.unpack('f', txn.get('max_subgraph_size'.encode()))
self.std_subgraph_size = struct.unpack('f', txn.get('std_subgraph_size'.encode()))
self.avg_enc_ratio = struct.unpack('f', txn.get('avg_enc_ratio'.encode()))
self.min_enc_ratio = struct.unpack('f', txn.get('min_enc_ratio'.encode()))
self.max_enc_ratio = struct.unpack('f', txn.get('max_enc_ratio'.encode()))
self.std_enc_ratio = struct.unpack('f', txn.get('std_enc_ratio'.encode()))
self.avg_num_pruned_nodes = struct.unpack('f', txn.get('avg_num_pruned_nodes'.encode()))
self.min_num_pruned_nodes = struct.unpack('f', txn.get('min_num_pruned_nodes'.encode()))
self.max_num_pruned_nodes = struct.unpack('f', txn.get('max_num_pruned_nodes'.encode()))
self.std_num_pruned_nodes = struct.unpack('f', txn.get('std_num_pruned_nodes'.encode()))
logging.info(f"Max distance from sub : {self.max_n_label[0]}, Max distance from obj : {self.max_n_label[1]}")
# logging.info('=====================')
# logging.info(f"Subgraph size stats: \n Avg size {self.avg_subgraph_size}, \n Min size {self.min_subgraph_size}, \n Max size {self.max_subgraph_size}, \n Std {self.std_subgraph_size}")
# logging.info('=====================')
# logging.info(f"Enclosed nodes ratio stats: \n Avg size {self.avg_enc_ratio}, \n Min size {self.min_enc_ratio}, \n Max size {self.max_enc_ratio}, \n Std {self.std_enc_ratio}")
# logging.info('=====================')
# logging.info(f"# of pruned nodes stats: \n Avg size {self.avg_num_pruned_nodes}, \n Min size {self.min_num_pruned_nodes}, \n Max size {self.max_num_pruned_nodes}, \n Std {self.std_num_pruned_nodes}")
with self.main_env.begin(db=self.db_pos) as txn:
self.num_graphs_pos = int.from_bytes(txn.get('num_graphs'.encode()), byteorder='little')
with self.main_env.begin(db=self.db_neg) as txn:
self.num_graphs_neg = int.from_bytes(txn.get('num_graphs'.encode()), byteorder='little')
self.__getitem__(0)
def __getitem__(self, index):
with self.main_env.begin(db=self.db_pos) as txn:
str_id = '{:08}'.format(index).encode('ascii')
nodes_pos, r_label_pos, g_label_pos, n_labels_pos = deserialize(txn.get(str_id)).values()
subgraph_pos = self._prepare_subgraphs(nodes_pos, r_label_pos, n_labels_pos)
subgraphs_neg = []
r_labels_neg = []
g_labels_neg = []
with self.main_env.begin(db=self.db_neg) as txn:
for i in range(self.num_neg_samples_per_link):
str_id = '{:08}'.format(index + i * (self.num_graphs_pos)).encode('ascii')
nodes_neg, r_label_neg, g_label_neg, n_labels_neg = deserialize(txn.get(str_id)).values()
subgraphs_neg.append(self._prepare_subgraphs(nodes_neg, r_label_neg, n_labels_neg))
r_labels_neg.append(r_label_neg)
g_labels_neg.append(g_label_neg)
return subgraph_pos, g_label_pos, r_label_pos, subgraphs_neg, g_labels_neg, r_labels_neg
def __len__(self):
return self.num_graphs_pos
def _prepare_subgraphs(self, nodes, r_label, n_labels):
if not isinstance(self.graph, dgl.DGLGraph):
subgraph = dgl.graph(self.graph.subgraph(nodes))
else:
subgraph = self.graph.subgraph(nodes)
#subgraph.edata['type'] = self.graph.edata['type'][self.graph.subgraph(nodes).parent_eid]
subgraph.edata['type'] = self.graph.edata['type'][subgraph.edata[dgl.EID]]
subgraph.edata['label'] = torch.tensor(r_label * np.ones(subgraph.edata['type'].shape), dtype=torch.long)
#print("请输出： ")
#print(subgraph)
#edges_btw_roots = subgraph.edge_id(0, 1, return_array=True)
#edges_btw_roots = subgraph.edge_ids(0, 1)
edges_btw_roots = torch.tensor([])
try:
edges_btw_roots = subgraph.edge_ids(torch.tensor([0]),torch.tensor([1]))
# edges_btw_roots = np.array([edges_btw_roots])
except:
#print("Error")
edges_btw_roots = torch.tensor([])
edges_btw_roots = edges_btw_roots.numpy()
rel_link = np.nonzero(subgraph.edata['type'][edges_btw_roots] == r_label)
if rel_link.squeeze().nelement() == 0:
subgraph = dgl.add_edges(subgraph, 0, 1)
subgraph.edata['type'][-1] = torch.tensor(r_label).type(torch.LongTensor)
subgraph.edata['label'][-1] = torch.tensor(r_label).type(torch.LongTensor)
# map the id read by GraIL to the entity IDs as registered by the KGE embeddings
kge_nodes = [self.kge_entity2id[self.id2entity[n]] for n in nodes] if self.kge_entity2id else None
n_feats = self.node_features[kge_nodes] if self.node_features is not None else None
subgraph = self._prepare_features_new(subgraph, n_labels, n_feats)
return subgraph
def _prepare_features(self, subgraph, n_labels, n_feats=None):
# One hot encode the node label feature and concat to n_featsure
n_nodes = subgraph.number_of_nodes()
label_feats = np.zeros((n_nodes, self.max_n_label[0] + 1))
label_feats[np.arange(n_nodes), n_labels] = 1
label_feats[np.arange(n_nodes), self.max_n_label[0] + 1 + n_labels[:, 1]] = 1
n_feats = np.concatenate((label_feats, n_feats), axis=1) if n_feats else label_feats
subgraph.ndata['feat'] = torch.FloatTensor(n_feats)
self.n_feat_dim = n_feats.shape[1] # Find cleaner way to do this -- i.e. set the n_feat_dim
return subgraph
def _prepare_features_new(self, subgraph, n_labels, n_feats=None):
# One hot encode the node label feature and concat to n_featsure
n_nodes = subgraph.number_of_nodes()
label_feats = np.zeros((n_nodes, self.max_n_label[0] + 1 + self.max_n_label[1] + 1))
label_feats[np.arange(n_nodes), n_labels[:, 0]] = 1
label_feats[np.arange(n_nodes), self.max_n_label[0] + 1 + n_labels[:, 1]] = 1
# label_feats = np.zeros((n_nodes, self.max_n_label[0] + 1 + self.max_n_label[1] + 1))
# label_feats[np.arange(n_nodes), 0] = 1
# label_feats[np.arange(n_nodes), self.max_n_label[0] + 1] = 1
n_feats = np.concatenate((label_feats, n_feats), axis=1) if n_feats is not None else label_feats
subgraph.ndata['feat'] = torch.FloatTensor(n_feats)
head_id = np.argwhere([label[0] == 0 and label[1] == 1 for label in n_labels])
tail_id = np.argwhere([label[0] == 1 and label[1] == 0 for label in n_labels])
n_ids = np.zeros(n_nodes)
n_ids[head_id] = 1 # head
n_ids[tail_id] = 2 # tail
subgraph.ndata['id'] = torch.FloatTensor(n_ids)
self.n_feat_dim = n_feats.shape[1] # Find cleaner way to do this -- i.e. set the n_feat_dim
return subgraph
@register_dataset('grail_link_prediction')
class Grail_LinkPrediction(LinkPredictionDataset):
def __init__(self, dataset_name, *args, **kwargs):
super(Grail_LinkPrediction, self).__init__(*args, **kwargs)
self.args = kwargs['args']
self.args.db_path = f'./openhgnn/dataset/data/{self.args.dataset}/subgraphs_en_{self.args.enclosing_sub_graph}_neg_{self.args.num_neg_samples_per_link}_hop_{self.args.hop}'
self.args.train_file = "train"
self.args.valid_file = "valid"
self.args.file_paths = {
'train': './openhgnn/dataset/data/{}/{}.txt'.format(self.args.dataset, self.args.train_file),
'valid': './openhgnn/dataset/data/{}/{}.txt'.format(self.args.dataset, self.args.valid_file)
}
relation2id_path = f'./openhgnn/dataset/data/{self.args.dataset}/relation2id.json'
self.data_folder = f'./openhgnn/dataset/data/{self.args.dataset}'
if not os.path.exists(self.data_folder):
os.makedirs(self.data_folder) # makedirs 创建文件时如果路径不存在会创建这个路径
url = f'https://github.com/kkteru/grail/blob/master/data/{self.args.dataset}'
self.download_folder(url,self.data_folder)
print("--- download data ---")
else:
print("--- There is data! ---")
if not os.path.exists(self.data_folder+'_ind'):
os.makedirs(self.data_folder+'_ind') # makedirs 创建文件时如果路径不存在会创建这个路径
url = f'https://github.com/kkteru/grail/blob/master/data/{self.args.dataset}_ind'
self.download_folder(url,self.data_folder+'_ind')
print("--- download data ---")
else:
print("--- There is data! ---")
if not os.path.isdir(self.args.db_path):
generate_subgraph_datasets(self.args, relation2id_path)
with open(relation2id_path) as f:
self.relation2id = json.load(f)
self.train = SubGraphDataset(self.args.db_path, 'train_pos', 'train_neg', self.args.file_paths,add_traspose_rels=self.args.add_traspose_rels,num_neg_samples_per_link=self.args.num_neg_samples_per_link,use_kge_embeddings=self.args.use_kge_embeddings, dataset=self.args.dataset,kge_model=self.args.kge_model, file_name=self.args.train_file)
self.valid = SubGraphDataset(self.args.db_path, 'valid_pos', 'valid_neg', self.args.file_paths,
add_traspose_rels=self.args.add_traspose_rels,
num_neg_samples_per_link=self.args.num_neg_samples_per_link,
use_kge_embeddings=self.args.use_kge_embeddings, dataset=self.args.dataset,
kge_model=self.args.kge_model, file_name= self.args.valid_file)
def download_folder(self,url, save_path):
response = requests.get(url)
if response.status_code == 200:
# 确保保存路径存在
os.makedirs(save_path, exist_ok=True)
# 解析响应内容
content = response.content.decode('utf-8')
lines = content.splitlines()
for line in lines:
# 提取文件名
file_name = line.split('/')[-1]
# 构建文件的完整URL
file_url = url + '/' + file_name
# 构建文件的保存路径
file_save_path = os.path.join(save_path, file_name)
# 下载文件
self.download_file(file_url, file_save_path)
def download_file(self,url, save_path):
response = requests.get(url)
if response.status_code == 200:
with open(save_path, 'wb') as file:
file.write(response.content)
class kg_sampler( ):
def __init__(self, ):
self.sampler = 'uniform'
return
def generate_sampled_graph_and_labels(self, triplets, sample_size, split_size,
num_rels, adj_list, degrees,
negative_rate, sampler="uniform"):
"""Get training graph and signals
First perform edge neighborhood sampling on graph, then perform negative
sampling to generate negative samples
"""
# perform edge neighbor sampling
if self.sampler == "uniform":
edges = sample_edge_uniform(adj_list, degrees, len(triplets), sample_size)
elif self.sampler == "neighbor":
edges = sample_edge_neighborhood(adj_list, degrees, len(triplets), sample_size)
else:
raise ValueError("Sampler type must be either 'uniform' or 'neighbor'.")
# relabel nodes to have consecutive node ids
edges = triplets[edges]
src, rel, dst = edges.transpose( )
uniq_v, edges = np.unique((src, dst), return_inverse=True)
src, dst = np.reshape(edges, (2, -1))
relabeled_edges = np.stack((src, rel, dst)).transpose( )
# negative sampling
samples, labels = negative_sampling(relabeled_edges, len(uniq_v),
negative_rate)
# further split graph, only half of the edges will be used as graph
# structure, while the rest half is used as unseen positive samples
split_size = int(sample_size * split_size)
graph_split_ids = np.random.choice(np.arange(sample_size),
size=split_size, replace=False)
src = src[graph_split_ids]
dst = dst[graph_split_ids]
rel = rel[graph_split_ids]
# build DGL graph
print("# sampled nodes: {}".format(len(uniq_v)))
print("# sampled edges: {}".format(len(src) * 2))
g, rel, norm = build_graph_from_triplets(len(uniq_v), num_rels,
(src, rel, dst))
return g, uniq_v, rel, norm, samples, labels
def sample_edge_neighborhood(adj_list, degrees, n_triplets, sample_size):
"""Sample edges by neighborhool expansion.
This guarantees that the sampled edges form a connected graph, which
may help deeper GNNs that require information from more than one hop.
"""
edges = np.zeros((sample_size), dtype=np.int32)
# initialize
sample_counts = np.array([d for d in degrees])
picked = np.array([False for _ in range(n_triplets)])
seen = np.array([False for _ in degrees])
for i in range(0, sample_size):
weights = sample_counts * seen
if np.sum(weights) == 0:
weights = np.ones_like(weights)
weights[np.where(sample_counts == 0)] = 0
probabilities = (weights) / np.sum(weights)
chosen_vertex = np.random.choice(np.arange(degrees.shape[0]),
p=probabilities)
chosen_adj_list = adj_list[chosen_vertex]
seen[chosen_vertex] = True
chosen_edge = np.random.choice(np.arange(chosen_adj_list.shape[0]))
chosen_edge = chosen_adj_list[chosen_edge]
edge_number = chosen_edge[0]
while picked[edge_number]:
chosen_edge = np.random.choice(np.arange(chosen_adj_list.shape[0]))
chosen_edge = chosen_adj_list[chosen_edge]
edge_number = chosen_edge[0]
edges[i] = edge_number
other_vertex = chosen_edge[1]
picked[edge_number] = True
sample_counts[chosen_vertex] -= 1
sample_counts[other_vertex] -= 1
seen[other_vertex] = True
return edges
def sample_edge_uniform(adj_list, degrees, n_triplets, sample_size):
"""Sample edges uniformly from all the edges."""
all_edges = np.arange(n_triplets)
return np.random.choice(all_edges, sample_size, replace=False)
# --- ExpressGNN ---
# grounded rule stats code
BAD = 0 # sample not valid
FULL_OBSERVERED = 1 # sample valid, but rule contains only observed vars and does not have negation for all atoms
GOOD = 2 # sample valid
@register_dataset('express_gnn')
class ExpressGNNDataset(BaseDataset):
def __init__(self, dataset_name, *args, **kwargs):
super( ).__init__(*args, **kwargs)
self.args = kwargs['args']
self.PRED_DICT = {}
self.dataset_name = dataset_name
self.const_dict = ConstantDict()
self.batchsize = self.args.batchsize
self.shuffle_sampling = self.args.shuffle_sampling
data_root = 'openhgnn'
data_root = os.path.join(data_root, 'dataset')
data_root = os.path.join(data_root, 'data')
data_root = os.path.join(data_root, self.dataset_name)
ext_rule_path = None
folder = os.path.exists(data_root)
print(data_root)
print('folder', folder)
if not folder: # 判断是否存在文件夹如果不存在则创建为文件夹
os.makedirs(data_root) # makedirs 创建文件时如果路径不存在会创建这个路径
# 下载数据
url = f"https://s3.cn-north-1.amazonaws.com.cn/dgl-data/dataset/openhgnn/{dataset_name}.zip"
response = requests.get(url)
with zipfile.ZipFile(io.BytesIO(response.content)) as myzip:
myzip.extractall(data_root)
print("--- download data ---")
else:
print("--- There is data! ---")
# Decide the way dataset will be load, set 1 to load FBWN dataset
load_method = 0
# print(dataset_name[0:13])
if dataset_name[0:13] == 'EXP_FB15k-237':
load_method = 1
else:
load_method = 0
guss_fb = 'EXP_FB15k' in data_root
if guss_fb != (load_method == 1):
print("WARNING: set load_method to 1 if you load Freebase dataset, otherwise 0")
# FBWN dataset
if load_method == 1:
fact_path_ls = [joinpath(data_root, 'facts.txt'),
joinpath(data_root, 'train.txt')]
query_path = joinpath(data_root, 'test.txt')
pred_path = joinpath(data_root, 'relations.txt')
const_path = joinpath(data_root, 'entities.txt')
valid_path = joinpath(data_root, 'valid.txt')
rule_path = joinpath(data_root, 'cleaned_rules_weight_larger_than_0.9.txt')
print(rule_path)
print(os.getcwd())
# print(fact_path_ls + [query_path, pred_path, const_path, valid_path, rule_path])
# assert all(map(isfile, fact_path_ls + [query_path, pred_path, const_path, valid_path, rule_path]))
# assuming only one type
TYPE_SET.update(['type'])
# add all const
for line in iterline(const_path):
self.const_dict.add_const('type', line)
# add all pred
for line in iterline(pred_path):
self.PRED_DICT[line] = Predicate(line, ['type', 'type'])
# add all facts
fact_ls = []
for fact_path in fact_path_ls:
for line in iterline(fact_path):
parts = line.split('\t')
assert len(parts) == 3, print(parts)
e1, pred_name, e2 = parts
assert self.const_dict.has_const('type', e1) and self.const_dict.has_const('type', e2)
assert pred_name in self.PRED_DICT
fact_ls.append(Fact(pred_name, [e1, e2], 1))
# add all validations
valid_ls = []
for line in iterline(valid_path):
parts = line.split('\t')
assert len(parts) == 3, print(parts)
e1, pred_name, e2 = parts
assert self.const_dict.has_const('type', e1) and self.const_dict.has_const('type', e2)
assert pred_name in self.PRED_DICT
valid_ls.append(Fact(pred_name, [e1, e2], 1))
# add all queries
query_ls = []
for line in iterline(query_path):
parts = line.split('\t')
assert len(parts) == 3, print(parts)
e1, pred_name, e2 = parts
assert self.const_dict.has_const('type', e1) and self.const_dict.has_const('type', e2)
assert pred_name in self.PRED_DICT
query_ls.append(Fact(pred_name, [e1, e2], 1))
# add all rules
rule_ls = []
strip_items = lambda ls: list(map(lambda x: x.strip( ), ls))
first_atom_reg = re.compile(r'([\d.]+) (!?)([^(]+)\((.*)\)')
atom_reg = re.compile(r'(!?)([^(]+)\((.*)\)')
for line in iterline(rule_path):
atom_str_ls = strip_items(line.split(' v '))
assert len(atom_str_ls) > 1, 'rule length must be greater than 1, but get %s' % line
atom_ls = []
rule_weight = 0.0
for i, atom_str in enumerate(atom_str_ls):
if i == 0:
m = first_atom_reg.match(atom_str)
assert m is not None, 'matching atom failed for %s' % atom_str
rule_weight = float(m.group(1))
neg = m.group(2) == '!'
pred_name = m.group(3).strip( )
var_name_ls = strip_items(m.group(4).split(','))
else:
m = atom_reg.match(atom_str)
assert m is not None, 'matching atom failed for %s' % atom_str
neg = m.group(1) == '!'
pred_name = m.group(2).strip( )
var_name_ls = strip_items(m.group(3).split(','))
atom = Atom(neg, pred_name, var_name_ls, self.PRED_DICT[pred_name].var_types)
atom_ls.append(atom)
rule = Formula(atom_ls, rule_weight)
rule_ls.append(rule)
else:
if dataset_name == 'Cora' or dataset_name == 'kinship':
data_root = joinpath(data_root, 'S' + str(self.args.load_s))
elif dataset_name == 'uw_cse':
if self.args.load_s == 1:
data_root = joinpath(data_root, 'ai')
elif self.args.load_s == 2:
data_root = joinpath(data_root, 'graphics')
elif self.args.load_s == 3:
data_root = joinpath(data_root, 'language')
elif self.args.load_s == 4:
data_root = joinpath(data_root, 'systems')
elif self.args.load_s == 5:
data_root = joinpath(data_root, 'theory')
else:
print('Warning: Invalid load_s')
else:
print('Warning: Invalid dataset for load_method = 0')
rpath = joinpath(data_root, 'rules') if ext_rule_path is None else ext_rule_path
fact_ls, rule_ls, query_ls = self.preprocess_kinship(joinpath(data_root, 'predicates'),
joinpath(data_root, 'facts'),
rpath,
joinpath(data_root, 'queries'))
valid_ls = []
self.const_sort_dict = dict(
[(type_name, sorted(list(self.const_dict[type_name]))) for type_name in self.const_dict.constants.keys( )])
if load_method == 1:
self.const2ind = dict([(const, i) for i, const in enumerate(self.const_sort_dict['type'])])
# linear in size of facts
self.fact_dict = dict((pred_name, set( )) for pred_name in self.PRED_DICT)
self.test_fact_dict = dict((pred_name, set( )) for pred_name in self.PRED_DICT)
self.valid_dict = dict((pred_name, set( )) for pred_name in self.PRED_DICT)
self.ht_dict = dict((pred_name, [dict( ), dict( )]) for pred_name in self.PRED_DICT)
self.ht_dict_train = dict((pred_name, [dict( ), dict( )]) for pred_name in self.PRED_DICT)
def add_ht(pn, c_ls, ht_dict):
if load_method == 0:
if c_ls[0] in ht_dict[pn][0]:
ht_dict[pn][0][c_ls[0]].add(c_ls[0])
else:
ht_dict[pn][0][c_ls[0]] = {c_ls[0]}
elif load_method == 1:
if c_ls[0] in ht_dict[pn][0]:
ht_dict[pn][0][c_ls[0]].add(c_ls[1])
else:
ht_dict[pn][0][c_ls[0]] = {c_ls[1]}
if c_ls[1] in ht_dict[pn][1]:
ht_dict[pn][1][c_ls[1]].add(c_ls[0])
else:
ht_dict[pn][1][c_ls[1]] = {c_ls[0]}
const_cnter = Counter()
for fact in fact_ls:
self.fact_dict[fact.pred_name].add((fact.val, tuple(fact.const_ls)))
add_ht(fact.pred_name, fact.const_ls, self.ht_dict)
add_ht(fact.pred_name, fact.const_ls, self.ht_dict_train)
const_cnter.update(fact.const_ls)
for fact in valid_ls:
self.valid_dict[fact.pred_name].add((fact.val, tuple(fact.const_ls)))
add_ht(fact.pred_name, fact.const_ls, self.ht_dict)
# the sorted list version
self.fact_dict_2 = dict((pred_name, sorted(list(self.fact_dict[pred_name])))
for pred_name in self.fact_dict.keys( ))
self.valid_dict_2 = dict((pred_name, sorted(list(self.valid_dict[pred_name])))
for pred_name in self.valid_dict.keys( ))
self.rule_ls = rule_ls
# pred_atom-key dict
self.atom_key_dict_ls = []
for rule in self.rule_ls:
atom_key_dict = dict( )
for atom in rule.atom_ls:
atom_dict = dict((var_name, dict( )) for var_name in atom.var_name_ls)
for i, var_name in enumerate(atom.var_name_ls):
if atom.pred_name not in self.fact_dict:
continue
for v in self.fact_dict[atom.pred_name]:
if v[1][i] not in atom_dict[var_name]:
atom_dict[var_name][v[1][i]] = [v]
else:
atom_dict[var_name][v[1][i]] += [v]
# happens if predicate occurs more than once in one rule then we merge the set
if atom.pred_name in atom_key_dict:
for k, v in atom_dict.items( ):
if k not in atom_key_dict[atom.pred_name]:
atom_key_dict[atom.pred_name][k] = v
else:
atom_key_dict[atom.pred_name] = atom_dict
self.atom_key_dict_ls.append(atom_key_dict)
self.test_fact_ls = []
self.valid_fact_ls = []
for fact in query_ls:
self.test_fact_ls.append((fact.val, fact.pred_name, tuple(fact.const_ls)))
self.test_fact_dict[fact.pred_name].add((fact.val, tuple(fact.const_ls)))
add_ht(fact.pred_name, fact.const_ls, self.ht_dict)
for fact in valid_ls:
self.valid_fact_ls.append((fact.val, fact.pred_name, tuple(fact.const_ls)))
self.num_rules = len(rule_ls)
self.rule_gens = None
self.reset( )
def generate_gnd_pred(self, pred_name):
"""
return a list of all instantiations of a predicate function, this can be extremely large
:param pred_name:
string
:return:
"""
assert pred_name in self.PRED_DICT
pred = self.PRED_DICT[pred_name]
subs = itertools.product(*[self.const_sort_dict[var_type] for var_type in pred.var_types])
return [(pred_name, sub) for sub in subs]
def generate_gnd_rule(self, rule):
subs = itertools.product(*[self.const_sort_dict[rule.rule_vars[k]] for k in rule.rule_vars.keys( )])
sub = next(subs, None)
while sub is not None:
latent_vars = []
latent_neg_mask = []
observed_neg_mask = []
for atom in rule.atom_ls:
grounding = tuple(sub[rule.key2ind[var_name]] for var_name in atom.var_name_ls)
pos_gnding, neg_gnding = (1, grounding), (0, grounding)
if pos_gnding in self.fact_dict[atom.pred_name]:
observed_neg_mask.append(0 if atom.neg else 1)
elif neg_gnding in self.fact_dict[atom.pred_name]:
observed_neg_mask.append(1 if atom.neg else 0)
else:
latent_vars.append((atom.pred_name, grounding))
latent_neg_mask.append(1 if atom.neg else 0)
isfullneg = (sum(latent_neg_mask) == len(latent_neg_mask)) and \
(sum(observed_neg_mask) > 0)
yield latent_vars, [latent_neg_mask, observed_neg_mask], isfullneg
sub = next(subs, None)
def get_batch(self, epoch_mode=False, filter_latent=True):
"""
return the ind-th batch of ground formula and latent variable indicators
:return:
Parameters
----------
filter_latent
epoch_mode
"""
batch_neg_mask = [[] for _ in range(len(self.rule_ls))]
batch_latent_var_inds = [[] for _ in range(len(self.rule_ls))]
observed_rule_cnts = [0.0 for _ in range(len(self.rule_ls))]
flat_latent_vars = dict( )
cnt = 0
inds = list(range(len(self.rule_ls)))
while cnt < self.batchsize:
if self.shuffle_sampling:
shuffle(inds)
hasdata = False
for ind in inds:
latent_vars, neg_mask, isfullneg = next(self.rule_gens[ind], (None, None, None))
if latent_vars is None:
if epoch_mode:
continue
else:
self.rule_gens[ind] = self.generate_gnd_rule(self.rule_ls[ind])
latent_vars, neg_mask, isfullneg = next(self.rule_gens[ind])
if epoch_mode:
hasdata = True
# if rule is fully latent
if (len(neg_mask[1]) == 0) and filter_latent:
continue
# if rule fully observed
if len(latent_vars) == 0:
observed_rule_cnts[ind] += 0 if isfullneg else 1
cnt += 1
if cnt >= self.batchsize:
break
else:
continue
batch_neg_mask[ind].append(neg_mask)
for latent_var in latent_vars:
if latent_var not in flat_latent_vars:
flat_latent_vars[latent_var] = len(flat_latent_vars)
batch_latent_var_inds[ind].append([flat_latent_vars[e] for e in latent_vars])
cnt += 1
if cnt >= self.batchsize:
break
if epoch_mode and (hasdata is False):
break
flat_list = sorted([(k, v) for k, v in flat_latent_vars.items( )], key=lambda x: x[1])
flat_list = [e[0] for e in flat_list]
return batch_neg_mask, flat_list, batch_latent_var_inds, observed_rule_cnts
def _instantiate_pred(self, atom, atom_dict, sub, rule, observed_prob):
key2ind = rule.key2ind
rule_vars = rule.rule_vars
# substitute with observed fact
if np.random.rand( ) < observed_prob:
fact_choice_set = None
for var_name in atom.var_name_ls:
const = sub[key2ind[var_name]]
if const is None:
choice_set = itertools.chain.from_iterable([v for k, v in atom_dict[var_name].items( )])
else:
if const in atom_dict[var_name]:
choice_set = atom_dict[var_name][const]
else:
choice_set = []
if fact_choice_set is None:
fact_choice_set = set(choice_set)
else:
fact_choice_set = fact_choice_set.intersection(set(choice_set))
if len(fact_choice_set) == 0:
break
if len(fact_choice_set) == 0:
for var_name in atom.var_name_ls:
if sub[key2ind[var_name]] is None:
sub[key2ind[var_name]] = choice(self.const_sort_dict[rule_vars[var_name]])
else:
val, const_ls = choice(sorted(list(fact_choice_set)))
for var_name, const in zip(atom.var_name_ls, const_ls):
sub[key2ind[var_name]] = const
# substitute with random facts
else:
for var_name in atom.var_name_ls:
if sub[key2ind[var_name]] is None:
sub[key2ind[var_name]] = choice(self.const_sort_dict[rule_vars[var_name]])
def _gen_mask(self, rule, sub, closed_world):
latent_vars = []
observed_vars = []
latent_neg_mask = []
observed_neg_mask = []
for atom in rule.atom_ls:
grounding = tuple(sub[rule.key2ind[var_name]] for var_name in atom.var_name_ls)
pos_gnding, neg_gnding = (1, grounding), (0, grounding)
if pos_gnding in self.fact_dict[atom.pred_name]:
observed_vars.append((1, atom.pred_name))
observed_neg_mask.append(0 if atom.neg else 1)
elif neg_gnding in self.fact_dict[atom.pred_name]:
observed_vars.append((0, atom.pred_name))
observed_neg_mask.append(1 if atom.neg else 0)
else:
if closed_world and (len(self.test_fact_dict[atom.pred_name]) == 0):
observed_vars.append((0, atom.pred_name))
observed_neg_mask.append(1 if atom.neg else 0)
else:
latent_vars.append((atom.pred_name, grounding))
latent_neg_mask.append(1 if atom.neg else 0)
return latent_vars, observed_vars, latent_neg_mask, observed_neg_mask
def _get_rule_stat(self, observed_vars, latent_vars, observed_neg_mask, filter_latent, filter_observed):
is_full_latent = len(observed_vars) == 0
is_full_observed = len(latent_vars) == 0
if is_full_latent and filter_latent:
return BAD
if is_full_observed:
if filter_observed:
return BAD
is_full_neg = sum(observed_neg_mask) == 0
if is_full_neg:
return BAD
else:
return FULL_OBSERVERED
# if observed var already yields 1
if sum(observed_neg_mask) > 0:
return BAD
return GOOD
def _inst_var(self, sub, var2ind, var2type, at, ht_dict, gen_latent):
if len(at.var_name_ls) != 2:
raise KeyError
must_latent = gen_latent
if must_latent:
tmp = [sub[var2ind[vn]] for vn in at.var_name_ls]
for i, subi in enumerate(tmp):
if subi is None:
tmp[i] = random.choice(self.const_sort_dict[var2type[at.var_name_ls[i]]])
islatent = (tmp[0] not in ht_dict[0]) or (tmp[1] not in ht_dict[0][tmp[0]])
for i, vn in enumerate(at.var_name_ls):
sub[var2ind[vn]] = tmp[i]
return [self.const2ind[subi] for subi in tmp], islatent, islatent or at.neg
vn0 = at.var_name_ls[0]
sub0 = sub[var2ind[vn0]]
vn1 = at.var_name_ls[1]
sub1 = sub[var2ind[vn1]]
if sub0 is None:
if sub1 is None:
if len(ht_dict[0]) > 0:
sub0 = random.choice(tuple(ht_dict[0].keys( )))
sub1 = random.choice(tuple(ht_dict[0][sub0]))
sub[var2ind[vn0]] = sub0
sub[var2ind[vn1]] = sub1
return [self.const2ind[sub0], self.const2ind[sub1]], False, at.neg
else:
if sub1 in ht_dict[1]:
sub0 = random.choice(tuple(ht_dict[1][sub1]))
sub[var2ind[vn0]] = sub0
return [self.const2ind[sub0], self.const2ind[sub1]], False, at.neg
else:
sub0 = random.choice(self.const_sort_dict[var2type[vn0]])
sub[var2ind[vn0]] = sub0
return [self.const2ind[sub0], self.const2ind[sub1]], True, True
else:
if sub1 is None:
if sub0 in ht_dict[0]:
sub1 = random.choice(tuple(ht_dict[0][sub0]))
sub[var2ind[vn1]] = sub1
return [self.const2ind[sub0], self.const2ind[sub1]], False, at.neg
else:
sub1 = random.choice(self.const_sort_dict[var2type[vn1]])
sub[var2ind[vn1]] = sub1
return [self.const2ind[sub0], self.const2ind[sub1]], True, True
else:
islatent = (sub0 not in ht_dict[0]) or (sub1 not in ht_dict[0][sub0])
return [self.const2ind[sub0], self.const2ind[sub1]], islatent, islatent or at.neg
def get_batch_fast(self, batchsize, observed_prob=0.9):
prob_decay = 0.5
for rule in self.rule_ls:
var2ind = rule.key2ind
var2type = rule.rule_vars
samples = [[atom.pred_name, []] for atom in rule.atom_ls]
neg_mask = [[atom.pred_name, []] for atom in rule.atom_ls]
latent_mask = [[atom.pred_name, []] for atom in rule.atom_ls]
obs_var = [[atom.pred_name, []] for atom in rule.atom_ls]
cnt = 0
while cnt <= batchsize:
sub = [None] * len(rule.rule_vars) # substitutions
sample_buff = [[] for _ in rule.atom_ls]
neg_mask_buff = [[] for _ in rule.atom_ls]
latent_mask_buff = [[] for _ in rule.atom_ls]
atom_inds = list(range(len(rule.atom_ls)))
shuffle(atom_inds)
succ = True
cur_threshold = observed_prob
obs_list = []
for atom_ind in atom_inds:
atom = rule.atom_ls[atom_ind]
pred_ht_dict = self.ht_dict_train[atom.pred_name]
gen_latent = np.random.rand( ) > cur_threshold
c_ls, islatent, atom_succ = self._inst_var(sub, var2ind, var2type,
atom, pred_ht_dict, gen_latent)
if not islatent:
obs_var[atom_ind][1].append(c_ls)
cur_threshold *= prob_decay
succ = succ and atom_succ
obs_list.append(not islatent)
if succ:
sample_buff[atom_ind].append(c_ls)
latent_mask_buff[atom_ind].append(1 if islatent else 0)
neg_mask_buff[atom_ind].append(0 if atom.neg else 1)
if succ and any(obs_list):
for i in range(len(rule.atom_ls)):
samples[i][1].extend(sample_buff[i])
latent_mask[i][1].extend(latent_mask_buff[i])
neg_mask[i][1].extend(neg_mask_buff[i])
cnt += 1
yield samples, neg_mask, latent_mask, obs_var
def get_batch_by_q(self, batchsize, observed_prob=1.0, validation=False):
samples_by_r = [[[atom.pred_name, []] for atom in rule.atom_ls] for rule in self.rule_ls]
neg_mask_by_r = [[[atom.pred_name, []] for atom in rule.atom_ls] for rule in self.rule_ls]
latent_mask_by_r = [[[atom.pred_name, []] for atom in rule.atom_ls] for rule in self.rule_ls]
obs_var_by_r = [[[atom.pred_name, []] for atom in rule.atom_ls] for rule in self.rule_ls]
neg_var_by_r = [[[atom.pred_name, []] for atom in rule.atom_ls] for rule in self.rule_ls]
cnt = 0
num_ents = len(self.const2ind)
ind2const = self.const_sort_dict['type']
def gen_fake(c1, c2, pn):
for _ in range(10):
c1_fake = random.randint(0, num_ents - 1)
c2_fake = random.randint(0, num_ents - 1)
if np.random.rand( ) > 0.5:
if ind2const[c1_fake] not in self.ht_dict_train[pn][1][ind2const[c2]]:
return c1_fake, c2
else:
if ind2const[c2_fake] not in self.ht_dict_train[pn][0][ind2const[c1]]:
return c1, c2_fake
return None, None
if validation:
fact_ls = self.valid_fact_ls
else:
fact_ls = self.test_fact_ls
for val, pred_name, consts in fact_ls:
for rule_i, rule in enumerate(self.rule_ls):
# find rule with pred_name as head
if rule.atom_ls[-1].pred_name != pred_name:
continue
samples = samples_by_r[rule_i]
neg_mask = neg_mask_by_r[rule_i]
latent_mask = latent_mask_by_r[rule_i]
obs_var = obs_var_by_r[rule_i]
neg_var = neg_var_by_r[rule_i]
var2ind = rule.key2ind
var2type = rule.rule_vars
sub = [None] * len(rule.rule_vars) # substitutions
vn0, vn1 = rule.atom_ls[-1].var_name_ls
sub[var2ind[vn0]] = consts[0]
sub[var2ind[vn1]] = consts[1]
sample_buff = [[] for _ in rule.atom_ls]
neg_mask_buff = [[] for _ in rule.atom_ls]
latent_mask_buff = [[] for _ in rule.atom_ls]
atom_inds = list(range(len(rule.atom_ls) - 1))
shuffle(atom_inds)
succ = True
obs_list = []
for atom_ind in atom_inds:
atom = rule.atom_ls[atom_ind]
pred_ht_dict = self.ht_dict_train[atom.pred_name]
gen_latent = np.random.rand( ) > observed_prob
c_ls, islatent, atom_succ = self._inst_var(sub, var2ind, var2type,
atom, pred_ht_dict, gen_latent)
assert atom_succ
if not islatent:
obs_var[atom_ind][1].append(c_ls)
c1, c2 = gen_fake(c_ls[0], c_ls[1], atom.pred_name)
if c1 is not None:
neg_var[atom_ind][1].append([c1, c2])
succ = succ and atom_succ
obs_list.append(not islatent)
sample_buff[atom_ind].append(c_ls)
latent_mask_buff[atom_ind].append(1 if islatent else 0)
neg_mask_buff[atom_ind].append(0 if atom.neg else 1)
if succ and any(obs_list):
for i in range(len(rule.atom_ls)):
samples[i][1].extend(sample_buff[i])
latent_mask[i][1].extend(latent_mask_buff[i])
neg_mask[i][1].extend(neg_mask_buff[i])
samples[-1][1].append([self.const2ind[consts[0]], self.const2ind[consts[1]]])
latent_mask[-1][1].append(1)
neg_mask[-1][1].append(1)
cnt += 1
if cnt >= batchsize:
yield samples_by_r, latent_mask_by_r, neg_mask_by_r, obs_var_by_r, neg_var_by_r
samples_by_r = [[[atom.pred_name, []] for atom in rule.atom_ls] for rule in self.rule_ls]
neg_mask_by_r = [[[atom.pred_name, []] for atom in rule.atom_ls] for rule in self.rule_ls]
latent_mask_by_r = [[[atom.pred_name, []] for atom in rule.atom_ls] for rule in self.rule_ls]
obs_var_by_r = [[[atom.pred_name, []] for atom in rule.atom_ls] for rule in self.rule_ls]
neg_var_by_r = [[[atom.pred_name, []] for atom in rule.atom_ls] for rule in self.rule_ls]
cnt = 0
yield samples_by_r, latent_mask_by_r, neg_mask_by_r, obs_var_by_r, neg_var_by_r
def get_batch_by_q_v2(self, batchsize, observed_prob=1.0):
samples_by_r = [[[atom.pred_name, []] for atom in rule.atom_ls] for rule in self.rule_ls]
neg_mask_by_r = [[[atom.pred_name, []] for atom in rule.atom_ls] for rule in self.rule_ls]
latent_mask_by_r = [[[atom.pred_name, []] for atom in rule.atom_ls] for rule in self.rule_ls]
obs_var_by_r = [[[atom.pred_name, []] for atom in rule.atom_ls] for rule in self.rule_ls]
neg_var_by_r = [[[atom.pred_name, []] for atom in rule.atom_ls] for rule in self.rule_ls]
cnt = 0
num_ents = len(self.const2ind)
ind2const = self.const_sort_dict['type']
def gen_fake(c1, c2, pn):
for _ in range(10):
c1_fake = random.randint(0, num_ents - 1)
c2_fake = random.randint(0, num_ents - 1)
if np.random.rand( ) > 0.5:
if ind2const[c1_fake] not in self.ht_dict_train[pn][1][ind2const[c2]]:
return c1_fake, c2
else:
if ind2const[c2_fake] not in self.ht_dict_train[pn][0][ind2const[c1]]:
return c1, c2_fake
return None, None
for val, pred_name, consts in self.test_fact_ls:
for rule_i, rule in enumerate(self.rule_ls):
# find rule with pred_name as head
if rule.atom_ls[-1].pred_name != pred_name:
continue
samples = samples_by_r[rule_i]
neg_mask = neg_mask_by_r[rule_i]
latent_mask = latent_mask_by_r[rule_i]
var2ind = rule.key2ind
var2type = rule.rule_vars
sub_ls = [[None for _ in range(len(rule.rule_vars))] for _ in range(2)] # substitutions
vn0, vn1 = rule.atom_ls[-1].var_name_ls
sub_ls[0][var2ind[vn0]] = consts[0]
sub_ls[0][var2ind[vn1]] = consts[1]
c1, c2 = gen_fake(self.const2ind[consts[0]], self.const2ind[consts[1]], pred_name)
if c1 is not None:
sub_ls[1][var2ind[vn0]] = ind2const[c1]
sub_ls[1][var2ind[vn1]] = ind2const[c2]
else:
sub_ls.pop(1)
pos_query_succ = False
for sub_ind, sub in enumerate(sub_ls):
sample_buff = [[] for _ in rule.atom_ls]
neg_mask_buff = [[] for _ in rule.atom_ls]
latent_mask_buff = [[] for _ in rule.atom_ls]
atom_inds = list(range(len(rule.atom_ls) - 1))
shuffle(atom_inds)
succ = True
obs_list = []
for atom_ind in atom_inds:
atom = rule.atom_ls[atom_ind]
pred_ht_dict = self.ht_dict_train[atom.pred_name]
gen_latent = np.random.rand( ) > observed_prob
if sub_ind == 1:
gen_latent = np.random.rand( ) > 0.5
c_ls, islatent, atom_succ = self._inst_var(sub, var2ind, var2type,
atom, pred_ht_dict, gen_latent)
assert atom_succ
succ = succ and atom_succ
obs_list.append(not islatent)
sample_buff[atom_ind].append(c_ls)
latent_mask_buff[atom_ind].append(1 if islatent else 0)
neg_mask_buff[atom_ind].append(0 if atom.neg else 1)
if succ:
if any(obs_list) or ((sub_ind == 1) and pos_query_succ):
for i in range(len(rule.atom_ls)):
samples[i][1].extend(sample_buff[i])
latent_mask[i][1].extend(latent_mask_buff[i])
neg_mask[i][1].extend(neg_mask_buff[i])
if sub_ind == 0:
samples[-1][1].append([self.const2ind[consts[0]], self.const2ind[consts[1]]])
latent_mask[-1][1].append(1)
neg_mask[-1][1].append(1)
pos_query_succ = True
cnt += 1
else:
samples[-1][1].append([c1, c2])
latent_mask[-1][1].append(0) # sample a negative fact at head
neg_mask[-1][1].append(1)
if cnt >= batchsize:
yield samples_by_r, latent_mask_by_r, neg_mask_by_r, obs_var_by_r, neg_var_by_r
samples_by_r = [[[atom.pred_name, []] for atom in rule.atom_ls] for rule in self.rule_ls]
neg_mask_by_r = [[[atom.pred_name, []] for atom in rule.atom_ls] for rule in self.rule_ls]
latent_mask_by_r = [[[atom.pred_name, []] for atom in rule.atom_ls] for rule in self.rule_ls]
obs_var_by_r = [[[atom.pred_name, []] for atom in rule.atom_ls] for rule in self.rule_ls]
neg_var_by_r = [[[atom.pred_name, []] for atom in rule.atom_ls] for rule in self.rule_ls]
cnt = 0
yield samples_by_r, latent_mask_by_r, neg_mask_by_r, obs_var_by_r, neg_var_by_r
def get_batch_rnd(self, observed_prob=0.7, filter_latent=True, closed_world=False, filter_observed=False):
"""
return a batch of gnd formulae by random sampling with controllable bias towards those containing
observed variables. The overall sampling logic is that:
1) rnd sample a rule from rule_ls
2) shuffle the predicates contained in the rule
3) for each of these predicates, with (observed_prob) it will be instantiated as observed variable, and
for (1-observed_prob) if will be simply uniformly instantiated.
3.1) if observed var, then sample from the knowledge base, which is self.fact_dict, if failed for any
reason, go to 3.2)
3.2) if uniformly sample, then for each logic variable in the predicate, instantiate it with a uniform
sample from the corresponding constant dict
:param observed_prob:
probability of instantiating a predicate as observed variable
:param filter_latent:
filter out ground formula containing only latent vars
:param closed_world:
if set True, reduce the sampling space of all predicates not in the test_dict to the set specified in
fact_dict
:param filter_observed:
filter out ground formula containing only observed vars
:return:
"""
batch_neg_mask = [[] for _ in range(len(self.rule_ls))]
batch_latent_var_inds = [[] for _ in range(len(self.rule_ls))]
batch_observed_vars = [[] for _ in range(len(self.rule_ls))]
observed_rule_cnts = [0.0 for _ in range(len(self.rule_ls))]
flat_latent_vars = dict( )
cnt = 0
inds = list(range(len(self.rule_ls)))
while cnt < self.batchsize:
# randomly sample a formula
if self.shuffle_sampling:
shuffle(inds)
for ind in inds:
rule = self.rule_ls[ind]
atom_key_dict = self.atom_key_dict_ls[ind]
sub = [None] * len(rule.rule_vars) # substitutions
# randomly sample an atom from the formula
atom_inds = list(range(len(rule.atom_ls)))
shuffle(atom_inds)
for atom_ind in atom_inds:
atom = rule.atom_ls[atom_ind]
atom_dict = atom_key_dict[atom.pred_name]
# instantiate the predicate
self._instantiate_pred(atom, atom_dict, sub, rule, observed_prob)
# if variable substitution is complete already then exit
if not (None in sub):
break
# generate latent and observed var labels and their negation masks
latent_vars, observed_vars, \
latent_neg_mask, observed_neg_mask = self._gen_mask(rule, sub, closed_world)
# check sampled ground rule status
stat_code = self._get_rule_stat(observed_vars, latent_vars, observed_neg_mask,
filter_latent, filter_observed)
# is a valid sample with only observed vars and does not have negation on all of them
if stat_code == FULL_OBSERVERED:
observed_rule_cnts[ind] += 1
cnt += 1
# is a valid sample
elif stat_code == GOOD:
batch_neg_mask[ind].append([latent_neg_mask, observed_neg_mask])
for latent_var in latent_vars:
if latent_var not in flat_latent_vars:
flat_latent_vars[latent_var] = len(flat_latent_vars)
batch_latent_var_inds[ind].append([flat_latent_vars[e] for e in latent_vars])
batch_observed_vars[ind].append(observed_vars)
cnt += 1
# not a valid sample
else:
continue
if cnt >= self.batchsize:
break
flat_list = sorted([(k, v) for k, v in flat_latent_vars.items( )], key=lambda x: x[1])
flat_list = [e[0] for e in flat_list]
return batch_neg_mask, flat_list, batch_latent_var_inds, observed_rule_cnts, batch_observed_vars
def reset(self):
self.rule_gens = [self.generate_gnd_rule(rule) for rule in self.rule_ls]
def get_stats(self):
num_ents = sum([len(v) for k, v in self.const_sort_dict.items( )])
num_rels = len(self.PRED_DICT)
num_facts = sum([len(v) for k, v in self.fact_dict.items( )])
num_queries = len(self.test_fact_ls)
num_gnd_atom = 0
for pred_name, pred in self.PRED_DICT.items( ):
cnt = 1
for var_type in pred.var_types:
cnt *= len(self.const_sort_dict[var_type])
num_gnd_atom += cnt
num_gnd_rule = 0
for rule in self.rule_ls:
cnt = 1
for var_type in rule.rule_vars.values( ):
cnt *= len(self.const_sort_dict[var_type])
num_gnd_rule += cnt
return num_ents, num_rels, num_facts, num_queries, num_gnd_atom, num_gnd_rule
def preprocess_kinship(self, ppath, fpath, rpath, qpath):
"""
:param ppath:
predicate file path
:param fpath:
facts file path
:param rpath:
rule file path
:param qpath:
query file path
:return:
"""
assert all(map(isfile, [ppath, fpath, rpath, qpath]))
strip_items = lambda ls: list(map(lambda x: x.strip( ), ls))
pred_reg = re.compile(r'(.*)\((.*)\)')
with open(ppath) as f:
for line in f:
# skip empty lines
if line.strip( ) == '':
continue
m = pred_reg.match(line.strip( ))
assert m is not None, 'matching predicate failed for %s' % line
name, var_types = m.group(1), m.group(2)
var_types = list(map(lambda x: x.strip( ), var_types.split(',')))
self.PRED_DICT[name] = Predicate(name, var_types)
TYPE_SET.update(var_types)
fact_ls = []
fact_reg = re.compile(r'(!?)(.*)\((.*)\)')
with open(fpath) as f:
for line in f:
# skip empty lines
if line.strip( ) == '':
continue
m = fact_reg.match(line.strip( ))
assert m is not None, 'matching fact failed for %s' % line
val = 0 if m.group(1) == '!' else 1
name, consts = m.group(2), m.group(3)
consts = strip_items(consts.split(','))
fact_ls.append(Fact(name, consts, val))
for var_type in self.PRED_DICT[name].var_types:
self.const_dict.add_const(var_type, consts.pop(0))
rule_ls = []
first_atom_reg = re.compile(r'([\d.]+) (!?)([\w\d]+)\((.*)\)')
atom_reg = re.compile(r'(!?)([\w\d]+)\((.*)\)')
with open(rpath) as f:
for line in f:
# skip empty lines
if line.strip( ) == '':
continue
atom_str_ls = strip_items(line.strip( ).split(' v '))
assert len(atom_str_ls) > 1, 'rule length must be greater than 1, but get %s' % line
atom_ls = []
rule_weight = 0.0
for i, atom_str in enumerate(atom_str_ls):
if i == 0:
m = first_atom_reg.match(atom_str)
assert m is not None, 'matching atom failed for %s' % atom_str
rule_weight = float(m.group(1))
neg = m.group(2) == '!'
pred_name = m.group(3).strip( )
var_name_ls = strip_items(m.group(4).split(','))
else:
m = atom_reg.match(atom_str)
assert m is not None, 'matching atom failed for %s' % atom_str
neg = m.group(1) == '!'
pred_name = m.group(2).strip( )
var_name_ls = strip_items(m.group(3).split(','))
atom = Atom(neg, pred_name, var_name_ls, self.PRED_DICT[pred_name].var_types)
atom_ls.append(atom)
rule = Formula(atom_ls, rule_weight)
rule_ls.append(rule)
query_ls = []
with open(qpath) as f:
for line in f:
# skip empty lines
if line.strip( ) == '':
continue
m = fact_reg.match(line.strip( ))
assert m is not None, 'matching fact failed for %s' % line
val = 0 if m.group(1) == '!' else 1
name, consts = m.group(2), m.group(3)
consts = strip_items(consts.split(','))
query_ls.append(Fact(name, consts, val))
for var_type in self.PRED_DICT[name].var_types:
self.const_dict.add_const(var_type, consts.pop(0))
return fact_ls, rule_ls, query_ls
TYPE_SET = set( )
def iterline(fpath):
with open(fpath) as f:
for line in f:
line = line.strip( )
if line == '':
continue
yield line
class ConstantDict:
def __init__(self):
self.constants = {}
def add_const(self, const_type, const):
"""
:param const_type:
string
:param const:
string
"""
# if const_type not in TYPE_DICT:
# TYPE_DICT[const_type] = len(TYPE_DICT)
if const_type in self.constants:
self.constants[const_type].add(const)
else:
self.constants[const_type] = {const}
def __getitem__(self, key):
return self.constants[key]
def has_const(self, key, const):
if key in self.constants:
return const in self[key]
else:
return False
class Predicate:
def __init__(self, name, var_types):
"""
:param name:
string
:param var_types:
list of strings
"""
self.name = name
self.var_types = var_types
self.num_args = len(var_types)
def __repr__(self):
return '%s(%s)' % (self.name, ','.join(self.var_types))
class Fact:
def __init__(self, pred_name, const_ls, val):
self.pred_name = pred_name
self.const_ls = deepcopy(const_ls)
self.val = val
def __repr__(self):
return self.pred_name + '(%s)' % ','.join(self.const_ls)
class Atom:
def __init__(self, neg, pred_name, var_name_ls, var_type_ls):
self.neg = neg
self.pred_name = pred_name
self.var_name_ls = var_name_ls
self.var_type_ls = var_type_ls
def __repr__(self):
return ('!' if self.neg else '') + self.pred_name + '(%s)' % ','.join(self.var_name_ls)
class Formula:
"""
only support clause form with disjunction, e.g. !
"""
def __init__(self, atom_ls, weight):
self.weight = weight
self.atom_ls = atom_ls
self.rule_vars = dict( )
for atom in self.atom_ls:
self.rule_vars.update(zip(atom.var_name_ls, atom.var_type_ls))
self.key2ind = dict(zip(self.rule_vars.keys( ), range(len(self.rule_vars.keys( )))))
def evaluate(self):
pass
def __repr__(self):
return ' v '.join(list(map(repr, self.atom_ls)))
class ConstantDict:
def __init__(self):
self.constants = {}
def add_const(self, const_type, const):
"""
:param const_type:
string
:param const:
string
"""
# if const_type not in TYPE_DICT:
# TYPE_DICT[const_type] = len(TYPE_DICT)
if const_type in self.constants:
self.constants[const_type].add(const)
else:
self.constants[const_type] = {const}
def __getitem__(self, key):
return self.constants[key]
def has_const(self, key, const):
if key in self.constants:
return const in self[key]
else:
return False
@register_dataset('NBF_link_prediction')
class NBF_LinkPrediction(LinkPredictionDataset):
r"""
The NBF dataset will be used in task *link prediction*.
"""
def __init__(self, dataset_name ,*args, **kwargs): # dataset_name in ['NBF_WN18RR','NBF_FB15k-237']
self.dataset = NBF_Dataset(root='./openhgnn/dataset/', name=dataset_name[4:], version="v1")
import os
import requests
import zipfile
import io
@register_dataset('DisenKGAT_link_prediction')
class DisenKGAT_LinkPrediction(LinkPredictionDataset):
def __init__(self, dataset ,*args, **kwargs): # dataset "DisenKGAT"
self.logger = kwargs.get("Logger")
self.args = kwargs.get("args")
self.current_dir = os.path.dirname(os.path.abspath(__file__))
self.dataset_name = dataset
self.raw_dir = os.path.join(self.current_dir, self.dataset_name ,"raw_dir" )
self.processed_dir = os.path.join(self.current_dir, self.dataset_name ,"processed_dir" )
if not os.path.exists(self.raw_dir):
os.makedirs(self.raw_dir)
self.download()
else:
print("raw_dir already exists")
def download(self):
url = "https://s3.cn-north-1.amazonaws.com.cn/dgl-data/dataset/openhgnn/{}.zip".format(self.dataset_name)
response = requests.get(url)
with zipfile.ZipFile(io.BytesIO(response.content)) as myzip:
myzip.extractall(self.raw_dir)
print("--- download finished---")