Transformer-Explainability/BERT_rationale_benchmark/models/pipeline/pipeline_utils.py

import itertools
import logging
from collections import defaultdict, namedtuple
from itertools import chain
from typing import Any, Dict, List, Tuple

import numpy as np
import torch
import torch.nn as nn
from rationale_benchmark.metrics import (PositionScoredDocument, Rationale,
                                         partial_match_score,
                                         score_hard_rationale_predictions,
                                         score_soft_tokens)
from rationale_benchmark.models.model_utils import PaddedSequence
from rationale_benchmark.utils import Annotation
from sklearn.metrics import accuracy_score, classification_report

SentenceEvidence = namedtuple(
    "SentenceEvidence", "kls ann_id query docid index sentence"
)


def token_annotations_to_evidence_classification(
    annotations: List[Annotation],
    documents: Dict[str, List[List[Any]]],
    class_interner: Dict[str, int],
) -> List[SentenceEvidence]:
    ret = []
    for ann in annotations:
        docid_to_ev = defaultdict(list)
        for evidence in ann.all_evidences():
            docid_to_ev[evidence.docid].append(evidence)
        for docid, evidences in docid_to_ev.items():
            evidences = sorted(evidences, key=lambda ev: ev.start_token)
            text = []
            covered_tokens = set()
            doc = list(chain.from_iterable(documents[docid]))
            for evidence in evidences:
                assert (
                    evidence.start_token >= 0
                    and evidence.end_token > evidence.start_token
                )
                assert evidence.start_token < len(doc) and evidence.end_token <= len(
                    doc
                )
                text.extend(evidence.text)
                new_tokens = set(range(evidence.start_token, evidence.end_token))
                if len(new_tokens & covered_tokens) > 0:
                    raise ValueError(
                        "Have overlapping token ranges covered in the evidence spans and the implementer was lazy; deal with it"
                    )
                covered_tokens |= new_tokens
            assert len(text) > 0
            ret.append(
                SentenceEvidence(
                    kls=class_interner[ann.classification],
                    query=ann.query,
                    ann_id=ann.annotation_id,
                    docid=docid,
                    index=-1,
                    sentence=tuple(text),
                )
            )
    return ret


def annotations_to_evidence_classification(
    annotations: List[Annotation],
    documents: Dict[str, List[List[Any]]],
    class_interner: Dict[str, int],
    include_all: bool,
) -> List[SentenceEvidence]:
    """Converts Corpus-Level annotations to Sentence Level relevance judgments.

    As this module is about a pipelined approach for evidence identification,
    inputs to both an evidence identifier and evidence classifier need to be to
    be on a sentence level, this module converts data to be that form.

    The return type is of the form
        annotation id -> docid -> [sentence level annotations]
    """
    ret = []
    for ann in annotations:
        ann_id = ann.annotation_id
        docids = set(ev.docid for ev in chain.from_iterable(ann.evidences))
        annotations_for_doc = defaultdict(list)
        for d in docids:
            for index, sent in enumerate(documents[d]):
                annotations_for_doc[d].append(
                    SentenceEvidence(
                        kls=class_interner[ann.classification],
                        query=ann.query,
                        ann_id=ann.annotation_id,
                        docid=d,
                        index=index,
                        sentence=tuple(sent),
                    )
                )
        if include_all:
            ret.extend(chain.from_iterable(annotations_for_doc.values()))
        else:
            contributes = set()
            for ev in chain.from_iterable(ann.evidences):
                for index in range(ev.start_sentence, ev.end_sentence):
                    contributes.add(annotations_for_doc[ev.docid][index])
            ret.extend(contributes)
    assert len(ret) > 0
    return ret


def annotations_to_evidence_identification(
    annotations: List[Annotation], documents: Dict[str, List[List[Any]]]
) -> Dict[str, Dict[str, List[SentenceEvidence]]]:
    """Converts Corpus-Level annotations to Sentence Level relevance judgments.

    As this module is about a pipelined approach for evidence identification,
    inputs to both an evidence identifier and evidence classifier need to be to
    be on a sentence level, this module converts data to be that form.

    The return type is of the form
        annotation id -> docid -> [sentence level annotations]
    """
    ret = defaultdict(dict)  # annotation id -> docid -> sentences
    for ann in annotations:
        ann_id = ann.annotation_id
        for ev_group in ann.evidences:
            for ev in ev_group:
                if len(ev.text) == 0:
                    continue
                if ev.docid not in ret[ann_id]:
                    ret[ann.annotation_id][ev.docid] = []
                    # populate the document with "not evidence"; to be filled in later
                    for index, sent in enumerate(documents[ev.docid]):
                        ret[ann.annotation_id][ev.docid].append(
                            SentenceEvidence(
                                kls=0,
                                query=ann.query,
                                ann_id=ann.annotation_id,
                                docid=ev.docid,
                                index=index,
                                sentence=sent,
                            )
                        )
                # define the evidence sections of the document
                for s in range(ev.start_sentence, ev.end_sentence):
                    ret[ann.annotation_id][ev.docid][s] = SentenceEvidence(
                        kls=1,
                        ann_id=ann.annotation_id,
                        query=ann.query,
                        docid=ev.docid,
                        index=ret[ann.annotation_id][ev.docid][s].index,
                        sentence=ret[ann.annotation_id][ev.docid][s].sentence,
                    )
    return ret


