def predict(args, model, tokenizer, prefix=""):
pred_output_dir = args.output_dir
if not os.path.exists(pred_output_dir) and args.local_rank in [-1, 0]:
os.makedirs(pred_output_dir)
test_dataset = load_and_cache_examples(args,
args.task_name,
tokenizer,
data_type='test')
print(len(test_dataset))
# Note that DistributedSampler samples randomly
test_sampler = SequentialSampler(
test_dataset) if args.local_rank == -1 else DistributedSampler(
test_dataset)
test_dataloader = DataLoader(test_dataset,
sampler=test_sampler,
batch_size=1,
collate_fn=collate_fn)
# Eval!
logger.info("***** Running prediction %s *****", prefix)
logger.info(" Num examples = %d", len(test_dataset))
logger.info(" Batch size = %d", 1)
results = []
output_submit_file = os.path.join(pred_output_dir, prefix,
"test_prediction.json")
pbar = ProgressBar(n_total=len(test_dataloader), desc="Predicting")
for step, batch in enumerate(test_dataloader):
model.eval()
batch = tuple(t.to(args.device) for t in batch)
with torch.no_grad():
inputs = {
"input_ids": batch[0],
"attention_mask": batch[1],
"start_positions": None,
"end_positions": None
}
if args.model_type != "distilbert":
# XLM and RoBERTa don"t use segment_ids
inputs["token_type_ids"] = (batch[2] if args.model_type
in ["bert", "xlnet"] else None)
outputs = model(**inputs)
start_logits, end_logits = outputs[:2]
R = bert_extract_item(start_logits, end_logits)
if R:
label_entities = [[args.id2label[x[0]], x[1], x[2]] for x in R]
else:
label_entities = []
json_d = {}
json_d['id'] = step
json_d['entities'] = label_entities
results.append(json_d)
pbar(step)
print(" ")
with open(output_submit_file, "w") as writer:
for record in results:
writer.write(json.dumps(record) + '\n')
def evaluate(args, model, tokenizer, prefix=""):
metric = SpanEntityScore(args.id2label)
eval_output_dir = args.output_dir
if not os.path.exists(eval_output_dir) and args.local_rank in [-1, 0]:
os.makedirs(eval_output_dir)
eval_features = load_and_cache_examples(args, args.task_name, tokenizer, data_type='dev')
args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)
# Eval!
logger.info("***** Running evaluation %s *****", prefix)
logger.info(" Num examples = %d", len(eval_features))
logger.info(" Batch size = %d", args.eval_batch_size)
eval_loss = 0.0
nb_eval_steps = 0
pbar = ProgressBar(n_total=len(eval_features), desc="Evaluating")
for step, f in enumerate(eval_features):
input_lens = f.input_len
input_ids = torch.tensor([f.input_ids[:input_lens]], dtype=torch.long).to(args.device)
input_mask = torch.tensor([f.input_mask[:input_lens]], dtype=torch.long).to(args.device)
segment_ids = torch.tensor([f.segment_ids[:input_lens]], dtype=torch.long).to(args.device)
start_ids = torch.tensor([f.start_ids[:input_lens]], dtype=torch.long).to(args.device)
end_ids = torch.tensor([f.end_ids[:input_lens]], dtype=torch.long).to(args.device)
subjects = f.subjects
model.eval()
with torch.no_grad():
inputs = {"input_ids": input_ids, "attention_mask": input_mask,
"start_positions": start_ids, "end_positions": end_ids}
if args.model_type != "distilbert":
# XLM and RoBERTa don"t use segment_ids
inputs["token_type_ids"] = (segment_ids if args.model_type in ["bert", "xlnet"] else None)
outputs = model(**inputs)
tmp_eval_loss, start_logits, end_logits = outputs[:3]
R = bert_extract_item(start_logits, end_logits)
T = subjects
metric.update(true_subject=T, pred_subject=R)
if args.n_gpu > 1:
tmp_eval_loss = tmp_eval_loss.mean() # mean() to average on multi-gpu parallel evaluating
eval_loss += tmp_eval_loss.item()
nb_eval_steps += 1
pbar(step)
logger.info("\n")
eval_loss = eval_loss / nb_eval_steps
eval_info, entity_info = metric.result()
results = {f'{key}': value for key, value in eval_info.items()}
results['loss'] = eval_loss
logger.info("***** Eval results %s *****", prefix)
info = "-".join([f' {key}: {value:.4f} ' for key, value in results.items()])
logger.info(info)
logger.info("***** Entity results %s *****", prefix)
for key in sorted(entity_info.keys()):
print("******* %s results ********" % key)
info = "-".join([f' {key}: {value:.4f} ' for key, value in entity_info[key].items()])
print(info)
return results
def predict(args, model, tokenizer, prefix=""):
pred_output_dir = args.output_dir
if not os.path.exists(pred_output_dir) and args.local_rank in [-1, 0]:
os.makedirs(pred_output_dir)
test_dataset = load_and_cache_examples(args,
args.task_name,
tokenizer,
data_type='test')
print(len(test_dataset))
# Note that DistributedSampler samples randomly
test_sampler = SequentialSampler(
test_dataset) if args.local_rank == -1 else DistributedSampler(
test_dataset)
test_dataloader = DataLoader(test_dataset,
sampler=test_sampler,
batch_size=1,
collate_fn=collate_fn)
# Eval!
logger.info("***** Running prediction %s *****", prefix)
logger.info(" Num examples = %d", len(test_dataset))
logger.info(" Batch size = %d", 1)
results = []
output_predict_file = os.path.join(pred_output_dir, prefix,
"test_predict.json")
pbar = ProgressBar(n_total=len(test_dataloader), desc="Predicting")
for step, batch in enumerate(test_dataloader):
model.eval()
batch = tuple(t.to(args.device) for t in batch)
with torch.no_grad():
inputs = {
"input_ids": batch[0],
"attention_mask": batch[1],
"start_positions": None,
"end_positions": None
}
if args.model_type != "distilbert":
# XLM and RoBERTa don"t use segment_ids
inputs["token_type_ids"] = (batch[2] if args.model_type
in ["bert", "xlnet"] else None)
outputs = model(**inputs)
start_logits, end_logits = outputs[:2]
R = bert_extract_item(start_logits, end_logits)
if R:
label_entities = [[args.id2label[x[0]], x[1], x[2]] for x in R]
else:
label_entities = []
json_d = {}
json_d['id'] = step
json_d['entities'] = label_entities
results.append(json_d)
pbar(step)
logger.info("\n")
with open(output_predict_file, "w") as writer:
for record in results:
writer.write(json.dumps(record) + '\n')
if args.task_name == "cluener":
output_submit_file = os.path.join(pred_output_dir, prefix,
"test_submit.json")
test_text = []
with open(os.path.join(args.data_dir, "test.json"), 'r') as fr:
for line in fr:
test_text.append(json.loads(line))
test_submit = []
for x, y in zip(test_text, results):
json_d = {}
json_d['id'] = x['id']
json_d['label'] = {}
entities = y['entities']
words = list(x['text'])
if len(entities) != 0:
for subject in entities:
tag = subject[0]
start = subject[1]
end = subject[2]
word = "".join(words[start:end + 1])
if tag in json_d['label']:
if word in json_d['label'][tag]:
json_d['label'][tag][word].append([start, end])
else:
json_d['label'][tag][word] = [[start, end]]
else:
json_d['label'][tag] = {}
json_d['label'][tag][word] = [[start, end]]
test_submit.append(json_d)
json_to_text(output_submit_file, test_submit)
