python dataset[trans_TestDatasetTF.py

import os

import pandas as pd

class KnowledgeGraph:

def __init__(self, data_dir):

# 考虑到tf的各项api使用，Python不能将Tensor类型直接转换成字符串类型，但是可以将TF类型转换成numpy类型

# 所以这里的训练三元组，测试三元组等等，都是id三元组，而不是字符串三元组

self.data_dir = data_dir

self.entity_dict = {}

self.entities = []

self.relation_dict = {}

self.n_entity = 0

self.n_relation = 0

self.training_triples = [] # list of triples in the form of (h, t, r)

self.validation_triples = []

self.test_triples = []

self.n_training_triple = 0

self.n_validation_triple = 0

self.n_test_triple = 0

'''load dicts and triples'''

self.load_dicts()

self.load_triples()

'''construct pools after loading'''

self.training_triple_pool = set(self.training_triples)

self.golden_triple_pool = set(

self.training_triples) | set(

self.validation_triples) | set(

self.test_triples)

def load_dicts(self):

entity_dict_file = 'entity2id.txt'

relation_dict_file = 'relation2id.txt'

print('-----Loading entity dict-----')

entity_df = pd.read_table(

os.path.join(

self.data_dir,

entity_dict_file),

header=None)

self.entity_dict = dict(zip(entity_df[0], entity_df[1]))

self.n_entity = len(self.entity_dict)

self.entities = list(self.entity_dict.values())

print('#entity: {}'.format(self.n_entity))

print('-----Loading relation dict-----')

relation_df = pd.read_table(

os.path.join(

self.data_dir,

relation_dict_file),

header=None)

self.relation_dict = dict(zip(relation_df[0], relation_df[1]))

self.n_relation = len(self.relation_dict)

print('#relation: {}'.format(self.n_relation))

def load_triples(self):

training_file = 'train.txt'

validation_file = 'valid.txt'

test_file = 'test.txt'

print('-----Loading training triples-----')

training_df = pd.read_table(

os.path.join(

self.data_dir,

training_file),

header=None)

self.training_triples = list(zip([self.entity_dict[h] for h in training_df[0]],

[self.entity_dict[t] for t in training_df[1]],

[self.relation_dict[r] for r in training_df[2]]))

self.n_training_triple = len(self.training_triples)

print('#training triple: {}'.format(self.n_training_triple))

print('-----Loading validation triples-----')

validation_df = pd.read_table(

os.path.join(

self.data_dir,

validation_file),

header=None)

self.validation_triples = list(zip([self.entity_dict[h] for h in validation_df[0]],

[self.entity_dict[t] for t in validation_df[1]],

[self.relation_dict[r] for r in validation_df[2]]))

self.n_validation_triple = len(self.validation_triples)

print('#validation triple: {}'.format(self.n_validation_triple))

print('-----Loading test triples------')

test_df = pd.read_table(

os.path.join(

self.data_dir,

test_file),

header=None)

self.test_triples = list(zip([self.entity_dict[h] for h in test_df[0]],

[self.entity_dict[t] for t in test_df[1]],

[self.relation_dict[r] for r in test_df[2]]))

self.n_test_triple = len(self.test_triples)

print('#test triple: {}'.format(self.n_test_triple))

一键复制

编辑

Web IDE

原始数据

按行查看

历史