def annotations_to_evidence_token_identification(
    annotations: List[Annotation],
    source_documents: Dict[str, List[List[str]]],
    interned_documents: Dict[str, List[List[int]]],
    token_mapping: Dict[str, List[List[Tuple[int, int]]]],
) -> Dict[str, Dict[str, List[SentenceEvidence]]]:
    # TODO document
    # TODO should we simplify to use only source text?
    ret = defaultdict(lambda: defaultdict(list))  # annotation id -> docid -> sentences
    positive_tokens = 0
    negative_tokens = 0
    for ann in annotations:
        annid = ann.annotation_id
        docids = set(ev.docid for ev in chain.from_iterable(ann.evidences))
        sentence_offsets = defaultdict(list)  # docid -> [(start, end)]
        classes = defaultdict(list)  # docid -> [token is yea or nay]
        for docid in docids:
            start = 0
            assert len(source_documents[docid]) == len(interned_documents[docid])
            for whole_token_sent, wordpiece_sent in zip(
                source_documents[docid], interned_documents[docid]
            ):
                classes[docid].extend([0 for _ in wordpiece_sent])
                end = start + len(wordpiece_sent)
                sentence_offsets[docid].append((start, end))
                start = end
        for ev in chain.from_iterable(ann.evidences):
            if len(ev.text) == 0:
                continue
            flat_token_map = list(chain.from_iterable(token_mapping[ev.docid]))
            if ev.start_token != -1:
                # start, end = token_mapping[ev.docid][ev.start_token][0], token_mapping[ev.docid][ev.end_token][1]
                start, end = (
                    flat_token_map[ev.start_token][0],
                    flat_token_map[ev.end_token - 1][1],
                )
            else:
                start = flat_token_map[sentence_offsets[ev.start_sentence][0]][0]
                end = flat_token_map[sentence_offsets[ev.end_sentence - 1][1]][1]
            for i in range(start, end):
                classes[ev.docid][i] = 1
        for docid, offsets in sentence_offsets.items():
            token_assignments = classes[docid]
            positive_tokens += sum(token_assignments)
            negative_tokens += len(token_assignments) - sum(token_assignments)
            for s, (start, end) in enumerate(offsets):
                sent = interned_documents[docid][s]
                ret[annid][docid].append(
                    SentenceEvidence(
                        kls=tuple(token_assignments[start:end]),
                        query=ann.query,
                        ann_id=ann.annotation_id,
                        docid=docid,
                        index=s,
                        sentence=sent,
                    )
                )
    logging.info(
        f"Have {positive_tokens} positive wordpiece tokens, {negative_tokens} negative wordpiece tokens"
    )
    return ret


def make_preds_batch(
    classifier: nn.Module,
    batch_elements: List[SentenceEvidence],
    device=None,
    criterion: nn.Module = None,
    tensorize_model_inputs: bool = True,
) -> Tuple[float, List[float], List[int], List[int]]:
    """Batch predictions

    Args:
        classifier: a module that looks like an AttentiveClassifier
        batch_elements: a list of elements to make predictions over. These must be SentenceEvidence objects.
        device: Optional; what compute device this should run on
        criterion: Optional; a loss function
        tensorize_model_inputs: should we convert our data to tensors before passing it to the model? Useful if we have a model that performs its own tokenization
    """
    # delete any "None" padding, if any (imposed by the use of the "grouper")
    batch_elements = filter(lambda x: x is not None, batch_elements)
    targets, queries, sentences = zip(
        *[(s.kls, s.query, s.sentence) for s in batch_elements]
    )
    ids = [(s.ann_id, s.docid, s.index) for s in batch_elements]
    targets = torch.tensor(targets, dtype=torch.long, device=device)
    if tensorize_model_inputs:
        queries = [torch.tensor(q, dtype=torch.long) for q in queries]
        sentences = [torch.tensor(s, dtype=torch.long) for s in sentences]
    preds = classifier(queries, ids, sentences)
    targets = targets.to(device=preds.device)
    if criterion:
        loss = criterion(preds, targets)
    else:
        loss = None
    # .float() because pytorch 1.3 introduces a bug where argmax is unsupported for float16
    hard_preds = torch.argmax(preds.float(), dim=-1)
    return loss, preds, hard_preds, targets


