def predict(self, batch_meta, batch_data):
self.network.eval()
task_id = batch_meta['task_id']
task_def = TaskDef.from_dict(batch_meta['task_def'])
task_type = task_def.task_type
task_obj = tasks.get_task_obj(task_def)
inputs = batch_data[:batch_meta['input_len']]
if len(inputs) == 3:
inputs.append(None)
inputs.append(None)
inputs.append(task_id)
score = self.mnetwork(*inputs)
if task_obj is not None:
score, predict = task_obj.test_predict(score)
elif task_type == TaskType.Ranking:
score = score.contiguous().view(-1, batch_meta['pairwise_size'])
assert task_type == TaskType.Ranking
score = F.softmax(score, dim=1)
score = score.data.cpu()
score = score.numpy()
predict = np.zeros(score.shape, dtype=int)
positive = np.argmax(score, axis=1)
for idx, pos in enumerate(positive):
predict[idx, pos] = 1
predict = predict.reshape(-1).tolist()
score = score.reshape(-1).tolist()
return score, predict, batch_meta['true_label']
elif task_type == TaskType.SeqenceLabeling:
mask = batch_data[batch_meta['mask']]
score = score.contiguous()
score = score.data.cpu()
score = score.numpy()
predict = np.argmax(score, axis=1).reshape(mask.size()).tolist()
valied_lenght = mask.sum(1).tolist()
final_predict = []
for idx, p in enumerate(predict):
final_predict.append(p[:valied_lenght[idx]])
score = score.reshape(-1).tolist()
return score, final_predict, batch_meta['label']
elif task_type == TaskType.Span:
start, end = score
predictions = []
if self.config['encoder_type'] == EncoderModelType.BERT:
import experiments.squad.squad_utils as mrc_utils
scores, predictions = mrc_utils.extract_answer(
batch_meta,
batch_data,
start,
end,
self.config.get('max_answer_len', 5),
do_lower_case=self.config.get('do_lower_case', False))
return scores, predictions, batch_meta['answer']
else:
raise ValueError("Unknown task_type: %s" % task_type)
return score, predict, batch_meta['label']
def load_model_for_viz_1(task_def_path,
checkpoint_path,
input_path,
model_type='bert-base-cased',
do_lower_case=False,
use_cuda=True):
# load task info
task = os.path.splitext(os.path.basename(task_def_path))[0]
task_defs = TaskDefs(task_def_path)
assert task in task_defs._task_type_map
assert task in task_defs._data_type_map
assert task in task_defs._metric_meta_map
prefix = task.split('_')[0]
task_def = task_defs.get_task_def(prefix)
data_type = task_defs._data_type_map[task]
task_type = task_defs._task_type_map[task]
metric_meta = task_defs._metric_meta_map[task]
# load model
assert os.path.exists(checkpoint_path)
state_dict = torch.load(checkpoint_path)
config = state_dict['config']
config["cuda"] = use_cuda
device = torch.device("cuda" if use_cuda else "cpu")
task_def = task_defs.get_task_def(prefix)
task_def_list = [task_def]
config['task_def_list'] = task_def_list
## temp fix
config['fp16'] = False
config['answer_opt'] = 0
config['adv_train'] = False
#del state_dict['optimizer']
config['output_attentions'] = True
config['local_rank'] = -1
model = MTDNNModel(config, device, state_dict=state_dict)
encoder_type = config.get('encoder_type', EncoderModelType.BERT)
root = os.path.basename(task_def_path)
literal_model_type = model_type.split('-')[0].upper()
encoder_model = EncoderModelType[literal_model_type]
literal_model_type = literal_model_type.lower()
mt_dnn_suffix = literal_model_type
if 'base' in model_type:
mt_dnn_suffix += "_base"
elif 'large' in model_type:
mt_dnn_suffix += "_large"
# load tokenizer
config_class, model_class, tokenizer_class = MODEL_CLASSES[
literal_model_type]
tokenizer = tokenizer_class.from_pretrained(model_type,
do_lower_case=do_lower_case)
# load data
prep_input = input_path
test_data_set = SingleTaskDataset(prep_input,
False,
maxlen=512,
task_id=0,
task_def=task_def)
collater = Collater(is_train=False, encoder_type=encoder_type)
test_data = DataLoader(test_data_set,
batch_size=1,
collate_fn=collater.collate_fn,
pin_memory=True)
idx = 0
results = []
for batch_meta, batch_data in tqdm(test_data):
if idx < 360:
idx += 1
continue
batch_meta, batch_data = Collater.patch_data(device, batch_meta,
batch_data)
model.network.eval()
task_id = batch_meta['task_id']
task_def = TaskDef.from_dict(batch_meta['task_def'])
task_type = task_def.task_type
task_obj = tasks.get_task_obj(task_def)
inputs = batch_data[:batch_meta['input_len']]
if len(inputs) == 3:
inputs.append(None)
inputs.append(None)
inputs.append(task_id)
input_ids = inputs[0]
token_type_ids = inputs[1]
attention = model.mnetwork.module.bert(
input_ids, token_type_ids=token_type_ids)[-1]
batch_size = batch_data[0].shape[0]
for i in range(batch_size):
attention = tuple([item[i:i + 1, :, :, :] for item in attention])
input_id_list = input_ids[i].tolist()
tokens = tokenizer.convert_ids_to_tokens(input_id_list)
idx_sep = listRightIndex(tokens, '[SEP]') + 1
tokens = tokens[:idx_sep]
attention = tuple(
[item[:, :, :idx_sep, :idx_sep] for item in attention])
results.append((attention, tokens))
idx += batch_size
return results
def predict(self, batch_meta, batch_data):
self.network.eval()
task_id = batch_meta["task_id"]
task_def = TaskDef.from_dict(batch_meta["task_def"])
task_type = task_def.task_type
task_obj = tasks.get_task_obj(task_def)
inputs = batch_data[:batch_meta["input_len"]]
if len(inputs) == 3:
inputs.append(None)
inputs.append(None)
inputs.append(task_id)
if task_type == TaskType.SeqenceGeneration:
# y_idx, #3 -> gen
inputs.append(None)
inputs.append(3)
score = self.mnetwork(*inputs)
if task_obj is not None:
score, predict = task_obj.test_predict(score)
elif task_type == TaskType.Ranking:
score = score.contiguous().view(-1, batch_meta["pairwise_size"])
assert task_type == TaskType.Ranking
score = F.softmax(score, dim=1)
score = score.data.cpu()
score = score.numpy()
predict = np.zeros(score.shape, dtype=int)
positive = np.argmax(score, axis=1)
for idx, pos in enumerate(positive):
predict[idx, pos] = 1
predict = predict.reshape(-1).tolist()
score = score.reshape(-1).tolist()
return score, predict, batch_meta["true_label"]
elif task_type == TaskType.SeqenceLabeling:
mask = batch_data[batch_meta["mask"]]
score = score.contiguous()
score = score.data.cpu()
score = score.numpy()
predict = np.argmax(score, axis=1).reshape(mask.size()).tolist()
valied_lenght = mask.sum(1).tolist()
final_predict = []
for idx, p in enumerate(predict):
final_predict.append(p[:valied_lenght[idx]])
score = score.reshape(-1).tolist()
return score, final_predict, batch_meta["label"]
elif task_type == TaskType.Span or task_type == TaskType.SpanYN:
predictions = []
features = []
for idx, offset in enumerate(batch_meta["offset_mapping"]):
token_is_max_context = (
batch_meta["token_is_max_context"][idx] if batch_meta.get(
"token_is_max_context", None) else None)
sample_id = batch_meta["uids"][idx]
if "label" in batch_meta:
feature = {
"offset_mapping": offset,
"token_is_max_context": token_is_max_context,
"uid": sample_id,
"context": batch_meta["context"][idx],
"answer": batch_meta["answer"][idx],
"label": batch_meta["label"][idx],
}
else:
feature = {
"offset_mapping": offset,
"token_is_max_context": token_is_max_context,
"uid": sample_id,
"context": batch_meta["context"][idx],
"answer": batch_meta["answer"][idx],
}
if "null_ans_index" in batch_meta:
feature["null_ans_index"] = batch_meta["null_ans_index"]
features.append(feature)
start, end = score
start = start.contiguous()
start = start.data.cpu()
start = start.numpy().tolist()
end = end.contiguous()
end = end.data.cpu()
end = end.numpy().tolist()
return (start, end), predictions, features
elif task_type == TaskType.SeqenceGeneration:
predicts = self.tokenizer.batch_decode(score,
skip_special_tokens=True)
predictions = {}
golds = {}
for idx, predict in enumerate(predicts):
sample_id = batch_meta["uids"][idx]
answer = batch_meta["answer"][idx]
predict = predict.strip()
if predict == DUMPY_STRING_FOR_EMPTY_ANS:
predict = ""
predictions[sample_id] = predict
golds[sample_id] = answer
score = score.contiguous()
score = score.data.cpu()
score = score.numpy().tolist()
return score, predictions, golds
elif task_type == TaskType.ClozeChoice:
score = score.contiguous().view(-1)
score = score.data.cpu()
score = score.numpy()
copy_score = score.tolist()
answers = batch_meta["answer"]
choices = batch_meta["choice"]
chunks = batch_meta["pairwise_size"]
uids = batch_meta["uids"]
predictions = {}
golds = {}
for chunk in chunks:
uid = uids[0]
answer = eval(answers[0])
choice = eval(choices[0])
answers = answers[chunk:]
choices = choices[chunk:]
current_p = score[:chunk]
score = score[chunk:]
positive = np.argmax(current_p)
predict = choice[positive]
predictions[uid] = predict
golds[uid] = answer
return copy_score, predictions, golds
else:
raise ValueError("Unknown task_type: %s" % task_type)
return score, predict, batch_meta["label"]