import os
import time
import json
import math
import copy
import collections
from typing import Optional, List, Dict, Tuple, Callable, Any, Union, NewType
import numpy as np
from tqdm import tqdm
import datasets
from transformers import AutoTokenizer
from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
from transformers.utils import logging
from transformers.trainer_utils import EvalLoopOutput, EvalPrediction
from .args import (
HfArgumentParser,
RetroArguments,
TrainingArguments,
)
from .base import BaseReader
from . import constants as C
from .preprocess import (
get_sketch_features,
get_intensive_features
)
from .metrics import (
compute_classification_metric,
compute_squad_v2
)
DataClassType = NewType("DataClassType", Any)
logger = logging.get_logger(__name__)
class SketchReader(BaseReader):
name: str = "sketch"
def postprocess(
self,
output: Union[np.ndarray, EvalLoopOutput],
eval_examples: datasets.Dataset,
eval_dataset: datasets.Dataset,
mode: str = "evaluate",
) -> Union[EvalPrediction, Dict[str, float]]:
# External Front Verification (E-FV)
if isinstance(output, EvalLoopOutput):
logits = output.predictions
else:
logits = output
example_id_to_index = {k: i for i, k in enumerate(eval_examples[C.ID_COLUMN_NAME])}
features_per_example = collections.defaultdict(list)
for i, feature in enumerate(eval_dataset):
features_per_example[example_id_to_index[feature["example_id"]]].append(i)
count_map = {k: len(v) for k, v in features_per_example.items()}
logits_ans = np.zeros(len(count_map))
logits_na = np.zeros(len(count_map))
for example_index, example in enumerate(tqdm(eval_examples)):
feature_indices = features_per_example[example_index]
n_strides = count_map[example_index]
logits_ans[example_index] += logits[example_index, 0] / n_strides
logits_na[example_index] += logits[example_index, 1] / n_strides
# Calculate E-FV score
score_ext = logits_ans - logits_na
# Save external front verification score
final_map = dict(zip(eval_examples[C.ID_COLUMN_NAME], score_ext.tolist()))
with open(os.path.join(self.args.output_dir, C.SCORE_EXT_FILE_NAME), "w") as writer:
writer.write(json.dumps(final_map, indent=4) + "\n")
if mode == "evaluate":
return EvalPrediction(
predictions=logits, label_ids=output.label_ids,
)
else:
return final_map
class IntensiveReader(BaseReader):
name: str = "intensive"
def postprocess(
self,
output: EvalLoopOutput,
eval_examples: datasets.Dataset,
eval_dataset: datasets.Dataset,
log_level: int = logging.WARNING,
mode: str = "evaluate",
) -> Union[List[Dict[str, Any]], EvalPrediction]:
# Internal Front Verification (I-FV)
# Verification is already done inside the model
# Post-processing: we match the start logits and end logits to answers in the original context.
predictions, nbest_json, scores_diff_json = self.compute_predictions(
eval_examples,
eval_dataset,
output.predictions,
version_2_with_negative=self.data_args.version_2_with_negative,
n_best_size=self.data_args.n_best_size,
max_answer_length=self.data_args.max_answer_length,
null_score_diff_threshold=self.data_args.null_score_diff_threshold,
output_dir=self.args.output_dir,
log_level=log_level,
n_tops=(self.data_args.start_n_top, self.data_args.end_n_top),
)
if mode == "retro_inference":
return nbest_json, scores_diff_json
# Format the result to the format the metric expects.
if self.data_args.version_2_with_negative:
formatted_predictions = [
{"id": k, "prediction_text": v, "no_answer_probability": scores_diff_json[k]}
for k, v in predictions.items()
]
else:
formatted_predictions = [
{"id": k, "prediction_text": v}
for k, v in predictions.items()
]
if mode == "predict":
return formatted_predictions
else:
references = [
{"id": ex[C.ID_COLUMN_NAME], "answers": ex[C.ANSWER_COLUMN_NAME]}
for ex in eval_examples
]
return EvalPrediction(
predictions=formatted_predictions, label_ids=references
)
def compute_predictions(
self,
examples: datasets.Dataset,
features: datasets.Dataset,
predictions: Tuple[np.ndarray, np.ndarray],
version_2_with_negative: bool = False,
n_best_size: int = 20,
max_answer_length: int = 30,
null_score_diff_threshold: float = 0.0,
output_dir: Optional[str] = None,
log_level: Optional[int] = logging.WARNING,
n_tops: Tuple[int, int] = (-1, -1),
use_choice_logits: bool = False,
):
# Threshold-based Answerable Verification (TAV)
if len(predictions) not in [2, 3]:
raise ValueError("`predictions` should be a tuple with two or three elements "
"(start_logits, end_logits, choice_logits).")
all_start_logits, all_end_logits = predictions[:2]
all_choice_logits = None
if len(predictions) == 3:
all_choice_logits = predictions[-1]
# Build a map example to its corresponding features.
example_id_to_index = {k: i for i, k in enumerate(examples[C.ID_COLUMN_NAME])}
features_per_example = collections.defaultdict(list)
for i, feature in enumerate(features):
features_per_example[example_id_to_index[feature["example_id"]]].append(i)
all_predictions = collections.OrderedDict()
all_nbest_json = collections.OrderedDict()
scores_diff_json = collections.OrderedDict() if version_2_with_negative else None
# Logging.
logger.setLevel(log_level)
logger.info(f"Post-processing {len(examples)} example predictions split into {len(features)} features.")
# Let's loop over all the examples!
for example_index, example in enumerate(tqdm(examples)):
# Those are the indices of the features associated to the current example.
feature_indices = features_per_example[example_index]
min_null_prediction = None
prelim_predictions = []
# Looping through all the features associated to the current example.
for feature_index in feature_indices:
# We grab the predictions of the model for this feature.
start_logits = all_start_logits[feature_index]
end_logits = all_end_logits[feature_index]
# score_null = s1 + e1
feature_null_score = start_logits[0] + end_logits[0]
if all_choice_logits is not None:
choice_logits = all_choice_logits[feature_index]
if use_choice_logits:
feature_null_score = choice_logits[1]
# This is what will allow us to map some the positions
# in our logits to span of texts in the original context.
offset_mapping = features[feature_index]["offset_mapping"]
# Optional `token_is_max_context`,
# if provided we will remove answers that do not have the maximum context
# available in the current feature.
token_is_max_context = features[feature_index].get("token_is_max_context", None)
# Update minimum null prediction.
if (
min_null_prediction is None or
min_null_prediction["score"] > feature_null_score
):
min_null_prediction = {
"offsets": (0, 0),
"score": feature_null_score,
"start_logit": start_logits[0],
"end_logit": end_logits[0],
}
# Go through all possibilities for the {top k} greater start and end logits
start_indexes = np.argsort(start_logits)[-1 : -n_best_size - 1 : -1].tolist()
end_indexes = np.argsort(end_logits)[-1 : -n_best_size - 1 : -1].tolist()
for start_index in start_indexes:
for end_index in end_indexes:
# Don't consider out-of-scope answers!
# either because the indices are out of bounds
# or correspond to part of the input_ids that are note in the context.
if (
start_index >= len(offset_mapping)
or end_index >= len(offset_mapping)
or not offset_mapping[start_index]
or not offset_mapping[end_index]
):
continue
# Don't consider answers with a length negative or > max_answer_length.
if end_index < start_index or end_index - start_index + 1 > max_answer_length:
continue
# Don't consider answer that don't have the maximum context available
# (if such information is provided).
if token_is_max_context is not None and not token_is_max_context.get(str(start_index), False):
continue
prelim_predictions.append(
{
"offsets": (offset_mapping[start_index][0], offset_mapping[end_index][1]),
"score": start_logits[start_index] + end_logits[end_index],
"start_logit": start_logits[start_index],
"end_logit": end_logits[end_index],
}
)
if version_2_with_negative:
# Add the minimum null prediction
prelim_predictions.append(min_null_prediction)
null_score = min_null_prediction["score"]
# Only keep the best `n_best_size` predictions
predictions = sorted(prelim_predictions, key=lambda x: x["score"], reverse=True)[:n_best_size]
# Add back the minimum null prediction if it was removed because of its low score.
if version_2_with_negative and not any(p["offsets"] == (0, 0) for p in predictions):
predictions.append(min_null_prediction)
# Use the offsets to gather the answer text in the original context.
context = example["context"]
for pred in predictions:
offsets = pred.pop("offsets")
pred["text"] = context[offsets[0] : offsets[1]]
# In the very rare edge case we have not a single non-null prediction,
# we create a fake prediction to avoid failure.
if len(predictions) == 0 or (len(predictions) == 1 and predictions[0]["text"] == ""):
predictions.insert(0, {"text": "", "start_logit": 0.0, "end_logit": 0.0, "score": 0.0,})
# Compute the softmax of all scores
# (we do it with numpy to stay independent from torch/tf) in this file,
# using the LogSum trick).
scores = np.array([pred.pop("score") for pred in predictions])
exp_scores = np.exp(scores - np.max(scores))
probs = exp_scores / exp_scores.sum()
# Include the probabilities in our predictions.
for prob, pred in zip(probs, predictions):
pred["probability"] = prob
# Pick the best prediction. If the null answer is not possible, this is easy.
if not version_2_with_negative:
all_predictions[example[C.ID_COLUMN_NAME]] = predictions[0]["text"]
else:
# Otherwise we first need to find the best non-empty prediction.
i = 0
try:
while predictions[i]["text"] == "":
i += 1
except:
i = 0
best_non_null_pred = predictions[i]
# Then we compare to the null prediction using the threshold.
score_diff = null_score - best_non_null_pred["start_logit"] - best_non_null_pred["end_logit"]
scores_diff_json[example[C.ID_COLUMN_NAME]] = float(score_diff) # To be JSON-serializable.
if score_diff > null_score_diff_threshold:
all_predictions[example[C.ID_COLUMN_NAME]] = ""
else:
all_predictions[example[C.ID_COLUMN_NAME]] = best_non_null_pred["text"]
# Make `predictions` JSON-serializable by casting np.float back to float.
all_nbest_json[example[C.ID_COLUMN_NAME]] = [
{k: (float(v) if isinstance(v, (np.float16, np.float32, np.float64)) else v) for k, v in pred.items()}
for pred in predictions
]
# If we have an output_dir, let's save all those dicts.
if output_dir is not None:
if not os.path.isdir(output_dir):
raise EnvironmentError(f"{output_dir} is not a directory.")
prediction_file = os.path.join(output_dir, C.INTENSIVE_PRED_FILE_NAME)
nbest_file = os.path.join(output_dir, C.NBEST_PRED_FILE_NAME)
if version_2_with_negative:
null_odds_file = os.path.join(output_dir, C.SCORE_DIFF_FILE_NAME)
logger.info(f"Saving predictions to {prediction_file}.")
with open(prediction_file, "w") as writer:
writer.write(json.dumps(all_predictions, indent=4) + "\n")
logger.info(f"Saving nbest_preds to {nbest_file}.")
with open(nbest_file, "w") as writer:
writer.write(json.dumps(all_nbest_json, indent=4) + "\n")
if version_2_with_negative:
logger.info(f"Saving null_odds to {null_odds_file}.")
with open(null_odds_file, "w") as writer:
writer.write(json.dumps(scores_diff_json, indent=4) + "\n")
return all_predictions, all_nbest_json, scores_diff_json
class RearVerifier:
def __init__(
self,
beta1: int = 1,
beta2: int = 1,
best_cof: int = 1,
thresh: float = 0.0,
):
self.beta1 = beta1
self.beta2 = beta2
self.best_cof = best_cof
self.thresh = thresh
def __call__(
self,
score_ext: Dict[str, float],
score_diff: Dict[str, float],
nbest_preds: Dict[str, Dict[int, Dict[str, float]]]
):
all_scores = collections.OrderedDict()
assert score_ext.keys() == score_diff.keys()
for key in score_ext.keys():
if key not in all_scores:
all_scores[key] = []
all_scores[key].extend(
[self.beta1 * score_ext[key],
self.beta2 * score_diff[key]]
)
output_scores = {}
for key, scores in all_scores.items():
mean_score = sum(scores) / float(len(scores))
output_scores[key] = mean_score
all_nbest = collections.OrderedDict()
for key, entries in nbest_preds.items():
if key not in all_nbest:
all_nbest[key] = collections.defaultdict(float)
for entry in entries:
prob = self.best_cof * entry["probability"]
all_nbest[key][entry["text"]] += prob
output_predictions = {}
for key, entry_map in all_nbest.items():
sorted_texts = sorted(
entry_map.keys(), key=lambda x: entry_map[x], reverse=True
)
best_text = sorted_texts[0]
output_predictions[key] = best_text
for qid in output_predictions.keys():
if output_scores[qid] > self.thresh:
output_predictions[qid] = ""
return output_predictions, output_scores
class RetroReader:
def __init__(
self,
args,
sketch_reader: SketchReader,
intensive_reader: IntensiveReader,
rear_verifier: RearVerifier,
prep_fn: Tuple[Callable, Callable],
):
self.args = args
# Set submodules
self.sketch_reader = sketch_reader
self.intensive_reader = intensive_reader
self.rear_verifier = rear_verifier
# Set prep function for inference
self.sketch_prep_fn, self.intensive_prep_fn = prep_fn
@classmethod
def load(
cls,
train_examples=None,
sketch_train_dataset=None,
intensive_train_dataset=None,
eval_examples=None,
sketch_eval_dataset=None,
intensive_eval_dataset=None,
config_file: str = C.DEFAULT_CONFIG_FILE,
):
# Get arguments from yaml files
parser = HfArgumentParser([RetroArguments, TrainingArguments])
retro_args, training_args = parser.parse_yaml_file(yaml_file=config_file)
if training_args.run_name is not None and "," in training_args.run_name:
sketch_run_name, intensive_run_name = training_args.run_name.split(",")
else:
sketch_run_name, intensive_run_name = None, None
if training_args.metric_for_best_model is not None and "," in training_args.metric_for_best_model:
sketch_best_metric, intensive_best_metric = training_args.metric_for_best_model.split(",")
else:
sketch_best_metric, intensive_best_metric = None, None
sketch_training_args = copy.deepcopy(training_args)
intensive_training_args = training_args
sketch_tokenizer = AutoTokenizer.from_pretrained(
pretrained_model_name_or_path=retro_args.sketch_tokenizer_name,
use_auth_token=retro_args.use_auth_token,
revision=retro_args.sketch_revision,
)
# If `train_examples` is feeded, perform preprocessing
if train_examples is not None and sketch_train_dataset is None:
sketch_prep_fn, is_batched = get_sketch_features(sketch_tokenizer, "train", retro_args)
sketch_train_dataset = train_examples.map(
sketch_prep_fn,
batched=is_batched,
remove_columns=train_examples.column_names,
num_proc=retro_args.preprocessing_num_workers,
load_from_cache_file=not retro_args.overwrite_cache,
)
# If `eval_examples` is feeded, perform preprocessing
if eval_examples is not None and sketch_eval_dataset is None:
sketch_prep_fn, is_batched = get_sketch_features(sketch_tokenizer, "eval", retro_args)
sketch_eval_dataset = eval_examples.map(
sketch_prep_fn,
batched=is_batched,
remove_columns=eval_examples.column_names,
num_proc=retro_args.preprocessing_num_workers,
load_from_cache_file=not retro_args.overwrite_cache,
)
# Get preprocessing function for inference
sketch_prep_fn, _ = get_sketch_features(sketch_tokenizer, "test", retro_args)
# Get model for sketch reader
sketch_model_cls = retro_args.sketch_model_cls
sketch_model = sketch_model_cls.from_pretrained(
pretrained_model_name_or_path=retro_args.sketch_model_name,
use_auth_token=retro_args.use_auth_token,
revision=retro_args.sketch_revision,
)
# Get sketch reader
sketch_training_args.run_name = sketch_run_name
sketch_training_args.output_dir += "/sketch"
sketch_training_args.metric_for_best_model = sketch_best_metric
sketch_reader = SketchReader(
model=sketch_model,
args=sketch_training_args,
train_dataset=sketch_train_dataset,
eval_dataset=sketch_eval_dataset,
eval_examples=eval_examples,
data_args=retro_args,
tokenizer=sketch_tokenizer,
compute_metrics=compute_classification_metric,
)
intensive_tokenizer = AutoTokenizer.from_pretrained(
pretrained_model_name_or_path=retro_args.intensive_tokenizer_name,
use_auth_token=retro_args.use_auth_token,
revision=retro_args.intensive_revision,
)
# If `train_examples` is feeded, perform preprocessing
if train_examples is not None and intensive_train_dataset is None:
intensive_prep_fn, is_batched = get_intensive_features(intensive_tokenizer, "train", retro_args)
intensive_train_dataset = train_examples.map(
intensive_prep_fn,
batched=is_batched,
remove_columns=train_examples.column_names,
num_proc=retro_args.preprocessing_num_workers,
load_from_cache_file=not retro_args.overwrite_cache,
)
# If `eval_examples` is feeded, perform preprocessing
if eval_examples is not None and intensive_eval_dataset is None:
intensive_prep_fn, is_batched = get_intensive_features(intensive_tokenizer, "eval", retro_args)
intensive_eval_dataset = eval_examples.map(
intensive_prep_fn,
batched=is_batched,
remove_columns=eval_examples.column_names,
num_proc=retro_args.preprocessing_num_workers,
load_from_cache_file=not retro_args.overwrite_cache,
)
# Get preprocessing function for inference
intensive_prep_fn, _ = get_intensive_features(intensive_tokenizer, "test", retro_args)
# Get model for intensive reader
intensive_model_cls = retro_args.intensive_model_cls
intensive_model = intensive_model_cls.from_pretrained(
pretrained_model_name_or_path=retro_args.intensive_model_name,
use_auth_token=retro_args.use_auth_token,
revision=retro_args.intensive_revision,
)
# Get intensive reader
intensive_training_args.run_name = intensive_run_name
intensive_training_args.output_dir += "/intensive"
intensive_training_args.metric_for_best_model = intensive_best_metric
intensive_reader = IntensiveReader(
model=intensive_model,
args=intensive_training_args,
train_dataset=intensive_train_dataset,
eval_dataset=intensive_eval_dataset,
eval_examples=eval_examples,
data_args=retro_args,
tokenizer=intensive_tokenizer,
compute_metrics=compute_squad_v2,
)
# Get rear verifier
rear_verifier = RearVerifier(
beta1=retro_args.beta1,
beta2=retro_args.beta2,
best_cof=retro_args.best_cof,
thresh=retro_args.rear_threshold,
)
return cls(
args=retro_args,
sketch_reader=sketch_reader,
intensive_reader=intensive_reader,
rear_verifier=rear_verifier,
prep_fn=(sketch_prep_fn, intensive_prep_fn),
)
def __call__(
self,
query: str,
context: Union[str, List[str]],
return_submodule_outputs: bool = False,
) -> Tuple[Any]:
if isinstance(context, list):
context = " ".join(context)
predict_examples = datasets.Dataset.from_dict({
"example_id": ["0"],
C.ID_COLUMN_NAME: ["id-01"],
C.QUESTION_COLUMN_NAME: [query],
C.CONTEXT_COLUMN_NAME: [context]
})
return self.inference(predict_examples)
def train(self, module: str = "all"):
def wandb_finish(module):
for callback in module.callback_handler.callbacks:
if "wandb" in str(type(callback)).lower():
callback._wandb.finish()
callback._initialized = False
# Train sketch reader
if module.lower() in ["all", "sketch"]:
self.sketch_reader.train()
self.sketch_reader.save_model()
self.sketch_reader.save_state()
self.sketch_reader.free_memory()
wandb_finish(self.sketch_reader)
# Train intensive reader
if module.lower() in ["all", "intensive"]:
self.intensive_reader.train()
self.intensive_reader.save_model()
self.intensive_reader.save_state()
self.intensive_reader.free_memory()
wandb_finish(self.intensive_reader)
def inference(self, predict_examples: datasets.Dataset) -> Tuple[Any]:
if "example_id" not in predict_examples.column_names:
test_dataset = predict_examples.map(
lambda _, i: {"example_id": str(i)},
with_indices=True,
)
sketch_features = predict_examples.map(
self.sketch_prep_fn,
batched=True,
remove_columns=predict_examples.column_names,
)
intensive_features = predict_examples.map(
self.intensive_prep_fn,
batched=True,
remove_columns=predict_examples.column_names,
)
# self.sketch_reader.to(self.sketch_reader.args.device)
score_ext = self.sketch_reader.predict(sketch_features, predict_examples)
# self.sketch_reader.to("cpu")
# self.intensive_reader.to(self.intensive_reader.args.device)
nbest_preds, score_diff = self.intensive_reader.predict(
intensive_features, predict_examples, mode="retro_inference")
# self.intensive_reader.to("cpu")
predictions, scores = self.rear_verifier(score_ext, score_diff, nbest_preds)
outputs = (predictions, scores)
# if self.return_submodule_outputs:
# outputs += (score_ext, nbest_preds, score_diff)
return outputs
@property
def null_score_diff_threshold(self):
return self.args.null_score_diff_threshold
@null_score_diff_threshold.setter
def null_score_diff_threshold(self, val):
self.args.null_score_diff_threshold = val
@property
def n_best_size(self):
return self.args.n_best_size
@n_best_size.setter
def n_best_size(self, val):
self.args.n_best_size = val
@property
def beta1(self):
return self.rear_verifier.beta1
@beta1.setter
def beta1(self, val):
self.rear_verifier.beta1 = val
@property
def beta2(self):
return self.rear_verifier.beta2
@beta2.setter
def beta2(self, val):
self.rear_verifier.beta2 = val
@property
def best_cof(self):
return self.rear_verifier.best_cof
@best_cof.setter
def best_cof(self, val):
self.rear_verifier.best_cof = val
@property
def rear_threshold(self):
return self.rear_verifier.thresh
@rear_threshold.setter
def rear_threshold(self, val):
self.rear_verifier.thresh = val