def make_preds_epoch(
    classifier: nn.Module,
    data: List[SentenceEvidence],
    batch_size: int,
    device=None,
    criterion: nn.Module = None,
    tensorize_model_inputs: bool = True,
):
    """Predictions for more than one batch.

    Args:
        classifier: a module that looks like an AttentiveClassifier
        data: a list of elements to make predictions over. These must be SentenceEvidence objects.
        batch_size: the biggest chunk we can fit in one batch.
        device: Optional; what compute device this should run on
        criterion: Optional; a loss function
        tensorize_model_inputs: should we convert our data to tensors before passing it to the model? Useful if we have a model that performs its own tokenization
    """
    epoch_loss = 0
    epoch_soft_pred = []
    epoch_hard_pred = []
    epoch_truth = []
    batches = _grouper(data, batch_size)
    classifier.eval()
    for batch in batches:
        loss, soft_preds, hard_preds, targets = make_preds_batch(
            classifier,
            batch,
            device,
            criterion=criterion,
            tensorize_model_inputs=tensorize_model_inputs,
        )
        if loss is not None:
            epoch_loss += loss.sum().item()
        epoch_hard_pred.extend(hard_preds)
        epoch_soft_pred.extend(soft_preds.cpu())
        epoch_truth.extend(targets)
    epoch_loss /= len(data)
    epoch_hard_pred = [x.item() for x in epoch_hard_pred]
    epoch_truth = [x.item() for x in epoch_truth]
    return epoch_loss, epoch_soft_pred, epoch_hard_pred, epoch_truth


def make_token_preds_batch(
    classifier: nn.Module,
    batch_elements: List[SentenceEvidence],
    token_mapping: Dict[str, List[List[Tuple[int, int]]]],
    device=None,
    criterion: nn.Module = None,
    tensorize_model_inputs: bool = True,
) -> Tuple[float, List[float], List[int], List[int]]:
    """Batch predictions

    Args:
        classifier: a module that looks like an AttentiveClassifier
        batch_elements: a list of elements to make predictions over. These must be SentenceEvidence objects.
        device: Optional; what compute device this should run on
        criterion: Optional; a loss function
        tensorize_model_inputs: should we convert our data to tensors before passing it to the model? Useful if we have a model that performs its own tokenization
    """
    # delete any "None" padding, if any (imposed by the use of the "grouper")
    batch_elements = filter(lambda x: x is not None, batch_elements)
    targets, queries, sentences = zip(
        *[(s.kls, s.query, s.sentence) for s in batch_elements]
    )
    ids = [(s.ann_id, s.docid, s.index) for s in batch_elements]
    targets = PaddedSequence.autopad(
        [torch.tensor(t, dtype=torch.long, device=device) for t in targets],
        batch_first=True,
        device=device,
    )
    aggregate_spans = [token_mapping[s.docid][s.index] for s in batch_elements]
    if tensorize_model_inputs:
        queries = [torch.tensor(q, dtype=torch.long) for q in queries]
        sentences = [torch.tensor(s, dtype=torch.long) for s in sentences]
    preds = classifier(queries, ids, sentences, aggregate_spans)
    targets = targets.to(device=preds.device)
    mask = targets.mask(on=1, off=0, device=preds.device, dtype=torch.float)
    if criterion:
        loss = criterion(
            preds, (targets.data.to(device=preds.device) * mask).squeeze()
        ).sum()
    else:
        loss = None
    hard_preds = [
        torch.round(x).to(dtype=torch.int).cpu() for x in targets.unpad(preds)
    ]
    targets = [[y.item() for y in x] for x in targets.unpad(targets.data.cpu())]
    return loss, preds, hard_preds, targets  # targets.unpad(targets.data.cpu())


# TODO fix the arguments
def make_token_preds_epoch(
    classifier: nn.Module,
    data: List[SentenceEvidence],
    token_mapping: Dict[str, List[List[Tuple[int, int]]]],
    batch_size: int,
    device=None,
    criterion: nn.Module = None,
    tensorize_model_inputs: bool = True,
):
    """Predictions for more than one batch.

    Args:
        classifier: a module that looks like an AttentiveClassifier
        data: a list of elements to make predictions over. These must be SentenceEvidence objects.
        batch_size: the biggest chunk we can fit in one batch.
        device: Optional; what compute device this should run on
        criterion: Optional; a loss function
        tensorize_model_inputs: should we convert our data to tensors before passing it to the model? Useful if we have a model that performs its own tokenization
    """
    epoch_loss = 0
    epoch_soft_pred = []
    epoch_hard_pred = []
    epoch_truth = []
    batches = _grouper(data, batch_size)
    classifier.eval()
    for batch in batches:
        loss, soft_preds, hard_preds, targets = make_token_preds_batch(
            classifier,
            batch,
            token_mapping,
            device,
            criterion=criterion,
            tensorize_model_inputs=tensorize_model_inputs,
        )
        if loss is not None:
            epoch_loss += loss.sum().item()
        epoch_hard_pred.extend(hard_preds)
        epoch_soft_pred.extend(soft_preds.cpu().tolist())
        epoch_truth.extend(targets)
    epoch_loss /= len(data)
    return epoch_loss, epoch_soft_pred, epoch_hard_pred, epoch_truth


