retro_reader/retro_reader.py

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