"""

def __init__(self, dim, gpus=[], cpu_per_gpu=auto, gpu_memory_limit=auto,

float_type=cfg.float_type, index_type=cfg.index_type):

self.dim = dim

self.gpus = gpus

self.cpu_per_gpu = cpu_per_gpu

self.gpu_memory_limit = gpu_memory_limit

self.float_type = float_type

self.index_type = index_type

self.set_format()

def get_graph(self, **kwargs):

raise NotImplementedError

def get_solver(self, **kwargs):

raise NotImplementedError

def set_format(self, delimiters=" \t\r\n", comment="#"):

"""

self.delimiters = delimiters

self.comment = comment

self.pattern = re.compile("[%s]" % self.delimiters)

@monitor.time

def load(self, **kwargs):

"""load(**kwargs)

self.graph = self.get_graph(**kwargs)

if "file_name" in kwargs or "vector_file" in "kwargs":

self.graph.load(delimiters=self.delimiters, comment=self.comment, **kwargs)

else:

self.graph.load(**kwargs)

@monitor.time

def build(self, **kwargs):

"""build(**kwargs)

self.solver = self.get_solver(**kwargs)

self.solver.build(self.graph, **kwargs)

@monitor.time

def train(self, **kwargs):

"""train(**kwargs)

self.solver.train(**kwargs)

@monitor.time

def evaluate(self, task, **kwargs):

"""evaluate(task, **kwargs)

func_name = task.replace(" ", "_")

if not hasattr(self, func_name):

raise ValueError("Unknown task `%s`" % task)

logger.info(lib.io.header(task))

result = getattr(self, func_name)(**kwargs)

if isinstance(result, dict):

for metric, value in sorted(result.items()):

logger.warning("%s: %g" % (metric, value))

return result

@monitor.time

def save(self, file_name):

"""

logger.warning("save embeddings and name mappings to `%s`" % file_name)

objects = EasyDict()

for name in dir(self.graph):

object = getattr(self.graph, name)

if isinstance(object, list) and len(object) > 0 and isinstance(object[0], str): # id2name

objects[name] = object

for name in dir(self.solver):

object = getattr(self.solver, name)

if isinstance(object, np.ndarray): # embedding

objects[name] = object

with open(file_name, "wb") as fout:

pickle.dump(objects, fout, protocol=pickle.HIGHEST_PROTOCOL)

def tokenize(self, str):

str = str.strip(self.delimiters)

comment_start = str.find(self.comment)

if comment_start != -1:

str = str[:comment_start]

return self.pattern.split(str)

def name_map(self, dicts, names):

assert len(dicts) == len(names), "The number of dictionaries and names must be equal"

indexes = [[] for _ in range(len(names))]

num_param = len(names)

num_sample = len(names[0])

for i in range(num_sample):

valid = True

for j in range(num_param):

if names[j][i] not in dicts[j]:

valid = False

break

if valid:

for j in range(num_param):

indexes[j].append(dicts[j][names[j][i]])

return indexes

def gpu_map(self, func, settings):

import torch

gpus = self.gpus if self.gpus else range(torch.cuda.device_count())

new_settings = []

for i, setting in enumerate(settings):

new_settings.append(setting + (gpus[i % len(gpus)],))

settings = new_settings

try:

start_method = multiprocessing.get_start_method()

# if there are other running processes, this could cause leakage of semaphores

multiprocessing.set_start_method("spawn", force=True)

pool = multiprocessing.Pool(len(gpus))

results = pool.map(func, settings, chunksize=1)

multiprocessing.set_start_method(start_method, force=True)

except AttributeError:

logger.info("Spawn mode is not supported by multiprocessing. Switch to serial execution.")

results = list(map(func, settings))

"""

def get_graph(self, **kwargs):

return graph.Graph(self.index_type)

def get_solver(self, **kwargs):

if self.cpu_per_gpu == auto:

num_sampler_per_worker = auto

else:

num_sampler_per_worker = self.cpu_per_gpu - 1

return solver.GraphSolver(self.dim, self.float_type, self.index_type, self.gpus, num_sampler_per_worker,

self.gpu_memory_limit)