# copied from https://docs.python.org/3/library/itertools.html#itertools-recipes
def _grouper(iterable, n, fillvalue=None):
    "Collect data into fixed-length chunks or blocks"
    # grouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx"
    args = [iter(iterable)] * n
    return itertools.zip_longest(*args, fillvalue=fillvalue)


def score_rationales(
    truth: List[Annotation],
    documents: Dict[str, List[List[int]]],
    input_data: List[SentenceEvidence],
    scores: List[float],
) -> dict:
    results = {}
    doc_to_sent_scores = dict()  # (annid, docid) -> [sentence scores]
    for sent, score in zip(input_data, scores):
        k = (sent.ann_id, sent.docid)
        if k not in doc_to_sent_scores:
            doc_to_sent_scores[k] = [0.0 for _ in range(len(documents[sent.docid]))]
        if not isinstance(score[1], float):
            score[1] = score[1].item()
        doc_to_sent_scores[(sent.ann_id, sent.docid)][sent.index] = score[1]
    # hard rationale scoring
    best_sentence = {k: np.argmax(np.array(v)) for k, v in doc_to_sent_scores.items()}
    predicted_rationales = []
    for (ann_id, docid), sent_idx in best_sentence.items():
        start_token = sum(len(s) for s in documents[docid][:sent_idx])
        end_token = start_token + len(documents[docid][sent_idx])
        predicted_rationales.append(Rationale(ann_id, docid, start_token, end_token))
    true_rationales = list(
        chain.from_iterable(Rationale.from_annotation(rat) for rat in truth)
    )

    results["hard_rationale_scores"] = score_hard_rationale_predictions(
        true_rationales, predicted_rationales
    )
    results["hard_rationale_partial_match_scores"] = partial_match_score(
        true_rationales, predicted_rationales, [0.5]
    )

    # soft rationale scoring
    instance_format = []
    for (ann_id, docid), sentences in doc_to_sent_scores.items():
        soft_token_predictions = []
        for sent_score, sent_text in zip(sentences, documents[docid]):
            soft_token_predictions.extend(sent_score for _ in range(len(sent_text)))
        instance_format.append(
            {
                "annotation_id": ann_id,
                "rationales": [
                    {
                        "docid": docid,
                        "soft_rationale_predictions": soft_token_predictions,
                        "soft_sentence_predictions": sentences,
                    }
                ],
            }
        )
    flattened_documents = {
        k: list(chain.from_iterable(v)) for k, v in documents.items()
    }
    token_scoring_format = PositionScoredDocument.from_results(
        instance_format, truth, flattened_documents, use_tokens=True
    )
    results["soft_token_scores"] = score_soft_tokens(token_scoring_format)
    sentence_scoring_format = PositionScoredDocument.from_results(
        instance_format, truth, documents, use_tokens=False
    )
    results["soft_sentence_scores"] = score_soft_tokens(sentence_scoring_format)
    return results


