directory = os.path.join('results', args.model)
if not os.path.exists(directory):
os.makedirs(directory)
args.out_dir = directory
args.perf_file = os.path.join(
directory, '_'.join([dataset, args.sample, args.update]) +
args.out_file_info + '.txt')
args.stat_file = os.path.join(
directory, '_'.join([dataset, args.sample, args.update]) + '.stat')
print('output file name:', args.perf_file, args.stat_file)
logger_init(args)
task_dir = args.task_dir
loader = DataLoader(task_dir, args.N_1)
n_ent, n_rel = loader.graph_size()
train_data = loader.load_data('train')
valid_data = loader.load_data('valid')
test_data = loader.load_data('test')
args.n_train = len(train_data[0])
print("Number of train:{}, valid:{}, test:{}.".format(
len(train_data[0]), len(valid_data[0]), len(test_data[0])))
plot_config(args)
heads, tails = loader.heads_tails()
head_idx, tail_idx, head_cache, tail_cache, head_pos, tail_pos = loader.get_cache_list(
)
parser.add_argument('--test_batch_size', type=int, default=100, help='test batch size')
parser.add_argument('--filter', type=bool, default=True, help='whether do filter in testing')
parser.add_argument('--mode', type=str, default='search', help='which mode this code is running for')
parser.add_argument('--out_file_info', type=str, default='', help='extra string for the output file name')
args = parser.parse_args()
dataset = args.task_dir.split('/')[-1]
directory = 'results'
if not os.path.exists(directory):
os.makedirs(directory)
args.out_dir = directory
loader = DataLoader(args.task_dir)
n_ent, n_rel = loader.graph_size()
train_data = loader.load_data('train')
valid_data = loader.load_data('valid')
test_data = loader.load_data('test')
n_train = len(train_data[0])
valid_head_filter, valid_tail_filter, test_head_filter, test_tail_filter = loader.get_filter()
train_data = [torch.LongTensor(vec) for vec in train_data]
valid_data = [torch.LongTensor(vec) for vec in valid_data]
test_data = [torch.LongTensor(vec) for vec in test_data]
def run_model(i, state):
print('newID:', i, state, len(state))
default='base.txt',
help='log prefix')
parser.add_argument('--log_to_file',
type=bool,
default=False,
help='log to file')
parser.add_argument('--log_dir',
type=str,
default='./log',
help='log save dir')
parser.add_argument('--log_prefix', type=str, default='', help='log prefix')
args = parser.parse_args()
filename = os.path.join(args.task_dir, 'TransE.mdl')
loader = DataLoader(args.task_dir, 50)
n_ent, n_rel = loader.graph_size()
train_data = loader.load_data('train')
test_data = loader.load_data('test')
train_data = [torch.LongTensor(vec) for vec in train_data]
test_data = [torch.LongTensor(vec) for vec in test_data]
heads, tails = loader.heads_tails()
model = TransEModule(n_ent, n_rel, args)
model.cuda()
model.load_state_dict(
torch.load(filename,
map_location=lambda storage, location: storage.cuda()))
head, tail, rela = train_data
n_train = len(head)