File size: 1,139 Bytes
4523bda | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 | import gc
import ray
import numpy as np
from openea.modules.utils.util import task_divide, merge_dic
ray.init()
def generate_neighbours(entity_embeds, entity_list, neighbors_num, frags_num=8):
ent_frags = task_divide(np.array(entity_list), frags_num)
ent_frag_indexes = task_divide(np.array(range(len(entity_list))), frags_num)
dic = dict()
rest = []
for i in range(len(ent_frags)):
res = find_neighbours.remote(ent_frags[i], np.array(entity_list), entity_embeds[ent_frag_indexes[i], :],
entity_embeds, neighbors_num)
rest.append(res)
for res in ray.get(rest):
dic = merge_dic(dic, res)
del entity_embeds
gc.collect()
return dic
@ray.remote(num_cpus=1)
def find_neighbours(frags, entity_list, sub_embed, embed, k):
dic = dict()
sim_mat = np.matmul(sub_embed, embed.T)
for i in range(sim_mat.shape[0]):
sort_index = np.argpartition(-sim_mat[i, :], k)
neighbors_index = sort_index[0:k]
neighbors = entity_list[neighbors_index].tolist()
dic[frags[i]] = neighbors
del sim_mat
return dic
|