def decode(
    evidence_identifier: nn.Module,
    evidence_classifier: nn.Module,
    train: List[Annotation],
    val: List[Annotation],
    test: List[Annotation],
    docs: Dict[str, List[List[int]]],
    class_interner: Dict[str, int],
    batch_size: int,
    tensorize_model_inputs: bool,
    decoding_docs: Dict[str, List[Any]] = None,
) -> dict:
    """Identifies and then classifies evidence

    Args:
        evidence_identifier: a module for identifying evidence statements
        evidence_classifier: a module for making a classification based on evidence statements
        train: A List of interned Annotations
        val: A List of interned Annotations
        test: A List of interned Annotations
        docs: A Dict of Documents, which are interned sentences.
        class_interner: Converts an Annotation's final class into ints
        batch_size: how big should our batches be?
        tensorize_model_inputs: should we convert our data to tensors before passing it to the model? Useful if we have a model that performs its own tokenization
    """
    device = None
    class_labels = [k for k, v in sorted(class_interner.items(), key=lambda x: x[1])]
    if decoding_docs is None:
        decoding_docs = docs

    def prep(data: List[Annotation]) -> List[Tuple[SentenceEvidence, SentenceEvidence]]:
        """Prepares data for evidence identification and classification.

        Creates paired evaluation data, wherein each (annotation, docid, sentence, kls)
        tuplet appears first as the kls determining if the sentence is evidence, and
        secondarily what the overall classification for the (annotation/docid) pair is.
        This allows selection based on model scores of the evidence_identifier for
        input to the evidence_classifier.
        """
        identification_data = annotations_to_evidence_identification(data, docs)
        classification_data = annotations_to_evidence_classification(
            data, docs, class_interner, include_all=True
        )
        ann_doc_sents = defaultdict(
            lambda: defaultdict(dict)
        )  # ann id -> docid -> sent idx -> sent data
        ret = []
        for sent_ev in classification_data:
            id_data = identification_data[sent_ev.ann_id][sent_ev.docid][sent_ev.index]
            ret.append((id_data, sent_ev))
            assert id_data.ann_id == sent_ev.ann_id
            assert id_data.docid == sent_ev.docid
            assert id_data.index == sent_ev.index
        assert len(ret) == len(classification_data)
        return ret

    def decode_batch(
        data: List[Tuple[SentenceEvidence, SentenceEvidence]],
        name: str,
        score: bool = False,
        annotations: List[Annotation] = None,
    ) -> dict:
        """Identifies evidence statements and then makes classifications based on it.

        Args:
            data: a paired list of SentenceEvidences, differing only in the kls field.
                  The first corresponds to whether or not something is evidence, and the second corresponds to an evidence class
            name: a name for a results dict
        """

        num_uniques = len(set((x.ann_id, x.docid) for x, _ in data))
        logging.info(
            f"Decoding dataset {name} with {len(data)} sentences, {num_uniques} annotations"
        )
        identifier_data, classifier_data = zip(*data)
        results = dict()
        IdentificationClassificationResult = namedtuple(
            "IdentificationClassificationResult",
            "identification_data classification_data soft_identification hard_identification soft_classification hard_classification",
        )
        with torch.no_grad():
            # make predictions for the evidence_identifier
            evidence_identifier.eval()
            evidence_classifier.eval()

            (
                _,
                soft_identification_preds,
                hard_identification_preds,
                _,
            ) = make_preds_epoch(
                evidence_identifier,
                identifier_data,
                batch_size,
                device,
                tensorize_model_inputs=tensorize_model_inputs,
            )
            assert len(soft_identification_preds) == len(data)
            identification_results = defaultdict(list)
            for id_data, cls_data, soft_id_pred, hard_id_pred in zip(
                identifier_data,
                classifier_data,
                soft_identification_preds,
                hard_identification_preds,
            ):
                res = IdentificationClassificationResult(
                    identification_data=id_data,
                    classification_data=cls_data,
                    # 1 is p(evidence|sent,query)
                    soft_identification=soft_id_pred[1].float().item(),
                    hard_identification=hard_id_pred,
                    soft_classification=None,
                    hard_classification=False,
                )
                identification_results[(id_data.ann_id, id_data.docid)].append(res)

            best_identification_results = {
                key: max(value, key=lambda x: x.soft_identification)
                for key, value in identification_results.items()
            }
            logging.info(
                f"Selected the best sentence for {len(identification_results)} examples from a total of {len(soft_identification_preds)} sentences"
            )
            ids, classification_data = zip(
                *[
                    (k, v.classification_data)
                    for k, v in best_identification_results.items()
                ]
            )
            (
                _,
                soft_classification_preds,
                hard_classification_preds,
                classification_truth,
            ) = make_preds_epoch(
                evidence_classifier,
                classification_data,
                batch_size,
                device,
                tensorize_model_inputs=tensorize_model_inputs,
            )
            classification_results = dict()
            for eyeD, soft_class, hard_class in zip(
                ids, soft_classification_preds, hard_classification_preds
            ):
                input_id_result = best_identification_results[eyeD]
                res = IdentificationClassificationResult(
                    identification_data=input_id_result.identification_data,
                    classification_data=input_id_result.classification_data,
                    soft_identification=input_id_result.soft_identification,
                    hard_identification=input_id_result.hard_identification,
                    soft_classification=soft_class,
                    hard_classification=hard_class,
                )
                classification_results[eyeD] = res

            if score:
                truth = []
                pred = []
                for res in classification_results.values():
                    truth.append(res.classification_data.kls)
                    pred.append(res.hard_classification)
                # results[f'{name}_f1'] = classification_report(classification_truth, pred, target_names=class_labels, output_dict=True)
                results[f"{name}_f1"] = classification_report(
                    classification_truth,
                    hard_classification_preds,
                    target_names=class_labels,
                    output_dict=True,
                )
                results[f"{name}_acc"] = accuracy_score(
                    classification_truth, hard_classification_preds
                )
                results[f"{name}_rationale"] = score_rationales(
                    annotations,
                    decoding_docs,
                    identifier_data,
                    soft_identification_preds,
                )

            # turn the above results into a format suitable for scoring via the rationale scorer
            # n.b. the sentence-level evidence predictions (hard and soft) are
            # broadcast to the token level for scoring. The comprehensiveness class
            # score is also a lie since the pipeline model above is faithful by
            # design.
            decoded = dict()
            decoded_scores = defaultdict(list)
            for (ann_id, docid), pred in classification_results.items():
                sentence_prediction_scores = [
                    x.soft_identification
                    for x in identification_results[(ann_id, docid)]
                ]
                sentence_start_token = sum(
                    len(s)
                    for s in decoding_docs[docid][: pred.identification_data.index]
                )
                sentence_end_token = sentence_start_token + len(
                    decoding_docs[docid][pred.classification_data.index]
                )
                hard_rationale_predictions = [
                    {
                        "start_token": sentence_start_token,
                        "end_token": sentence_end_token,
                    }
                ]
                soft_rationale_predictions = []
                for sent_result in sorted(
                    identification_results[(ann_id, docid)],
                    key=lambda x: x.identification_data.index,
                ):
                    soft_rationale_predictions.extend(
                        sent_result.soft_identification
                        for _ in range(
                            len(
                                decoding_docs[sent_result.identification_data.docid][
                                    sent_result.identification_data.index
                                ]
                            )
                        )
                    )
                if ann_id not in decoded:
                    decoded[ann_id] = {
                        "annotation_id": ann_id,
                        "rationales": [],
                        "classification": class_labels[pred.hard_classification],
                        "classification_scores": {
                            class_labels[i]: s.item()
                            for i, s in enumerate(pred.soft_classification)
                        },
                        # TODO this should turn into the data distribution for the predicted class
                        # "comprehensiveness_classification_scores": 0.0,
                        "truth": pred.classification_data.kls,
                    }
                decoded[ann_id]["rationales"].append(
                    {
                        "docid": docid,
                        "hard_rationale_predictions": hard_rationale_predictions,
                        "soft_rationale_predictions": soft_rationale_predictions,
                        "soft_sentence_predictions": sentence_prediction_scores,
                    }
                )
                decoded_scores[ann_id].append(pred.soft_classification)

            # in practice, this is always a single element operation:
            # in evidence inference (prompt is really a prompt + document), fever (we split documents into two classifications), movies (you only have one opinion about a movie), or boolQ (single document prompts)
            # this exists to support weird models we *might* implement for cose/esnli
            for ann_id, scores_list in decoded_scores.items():
                scores = torch.stack(scores_list)
                score_avg = torch.mean(scores, dim=0)
                # .float() because pytorch 1.3 introduces a bug where argmax is unsupported for float16
                hard_pred = torch.argmax(score_avg.float()).item()
                decoded[ann_id]["classification"] = class_labels[hard_pred]
                decoded[ann_id]["classification_scores"] = {
                    class_labels[i]: s.item() for i, s in enumerate(score_avg)
                }
            return results, list(decoded.values())

    test_results, test_decoded = decode_batch(prep(test), "test", score=False)
    val_results, val_decoded = dict(), []
    train_results, train_decoded = dict(), []
    # val_results, val_decoded = decode_batch(prep(val), 'val', score=True, annotations=val)
    # train_results, train_decoded = decode_batch(prep(train), 'train', score=True, annotations=train)
    return (
        dict(**train_results, **val_results, **test_results),
        train_decoded,
        val_decoded,
        test_decoded,
    )


