def _save_raw_predictions(checkpoint: str,
raw_predictions: Sequence[Mapping[str, np.ndarray]],
use_wordpiece: bool) -> None:
"""Save raw prediction to file as tf.Examples."""
output_file = f"{checkpoint}.predicted-tfrecords"
with tf.python_io.TFRecordWriter(output_file) as writer:
for raw_prediction in raw_predictions:
features = collections.OrderedDict()
for output_name in ["unique_ids", "type", "level"]:
features[output_name] = input_utils.create_bytes_feature(
[raw_prediction[output_name]])
for output_name in [
"long_token_ids", "long_sentence_ids",
"long_token_type_ids", "global_token_ids",
"global_sentence_ids", "global_paragraph_ids",
"answer_begin_top_indices", "answer_end_top_indices",
"answer_types"
]:
features[output_name] = input_utils.create_int_feature(
raw_prediction[output_name])
for output_name in [
"supporting_facts_probs",
"answer_begin_top_probs",
"answer_end_top_probs",
]:
features[output_name] = input_utils.create_float_feature(
raw_prediction[output_name])
if use_wordpiece:
features[
"long_tokens_to_unigrams"] = input_utils.create_int_feature(
raw_prediction["long_tokens_to_unigrams"])
writer.write(
tf.train.Example(features=tf.train.Features(
feature=features)).SerializeToString())
def to_tf_example(self) -> tf.train.Example:
"""Returns a TF Example."""
# All features are int features except for these float features.
float_features = {
'global_x_coords', 'global_y_coords', 'global_widths',
'global_heights', 'global_parent_x_coords',
'global_parent_y_coords', 'global_parent_widths',
'global_parent_heights'
}
fields = attr.asdict(self)
assert all(x in fields for x in float_features)
features = collections.OrderedDict()
for name in attr.fields_dict(OpenKpEtcFeatures).keys():
values = getattr(self, name)
if name in float_features:
features[name] = input_utils.create_float_feature(values)
else:
features[name] = input_utils.create_int_feature(values)
return tf.train.Example(features=tf.train.Features(feature=features))
def _process_prediction(
prediction: Mapping[Text, np.ndarray],
text_examples: Mapping[Text, eval_utils.OpenKpTextExample],
writer_tfrecord,
writer_jsonl,
metrics: Optional[Mapping[int, _MetricAverager]] = None) -> None:
"""Processes a single TF `Estimator.predict` prediction.
This function assumes that `Estimator.predict` was called with
`yield_single_examples=True`.
Args:
prediction: Prediction from `Estimator.predict` for a single example.
text_examples: A dictionary of `OpenKpTextExample` objects, keyed by URL.
This is used to generate the KeyPhrase predictions based on the ngram
logits in the prediction.
writer_tfrecord: An open `tf.python_io.TFRecordWriter` to write to.
writer_jsonl: An open text file writer to write JSON Lines to.
metrics: Optional `_MetricAverager`s to update with this prediction. If
None, metric calculation is skipped completely. None is appropriate for
example if we're just running inference for unlabeled examples.
"""
# [kp_max_length, long_max_length] shape.
ngram_logits = prediction['ngram_logits']
features = collections.OrderedDict()
features['ngram_logits'] = input_utils.create_float_feature(
ngram_logits.flatten())
position_predictions = eval_utils.logits_to_predictions(
ngram_logits, max_predictions=FLAGS.max_position_predictions)
# Sort predictions for convenience.
position_predictions.sort(key=lambda x: x.logit, reverse=True)
features['top_pos_logit'] = input_utils.create_float_feature(
x.logit for x in position_predictions)
features['top_pos_start_idx'] = input_utils.create_int_feature(
x.start_idx for x in position_predictions)
features['top_pos_phrase_len'] = input_utils.create_int_feature(
x.phrase_len for x in position_predictions)
url = ''.join(chr(x) for x in prediction['url_code_points'] if x != -1)
features['url'] = input_utils.create_bytes_feature([url])
if url in text_examples:
text_example = text_examples[url]
kp_predictions = text_example.get_key_phrase_predictions(
position_predictions, max_predictions=FLAGS.max_kp_predictions)
if len(kp_predictions) < FLAGS.max_kp_predictions:
tf.logging.warn(
f'Made fewer than `max_kp_predictions` for URL: {url}')
writer_jsonl.write(
json.dumps({
'url': url,
'KeyPhrases': [[kp] for kp in kp_predictions]
}) + '\n')
features['kp_predictions'] = input_utils.create_bytes_feature(
kp_predictions)
if metrics is not None:
precision, recall, f1 = text_example.get_score_full(kp_predictions)
for i in (1, 3, 5):
p = precision[i - 1]
r = recall[i - 1]
f = f1[i - 1]
features[f'p_at_{i}'] = input_utils.create_float_feature([p])
features[f'r_at_{i}'] = input_utils.create_float_feature([r])
features[f'f1_at_{i}'] = input_utils.create_float_feature([f])
metrics[i].add_example(precision=p, recall=r, f1=f)
else:
tf.logging.error(f'No text example found for URL: {url}')
writer_tfrecord.write(
tf.train.Example(features=tf.train.Features(
feature=features)).SerializeToString())