def node_classification(self, X=None, Y=None, file_name=None, portions=(0.02,), normalization=False, times=1,

patience=100):

"""

import scipy.sparse as sp

self.solver.clear()

if file_name:

if not (X is None and Y is None):

raise ValueError("Evaluation data and file should not be provided at the same time")

X = []

Y = []

with open(file_name, "r") as fin:

for line in fin:

tokens = self.tokenize(line)

if len(tokens) == 0:

continue

x, y = tokens

X.append(x)

Y.append(y)

if X is None or Y is None:

raise ValueError("Either evaluataion data (X, Y) or a file name should be provided")

name2id = self.graph.name2id

class2id = {c:i for i, c in enumerate(np.unique(Y))}

new_X, new_Y = self.name_map((name2id, class2id), (X, Y))

logger.info("effective labels: %d / %d" % (len(new_X), len(X)))

X = np.asarray(new_X)

Y = np.asarray(new_Y)

labels = sp.coo_matrix((np.ones_like(X), (X, Y)), dtype=np.int32).todense()

indexes, _ = np.where(np.sum(labels, axis=1) > 0)

# discard non-labeled nodes

labels = labels[indexes]

vertex_embeddings = SharedNDArray(self.solver.vertex_embeddings[indexes])

settings = []

for portion in portions:

settings.append((vertex_embeddings, labels, portion, normalization, times, patience))

results = self.gpu_map(linear_classification, settings)

metrics = {}

for result in results:

metrics.update(result)

return metrics

def link_prediction(self, H=None, T=None, Y=None, file_name=None, filter_H=None, filter_T=None, filter_file=None):

"""

import torch

from .network import LinkPredictor

self.solver.clear()

if file_name:

if not (H is None and T is None and Y is None):

raise ValueError("Evaluation data and file should not be provided at the same time")

H = []

T = []

Y = []

with open(file_name, "r") as fin:

for line in fin:

tokens = self.tokenize(line)

if len(tokens) == 0:

continue

h, t, y = tokens

H.append(h)

T.append(t)

Y.append(y)

if H is None or T is None or Y is None:

raise ValueError("Either evaluation data or file should be provided")

if filter_file:

if not (filter_H is None and filter_T is None):

raise ValueError("Filter data and file should not be provided at the same time")

filter_H = []

filter_T = []

with open(filter_file, "r") as fin:

for line in fin:

tokens = self.tokenize(line)

if len(tokens) == 0:

continue

h, t = tokens

filter_H.append(h)

filter_T.append(t)

elif filter_H is None:

filter_H = []

filter_T = []

name2id = self.graph.name2id

Y = [int(y) for y in Y]

new_H, new_T, new_Y = self.name_map((name2id, name2id, {0:0, 1:1}), (H, T, Y))

logger.info("effective edges: %d / %d" % (len(new_H), len(H)))

H = new_H

T = new_T

Y = new_Y

new_H, new_T = self.name_map((name2id, name2id), (filter_H, filter_T))

logger.info("effective filter edges: %d / %d" % (len(new_H), len(filter_H)))

filters = set(zip(new_H, new_T))

new_H = []

new_T = []

new_Y = []

for h, t, y in zip(H, T, Y):

if (h, t) not in filters:

new_H.append(h)

new_T.append(t)

new_Y.append(y)

logger.info("remaining edges: %d / %d" % (len(new_H), len(H)))

H = np.asarray(new_H)

T = np.asarray(new_T)

Y = np.asarray(new_Y)

vertex_embeddings = self.solver.vertex_embeddings

context_embeddings = self.solver.context_embeddings

model = LinkPredictor(self.solver.model, vertex_embeddings, context_embeddings)

model = model.cuda()

H = torch.as_tensor(H)

T = torch.as_tensor(T)

Y = torch.as_tensor(Y)

H = H.cuda()

T = T.cuda()

Y = Y.cuda()

score = model(H, T)

order = torch.argsort(score, descending=True)

Y = Y[order]

hit = torch.cumsum(Y, dim=0)

all = torch.sum(Y == 0) * torch.sum(Y == 1)

auc = torch.sum(hit[Y == 0]).item() / all.item()

