| 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 |
|
|
