def read_data(split, params, logger):
samples = utils.read_dataset(split, params["data_path"])
# Check if dataset has multiple ground-truth labels
has_mult_labels = "labels" in samples[0].keys()
if params["filter_unlabeled"]:
# Filter samples without gold entities
samples = list(
filter(
lambda sample: (len(sample["labels"]) > 0)
if has_mult_labels else (sample["label"] is not None),
samples))
logger.info(f"Read {len(samples)} {split} samples.")
return samples, has_mult_labels
def main(params):
output_path = params["output_path"]
if not os.path.exists(output_path):
os.makedirs(output_path)
logger = utils.get_logger(params["output_path"])
# Init model
reranker = BiEncoderRanker(params)
tokenizer = reranker.tokenizer
model = reranker.model
device = reranker.device
cand_encode_path = params.get("cand_encode_path", None)
# candidate encoding is not pre-computed.
# load/generate candidate pool to compute candidate encoding.
cand_pool_path = params.get("cand_pool_path", None)
candidate_pool = load_or_generate_candidate_pool(
tokenizer,
params,
logger,
cand_pool_path,
)
candidate_encoding = None
if cand_encode_path is not None:
# try to load candidate encoding from path
# if success, avoid computing candidate encoding
try:
logger.info("Loading pre-generated candidate encode path.")
candidate_encoding = torch.load(cand_encode_path)
except:
logger.info("Loading failed. Generating candidate encoding.")
if candidate_encoding is None:
candidate_encoding = encode_candidate_zeshel(
reranker,
candidate_pool,
params["encode_batch_size"],
silent=params["silent"],
logger=logger,
)
if cand_encode_path is not None:
# Save candidate encoding to avoid re-compute
logger.info("Saving candidate encoding to file " +
cand_encode_path)
torch.save(cand_encode_path, candidate_encoding)
test_samples = utils.read_dataset(params["mode"], params["data_path"])
logger.info("Read %d test samples." % len(test_samples))
test_data, test_tensor_data = data.process_mention_data(
test_samples,
tokenizer,
params["max_context_length"],
params["max_cand_length"],
context_key=params['context_key'],
silent=params["silent"],
logger=logger,
debug=params["debug"],
)
test_sampler = SequentialSampler(test_tensor_data)
test_dataloader = DataLoader(test_tensor_data,
sampler=test_sampler,
batch_size=params["encode_batch_size"])
save_results = params.get("save_topk_result")
new_data = nnquery.get_topk_predictions(
reranker,
test_dataloader,
candidate_pool,
candidate_encoding,
params["silent"],
logger,
params["top_k"],
params.get("zeshel", None),
save_results,
)
if save_results:
save_data_path = os.path.join(
params['output_path'],
'candidates_%s_top%d.t7' % (params['mode'], params['top_k']))
torch.save(new_data, save_data_path)
def main(params):
output_path = params["output_path"]
if not os.path.exists(output_path):
os.makedirs(output_path)
logger = utils.get_logger(params["output_path"])
# Init model
reranker = BiEncoderRanker(params)
tokenizer = reranker.tokenizer
model = reranker.model
device = reranker.device
cand_encode_path = params.get("cand_encode_path", None)
# candidate encoding is not pre-computed.
# load/generate candidate pool to compute candidate encoding.
cand_pool_path = params.get("cand_pool_path", None)
candidate_pool = load_or_generate_candidate_pool(
tokenizer,
params,
logger,
cand_pool_path,
)
candidate_encoding = None
if cand_encode_path is not None:
# try to load candidate encoding from path
# if success, avoid computing candidate encoding
try:
logger.info("Loading pre-generated candidate encode path.")
candidate_encoding = torch.load(cand_encode_path)
except:
logger.info("Loading failed. Generating candidate encoding.")
if candidate_encoding is None:
candidate_encoding = encode_candidate(reranker,
candidate_pool,
params["encode_batch_size"],
silent=params["silent"],
logger=logger,
is_zeshel=params.get(
"zeshel", None))
if cand_encode_path is not None:
# Save candidate encoding to avoid re-compute
logger.info("Saving candidate encoding to file " +
cand_encode_path)
torch.save(candidate_encoding, cand_encode_path)
test_samples = utils.read_dataset(params["mode"], params["data_path"])
# test_samples_custom = utils.read_dataset(params["mode"], 'data/zeshel/processed')
# # Copy custom dataset to original except 'label_id'
# for i in range(len(test_samples)):
# for k in test_samples[i]:
# if k == 'label_id':
# continue
# k_custom = 'type' if k == 'world' else k
# test_samples[i][k] = test_samples_custom[i][k_custom]
logger.info("Read %d test samples." % len(test_samples))
test_data, test_tensor_data = data.process_mention_data(
test_samples,
tokenizer,
params["max_context_length"],
params["max_cand_length"],
context_key=params['context_key'],
silent=params["silent"],
logger=logger,
debug=params["debug"],
)
# custom_embeds = torch.load('models/trained/zeshel_og/eval/data_og/custom_embed.t7')
# dict_idxs_by_type = torch.load('models/trained/zeshel_og/eval/data_og/dict_idx_mapping.t7')
# test_dict_vecs = torch.load('models/trained/zeshel_og/eval/data_og/custom_test_ent_vecs.t7')
# test_men_vecs = torch.load('models/trained/zeshel_og/eval/data_og/custom_test_men_vecs.t7')
# world_to_type = {12:'forgotten_realms', 13:'lego', 14:'star_trek', 15:'yugioh'}
# embed()
test_sampler = SequentialSampler(test_tensor_data)
test_dataloader = DataLoader(test_tensor_data,
sampler=test_sampler,
batch_size=params["encode_batch_size"])
save_results = params.get("save_topk_result")
new_data = nnquery.get_topk_predictions(
reranker,
test_dataloader,
candidate_pool,
candidate_encoding,
params["silent"],
logger,
params["top_k"],
params.get("zeshel", None),
save_results,
)
if save_results:
save_data_path = os.path.join(
params['output_path'],
'candidates_%s_top%d.t7' % (params['mode'], params['top_k']))
torch.save(new_data, save_data_path)
def main(params):
output_path = params["output_path"]
if not os.path.exists(output_path):
os.makedirs(output_path)
logger = utils.get_logger(params["output_path"], 'log-eval')
pickle_src_path = params["pickle_src_path"]
if pickle_src_path is None or not os.path.exists(pickle_src_path):
pickle_src_path = output_path
embed_data_path = params["embed_data_path"]
if embed_data_path is None or not os.path.exists(embed_data_path):
embed_data_path = output_path
cands_path = os.path.join(output_path, 'cands.pickle')
# Init model
reranker = BiEncoderRanker(params)
reranker.model.eval()
tokenizer = reranker.tokenizer
n_gpu = reranker.n_gpu
knn = params["knn"]
use_types = params["use_types"]
within_doc = params["within_doc"]
data_split = params["data_split"] # Default = "test"
# Load test data
test_samples = None
entity_dictionary_loaded = False
test_dictionary_pkl_path = os.path.join(pickle_src_path,
'test_dictionary.pickle')
test_tensor_data_pkl_path = os.path.join(pickle_src_path,
'test_tensor_data.pickle')
test_mention_data_pkl_path = os.path.join(pickle_src_path,
'test_mention_data.pickle')
if params['transductive']:
train_tensor_data_pkl_path = os.path.join(pickle_src_path,
'train_tensor_data.pickle')
train_mention_data_pkl_path = os.path.join(
pickle_src_path, 'train_mention_data.pickle')
if os.path.isfile(test_dictionary_pkl_path):
print("Loading stored processed entity dictionary...")
with open(test_dictionary_pkl_path, 'rb') as read_handle:
test_dictionary = pickle.load(read_handle)
entity_dictionary_loaded = True
if os.path.isfile(test_tensor_data_pkl_path) and os.path.isfile(
test_mention_data_pkl_path):
print("Loading stored processed test data...")
with open(test_tensor_data_pkl_path, 'rb') as read_handle:
test_tensor_data = pickle.load(read_handle)
with open(test_mention_data_pkl_path, 'rb') as read_handle:
mention_data = pickle.load(read_handle)
else:
test_samples = utils.read_dataset(data_split, params["data_path"])
if not entity_dictionary_loaded:
with open(os.path.join(params["data_path"], 'dictionary.pickle'),
'rb') as read_handle:
test_dictionary = pickle.load(read_handle)
# Check if dataset has multiple ground-truth labels
mult_labels = "labels" in test_samples[0].keys()
if params["filter_unlabeled"]:
# Filter samples without gold entities
test_samples = list(
filter(
lambda sample: (len(sample["labels"]) > 0)
if mult_labels else (sample["label"] is not None),
test_samples))
logger.info("Read %d test samples." % len(test_samples))
mention_data, test_dictionary, test_tensor_data = data_process.process_mention_data(
test_samples,
test_dictionary,
tokenizer,
params["max_context_length"],
params["max_cand_length"],
multi_label_key="labels" if mult_labels else None,
context_key=params["context_key"],
silent=params["silent"],
logger=logger,
debug=params["debug"],
knn=knn,
dictionary_processed=entity_dictionary_loaded)
print("Saving processed test data...")
if not entity_dictionary_loaded:
with open(test_dictionary_pkl_path, 'wb') as write_handle:
pickle.dump(test_dictionary,
write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
entity_dictionary_loaded = True
with open(test_tensor_data_pkl_path, 'wb') as write_handle:
pickle.dump(test_tensor_data,
write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
with open(test_mention_data_pkl_path, 'wb') as write_handle:
pickle.dump(mention_data,
write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
# Store test dictionary token ids
test_dict_vecs = torch.tensor(list(map(lambda x: x['ids'],
test_dictionary)),
dtype=torch.long)
# Store test mention token ids
test_men_vecs = test_tensor_data[:][0]
n_entities = len(test_dict_vecs)
n_mentions = len(test_tensor_data)
if within_doc:
if test_samples is None:
test_samples, _ = read_data(data_split, params, logger)
test_context_doc_ids = [s['context_doc_id'] for s in test_samples]
if params["transductive"]:
if os.path.isfile(train_tensor_data_pkl_path) and os.path.isfile(
train_mention_data_pkl_path):
print("Loading stored processed train data...")
with open(train_tensor_data_pkl_path, 'rb') as read_handle:
train_tensor_data = pickle.load(read_handle)
with open(train_mention_data_pkl_path, 'rb') as read_handle:
train_mention_data = pickle.load(read_handle)
else:
train_samples = utils.read_dataset('train', params["data_path"])
# Check if dataset has multiple ground-truth labels
mult_labels = "labels" in train_samples[0].keys()
logger.info("Read %d test samples." % len(test_samples))
train_mention_data, _, train_tensor_data = data_process.process_mention_data(
train_samples,
test_dictionary,
tokenizer,
params["max_context_length"],
params["max_cand_length"],
multi_label_key="labels" if mult_labels else None,
context_key=params["context_key"],
silent=params["silent"],
logger=logger,
debug=params["debug"],
knn=knn,
dictionary_processed=entity_dictionary_loaded)
print("Saving processed train data...")
with open(train_tensor_data_pkl_path, 'wb') as write_handle:
pickle.dump(train_tensor_data,
write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
with open(train_mention_data_pkl_path, 'wb') as write_handle:
pickle.dump(train_mention_data,
write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
# Store train mention token ids
train_men_vecs = train_tensor_data[:][0]
n_mentions += len(train_tensor_data)
n_train_mentions = len(train_tensor_data)
# if os.path.isfile(cands_path):
# print("Loading stored candidates...")
# with open(cands_path, 'rb') as read_handle:
# men_cands = pickle.load(read_handle)
# else:
# Check and load stored embedding data
embed_data_path = os.path.join(embed_data_path, 'embed_data.t7')
embed_data = None
if os.path.isfile(embed_data_path):
embed_data = torch.load(embed_data_path)
if use_types:
if embed_data is not None:
logger.info('Loading stored embeddings and computing indexes')
dict_embeds = embed_data['dict_embeds']
if 'dict_idxs_by_type' in embed_data:
dict_idxs_by_type = embed_data['dict_idxs_by_type']
else:
dict_idxs_by_type = data_process.get_idxs_by_type(
test_dictionary)
dict_indexes = data_process.get_index_from_embeds(
dict_embeds,
dict_idxs_by_type,
force_exact_search=params['force_exact_search'],
probe_mult_factor=params['probe_mult_factor'])
men_embeds = embed_data['men_embeds']
if 'men_idxs_by_type' in embed_data:
men_idxs_by_type = embed_data['men_idxs_by_type']
else:
men_idxs_by_type = data_process.get_idxs_by_type(mention_data)
men_indexes = data_process.get_index_from_embeds(
men_embeds,
men_idxs_by_type,
force_exact_search=params['force_exact_search'],
probe_mult_factor=params['probe_mult_factor'])
else:
logger.info("Dictionary: Embedding and building index")
dict_embeds, dict_indexes, dict_idxs_by_type = data_process.embed_and_index(
reranker,
test_dict_vecs,
encoder_type="candidate",
n_gpu=n_gpu,
corpus=test_dictionary,
force_exact_search=params['force_exact_search'],
batch_size=params['embed_batch_size'],
probe_mult_factor=params['probe_mult_factor'])
logger.info("Queries: Embedding and building index")
vecs = test_men_vecs
men_data = mention_data
if params['transductive']:
vecs = torch.cat((train_men_vecs, vecs), dim=0)
men_data = train_mention_data + mention_data
men_embeds, men_indexes, men_idxs_by_type = data_process.embed_and_index(
reranker,
vecs,
encoder_type="context",
n_gpu=n_gpu,
corpus=men_data,
force_exact_search=params['force_exact_search'],
batch_size=params['embed_batch_size'],
probe_mult_factor=params['probe_mult_factor'])
else:
if embed_data is not None:
logger.info('Loading stored embeddings and computing indexes')
dict_embeds = embed_data['dict_embeds']
dict_index = data_process.get_index_from_embeds(
dict_embeds,
force_exact_search=params['force_exact_search'],
probe_mult_factor=params['probe_mult_factor'])
men_embeds = embed_data['men_embeds']
men_index = data_process.get_index_from_embeds(
men_embeds,
force_exact_search=params['force_exact_search'],
probe_mult_factor=params['probe_mult_factor'])
else:
logger.info("Dictionary: Embedding and building index")
dict_embeds, dict_index = data_process.embed_and_index(
reranker,
test_dict_vecs,
'candidate',
n_gpu=n_gpu,
force_exact_search=params['force_exact_search'],
batch_size=params['embed_batch_size'],
probe_mult_factor=params['probe_mult_factor'])
logger.info("Queries: Embedding and building index")
vecs = test_men_vecs
if params['transductive']:
vecs = torch.cat((train_men_vecs, vecs), dim=0)
men_embeds, men_index = data_process.embed_and_index(
reranker,
vecs,
'context',
n_gpu=n_gpu,
force_exact_search=params['force_exact_search'],
batch_size=params['embed_batch_size'],
probe_mult_factor=params['probe_mult_factor'])
# Save computed embedding data if not loaded from disk
if embed_data is None:
embed_data = {}
embed_data['dict_embeds'] = dict_embeds
embed_data['men_embeds'] = men_embeds
if use_types:
embed_data['dict_idxs_by_type'] = dict_idxs_by_type
embed_data['men_idxs_by_type'] = men_idxs_by_type
# NOTE: Cannot pickle faiss index because it is a SwigPyObject
torch.save(embed_data,
embed_data_path,
pickle_protocol=pickle.HIGHEST_PROTOCOL)
logger.info("Starting KNN search...")
# Fetch (k+1) NN mention candidates
if not use_types:
_men_embeds = men_embeds
if params['transductive']:
_men_embeds = _men_embeds[n_train_mentions:]
n_mens_to_fetch = len(_men_embeds) if within_doc else knn + 1
nn_men_dists, nn_men_idxs = men_index.search(_men_embeds,
n_mens_to_fetch)
else:
query_len = len(men_embeds) - (n_train_mentions
if params['transductive'] else 0)
nn_men_idxs = -1 * np.ones((query_len, query_len), dtype=int)
nn_men_dists = -1 * np.ones((query_len, query_len), dtype='float64')
for entity_type in men_indexes:
men_embeds_by_type = men_embeds[men_idxs_by_type[entity_type][
men_idxs_by_type[entity_type] >=
n_train_mentions]] if params['transductive'] else men_embeds[
men_idxs_by_type[entity_type]]
n_mens_to_fetch = len(
men_embeds_by_type) if within_doc else knn + 1
nn_men_dists_by_type, nn_men_idxs_by_type = men_indexes[
entity_type].search(
men_embeds_by_type,
min(n_mens_to_fetch, len(men_embeds_by_type)))
nn_men_idxs_by_type = np.array(
list(
map(lambda x: men_idxs_by_type[entity_type][x],
nn_men_idxs_by_type)))
i = -1
for idx in men_idxs_by_type[entity_type]:
if params['transductive']:
idx -= n_train_mentions
if idx < 0:
continue
i += 1
nn_men_idxs[idx][:len(nn_men_idxs_by_type[i]
)] = nn_men_idxs_by_type[i]
nn_men_dists[idx][:len(nn_men_dists_by_type[i]
)] = nn_men_dists_by_type[i]
logger.info("Search finished")
logger.info(f'Calculating mention recall@{knn}')
# Find the most similar entity and k-nn mentions for each mention query
men_recall_knn = []
men_cands = []
cui_sums = defaultdict(int)
for m in mention_data:
cui_sums[m['label_cuis'][0]] += 1
for idx in range(len(nn_men_idxs)):
filter_mask_neg1 = nn_men_idxs[idx] != -1
men_cand_idxs = nn_men_idxs[idx][filter_mask_neg1]
men_cand_scores = nn_men_dists[idx][filter_mask_neg1]
if within_doc:
men_cand_idxs, wd_mask = filter_by_context_doc_id(
men_cand_idxs,
test_context_doc_ids[idx],
test_context_doc_ids,
return_numpy=True)
men_cand_scores = men_cand_scores[wd_mask]
# Filter candidates to remove mention query and keep only the top k candidates
filter_mask = men_cand_idxs != idx
men_cand_idxs, men_cand_scores = men_cand_idxs[
filter_mask][:knn], men_cand_scores[filter_mask][:knn]
assert len(men_cand_idxs) == knn
men_cands.append(men_cand_idxs)
# Calculate mention recall@k
gold_cui = mention_data[idx]['label_cuis'][0]
if cui_sums[gold_cui] > 1:
recall_hit = [
1. for midx in men_cand_idxs
if mention_data[midx]['label_cuis'][0] == gold_cui
]
recall_knn = sum(recall_hit) / min(cui_sums[gold_cui] - 1, knn)
men_recall_knn.append(recall_knn)
logger.info('Done')
# assert len(men_recall_knn) == len(nn_men_idxs)
# Pickle the graphs
print(f"Saving top-{knn} candidates")
with open(cands_path, 'wb') as write_handle:
pickle.dump(men_cands, write_handle, protocol=pickle.HIGHEST_PROTOCOL)
# Output final results
mean_recall_knn = np.mean(men_recall_knn)
logger.info(f"Result: Final mean recall@{knn} = {mean_recall_knn}")
def main(params):
output_path = params["output_path"]
if not os.path.exists(output_path):
os.makedirs(output_path)
logger = utils.get_logger(params["output_path"], 'log')
pickle_src_path = params["pickle_src_path"]
if pickle_src_path is None or not os.path.exists(pickle_src_path):
pickle_src_path = output_path
embed_data_path = params["embed_data_path"]
if embed_data_path is None or not os.path.exists(embed_data_path):
embed_data_path = output_path
# Init model
reranker = BiEncoderRanker(params)
reranker.model.eval()
tokenizer = reranker.tokenizer
n_gpu = reranker.n_gpu
knn = params["knn"] # Use as the max-knn value for the graph construction
use_types = params["use_types"]
# within_doc = params["within_doc"]
data_split = params["data_split"] # Default = "test"
# Load test data
entity_dictionary_loaded = False
test_dictionary_pkl_path = os.path.join(pickle_src_path,
'test_dictionary.pickle')
test_tensor_data_pkl_path = os.path.join(pickle_src_path,
'test_tensor_data.pickle')
test_mention_data_pkl_path = os.path.join(pickle_src_path,
'test_mention_data.pickle')
# if params['transductive']:
# train_tensor_data_pkl_path = os.path.join(pickle_src_path, 'train_tensor_data.pickle')
# train_mention_data_pkl_path = os.path.join(pickle_src_path, 'train_mention_data.pickle')
if os.path.isfile(test_dictionary_pkl_path):
print("Loading stored processed entity dictionary...")
with open(test_dictionary_pkl_path, 'rb') as read_handle:
test_dictionary = pickle.load(read_handle)
entity_dictionary_loaded = True
if os.path.isfile(test_tensor_data_pkl_path) and os.path.isfile(
test_mention_data_pkl_path):
print("Loading stored processed test data...")
with open(test_tensor_data_pkl_path, 'rb') as read_handle:
test_tensor_data = pickle.load(read_handle)
with open(test_mention_data_pkl_path, 'rb') as read_handle:
mention_data = pickle.load(read_handle)
else:
test_samples = utils.read_dataset(data_split, params["data_path"])
if not entity_dictionary_loaded:
with open(os.path.join(params["data_path"], 'dictionary.pickle'),
'rb') as read_handle:
test_dictionary = pickle.load(read_handle)
# Check if dataset has multiple ground-truth labels
mult_labels = "labels" in test_samples[0].keys()
# Filter samples without gold entities
test_samples = list(
filter(
lambda sample: (len(sample["labels"]) > 0) if mult_labels else
(sample["label"] is not None), test_samples))
logger.info("Read %d test samples." % len(test_samples))
mention_data, test_dictionary, test_tensor_data = data_process.process_mention_data(
test_samples,
test_dictionary,
tokenizer,
params["max_context_length"],
params["max_cand_length"],
multi_label_key="labels" if mult_labels else None,
context_key=params["context_key"],
silent=params["silent"],
logger=logger,
debug=params["debug"],
knn=knn,
dictionary_processed=entity_dictionary_loaded)
print("Saving processed test data...")
if not entity_dictionary_loaded:
with open(test_dictionary_pkl_path, 'wb') as write_handle:
pickle.dump(test_dictionary,
write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
with open(test_tensor_data_pkl_path, 'wb') as write_handle:
pickle.dump(test_tensor_data,
write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
with open(test_mention_data_pkl_path, 'wb') as write_handle:
pickle.dump(mention_data,
write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
# Reducing the entity dictionary to only the ground truth of the mention queries
# Combining the entities and mentions into one structure for joint embedding and indexing
new_ents = {}
new_ents_arr = []
men_labels = []
for men in mention_data:
ent = men['label_idxs'][0]
if ent not in new_ents:
new_ents[ent] = len(new_ents_arr)
new_ents_arr.append(ent)
men_labels.append(new_ents[ent])
ent_labels = [i for i in range(len(new_ents_arr))]
new_ent_vecs = torch.tensor(
list(map(lambda x: test_dictionary[x]['ids'], new_ents_arr)))
new_ent_types = list(
map(lambda x: {"type": test_dictionary[x]['type']}, new_ents_arr))
test_men_vecs = test_tensor_data[:][0]
n_mentions = len(test_tensor_data)
n_entities = len(new_ent_vecs)
n_embeds = n_mentions + n_entities
leaf_labels = np.array(ent_labels + men_labels, dtype=int)
all_vecs = torch.cat((new_ent_vecs, test_men_vecs))
all_types = new_ent_types + mention_data # Array of dicts containing key "type" for selected ents and all mentions
# Values of k to run the evaluation against
knn_vals = [25 * 2**i for i in range(int(math.log(knn / 25, 2)) + 1)
] if params["exact_knn"] is None else [params["exact_knn"]]
# Store the maximum evaluation k
max_knn = knn_vals[-1]
time_start = time.time()
# Check if graphs are already built
graph_path = os.path.join(output_path, 'graphs.pickle')
if os.path.isfile(graph_path):
print("Loading stored joint graphs...")
with open(graph_path, 'rb') as read_handle:
joint_graphs = pickle.load(read_handle)
else:
# Initialize graphs to store mention-mention and mention-entity similarity score edges;
# Keyed on k, the number of nearest mentions retrieved
joint_graphs = {}
for k in knn_vals:
joint_graphs[k] = {
'rows': np.array([]),
'cols': np.array([]),
'data': np.array([]),
'shape': (n_embeds, n_embeds)
}
# Check and load stored embedding data
embed_data_path = os.path.join(embed_data_path, 'embed_data.t7')
embed_data = None
if os.path.isfile(embed_data_path):
embed_data = torch.load(embed_data_path)
if use_types:
if embed_data is not None:
logger.info('Loading stored embeddings')
embeds = embed_data['embeds']
if 'idxs_by_type' in embed_data:
idxs_by_type = embed_data['idxs_by_type']
else:
idxs_by_type = data_process.get_idxs_by_type(all_types)
else:
logger.info("Embedding data")
dict_embeds = data_process.embed_and_index(
reranker,
all_vecs[:n_entities],
encoder_type='candidate',
only_embed=True,
n_gpu=n_gpu,
batch_size=params['embed_batch_size'])
men_embeds = data_process.embed_and_index(
reranker,
all_vecs[n_entities:],
encoder_type='context',
only_embed=True,
n_gpu=n_gpu,
batch_size=params['embed_batch_size'])
embeds = np.concatenate((dict_embeds, men_embeds), axis=0)
idxs_by_type = data_process.get_idxs_by_type(all_types)
search_indexes = data_process.get_index_from_embeds(
embeds, corpus_idxs=idxs_by_type, force_exact_search=True)
else:
if embed_data is not None:
logger.info('Loading stored embeddings')
embeds = embed_data['embeds']
else:
logger.info("Embedding data")
dict_embeds = data_process.embed_and_index(
reranker,
all_vecs[:n_entities],
encoder_type='candidate',
only_embed=True,
n_gpu=n_gpu,
batch_size=params['embed_batch_size'])
men_embeds = data_process.embed_and_index(
reranker,
all_vecs[n_entities:],
encoder_type='context',
only_embed=True,
n_gpu=n_gpu,
batch_size=params['embed_batch_size'])
embeds = np.concatenate((dict_embeds, men_embeds), axis=0)
search_index = data_process.get_index_from_embeds(
embeds, force_exact_search=True)
# Save computed embedding data if not loaded from disk
if embed_data is None:
embed_data = {}
embed_data['embeds'] = embeds
if use_types:
embed_data['idxs_by_type'] = idxs_by_type
# NOTE: Cannot pickle faiss index because it is a SwigPyObject
torch.save(embed_data,
embed_data_path,
pickle_protocol=pickle.HIGHEST_PROTOCOL)
# Build faiss search index
if params["normalize_embeds"]:
embeds = normalize(embeds, axis=0)
logger.info("Building KNN index...")
if use_types:
search_indexes = data_process.get_index_from_embeds(
embeds, corpus_idxs=idxs_by_type, force_exact_search=True)
else:
search_index = data_process.get_index_from_embeds(
embeds, force_exact_search=True)
logger.info("Starting KNN search...")
if not use_types:
faiss_dists, faiss_idxs = search_index.search(embeds, max_knn + 1)
else:
query_len = n_embeds
faiss_idxs = np.zeros((query_len, max_knn + 1))
faiss_dists = np.zeros((query_len, max_knn + 1), dtype=float)
for entity_type in search_indexes:
embeds_by_type = embeds[idxs_by_type[entity_type]]
nn_dists_by_type, nn_idxs_by_type = search_indexes[
entity_type].search(embeds_by_type, max_knn + 1)
for i, idx in enumerate(idxs_by_type[entity_type]):
faiss_idxs[idx] = nn_idxs_by_type[i]
faiss_dists[idx] = nn_dists_by_type[i]
logger.info("Search finished")
logger.info('Building graphs')
# Find the most similar nodes for each mention and node in the set (minus self)
for idx in trange(n_embeds):
# Compute adjacent node edge weight
if idx != 0:
adj_idx = idx - 1
adj_data = embeds[adj_idx] @ embeds[idx]
nn_idxs = faiss_idxs[idx]
nn_scores = faiss_dists[idx]
# Filter candidates to remove mention query and keep only the top k candidates
filter_mask = nn_idxs != idx
nn_idxs, nn_scores = nn_idxs[filter_mask][:max_knn], nn_scores[
filter_mask][:max_knn]
# Add edges to the graphs
for k in joint_graphs:
# Add edge to adjacent node to force the graph to be connected
if idx != 0:
joint_graph['rows'] = np.append(joint_graph['rows'],
adj_idx)
joint_graph['cols'] = np.append(joint_graph['cols'], idx)
joint_graph['data'] = np.append(joint_graph['data'],
adj_data)
joint_graph = joint_graphs[k]
# Add mention-mention edges
joint_graph['rows'] = np.append(joint_graph['rows'], [idx] * k)
joint_graph['cols'] = np.append(joint_graph['cols'],
nn_idxs[:k])
joint_graph['data'] = np.append(joint_graph['data'],
nn_scores[:k])
knn_fetch_time = time.time() - time_start
# Pickle the graphs
print("Saving joint graphs...")
with open(graph_path, 'wb') as write_handle:
pickle.dump(joint_graphs,
write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
if params['only_embed_and_build']:
logger.info(f"Saved embedding data at: {embed_data_path}")
logger.info(f"Saved graphs at: {graph_path}")
exit()
results = {
'n_leaves': n_embeds,
'n_entities': n_entities,
'n_mentions': n_mentions
}
graph_processing_time = time.time()
n_graphs_processed = 0.
linkage_fns = ["single", "complete", "average"] if params["linkage"] is None \
else [params["linkage"]] # Different HAC linkage functions to run the analyses over
for fn in linkage_fns:
logger.info(f"Linkage function: {fn}")
purities = []
fn_result = {}
for k in joint_graphs:
graph = hg.UndirectedGraph(n_embeds)
graph.add_edges(joint_graphs[k]['rows'], joint_graphs[k]['cols'])
weights = -joint_graphs[k][
'data'] # Since Higra expects weights as distances, not similarity
tree = get_hac_tree(graph, weights, linkage=fn)
purity = hg.dendrogram_purity(tree, leaf_labels)
fn_result[f"purity@{k}nn"] = purity
logger.info(f"purity@{k}nn = {purity}")
purities.append(purity)
n_graphs_processed += 1
fn_result["average"] = round(np.mean(purities), 4)
logger.info(f"average = {fn_result['average']}")
results[fn] = fn_result
avg_graph_processing_time = (time.time() -
graph_processing_time) / n_graphs_processed
avg_per_graph_time = (knn_fetch_time + avg_graph_processing_time) / 60
execution_time = (time.time() - time_start) / 60
# Store results
output_file_name = os.path.join(
output_path,
f"results_{__import__('calendar').timegm(__import__('time').gmtime())}"
)
logger.info(f"Results: \n {results}")
with open(f'{output_file_name}.json', 'w') as f:
json.dump(results, f, indent=2)
print(f"\nResults saved at: {output_file_name}.json")
logger.info("\nThe avg. per graph evaluation time is {} minutes\n".format(
avg_per_graph_time))
logger.info(
"\nThe total evaluation took {} minutes\n".format(execution_time))
def load_data(data_split,
bi_tokenizer,
max_context_length,
max_cand_length,
knn,
pickle_src_path,
params,
logger,
return_dict_only=False):
entity_dictionary_loaded = False
entity_dictionary_pkl_path = os.path.join(pickle_src_path, 'entity_dictionary.pickle')
if os.path.isfile(entity_dictionary_pkl_path):
print("Loading stored processed entity dictionary...")
with open(entity_dictionary_pkl_path, 'rb') as read_handle:
entity_dictionary = pickle.load(read_handle)
entity_dictionary_loaded = True
if return_dict_only and entity_dictionary_loaded:
return entity_dictionary
# Load data
tensor_data_pkl_path = os.path.join(pickle_src_path, f'{data_split}_tensor_data.pickle')
processed_data_pkl_path = os.path.join(pickle_src_path, f'{data_split}_processed_data.pickle')
if os.path.isfile(tensor_data_pkl_path) and os.path.isfile(processed_data_pkl_path):
print("Loading stored processed data...")
with open(tensor_data_pkl_path, 'rb') as read_handle:
tensor_data = pickle.load(read_handle)
with open(processed_data_pkl_path, 'rb') as read_handle:
processed_data = pickle.load(read_handle)
else:
data_samples = utils.read_dataset(data_split, params["data_path"])
if not entity_dictionary_loaded:
with open(os.path.join(params["data_path"], 'dictionary.pickle'), 'rb') as read_handle:
entity_dictionary = pickle.load(read_handle)
# Check if dataset has multiple ground-truth labels
mult_labels = "labels" in data_samples[0].keys()
# Filter samples without gold entities
data_samples = list(
filter(lambda sample: (len(sample["labels"]) > 0) if mult_labels else (sample["label"] is not None),
data_samples))
logger.info("Read %d data samples." % len(data_samples))
processed_data, entity_dictionary, tensor_data = data_process.process_mention_data(
data_samples,
entity_dictionary,
bi_tokenizer,
max_context_length,
max_cand_length,
context_key=params["context_key"],
multi_label_key="labels" if mult_labels else None,
silent=params["silent"],
logger=logger,
debug=params["debug"],
knn=knn,
dictionary_processed=entity_dictionary_loaded
)
print("Saving processed data...")
if not entity_dictionary_loaded:
with open(entity_dictionary_pkl_path, 'wb') as write_handle:
pickle.dump(entity_dictionary, write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
with open(tensor_data_pkl_path, 'wb') as write_handle:
pickle.dump(tensor_data, write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
with open(processed_data_pkl_path, 'wb') as write_handle:
pickle.dump(processed_data, write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
if return_dict_only:
return entity_dictionary
return entity_dictionary, tensor_data, processed_data
def main(params):
model_output_path = params["output_path"]
if not os.path.exists(model_output_path):
os.makedirs(model_output_path)
logger = utils.get_logger(params["output_path"])
# Init model
reranker = BiEncoderRanker(params)
tokenizer = reranker.tokenizer
model = reranker.model
# utils.save_model(model, tokenizer, model_output_path)
device = reranker.device
n_gpu = reranker.n_gpu
if params["gradient_accumulation_steps"] < 1:
raise ValueError(
"Invalid gradient_accumulation_steps parameter: {}, should be >= 1"
.format(params["gradient_accumulation_steps"]))
# An effective batch size of `x`, when we are accumulating the gradient accross `y` batches will be achieved by having a batch size of `z = x / y`
# args.gradient_accumulation_steps = args.gradient_accumulation_steps // n_gpu
params["train_batch_size"] = (params["train_batch_size"] //
params["gradient_accumulation_steps"])
train_batch_size = params["train_batch_size"]
eval_batch_size = params["eval_batch_size"]
grad_acc_steps = params["gradient_accumulation_steps"]
# Fix the random seeds
seed = params["seed"]
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if reranker.n_gpu > 0:
torch.cuda.manual_seed_all(seed)
# Load train data
train_samples = utils.read_dataset("train", params["data_path"])
logger.info("Read %d train samples." % len(train_samples))
train_data, train_tensor_data = data.process_mention_data(
train_samples,
tokenizer,
params["max_context_length"],
params["max_cand_length"],
context_key=params["context_key"],
silent=params["silent"],
logger=logger,
debug=params["debug"],
)
if params["shuffle"]:
train_sampler = RandomSampler(train_tensor_data)
else:
train_sampler = SequentialSampler(train_tensor_data)
train_dataloader = DataLoader(train_tensor_data,
sampler=train_sampler,
batch_size=train_batch_size)
# Load eval data
# TODO: reduce duplicated code here
valid_samples = utils.read_dataset("valid", params["data_path"])
logger.info("Read %d valid samples." % len(valid_samples))
valid_data, valid_tensor_data = data.process_mention_data(
valid_samples,
tokenizer,
params["max_context_length"],
params["max_cand_length"],
context_key=params["context_key"],
silent=params["silent"],
logger=logger,
debug=params["debug"],
)
valid_sampler = SequentialSampler(valid_tensor_data)
valid_dataloader = DataLoader(valid_tensor_data,
sampler=valid_sampler,
batch_size=eval_batch_size)
# evaluate before training
results = evaluate(
reranker,
valid_dataloader,
params,
device=device,
logger=logger,
)
number_of_samples_per_dataset = {}
time_start = time.time()
utils.write_to_file(os.path.join(model_output_path, "training_params.txt"),
str(params))
logger.info("Starting training")
logger.info("device: {} n_gpu: {}, distributed training: {}".format(
device, n_gpu, False))
optimizer = get_optimizer(model, params)
scheduler = get_scheduler(params, optimizer, len(train_tensor_data),
logger)
model.train()
best_epoch_idx = -1
best_score = -1
num_train_epochs = params["num_train_epochs"]
for epoch_idx in trange(int(num_train_epochs), desc="Epoch"):
tr_loss = 0
results = None
if params["silent"]:
iter_ = train_dataloader
else:
iter_ = tqdm(train_dataloader, desc="Batch")
for step, batch in enumerate(iter_):
batch = tuple(t.to(device) for t in batch)
if params["zeshel"]:
context_input, candidate_input, _, _ = batch
else:
context_input, candidate_input, _ = batch
loss, _ = reranker(context_input, candidate_input)
# if n_gpu > 1:
# loss = loss.mean() # mean() to average on multi-gpu.
if grad_acc_steps > 1:
loss = loss / grad_acc_steps
tr_loss += loss.item()
if (step + 1) % (params["print_interval"] * grad_acc_steps) == 0:
logger.info("Step {} - epoch {} average loss: {}\n".format(
step,
epoch_idx,
tr_loss / (params["print_interval"] * grad_acc_steps),
))
tr_loss = 0
loss.backward()
if (step + 1) % grad_acc_steps == 0:
torch.nn.utils.clip_grad_norm_(model.parameters(),
params["max_grad_norm"])
optimizer.step()
scheduler.step()
optimizer.zero_grad()
if (step + 1) % (params["eval_interval"] * grad_acc_steps) == 0:
logger.info("Evaluation on the development dataset")
evaluate(
reranker,
valid_dataloader,
params,
device=device,
logger=logger,
)
model.train()
logger.info("\n")
logger.info("***** Saving fine - tuned model *****")
epoch_output_folder_path = os.path.join(model_output_path,
"epoch_{}".format(epoch_idx))
utils.save_model(model, tokenizer, epoch_output_folder_path)
output_eval_file = os.path.join(epoch_output_folder_path,
"eval_results.txt")
results = evaluate(
reranker,
valid_dataloader,
params,
device=device,
logger=logger,
)
ls = [best_score, results["normalized_accuracy"]]
li = [best_epoch_idx, epoch_idx]
best_score = ls[np.argmax(ls)]
best_epoch_idx = li[np.argmax(ls)]
logger.info("\n")
execution_time = (time.time() - time_start) / 60
utils.write_to_file(
os.path.join(model_output_path, "training_time.txt"),
"The training took {} minutes\n".format(execution_time),
)
logger.info("The training took {} minutes\n".format(execution_time))
# save the best model in the parent_dir
logger.info("Best performance in epoch: {}".format(best_epoch_idx))
params["path_to_model"] = os.path.join(model_output_path,
"epoch_{}".format(best_epoch_idx))
utils.save_model(reranker.model, tokenizer, model_output_path)
if params["evaluate"]:
params["path_to_model"] = model_output_path
results = evaluate(
reranker,
valid_dataloader,
params,
device=device,
logger=logger,
)
def main(params):
output_path = params["output_path"]
if not os.path.exists(output_path):
os.makedirs(output_path)
logger = utils.get_logger(params["output_path"])
pickle_src_path = params["pickle_src_path"]
if pickle_src_path is None or not os.path.exists(pickle_src_path):
pickle_src_path = output_path
# Init model
reranker = BiEncoderRanker(params)
reranker.model.eval()
tokenizer = reranker.tokenizer
model = reranker.model
device = reranker.device
n_gpu = reranker.n_gpu
knn = params["knn"]
directed_graph = params["directed_graph"]
use_types = params["use_types"]
data_split = params["data_split"] # Parameter default is "test"
# Load test data
entity_dictionary_loaded = False
test_dictionary_pkl_path = os.path.join(pickle_src_path, 'test_dictionary.pickle')
test_tensor_data_pkl_path = os.path.join(pickle_src_path, 'test_tensor_data.pickle')
test_mention_data_pkl_path = os.path.join(pickle_src_path, 'test_mention_data.pickle')
if os.path.isfile(test_dictionary_pkl_path):
print("Loading stored processed entity dictionary...")
with open(test_dictionary_pkl_path, 'rb') as read_handle:
test_dictionary = pickle.load(read_handle)
entity_dictionary_loaded = True
if os.path.isfile(test_tensor_data_pkl_path) and os.path.isfile(test_mention_data_pkl_path):
print("Loading stored processed test data...")
with open(test_tensor_data_pkl_path, 'rb') as read_handle:
test_tensor_data = pickle.load(read_handle)
with open(test_mention_data_pkl_path, 'rb') as read_handle:
mention_data = pickle.load(read_handle)
else:
test_samples = utils.read_dataset(data_split, params["data_path"])
if not entity_dictionary_loaded:
with open(os.path.join(params["data_path"], 'dictionary.pickle'), 'rb') as read_handle:
test_dictionary = pickle.load(read_handle)
# Check if dataset has multiple ground-truth labels
mult_labels = "labels" in test_samples[0].keys()
if params["filter_unlabeled"]:
# Filter samples without gold entities
test_samples = list(filter(lambda sample: (len(sample["labels"]) > 0) if mult_labels else (sample["label"] is not None), test_samples))
logger.info("Read %d test samples." % len(test_samples))
mention_data, test_dictionary, test_tensor_data = data_process.process_mention_data(
test_samples,
test_dictionary,
tokenizer,
params["max_context_length"],
params["max_cand_length"],
multi_label_key="labels" if mult_labels else None,
context_key=params["context_key"],
silent=params["silent"],
logger=logger,
debug=params["debug"],
knn=knn,
dictionary_processed=entity_dictionary_loaded
)
print("Saving processed test data...")
if not entity_dictionary_loaded:
with open(test_dictionary_pkl_path, 'wb') as write_handle:
pickle.dump(test_dictionary, write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
with open(test_tensor_data_pkl_path, 'wb') as write_handle:
pickle.dump(test_tensor_data, write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
with open(test_mention_data_pkl_path, 'wb') as write_handle:
pickle.dump(mention_data, write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
# Store test dictionary token ids
test_dict_vecs = torch.tensor(
list(map(lambda x: x['ids'], test_dictionary)), dtype=torch.long)
# Store test mention token ids
test_men_vecs = test_tensor_data[:][0]
n_entities = len(test_dict_vecs)
n_mentions = len(test_tensor_data)
# Values of k to run the evaluation against
knn_vals = [0] + [2**i for i in range(int(math.log(knn, 2)) + 1)]
# Store the maximum evaluation k
max_knn = knn_vals[-1]
# Check if graphs are already built
graph_path = os.path.join(output_path, 'graphs.pickle')
if not params['only_recall'] and os.path.isfile(graph_path):
print("Loading stored joint graphs...")
with open(graph_path, 'rb') as read_handle:
joint_graphs = pickle.load(read_handle)
else:
# Initialize graphs to store mention-mention and mention-entity similarity score edges;
# Keyed on k, the number of nearest mentions retrieved
joint_graphs = {}
for k in knn_vals:
joint_graphs[k] = {
'rows': np.array([]),
'cols': np.array([]),
'data': np.array([]),
'shape': (n_entities+n_mentions, n_entities+n_mentions)
}
if use_types:
print("Dictionary: Embedding and building index")
dict_embeds, dict_indexes, dict_idxs_by_type = data_process.embed_and_index(reranker, test_dict_vecs, encoder_type="candidate", n_gpu=n_gpu, corpus=test_dictionary, force_exact_search=True)
# Verifiy embeddings
og_embeds = torch.load('models/trained/zeshel_og/eval/data_og/cand_encodes.t7')
world_to_type = {12:'forgotten_realms', 13:'lego', 14:'star_trek', 15:'yugioh'}
for world in og_embeds:
for i,oge in enumerate(tqdm(og_embeds[world])):
dict_embed_idx = dict_idxs_by_type[world_to_type[world]][i]
try:
assert torch.eq(oge, torch.tensor(dict_embeds[dict_embed_idx]))
except:
embed()
exit()
print('PASS')
exit()
print("Queries: Embedding and building index")
men_embeds, men_indexes, men_idxs_by_type = data_process.embed_and_index(reranker, test_men_vecs, encoder_type="context", n_gpu=n_gpu, corpus=mention_data, force_exact_search=True)
else:
print("Dictionary: Embedding and building index")
dict_embeds, dict_index = data_process.embed_and_index(
reranker, test_dict_vecs, 'candidate', n_gpu=n_gpu)
print("Queries: Embedding and building index")
men_embeds, men_index = data_process.embed_and_index(
reranker, test_men_vecs, 'context', n_gpu=n_gpu)
recall_accuracy = {2**i: 0 for i in range(int(math.log(params['recall_k'], 2)) + 1)}
recall_idxs = [0.]*params['recall_k']
# Find the most similar entity and k-nn mentions for each mention query
for men_query_idx, men_embed in enumerate(tqdm(men_embeds, total=len(men_embeds), desc="Fetching k-NN")):
men_embed = np.expand_dims(men_embed, axis=0)
dict_type_idx_mapping, men_type_idx_mapping = None, None
if use_types:
entity_type = mention_data[men_query_idx]['type']
dict_index = dict_indexes[entity_type]
men_index = men_indexes[entity_type]
dict_type_idx_mapping = dict_idxs_by_type[entity_type]
men_type_idx_mapping = men_idxs_by_type[entity_type]
# Fetch nearest entity candidate
gold_idxs = mention_data[men_query_idx]["label_idxs"][:mention_data[men_query_idx]["n_labels"]]
dict_cand_idx, dict_cand_score, recall_idx = get_query_nn(
1, dict_embeds, dict_index, men_embed, searchK=params['recall_k'], gold_idxs=gold_idxs, type_idx_mapping=dict_type_idx_mapping)
# Compute recall metric
if recall_idx > -1:
recall_idxs[recall_idx] += 1.
for recall_k in recall_accuracy:
if recall_idx < recall_k:
recall_accuracy[recall_k] += 1.
if not params['only_recall']:
# Fetch (k+1) NN mention candidates
men_cand_idxs, men_cand_scores = get_query_nn(
max_knn + 1, men_embeds, men_index, men_embed, type_idx_mapping=men_type_idx_mapping)
# Filter candidates to remove mention query and keep only the top k candidates
filter_mask = men_cand_idxs != men_query_idx
men_cand_idxs, men_cand_scores = men_cand_idxs[filter_mask][:max_knn], men_cand_scores[filter_mask][:max_knn]
# Add edges to the graphs
for k in joint_graphs:
joint_graph = joint_graphs[k]
# Add mention-entity edge
joint_graph['rows'] = np.append(
joint_graph['rows'], [n_entities+men_query_idx]) # Mentions added at an offset of maximum entities
joint_graph['cols'] = np.append(
joint_graph['cols'], dict_cand_idx)
joint_graph['data'] = np.append(
joint_graph['data'], dict_cand_score)
if k > 0:
# Add mention-mention edges
joint_graph['rows'] = np.append(
joint_graph['rows'], [n_entities+men_query_idx]*len(men_cand_idxs[:k]))
joint_graph['cols'] = np.append(
joint_graph['cols'], n_entities+men_cand_idxs[:k])
joint_graph['data'] = np.append(
joint_graph['data'], men_cand_scores[:k])
# Compute and print recall metric
recall_idx_mode = np.argmax(recall_idxs)
recall_idx_mode_prop = recall_idxs[recall_idx_mode]/np.sum(recall_idxs)
logger.info(f"""
Recall metrics (for {len(men_embeds)} queries):
---------------""")
logger.info(f"highest recall idx = {recall_idx_mode} ({recall_idxs[recall_idx_mode]}/{np.sum(recall_idxs)} = {recall_idx_mode_prop})")
for recall_k in recall_accuracy:
recall_accuracy[recall_k] /= len(men_embeds)
logger.info(f"recall@{recall_k} = {recall_accuracy[recall_k]}")
if params['only_recall']:
exit()
# Pickle the graphs
print("Saving joint graphs...")
with open(graph_path, 'wb') as write_handle:
pickle.dump(joint_graphs, write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
results = []
for k in joint_graphs:
print(f"\nGraph (k={k}):")
# Partition graph based on cluster-linking constraints
partitioned_graph, clusters = partition_graph(
joint_graphs[k], n_entities, directed_graph, return_clusters=True)
# Infer predictions from clusters
result = analyzeClusters(clusters, test_dictionary, mention_data, k)
# Store result
results.append(result)
# Store results
output_file_name = os.path.join(
output_path, f"eval_results_{__import__('calendar').timegm(__import__('time').gmtime())}")
result_overview = {
'n_entities': results[0]['n_entities'],
'n_mentions': results[0]['n_mentions'],
'directed': directed_graph
}
try:
for recall_k in recall_accuracy:
result_overview[f'recall@{recall_k}'] = recall_accuracy[recall_k]
except:
logger.info("Recall data not available since graphs were loaded from disk")
for r in results:
k = r['knn_mentions']
result_overview[f'accuracy@knn{k}'] = r['accuracy']
logger.info(f"accuracy@knn{k} = {r['accuracy']}")
output_file = f'{output_file_name}-{k}.json'
with open(output_file, 'w') as f:
json.dump(r, f, indent=2)
print(f"\nPredictions @knn{k} saved at: {output_file}")
with open(f'{output_file_name}.json', 'w') as f:
json.dump(result_overview, f, indent=2)
print(f"\nPredictions overview saved at: {output_file_name}.json")
def main(params):
model_output_path = params["output_path"]
if not os.path.exists(model_output_path):
os.makedirs(model_output_path)
logger = utils.get_logger(params["output_path"])
pickle_src_path = params["pickle_src_path"]
if pickle_src_path is None or not os.path.exists(pickle_src_path):
pickle_src_path = model_output_path
knn = params["knn"]
use_types = params["use_types"]
# Init model
reranker = BiEncoderRanker(params)
tokenizer = reranker.tokenizer
model = reranker.model
device = reranker.device
n_gpu = reranker.n_gpu
if params["gradient_accumulation_steps"] < 1:
raise ValueError(
"Invalid gradient_accumulation_steps parameter: {}, should be >= 1"
.format(params["gradient_accumulation_steps"]))
# An effective batch size of `x`, when we are accumulating the gradient accross `y` batches will be achieved by having a batch size of `z = x / y`
params["train_batch_size"] = (params["train_batch_size"] //
params["gradient_accumulation_steps"])
train_batch_size = params["train_batch_size"]
eval_batch_size = params["eval_batch_size"]
grad_acc_steps = params["gradient_accumulation_steps"]
# Fix the random seeds
seed = params["seed"]
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if reranker.n_gpu > 0:
torch.cuda.manual_seed_all(seed)
entity_dictionary_loaded = False
entity_dictionary_pkl_path = os.path.join(pickle_src_path,
'entity_dictionary.pickle')
if os.path.isfile(entity_dictionary_pkl_path):
print("Loading stored processed entity dictionary...")
with open(entity_dictionary_pkl_path, 'rb') as read_handle:
entity_dictionary = pickle.load(read_handle)
entity_dictionary_loaded = True
if not params["only_evaluate"]:
# Load train data
train_tensor_data_pkl_path = os.path.join(pickle_src_path,
'train_tensor_data.pickle')
train_processed_data_pkl_path = os.path.join(
pickle_src_path, 'train_processed_data.pickle')
if os.path.isfile(train_tensor_data_pkl_path) and os.path.isfile(
train_processed_data_pkl_path):
print("Loading stored processed train data...")
with open(train_tensor_data_pkl_path, 'rb') as read_handle:
train_tensor_data = pickle.load(read_handle)
with open(train_processed_data_pkl_path, 'rb') as read_handle:
train_processed_data = pickle.load(read_handle)
else:
train_samples = utils.read_dataset("train", params["data_path"])
if not entity_dictionary_loaded:
with open(
os.path.join(params["data_path"], 'dictionary.pickle'),
'rb') as read_handle:
entity_dictionary = pickle.load(read_handle)
# Check if dataset has multiple ground-truth labels
mult_labels = "labels" in train_samples[0].keys()
if params["filter_unlabeled"]:
# Filter samples without gold entities
train_samples = list(
filter(
lambda sample: (len(sample["labels"]) > 0)
if mult_labels else (sample["label"] is not None),
train_samples))
logger.info("Read %d train samples." % len(train_samples))
# For discovery experiment: Drop entities used in training that were dropped randomly from dev/test set
if params["drop_entities"]:
assert entity_dictionary_loaded
drop_set_pkl_path = os.path.join(
pickle_src_path, 'drop_set_mention_data.pickle')
with open(drop_set_pkl_path, 'rb') as read_handle:
drop_set_data = pickle.load(read_handle)
drop_set_mention_gold_cui_idxs = list(
map(lambda x: x['label_idxs'][0], drop_set_data))
ents_in_data = np.unique(drop_set_mention_gold_cui_idxs)
ent_drop_prop = 0.1
logger.info(
f"Dropping {ent_drop_prop*100}% of {len(ents_in_data)} entities found in drop set"
)
# Get entity indices to drop
n_ents_dropped = int(ent_drop_prop * len(ents_in_data))
rng = np.random.default_rng(seed=17)
dropped_ent_idxs = rng.choice(ents_in_data,
size=n_ents_dropped,
replace=False)
# Drop entities from dictionary (subsequent processing will automatically drop corresponding mentions)
keep_mask = np.ones(len(entity_dictionary), dtype='bool')
keep_mask[dropped_ent_idxs] = False
entity_dictionary = np.array(entity_dictionary)[keep_mask]
train_processed_data, entity_dictionary, train_tensor_data = data_process.process_mention_data(
train_samples,
entity_dictionary,
tokenizer,
params["max_context_length"],
params["max_cand_length"],
context_key=params["context_key"],
multi_label_key="labels" if mult_labels else None,
silent=params["silent"],
logger=logger,
debug=params["debug"],
knn=knn,
dictionary_processed=entity_dictionary_loaded)
print("Saving processed train data...")
if not entity_dictionary_loaded:
with open(entity_dictionary_pkl_path, 'wb') as write_handle:
pickle.dump(entity_dictionary,
write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
with open(train_tensor_data_pkl_path, 'wb') as write_handle:
pickle.dump(train_tensor_data,
write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
with open(train_processed_data_pkl_path, 'wb') as write_handle:
pickle.dump(train_processed_data,
write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
# Store the query mention vectors
train_men_vecs = train_tensor_data[:][0]
if params["shuffle"]:
train_sampler = RandomSampler(train_tensor_data)
else:
train_sampler = SequentialSampler(train_tensor_data)
train_dataloader = DataLoader(train_tensor_data,
sampler=train_sampler,
batch_size=train_batch_size)
# Store the entity dictionary vectors
entity_dict_vecs = torch.tensor(list(
map(lambda x: x['ids'], entity_dictionary)),
dtype=torch.long)
# Load eval data
valid_tensor_data_pkl_path = os.path.join(pickle_src_path,
'valid_tensor_data.pickle')
valid_processed_data_pkl_path = os.path.join(
pickle_src_path, 'valid_processed_data.pickle')
if os.path.isfile(valid_tensor_data_pkl_path) and os.path.isfile(
valid_processed_data_pkl_path):
print("Loading stored processed valid data...")
with open(valid_tensor_data_pkl_path, 'rb') as read_handle:
valid_tensor_data = pickle.load(read_handle)
with open(valid_processed_data_pkl_path, 'rb') as read_handle:
valid_processed_data = pickle.load(read_handle)
else:
valid_samples = utils.read_dataset("valid", params["data_path"])
# Check if dataset has multiple ground-truth labels
mult_labels = "labels" in valid_samples[0].keys()
# Filter samples without gold entities
valid_samples = list(
filter(
lambda sample: (len(sample["labels"]) > 0) if mult_labels else
(sample["label"] is not None), valid_samples))
logger.info("Read %d valid samples." % len(valid_samples))
valid_processed_data, _, valid_tensor_data = data_process.process_mention_data(
valid_samples,
entity_dictionary,
tokenizer,
params["max_context_length"],
params["max_cand_length"],
context_key=params["context_key"],
multi_label_key="labels" if mult_labels else None,
silent=params["silent"],
logger=logger,
debug=params["debug"],
knn=knn,
dictionary_processed=True)
print("Saving processed valid data...")
with open(valid_tensor_data_pkl_path, 'wb') as write_handle:
pickle.dump(valid_tensor_data,
write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
with open(valid_processed_data_pkl_path, 'wb') as write_handle:
pickle.dump(valid_processed_data,
write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
# Store the query mention vectors
valid_men_vecs = valid_tensor_data[:][0]
# valid_sampler = SequentialSampler(valid_tensor_data)
# valid_dataloader = DataLoader(
# valid_tensor_data, sampler=valid_sampler, batch_size=eval_batch_size
# )
if params["only_evaluate"]:
evaluate(reranker,
entity_dict_vecs,
valid_men_vecs,
device=device,
logger=logger,
knn=knn,
n_gpu=n_gpu,
entity_data=entity_dictionary,
query_data=valid_processed_data,
silent=params["silent"],
use_types=use_types or params["use_types_for_eval"],
embed_batch_size=params['embed_batch_size'],
force_exact_search=use_types or params["use_types_for_eval"]
or params["force_exact_search"],
probe_mult_factor=params['probe_mult_factor'])
exit()
time_start = time.time()
utils.write_to_file(os.path.join(model_output_path, "training_params.txt"),
str(params))
logger.info("Starting training")
logger.info("device: {} n_gpu: {}, data_parallel: {}".format(
device, n_gpu, params["data_parallel"]))
# Set model to training mode
optimizer = get_optimizer(model, params)
scheduler = get_scheduler(params, optimizer, len(train_tensor_data),
logger)
best_epoch_idx = -1
best_score = -1
num_train_epochs = params["num_train_epochs"]
init_base_model_run = True if params.get("path_to_model",
None) is None else False
init_run_pkl_path = os.path.join(
pickle_src_path, f'init_run_{"type" if use_types else "notype"}.t7')
dict_embed_data = None
for epoch_idx in trange(int(num_train_epochs), desc="Epoch"):
model.train()
torch.cuda.empty_cache()
tr_loss = 0
results = None
# Check if embeddings and index can be loaded
init_run_data_loaded = False
if init_base_model_run:
if os.path.isfile(init_run_pkl_path):
logger.info('Loading init run data')
init_run_data = torch.load(init_run_pkl_path)
init_run_data_loaded = True
load_stored_data = init_base_model_run and init_run_data_loaded
# Compute mention and entity embeddings at the start of each epoch
if use_types:
if load_stored_data:
train_dict_embeddings, dict_idxs_by_type = init_run_data[
'train_dict_embeddings'], init_run_data[
'dict_idxs_by_type']
train_dict_indexes = data_process.get_index_from_embeds(
train_dict_embeddings,
dict_idxs_by_type,
force_exact_search=params['force_exact_search'],
probe_mult_factor=params['probe_mult_factor'])
train_men_embeddings, men_idxs_by_type = init_run_data[
'train_men_embeddings'], init_run_data['men_idxs_by_type']
train_men_indexes = data_process.get_index_from_embeds(
train_men_embeddings,
men_idxs_by_type,
force_exact_search=params['force_exact_search'],
probe_mult_factor=params['probe_mult_factor'])
else:
logger.info('Embedding and indexing')
if dict_embed_data is not None:
train_dict_embeddings, train_dict_indexes, dict_idxs_by_type = dict_embed_data[
'dict_embeds'], dict_embed_data[
'dict_indexes'], dict_embed_data[
'dict_idxs_by_type']
else:
train_dict_embeddings, train_dict_indexes, dict_idxs_by_type = data_process.embed_and_index(
reranker,
entity_dict_vecs,
encoder_type="candidate",
n_gpu=n_gpu,
corpus=entity_dictionary,
force_exact_search=params['force_exact_search'],
batch_size=params['embed_batch_size'],
probe_mult_factor=params['probe_mult_factor'])
train_men_embeddings, train_men_indexes, men_idxs_by_type = data_process.embed_and_index(
reranker,
train_men_vecs,
encoder_type="context",
n_gpu=n_gpu,
corpus=train_processed_data,
force_exact_search=params['force_exact_search'],
batch_size=params['embed_batch_size'],
probe_mult_factor=params['probe_mult_factor'])
else:
if load_stored_data:
train_dict_embeddings = init_run_data['train_dict_embeddings']
train_dict_index = data_process.get_index_from_embeds(
train_dict_embeddings,
force_exact_search=params['force_exact_search'],
probe_mult_factor=params['probe_mult_factor'])
train_men_embeddings = init_run_data['train_men_embeddings']
train_men_index = data_process.get_index_from_embeds(
train_men_embeddings,
force_exact_search=params['force_exact_search'],
probe_mult_factor=params['probe_mult_factor'])
else:
logger.info('Embedding and indexing')
if dict_embed_data is not None:
train_dict_embeddings, train_dict_index = dict_embed_data[
'dict_embeds'], dict_embed_data['dict_index']
else:
train_dict_embeddings, train_dict_index = data_process.embed_and_index(
reranker,
entity_dict_vecs,
encoder_type="candidate",
n_gpu=n_gpu,
force_exact_search=params['force_exact_search'],
batch_size=params['embed_batch_size'],
probe_mult_factor=params['probe_mult_factor'])
train_men_embeddings, train_men_index = data_process.embed_and_index(
reranker,
train_men_vecs,
encoder_type="context",
n_gpu=n_gpu,
force_exact_search=params['force_exact_search'],
batch_size=params['embed_batch_size'],
probe_mult_factor=params['probe_mult_factor'])
# Save the initial embeddings and index if this is the first run and data isn't persistent
if init_base_model_run and not load_stored_data:
init_run_data = {}
init_run_data['train_dict_embeddings'] = train_dict_embeddings
init_run_data['train_men_embeddings'] = train_men_embeddings
if use_types:
init_run_data['dict_idxs_by_type'] = dict_idxs_by_type
init_run_data['men_idxs_by_type'] = men_idxs_by_type
# NOTE: Cannot pickle faiss index because it is a SwigPyObject
torch.save(init_run_data,
init_run_pkl_path,
pickle_protocol=pickle.HIGHEST_PROTOCOL)
init_base_model_run = False
if params["silent"]:
iter_ = train_dataloader
else:
iter_ = tqdm(train_dataloader, desc="Batch")
logger.info("Starting KNN search...")
if not use_types:
_, dict_nns = train_dict_index.search(train_men_embeddings, knn)
else:
dict_nns = np.zeros((len(train_men_embeddings), knn))
for entity_type in train_men_indexes:
men_embeds_by_type = train_men_embeddings[
men_idxs_by_type[entity_type]]
_, dict_nns_by_type = train_dict_indexes[entity_type].search(
men_embeds_by_type, knn)
dict_nns_idxs = np.array(
list(
map(lambda x: dict_idxs_by_type[entity_type][x],
dict_nns_by_type)))
for i, idx in enumerate(men_idxs_by_type[entity_type]):
dict_nns[idx] = dict_nns_idxs[i]
logger.info("Search finished")
for step, batch in enumerate(iter_):
batch = tuple(t.to(device) for t in batch)
context_inputs, candidate_idxs, n_gold, mention_idxs = batch
mention_embeddings = train_men_embeddings[mention_idxs.cpu()]
if len(mention_embeddings.shape) == 1:
mention_embeddings = np.expand_dims(mention_embeddings, axis=0)
# context_inputs: Shape: batch x token_len
candidate_inputs = np.array(
[], dtype=np.int) # Shape: (batch*knn) x token_len
label_inputs = torch.tensor(
[[1] + [0] * (knn - 1)] * n_gold.sum(),
dtype=torch.float32) # Shape: batch(with split rows) x knn
context_inputs_split = torch.zeros(
(label_inputs.size(0), context_inputs.size(1)),
dtype=torch.long) # Shape: batch(with split rows) x token_len
# label_inputs = (candidate_idxs >= 0).type(torch.float32) # Shape: batch x knn
for i, m_embed in enumerate(mention_embeddings):
knn_dict_idxs = dict_nns[mention_idxs[i]]
knn_dict_idxs = knn_dict_idxs.astype(np.int64).flatten()
gold_idxs = candidate_idxs[i][:n_gold[i]].cpu()
for ng, gold_idx in enumerate(gold_idxs):
context_inputs_split[i + ng] = context_inputs[i]
candidate_inputs = np.concatenate(
(candidate_inputs,
np.concatenate(
([gold_idx],
knn_dict_idxs[~np.isin(knn_dict_idxs, gold_idxs)]
))[:knn]))
candidate_inputs = torch.tensor(
list(
map(lambda x: entity_dict_vecs[x].numpy(),
candidate_inputs))).cuda()
context_inputs_split = context_inputs_split.cuda()
label_inputs = label_inputs.cuda()
loss, _ = reranker(context_inputs_split,
candidate_inputs,
label_inputs,
pos_neg_loss=params["pos_neg_loss"])
if grad_acc_steps > 1:
loss = loss / grad_acc_steps
tr_loss += loss.item()
if (step + 1) % (params["print_interval"] * grad_acc_steps) == 0:
logger.info("Step {} - epoch {} average loss: {}\n".format(
step,
epoch_idx,
tr_loss / (params["print_interval"] * grad_acc_steps),
))
tr_loss = 0
loss.backward()
if (step + 1) % grad_acc_steps == 0:
torch.nn.utils.clip_grad_norm_(model.parameters(),
params["max_grad_norm"])
optimizer.step()
scheduler.step()
optimizer.zero_grad()
if (step + 1) % (params["eval_interval"] * grad_acc_steps) == 0:
logger.info("Evaluation on the development dataset")
evaluate(reranker,
entity_dict_vecs,
valid_men_vecs,
device=device,
logger=logger,
knn=knn,
n_gpu=n_gpu,
entity_data=entity_dictionary,
query_data=valid_processed_data,
silent=params["silent"],
use_types=use_types or params["use_types_for_eval"],
embed_batch_size=params['embed_batch_size'],
force_exact_search=use_types
or params["use_types_for_eval"]
or params["force_exact_search"],
probe_mult_factor=params['probe_mult_factor'])
model.train()
logger.info("\n")
logger.info("***** Saving fine-tuned model *****")
epoch_output_folder_path = os.path.join(model_output_path,
"epoch_{}".format(epoch_idx))
utils.save_model(model, tokenizer, epoch_output_folder_path)
logger.info(f"Model saved at {epoch_output_folder_path}")
normalized_accuracy, dict_embed_data = evaluate(
reranker,
entity_dict_vecs,
valid_men_vecs,
device=device,
logger=logger,
knn=knn,
n_gpu=n_gpu,
entity_data=entity_dictionary,
query_data=valid_processed_data,
silent=params["silent"],
use_types=use_types or params["use_types_for_eval"],
embed_batch_size=params['embed_batch_size'],
force_exact_search=use_types or params["use_types_for_eval"]
or params["force_exact_search"],
probe_mult_factor=params['probe_mult_factor'])
ls = [best_score, normalized_accuracy]
li = [best_epoch_idx, epoch_idx]
best_score = ls[np.argmax(ls)]
best_epoch_idx = li[np.argmax(ls)]
logger.info("\n")
execution_time = (time.time() - time_start) / 60
utils.write_to_file(
os.path.join(model_output_path, "training_time.txt"),
"The training took {} minutes\n".format(execution_time),
)
logger.info("The training took {} minutes\n".format(execution_time))
# save the best model in the parent_dir
logger.info("Best performance in epoch: {}".format(best_epoch_idx))
params["path_to_model"] = os.path.join(model_output_path,
"epoch_{}".format(best_epoch_idx))
utils.save_model(reranker.model, tokenizer, model_output_path)
logger.info(f"Best model saved at {model_output_path}")
def main(params):
output_path = params["output_path"]
if not os.path.exists(output_path):
os.makedirs(output_path)
logger = utils.get_logger(params["output_path"])
# Init model
reranker = BiEncoderRanker(params)
tokenizer = reranker.tokenizer
# laod entities
entity_dict, entity_json = load_entity_dict(logger, params)
# load tfidf candidates
tfidf_cand_dict = read_tfidf_cands(params["data_path"], params["mode"])
# load mentions
test_samples = utils.read_dataset(params["mode"], params["data_path"])
logger.info("Read %d test samples." % len(test_samples))
# get only the cands we need to tokenize
cand_ids = [c for l in tfidf_cand_dict.values() for c in l]
cand_ids.extend([x["label_umls_cuid"] for x in test_samples])
cand_ids = list(set(cand_ids))
num_cands = len(cand_ids)
# tokenize the candidates
cand_uid_map = {c : i for i, c in enumerate(cand_ids)}
candidate_pool = get_candidate_pool_tensor(
[entity_dict[c] for c in cand_ids],
tokenizer,
params["max_cand_length"],
logger
)
# create mention maps
ctxt_uid_map = {x["mm_mention_id"] : i + num_cands
for i, x in enumerate(test_samples)}
ctxt_cand_map = {x["mm_mention_id"] : x["label_umls_cuid"]
for x in test_samples}
ctxt_doc_map = {x["mm_mention_id"] : x["context_doc_id"]
for x in test_samples}
doc_ctxt_map = defaultdict(list)
for c, d in ctxt_doc_map.items():
doc_ctxt_map[d].append(c)
# create text maps for investigative evaluation
uid_to_json = {
uid : entity_json[cuid] for cuid, uid in cand_uid_map.items()
}
uid_to_json.update({i+num_cands : x for i, x in enumerate(test_samples)})
# tokenize the contexts
test_data, test_tensor_data = data.process_mention_data(
test_samples,
tokenizer,
params["max_context_length"],
params["max_cand_length"],
context_key=params['context_key'],
silent=params["silent"],
logger=logger,
debug=params["debug"],
)
context_pool = test_data["context_vecs"]
# create output variables
contexts = context_pool
context_uids = torch.LongTensor(list(ctxt_uid_map.values()))
pos_coref_ctxts = []
pos_coref_ctxt_uids = []
for i, c in enumerate(ctxt_doc_map.keys()):
assert ctxt_uid_map[c] == i + num_cands
doc = ctxt_doc_map[c]
coref_ctxts = [x for x in doc_ctxt_map[doc]
if x != c and ctxt_cand_map[x] == ctxt_cand_map[c]]
coref_ctxt_uids = [ctxt_uid_map[x] for x in coref_ctxts]
coref_ctxt_idxs = [x - num_cands for x in coref_ctxt_uids]
pos_coref_ctxts.append(context_pool[coref_ctxt_idxs])
pos_coref_ctxt_uids.append(torch.LongTensor(coref_ctxt_uids))
knn_ctxts = []
knn_ctxt_uids = []
for i, c in enumerate(ctxt_doc_map.keys()):
assert ctxt_uid_map[c] == i + num_cands
doc = ctxt_doc_map[c]
wdoc_ctxts = [x for x in doc_ctxt_map[doc] if x != c]
wdoc_ctxt_uids = [ctxt_uid_map[x] for x in wdoc_ctxts]
wdoc_ctxt_idxs = [x - num_cands for x in wdoc_ctxt_uids]
knn_ctxts.append(context_pool[wdoc_ctxt_idxs])
knn_ctxt_uids.append(torch.LongTensor(wdoc_ctxt_uids))
pos_cands = []
pos_cand_uids = []
for i, c in enumerate(ctxt_cand_map.keys()):
assert ctxt_uid_map[c] == i + num_cands
pos_cands.append(candidate_pool[cand_uid_map[ctxt_cand_map[c]]])
pos_cand_uids.append(torch.LongTensor([cand_uid_map[ctxt_cand_map[c]]]))
knn_cands = []
knn_cand_uids = []
for i, c in enumerate(ctxt_cand_map.keys()):
assert ctxt_uid_map[c] == i + num_cands
tfidf_cands = tfidf_cand_dict.get(c, [])
tfidf_cand_uids = [cand_uid_map[x] for x in tfidf_cands]
knn_cands.append(candidate_pool[tfidf_cand_uids])
knn_cand_uids.append(torch.LongTensor(tfidf_cand_uids))
tfidf_data = {
"contexts" : contexts,
"context_uids": context_uids,
"pos_coref_ctxts": pos_coref_ctxts,
"pos_coref_ctxt_uids": pos_coref_ctxt_uids,
"knn_ctxts": knn_ctxts,
"knn_ctxt_uids": knn_ctxt_uids,
"pos_cands": pos_cands,
"pos_cand_uids": pos_cand_uids,
"knn_cands": knn_cands,
"knn_cand_uids": knn_cand_uids,
"uid_to_json": uid_to_json,
}
save_data_path = os.path.join(
params['output_path'],
'joint_candidates_%s_tfidf.t7' % (params['mode'])
)
torch.save(tfidf_data, save_data_path)
def main(params):
output_path = params["output_path"]
if not os.path.exists(output_path):
os.makedirs(output_path)
logger = utils.get_logger(params["output_path"], 'log-eval')
pickle_src_path = params["pickle_src_path"]
if pickle_src_path is None or not os.path.exists(pickle_src_path):
pickle_src_path = output_path
embed_data_path = params["embed_data_path"]
if embed_data_path is None or not os.path.exists(embed_data_path):
embed_data_path = output_path
# Init model
reranker = BiEncoderRanker(params)
reranker.model.eval()
tokenizer = reranker.tokenizer
n_gpu = reranker.n_gpu
knn = params["knn"]
use_types = params["use_types"]
within_doc = params["within_doc"]
data_split = params["data_split"] # Default = "test"
# Load test data
test_samples = None
entity_dictionary_loaded = False
test_dictionary_pkl_path = os.path.join(pickle_src_path,
'test_dictionary.pickle')
test_tensor_data_pkl_path = os.path.join(pickle_src_path,
'test_tensor_data.pickle')
test_mention_data_pkl_path = os.path.join(pickle_src_path,
'test_mention_data.pickle')
if params['transductive']:
train_tensor_data_pkl_path = os.path.join(pickle_src_path,
'train_tensor_data.pickle')
train_mention_data_pkl_path = os.path.join(
pickle_src_path, 'train_mention_data.pickle')
if os.path.isfile(test_dictionary_pkl_path):
print("Loading stored processed entity dictionary...")
with open(test_dictionary_pkl_path, 'rb') as read_handle:
test_dictionary = pickle.load(read_handle)
entity_dictionary_loaded = True
if os.path.isfile(test_tensor_data_pkl_path) and os.path.isfile(
test_mention_data_pkl_path):
print("Loading stored processed test data...")
with open(test_tensor_data_pkl_path, 'rb') as read_handle:
test_tensor_data = pickle.load(read_handle)
with open(test_mention_data_pkl_path, 'rb') as read_handle:
mention_data = pickle.load(read_handle)
else:
test_samples = utils.read_dataset(data_split, params["data_path"])
if not entity_dictionary_loaded:
with open(os.path.join(params["data_path"], 'dictionary.pickle'),
'rb') as read_handle:
test_dictionary = pickle.load(read_handle)
# Check if dataset has multiple ground-truth labels
mult_labels = "labels" in test_samples[0].keys()
if params["filter_unlabeled"]:
# Filter samples without gold entities
test_samples = list(
filter(
lambda sample: (len(sample["labels"]) > 0)
if mult_labels else (sample["label"] is not None),
test_samples))
logger.info("Read %d test samples." % len(test_samples))
mention_data, test_dictionary, test_tensor_data = data_process.process_mention_data(
test_samples,
test_dictionary,
tokenizer,
params["max_context_length"],
params["max_cand_length"],
multi_label_key="labels" if mult_labels else None,
context_key=params["context_key"],
silent=params["silent"],
logger=logger,
debug=params["debug"],
knn=knn,
dictionary_processed=entity_dictionary_loaded)
print("Saving processed test data...")
if not entity_dictionary_loaded:
with open(test_dictionary_pkl_path, 'wb') as write_handle:
pickle.dump(test_dictionary,
write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
entity_dictionary_loaded = True
with open(test_tensor_data_pkl_path, 'wb') as write_handle:
pickle.dump(test_tensor_data,
write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
with open(test_mention_data_pkl_path, 'wb') as write_handle:
pickle.dump(mention_data,
write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
# Store test dictionary token ids
test_dict_vecs = torch.tensor(list(map(lambda x: x['ids'],
test_dictionary)),
dtype=torch.long)
# Store test mention token ids
test_men_vecs = test_tensor_data[:][0]
n_entities = len(test_dict_vecs)
n_mentions = len(test_tensor_data)
if within_doc:
if test_samples is None:
test_samples, _ = read_data(data_split, params, logger)
test_context_doc_ids = [s['context_doc_id'] for s in test_samples]
if params["transductive"]:
if os.path.isfile(train_tensor_data_pkl_path) and os.path.isfile(
train_mention_data_pkl_path):
print("Loading stored processed train data...")
with open(train_tensor_data_pkl_path, 'rb') as read_handle:
train_tensor_data = pickle.load(read_handle)
with open(train_mention_data_pkl_path, 'rb') as read_handle:
train_mention_data = pickle.load(read_handle)
else:
train_samples = utils.read_dataset('train', params["data_path"])
# Check if dataset has multiple ground-truth labels
mult_labels = "labels" in train_samples[0].keys()
logger.info("Read %d test samples." % len(test_samples))
train_mention_data, _, train_tensor_data = data_process.process_mention_data(
train_samples,
test_dictionary,
tokenizer,
params["max_context_length"],
params["max_cand_length"],
multi_label_key="labels" if mult_labels else None,
context_key=params["context_key"],
silent=params["silent"],
logger=logger,
debug=params["debug"],
knn=knn,
dictionary_processed=entity_dictionary_loaded)
print("Saving processed train data...")
with open(train_tensor_data_pkl_path, 'wb') as write_handle:
pickle.dump(train_tensor_data,
write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
with open(train_mention_data_pkl_path, 'wb') as write_handle:
pickle.dump(train_mention_data,
write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
# Store train mention token ids
train_men_vecs = train_tensor_data[:][0]
n_mentions += len(train_tensor_data)
n_train_mentions = len(train_tensor_data)
# Values of k to run the evaluation against
knn_vals = [0] + [2**i for i in range(int(math.log(knn, 2)) + 1)]
# Store the maximum evaluation k
max_knn = knn_vals[-1]
time_start = time.time()
# Check if graphs are already built
graph_path = os.path.join(output_path, 'graphs.pickle')
if not params['only_recall'] and os.path.isfile(graph_path):
print("Loading stored joint graphs...")
with open(graph_path, 'rb') as read_handle:
joint_graphs = pickle.load(read_handle)
else:
# Initialize graphs to store mention-mention and mention-entity similarity score edges;
# Keyed on k, the number of nearest mentions retrieved
joint_graphs = {}
for k in knn_vals:
joint_graphs[k] = {
'rows': np.array([]),
'cols': np.array([]),
'data': np.array([]),
'shape': (n_entities + n_mentions, n_entities + n_mentions)
}
# Check and load stored embedding data
embed_data_path = os.path.join(embed_data_path, 'embed_data.t7')
embed_data = None
if os.path.isfile(embed_data_path):
embed_data = torch.load(embed_data_path)
if use_types:
if embed_data is not None:
logger.info('Loading stored embeddings and computing indexes')
dict_embeds = embed_data['dict_embeds']
if 'dict_idxs_by_type' in embed_data:
dict_idxs_by_type = embed_data['dict_idxs_by_type']
else:
dict_idxs_by_type = data_process.get_idxs_by_type(
test_dictionary)
dict_indexes = data_process.get_index_from_embeds(
dict_embeds,
dict_idxs_by_type,
force_exact_search=params['force_exact_search'],
probe_mult_factor=params['probe_mult_factor'])
men_embeds = embed_data['men_embeds']
if 'men_idxs_by_type' in embed_data:
men_idxs_by_type = embed_data['men_idxs_by_type']
else:
men_idxs_by_type = data_process.get_idxs_by_type(
mention_data)
men_indexes = data_process.get_index_from_embeds(
men_embeds,
men_idxs_by_type,
force_exact_search=params['force_exact_search'],
probe_mult_factor=params['probe_mult_factor'])
else:
logger.info("Dictionary: Embedding and building index")
dict_embeds, dict_indexes, dict_idxs_by_type = data_process.embed_and_index(
reranker,
test_dict_vecs,
encoder_type="candidate",
n_gpu=n_gpu,
corpus=test_dictionary,
force_exact_search=params['force_exact_search'],
batch_size=params['embed_batch_size'],
probe_mult_factor=params['probe_mult_factor'])
logger.info("Queries: Embedding and building index")
vecs = test_men_vecs
men_data = mention_data
if params['transductive']:
vecs = torch.cat((train_men_vecs, vecs), dim=0)
men_data = train_mention_data + mention_data
men_embeds, men_indexes, men_idxs_by_type = data_process.embed_and_index(
reranker,
vecs,
encoder_type="context",
n_gpu=n_gpu,
corpus=men_data,
force_exact_search=params['force_exact_search'],
batch_size=params['embed_batch_size'],
probe_mult_factor=params['probe_mult_factor'])
else:
if embed_data is not None:
logger.info('Loading stored embeddings and computing indexes')
dict_embeds = embed_data['dict_embeds']
dict_index = data_process.get_index_from_embeds(
dict_embeds,
force_exact_search=params['force_exact_search'],
probe_mult_factor=params['probe_mult_factor'])
men_embeds = embed_data['men_embeds']
men_index = data_process.get_index_from_embeds(
men_embeds,
force_exact_search=params['force_exact_search'],
probe_mult_factor=params['probe_mult_factor'])
else:
logger.info("Dictionary: Embedding and building index")
dict_embeds, dict_index = data_process.embed_and_index(
reranker,
test_dict_vecs,
'candidate',
n_gpu=n_gpu,
force_exact_search=params['force_exact_search'],
batch_size=params['embed_batch_size'],
probe_mult_factor=params['probe_mult_factor'])
logger.info("Queries: Embedding and building index")
vecs = test_men_vecs
if params['transductive']:
vecs = torch.cat((train_men_vecs, vecs), dim=0)
men_embeds, men_index = data_process.embed_and_index(
reranker,
vecs,
'context',
n_gpu=n_gpu,
force_exact_search=params['force_exact_search'],
batch_size=params['embed_batch_size'],
probe_mult_factor=params['probe_mult_factor'])
# Save computed embedding data if not loaded from disk
if embed_data is None:
embed_data = {}
embed_data['dict_embeds'] = dict_embeds
embed_data['men_embeds'] = men_embeds
if use_types:
embed_data['dict_idxs_by_type'] = dict_idxs_by_type
embed_data['men_idxs_by_type'] = men_idxs_by_type
# NOTE: Cannot pickle faiss index because it is a SwigPyObject
torch.save(embed_data,
embed_data_path,
pickle_protocol=pickle.HIGHEST_PROTOCOL)
recall_accuracy = {
2**i: 0
for i in range(int(math.log(params['recall_k'], 2)) + 1)
}
recall_idxs = [0.] * params['recall_k']
logger.info("Starting KNN search...")
# Fetch recall_k (default 16) knn entities for all mentions
# Fetch (k+1) NN mention candidates
if not use_types:
_men_embeds = men_embeds
if params['transductive']:
_men_embeds = _men_embeds[n_train_mentions:]
nn_ent_dists, nn_ent_idxs = dict_index.search(
_men_embeds, params['recall_k'])
n_mens_to_fetch = len(_men_embeds) if within_doc else max_knn + 1
nn_men_dists, nn_men_idxs = men_index.search(
_men_embeds, n_mens_to_fetch)
else:
query_len = len(men_embeds) - (n_train_mentions
if params['transductive'] else 0)
nn_ent_idxs = np.zeros((query_len, params['recall_k']))
nn_ent_dists = np.zeros((query_len, params['recall_k']),
dtype='float64')
nn_men_idxs = -1 * np.ones((query_len, query_len), dtype=int)
nn_men_dists = -1 * np.ones(
(query_len, query_len), dtype='float64')
for entity_type in men_indexes:
men_embeds_by_type = men_embeds[men_idxs_by_type[entity_type][
men_idxs_by_type[entity_type] >=
n_train_mentions]] if params[
'transductive'] else men_embeds[
men_idxs_by_type[entity_type]]
nn_ent_dists_by_type, nn_ent_idxs_by_type = dict_indexes[
entity_type].search(men_embeds_by_type, params['recall_k'])
nn_ent_idxs_by_type = np.array(
list(
map(lambda x: dict_idxs_by_type[entity_type][x],
nn_ent_idxs_by_type)))
n_mens_to_fetch = len(
men_embeds_by_type) if within_doc else max_knn + 1
nn_men_dists_by_type, nn_men_idxs_by_type = men_indexes[
entity_type].search(
men_embeds_by_type,
min(n_mens_to_fetch, len(men_embeds_by_type)))
nn_men_idxs_by_type = np.array(
list(
map(lambda x: men_idxs_by_type[entity_type][x],
nn_men_idxs_by_type)))
i = -1
for idx in men_idxs_by_type[entity_type]:
if params['transductive']:
idx -= n_train_mentions
if idx < 0:
continue
i += 1
nn_ent_idxs[idx] = nn_ent_idxs_by_type[i]
nn_ent_dists[idx] = nn_ent_dists_by_type[i]
nn_men_idxs[idx][:len(nn_men_idxs_by_type[i]
)] = nn_men_idxs_by_type[i]
nn_men_dists[idx][:len(nn_men_dists_by_type[i]
)] = nn_men_dists_by_type[i]
logger.info("Search finished")
logger.info('Building graphs')
# Find the most similar entity and k-nn mentions for each mention query
for idx in range(len(nn_ent_idxs)):
# Get nearest entity candidate
dict_cand_idx = nn_ent_idxs[idx][0]
dict_cand_score = nn_ent_dists[idx][0]
# Compute recall metric
gold_idxs = mention_data[idx][
"label_idxs"][:mention_data[idx]["n_labels"]]
recall_idx = np.argwhere(nn_ent_idxs[idx] == gold_idxs[0])
if len(recall_idx) != 0:
recall_idx = int(recall_idx)
recall_idxs[recall_idx] += 1.
for recall_k in recall_accuracy:
if recall_idx < recall_k:
recall_accuracy[recall_k] += 1.
if not params['only_recall']:
filter_mask_neg1 = nn_men_idxs[idx] != -1
men_cand_idxs = nn_men_idxs[idx][filter_mask_neg1]
men_cand_scores = nn_men_dists[idx][filter_mask_neg1]
if within_doc:
men_cand_idxs, wd_mask = filter_by_context_doc_id(
men_cand_idxs,
test_context_doc_ids[idx],
test_context_doc_ids,
return_numpy=True)
men_cand_scores = men_cand_scores[wd_mask]
# Filter candidates to remove mention query and keep only the top k candidates
filter_mask = men_cand_idxs != idx
men_cand_idxs, men_cand_scores = men_cand_idxs[
filter_mask][:max_knn], men_cand_scores[
filter_mask][:max_knn]
if params['transductive']:
idx += n_train_mentions
# Add edges to the graphs
for k in joint_graphs:
joint_graph = joint_graphs[k]
# Add mention-entity edge
joint_graph['rows'] = np.append(
joint_graph['rows'],
[n_entities + idx
]) # Mentions added at an offset of maximum entities
joint_graph['cols'] = np.append(joint_graph['cols'],
dict_cand_idx)
joint_graph['data'] = np.append(joint_graph['data'],
dict_cand_score)
if k > 0:
# Add mention-mention edges
joint_graph['rows'] = np.append(
joint_graph['rows'],
[n_entities + idx] * len(men_cand_idxs[:k]))
joint_graph['cols'] = np.append(
joint_graph['cols'],
n_entities + men_cand_idxs[:k])
joint_graph['data'] = np.append(
joint_graph['data'], men_cand_scores[:k])
if params['transductive']:
# Add positive infinity mention-entity edges from training queries to labeled entities
for idx, train_men in enumerate(train_mention_data):
dict_cand_idx = train_men["label_idxs"][0]
for k in joint_graphs:
joint_graph = joint_graphs[k]
joint_graph['rows'] = np.append(
joint_graph['rows'],
[n_entities + idx
]) # Mentions added at an offset of maximum entities
joint_graph['cols'] = np.append(joint_graph['cols'],
dict_cand_idx)
joint_graph['data'] = np.append(joint_graph['data'],
float('inf'))
# Compute and print recall metric
recall_idx_mode = np.argmax(recall_idxs)
recall_idx_mode_prop = recall_idxs[recall_idx_mode] / np.sum(
recall_idxs)
logger.info(f"""
Recall metrics (for {len(nn_ent_idxs)} queries):
---------------""")
logger.info(
f"highest recall idx = {recall_idx_mode} ({recall_idxs[recall_idx_mode]}/{np.sum(recall_idxs)} = {recall_idx_mode_prop})"
)
for recall_k in recall_accuracy:
recall_accuracy[recall_k] /= len(nn_ent_idxs)
logger.info(f"recall@{recall_k} = {recall_accuracy[recall_k]}")
if params['only_recall']:
exit()
# Pickle the graphs
print("Saving joint graphs...")
with open(graph_path, 'wb') as write_handle:
pickle.dump(joint_graphs,
write_handle,
protocol=pickle.HIGHEST_PROTOCOL)
if params['only_embed_and_build']:
logger.info(f"Saved embedding data at: {embed_data_path}")
logger.info(f"Saved graphs at: {graph_path}")
exit()
graph_mode = params.get('graph_mode', None)
result_overview = {
'n_entities':
n_entities,
'n_mentions':
n_mentions - (n_train_mentions if params['transductive'] else 0)
}
results = {}
if graph_mode is None or graph_mode not in ['directed', 'undirected']:
results['directed'] = []
results['undirected'] = []
else:
results[graph_mode] = []
knn_fetch_time = time.time() - time_start
graph_processing_time = time.time()
n_graphs_processed = 0.
for mode in results:
print(f'\nEvaluation mode: {mode.upper()}')
for k in joint_graphs:
if k <= knn:
print(f"\nGraph (k={k}):")
# Partition graph based on cluster-linking constraints
partitioned_graph, clusters = partition_graph(
joint_graphs[k],
n_entities,
mode == 'directed',
return_clusters=True)
# Infer predictions from clusters
result = analyzeClusters(
clusters, test_dictionary, mention_data, k,
n_train_mentions if params['transductive'] else 0)
# Store result
results[mode].append(result)
n_graphs_processed += 1
avg_graph_processing_time = (time.time() -
graph_processing_time) / n_graphs_processed
avg_per_graph_time = (knn_fetch_time + avg_graph_processing_time) / 60
execution_time = (time.time() - time_start) / 60
# Store results
output_file_name = os.path.join(
output_path,
f"eval_results_{__import__('calendar').timegm(__import__('time').gmtime())}"
)
try:
for recall_k in recall_accuracy:
result_overview[f'recall@{recall_k}'] = recall_accuracy[recall_k]
except:
logger.info(
"Recall data not available since graphs were loaded from disk")
for mode in results:
mode_results = results[mode]
result_overview[mode] = {}
for r in mode_results:
k = r['knn_mentions']
result_overview[mode][f'accuracy@knn{k}'] = r['accuracy']
logger.info(f"{mode} accuracy@knn{k} = {r['accuracy']}")
output_file = f'{output_file_name}-{mode}-{k}.json'
with open(output_file, 'w') as f:
json.dump(r, f, indent=2)
print(
f"\nPredictions ({mode}) @knn{k} saved at: {output_file}")
with open(f'{output_file_name}.json', 'w') as f:
json.dump(result_overview, f, indent=2)
print(f"\nPredictions overview saved at: {output_file_name}.json")
logger.info("\nThe avg. per graph evaluation time is {} minutes\n".format(
avg_per_graph_time))
logger.info(
"\nThe total evaluation took {} minutes\n".format(execution_time))