return {

"AUC": auc

import torch

from torch import optim

from torch.nn import functional as F

from .network import NodeClassifier

def generate_one_vs_rest(indexes, labels):

new_indexes = []

new_labels = []

num_class = labels.shape[1]

for index, sample_labels in zip(indexes, labels):

for cls in np.where(sample_labels)[0]:

new_indexes.append(index)

new_label = np.zeros(num_class, dtype=np.int)

new_label[cls] = 1

new_labels.append(new_label)

return torch.as_tensor(new_indexes), torch.as_tensor(new_labels)

embeddings, labels, portion, normalization, times, patience, gpu = args

embeddings = np.asarray(embeddings)

num_sample, num_class = labels.shape

num_train = int(num_sample * portion)

macro_f1s = []

micro_f1s = []

for t in range(times):

samples = np.random.permutation(num_sample)

train_samples = samples[:num_train]

train_labels = np.asarray(labels[train_samples])

train_samples, train_labels = generate_one_vs_rest(train_samples, train_labels)

test_samples = torch.as_tensor(samples[num_train:])

test_labels = torch.as_tensor(labels[test_samples])

model = NodeClassifier(embeddings, num_class, normalization=normalization)

train_samples = train_samples.cuda(gpu)

train_labels = train_labels.cuda(gpu)

test_samples = test_samples.cuda(gpu)

test_labels = test_labels.cuda(gpu)

model = model.cuda(gpu)

# train

optimizer = optim.SGD(model.parameters(), lr=1, weight_decay=2e-5, momentum=0.9)

best_loss = float("inf")

best_epoch = -1

for epoch in range(100000):

optimizer.zero_grad()

logits = model(train_samples)

loss = F.binary_cross_entropy_with_logits(logits, train_labels.float())

loss.backward()

optimizer.step()

loss = loss.item()

if loss < best_loss:

best_epoch = epoch

best_loss = loss

if epoch == best_epoch + patience:

break

# test

logits = model(test_samples)

num_labels = test_labels.sum(dim=1, keepdim=True)

sorted, _ = logits.sort(dim=1, descending=True)

thresholds = sorted.gather(dim=1, index=num_labels-1)

predictions = (logits >= thresholds).int()

# compute metric

num_TP_per_class = (predictions & test_labels).sum(dim=0).float()

num_T_per_class = test_labels.sum(dim=0).float()

num_P_per_class = predictions.sum(dim=0).float()

macro_f1s.append((2 * num_TP_per_class / (num_T_per_class + num_P_per_class)).mean().item())

num_TP = (predictions & test_labels).sum().float()

num_T = test_labels.sum().float()

num_P = predictions.sum().float()

micro_f1s.append((2 * num_TP / (num_T + num_P)).item())

return {

"macro-F1@%g%%" % (portion * 100): np.mean(macro_f1s),

"micro-F1@%g%%" % (portion * 100): np.mean(micro_f1s)

"""

def get_graph(self, **kwargs):

return graph.WordGraph(self.index_type)

def get_solver(self, **kwargs):

if self.cpu_per_gpu == auto:

num_sampler_per_worker = auto

else:

num_sampler_per_worker = self.cpu_per_gpu - 1

return solver.GraphSolver(self.dim, self.float_type, self.index_type, self.gpus, num_sampler_per_worker,

"""

SAMPLE_PER_DIMENSION = 7

MEMORY_SCALE_FACTOR = 1.5

def get_graph(self, **kwargs):

return graph.KnowledgeGraph(self.index_type)

def get_solver(self, **kwargs):

if self.cpu_per_gpu == auto:

num_sampler_per_worker = auto

else:

num_sampler_per_worker = self.cpu_per_gpu - 1

return solver.KnowledgeGraphSolver(self.dim, self.float_type, self.index_type, self.gpus, num_sampler_per_worker,

self.gpu_memory_limit)

def entity_prediction(self, H=None, R=None, T=None, file_name=None, save_file=None, target="tail", k=10,

backend=cfg.backend):

"""

def torch_predict():

import torch

entity_embeddings = SharedNDArray(self.solver.entity_embeddings)

relation_embeddings = SharedNDArray(self.solver.relation_embeddings)

num_gpu = len(self.gpus) if self.gpus else torch.cuda.device_count()

work_load = (num_sample + num_gpu - 1) // num_gpu

settings = []

for i in range(num_gpu):

work_H = H[work_load * i: work_load * (i+1)]

work_R = R[work_load * i: work_load * (i+1)]

work_T = T[work_load * i: work_load * (i+1)]

settings.append((entity_embeddings, relation_embeddings, work_H, work_R, work_T,

None, None, target, k, self.solver.model, self.solver.margin))

results = self.gpu_map(triplet_prediction, settings)

return sum(results, [])

def graphvite_predict():

num_entity = len(entity2id)

batch_size = self.get_batch_size(num_entity)

recalls = []

for i in range(0, num_sample, batch_size):

batch_h = H[i: i + batch_size]

batch_r = R[i: i + batch_size]

batch_t = T[i: i + batch_size]

batch = self.generate_one_vs_rest(batch_h, batch_r, batch_t, num_entity, target)

scores = self.solver.predict(batch)

scores = scores.reshape(-1, num_entity)

indexes = np.argpartition(scores, num_entity - k, axis=-1)

for index, score in zip(indexes, scores):

index = index[-k:]

score = score[index]

order = np.argsort(score)[::-1]

recall = list(zip(index[order], score[order]))

recalls.append(recall)

return recalls

assert_in(["head", "tail"], target=target)

assert_in(["graphvite", "torch"], backend=backend)

if backend == "torch":

self.solver.clear()

if file_name:

if not (H is None and R is None and T is None):

raise ValueError("Evaluation data and file should not be provided at the same time")

H = []

R = []

T = []

with open(file_name, "r") as fin:

for line in fin:

tokens = self.tokenize(line)

if len(tokens) == 0:

continue

if len(tokens) == 3:

h, r, t = tokens

elif target == "head":

r, t = tokens

h = None

else:

h, r = tokens

t = None

H.append(h)

R.append(r)

T.append(t)

if (H is None and T is None) or R is None:

raise ValueError("Either evaluation data or file should be provided")

if H is None:

target = "head"

if T is None:

target = "tail"

entity2id = self.graph.entity2id

relation2id = self.graph.relation2id

num_sample = len(R)

new_H = np.zeros(num_sample, dtype=np.uint32)

new_T = np.zeros(num_sample, dtype=np.uint32)

if target == "head":

new_R, new_T = self.name_map((relation2id, entity2id), (R, T))

if target == "tail":

new_H, new_R = self.name_map((entity2id, relation2id), (H, R))

assert len(new_R) == len(R), "Can't recognize some entities or relations"

H = np.asarray(new_H, dtype=np.uint32)

R = np.asarray(new_R, dtype=np.uint32)

T = np.asarray(new_T, dtype=np.uint32)

if backend == "graphvite":

recalls = graphvite_predict()

else:

recalls = torch_predict()

id2entity = self.graph.id2entity

new_recalls = []

for recall in recalls:

new_recall = [(id2entity[e], s) for e, s in recall]

new_recalls.append(new_recall)

recalls = new_recalls

if save_file:

extension = os.path.splitext(save_file)[1]

if extension == ".txt":

with open(save_file, "w") as fout:

for recall in recalls:

tokens = ["%s: %g" % x for x in recall]

fout.write("%s\n" % "\t".join(tokens))

elif extension == ".pkl":

with open(save_file, "wb") as fout:

pickle.dump(recalls, fout, protocol=pickle.HIGHEST_PROTOCOL)

else:

raise ValueError("Unknown file extension `%s`" % extension)

else:

return recalls

def link_prediction(self, H=None, R=None, T=None, filter_H=None, filter_R=None, filter_T=None, file_name=None,

filter_files=None, target="both", fast_mode=None, backend=cfg.backend):

"""

def torch_predict():

import torch

entity_embeddings = SharedNDArray(self.solver.entity_embeddings)

relation_embeddings = SharedNDArray(self.solver.relation_embeddings)

num_gpu = len(self.gpus) if self.gpus else torch.cuda.device_count()

work_load = (fast_mode + num_gpu - 1) // num_gpu

settings = []

for i in range(num_gpu):

work_H = H[work_load * i: work_load * (i+1)]

work_R = R[work_load * i: work_load * (i+1)]

work_T = T[work_load * i: work_load * (i+1)]

settings.append((entity_embeddings, relation_embeddings, work_H, work_R, work_T,

exclude_H, exclude_T, target, None, self.solver.model, self.solver.margin))

results = self.gpu_map(triplet_prediction, settings)

return np.concatenate(results)

def graphvite_predict():

num_entity = len(entity2id)

if target == "both":

batch_size = self.get_batch_size(num_entity * 2)

else:

batch_size = self.get_batch_size(num_entity)

rankings = []

for i in range(0, fast_mode, batch_size):

batch_h = H[i: i + batch_size]

batch_r = R[i: i + batch_size]

batch_t = T[i: i + batch_size]

batch = self.generate_one_vs_rest(batch_h, batch_r, batch_t, num_entity, target)

masks = self.generate_mask(batch_h, batch_r, batch_t, exclude_H, exclude_T, num_entity, target)

if target == "head":

positives = batch_h

if target == "tail":

positives = batch_t

if target == "both":

positives = np.asarray([batch_h, batch_t]).transpose()

positives = positives.ravel()

scores = self.solver.predict(batch)

scores = scores.reshape(-1, num_entity)

truths = scores[range(len(positives)), positives]

ranking = np.sum((scores >= truths[:, np.newaxis]) * masks, axis=1)

rankings.append(ranking)

return np.concatenate(rankings)

assert_in(["head", "tail", "both"], target=target)

assert_in(["graphvite", "torch"], backend=backend)

if backend == "torch":

self.solver.clear()

if file_name:

if not (H is None and R is None and T is None):

raise ValueError("Evaluation data and file should not be provided at the same time")

H = []

R = []

T = []

with open(file_name, "r") as fin:

for line in fin:

tokens = self.tokenize(line)

if len(tokens) == 0:

continue

h, r, t = tokens

H.append(h)

R.append(r)

T.append(t)

if H is None or R is None or T is None:

raise ValueError("Either evaluation data or file should be provided")

if filter_files:

if not (filter_H is None and filter_R is None and filter_T is None):

raise ValueError("Filter data and file should not be provided at the same time")

filter_H = []

filter_R = []

filter_T = []

for filter_file in filter_files:

with open(filter_file, "r") as fin:

for line in fin:

tokens = self.tokenize(line)

if len(tokens) == 0:

continue

h, r, t = tokens

filter_H.append(h)

filter_R.append(r)

filter_T.append(t)

elif filter_H is None:

filter_H = []

filter_R = []

filter_T = []

entity2id = self.graph.entity2id

relation2id = self.graph.relation2id

new_H, new_R, new_T = self.name_map((entity2id, relation2id, entity2id), (H, R, T))

logger.info("effective triplets: %d / %d" % (len(new_H), len(H)))

H = np.asarray(new_H, dtype=np.uint32)

R = np.asarray(new_R, dtype=np.uint32)

T = np.asarray(new_T, dtype=np.uint32)

new_H, new_R, new_T = self.name_map((entity2id, relation2id, entity2id), (filter_H, filter_R, filter_T))

logger.info("effective filter triplets: %d / %d" % (len(new_H), len(filter_H)))

filter_H = np.asarray(new_H, dtype=np.uint32)

filter_R = np.asarray(new_R, dtype=np.uint32)

filter_T = np.asarray(new_T, dtype=np.uint32)

exclude_H = defaultdict(set)

exclude_T = defaultdict(set)

for h, r, t in zip(filter_H, filter_R, filter_T):

exclude_H[(t, r)].add(h)

exclude_T[(h, r)].add(t)

num_sample = len(H)

fast_mode = fast_mode or num_sample

indexes = np.random.permutation(num_sample)[:fast_mode]

H = H[indexes]

R = R[indexes]

T = T[indexes]

if backend == "graphvite":

rankings = graphvite_predict()

elif backend == "torch":

rankings = torch_predict()

return {

"MR": np.mean(rankings),

"MRR": np.mean(1 / rankings),

"HITS@1": np.mean(rankings <= 1),

"HITS@3": np.mean(rankings <= 3),

"HITS@10": np.mean(rankings <= 10)

}

def get_batch_size(self, sample_size):

import psutil

memory = psutil.virtual_memory()

batch_size = int(self.SAMPLE_PER_DIMENSION * self.dim * self.graph.num_vertex

* self.solver.num_partition / self.solver.num_worker / sample_size)

# 2 triplet (Python, C++ sample pool) + 1 sample index

mem_per_sample = sample_size * (2 * 3 * np.uint32().itemsize + 1 * np.uint64().itemsize)

max_batch_size = int(memory.available / mem_per_sample / self.MEMORY_SCALE_FACTOR)

if max_batch_size < batch_size:

logger.info("Memory is not enough for optimal prediction batch size. "

"Use the maximal possible size instead.")

batch_size = max_batch_size

return batch_size

def generate_one_vs_rest(self, H, R, T, num_entity, target="both"):

one = np.ones(num_entity, dtype=np.bool)

all = np.arange(num_entity, dtype=np.uint32)

batches = []

for h, r, t in zip(H, R, T):

if target == "head" or target == "both":

batch = np.asarray([all, t * one, r * one]).transpose()

batches.append(batch)

if target == "tail" or target == "both":

batch = np.asarray([h * one, all, r * one]).transpose()

batches.append(batch)

batches = np.concatenate(batches)

return batches

def generate_mask(self, H, R, T, exclude_H, exclude_T, num_entity, target="both"):

one = np.ones(num_entity, dtype=np.bool)

masks = []

for h, r, t in zip(H, R, T):

if target == "head" or target == "both":

mask = one.copy()

mask[list(exclude_H[(t, r)])] = 0

mask[h] = 1

masks.append(mask)

if target == "tail" or target == "both":

mask = one.copy()

mask[list(exclude_T[(h, r)])] = 0

mask[t] = 1

masks.append(mask)

masks = np.asarray(masks)

import torch

from .network import LinkPredictor

torch.set_grad_enabled(False)

entity_embeddings, relation_embeddings, H, R, T, \

exclude_H, exclude_T, target, k, score_function, margin, device = args

entity_embeddings = np.asarray(entity_embeddings)

relation_embeddings = np.asarray(relation_embeddings)

num_entity = len(entity_embeddings)

score_function = LinkPredictor(score_function, entity_embeddings, relation_embeddings, entity_embeddings,

margin=margin)

if device != "cpu":

try:

score_function = score_function.to(device)

except RuntimeError:

logger.info("Model is too large for GPU evaluation with PyTorch. Switch to CPU evaluation.")

device = "cpu"

if device == "cpu":

del score_function

torch.cuda.empty_cache()

score_function = LinkPredictor(score_function, entity_embeddings, relation_embeddings, entity_embeddings,

margin=margin)

one = torch.ones(num_entity, dtype=torch.long, device=device)

all = torch.arange(num_entity, dtype=torch.long, device=device)

results = [] # rankings or top-k recalls

for h, r, t in zip(H, R, T):

if target == "head" or target == "both":

batch_h = all

batch_r = r * one

batch_t = t * one

score = score_function(batch_h, batch_r, batch_t)

if k: # top-k recalls

score, index = torch.topk(score, k)

score = score.cpu().numpy()

index = index.cpu().numpy()

recall = list(zip(index, score))

results.append(recall)

else: # ranking

mask = torch.ones(num_entity, dtype=torch.uint8, device=device)

index = torch.tensor(list(exclude_H[(t, r)]), dtype=torch.long, device=device)

mask[index] = 0

mask[h] = 1

ranking = torch.sum((score >= score[h]) * mask).item()

results.append(ranking)

if target == "tail" or target == "both":

batch_h = h * one

batch_r = r * one

batch_t = all

score = score_function(batch_h, batch_r, batch_t)

if k: # top-k recalls

score, index = torch.topk(score, k)

score = score.cpu().numpy()

index = index.cpu().numpy()

recall = list(zip(index, score))

results.append(recall)

else: # ranking

mask = torch.ones(num_entity, dtype=torch.uint8, device=device)

index = torch.tensor(list(exclude_T[(h, r)]), dtype=torch.long, device=device)