def decode_evidence_tokens_and_classify(
    evidence_token_identifier: nn.Module,
    evidence_classifier: nn.Module,
    train: List[Annotation],
    val: List[Annotation],
    test: List[Annotation],
    docs: Dict[str, List[List[int]]],
    source_documents: Dict[str, List[List[str]]],
    token_mapping: Dict[str, List[List[Tuple[int, int]]]],
    class_interner: Dict[str, int],
    batch_size: int,
    decoding_docs: Dict[str, List[Any]],
    use_cose_hack: bool = False,
) -> dict:
    """Identifies and then classifies evidence

    Args:
        evidence_token_identifier: a module for identifying evidence statements
        evidence_classifier: a module for making a classification based on evidence statements
        train: A List of interned Annotations
        val: A List of interned Annotations
        test: A List of interned Annotations
        docs: A Dict of Documents, which are interned sentences.
        class_interner: Converts an Annotation's final class into ints
        batch_size: how big should our batches be?
    """
    device = None
    class_labels = [k for k, v in sorted(class_interner.items(), key=lambda x: x[1])]
    if decoding_docs is None:
        decoding_docs = docs

    def prep(data: List[Annotation]) -> List[Tuple[SentenceEvidence, SentenceEvidence]]:
        """Prepares data for evidence identification and classification.

        Creates paired evaluation data, wherein each (annotation, docid, sentence, kls)
        tuplet appears first as the kls determining if the sentence is evidence, and
        secondarily what the overall classification for the (annotation/docid) pair is.
        This allows selection based on model scores of the evidence_token_identifier for
        input to the evidence_classifier.
        """
        # identification_data = annotations_to_evidence_identification(data, docs)
        classification_data = token_annotations_to_evidence_classification(
            data, docs, class_interner
        )
        # annotation id -> docid -> [SentenceEvidence])
        identification_data = annotations_to_evidence_token_identification(
            data,
            source_documents=decoding_docs,
            interned_documents=docs,
            token_mapping=token_mapping,
        )
        ann_doc_sents = defaultdict(
            lambda: defaultdict(dict)
        )  # ann id -> docid -> sent idx -> sent data
        ret = []
        for sent_ev in classification_data:
            id_data = identification_data[sent_ev.ann_id][sent_ev.docid][sent_ev.index]
            ret.append((id_data, sent_ev))
            assert id_data.ann_id == sent_ev.ann_id
            assert id_data.docid == sent_ev.docid
            # assert id_data.index == sent_ev.index
        assert len(ret) == len(classification_data)
        return ret

    def decode_batch(
        data: List[Tuple[SentenceEvidence, SentenceEvidence]],
        name: str,
        score: bool = False,
        annotations: List[Annotation] = None,
        class_labels: dict = class_labels,
    ) -> dict:
        """Identifies evidence statements and then makes classifications based on it.

        Args:
            data: a paired list of SentenceEvidences, differing only in the kls field.
                  The first corresponds to whether or not something is evidence, and the second corresponds to an evidence class
            name: a name for a results dict
        """

        num_uniques = len(set((x.ann_id, x.docid) for x, _ in data))
        logging.info(
            f"Decoding dataset {name} with {len(data)} sentences, {num_uniques} annotations"
        )
        identifier_data, classifier_data = zip(*data)
        results = dict()
        with torch.no_grad():
            # make predictions for the evidence_token_identifier
            evidence_token_identifier.eval()
            evidence_classifier.eval()

            (
                _,
                soft_identification_preds,
                hard_identification_preds,
                id_preds_truth,
            ) = make_token_preds_epoch(
                evidence_token_identifier,
                identifier_data,
                token_mapping,
                batch_size,
                device,
                tensorize_model_inputs=True,
            )
            assert len(soft_identification_preds) == len(data)
            evidence_only_cls = []
            for id_data, cls_data, soft_id_pred, hard_id_pred in zip(
                identifier_data,
                classifier_data,
                soft_identification_preds,
                hard_identification_preds,
            ):
                assert cls_data.ann_id == id_data.ann_id
                sent = []
                for start, end in token_mapping[cls_data.docid][0]:
                    if bool(hard_id_pred[start]):
                        sent.extend(id_data.sentence[start:end])
                # assert len(sent) > 0
                new_cls_data = SentenceEvidence(
                    cls_data.kls,
                    cls_data.ann_id,
                    cls_data.query,
                    cls_data.docid,
                    cls_data.index,
                    tuple(sent),
                )
                evidence_only_cls.append(new_cls_data)
            (
                _,
                soft_classification_preds,
                hard_classification_preds,
                classification_truth,
            ) = make_preds_epoch(
                evidence_classifier,
                evidence_only_cls,
                batch_size,
                device,
                tensorize_model_inputs=True,
            )

            if use_cose_hack:
                logging.info(
                    "Reformatting identification and classification results to fit COS-E"
                )
                grouping = 5
                new_soft_identification_preds = []
                new_hard_identification_preds = []
                new_id_preds_truth = []
                new_soft_classification_preds = []
                new_hard_classification_preds = []
                new_classification_truth = []
                new_identifier_data = []
                class_labels = []

                # TODO fix the labels for COS-E
                for i in range(0, len(soft_identification_preds), grouping):
                    cls_scores = torch.stack(
                        soft_classification_preds[i : i + grouping]
                    )
                    cls_scores = nn.functional.softmax(cls_scores, dim=-1)
                    cls_scores = cls_scores[:, 1]
                    choice = torch.argmax(cls_scores)
                    cls_labels = [
                        x.ann_id.split("_")[-1]
                        for x in evidence_only_cls[i : i + grouping]
                    ]
                    class_labels = cls_labels  # we need to update the class labels because of the terrible hackery used to train this
                    cls_truths = [x.kls for x in evidence_only_cls[i : i + grouping]]
                    # cls_choice = evidence_only_cls[i + choice].ann_id.split('_')[-1]
                    cls_truth = np.argmax(cls_truths)
                    new_soft_identification_preds.append(
                        soft_identification_preds[i + choice]
                    )
                    new_hard_identification_preds.append(
                        hard_identification_preds[i + choice]
                    )
                    new_id_preds_truth.append(id_preds_truth[i + choice])
                    new_soft_classification_preds.append(
                        soft_classification_preds[i + choice]
                    )
                    new_hard_classification_preds.append(choice)
                    new_identifier_data.append(identifier_data[i + choice])
                    # new_hard_classification_preds.append(hard_classification_preds[i + choice])
                    # new_classification_truth.append(classification_truth[i + choice])
                    new_classification_truth.append(cls_truth)

                soft_identification_preds = new_soft_identification_preds
                hard_identification_preds = new_hard_identification_preds
                id_preds_truth = new_id_preds_truth
                soft_classification_preds = new_soft_classification_preds
                hard_classification_preds = new_hard_classification_preds
                classification_truth = new_classification_truth
                identifier_data = new_identifier_data
            if score:
                results[f"{name}_f1"] = classification_report(
                    classification_truth,
                    hard_classification_preds,
                    target_names=class_labels,
                    output_dict=True,
                )
                results[f"{name}_acc"] = accuracy_score(
                    classification_truth, hard_classification_preds
                )
                results[f"{name}_token_pred_acc"] = accuracy_score(
                    list(chain.from_iterable(id_preds_truth)),
                    list(chain.from_iterable(hard_identification_preds)),
                )
                results[f"{name}_token_pred_f1"] = classification_report(
                    list(chain.from_iterable(id_preds_truth)),
                    list(chain.from_iterable(hard_identification_preds)),
                    output_dict=True,
                )
                # TODO for token level stuff!
                soft_id_scores = [
                    [1 - x, x] for x in chain.from_iterable(soft_identification_preds)
                ]
                results[f"{name}_rationale"] = score_rationales(
                    annotations, decoding_docs, identifier_data, soft_id_scores
                )
                logging.info(f"Results: {results}")

            # turn the above results into a format suitable for scoring via the rationale scorer
            # n.b. the sentence-level evidence predictions (hard and soft) are
            # broadcast to the token level for scoring. The comprehensiveness class
            # score is also a lie since the pipeline model above is faithful by
            # design.
            decoded = dict()
            scores = []
            assert len(identifier_data) == len(soft_identification_preds)
            for (
                id_data,
                soft_id_pred,
                hard_id_pred,
                soft_cls_preds,
                hard_cls_pred,
            ) in zip(
                identifier_data,
                soft_identification_preds,
                hard_identification_preds,
                soft_classification_preds,
                hard_classification_preds,
            ):
                docid = id_data.docid
                if use_cose_hack:
                    docid = "_".join(docid.split("_")[0:-1])
                assert len(docid) > 0
                rationales = {
                    "docid": docid,
                    "hard_rationale_predictions": [],
                    # token level classifications, a value must be provided per-token
                    # in an ideal world, these correspond to the hard-decoding above.
                    "soft_rationale_predictions": [],
                    # sentence level classifications, a value must be provided for every
                    # sentence in each document, or not at all
                    "soft_sentence_predictions": [1.0],
                }
                last = -1
                start_span = -1
                for pos, (start, _) in enumerate(token_mapping[id_data.docid][0]):
                    rationales["soft_rationale_predictions"].append(soft_id_pred[start])
                    if bool(hard_id_pred[start]):
                        if start_span == -1:
                            start_span = pos
                        last = pos
                    else:
                        if start_span != -1:
                            rationales["hard_rationale_predictions"].append(
                                {
                                    "start_token": start_span,
                                    "end_token": last + 1,
                                }
                            )
                        last = -1
                        start_span = -1
                if start_span != -1:
                    rationales["hard_rationale_predictions"].append(
                        {
                            "start_token": start_span,
                            "end_token": last + 1,
                        }
                    )

                ann_id = id_data.ann_id
                if use_cose_hack:
                    ann_id = "_".join(ann_id.split("_")[0:-1])
                soft_cls_preds = nn.functional.softmax(soft_cls_preds)
                decoded[id_data.ann_id] = {
                    "annotation_id": ann_id,
                    "rationales": [rationales],
                    "classification": class_labels[hard_cls_pred],
                    "classification_scores": {
                        class_labels[i]: score.item()
                        for i, score in enumerate(soft_cls_preds)
                    },
                }
            return results, list(decoded.values())

    # test_results, test_decoded = dict(), []
    # val_results, val_decoded = dict(), []
    train_results, train_decoded = dict(), []
    val_results, val_decoded = decode_batch(
        prep(val), "val", score=True, annotations=val, class_labels=class_labels
    )
    test_results, test_decoded = decode_batch(
        prep(test), "test", score=False, class_labels=class_labels
    )
    # train_results, train_decoded = decode_batch(prep(train), 'train', score=True, annotations=train, class_labels=class_labels)
    return (
        dict(**train_results, **val_results, **test_results),
        train_decoded,
        val_decoded,
        test_decoded,
    )