Enhancing small language models with LLM generated counterfactuals.
import ast
import numpy as np
import pandas as pd
from collections import OrderedDict
from tqdm import tqdm
import math
import re
import spacy
from spacy.tokens import Doc
import torch

# from supar import Parser

from src.calibration.baseline import common, evaluate, utils

import warnings

# ignore warning
warnings.filterwarnings("ignore")


class CreateFeatures:
    def __init__(self, path, model, dataset, method):
        self.path = path
        self.cls_tokens = ["[CLS]", "<s>"]
        self.sep_tokens = ["[SEP]", "</s>"]
        self.count = 0
        self.nlp = spacy.load("en_core_web_sm")
        self.model = model
        self.dataset = dataset.split("_")[0]
        self.method = method

    def _load_data(self):
        self.tagger_info = utils.load_bin(f"{self.path}/pos_info.bin")
        # self.ent_info = utils.load_bin(f"{self.path}/ent_info.bin")
        self.states = utils.load_bin(
            f"{self.path}/dense_repr_pca_10_info_{self.model}.bin"
        )
        self.attributions = utils.load_bin(
            f"{self.path}/{self.method}_info_{self.model}.bin"
        )
        pred_df = pd.read_csv(f"{self.path}/outputs_{self.model}.csv")
        self.preds_info = pred_df.set_index("id").T.to_dict("dict")
        # self.preds_info = utils.read_json('outputs/addsent-dev_squad_predictions.json')

    def featurize(self):
        self._load_data()
        processed_instances = OrderedDict()
        c = 0
        for q_id in tqdm(
            self.tagger_info, total=len(self.tagger_info), desc="transforming..."
        ):
            tags = self.tagger_info[q_id]
            # entities = self.ent_info[q_id]
            attributions = self.attributions[q_id]
            states = self.states[q_id]
            preds = self.preds_info[q_id]
            processed_instances[q_id] = self.extract_feature_for_instance(
                attributions, tags, preds, states
            )
            c += 1
        utils.dump_to_bin(
            processed_instances,
            f"{self.path}/calib_data_{self.method}_{self.model}.bin",
        )

    def lemmatize_pos_tag(self, token_tags):
        """
        limit the amount of POS tags to reduce the number of features
        :param token_tags: tagged tokens
        :return:
        """
        tok, pos, tag = token_tags
        if tag == "NNS":
            tag = "NN"
        if tag == "NNPS":
            tag = "NNP"
        if tag.startswith("JJ"):
            tag = "JJ"
        if tag.startswith("RB"):
            tag = "RB"
        if tag.startswith("W"):
            tag = "W"
        if tag.startswith("PRP"):
            tag = "PRP"
        if tag.startswith("VB"):
            tag = "VB"
        if pos == "PUNCT":
            tag = "PUNCT"
        return tok, pos, tag

    def extract_baseline_feature(self, preds):
        # print(preds)
        predictions = ast.literal_eval(preds["pred_text"])[0]
        # print(predictions)
        feat = common.IndexedFeature()
        # base features
        for rank, p in enumerate(predictions[:5]):
            feat.add(f"BASELINE_PROB_{rank}", p["score"])
        feat.add("BASELINE_CONTEXT_LENGTH", len(preds["context"]))
        feat.add("BASELINE_PRED_ANS_LENGTH", len(predictions[0]["answer"]))

        top_pred = predictions[0]["answer"]
        first_distinct_prob = 0
        for i, p in enumerate(predictions[1:]):
            overlapping = evaluate.f1_score(p["answer"], top_pred)
            if overlapping > 0:
                continue
            first_distinct_prob = p["score"]
            # print(i+1, top_pred, p['answer'], first_distinct_prob)
        feat.add("FIRST_DISTINCT_PROB", first_distinct_prob)
        return feat

    def source_of_token(self, idx, tok, context_start, ans_range):
        if tok in self.cls_tokens or tok in self.sep_tokens:
            return "S"
        if idx >= 1 and idx < context_start:
            return "Q"
        if idx >= ans_range[0] and idx <= ans_range[1]:
            # print("Answer token:", tok)
            return "A"
        return "C"

    def nan_check(self, result):
        for i, (key, value) in enumerate(result):
            if math.isnan(float(value)):
                result[i][1] = 0
        return result

    def count_syllables(self, word: str):
        return len(
            re.findall("(?!e$)[aeiouy]+", word, re.I)
            + re.findall("^[^aeiouy]*e$", word, re.I)
        )

    def handle_zero_division(self, n, d):
        return n / d if d else 0

    def extract_state_features(self, attr, tokens, ans_range, states):
        """
        Extract Q, C, and A features from model hidden states
        """
        feat = common.IndexedFeature()
        context_start = tokens.index(self.sep_tokens[1])
        start, end = ans_range
        ans_indices = list(range(start, end + 1))
        q_states = states[:, 1:context_start, :]
        c_states = states[:, context_start + 2 : -1, :]
        torch.set_printoptions(threshold=10_000)
        # remove answer states from context states
        c_states_before = c_states[:, : start - (context_start + 2), :]
        c_states_after = c_states[:, end + 1 - (context_start + 2) :, :]
        final_c_states = torch.cat((c_states_before, c_states_after), dim=1)
        a_states = states[:, start : end + 1, :]
        attribution = np.array(attr["attributions"])
        question_attr = attribution[1:context_start]
        context_attr = attribution[context_start + 2 : -1]
        c_attr_before = context_attr[: start - (context_start + 2)]
        c_attr_after = context_attr[end + 1 - (context_start + 2) :]
        final_c_attr = np.concatenate((c_attr_before, c_attr_after), axis=0)
        ans_attr = attribution[start : end + 1]

        ques_sorted_indices = np.argsort(question_attr)
        # take topk indices
        topk_ques_indices = ques_sorted_indices[
            -int(ques_sorted_indices.shape[0] * 0.10) :
        ]
        ques_topk_states = np.take(q_states, topk_ques_indices, axis=1)

        topc_percent = 0.10
        cxt_sorted_indices = np.argsort(final_c_attr)
        cxt_top_k_indices = cxt_sorted_indices[
            -int(cxt_sorted_indices.shape[0] * topc_percent) :
        ]
        cxt_topk_states = np.take(final_c_states, cxt_top_k_indices, axis=1)

        topa_percent = 0.20
        ans_sorted_indices = np.argsort(ans_attr)
        ans_top_k_indices = ans_sorted_indices[
            -int(ans_sorted_indices.shape[0] * topa_percent) :
        ]
        ans_topk_states = np.take(a_states, ans_top_k_indices, axis=1)

        q_rep = torch.mean(ques_topk_states, dim=1).tolist()[0]
        # print(q_rep.shape)
        c_rep = torch.mean(cxt_topk_states, dim=1).tolist()[0]
        # print(c_rep.shape)
        a_rep = torch.mean(ans_topk_states, dim=1).tolist()[0]

        # for i, value in enumerate(q_rep):
        #     if math.isnan(value):
        #         continue
        #     feat.add_new(f"REPR_TOPK_Q_{i}", value)
        for i, value in enumerate(c_rep):
            if math.isnan(value):
                continue
            feat.add_new(f"REPR_TOPK_C_{i}", value)
            # print(value)
        for i, value in enumerate(a_rep):
            if math.isnan(value):
                continue
            feat.add_new(f"REPR_TOPK_A_{i}", value)

        return feat

    def extract_bow_feature(self, words, tags, ans_range):
        feat = common.IndexedFeature()
        context_start = words.index(self.sep_tokens[1])
        for i, (i_token, i_pos, i_tag) in enumerate(tags):
            i_src = self.source_of_token(i, i_token, context_start, ans_range)
            if i_src == "Q" or i_src == "A" or i_src == "C":
                # print('BOW_{}_{}'.format(i_src, i_tag))
                feat.add("BOW_{}_{}".format(i_src, i_tag))
                feat.add("BOW_IN_{}".format(i_tag))
        return feat

    def merge_attribution_by_segments(self, attribution, segments):
        new_val = []
        for a, b in segments:
            new_val.append(np.sum(attribution[a:b], axis=0))
        importance = np.stack(new_val, axis=0)
        return importance

    def aggregate_token_attribution(self, attr, tags, polarity):
        # attr = ast.literal_eval(attr)
        # print(attr)
        attribution_val = np.array(attr["attributions"])
        if polarity == "POS":
            attribution_val[attribution_val < 0] = 0
        elif polarity == "NEG":
            attribution_val[attribution_val > 0] = 0
        elif polarity == "NEU":
            pass
        else:
            raise RuntimeError("Invalid polarity")

        attribution_val = self.merge_attribution_by_segments(
            attribution_val, tags["segments"]
        )
        assert attribution_val.shape[0] == len(tags["segments"])
        # normalize
        attribution_val = attribution_val / np.sum(attribution_val)
        # print(attribution_val)
        return attribution_val

    def normalize_token_attr(self, feat, attributions, norm_method=None):
        if norm_method is None:
            return feat
        if norm_method == "all":
            sum_v = np.sum(attributions)
            for k in feat.data:
                feat.data[k] = feat.data[k] / sum_v if sum_v != 0 else 0
            return feat
        if norm_method == "counted":
            sum_v = sum(feat.data.values())
            for k in feat.data:
                feat.data[k] = feat.data[k] / sum_v if sum_v != 0 else 0
            return feat
        raise RuntimeError(norm_method)

    def extract_token_attr_feature_in_question(
        self, words, tags, attributions, include_punct=False
    ):
        context_start = words.index(self.sep_tokens[1])
        tags = tags[1:context_start]
        attributions = attributions[1:context_start]

        feat = common.IndexedFeature()
        unnorm = common.IndexedFeature()
        sum_v = 0
        for i, (token, pos, tag) in enumerate(tags):
            if tag:
                v = attributions[i]
                if pos == "PUNCT" and not include_punct:
                    continue
                feat.add("NORMED_TOK_Q_" + tag, v)
                unnorm.add("UNNORM_TOK_Q_" + tag, v)
                sum_v += v

        feat = self.normalize_token_attr(feat, attributions)
        feat.add_set(unnorm)
        feat.add("SUM_TOK_Q", sum_v)
        return feat

    def extract_token_attr_feature_in_context(
        self, words, tags, attributions, include_punct=False
    ):
        context_start = words.index(self.sep_tokens[1])
        tags = tags[context_start + 2 : -1]
        attributions = attributions[context_start + 2 : -1]

        feat = common.IndexedFeature()
        unnorm = common.IndexedFeature()
        sum_v = 0
        for i, (token, pos, tag) in enumerate(tags):
            if tag:
                v = attributions[i]
                if pos == "PUNCT" and not include_punct:
                    continue
                feat.add("NORMED_TOK_C_" + tag, v)
                feat.add("UNNORM_TOK_C_" + tag, v)
                sum_v += v

        feat = self.normalize_token_attr(feat, attributions)
        feat.add_set(unnorm)
        feat.add("SUM_TOK_C", sum_v)
        return feat

    def extract_token_attr_feature_in_input(
        self, words, tags, attributions, ans_range, include_punct=False
    ):
        feat = common.IndexedFeature()
        context_start = words.index(self.sep_tokens[1])
        normed_a_feat = common.IndexedFeature()
        unnormed_a_feat = common.IndexedFeature()
        for i, (token, pos, tag) in enumerate(tags):
            if tag:
                v = attributions[i]
                if pos == "PUNCT" and not include_punct:
                    continue
                feat.add("TOK_IN_" + tag, v)

                if i >= ans_range[0] and i <= ans_range[1]:
                    unnormed_a_feat.add("UNNORM_TOK_A_" + tag, v)
                    normed_a_feat.add("NORMED_TOK_A_" + tag, v)

        feat.add_set(self.normalize_token_attr(normed_a_feat, [], "counted"))
        feat.add_set(unnormed_a_feat)
        return feat

    def extract_token_attr_stats_in_input(self, words, tags, attributions, part):
        feat = common.IndexedFeature()
        context_start = words.index(self.sep_tokens[1])
        if part == "Q":
            tags = tags[1:context_start]
            attributions = attributions[1:context_start]
        if part == "C":
            tags = tags[context_start + 2 : -1]
            attributions = attributions[context_start + 2 : -1]

        feat.add("STAT_MEAN_" + part, attributions.mean())
        feat.add("STAT_STD_" + part, attributions.std())
        return feat

    def extract_polarity_feature(
        self,
        attr,
        tags,
        words,
        tags_for_tok,
        ans_index,
        polarity,
        include_basic=True,
        include_stats=False,
    ):
        named_feat = common.IndexedFeature()
        token_attribution = self.aggregate_token_attribution(attr, tags, polarity)

        if include_basic:
            named_feat.add_set(
                self.extract_token_attr_feature_in_question(
                    words, tags_for_tok, token_attribution
                )
            )
            named_feat.add_set(
                self.extract_token_attr_feature_in_context(
                    words, tags_for_tok, token_attribution
                )
            )
            named_feat.add_set(
                self.extract_token_attr_feature_in_input(
                    words, tags_for_tok, token_attribution, ans_index
                )
            )

        if include_stats:
            named_feat.add_set(
                self.extract_token_attr_stats_in_input(
                    words, tags_for_tok, token_attribution, "Q"
                )
            )
            named_feat.add_set(
                self.extract_token_attr_stats_in_input(
                    words, tags_for_tok, token_attribution, "C"
                )
            )
            named_feat.add_set(
                self.extract_token_attr_stats_in_input(
                    words, tags_for_tok, token_attribution, "IN"
                )
            )
        named_feat.add_prefix(polarity + "_")
        return named_feat

    def extract_feature_for_instance(self, attr, tags, preds, states):
        predictions = ast.literal_eval(preds["pred_text"])[0]
        pred_text = predictions[0]["answer"]
        if self.dataset == "squad":
            gold_text = ast.literal_eval(preds["gold_text"])
        else:
            gold_text = preds["gold_text"]

        exact_match = evaluate.get_score(
            metric="exact_match", pred_text=pred_text, gold_text=gold_text
        )
        f1 = evaluate.get_score(metric="f1", pred_text=pred_text, gold_text=gold_text)
        calib_label = 1 if exact_match > 0 else 0
        # print(calib_label)

        # baseline features
        named_feat = common.IndexedFeature()
        named_feat.add_set(self.extract_baseline_feature(preds))

        start_index, end_index = (
            ast.literal_eval(preds["answer_start"])[0],
            ast.literal_eval(preds["answer_end"])[0],
        )
        segments = tags["segments"]
        transformed_start_idx = 0
        while segments[transformed_start_idx + 1][0] <= start_index:
            transformed_start_idx += 1
        transformed_end_idx = 0
        while segments[transformed_end_idx][1] < (end_index + 1):
            transformed_end_idx += 1
        ans_range = (transformed_start_idx, transformed_end_idx)
        # syntactic features
        words, tags_for_tok = tags["words"], tags["tags"]
        # words = [w for w in words if w != " "]
        tags_for_tok = [self.lemmatize_pos_tag(x) for x in tags_for_tok]
        # ents_for_tok = entities['ents']

        named_feat.add_set(self.extract_state_features(attr, words, ans_range, states))
        named_feat.add_set(self.extract_bow_feature(words, tags_for_tok, ans_range))
        named_feat.add_set(
            self.extract_polarity_feature(
                attr, tags, words, tags_for_tok, ans_range, "NEU"
            )
        )
        # print({'feature': named_feat, 'label': calib_label, 'f1_score': f1})
        return {"feature": named_feat, "label": calib_label, "f1_score": f1}


if __name__ == "__main__":

    import argparse

    parser = argparse.ArgumentParser(
        description="Passing arguments for model, tokenizer, and dataset."
    )

    parser.add_argument(
        "--model_name", type=str, required=False, help="Specify the model to use."
    )
    parser.add_argument(
        "--dataset", type=str, required=True, help="Specify the dataset to use."
    )
    parser.add_argument(
        "--method",
        type=str,
        required=True,
        help="Specify the explanation method to use.",
    )
    args = parser.parse_args()

    data_path = f"./src/data/{args.dataset}"
    feat = CreateFeatures(
        path=data_path, model=args.model_name, dataset=args.dataset, method=args.method
    )
    feat.featurize()

src/calibration/baseline/calib_dataset.py
UKPLab/CATfOOD
import os
import torch
import numpy as np
import pandas as pd
from collections import OrderedDict
from tqdm import tqdm
import math
import re
import spacy
from spacy.tokens import Doc

# from supar import Parser
import traceback

from src.calibration.baseline import common, evaluate, utils

import warnings

# ignore warning
warnings.filterwarnings("ignore")


class CreateFeatures:
    def __init__(self, path, model, dataset, method):
        self.path = path
        self.cls_tokens = ["[CLS]", "<s>"]
        self.sep_tokens = ["[SEP]", "</s>"]
        self.count = 0
        self.nlp = spacy.load("en_core_web_sm")
        self.model = model
        self.dataset = dataset  # check the naming
        self.method = method
        # self.supar = Parser.load('crf-con-en')

    def _load_data(self):
        if self.dataset == "squad":
            pos_file = "pos_info.bin"
        else:
            pos_file = "pos_info_wo_space.bin"

        self.tagger_info = utils.load_bin(f"{self.path}/{pos_file}")
        self.attributions = self.build_file_dict(
            dataset=self.dataset, split="addsent-dev", method=self.method
        )
        self.states = utils.load_bin(
            f"{self.path}/dense_repr_pca_10_info_{self.model}.bin"
        )
        self.preds_info = utils.read_json(
            f"{self.path}/nbest_predictions_{self.model}.json"
        )

    def load_interp_info(self, file_dict, qas_id):
        return torch.load(file_dict[qas_id])

    def build_file_dict(self, dataset, split, method):
        # hard-coded path here: be careful
        # prefix = 'squad_sample-addsent_roberta-base'
        prefix = f"{dataset}/dev/roberta"
        if self.model.__contains__("gpt-neox"):
            exp_file_name = f"exp_roberta_gpt_neox_context_rel"
        if self.model.__contains__("llama"):
            exp_file_name = f"exp_roberta_llama_context_rel"
        elif self.model.__contains__("flan-ul2"):
            exp_file_name = f"exp_roberta_flan_ul2_context_noise_rel"
        elif self.model.__contains__("t5"):
            exp_file_name = f"exp_roberta_rag"
        elif self.model == "roberta-squad":
            exp_file_name = f"exp_roberta_base"
        fnames = os.listdir(os.path.join(exp_file_name, method, prefix))
        # print(fnames)
        qa_ids = [x.split(".")[0] for x in fnames]
        # exp_roberta_flan_ul2_context_noise_rel
        fullnames = [os.path.join(exp_file_name, method, prefix, x) for x in fnames]
        return dict(zip(qa_ids, fullnames))

    def featurize(self):
        self._load_data()
        processed_instances = OrderedDict()
        c = 0
        total_count = 0
        # print(self.attributions)
        for q_id in tqdm(
            self.tagger_info, total=len(self.tagger_info), desc="transforming..."
        ):
            total_count += 1
            try:
                tags = self.tagger_info[q_id]
                # print(tags)
                interp = self.load_interp_info(self.attributions, q_id)
                # print(interp)
                states = self.states[q_id]
                preds = self.preds_info[q_id]
                processed_instances[q_id] = self.extract_feature_for_instance(
                    interp, tags, preds, states
                )
            except Exception as e:
                c += 1
                print(f"An exception occurred: {e}")
                print("Total", total_count)
                print(q_id)
                traceback.print_exc()
                continue
            # if total_count==1:
            #     break
        print("Total instances not processed: ", c)
        utils.dump_to_bin(
            processed_instances,
            f"{self.path}/calib_data_{self.method}_{self.model}.bin",
        )

    def lemmatize_pos_tag(self, token_tags):
        """
        limit the amount of POS tags to reduce the number of features
        :param token_tags: tagged tokens
        :return:
        """
        tok, pos, tag = token_tags
        if tag == "NNS":
            tag = "NN"
        if tag == "NNPS":
            tag = "NNP"
        if tag.startswith("JJ"):
            tag = "JJ"
        if tag.startswith("RB"):
            tag = "RB"
        if tag.startswith("W"):
            tag = "W"
        if tag.startswith("PRP"):
            tag = "PRP"
        if tag.startswith("VB"):
            tag = "VB"
        if pos == "PUNCT":  # or tag in [":"]:
            tag = "PUNCT"
        # if tag == '$':
        #     tag = 'SYM'
        # if tag == "``" or tag == "''":
        #     tag = "AUX"
        return tok, pos, tag

    def extract_baseline_feature(self, preds, attr):
        feat = common.IndexedFeature()
        # base features
        for rank, p in enumerate(preds[:5]):
            feat.add(f"BASELINE_PROB_{rank}", p["probability"])
        feat.add("BASELINE_CONTEXT_LENGTH", len(attr["example"][1]))
        feat.add("BASELINE_PRED_ANS_LENGTH", len(preds[0]["text"]))

        top_pred = preds[0]["text"]
        first_distinct_prob = 0
        for i, p in enumerate(preds[1:]):
            overlapping = evaluate.f1_score(p["text"], top_pred)
            if overlapping > 0:
                continue
            first_distinct_prob = p["probability"]
            # print(i+1, top_pred, p['text'], first_distinct_prob)
        feat.add("FIRST_DISTINCT_PROB", first_distinct_prob)
        return feat

    def source_of_token(self, idx, tok, context_start, ans_range):
        if tok in self.cls_tokens or tok in self.sep_tokens:
            return "S"
        if idx >= 1 and idx < context_start:
            return "Q"
        if idx >= ans_range[0] and idx <= ans_range[1]:
            # print("Answer token:", tok)
            return "A"
        return "C"

    def nan_check(self, result):
        for i, (key, value) in enumerate(result):
            if math.isnan(float(value)):
                result[i][1] = 0
        return result

    def count_syllables(self, word: str):
        return len(
            re.findall("(?!e$)[aeiouy]+", word, re.I)
            + re.findall("^[^aeiouy]*e$", word, re.I)
        )

    def handle_zero_division(self, n, d):
        return n / d if d else 0

    def extract_bow_feature(self, words, tags, ans_range):
        feat = common.IndexedFeature()
        context_start = words.index(self.sep_tokens[1])
        for i, (i_token, i_pos, i_tag) in enumerate(tags):
            i_src = self.source_of_token(i, i_token, context_start, ans_range)
            if i_src == "Q" or i_src == "A" or i_src == "C":
                # print('BOW_{}_{}'.format(i_src, i_tag))
                feat.add("BOW_{}_{}".format(i_src, i_tag))
                feat.add("BOW_IN_{}".format(i_tag))

        return feat

    def merge_attribution_by_segments(self, attribution, segments):
        new_val = []
        for a, b in segments:
            new_val.append(np.sum(attribution[a:b], axis=0))
        importance = np.stack(new_val, axis=0)
        return importance

    def aggregate_token_attribution(self, attr, tags, polarity):
        # attr = ast.literal_eval(attr)
        # print(attr)
        # print(np.array(list(attr["attributions"])))
        # attri = [i[1] for i in attr["attributions"]]
        # print(attri)
        # attribution_val = np.array(attr["attributions"])
        attribution_val = attr["attribution"].numpy().copy()
        if polarity == "POS":
            attribution_val[attribution_val < 0] = 0
        elif polarity == "NEG":
            attribution_val[attribution_val > 0] = 0
        elif polarity == "NEU":
            pass
        else:
            raise RuntimeError("Invalid polarity")

        attribution_val = self.merge_attribution_by_segments(
            attribution_val, tags["segments"]
        )
        assert attribution_val.shape[0] == len(tags["segments"])
        # normalize
        attribution_val = attribution_val / np.sum(attribution_val)
        # print(attribution_val)
        return attribution_val

    def normalize_token_attr(self, feat, attributions, norm_method=None):
        if norm_method is None:
            return feat
        if norm_method == "all":
            sum_v = np.sum(attributions)
            for k in feat.data:
                feat.data[k] = feat.data[k] / sum_v if sum_v != 0 else 0
            return feat
        if norm_method == "counted":
            sum_v = sum(feat.data.values())
            for k in feat.data:
                feat.data[k] = feat.data[k] / sum_v if sum_v != 0 else 0
            return feat
        raise RuntimeError(norm_method)

    def extract_token_attr_feature_in_question(
        self, words, tags, attributions, include_punct=False
    ):
        context_start = words.index(self.sep_tokens[1])
        tags = tags[1:context_start]
        attributions = attributions[1:context_start]

        feat = common.IndexedFeature()
        unnorm = common.IndexedFeature()
        sum_v = 0
        for i, (token, pos, tag) in enumerate(tags):
            if tag:
                v = attributions[i]
                if pos == "PUNCT" and not include_punct:
                    continue
                feat.add("NORMED_TOK_Q_" + tag, v)
                unnorm.add("UNNORM_TOK_Q_" + tag, v)
                sum_v += v

        feat = self.normalize_token_attr(feat, attributions)
        feat.add_set(unnorm)
        feat.add("SUM_TOK_Q", sum_v)
        return feat

    def extract_token_attr_feature_in_context(
        self, words, tags, attributions, include_punct=False
    ):
        context_start = words.index(self.sep_tokens[1])
        tags = tags[context_start + 2 : -1]
        attributions = attributions[context_start + 2 : -1]

        feat = common.IndexedFeature()
        unnorm = common.IndexedFeature()
        sum_v = 0
        for i, (token, pos, tag) in enumerate(tags):
            if tag:
                v = attributions[i]
                if pos == "PUNCT" and not include_punct:
                    continue
                feat.add("NORMED_TOK_C_" + tag, v)
                feat.add("UNNORM_TOK_C_" + tag, v)
                sum_v += v

        feat = self.normalize_token_attr(feat, attributions)
        feat.add_set(unnorm)
        feat.add("SUM_TOK_C", sum_v)
        return feat

    def extract_token_attr_feature_in_input(
        self, words, tags, attributions, ans_range, include_punct=False
    ):
        feat = common.IndexedFeature()
        context_start = words.index(self.sep_tokens[1])
        normed_a_feat = common.IndexedFeature()
        unnormed_a_feat = common.IndexedFeature()
        for i, (token, pos, tag) in enumerate(tags):
            if tag:
                v = attributions[i]
                if pos == "PUNCT" and not include_punct:
                    continue
                feat.add("TOK_IN_" + tag, v)

                if i >= ans_range[0] and i <= ans_range[1]:
                    unnormed_a_feat.add("UNNORM_TOK_A_" + tag, v)
                    normed_a_feat.add("NORMED_TOK_A_" + tag, v)

        feat.add_set(self.normalize_token_attr(normed_a_feat, [], "counted"))
        feat.add_set(unnormed_a_feat)
        return feat

    def extract_token_attr_stats_in_input(self, words, tags, attributions, part):
        feat = common.IndexedFeature()
        context_start = words.index(self.sep_tokens[1])
        if part == "Q":
            tags = tags[1:context_start]
            attributions = attributions[1:context_start]
        if part == "C":
            tags = tags[context_start + 2 : -1]
            attributions = attributions[context_start + 2 : -1]

        feat.add("STAT_MEAN_" + part, attributions.mean())
        feat.add("STAT_STD_" + part, attributions.std())
        return feat

    def extract_state_features(self, attr, tokens, ans_range, states):
        feat = common.IndexedFeature()
        context_start = tokens.index(self.sep_tokens[1])
        start, end = ans_range
        ans_indices = list(range(start, end + 1))
        # print(ans_indices)
        # print(states.shape)
        q_states = states[:, 1:context_start, :]
        # print(q_states.shape)
        # print(q_states)
        c_states = states[:, context_start + 2 : -1, :]
        torch.set_printoptions(threshold=10_000)
        # print(c_states)
        # print(c_states[:, :start-(context_start+2), :])
        # remove answer states from context states
        c_states_before = c_states[:, : start - (context_start + 2), :]
        c_states_after = c_states[:, end + 1 - (context_start + 2) :, :]
        final_c_states = torch.cat((c_states_before, c_states_after), dim=1)
        # print(final_c_states.shape)

        a_states = states[:, start : end + 1, :]
        # print(a_states)

        attribution = np.array(attr["attribution"])
        # print(len(attribution))
        question_attr = attribution[1:context_start]
        # print(len(question_attr))
        context_attr = attribution[context_start + 2 : -1]
        c_attr_before = context_attr[: start - (context_start + 2)]
        c_attr_after = context_attr[end + 1 - (context_start + 2) :]
        final_c_attr = np.concatenate((c_attr_before, c_attr_after), axis=0)
        # print(len(context_attr))
        ans_attr = attribution[start : end + 1]
        # print(len(ans_attr))

        ques_sorted_indices = np.argsort(question_attr)
        # take topk indices
        topk_ques_indices = ques_sorted_indices[
            -int(ques_sorted_indices.shape[0] * 0.10) :
        ]
        # print(ques_sorted_indices.shape)
        ques_topk_states = np.take(q_states, topk_ques_indices, axis=1)  # .tolist()[0]

        topc_percent = 0.10
        cxt_sorted_indices = np.argsort(final_c_attr)
        cxt_top_k_indices = cxt_sorted_indices[
            -int(cxt_sorted_indices.shape[0] * topc_percent) :
        ]
        cxt_topk_states = np.take(
            final_c_states, cxt_top_k_indices, axis=1
        )  # .tolist()[0]

        topa_percent = 0.20
        ans_sorted_indices = np.argsort(ans_attr)
        ans_top_k_indices = ans_sorted_indices[
            -int(ans_sorted_indices.shape[0] * topa_percent) :
        ]
        ans_topk_states = np.take(a_states, ans_top_k_indices, axis=1)

        q_rep = torch.mean(ques_topk_states, dim=1).tolist()[0]
        # print(q_rep.shape)
        c_rep = torch.mean(cxt_topk_states, dim=1).tolist()[0]
        # print(c_rep.shape)
        a_rep = torch.mean(ans_topk_states, dim=1).tolist()[0]

        # for i, value in enumerate(q_rep):
        #     if math.isnan(value):
        #         continue
        #     feat.add_new(f"REPR_TOPK_Q_{i}", value)
        for i, value in enumerate(c_rep):
            if math.isnan(value):
                continue
            feat.add_new(f"REPR_TOPK_C_{i}", value)
            # print(value)
        for i, value in enumerate(a_rep):
            if math.isnan(value):
                continue
            feat.add_new(f"REPR_TOPK_A_{i}", value)

        return feat

    def extract_polarity_feature(
        self,
        attr,
        tags,
        words,
        tags_for_tok,
        ans_index,
        polarity,
        include_basic=True,
        include_stats=False,
    ):
        named_feat = common.IndexedFeature()
        token_attribution = self.aggregate_token_attribution(attr, tags, polarity)
        assert token_attribution.size == len(words)
        # if link_attribution is not None:
        #     assert link_attribution.shape == (len(words), len(words))
        if include_basic:
            named_feat.add_set(
                self.extract_token_attr_feature_in_question(
                    words, tags_for_tok, token_attribution
                )
            )
            named_feat.add_set(
                self.extract_token_attr_feature_in_context(
                    words, tags_for_tok, token_attribution
                )
            )
            named_feat.add_set(
                self.extract_token_attr_feature_in_input(
                    words, tags_for_tok, token_attribution, ans_index
                )
            )

        if include_stats:
            named_feat.add_set(
                self.extract_token_attr_stats_in_input(
                    words, tags_for_tok, token_attribution, "Q"
                )
            )
            named_feat.add_set(
                self.extract_token_attr_stats_in_input(
                    words, tags_for_tok, token_attribution, "C"
                )
            )
            named_feat.add_set(
                self.extract_token_attr_stats_in_input(
                    words, tags_for_tok, token_attribution, "IN"
                )
            )
        named_feat.add_prefix(polarity + "_")
        return named_feat

    def extract_feature_for_instance(self, attr, tags, preds, states):
        # print(self.tagger_info)
        # print('extracting features for instance', attr)
        # print(preds)
        pred_text = preds[0]["text"]
        # print(pred_text)
        example = attr["example"]
        # print("------------------")
        # print(preds)
        # print("example:", example)
        if self.dataset == "squad":
            gold_text = example[2]["text"]
        else:
            gold_text = example[2]
        # print('gold text', gold_text)
        # print(gold_text)

        exact_match = evaluate.get_score(
            metric="exact_match", pred_text=pred_text, gold_text=gold_text
        )
        f1 = evaluate.get_score(metric="f1", pred_text=pred_text, gold_text=gold_text)
        calib_label = 1 if exact_match > 0 else 0
        # print(calib_label)

        # baseline features
        named_feat = common.IndexedFeature()
        named_feat.add_set(self.extract_baseline_feature(preds, attr))

        # print("attr:", attr["prediction"])
        start_index, end_index = attr["prelim_result"][0], attr["prelim_result"][1]
        # print("start_index:", start_index)
        # print("end_index:", end_index)
        segments = tags["segments"]

        transformed_start_idx = 0
        while segments[transformed_start_idx + 1][0] <= start_index:
            transformed_start_idx += 1
        transformed_end_idx = 0
        while segments[transformed_end_idx][1] < (end_index + 1):
            transformed_end_idx += 1

        ans_range = (transformed_start_idx, transformed_end_idx)
        # print("ans range:", ans_range)
        # syntactic features
        words, tags_for_tok = tags["words"], tags["tags"]
        # words = [w for w in words if w != " "]
        tags_for_tok = [self.lemmatize_pos_tag(x) for x in tags_for_tok]

        named_feat.add_set(self.extract_state_features(attr, words, ans_range, states))
        named_feat.add_set(self.extract_bow_feature(words, tags_for_tok, ans_range))
        named_feat.add_set(
            self.extract_polarity_feature(
                attr, tags, words, tags_for_tok, ans_range, "NEU"
            )
        )
        # print({'feature': named_feat, 'label': calib_label, 'f1_score': f1})
        return {"feature": named_feat, "label": calib_label, "f1_score": f1}


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        description="Passing arguments for model, tokenizer, and dataset."
    )

    parser.add_argument(
        "--model_name", type=str, required=False, help="Specify the model to use."
    )
    parser.add_argument(
        "--dataset", type=str, required=True, help="Specify the dataset to use."
    )
    parser.add_argument(
        "--method",
        type=str,
        default="shap",
        help="Specify the explanation method to use.",
    )
    args = parser.parse_args()

    data_path = f"./src/data/{args.dataset}"
    feat = CreateFeatures(
        path=data_path, model=args.model_name, dataset=args.dataset, method=args.method
    )
    feat.featurize()

src/calibration/baseline/calibration_dataset.py
UKPLab/CATfOOD
import pandas as pd
import string
import re
from collections import Counter

import src.calibration.baseline.dataloader as dataloader


# BASE_PATH = "/home/sachdeva/projects/ukp/exp_calibration/"
BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


def normalize_answer(s):
    """
    Lower text and remove punctuation, articles and extra whitespace.
    """

    def remove_articles(text):
        return re.sub(r"\b(a|an|the)\b", " ", text)

    def white_space_fix(text):
        return " ".join(text.split())

    def remove_punc(text):
        exclude = set(string.punctuation)
        return "".join(ch for ch in text if ch not in exclude)

    def lower(text):
        return text.lower()

    return white_space_fix(remove_articles(remove_punc(lower(s))))


def f1_score(prediction, ground_truth):
    """
    Compute f1 score of prediction
    :param prediction:
    :param ground_truth:
    :return:
    """
    prediction_tokens = normalize_answer(prediction).split()
    ground_truth_tokens = normalize_answer(ground_truth).split()
    common = Counter(prediction_tokens) & Counter(ground_truth_tokens)
    num_same = sum(common.values())
    if num_same == 0:
        return 0
    precision = 1.0 * num_same / len(prediction_tokens)
    recall = 1.0 * num_same / len(ground_truth_tokens)
    f1 = (2 * precision * recall) / (precision + recall)
    return f1


def exact_match_score(prediction, ground_truth):
    """
    Exact match for Squad evaluation
    :param prediction:
    :param ground_truth:
    :return:
    """
    # print("-------------------")
    # print(normalize_answer(prediction), normalize_answer(ground_truth))
    if normalize_answer(prediction) == normalize_answer(ground_truth):
        return 1
    else:
        return 0


def metric_max_over_ground_truths(metric_fn, prediction, ground_truths):
    scores_for_ground_truths = []
    for ground_truth in ground_truths:
        score = metric_fn(prediction, ground_truth)
        scores_for_ground_truths.append(score)
    # print(scores_for_ground_truths)
    return max(scores_for_ground_truths)


def get_score(metric, pred_text, gold_text):
    # print(pred_text, gold_text)
    # print(type(pred_text[0][0]["answer"]))
    if not isinstance(pred_text, str):
        pred_text = pred_text[0][0]["answer"]
    # print(gold_text, pred_text)
    if isinstance(gold_text, str):
        gold_text = [gold_text]
    if metric == "exact_match":
        score = metric_max_over_ground_truths(exact_match_score, pred_text, gold_text)
    elif metric == "f1":
        score = metric_max_over_ground_truths(f1_score, pred_text, gold_text)
    return score


def macro_ce(preds, golds, metric):
    """Macro CE"""
    ice_pos, ice_neg = 0, 0
    num_pos, num_neg = 0, 0
    for i in range(len(preds)):
        pred_text = preds[i]["text"]
        pred_proba = preds[i]["probability"]
        gold_text = golds[i]
        score = get_score(metric, pred_text, gold_text)
        if score == 1:
            ice_pos += 1 - pred_proba
            num_pos += 1
        else:
            ice_neg += pred_proba - 0
            num_neg += 1
    ice_pos /= num_pos
    ice_neg /= num_neg
    macro_ce = (ice_pos + ice_neg) / 2
    print("Macro CE: ", macro_ce)
    return macro_ce


if __name__ == "__main__":
    pred_df = (
        pd.read_json(
            BASE_PATH
            + "src/data/squad_adv_new/nbest_predictions_roberta_amb_0.75_top20.json"
        )
        .T.rename_axis("id")
        .reset_index()
    )
    answers = [f"answer_{i}" for i in range(20)]
    pred_df.columns = ["id"] + answers
    squad_data = dataloader.PreprocessData(
        "squad_adversarial", "AddSent", save_data=False, save_path="../../../../"
    )
    gold_text = [
        (sample["id"], sample["answers"]["text"])
        for sample in squad_data.processed_val_set()
    ]
    gold_df = pd.DataFrame(gold_text, columns=["id", "gold_answers"])
    data = pd.merge(pred_df, gold_df, on="id")
    pred_data = data.answer_0.values
    gold_data = data.gold_answers.values
    macro_ce(pred_data, gold_data, "exact_match")

src/calibration/baseline/calibration_metrics.py
UKPLab/CATfOOD
from collections import OrderedDict


class IndexedFeature:
    def __init__(self):
        self.data = OrderedDict()

    def add(self, k, v=1.0):
        if k in self.data:
            self.data[k] = self.data[k] + v
        else:
            self.data[k] = v

    def add_new(self, k, v):
        self.data[k] = v

    def add_if_absent(self, k, v=1.0):
        if k not in self.data:
            self.data[k] = v

    def add_set(self, other):
        for k, v in other.data.items():
            self.data[k] = v

    def __getitem__(self, k):
        return self.data.get(k, 0.0)

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        return str(self.data)

    def __len__(self):
        return len(self.data)

    def add_prefix(self, prefix):
        new_data = OrderedDict([(prefix + k, v) for (k, v) in self.data.items()])
        self.data = new_data


class FeatureVocab:
    def __init__(self):
        self.feat_to_id = {}
        self.id_to_feat = {}

    def __getitem__(self, word):
        return self.feat_to_id.get(word, -1)

    def __contains__(self, word):
        return word in self.feat_to_id

    def __len__(self):
        return len(self.feat_to_id)

    def size(self):
        return len(self.feat_to_id)

    def get_word(self, wid):
        return self.id_to_feat[wid]

    def get_names(self):
        return [self.id_to_feat[i] for i in range(len(self))]

    def add(self, word):
        if word not in self:
            wid = self.feat_to_id[word] = len(self)
            self.id_to_feat[wid] = word
            return wid

src/calibration/baseline/common.py
UKPLab/CATfOOD
import itertools
from pprint import pprint
import os
import json
import jsonlines
import re
from tqdm import tqdm

from typing import List

from collections import OrderedDict

from datasets import load_dataset, concatenate_datasets, Features, Value, Sequence


BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"
# BASE_PATH = "/home/sachdeva/projects/ukp/exp_calibration/"


class PreprocessData:
    def __init__(
        self,
        dataset_name: str,
        dataset_config: str,
        cf_path: str = None,
        save_data: bool = False,
        save_path: str = "",
    ):
        """
        load datasets
        """
        dataset = load_dataset(dataset_name, dataset_config)
        if dataset_name == "squad":
            self.train_set = dataset["train"]
        self.val_set = dataset["validation"]
        self.save_data = save_data
        self.save_path = save_path
        self.cf_path = cf_path
        self.count = 0

    def get_data_info(self):
        """
        data statistics
        """
        train_samples = self.train_set.shape[0]
        val_samples = self.val_set.shape[0]
        column_names = self.train_set.column_names
        train_ex = self.train_set.select([0])[0]
        val_ex = self.val_set.select([0])[0]
        info = {
            "train_samples": train_samples,
            "val_samples": val_samples,
            "column_names": column_names,
            "train_ex": train_ex,
            "val_ex": val_ex,
        }
        return OrderedDict(info)

    def remove_trailing_space(self, example):
        # for questions in squad that have extra starting end space
        example["question"] = example["question"].strip()
        return example

    def remove_white_space(self, example):
        example["question"] = " ".join(example["question"].split())
        example["context"] = " ".join(example["context"].split())
        return example

    def processed_train_val_set(self):
        # filter unanswerable questions
        # print(self.train_set.filter(lambda x: len(x["answers"]["text"]) != 1))
        # self.train_set = self.train_set.select(range(1000))
        # self.val_set = self.val_set.select(range(10))
        if self.cf_path is not None:
            self.train_set = self._add_counterfactuals()
        answerable_train_set = self.train_set.filter(
            lambda x: len(x["answers"]["text"]) != 0
        )
        answerable_train_set = answerable_train_set.map(self.remove_trailing_space)
        answerable_train_set = answerable_train_set.map(self.remove_white_space)

        answerable_val_set = self.val_set.filter(
            lambda x: len(x["answers"]["text"]) != 0
        )
        answerable_val_set = answerable_val_set.map(self.remove_trailing_space)
        answerable_val_set = answerable_val_set.map(self.remove_white_space)
        return answerable_train_set, answerable_val_set

    def processed_val_set(self):
        # filter unanswerable questions
        # print(self.train_set.filter(lambda x: len(x["answers"]["text"]) != 1))
        answerable_val_set = self.val_set.filter(
            lambda x: len(x["answers"]["text"]) != 0
        )
        answerable_val_set = answerable_val_set.map(self.remove_trailing_space)
        answerable_val_set = answerable_val_set.map(self.remove_white_space)
        return answerable_val_set

    def sanitize_ptb_tokenized_string(self, text: str) -> str:
        """
        Sanitizes string that was tokenized using PTBTokenizer
        """
        tokens = text.split(" ")
        if len(tokens) == 0:
            return text

        # Replace quotation marks and parentheses
        token_map = {
            "``": '"',
            "''": '"',
            "-lrb-": "(",
            "-rrb-": ")",
            "-lsb-": "[",
            "-rsb-": "]",
            "-lcb-": "{",
            "-rcb-": "}",
            "<s>": "",
            "</s>": "",
        }

        # Merge punctuation with previous tokens
        punct_forward = {"`", "$", "#"}
        punct_backward = {".", ",", "!", "?", ":", ";", "%", "'"}

        # Exact matches that get merged forward or backward
        em_forward = {"(", "[", "{"}
        em_backward = {"n't", "na", ")", "]", "}"}

        new_tokens: List[str] = []

        merge_fwd = False
        for i, orig_token in enumerate(tokens):
            tokens[i] = (
                token_map[orig_token.lower()]
                if orig_token.lower() in token_map
                else orig_token
            )
            new_token = tokens[i].lower()

            # merge_fwd was set by previous token, so it should be prepended to current token
            if merge_fwd:
                tokens[i] = tokens[i - 1] + tokens[i]

            if len(tokens[i]) == 0:
                continue

            # Special cases for `` and '', those tells us if " is the start or end of a quotation.
            # Also always merge tokens starting with ' backward and don't merge back if we just merged forward
            merge_bckwd = not merge_fwd and (
                orig_token == "''"
                or new_token in em_backward
                or new_token.startswith("'")
                or all(c in punct_backward for c in new_token)
            )
            merge_fwd = (
                orig_token == "``"
                or new_token in em_forward
                or all(c in punct_forward for c in new_token)
            )

            if merge_bckwd and new_tokens:
                new_tokens[-1] += tokens[i]
            elif not new_tokens or not merge_fwd or i == len(tokens) - 1:
                new_tokens.append(tokens[i])

        return " ".join(new_tokens)

    def _convert_answer_column(self, example):
        """
        convert answer column to squad format
        """

        context = example["context"].strip()
        # remove extra white space
        context = re.sub(r"\s+", " ", context)
        example["context"] = context

        if "answers" in example:
            answer = example["answers"]["text"][0]
        else:
            answer = example["answer"]["text"]
        # handle spaces between special characters
        special_chars = ["$", "@"]
        if any(char in answer for char in special_chars):
            for char in special_chars:
                answer = answer.replace(f"{char} ", char)
        # remove spaces between special characters (quotations)
        answer = re.sub(r'"\s*([^"]*?)\s*"', r'"\1"', answer)
        answer = re.sub(r" - ", "-", answer)

        # answer = answer.replace("$ ", "$")
        answer_start = str(context.lower()).find(str(answer.lower()))
        if answer_start == -1:
            self.count += 1
            answer_start = str(context.lower()).find(
                self.sanitize_ptb_tokenized_string(answer).lower()
            )
            if answer_start != -1:
                answer = str(self.sanitize_ptb_tokenized_string(answer))
                self.count -= 1
            else:
                # print("=====================================")
                # print(context)
                # print(answer)
                # print("------")
                # print(context.lower())
                # print(self.sanitize_ptb_tokenized_string(answer))
                # print(context.lower().find(self.sanitize_ptb_tokenized_string(answer).lower()))
                # # print(context.lower().index(self.sanitize_ptb_tokenized_string(answer).lower()))
                # print("=====================================")
                answer = ""

        example["answers"] = {"answer_start": [answer_start], "text": [answer]}
        return example

    def _add_counterfactuals(self, process=True):
        """
        add counterfactuals to data
        """
        counterfactuals_dataset = load_dataset(
            "json", data_files=self.cf_path, split="train"
        )
        # counterfactuals_dataset = counterfactuals_dataset.remove_columns("similarity")
        # print(len(counterfactuals_dataset[]))
        # print(counterfactuals_dataset.column_names["train"])
        if "title" not in counterfactuals_dataset.column_names:
            new_title_column = [""] * len(counterfactuals_dataset)
            counterfactuals_dataset = counterfactuals_dataset.add_column(
                "title", new_title_column
            )

        if process:
            counterfactuals_dataset = counterfactuals_dataset.map(
                self._convert_answer_column
            )
            if "answer" in counterfactuals_dataset.column_names:
                counterfactuals_dataset = counterfactuals_dataset.remove_columns(
                    "answer"
                )

        # remove samples with no answer
        counterfactuals_dataset = counterfactuals_dataset.filter(
            lambda x: x["answers"]["answer_start"][0] != -1
        )
        counterfactuals_dataset = counterfactuals_dataset.cast(
            Features(
                {
                    "id": Value("string"),
                    "title": Value("string"),
                    "context": Value("string"),
                    "question": Value("string"),
                    "answers": Sequence(
                        {"answer_start": Value("int32"), "text": Value("string"),}
                    ),
                }
            )
        )
        self.train_set = self.train_set.cast(
            Features(
                {
                    "id": Value("string"),
                    "title": Value("string"),
                    "context": Value("string"),
                    "question": Value("string"),
                    "answers": Sequence(
                        {"answer_start": Value("int32"), "text": Value("string"),}
                    ),
                }
            )
        )

        # ex = counterfactuals_dataset[0]
        # print("CF: ", ex)
        # print(ex["context"].strip().lower())
        # print(ex["context"].strip().lower().find(ex["answers"]["text"][0]))
        # ex = self.train_set[0]
        # print("Train: ", ex)
        # print(ex["context"].find(ex["answers"]["text"][0]))
        # add counterfactuals to example data
        merged_dataset = concatenate_datasets([self.train_set, counterfactuals_dataset])
        return merged_dataset


def get_dev_examples(data_dir, file_name):
    """
    Data loader for dev set of trivia qa and hotpot qa
    """

    with open(os.path.join(data_dir, file_name), "r", encoding="utf-8") as reader:
        input_data = json.load(reader)["data"]
    sum = 0
    for i in input_data:
        sum += len(i["paragraphs"])

    examples = []
    for entry in tqdm(input_data):
        title = entry["title"]
        for paragraph in entry["paragraphs"]:
            context_text = paragraph["context"]

            for qa in paragraph["qas"]:
                qas_id = qa["id"]
                question_text = qa["question"]
                start_position_character = None
                answer_text = None
                answers = []
                question_type = qa.get("question_type", "none")
                is_yesno = qa.get("is_yesno", False)
                is_impossible = qa.get("is_impossible", False)
                if not is_impossible:
                    answer = qa["answers"][0]
                    answer_text = answer["text"]
                    start_position_character = answer["answer_start"]
                else:
                    answers = qa["answers"]
                example = {
                    "qas_id": qas_id,
                    "question_text": question_text,
                    "context_text": context_text,
                    "answer_text": answer_text,
                    "start_position_character": start_position_character,
                    "answers": answers,
                }
                examples.append(example)
    return examples


def get_dev_examples_hf():
    """
    Data loader for dev set of trivia qa and hotpot qa
    """
    with open(
        os.path.join(BASE_PATH + "src/data/dev_trivia.json"), "r", encoding="utf-8"
    ) as reader:
        input_data = json.load(reader)["data"]
    sum = 0
    for i in input_data:
        sum += len(i["paragraphs"])
    print(sum)

    # input_data = input_data.select(range(0, 2000))

    examples = []
    for entry in tqdm(input_data):
        title = entry["title"]
        for paragraph in entry["paragraphs"]:
            context_text = paragraph["context"]

            for qa in paragraph["qas"]:
                qas_id = qa["id"]
                question_text = qa["question"]
                start_position_character = None
                answer_text = None
                answers = []
                question_type = qa.get("question_type", "none")
                is_yesno = qa.get("is_yesno", False)
                is_impossible = qa.get("is_impossible", False)
                if not is_impossible:
                    answer = qa["answers"][0]
                    answer_text = answer["text"]
                    start_position_character = answer["answer_start"]
                else:
                    answers = qa["answers"]
                yield {
                    "id": qas_id,
                    "question": question_text,
                    "context": context_text,
                    "answers": [answer_text],
                    "start_position_character": start_position_character,
                    # "answers": answers,
                }


def get_dev_samples_mrqa(data_path):
    examples = []
    with jsonlines.open(data_path) as reader:
        for example in tqdm(reader):
            # skip first header row
            if "header" in example:
                continue
            context_text = example["context"]
            questions = example["qas"]
            for question in questions:
                id = question["qid"]
                question_text = question["question"]
                answers = question["detected_answers"][0]
                answer_text = answers["text"]
                char_spans = answers["char_spans"]
                token_spans = answers["token_spans"]
                example = {
                    "qas_id": id,
                    "question_text": question_text,
                    "context_text": context_text,
                    "answer_text": answer_text,
                    "char_spans": char_spans,
                    "token_spans": token_spans,
                    "answers": [answer_text],
                }
                examples.append(example)
        return examples


if __name__ == "__main__":
    save_dir = "../data/squad/train_data/"
    BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"
    # BASE_PATH = "/home/sachdeva/projects/ukp/exp_calibration/"
    dataloader = PreprocessData(
        "squad",
        "plain_text",
        # cf_path=None,
        cf_path=BASE_PATH
        + "src/data/squad/counterfactual_data_GPT-NeoXT-Chat-Base-20B_qg_pipeline_final_cleaned.jsonl",
        save_data=False,
        save_path=save_dir,
    )
    dataset = dataloader._add_counterfactuals()
    print(len(dataset))

src/calibration/baseline/dataloader.py
UKPLab/CATfOOD
from collections import Counter
import re
import string
import pandas as pd
import ast


def normalize_answer(s):
    """
    Lower text and remove punctuation, articles and extra whitespace.
    """

    def remove_articles(text):
        return re.sub(r"\b(a|an|the)\b", " ", text)

    def white_space_fix(text):
        return " ".join(text.split())

    def remove_punc(text):
        exclude = set(string.punctuation)
        return "".join(ch for ch in text if ch not in exclude)

    def lower(text):
        return text.lower()

    return white_space_fix(remove_articles(remove_punc(lower(s))))


def f1_score(prediction, ground_truth):
    """
    Compute f1 score of prediction
    :param prediction:
    :param ground_truth:
    :return:
    """
    prediction_tokens = normalize_answer(prediction).split()
    ground_truth_tokens = normalize_answer(ground_truth).split()
    common = Counter(prediction_tokens) & Counter(ground_truth_tokens)
    num_same = sum(common.values())
    if num_same == 0:
        return 0
    precision = 1.0 * num_same / len(prediction_tokens)
    recall = 1.0 * num_same / len(ground_truth_tokens)
    f1 = (2 * precision * recall) / (precision + recall)
    return f1


def exact_match_score(prediction, ground_truth):
    """
    Exact match for Squad evaluation
    :param prediction:
    :param ground_truth:
    :return:
    """
    # print("-------------------")
    # print(normalize_answer(prediction), normalize_answer(ground_truth))
    if normalize_answer(prediction) == normalize_answer(ground_truth):
        return 1
    else:
        return 0


def metric_max_over_ground_truths(metric_fn, prediction, ground_truths):
    scores_for_ground_truths = []
    for ground_truth in ground_truths:
        score = metric_fn(prediction, ground_truth)
        scores_for_ground_truths.append(score)
    # print(scores_for_ground_truths)
    return max(scores_for_ground_truths)


def get_score(metric, pred_text, gold_text):
    # print(pred_text, gold_text)
    # print(type(pred_text[0][0]["answer"]))
    if not isinstance(pred_text, str):
        pred_text = pred_text[0][0]["answer"]
    # print(gold_text, pred_text)
    if isinstance(gold_text, str):
        gold_text = [gold_text]
    if metric == "exact_match":
        score = metric_max_over_ground_truths(exact_match_score, pred_text, gold_text)
    elif metric == "f1":
        score = metric_max_over_ground_truths(f1_score, pred_text, gold_text)
    return score


if __name__ == "__main__":
    data = pd.read_csv("src/data/squad_adv_new/outputs.csv")
    # print(data.head())
    # print(data.columns)
    # data["exact_match"] = data.apply(lambda x: get_score("exact_match",
    #                                                      ast.literal_eval(x["pred_text"]),
    #                                                      ast.literal_eval(x["gold_text"])
    #                                                      ),
    #                                  axis=1)

    data["exact_match"] = data.apply(
        lambda x: get_score(
            "exact_match",
            ast.literal_eval(x["pred_text"]),
            ast.literal_eval(x["gold_text"]),
        ),
        axis=1,
    )

    print(data["exact_match"])
    print(data["exact_match"].value_counts())

src/calibration/baseline/evaluate.py
UKPLab/CATfOOD
import string
import pandas as pd
from typing import List, Tuple
from collections import OrderedDict

import torch
import spacy
from spacy.tokens import Doc
from tqdm import tqdm
from supar import Parser

from transformers import (
    AutoTokenizer,
    # RobertaAdapterModel,
)

from transformers.models.gpt2.tokenization_gpt2 import bytes_to_unicode
from src.calibration.baseline import dataloader


class FeatureExtractor:
    def __init__(self, tokenizer):
        self.tokenizer = tokenizer
        self.cls_token = self.tokenizer.cls_token
        self.eos_token = self.tokenizer.eos_token
        self.supar = Parser.load("crf-con-en")

    def encode(
        self, inputs: list = None, add_special_tokens: bool = True, return_tensors=None
    ):
        """
        Encode inputs using the model tokenizer
        Args:
            inputs: question, context pair as a list
            add_special_tokens: where to add CLS, SEP tokens
            return_tensors: whether to return tensors
        Return:
            tokenized inputs
        """
        return self.tokenizer(
            inputs,
            add_special_tokens=add_special_tokens,
            return_tensors=return_tensors,
            padding=True,
            truncation=True,
            max_length=512,
        )

    def decode(self, input_ids: torch.Tensor, skip_special_tokens: bool) -> List[str]:
        """
        Decode received input_ids into a list of word tokens.
        Args:
            input_ids (torch.Tensor): Input ids representing
            word tokens for a sentence/document.
        """
        return self.tokenizer.convert_ids_to_tokens(
            input_ids[0], skip_special_tokens=skip_special_tokens
        )

    def _bpe_decode(self, tokens: List[str],) -> Tuple[List[str], List]:

        byte_encoder = bytes_to_unicode()
        byte_decoder = {v: k for k, v in byte_encoder.items()}
        decoded_each_tok = [
            bytearray([byte_decoder[c] for c in t]).decode(
                encoding="utf-8", errors="replace"
            )
            for t in tokens
        ]

        end_points = []
        force_break = False
        for idx, token in enumerate(decoded_each_tok):
            # special token, punctuation, alphanumeric
            if (
                token in self.tokenizer.all_special_tokens
                or token in string.punctuation
                or not any([x.isalnum() for x in token.lstrip()])
                or token.lstrip == "'s"
            ):
                end_points.append(idx)
                force_break = True
                continue

            if force_break:
                end_points.append(idx)
                force_break = False
                continue

            if token[0] == " ":
                tokens[idx] = token[:]
                end_points.append(idx)

        end_points.append(len(tokens))

        segments = []
        for i in range(1, len(end_points)):
            if end_points[i - 1] == end_points[i]:
                continue
            segments.append((end_points[i - 1], end_points[i]))

        filtered_tokens = []
        for s0, s1 in segments:
            filtered_tokens.append("".join(decoded_each_tok[s0:s1]))

        return filtered_tokens, segments

    def process_input(self, request):
        inputs = [[request["question"], request["context"]]]
        encoded_inputs = self.encode(
            inputs, add_special_tokens=True, return_tensors="pt"
        )

        decoded_text = self.decode(
            encoded_inputs["input_ids"], skip_special_tokens=False
        )
        filtered_tokens, segments = self._bpe_decode(decoded_text)
        return filtered_tokens, segments

    def entity_features(self, tokens, nlp):
        words = [x.lstrip() for x in tokens]
        spaces = [
            False if i == len(tokens) - 1 else tokens[i + 1][0] == " "
            for i in range(len(tokens))
        ]

        valid_idx = [i for i, w in enumerate(words) if len(w)]
        words = [words[i] for i in valid_idx]
        spaces = [spaces[i] for i in valid_idx]
        doc = Doc(nlp.vocab, words=words, spaces=spaces)
        processed_tokens = nlp(doc)
        n_sent, n_token = 0, 0
        for sent in processed_tokens.sents:
            n_sent += 1
            for token in sent:
                n_token += 1

        to_EntiM_C = 0
        to_UEnti_C = 0
        ent_list = []
        unique_ent_list = []

        for ent in processed_tokens.ents:
            to_EntiM_C += 1
            ent_list.append(ent.text)

        for ent in ent_list:
            if ent_list.count(ent) == 1:
                to_UEnti_C += 1
                unique_ent_list.append(ent)

        result = {
            # total number of Entities Mentions counts
            "to_EntiM_C": to_EntiM_C,
            # average number of Entities Mentions counts per sentence
            "as_EntiM_C": to_EntiM_C / n_sent,
            # average number of Entities Mentions counts per token (word)
            "at_EntiM_C": to_EntiM_C / n_token,
            # unique ents...
            "to_UEnti_C": to_UEnti_C,
            "as_UEnti_C": to_UEnti_C / n_sent,
            "at_UEnti_C": to_UEnti_C / n_token,
        }
        print(result)

    def retrieve(self, tokens, nlp):

        words = [x.lstrip() for x in tokens]
        spaces = [
            False if i == len(tokens) - 1 else tokens[i + 1][0] == " "
            for i in range(len(tokens))
        ]

        valid_idx = [i for i, w in enumerate(words) if len(w)]
        words = [words[i] for i in valid_idx]
        spaces = [spaces[i] for i in valid_idx]
        doc = Doc(nlp.vocab, words=words, spaces=spaces)
        processed_tokens = nlp(doc)
        n_sent, n_token = 0, 0
        for sent in processed_tokens.sents:
            n_sent += 1
            for token in sent:
                n_token += 1

        to_AAKuW_C = 0
        to_AAKuL_C = 0
        to_AABiL_C = 0
        to_AABrL_C = 0
        to_AACoL_C = 0

        DB = pd.read_csv("./src/resources/AoAKuperman.csv")
        DB.set_index("Word", inplace=True, drop=True)
        for token in words:
            if token in DB.index:
                scores_for_this_token = list(DB.loc[token, :])
                for i, score in enumerate(scores_for_this_token):
                    scores_for_this_token[i] = (
                        0 if str(score) == "none" else scores_for_this_token[i]
                    )
                to_AAKuW_C += float(scores_for_this_token[7])
                to_AAKuL_C += float(scores_for_this_token[9])
                to_AABiL_C += float(scores_for_this_token[11])
                to_AABrL_C += float(scores_for_this_token[12])
                to_AACoL_C += float(scores_for_this_token[13])

        result = {
            "to_AAKuW_C": to_AAKuW_C,
            "as_AAKuW_C": to_AAKuW_C / n_sent,
            "at_AAKuW_C": to_AAKuW_C / n_token,
            "to_AAKuL_C": to_AAKuL_C,
            "as_AAKuL_C": to_AAKuL_C / n_sent,
            "at_AAKuL_C": to_AAKuL_C / n_token,
            "to_AABiL_C": to_AABiL_C,
            "as_AABiL_C": to_AABiL_C / n_sent,
            "at_AABiL_C": to_AABiL_C / n_token,
            "to_AABrL_C": to_AABrL_C,
            "as_AABrL_C": to_AABrL_C / n_sent,
            "at_AABrL_C": to_AABrL_C / n_token,
            "to_AACoL_C": to_AABrL_C,
            "as_AACoL_C": to_AABrL_C / n_sent,
            "at_AACoL_C": to_AABrL_C / n_token,
        }
        return result

    def handle_zero_division(self, n, d):
        return n / d if d else 0

    def phrF(self, tokens, nlp):

        words = [x.lstrip() for x in tokens]
        spaces = [
            False if i == len(tokens) - 1 else tokens[i + 1][0] == " "
            for i in range(len(tokens))
        ]

        valid_idx = [i for i, w in enumerate(words) if len(w)]
        words = [words[i] for i in valid_idx]
        spaces = [spaces[i] for i in valid_idx]
        doc = Doc(nlp.vocab, words=words, spaces=spaces)
        processed_tokens = nlp(doc)
        n_sent, n_token = 0, 0
        sent_token_list = []
        for sent in processed_tokens.sents:
            n_sent += 1
            temp_list = []
            for token in sent:
                n_token += 1
                temp_list.append(token.text)
            sent_token_list.append(temp_list)

        to_NoPhr_C = 0
        to_VePhr_C = 0
        to_SuPhr_C = 0
        to_PrPhr_C = 0
        to_AjPhr_C = 0
        to_AvPhr_C = 0
        for sent in sent_token_list:
            dataset = self.supar.predict([sent], prob=True, verbose=False)
            parsed_tree = str(dataset.sentences)
            to_NoPhr_C += parsed_tree.count("NP")
            to_VePhr_C += parsed_tree.count("VP")
            to_SuPhr_C += parsed_tree.count("SBAR")
            to_PrPhr_C += parsed_tree.count("PP")
            to_AjPhr_C += parsed_tree.count("ADJP")
            to_AvPhr_C += parsed_tree.count("ADVP")
        result = {
            "to_NoPhr_C": to_NoPhr_C,
            "as_NoPhr_C": float(self.handle_zero_division(to_NoPhr_C, n_sent)),
            "at_NoPhr_C": float(self.handle_zero_division(to_NoPhr_C, n_token)),
            "ra_NoVeP_C": float(self.handle_zero_division(to_NoPhr_C, to_VePhr_C)),
            "ra_NoSuP_C": float(self.handle_zero_division(to_NoPhr_C, to_SuPhr_C)),
            "ra_NoPrP_C": float(self.handle_zero_division(to_NoPhr_C, to_PrPhr_C)),
            "ra_NoAjP_C": float(self.handle_zero_division(to_NoPhr_C, to_AjPhr_C)),
            "ra_NoAvP_C": float(self.handle_zero_division(to_NoPhr_C, to_AvPhr_C)),
            "to_VePhr_C": to_VePhr_C,
            "as_VePhr_C": float(self.handle_zero_division(to_VePhr_C, n_sent)),
            "at_VePhr_C": float(self.handle_zero_division(to_VePhr_C, n_token)),
            "ra_VeNoP_C": float(self.handle_zero_division(to_VePhr_C, to_NoPhr_C)),
            "ra_VeSuP_C": float(self.handle_zero_division(to_VePhr_C, to_SuPhr_C)),
            "ra_VePrP_C": float(self.handle_zero_division(to_VePhr_C, to_PrPhr_C)),
            "ra_VeAjP_C": float(self.handle_zero_division(to_VePhr_C, to_AjPhr_C)),
            "ra_VeAvP_C": float(self.handle_zero_division(to_VePhr_C, to_AvPhr_C)),
            "to_SuPhr_C": to_SuPhr_C,
            "as_SuPhr_C": float(self.handle_zero_division(to_SuPhr_C, n_sent)),
            "at_SuPhr_C": float(self.handle_zero_division(to_SuPhr_C, n_token)),
            "ra_SuNoP_C": float(self.handle_zero_division(to_SuPhr_C, to_NoPhr_C)),
            "ra_SuVeP_C": float(self.handle_zero_division(to_SuPhr_C, to_VePhr_C)),
            "ra_SuPrP_C": float(self.handle_zero_division(to_SuPhr_C, to_PrPhr_C)),
            "ra_SuAjP_C": float(self.handle_zero_division(to_SuPhr_C, to_AjPhr_C)),
            "ra_SuAvP_C": float(self.handle_zero_division(to_SuPhr_C, to_AvPhr_C)),
            "to_PrPhr_C": to_PrPhr_C,
            "as_PrPhr_C": float(self.handle_zero_division(to_PrPhr_C, n_sent)),
            "at_PrPhr_C": float(self.handle_zero_division(to_PrPhr_C, n_token)),
            "ra_PrNoP_C": float(self.handle_zero_division(to_PrPhr_C, to_NoPhr_C)),
            "ra_PrVeP_C": float(self.handle_zero_division(to_PrPhr_C, to_VePhr_C)),
            "ra_PrSuP_C": float(self.handle_zero_division(to_PrPhr_C, to_SuPhr_C)),
            "ra_PrAjP_C": float(self.handle_zero_division(to_PrPhr_C, to_AjPhr_C)),
            "ra_PrAvP_C": float(self.handle_zero_division(to_PrPhr_C, to_AvPhr_C)),
            "to_AjPhr_C": to_AjPhr_C,
            "as_AjPhr_C": float(self.handle_zero_division(to_AjPhr_C, n_sent)),
            "at_AjPhr_C": float(self.handle_zero_division(to_AjPhr_C, n_token)),
            "ra_AjNoP_C": float(self.handle_zero_division(to_AjPhr_C, to_NoPhr_C)),
            "ra_AjVeP_C": float(self.handle_zero_division(to_AjPhr_C, to_VePhr_C)),
            "ra_AjSuP_C": float(self.handle_zero_division(to_AjPhr_C, to_SuPhr_C)),
            "ra_AjPrP_C": float(self.handle_zero_division(to_AjPhr_C, to_PrPhr_C)),
            "ra_AjAvP_C": float(self.handle_zero_division(to_AjPhr_C, to_AvPhr_C)),
            "to_AvPhr_C": to_AvPhr_C,
            "as_AvPhr_C": float(self.handle_zero_division(to_AvPhr_C, n_sent)),
            "at_AvPhr_C": float(self.handle_zero_division(to_AvPhr_C, n_token)),
            "ra_AvNoP_C": float(self.handle_zero_division(to_AvPhr_C, to_NoPhr_C)),
            "ra_AvVeP_C": float(self.handle_zero_division(to_AvPhr_C, to_VePhr_C)),
            "ra_AvSuP_C": float(self.handle_zero_division(to_AvPhr_C, to_SuPhr_C)),
            "ra_AvPrP_C": float(self.handle_zero_division(to_AvPhr_C, to_PrPhr_C)),
            "ra_AvAjP_C": float(self.handle_zero_division(to_AvPhr_C, to_AjPhr_C)),
        }
        return result

    def extract_features(self, request, nlp):
        words, segments = self.process_input(request)
        # tokens = list(filter(None, tokens))

        context_start = words.index(self.tokenizer.eos_token)
        question_tokens = words[1:context_start]
        context_tokens = words[context_start + 2 : -1]
        question_ent_info = self.phrF(question_tokens, nlp)
        print(question_ent_info)
        # context_ent_info = self.entity_features(context_tokens, nlp)
        # tag_info = [(self.cls_token, 'SOS', 'SOS')] + \
        #            question_tag_info +\
        #            [(self.eos_token, 'EOS', 'EOS'),
        #             (self.eos_token, 'EOS', 'EOS')] + \
        #            context_tag_info + \
        #            [(self.eos_token, 'EOS', 'EOS')]
        #
        # assert len(tag_info) == len(words), \
        #     "tags and words not equal"
        # # print(tag_info)
        # instance_info = {'words': words, 'segments': segments, 'tags': tag_info}
        # # print(instance_info)
        # return instance_info


if __name__ == "__main__":
    model_path = "./src/checkpoints/roberta_squad"
    # model = RobertaModel.from_pretrained(model_path)
    tokenizer = AutoTokenizer.from_pretrained(model_path)

    data = dataloader.PreprocessData(
        "squad_adversarial", "AddSent", save_data=False, save_path="../../../../"
    )
    outputs = list()
    tagger = FeatureExtractor(tokenizer=tokenizer)
    nlp = spacy.load("en_core_web_sm")

    c = 0
    processed_instances = OrderedDict()
    for ex in tqdm(data.processed_val_set()):
        # ex = remove_white_space(ex)
        try:
            # if ex["id"] == "56e1239acd28a01900c67641":
            #     print(ex)
            tag_info = tagger.extract_features(
                request={
                    "id": ex["id"],
                    "question": ex["question"],
                    "context": ex["context"],
                },
                nlp=nlp,
            )
            processed_instances[ex["id"]] = tag_info
            # print(tag_info)
            c += 1
        except Exception as e:
            print(f"Unable to get tags: {e}")
            print(ex)
        # if c == 1:
        break
    # utils.dump_to_bin(processed_instances,
    #                   "./src/results/squad_tag_info.bin")
    #
    # print(f"Saved instances: {c}")

src/calibration/baseline/feature_extraction.py
UKPLab/CATfOOD
import numpy as np
import pandas as pd
import argparse
from tqdm import tqdm
from typing import List, Tuple

import torch
from transformers import (
    PreTrainedModel,
    PreTrainedTokenizer,
    AutoTokenizer,
    AutoModelForQuestionAnswering,
    # BertAdapterModel,
    # RobertaAdapterModel
)
from src.calibration.baseline import dataloader

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"
# BASE_PATH = "/home/sachdeva/projects/ukp/exp_calibration/"


class Inference:
    def __init__(
        self, model: PreTrainedModel, tokenizer: PreTrainedTokenizer,
    ):
        self.model = model
        self.tokenizer = tokenizer

        if self.model.config.model_type == "gpt2":
            self.ref_token_id = self.tokenizer.eos_token_id
        else:
            self.ref_token_id = self.tokenizer.pad_token_id

        self.sep_token_id = (
            self.tokenizer.sep_token_id
            if self.tokenizer.sep_token_id is not None
            else self.tokenizer.eos_token_id
        )
        self.cls_token_id = (
            self.tokenizer.cls_token_id
            if self.tokenizer.cls_token_id is not None
            else self.tokenizer.bos_token_id
        )

        self.model_prefix = model.base_model_prefix

        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)

    def _ensure_tensor_on_device(self, **inputs):
        """
        Ensure PyTorch tensors are on the specified device.

        Args:
            inputs (keyword arguments that should be :obj:`torch.Tensor`): The tensors to place on :obj:`self.device`.

        Return:
            :obj:`Dict[str, torch.Tensor]`: The same as :obj:`inputs` but on the proper device.
        """
        return {name: tensor.to(self.model.device) for name, tensor in inputs.items()}

    def encode(
        self, inputs: list = None, add_special_tokens: bool = True, return_tensors=None
    ):
        """
        Encode inputs using the model tokenizer
        Args:
            inputs: question, context pair as a list
            add_special_tokens: where to add CLS, SEP tokens
            return_tensors: whether to return tensors
        Return:
            tokenized inputs
        """
        return self.tokenizer(
            inputs,
            add_special_tokens=add_special_tokens,
            return_tensors=return_tensors,
            padding=True,
            truncation=True,
            max_length=512,
        )

    def decode(self, input_ids: torch.Tensor, skip_special_tokens: bool) -> List[str]:
        """
        Decode received input_ids into a list of word tokens.
        Args:
            input_ids (torch.Tensor): Input ids representing
            word tokens for a sentence/document.
        """
        return self.tokenizer.convert_ids_to_tokens(
            input_ids[0], skip_special_tokens=skip_special_tokens
        )

    def _predict(self, request, batch_size: int = 1) -> tuple:
        """
        Inference on the input.
        Args:
            inputs: list of question, context input
        Returns:
             The model outputs and optionally the input features
        """
        all_predictions = []
        # self.model.to("cuda" if torch.cuda.is_available() else "cpu")

        features = self.encode(request, add_special_tokens=True, return_tensors="pt",)

        for start_idx in range(0, len(request), batch_size):
            with torch.no_grad():
                input_features = {
                    k: features[k][start_idx : start_idx + batch_size]
                    for k in features.keys()
                }
                input_features = self._ensure_tensor_on_device(**input_features)
                predictions = self.model(**input_features)
                all_predictions.append(predictions)
        keys = all_predictions[0].keys()
        final_prediction = {}
        for key in keys:
            if isinstance(all_predictions[0][key], tuple):
                tuple_of_lists = list(
                    zip(
                        *[
                            [
                                torch.stack(p).to(self.device)
                                if isinstance(p, tuple)
                                else p.to(self.device)
                                for p in tpl[key]
                            ]
                            for tpl in all_predictions
                        ]
                    )
                )
                final_prediction[key] = tuple(torch.cat(l) for l in tuple_of_lists)
            else:
                final_prediction[key] = torch.cat(
                    [p[key].to(self.device) for p in all_predictions]
                )
        # print(final_prediction)
        # print(features)
        return final_prediction, features

    def question_answering(self, request, top_k: int = 1):
        """
        Span-based question answering for a given question and context.
        We expect the input to use the (question, context) format for the text pairs.
        Args:
          request: the prediction request
          top_k: num preds

        """

        def decode(
            start_: np.ndarray,
            end_: np.ndarray,
            topk: int,
            max_answer_len: int,
            undesired_tokens_: np.ndarray,
        ) -> Tuple:
            """
            Take the output of any :obj:`ModelForQuestionAnswering` and will generate probabilities
            for each span to be the actual answer.

            In addition, it filters out some unwanted/impossible cases like answer len being greater
            than max_answer_len or answer end position being before the starting position. The method
            supports output the k-best answer through the topk argument.

            Args:
                start_ (:obj:`np.ndarray`): Individual start probabilities for each token.
                end (:obj:`np.ndarray`): Individual end_ probabilities for each token.
                topk (:obj:`int`): Indicates how many possible answer span(s) to extract from the model output.
                max_answer_len (:obj:`int`): Maximum size of the answer to extract from the model's output.
                undesired_tokens_ (:obj:`np.ndarray`): Mask determining tokens that can be part of the answer
            """
            # Ensure we have batch axis
            if start_.ndim == 1:
                start_ = start_[None]

            if end_.ndim == 1:
                end_ = end_[None]

            # Compute the score of each tuple(start_, end_) to be the real answer
            outer = np.matmul(np.expand_dims(start_, -1), np.expand_dims(end_, 1))

            # Remove candidate with end_ < start_ and end_ - start_ > max_answer_len
            candidates = np.tril(np.triu(outer), max_answer_len - 1)

            #  Inspired by Chen & al. (https://github.com/facebookresearch/DrQA)
            scores_flat = candidates.flatten()
            if topk == 1:
                idx_sort = [np.argmax(scores_flat)]
            elif len(scores_flat) < topk:
                idx_sort = np.argsort(-scores_flat)
            else:
                idx = np.argpartition(-scores_flat, topk)[0:topk]
                idx_sort = idx[np.argsort(-scores_flat[idx])]

            starts_, ends_ = np.unravel_index(idx_sort, candidates.shape)[1:]
            desired_spans = np.isin(starts_, undesired_tokens_.nonzero()) & np.isin(
                ends_, undesired_tokens_.nonzero()
            )
            starts_ = starts_[desired_spans]
            ends_ = ends_[desired_spans]
            scores_ = candidates[0, starts_, ends_]

            return starts_, ends_, scores_

        predictions, features = self._predict(request)
        start_idx = torch.argmax(predictions["start_logits"])
        end_idx = torch.argmax(predictions["end_logits"])
        answer_start = torch.tensor([start_idx])
        answer_end = torch.tensor([end_idx])

        # print(features)
        task_outputs = {"answers": []}
        for idx, (start, end, (_, context)) in enumerate(
            zip(predictions["start_logits"], predictions["end_logits"], request)
        ):
            start = start.cpu().detach().numpy()
            end = end.cpu().detach().numpy()
            # Ensure padded tokens & question tokens cannot belong to the set of candidate answers.
            question_tokens = np.abs(
                np.array([s != 1 for s in features.sequence_ids(idx)]) - 1
            )
            # Unmask CLS token for 'no answer'
            question_tokens[0] = 1
            undesired_tokens = question_tokens & features["attention_mask"][idx].numpy()

            # Generate mask
            undesired_tokens_mask = undesired_tokens == 0.0

            # Make sure non-context indexes in the tensor cannot contribute to the softmax
            start = np.where(undesired_tokens_mask, -10000.0, start)
            end = np.where(undesired_tokens_mask, -10000.0, end)

            start = np.exp(
                start - np.log(np.sum(np.exp(start), axis=-1, keepdims=True))
            )
            end = np.exp(end - np.log(np.sum(np.exp(end), axis=-1, keepdims=True)))

            # Get score for 'no answer' then mask for decoding step (CLS token
            no_answer_score = (start[0] * end[0]).item()
            start[0] = end[0] = 0.0

            starts, ends, scores = decode(start, end, top_k, 128, undesired_tokens)
            enc = features[idx]
            answers = [
                {
                    "score": score.item(),
                    "start": enc.word_to_chars(enc.token_to_word(s), sequence_index=1)[
                        0
                    ],
                    "end": enc.word_to_chars(enc.token_to_word(e), sequence_index=1)[1],
                    "answer": context[
                        enc.word_to_chars(enc.token_to_word(s), sequence_index=1)[
                            0
                        ] : enc.word_to_chars(enc.token_to_word(e), sequence_index=1)[1]
                    ],
                }
                for s, e, score in zip(starts, ends, scores)
            ]

            answers.append(
                {"score": no_answer_score, "start": 0, "end": 0, "answer": ""}
            )
            answers = sorted(answers, key=lambda x: x["score"], reverse=True)[:top_k]
            task_outputs["answers"].append(answers)
            task_outputs["answer_start"] = answer_start.cpu().detach().numpy()
            task_outputs["answer_end"] = answer_end.cpu().detach().numpy()
        return task_outputs


if __name__ == "__main__":

    parser = argparse.ArgumentParser(
        description="Passing arguments for model, tokenizer, and dataset."
    )

    parser.add_argument(
        "--model_name",
        default="roberta-squad-flan-ul2-cfs-cleaned-qa",
        type=str,
        required=False,
        help="Specify the model to use.",
    )
    parser.add_argument(
        "--tokenizer",
        default="roberta-base",
        type=str,
        required=False,
        help="Specify the tokenizer to use.",
    )
    parser.add_argument(
        "--dataset", type=str, required=True, help="Specify the dataset to use."
    )

    args = parser.parse_args()

    model = AutoModelForQuestionAnswering.from_pretrained(BASE_PATH + args.model_name)
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)

    inf = Inference(model=model, tokenizer=tokenizer)

    if args.dataset == "squad":
        loader = dataloader.PreprocessData(
            "squad", "plain_text", save_data=False, save_path="../../../../"
        )
        data = loader.processed_val_set()
    elif args.dataset == "squad_adversarial":
        loader = dataloader.PreprocessData(
            "squad_adversarial", "AddSent", save_data=False, save_path="../../../../"
        )
        data = loader.processed_val_set()
    elif args.dataset == "trivia_qa":
        data = dataloader.get_dev_examples("./src/data", "dev_trivia.json")
    elif args.dataset == "hotpot_qa":
        data = dataloader.get_dev_examples("./src/data", "dev_hotpot.json")
    elif args.dataset == "news_qa":
        data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/NewsQA.jsonl")
    elif args.dataset == "bioasq":
        data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/BioASQ-dev.jsonl")
    elif args.dataset == "natural_questions":
        data = dataloader.get_dev_samples_mrqa(
            BASE_PATH + "src/data/NaturalQuestionsShort.jsonl"
        )
    else:
        raise ValueError("Dataset not supported.")

    def remove_white_space(example):
        example["question_text"] = " ".join(example["question_text"].split())
        example["context_text"] = " ".join(example["context_text"].split())
        return example

    outputs = list()
    c = 0
    for ex in tqdm(data):
        # print(ex)
        if args.dataset not in ["squad", "squad_adversarial"]:
            ex = remove_white_space(ex)
        try:
            if args.dataset not in ["squad", "squad_adversarial"]:
                id = ex["qas_id"]
                question = ex["question_text"]
                context = ex["context_text"]
                answer = ex["answer_text"]
            else:
                id = ex["id"]
                question = ex["question"]
                context = ex["context"]
                answer = ex["answers"]["text"]

            prediction = inf.question_answering(
                request=[[question, context,]], top_k=20
            )
            result = {
                "id": id,
                "question": question,
                "context": context,
                "gold_text": answer,
                "pred_text": prediction["answers"],
                "answer_start": prediction["answer_start"],
                "answer_end": prediction["answer_end"],
            }
            outputs.append(result)
            c += 1
        except Exception as e:
            print(f"Unable to get prediction: {e}")

    df = pd.DataFrame(outputs)
    df.to_csv(f"{BASE_PATH}src/data/{args.dataset}/outputs_{args.model_name}.csv")

src/calibration/baseline/inference.py
UKPLab/CATfOOD
# class to perturb input until the label changes

import numpy as np
import torch
from typing import Tuple
from transformers import (
    BertAdapterModel,
    AutoTokenizer,
    PreTrainedTokenizer,
    PreTrainedModel,
    BatchEncoding,
)

from src.calibration.explainers import random_attribution
from src import shap_attributions
from src.calibration.explainers.gradients import simple_grads, integrated_grads
from src.calibration.explainers.attentions import attention, scaled_attention

torch.manual_seed(4)
torch.cuda.manual_seed(4)
np.random.seed(4)


class InputReduction:
    def __init__(
        self,
        model: PreTrainedModel,
        tokenizer: PreTrainedTokenizer,
        method: str,
        request: dict,
    ):

        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.model = model
        self.model.to(self.device)
        self.tokenizer = tokenizer
        self.question = request["question"]
        self.context = request["context"]
        self.top_k = request["top_k"]
        self.mode = request["mode"]
        self.method = method

        self.sep_token = (
            self.tokenizer.sep_token
            if self.tokenizer.sep_token is not None
            else self.tokenizer.eos_token
        )
        self.cls_token = (
            self.tokenizer.cls_token
            if self.tokenizer.cls_token is not None
            else self.tokenizer.bos_token
        )

    def compute_attributions(self):
        if self.method == "random":
            imp_scores = random_attribution.RandomAttributions(
                self.model, self.tokenizer
            )
            attributions = imp_scores.interpret(
                [[self.question, self.context]], self.top_k, output="raw"
            )
        elif self.method == "attention":
            grads = attention.AttnAttribution(self.model, self.tokenizer)
            attributions = grads.interpret(
                [[self.question, self.context]], self.top_k, output="raw"
            )
        elif self.method == "scaled_attn":
            grads = scaled_attention.ScaledAttention(self.model, self.tokenizer)
            attributions = grads.interpret(
                [[self.question, self.context]], self.top_k, output="raw"
            )
        elif self.method == "simple_grads":
            grads = simple_grads.SimpleGradients(self.model, self.tokenizer)
            attributions = grads.interpret(
                [[self.question, self.context]], self.top_k, output="raw"
            )
        elif self.method == "integrated_grads":
            grads = integrated_grads.IntegratedGradients(self.model, self.tokenizer)
            attributions = grads.interpret(
                [[self.question, self.context]], self.top_k, output="raw"
            )
        # elif self.method == "smooth_grads":
        #     grads = smooth_grads.SmoothGradients(self.model, self.tokenizer)
        #     attributions = grads.interpret([[self.question, self.context]], self.top_k, output="raw")
        elif self.method == "shap":
            grads = shap_attributions.SHAPAttributions(
                model=self.model,
                tokenizer=self.tokenizer,
                request=[[self.question, self.context]],
                visualize=False,
                topk=self.top_k,
            )
            attributions = grads.run(output="raw")
        else:
            raise Exception("Attribution method not allowed")
        # print(attributions)
        return attributions

    def get_predictions(self):
        enc_inputs, attributions, answer_start, answer_end = self.compute_attributions()
        enc_inputs.to("cpu")
        # print(len(enc_inputs["input_ids"][0]), len(attributions))
        token_ids = enc_inputs["input_ids"]
        sep_idx = token_ids.tolist()[0].index(self.tokenizer.sep_token_id)
        # print(sep_idx)
        # don't mask sep and pad tokens
        mask = (token_ids != self.tokenizer.sep_token_id).long() & (
            token_ids != self.tokenizer.pad_token_id
        ).long()
        instance_attribution = (
            torch.tensor(attributions, dtype=torch.float32).to(self.device)
            + (1 - mask.to(self.device)) * 1e-10
        )
        # get top_k % words
        num_words_to_mask = int((self.top_k / 100) * instance_attribution.size()[1])
        if self.mode == "max":
            # get top_k attributions from context; can be changed later to include question too
            topk_idx = torch.topk(
                instance_attribution[0][sep_idx:], num_words_to_mask
            ).indices
        elif self.mode == "min":
            topk_idx = torch.topk(
                instance_attribution[0][sep_idx:], num_words_to_mask, largest=False
            ).indices
        # mask top_k tokens from context
        words_to_mask = torch.zeros(instance_attribution.size())
        words_to_mask[0][topk_idx + sep_idx] = 1
        tmp_mask = words_to_mask
        sum_before_mask = torch.sum(tmp_mask[0])
        # unmask answer and sep tokens
        words_to_mask[0][answer_start : answer_end + 1] = 0
        sum_after_mask = torch.sum(words_to_mask[0])
        mask_diff = (sum_before_mask - sum_after_mask).cpu().detach().numpy()
        # if answer tokens are masked, we need to unmask them and mask other tokens
        if mask_diff > 0:
            topk_idx = torch.topk(
                instance_attribution[0][sep_idx:], num_words_to_mask + int(mask_diff)
            ).indices
            # mask top_k tokens from context
            words_to_mask = torch.zeros(instance_attribution.size())
            words_to_mask[0][topk_idx + sep_idx] = 1
            # unmask answer and sep tokens
            words_to_mask[0][answer_start : answer_end + 1] = 0

        words_to_mask[0][sep_idx] = 0
        # words_to_mask.to(self.device)
        # mask features
        # print(token_ids, words_to_mask)
        inputs_masked = torch.tensor(
            token_ids * (1 - words_to_mask)
            + words_to_mask * self.tokenizer.mask_token_id,
            dtype=torch.int,
        ).to(self.device)
        # print(inputs_masked)
        dec_inputs = self.tokenizer.decode(
            torch.tensor(inputs_masked[0], dtype=torch.int)
        )
        # print(dec_inputs)
        inputs = dec_inputs.split()
        filter_idx = list(inputs).index(self.sep_token)
        inputs = [
            value if self.sep_token not in value else value.replace(self.sep_token, "")
            for value in inputs
        ]
        masked_context = " ".join(inputs[filter_idx:])
        # print(masked_context)
        output_new = self.question_answering([[self.question, masked_context]])
        outputs_old = self.question_answering([[self.question, self.context]])

        return outputs_old, output_new

    def _ensure_tensor_on_device(self, **inputs):
        """
        Ensure PyTorch tensors are on the specified device.

        Args:
            inputs (keyword arguments that should be :obj:`torch.Tensor`): The tensors to place on :obj:`self.device`.

        Return:
            :obj:`Dict[str, torch.Tensor]`: The same as :obj:`inputs` but on the proper device.
        """
        return {name: tensor.to(self.model.device) for name, tensor in inputs.items()}

    def _predict(self, request, batch_size=1) -> Tuple[dict, BatchEncoding]:
        """
        Inference on the input.

        Args:
         request: the request with the input and optional kwargs
         batch_size: input batch size

        Returns:
             The model outputs and optionally the input features
        """
        all_predictions = []
        # self.model.to("cuda" if torch.cuda.is_available() else "cpu")

        features = self.tokenizer(
            request, return_tensors="pt", padding=True, truncation=True, max_length=512
        )
        for start_idx in range(0, len(request), batch_size):
            with torch.no_grad():
                input_features = {
                    k: features[k][start_idx : start_idx + batch_size]
                    for k in features.keys()
                }
                input_features = self._ensure_tensor_on_device(**input_features)
                predictions = self.model(**input_features)
                all_predictions.append(predictions)
        keys = all_predictions[0].keys()
        final_prediction = {}
        for key in keys:
            if isinstance(all_predictions[0][key], tuple):
                tuple_of_lists = list(
                    zip(
                        *[
                            [
                                torch.stack(p).to(self.device)
                                if isinstance(p, tuple)
                                else p.to(self.device)
                                for p in tpl[key]
                            ]
                            for tpl in all_predictions
                        ]
                    )
                )
                final_prediction[key] = tuple(torch.cat(l) for l in tuple_of_lists)
            else:
                final_prediction[key] = torch.cat(
                    [p[key].to(self.device) for p in all_predictions]
                )
        # print(final_prediction)
        # print(features)
        return final_prediction, features

    def question_answering(self, request):
        """
        Span-based question answering for a given question and context.
        We expect the input to use the (question, context) format for the text pairs.

        Args:
          request: the prediction request

        """

        def decode(
            start_: np.ndarray,
            end_: np.ndarray,
            topk: int,
            max_answer_len: int,
            undesired_tokens_: np.ndarray,
        ) -> Tuple:
            """
                Take the output of any :obj:`ModelForQuestionAnswering` and will generate probabilities
                for each span to be the actual answer.

                In addition, it filters out some unwanted/impossible cases like answer len being greater
                than max_answer_len or answer end position being before the starting position. The method
                supports output the k-best answer through the topk argument.

                Args:
                    start_ (:obj:`np.ndarray`): Individual start probabilities for each token.
                    end (:obj:`np.ndarray`): Individual end_ probabilities for each token.
                    topk (:obj:`int`): Indicates how many possible answer span(s) to extract from the model output.
                    max_answer_len (:obj:`int`): Maximum size of the answer to extract from the model's output.
                    undesired_tokens_ (:obj:`np.ndarray`): Mask determining tokens that can be part of the answer
                """
            # Ensure we have batch axis
            if start_.ndim == 1:
                start_ = start_[None]

            if end_.ndim == 1:
                end_ = end_[None]

            # Compute the score of each tuple(start_, end_) to be the real answer
            outer = np.matmul(np.expand_dims(start_, -1), np.expand_dims(end_, 1))

            # Remove candidate with end_ < start_ and end_ - start_ > max_answer_len
            candidates = np.tril(np.triu(outer), max_answer_len - 1)

            #  Inspired by Chen & al. (https://github.com/facebookresearch/DrQA)
            scores_flat = candidates.flatten()
            if topk == 1:
                idx_sort = [np.argmax(scores_flat)]
            elif len(scores_flat) < topk:
                idx_sort = np.argsort(-scores_flat)
            else:
                idx = np.argpartition(-scores_flat, topk)[0:topk]
                idx_sort = idx[np.argsort(-scores_flat[idx])]

            starts_, ends_ = np.unravel_index(idx_sort, candidates.shape)[1:]
            desired_spans = np.isin(starts_, undesired_tokens_.nonzero()) & np.isin(
                ends_, undesired_tokens_.nonzero()
            )
            starts_ = starts_[desired_spans]
            ends_ = ends_[desired_spans]
            scores_ = candidates[0, starts_, ends_]

            return starts_, ends_, scores_

        predictions, features = self._predict(request)
        # print(predictions, request)
        task_outputs = {"answers": []}
        for idx, (start, end, (_, context)) in enumerate(
            zip(predictions["start_logits"], predictions["end_logits"], request)
        ):
            start = start.cpu().detach().numpy()
            end = end.cpu().detach().numpy()
            # Ensure padded tokens & question tokens cannot belong to the set of candidate answers.
            question_tokens = np.abs(
                np.array([s != 1 for s in features.sequence_ids(idx)]) - 1
            )
            # Unmask CLS token for 'no answer'
            question_tokens[0] = 1
            undesired_tokens = question_tokens & features["attention_mask"][idx].numpy()

            # Generate mask
            undesired_tokens_mask = undesired_tokens == 0.0

            # Make sure non-context indexes in the tensor cannot contribute to the softmax
            start = np.where(undesired_tokens_mask, -10000.0, start)
            end = np.where(undesired_tokens_mask, -10000.0, end)

            start = np.exp(
                start - np.log(np.sum(np.exp(start), axis=-1, keepdims=True))
            )
            end = np.exp(end - np.log(np.sum(np.exp(end), axis=-1, keepdims=True)))

            # print(start, end)

            # Get score for 'no answer' then mask for decoding step (CLS token
            no_answer_score = (start[0] * end[0]).item()
            start[0] = end[0] = 0.0

            starts, ends, scores = decode(start, end, 1, 128, undesired_tokens)
            enc = features[idx]
            answers = [
                {
                    "score": score.item(),
                    "start": enc.word_to_chars(enc.token_to_word(s), sequence_index=1)[
                        0
                    ],
                    "end": enc.word_to_chars(enc.token_to_word(e), sequence_index=1)[1],
                    "answer": context[
                        enc.word_to_chars(enc.token_to_word(s), sequence_index=1)[
                            0
                        ] : enc.word_to_chars(enc.token_to_word(e), sequence_index=1)[1]
                    ].lower(),
                }
                for s, e, score in zip(starts, ends, scores)
            ]
            answers.append(
                {"score": no_answer_score, "start": 0, "end": 0, "answer": ""}
            )
            answers = sorted(answers, key=lambda x: x["score"], reverse=True)[:1]
            task_outputs["answers"].append(answers)
        return task_outputs


if __name__ == "__main__":
    base_model = "bert-base-uncased"
    adapter_model = "AdapterHub/bert-base-uncased-pf-squad_v2"
    model = BertAdapterModel.from_pretrained(base_model)
    tokenizer = AutoTokenizer.from_pretrained(base_model)
    adapter_name = model.load_adapter(adapter_model, source="hf")
    model.active_adapters = adapter_name

    question, context = (
        "Who patronized the monks in Italy?",
        "At Saint Evroul, a tradition of singing had developed and the choir achieved fame in Normandy. Under the Norman abbot Robert de Grantmesnil, several monks of Saint-Evroul fled to southern Italy, where they were patronised by Robert Guiscard and established a Latin monastery at Sant'Eufemia. There they continued the tradition of singing.",
    )
    reduce = InputReduction(
        model,
        tokenizer,
        request={"question": question, "context": context, "top_k": 5, "mode": "max"},
        method="scaled_attn",
    )
    resp = reduce.get_predictions()
    print(resp)

    # dataloader = PreprocessData("squad_v2", "squad_v2", save_data=False, save_path="../../")
    # outputs = list()
    # count = 0
    # for ex in tqdm(dataloader.processed_val_set()):
    #     # print(ex)
    #     result = dict()
    #     reduce = InputReduction(model,
    #                             tokenizer,
    #                             request={
    #                                 "question": ex["question"],
    #                                 "context": ex["context"],
    #                                 "top_k": 5,
    #                                 "mode": "max"
    #                             },
    #                             method="shap")
    #     resp = reduce.get_predictions()
    #     print(resp)
    #     count += 1
    #     if count == 2:
    #         break

src/calibration/baseline/input_reduction.py
UKPLab/CATfOOD
import ast
import os
import torch
import numpy as np
import pandas as pd
from collections import OrderedDict
from tqdm import tqdm
import math
import re
import spacy
from spacy.tokens import Doc

# from supar import Parser

from src.calibration.baseline import common, evaluate, utils

import warnings

# ignore warning
warnings.filterwarnings("ignore")


class CreateFeatures:
    def __init__(self, path):
        self.path = path
        self.cls_tokens = ["[CLS]", "<s>"]
        self.sep_tokens = ["[SEP]", "</s>"]
        self.count = 0
        self.nlp = spacy.load("en_core_web_sm")
        # self.supar = Parser.load('crf-con-en')

    def _load_data(self):

        self.tagger_info = utils.load_bin(f"{self.path}/pos_info.bin")
        # self.attributions = utils.load_bin(f"{self.path}/ig_info_with_cf.bin")
        self.attributions = self.build_file_dict(
            dataset="squad", split="addsent-dev", method="shap"
        )
        pred_df = pd.read_csv(f"{self.path}/outputs.csv")
        self.preds_info = pred_df.set_index("id").T.to_dict("dict")
        # self.preds_info = utils.read_json('outputs/addsent-dev_squad_predictions.json')

    def load_interp_info(self, file_dict, qas_id):
        return torch.load(file_dict[qas_id])

    def build_file_dict(self, dataset, split, method):
        # hard-coded path here: be careful
        # prefix = 'squad_sample-addsent_roberta-base'
        prefix = "{}_{}_roberta".format(dataset, split)
        fnames = os.listdir(os.path.join("explanations", method, prefix))
        qa_ids = [x.split("-", 1)[1].split(".")[0] for x in fnames]
        fullnames = [os.path.join("explanations", method, prefix, x) for x in fnames]
        return dict(zip(qa_ids, fullnames))

    def featurize(self):
        self._load_data()
        processed_instances = OrderedDict()
        c = 0
        # print(self.attributions)
        for q_id in tqdm(
            self.tagger_info, total=len(self.tagger_info), desc="transforming..."
        ):
            # print(q_id)
            tags = self.tagger_info[q_id]
            # print(tags)
            interp = self.load_interp_info(self.attributions, q_id)
            # print(interp)
            attr_len = len(interp["attribution"])
            interp["attribution"] = torch.rand(attr_len)
            # print(len(ast.literal_eval(attributions)["attributions"]))
            preds = self.preds_info[q_id]
            # print(preds)
            # print(attributions)
            # print(tags)
            # print(len(tags["words"]))
            # print(len(ast.literal_eval(attributions)))
            # print(attributions)
            # print(ast.literal_eval(preds["pred_text"])[0])
            processed_instances[q_id] = self.extract_feature_for_instance(
                interp, tags, preds
            )
            c += 1
            # if c==2:
            #     break
            # print(self.count)
        # print(processed_instances)
        utils.dump_to_bin(processed_instances, f"{self.path}/calibration_data.bin")

    def lemmatize_pos_tag(self, token_tags):
        """
        limit the amount of POS tags to reduce the number of features
        :param token_tags: tagged tokens
        :return:
        """
        tok, pos, tag = token_tags
        if tag == "NNS":
            tag = "NN"
        if tag == "NNPS":
            tag = "NNP"
        if tag.startswith("JJ"):
            tag = "JJ"
        if tag.startswith("RB"):
            tag = "RB"
        if tag.startswith("W"):
            tag = "W"
        if tag.startswith("PRP"):
            tag = "PRP"
        if tag.startswith("VB"):
            tag = "VB"
        if pos == "PUNCT":  # or tag in [":"]:
            tag = "PUNCT"
        # if tag == '$':
        #     tag = 'SYM'
        # if tag == "``" or tag == "''":
        #     tag = "AUX"
        return tok, pos, tag

    def extract_baseline_feature(self, preds):
        # print(preds)
        predictions = ast.literal_eval(preds["pred_text"])[0]
        # print(predictions)
        feat = common.IndexedFeature()
        # base features
        for rank, p in enumerate(predictions[:5]):
            feat.add(f"BASELINE_PROB_{rank}", p["score"])
        feat.add("BASELINE_CONTEXT_LENGTH", len(preds["context"]))
        feat.add("BASELINE_PRED_ANS_LENGTH", len(predictions[0]["answer"]))

        top_pred = predictions[0]["answer"]
        first_distinct_prob = 0
        for i, p in enumerate(predictions[1:]):
            overlapping = evaluate.f1_score(p["answer"], top_pred)
            if overlapping > 0:
                continue
            first_distinct_prob = p["score"]
            # print(i+1, top_pred, p['answer'], first_distinct_prob)
        feat.add("FIRST_DISTINCT_PROB", first_distinct_prob)

        return feat

    def source_of_token(self, idx, tok, context_start, ans_range):
        if tok in self.cls_tokens or tok in self.sep_tokens:
            return "S"
        if idx >= 1 and idx < context_start:
            return "Q"
        if idx >= ans_range[0] and idx <= ans_range[1]:
            # print(tok)
            return "A"
        return "C"

    def nan_check(self, result):
        for i, (key, value) in enumerate(result):
            if math.isnan(float(value)):
                result[i][1] = 0
        return result

    def count_syllables(self, word: str):
        return len(
            re.findall("(?!e$)[aeiouy]+", word, re.I)
            + re.findall("^[^aeiouy]*e$", word, re.I)
        )

    def handle_zero_division(self, n, d):
        return n / d if d else 0

    def worF(self, tokens):
        words = [x.lstrip() for x in tokens]
        spaces = [
            False if i == len(tokens) - 1 else tokens[i + 1][0] == " "
            for i in range(len(tokens))
        ]

        valid_idx = [i for i, w in enumerate(words) if len(w)]
        words = [words[i] for i in valid_idx]
        spaces = [spaces[i] for i in valid_idx]
        doc = Doc(self.nlp.vocab, words=words, spaces=spaces)
        processed_tokens = self.nlp(doc)
        n_sent, n_token = 0, 0
        token_list = []
        for sent in processed_tokens.sents:
            n_sent += 1
            for token in sent:
                token_list.append(token.lemma_.lower())
                n_token += 1

        to_SbFrQ_C = 0
        to_SbCDC_C = 0
        to_SbFrL_C = 0
        to_SbCDL_C = 0
        to_SbSBW_C = 0
        to_SbL1W_C = 0
        to_SbSBC_C = 0
        to_SbL1C_C = 0

        DB = pd.read_csv("./src/resources/SUBTLEXus.csv")
        DB.set_index("Word_lowercased", inplace=True, drop=True)
        for token in token_list:
            if token in DB.index:
                scores_for_this_token = list(DB.loc[token, :])
                for i, score in enumerate(scores_for_this_token):
                    scores_for_this_token[i] = (
                        0 if str(score) == "none" else scores_for_this_token[i]
                    )
                to_SbFrQ_C += float(scores_for_this_token[1])
                to_SbCDC_C += float(scores_for_this_token[2])
                to_SbFrL_C += float(scores_for_this_token[3])
                to_SbCDL_C += float(scores_for_this_token[4])
                to_SbSBW_C += float(scores_for_this_token[5])
                to_SbL1W_C += float(scores_for_this_token[6])
                to_SbSBC_C += float(scores_for_this_token[7])
                to_SbL1C_C += float(scores_for_this_token[8])

        result = {
            "to_SbFrQ_C": to_SbFrQ_C,
            "as_SbFrQ_C": to_SbFrQ_C / n_sent,
            "at_SbFrQ_C": to_SbFrQ_C / n_token,
            "to_SbCDC_C": to_SbCDC_C,
            "as_SbCDC_C": to_SbCDC_C / n_sent,
            "at_SbCDC_C": to_SbCDC_C / n_token,
            "to_SbFrL_C": to_SbFrL_C,
            "as_SbFrL_C": to_SbFrL_C / n_sent,
            "at_SbFrL_C": to_SbFrL_C / n_token,
            "to_SbCDL_C": to_SbCDL_C,
            "as_SbCDL_C": to_SbCDL_C / n_sent,
            "at_SbCDL_C": to_SbCDL_C / n_token,
            "to_SbSBW_C": to_SbSBW_C,
            "as_SbSBW_C": to_SbSBW_C / n_sent,
            "at_SbSBW_C": to_SbSBW_C / n_token,
            "to_SbL1W_C": to_SbL1W_C,
            "as_SbL1W_C": to_SbL1W_C / n_sent,
            "at_SbL1W_C": to_SbL1W_C / n_token,
            "to_SbSBC_C": to_SbSBC_C,
            "as_SbSBC_C": to_SbSBC_C / n_sent,
            "at_SbSBC_C": to_SbSBC_C / n_token,
            "to_SbL1C_C": to_SbL1C_C,
            "as_SbL1C_C": to_SbL1C_C / n_sent,
            "at_SbL1C_C": to_SbL1C_C / n_token,
        }
        return result

    def extract_worf_features(self, tokens, ans_range):
        feat = common.IndexedFeature()
        context_start = tokens.index(self.sep_tokens[1])
        question_tokens = tokens[1:context_start]
        context_tokens = tokens[context_start + 2 : -1]

        q_ents = self.worF(question_tokens)
        c_ents = self.worF(context_tokens)
        # print(q_ents)
        # print(c_ents)
        for k, v in q_ents.items():
            feat.add_new(f"WOR_{k}_Q", v)
        for k, v in c_ents.items():
            feat.add_new(f"WOR_{k}_C", v)
        return feat

    def ttrF(self, tokens):
        words = [x.lstrip() for x in tokens]
        spaces = [
            False if i == len(tokens) - 1 else tokens[i + 1][0] == " "
            for i in range(len(tokens))
        ]

        valid_idx = [i for i, w in enumerate(words) if len(w)]
        words = [words[i] for i in valid_idx]
        spaces = [spaces[i] for i in valid_idx]
        doc = Doc(self.nlp.vocab, words=words, spaces=spaces)
        processed_tokens = self.nlp(doc)
        n_sent, n_token = 0, 0
        token_list = []
        for sent in processed_tokens.sents:
            n_sent += 1
            for token in sent:
                token_list.append(token.lemma_.lower())
                n_token += 1

        n_utoken = 1
        default_MTLD = 0.72
        MTLD_count = 0
        for token in token_list:
            if token_list.count(token) == 1:
                n_utoken += 1
            if float(n_utoken / n_token) >= 0.72:
                MTLD_count += 1

        result = {
            "SimpTTR_S": float(n_utoken / n_token),
            "CorrTTR_S": float(n_utoken / math.sqrt(2 * n_token)),
            "BiLoTTR_S": float(math.log(n_utoken) / math.log(n_token)),
            "UberTTR_S": float(
                self.handle_zero_division(
                    ((math.log(n_utoken)) ** 2), (math.log(n_token / n_utoken))
                )
            ),
            "MTLDTTR_S": float(MTLD_count),
        }
        return result

    def extract_ttrf_features(self, tokens, ans_range):
        feat = common.IndexedFeature()
        context_start = tokens.index(self.sep_tokens[1])
        question_tokens = tokens[1:context_start]
        context_tokens = tokens[context_start + 2 : -1]

        q_ents = self.ttrF(question_tokens)
        c_ents = self.ttrF(context_tokens)
        # print(q_ents)
        # print(c_ents)
        for k, v in q_ents.items():
            feat.add_new(f"TTR_{k}_Q", v)
        for k, v in c_ents.items():
            feat.add_new(f"TTR_{k}_C", v)
        return feat

    def varF(self, tokens):

        words = [x.lstrip() for x in tokens]
        spaces = [
            False if i == len(tokens) - 1 else tokens[i + 1][0] == " "
            for i in range(len(tokens))
        ]

        valid_idx = [i for i, w in enumerate(words) if len(w)]
        words = [words[i] for i in valid_idx]
        spaces = [spaces[i] for i in valid_idx]
        doc = Doc(self.nlp.vocab, words=words, spaces=spaces)
        processed_tokens = self.nlp(doc)
        n_sent, n_token = 0, 0
        for sent in processed_tokens.sents:
            n_sent += 1
            for token in sent:
                n_token += 1

        noun_list = []
        verb_list = []
        adje_list = []
        adve_list = []
        n_unoun = 0
        n_uverb = 0
        n_uadje = 0
        n_uadve = 0

        for token in processed_tokens:
            if token.pos_ == "NOUN":
                noun_list.append(token.lemma)
            if token.pos_ == "VERB":
                verb_list.append(token.lemma)
            if token.pos_ == "ADJ":
                adje_list.append(token.lemma)
            if token.pos_ == "ADV":
                adve_list.append(token.lemma)
        for noun in noun_list:
            if noun_list.count(noun) == 1:
                n_unoun += 1
        for verb in verb_list:
            if verb_list.count(verb) == 1:
                n_uverb += 1
        for adje in adje_list:
            if adje_list.count(adje) == 1:
                n_uadje += 1
        for adve in adve_list:
            if adve_list.count(adve) == 1:
                n_uadve += 1

        result = {
            "SimpNoV_S": float(self.handle_zero_division(n_unoun, (len(noun_list)))),
            "SquaNoV_S": float(
                self.handle_zero_division((n_unoun) ** 2, (len(noun_list)))
            ),
            "CorrNoV_S": float(
                self.handle_zero_division(n_unoun, (math.sqrt(2 * len(noun_list))))
            ),
            "SimpVeV_S": float(self.handle_zero_division(n_uverb, (len(verb_list)))),
            "SquaVeV_S": float(
                self.handle_zero_division((n_uverb) ** 2, (len(verb_list)))
            ),
            "CorrVeV_S": float(
                self.handle_zero_division(n_uverb, (math.sqrt(2 * len(verb_list))))
            ),
            "SimpAjV_S": float(self.handle_zero_division(n_uadje, (len(adje_list)))),
            "SquaAjV_S": float(
                self.handle_zero_division((n_uadje) ** 2, (len(adje_list)))
            ),
            "CorrAjV_S": float(
                self.handle_zero_division(n_uadje, (math.sqrt(2 * len(adje_list))))
            ),
            "SimpAvV_S": float(self.handle_zero_division(n_uadve, (len(adve_list)))),
            "SquaAvV_S": float(
                self.handle_zero_division((n_uadve) ** 2, (len(adve_list)))
            ),
            "CorrAvV_S": float(
                self.handle_zero_division(n_uadve, (math.sqrt(2 * len(adve_list))))
            ),
        }

        return result

    def extract_varf_features(self, tokens, ans_range):
        feat = common.IndexedFeature()
        context_start = tokens.index(self.sep_tokens[1])
        question_tokens = tokens[1:context_start]
        context_tokens = tokens[context_start + 2 : -1]

        q_ents = self.varF(question_tokens)
        c_ents = self.varF(context_tokens)
        # print(q_ents)
        # print(c_ents)
        for k, v in q_ents.items():
            feat.add_new(f"VAR_{k}_Q", v)
        for k, v in c_ents.items():
            feat.add_new(f"VAR_{k}_C", v)
        return feat

    def pos_tags(self, tokens):
        words = [x.lstrip() for x in tokens]
        spaces = [
            False if i == len(tokens) - 1 else tokens[i + 1][0] == " "
            for i in range(len(tokens))
        ]

        valid_idx = [i for i, w in enumerate(words) if len(w)]
        words = [words[i] for i in valid_idx]
        spaces = [spaces[i] for i in valid_idx]
        doc = Doc(self.nlp.vocab, words=words, spaces=spaces)
        processed_tokens = self.nlp(doc)
        n_sent, n_token = 0, 0
        for sent in processed_tokens.sents:
            n_sent += 1
            for token in sent:
                n_token += 1

        to_NoTag_C = 0
        to_VeTag_C = 0
        to_AjTag_C = 0
        to_AvTag_C = 0
        to_SuTag_C = 0
        to_CoTag_C = 0
        to_ContW_C = 0
        to_FuncW_C = 0

        for token in processed_tokens:
            if (
                token.pos_ == "NOUN"
                or token.pos_ == "VERB"
                or token.pos_ == "NUM"
                or token.pos_ == "ADJ"
                or token.pos_ == "ADV"
            ):
                to_ContW_C += 1
            else:
                to_FuncW_C += 1

            if token.pos_ == "NOUN":
                to_NoTag_C += 1
            if token.pos_ == "VERB":
                to_VeTag_C += 1
            if token.pos_ == "ADJ":
                to_AjTag_C += 1
            if token.pos_ == "ADV":
                to_AvTag_C += 1
            if token.pos_ == "SCONJ":
                to_SuTag_C += 1
            if token.pos_ == "CCONJ":
                to_CoTag_C += 1

        result = {
            "to_NoTag_C": float(to_NoTag_C),
            "as_NoTag_C": float(self.handle_zero_division(to_NoTag_C, n_sent)),
            "at_NoTag_C": float(self.handle_zero_division(to_NoTag_C, n_token)),
            "ra_NoAjT_C": float(self.handle_zero_division(to_NoTag_C, to_AjTag_C)),
            "ra_NoVeT_C": float(self.handle_zero_division(to_NoTag_C, to_VeTag_C)),
            "ra_NoAvT_C": float(self.handle_zero_division(to_NoTag_C, to_AvTag_C)),
            "ra_NoSuT_C": float(self.handle_zero_division(to_NoTag_C, to_SuTag_C)),
            "ra_NoCoT_C": float(self.handle_zero_division(to_NoTag_C, to_CoTag_C)),
            "to_VeTag_C": float(to_VeTag_C),
            "as_VeTag_C": float(self.handle_zero_division(to_VeTag_C, n_sent)),
            "at_VeTag_C": float(self.handle_zero_division(to_VeTag_C, n_token)),
            "ra_VeAjT_C": float(self.handle_zero_division(to_VeTag_C, to_AjTag_C)),
            "ra_VeNoT_C": float(self.handle_zero_division(to_VeTag_C, to_NoTag_C)),
            "ra_VeAvT_C": float(self.handle_zero_division(to_VeTag_C, to_AvTag_C)),
            "ra_VeSuT_C": float(self.handle_zero_division(to_VeTag_C, to_SuTag_C)),
            "ra_VeCoT_C": float(self.handle_zero_division(to_VeTag_C, to_CoTag_C)),
            "to_AjTag_C": float(to_AjTag_C),
            "as_AjTag_C": float(self.handle_zero_division(to_AjTag_C, n_sent)),
            "at_AjTag_C": float(self.handle_zero_division(to_AjTag_C, n_token)),
            "ra_AjNoT_C": float(self.handle_zero_division(to_AjTag_C, to_NoTag_C)),
            "ra_AjVeT_C": float(self.handle_zero_division(to_AjTag_C, to_VeTag_C)),
            "ra_AjAvT_C": float(self.handle_zero_division(to_AjTag_C, to_AvTag_C)),
            "ra_AjSuT_C": float(self.handle_zero_division(to_AjTag_C, to_SuTag_C)),
            "ra_AjCoT_C": float(self.handle_zero_division(to_AjTag_C, to_CoTag_C)),
            "to_AvTag_C": float(to_AvTag_C),
            "as_AvTag_C": float(self.handle_zero_division(to_AvTag_C, n_sent)),
            "at_AvTag_C": float(self.handle_zero_division(to_AvTag_C, n_token)),
            "ra_AvAjT_C": float(self.handle_zero_division(to_AvTag_C, to_AjTag_C)),
            "ra_AvNoT_C": float(self.handle_zero_division(to_AvTag_C, to_NoTag_C)),
            "ra_AvVeT_C": float(self.handle_zero_division(to_AvTag_C, to_VeTag_C)),
            "ra_AvSuT_C": float(self.handle_zero_division(to_AvTag_C, to_SuTag_C)),
            "ra_AvCoT_C": float(self.handle_zero_division(to_AvTag_C, to_CoTag_C)),
            "to_SuTag_C": float(to_SuTag_C),
            "as_SuTag_C": float(self.handle_zero_division(to_SuTag_C, n_sent)),
            "at_SuTag_C": float(self.handle_zero_division(to_SuTag_C, n_token)),
            "ra_SuAjT_C": float(self.handle_zero_division(to_SuTag_C, to_AjTag_C)),
            "ra_SuNoT_C": float(self.handle_zero_division(to_SuTag_C, to_NoTag_C)),
            "ra_SuVeT_C": float(self.handle_zero_division(to_SuTag_C, to_VeTag_C)),
            "ra_SuAvT_C": float(self.handle_zero_division(to_SuTag_C, to_AvTag_C)),
            "ra_SuCoT_C": float(self.handle_zero_division(to_SuTag_C, to_CoTag_C)),
            "to_CoTag_C": float(to_CoTag_C),
            "as_CoTag_C": float(self.handle_zero_division(to_CoTag_C, n_sent)),
            "at_CoTag_C": float(self.handle_zero_division(to_CoTag_C, n_token)),
            "ra_CoAjT_C": float(self.handle_zero_division(to_CoTag_C, to_AjTag_C)),
            "ra_CoNoT_C": float(self.handle_zero_division(to_CoTag_C, to_NoTag_C)),
            "ra_CoVeT_C": float(self.handle_zero_division(to_CoTag_C, to_VeTag_C)),
            "ra_CoAvT_C": float(self.handle_zero_division(to_CoTag_C, to_AvTag_C)),
            "ra_CoSuT_C": float(self.handle_zero_division(to_CoTag_C, to_SuTag_C)),
            "to_ContW_C": float(to_ContW_C),
            "as_ContW_C": float(self.handle_zero_division(to_ContW_C, n_sent)),
            "at_ContW_C": float(self.handle_zero_division(to_ContW_C, n_token)),
            "to_FuncW_C": float(to_FuncW_C),
            "as_FuncW_C": float(self.handle_zero_division(to_FuncW_C, n_sent)),
            "at_FuncW_C": float(self.handle_zero_division(to_FuncW_C, n_token)),
            "ra_CoFuW_C": float(self.handle_zero_division(to_ContW_C, to_FuncW_C)),
        }
        return result

    def extract_pos(self, tokens, ans_range):
        feat = common.IndexedFeature()
        context_start = tokens.index(self.sep_tokens[1])
        question_tokens = tokens[1:context_start]
        context_tokens = tokens[context_start + 2 : -1]

        q_ents = self.pos_tags(question_tokens)
        c_ents = self.pos_tags(context_tokens)
        # print(q_ents)
        # print(c_ents)
        for k, v in q_ents.items():
            feat.add_new(f"_POS_{k}_Q", v)
        for k, v in c_ents.items():
            feat.add_new(f"_POS_{k}_C", v)
        return feat

    def shaf(self, tokens, source):
        words = [x.lstrip() for x in tokens]
        spaces = [
            False if i == len(tokens) - 1 else tokens[i + 1][0] == " "
            for i in range(len(tokens))
        ]

        valid_idx = [i for i, w in enumerate(words) if len(w)]
        words = [words[i] for i in valid_idx]
        spaces = [spaces[i] for i in valid_idx]
        doc = Doc(self.nlp.vocab, words=words, spaces=spaces)
        processed_tokens = self.nlp(doc)
        n_sent, n_token = 0, 0
        for sent in processed_tokens.sents:
            n_sent += 1
            for token in sent:
                n_token += 1

        total_count_char = len(words)
        total_count_tokn = n_token
        total_count_syll = 0
        for token in words:
            total_count_syll += self.count_syllables(token)
        result = [
            (f"SHA_TokSenM_S_{source}", float(n_token * n_sent)),
            (f"SHA_TokSenS_S_{source}", float(math.sqrt(n_token * n_sent))),
            (
                f"SHA_TokSenL_S_{source}",
                float(self.handle_zero_division(math.log(n_token), math.log(n_sent))),
            ),
            (
                f"SHA_as_Token_C_{source}",
                float(self.handle_zero_division(total_count_tokn, n_sent)),
            ),
            (
                f"SHA_as_Sylla_C_{source}",
                float(self.handle_zero_division(total_count_syll, n_sent)),
            ),
            (f"SHA_at_Sylla_C_{source}", float(total_count_syll / n_token)),
            (
                f"SHA_as_Chara_C_{source}",
                float(self.handle_zero_division(total_count_char, n_sent)),
            ),
            (f"SHA_at_Chara_C_{source}", float(total_count_char / n_token)),
        ]
        # result = self.nan_check(result)
        return result

    def psyF(self, tokens, source):
        words = [x.lstrip() for x in tokens]
        spaces = [
            False if i == len(tokens) - 1 else tokens[i + 1][0] == " "
            for i in range(len(tokens))
        ]

        valid_idx = [i for i, w in enumerate(words) if len(w)]
        words = [words[i] for i in valid_idx]
        spaces = [spaces[i] for i in valid_idx]
        doc = Doc(self.nlp.vocab, words=words, spaces=spaces)
        processed_tokens = self.nlp(doc)
        n_sent, n_token = 0, 0
        for sent in processed_tokens.sents:
            n_sent += 1
            for token in sent:
                n_token += 1

        to_AAKuW_C = 0
        to_AAKuL_C = 0
        to_AABiL_C = 0
        to_AABrL_C = 0
        to_AACoL_C = 0

        DB = pd.read_csv("./src/resources/AoAKuperman.csv")
        DB.set_index("Word", inplace=True, drop=True)
        for token in words:
            if token in DB.index:
                scores_for_this_token = list(DB.loc[token, :])
                for i, score in enumerate(scores_for_this_token):
                    scores_for_this_token[i] = (
                        0 if str(score) == "none" else scores_for_this_token[i]
                    )
                to_AAKuW_C += float(scores_for_this_token[7])
                to_AAKuL_C += float(scores_for_this_token[9])
                to_AABiL_C += float(scores_for_this_token[11])
                to_AABrL_C += float(scores_for_this_token[12])
                to_AACoL_C += float(scores_for_this_token[13])

        result = [
            (f"PSY_to_AAKuW_C_{source}", to_AAKuW_C),
            (f"PSY_as_AAKuW_C_{source}", to_AAKuW_C / n_sent),
            (f"PSY_at_AAKuW_C_{source}", to_AAKuW_C / n_token),
            (f"PSY_to_AAKuL_C_{source}", to_AAKuL_C),
            (f"PSY_as_AAKuL_C_{source}", to_AAKuL_C / n_sent),
            (f"PSY_at_AAKuL_C_{source}", to_AAKuL_C / n_token),
            (f"PSY_to_AABiL_C_{source}", to_AABiL_C),
            (f"PSY_as_AABiL_C_{source}", to_AABiL_C / n_sent),
            (f"PSY_at_AABiL_C_{source}", to_AABiL_C / n_token),
            (f"PSY_to_AABrL_C_{source}", to_AABrL_C),
            (f"PSY_as_AABrL_C_{source}", to_AABrL_C / n_sent),
            (f"PSY_at_AABrL_C_{source}", to_AABrL_C / n_token),
            (f"PSY_to_AACoL_C_{source}", to_AABrL_C),
            (f"PSY_as_AACoL_C_{source}", to_AABrL_C / n_sent),
            (f"PSY_at_AACoL_C_{source}", to_AABrL_C / n_token),
        ]
        return result

    def endf(self, tokens, source):

        words = [x.lstrip() for x in tokens]
        spaces = [
            False if i == len(tokens) - 1 else tokens[i + 1][0] == " "
            for i in range(len(tokens))
        ]

        valid_idx = [i for i, w in enumerate(words) if len(w)]
        words = [words[i] for i in valid_idx]
        spaces = [spaces[i] for i in valid_idx]
        doc = Doc(self.nlp.vocab, words=words, spaces=spaces)
        processed_tokens = self.nlp(doc)
        n_sent, n_token = 0, 0
        for sent in processed_tokens.sents:
            n_sent += 1
            for token in sent:
                n_token += 1

        to_EntiM_C = 0
        to_UEnti_C = 0
        ent_list = []
        unique_ent_list = []

        for ent in processed_tokens.ents:
            to_EntiM_C += 1
            ent_list.append(ent.text)

        for ent in ent_list:
            if ent_list.count(ent) == 1:
                to_UEnti_C += 1
                unique_ent_list.append(ent)

        result = [
            # total number of Entities Mentions counts
            (f"ENT_to_EntiM_{source}", to_EntiM_C),
            # average number of Entities Mentions counts per sentence
            (f"ENT_as_EntiM_{source}", to_EntiM_C / n_sent),
            # average number of Entities Mentions counts per token (word)
            (f"ENT_at_EntiM_{source}", to_EntiM_C / n_token),
            # unique ents...
            (f"ENT_to_UEnti_{source}", to_UEnti_C),
            (f"ENT_as_UEnti_{source}", to_UEnti_C / n_sent),
            (f"ENT_at_UEnti_{source}", to_UEnti_C / n_token),
        ]

        return result

    def extract_entity_features(self, tokens, ans_range):
        feat = common.IndexedFeature()
        context_start = tokens.index(self.sep_tokens[1])
        question_tokens = tokens[1:context_start]
        context_tokens = tokens[context_start + 2 : -1]

        q_ents = self.endf(question_tokens, source="Q")
        c_ents = self.endf(context_tokens, source="C")
        # print(q_ents)
        # print(c_ents)
        for f in q_ents:
            feat.add_new(f[0], f[1])
        for f in c_ents:
            feat.add_new(f[0], f[1])
        return feat

    def extract_shallow_features(self, tokens, ans_range):
        feat = common.IndexedFeature()
        context_start = tokens.index(self.sep_tokens[1])
        question_tokens = tokens[1:context_start]
        context_tokens = tokens[context_start + 2 : -1]

        q_ents = self.shaf(question_tokens, source="Q")
        c_ents = self.shaf(context_tokens, source="C")
        # print(q_ents)
        # print(c_ents)
        for f in q_ents:
            feat.add_new(f[0], f[1])
        for f in c_ents:
            feat.add_new(f[0], f[1])
        return feat

    def extract_psy_features(self, tokens, ans_range):
        feat = common.IndexedFeature()
        context_start = tokens.index(self.sep_tokens[1])
        question_tokens = tokens[1:context_start]
        context_tokens = tokens[context_start + 2 : -1]

        q_ents = self.psyF(question_tokens, source="Q")
        c_ents = self.psyF(context_tokens, source="C")
        # print(q_ents)
        # print(c_ents)
        for f in q_ents:
            feat.add_new(f[0], f[1])
        for f in c_ents:
            feat.add_new(f[0], f[1])
        return feat

    def extract_bow_feature(self, words, tags, ans_range):
        feat = common.IndexedFeature()
        context_start = words.index(self.sep_tokens[1])
        # print(words)
        # print(words[ans_range[0]: ans_range[1]])
        # # print(words[ans_range[1]])
        # print(context_start)
        # print(ans_range)
        for i, (i_token, i_pos, i_tag) in enumerate(tags):
            i_src = self.source_of_token(i, i_token, context_start, ans_range)
            if i_src == "Q" or i_src == "A" or i_src == "C":
                # print('BOW_{}_{}'.format(i_src, i_tag))
                feat.add("BOW_{}_{}".format(i_src, i_tag))
                feat.add("BOW_IN_{}".format(i_tag))

        return feat

    def merge_attribution_by_segments(self, attribution, segments):
        new_val = []
        for a, b in segments:
            new_val.append(np.sum(attribution[a:b], axis=0))
        importance = np.stack(new_val, axis=0)
        return importance

    def aggregate_token_attribution(self, attr, tags, polarity):
        # attr = ast.literal_eval(attr)
        # print(attr)
        # print(np.array(list(attr["attributions"])))
        # attri = [i[1] for i in attr["attributions"]]
        # print(attri)
        # attribution_val = np.array(attr["attributions"])
        attribution_val = attr["attribution"].numpy().copy()
        if polarity == "POS":
            attribution_val[attribution_val < 0] = 0
        elif polarity == "NEG":
            attribution_val[attribution_val > 0] = 0
        elif polarity == "NEU":
            pass
        else:
            raise RuntimeError("Invalid polarity")

        attribution_val = self.merge_attribution_by_segments(
            attribution_val, tags["segments"]
        )
        assert attribution_val.shape[0] == len(tags["segments"])
        # normalize
        attribution_val = attribution_val / np.sum(attribution_val)
        # print(attribution_val)
        return attribution_val

    def normalize_token_attr(self, feat, attributions, norm_method=None):
        if norm_method is None:
            return feat
        if norm_method == "all":
            sum_v = np.sum(attributions)
            for k in feat.data:
                feat.data[k] = feat.data[k] / sum_v if sum_v != 0 else 0
            return feat
        if norm_method == "counted":
            sum_v = sum(feat.data.values())
            for k in feat.data:
                feat.data[k] = feat.data[k] / sum_v if sum_v != 0 else 0
            return feat
        raise RuntimeError(norm_method)

    def extract_token_attr_feature_in_question(
        self, words, tags, attributions, include_punct=False
    ):
        context_start = words.index(self.sep_tokens[1])
        tags = tags[1:context_start]
        attributions = attributions[1:context_start]

        feat = common.IndexedFeature()
        unnorm = common.IndexedFeature()
        sum_v = 0
        for i, (token, pos, tag) in enumerate(tags):
            if tag:
                v = attributions[i]
                if pos == "PUNCT" and not include_punct:
                    continue
                feat.add("NORMED_TOK_Q_" + tag, v)
                unnorm.add("UNNORM_TOK_Q_" + tag, v)
                sum_v += v

        feat = self.normalize_token_attr(feat, attributions)
        feat.add_set(unnorm)
        feat.add("SUM_TOK_Q", sum_v)
        return feat

    def extract_token_attr_feature_in_context(
        self, words, tags, attributions, include_punct=False
    ):
        context_start = words.index(self.sep_tokens[1])
        tags = tags[context_start + 2 : -1]
        attributions = attributions[context_start + 2 : -1]

        feat = common.IndexedFeature()
        unnorm = common.IndexedFeature()
        sum_v = 0
        for i, (token, pos, tag) in enumerate(tags):
            if tag:
                v = attributions[i]
                if pos == "PUNCT" and not include_punct:
                    continue
                feat.add("NORMED_TOK_C_" + tag, v)
                feat.add("UNNORM_TOK_C_" + tag, v)
                sum_v += v

        feat = self.normalize_token_attr(feat, attributions)
        feat.add_set(unnorm)
        feat.add("SUM_TOK_C", sum_v)
        return feat

    def extract_token_attr_feature_in_input(
        self, words, tags, attributions, ans_range, include_punct=False
    ):
        feat = common.IndexedFeature()
        context_start = words.index(self.sep_tokens[1])
        normed_a_feat = common.IndexedFeature()
        unnormed_a_feat = common.IndexedFeature()
        for i, (token, pos, tag) in enumerate(tags):
            if tag:
                v = attributions[i]
                if pos == "PUNCT" and not include_punct:
                    continue
                feat.add("TOK_IN_" + tag, v)

                if i >= ans_range[0] and i <= ans_range[1]:
                    unnormed_a_feat.add("UNNORM_TOK_A_" + tag, v)
                    normed_a_feat.add("NORMED_TOK_A_" + tag, v)

        feat.add_set(self.normalize_token_attr(normed_a_feat, [], "counted"))
        feat.add_set(unnormed_a_feat)
        return feat

    def extract_token_attr_stats_in_input(self, words, tags, attributions, part):
        feat = common.IndexedFeature()
        context_start = words.index(self.sep_tokens[1])
        if part == "Q":
            tags = tags[1:context_start]
            attributions = attributions[1:context_start]
        if part == "C":
            tags = tags[context_start + 2 : -1]
            attributions = attributions[context_start + 2 : -1]

        feat.add("STAT_MEAN_" + part, attributions.mean())
        feat.add("STAT_STD_" + part, attributions.std())
        return feat

    def extract_polarity_feature(
        self,
        attr,
        tags,
        words,
        tags_for_tok,
        ans_index,
        polarity,
        include_basic=True,
        include_stats=False,
    ):
        named_feat = common.IndexedFeature()
        token_attribution = self.aggregate_token_attribution(attr, tags, polarity)
        # link_attribution = None
        # elif args.method in ['probe']:
        # link_attribution = self.aggregate_link_attribution(attr, polarity)
        # print(token_attribution.size, len(words))
        # print(words)
        # print(attr)
        #
        # if token_attribution.size != len(words):
        #     print(token_attribution.size, len(words))
        #     print(words)
        #     self.count += 1
        # assert token_attribution.size == len(words)
        # if link_attribution is not None:
        #     assert link_attribution.shape == (len(words), len(words))
        if include_basic:
            named_feat.add_set(
                self.extract_token_attr_feature_in_question(
                    words, tags_for_tok, token_attribution
                )
            )
            named_feat.add_set(
                self.extract_token_attr_feature_in_context(
                    words, tags_for_tok, token_attribution
                )
            )
            named_feat.add_set(
                self.extract_token_attr_feature_in_input(
                    words, tags_for_tok, token_attribution, ans_index
                )
            )

            # for entites TODO: change format later
            # named_feat.add_set(self.extract_token_attr_feature_in_question(
            #     words,
            #     ents_for_tok,
            #     token_attribution))
            # named_feat.add_set(self.extract_token_attr_feature_in_context(
            #     words,
            #     ents_for_tok,
            #     token_attribution))
            # named_feat.add_set(
            #     self.extract_token_attr_feature_in_input(
            #         words,
            #         ents_for_tok,
            #         token_attribution,
            #         ans_index))

            # if link_attribution is not None:
            #     named_feat.add_set(self.extract_link_attr_feature(tags_for_tok, link_attribution, ans_index))
        if include_stats:
            named_feat.add_set(
                self.extract_token_attr_stats_in_input(
                    words, tags_for_tok, token_attribution, "Q"
                )
            )
            named_feat.add_set(
                self.extract_token_attr_stats_in_input(
                    words, tags_for_tok, token_attribution, "C"
                )
            )
            named_feat.add_set(
                self.extract_token_attr_stats_in_input(
                    words, tags_for_tok, token_attribution, "IN"
                )
            )
        named_feat.add_prefix(polarity + "_")
        return named_feat

    def extract_feature_for_instance(self, attr, tags, preds):
        # print(self.tagger_info)
        # print(preds)
        predictions = ast.literal_eval(preds["pred_text"])[0]
        # print(predictions)
        # create labels
        pred_text = predictions[0]["answer"]
        # print(pred_text)
        gold_text = ast.literal_eval(preds["gold_text"])  # preds['gold_text']
        # print(gold_text)

        exact_match = evaluate.get_score(
            metric="exact_match", pred_text=pred_text, gold_text=gold_text
        )
        f1 = evaluate.get_score(metric="f1", pred_text=pred_text, gold_text=gold_text)
        calib_label = 1 if exact_match > 0 else 0
        print(calib_label)

        # baseline features
        named_feat = common.IndexedFeature()
        named_feat.add_set(self.extract_baseline_feature(preds))

        start_index, end_index = (
            ast.literal_eval(preds["answer_start"])[0],
            ast.literal_eval(preds["answer_end"])[0],
        )
        segments = tags["segments"]
        transformed_start_idx = 0
        while segments[transformed_start_idx + 1][0] <= start_index:
            transformed_start_idx += 1
        transformed_end_idx = 0
        while segments[transformed_end_idx][1] < (end_index + 1):
            transformed_end_idx += 1
        ans_range = (transformed_start_idx, transformed_end_idx)
        # syntactic features
        words, tags_for_tok = tags["words"], tags["tags"]
        # words = [w for w in words if w != " "]
        tags_for_tok = [self.lemmatize_pos_tag(x) for x in tags_for_tok]

        # named_feat.add_set(self.extract_entity_features(words, ans_range))
        # named_feat.add_set(self.extract_shallow_features(words, ans_range))
        # named_feat.add_set(self.extract_psy_features(words, ans_range))
        # named_feat.add_set(self.extract_varf_features(words, ans_range))
        # named_feat.add_set(self.extract_ttrf_features(words, ans_range))
        # named_feat.add_set(self.extract_worf_features(words, ans_range))
        named_feat.add_set(self.extract_bow_feature(words, tags_for_tok, ans_range))
        named_feat.add_set(
            self.extract_polarity_feature(
                attr, tags, words, tags_for_tok, ans_range, "POS"
            )
        )
        # print({'feature': named_feat, 'label': calib_label, 'f1_score': f1})
        return {"feature": named_feat, "label": calib_label, "f1_score": f1}


if __name__ == "__main__":
    # base_model = "roberta-base"
    # tokenizer = AutoTokenizer.from_pretrained(base_model)
    data_path = "./src/data/squad_adv_new"
    feat = CreateFeatures(path=data_path)
    feat.featurize()

src/calibration/baseline/make_calib_dataset.py
UKPLab/CATfOOD
from collections import OrderedDict
import argparse
import collections
from collections import OrderedDict
from itertools import chain
from itertools import islice

import numpy as np
import shap
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import auc
from sklearn.svm import SVC
from tqdm import tqdm

from src.calibration.baseline.common import IndexedFeature, FeatureVocab
from src.calibration.baseline.utils import load_bin


def _f1auc_score(x, y):
    x = np.ravel(x)
    y = np.ravel(y)
    desc_score_indices = np.argsort(x, kind="mergesort")[::-1]
    x = x[desc_score_indices]
    y = y[desc_score_indices]

    distinct_value_indices = np.where(np.diff(x, append=0))[0]
    # accumulate the true positives with decreasing threshold
    threshold_values = x[distinct_value_indices]

    threshold_f1 = np.array([np.mean(y[: (t + 1)]) for t in distinct_value_indices])
    return auc(threshold_values, threshold_f1)


def f1auc_score(score, f1):
    score = np.ravel(score)
    f1 = np.ravel(f1)
    sorted_idx = np.argsort(-score)
    score = score[sorted_idx]
    f1 = f1[sorted_idx]
    num_test = f1.size
    segment = min(1000, score.size - 1)
    T = np.arange(segment) + 1
    T = T / segment
    results = np.array([np.mean(f1[: int(num_test * t)]) for t in T])
    return np.mean(results)


def _parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_name", type=str, required=True)
    parser.add_argument("--dataset", type=str, required=True)
    parser.add_argument("--method", type=str, required=True)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--do_baseline", default=False, action="store_true")
    parser.add_argument("--dense_features", default=False, action="store_true")
    parser.add_argument("--do_maxprob", default=False, action="store_true")
    parser.add_argument("--do_bow", default=False, action="store_true")
    parser.add_argument("--rm_bow", default=False, action="store_true")
    parser.add_argument("--rm_link", default=False, action="store_true")
    parser.add_argument("--rm_baseline", default=False, action="store_true")
    parser.add_argument("--rm_func", default=False, action="store_true")
    parser.add_argument("--do_unnorm", default=False, action="store_true")
    parser.add_argument("--n_run", type=int, default=20)
    parser.add_argument("--train_size", type=int, default=500)
    parser.add_argument("--force_dev_size", type=int, default=0)
    parser.add_argument("--model", type=str, default="rf")
    parser.add_argument("--show_imp", default=False, action="store_true")
    parser.add_argument("--split", type=str, default=None)
    parser.add_argument("--arg_n_tree", type=int, default=100)
    parser.add_argument("--arg_max_depth", type=int, default=None)
    parser.add_argument("--quant_size", type=int, default=0)
    parser.add_argument("--quant_type", type=str, default="val")
    parser.add_argument("--include_cf", default=False, action="store_true")
    args = parser.parse_args()
    if args.split == None:
        args.split = "addsent-dev" if args.dataset == "squad" else "dev"
    return args


def f1_prob_curve(f1, score):
    sorted_idx = np.argsort(-score)
    score = score[sorted_idx]
    f1 = f1[sorted_idx]
    num_test = f1.size
    # T = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
    T = [0.25, 0.5, 0.75, 1.0]
    results = np.array([np.mean(f1[: int(num_test * t)]) for t in T])
    return results


def train_max_accuracy(x, y):
    x = x.flatten()
    best_acc = 0
    best_v = 0
    for v in x:
        p = x > v
        ac = np.sum(p == y) / y.size
        if ac > best_acc:
            best_acc = ac
            best_v = v
    # print("Best threshold: ", best_v, " acc: ", best_acc)
    return best_acc, best_v


def test_max_accuracy(x, y, v):
    x = x.flatten()
    p = x > v
    ac = np.sum(p == y) / y.size
    # convert true false to 1 0 in p
    preds = p.astype(int)
    # print("p", preds)
    return ac, preds


def feat_to_list(indexed_feat, vocab):
    val_feat = [0.0] * len(vocab)
    for f, v in indexed_feat.data.items():
        val_feat[vocab[f]] = v
    return val_feat


def make_np_dataset(args, indexed_data):
    vocab = FeatureVocab()
    for k, v in indexed_data.items():
        feat = v["feature"]
        for f in feat.data:
            # if f not in ["BOW_Q_XX", "BOW_A_XX"]:   #, "POS_NORMED_TOK_A_XX", "POS_NORMED_TOK_Q_XX"]:
            if not args.dense_features:
                if not f.startswith("REPR_"):
                    vocab.add(f)
            else:
                vocab.add(f)
    print("Total Num of Feature", len(vocab))
    print(vocab.get_names())
    y = np.array([v["label"] for v in indexed_data.values()])
    f1 = np.array([v["f1_score"] for v in indexed_data.values()])
    x = np.array([feat_to_list(v["feature"], vocab) for v in indexed_data.values()])

    for index, name in vocab.id_to_feat.items():
        if "LENGTH" in name:
            x[:, index] = x[:, index] / np.max(x[:, index])
    return x, y, f1, vocab


def train_test_split(X, Y, F1, permuted_idx=None, ratio=0.75):
    num_data = X.shape[0]
    perm_idx = np.random.permutation(num_data) if permuted_idx is None else permuted_idx
    num_train = int(num_data * ratio)
    train_idx = perm_idx[:num_train]
    dev_idx = perm_idx[num_train:]
    train_x, train_y, train_F1 = X[train_idx, :], Y[train_idx], F1[train_idx]
    dev_x, dev_y, dev_F1 = X[dev_idx, :], Y[dev_idx], F1[dev_idx]

    return train_x, train_y, dev_x, dev_y, train_F1, dev_F1


def apply_train_test_split(X, Y, F1, train_test_split, force_dev_size=0):
    train_idx, dev_idx = train_test_split

    # for hyper par search
    if force_dev_size:
        train_size = len(train_idx)
        dev_idx = train_idx[(train_size - force_dev_size) :]
        train_idx = train_idx[: (train_size - force_dev_size)]
    train_x, train_y, train_F1 = X[train_idx, :], Y[train_idx], F1[train_idx]
    dev_x, dev_y, dev_F1 = X[dev_idx, :], Y[dev_idx], F1[dev_idx]

    return train_x, train_y, dev_x, dev_y, train_F1, dev_F1


def interp_calibrator_model(cls, vocab):
    if isinstance(cls, LogisticRegression):
        feat_imp = cls.coef_.flatten()
    elif isinstance(cls, RandomForestClassifier):
        feat_imp = cls.feature_importances_
    elif isinstance(cls, GradientBoostingClassifier):
        feat_imp = cls.feature_importances_
    else:
        print("UnInterpretable model")
        return

    imp_idx = np.argsort(-np.abs(feat_imp))
    for i in imp_idx[:100]:
        print(vocab.get_word(i), feat_imp[i])


def proc_input_data(args, data):
    if args.dataset == "squad":
        data = OrderedDict([(k, v) for (k, v) in data.items() if "-" in k])
    new_data = OrderedDict()
    for qas_id, ex in data.items():
        new_feat = IndexedFeature()
        for f, val in ex["feature"].data.items():
            if args.do_unnorm and "NORMED" in f:
                continue
            if not args.do_unnorm and "UNNORM" in f:
                continue
            if "TOK_IN_" in f:
                continue
            if "BOW_IN_" in f:
                continue
            if not args.do_bow:
                if "BOW_C" in f:
                    continue
            if f.startswith("ENT"):
                continue
            new_feat.add(f, val)

        new_data[qas_id] = {
            "label": ex["label"],
            "f1_score": ex["f1_score"],
            "feature": new_feat,
        }
    data = new_data
    if args.rm_baseline:
        new_data = OrderedDict()
        for qas_id, ex in data.items():
            new_feat = IndexedFeature()
            for f, val in ex["feature"].data.items():
                if f.startswith("BASELINE") and not f.startswith("BASELINE_PROB_0"):
                    continue
                new_feat.add(f, val)
            new_data[qas_id] = {
                "label": ex["label"],
                "f1_score": ex["f1_score"],
                "feature": new_feat,
            }
        data = new_data
    if args.do_baseline:
        new_data = OrderedDict()
        for qas_id, ex in data.items():
            new_feat = IndexedFeature()
            for f, val in ex["feature"].data.items():
                if not f.startswith("BASELINE" or "REPR"):
                    continue
                # if not ('BASELINE' in f or 'REPR' in f):
                #     continue
                new_feat.add(f, val)
            new_data[qas_id] = {
                "label": ex["label"],
                "f1_score": ex["f1_score"],
                "feature": new_feat,
            }
        data = new_data
    if args.do_maxprob:
        new_data = OrderedDict()
        for qas_id, ex in data.items():
            new_feat = IndexedFeature()
            for f, val in ex["feature"].data.items():
                if not f.startswith("BASELINE_PROB_0"):
                    continue
                # if not 'REPR' in f:
                #     continue
                new_feat.add(f, val)
            new_data[qas_id] = {
                "label": ex["label"],
                "f1_score": ex["f1_score"],
                "feature": new_feat,
            }
        data = new_data
    if args.do_bow:
        new_data = OrderedDict()
        for qas_id, ex in data.items():
            new_feat = IndexedFeature()
            for f, val in ex["feature"].data.items():
                if not ("BASELINE" in f or "BOW" in f):
                    continue

                new_feat.add(f, val)
            new_data[qas_id] = {
                "label": ex["label"],
                "f1_score": ex["f1_score"],
                "feature": new_feat,
            }
        data = new_data
    if args.rm_bow:
        new_data = OrderedDict()
        for qas_id, ex in data.items():
            new_feat = IndexedFeature()
            for f, val in ex["feature"].data.items():
                if "BOW" in f:
                    continue
                new_feat.add(f, val)
            new_data[qas_id] = {
                "label": ex["label"],
                "f1_score": ex["f1_score"],
                "feature": new_feat,
            }
        data = new_data
    if args.rm_link:
        new_data = OrderedDict()
        for qas_id, ex in data.items():
            new_feat = IndexedFeature()
            for f, val in ex["feature"].data.items():
                if "LINK" in f and "AGG" not in f:
                    continue
                new_feat.add(f, val)
            new_data[qas_id] = {
                "label": ex["label"],
                "f1_score": ex["f1_score"],
                "feature": new_feat,
            }
        data = new_data
    # print(data)
    return data


def get_feature_importances(cls):
    if isinstance(cls, LogisticRegression):
        feat_imp = cls.coef_.flatten()
    elif isinstance(cls, RandomForestClassifier):
        feat_imp = cls.feature_importances_
    elif isinstance(cls, GradientBoostingClassifier):
        feat_imp = cls.feature_importances_
    else:
        return None
    return feat_imp


def selection_based_rf(train_x, train_y, dev_x, n_feat=80):
    clf = RandomForestClassifier(n_estimators=300).fit(train_x, train_y)
    feat_imp = clf.feature_importances_
    imp_idx = np.argsort(-np.abs(feat_imp))
    train_x = train_x[
        :, imp_idx[:n_feat],
    ]
    dev_x = dev_x[:, imp_idx[:n_feat]]
    return (
        RandomForestClassifier(n_estimators=300).fit(train_x, train_y),
        train_x,
        dev_x,
    )


def one_pass_exp(args, X, Y, F1, vocab, train_test_split):
    train_x, train_y, dev_x, dev_y, train_F1, dev_F1 = apply_train_test_split(
        X, Y, F1, train_test_split, force_dev_size=args.force_dev_size
    )
    # print(train_x.shape)
    # print(train_x)
    # print(train_x[0])

    majority_acc = max(np.sum(dev_y == 0), np.sum(dev_y == 1)) / dev_y.size

    if args.do_maxprob:
        train_acc, best_threshold = train_max_accuracy(train_x, train_y)
        # print("dev x", dev_x)
        # print("dev y", dev_y)
        dev_acc, preds = test_max_accuracy(dev_x, dev_y, best_threshold)
        dev_auc = f1auc_score(dev_x, dev_F1)
        _macro_ce = macro_ce(preds, dev_y, dev_x)
        # print("macro ce", _macro_ce)
        return (
            majority_acc,
            train_acc,
            dev_acc,
            dev_auc,
            f1_prob_curve(dev_F1, dev_x.flatten()),
            None,
            _macro_ce,
        )

    if args.model == "lr":
        clf = LogisticRegression(C=1000.0, max_iter=1000, solver="saga").fit(
            train_x, train_y
        )
    elif args.model == "rf":
        # print("in rf")
        # clf = KNeighborsClassifier(n_neighbors=10).fit(train_x, train_y)
        # clf = RandomForestClassifier(n_estimators=200, max_depth=5).fit(train_x, train_y)
        feature_names = [vocab.get_word(i) for i in range(len(dev_x[0]))]
        # print(feature_names)

        clf = RandomForestClassifier(
            n_estimators=args.arg_n_tree,
            # min_samples_split=3,
            max_depth=args.arg_max_depth,
        ).fit(train_x, train_y)

    elif args.model == "svm":
        clf = SVC(C=10.0).fit(train_x, train_y)
    elif args.model == "gdbt":
        clf = GradientBoostingClassifier(n_estimators=100, max_depth=10).fit(
            train_x, train_y
        )
    else:
        raise RuntimeError("Model not supported")
    train_pred = clf.predict(train_x)
    dev_pred = clf.predict(dev_x)

    train_acc = np.sum(train_pred == train_y) / train_pred.size
    dev_acc = np.sum(dev_pred == dev_y) / dev_pred.size
    dev_score = clf.predict_proba(dev_x)[:, 1]

    dev_auc = f1auc_score(dev_score, dev_F1)
    f1_curve = f1_prob_curve(dev_F1, clf.predict_proba(dev_x)[:, 1])

    _macro_ce = macro_ce(dev_pred, dev_y, clf.predict_proba(dev_x)[:, 1])
    if args.show_imp:
        outputs = []
        selected_indices = [i for i in range(len(dev_pred)) if dev_pred[i] == dev_y[i]]
        # Create Tree Explainer object that can calculate shap values
        explainer = shap.TreeExplainer(clf)
        shap_values_all = explainer.shap_values(dev_x)[1]
        shap.summary_plot(
            shap_values_all[selected_indices, :],
            dev_x[selected_indices, :],
            feature_names=feature_names,
        )

        for i, idx in tqdm(enumerate(selected_indices)):
            # Calculate Shap values
            choosen_instance = dev_x[idx]
            shap_values = explainer.shap_values(choosen_instance)[1]

            # identify importances
            # shap.initjs()
            # shap.force_plot(explainer.expected_value[1], shap_values[1], choosen_instance)
            # break

            # print(dev_x[incorrect_indices[0]])
            # X_incorrect = X[incorrect_indices, :]
            # feature_importances = clf.feature_importances_
            # print(feature_importances)
            # importances_incorrect = feature_importances[incorrect_indices]
            # print(importances_incorrect)

            # train features
            # imp = permutation_importance(clf, dev_x, dev_y, n_repeats=10, random_state=42, n_jobs=2)

            # Identify the misclassified samples
            # misclassified = dev_x[i]

            # Estimate the feature importance for the misclassified samples using permutation importance
            # imp = permutation_importance(
            #     clf, dev_x[idx].reshape(1, -1), dev_y[idx].reshape(1, -1), n_repeats=10, random_state=42, n_jobs=2
            # )
            #
            # # # Extract the feature importance scores
            # misclassified_importance = imp.importances_mean
            # # print(misclassified_importance)
            # # imp_idx = np.argsort(-np.abs(misclassified_importance))
            # # Sort the feature importances in descending order
            sorted_indices = np.argsort(shap_values)[::-1]
            importance = []
            for rank, j in enumerate(sorted_indices[:20]):
                importance.append((rank, vocab.get_word(j), shap_values[j]))

            output = collections.OrderedDict()
            output["idx"] = idx
            output["importance"] = importance
            outputs.append(output)

    return (
        majority_acc,
        train_acc,
        dev_acc,
        dev_auc,
        f1_curve,
        get_feature_importances(clf),
        _macro_ce,
    )


def macro_ce(preds, golds, probs):
    """Macro CE"""
    ice_pos, ice_neg = 0, 0
    num_pos, num_neg = 0, 0
    for i in range(len(preds)):
        pred = preds[i]
        pred_proba = probs[i]
        gold = golds[i]
        if pred == gold:
            ice_pos += 1 - pred_proba
            num_pos += 1
        else:
            ice_neg += pred_proba - 0
            num_neg += 1
    ice_pos /= num_pos
    ice_neg /= num_neg
    macro_ce = (ice_pos + ice_neg) / 2
    print("Macro CE: ", macro_ce)
    return macro_ce


def gen_predefined_train_test_splits(baseids, n_run, train_size):
    num_data = len(baseids)

    baseid_indexer = OrderedDict()
    for i, b in enumerate(baseids):
        if b in baseid_indexer:
            baseid_indexer[b].append(i)
        else:
            baseid_indexer[b] = [i]
    num_base = len(baseid_indexer)
    print("Number of Exs", num_data, "Number of Unique Base", num_base)

    predefined_permutations = [np.random.permutation(num_base) for _ in range(n_run)]
    baseid_groups = list(baseid_indexer.values())
    num_train = train_size

    splits = []
    for perm in predefined_permutations:
        flat_idx = list(chain(*[baseid_groups[i] for i in perm]))
        # flat_idx_cf = list(chain(*[baseid_groups_cf[i] for i in perm_cf]))
        train_index = flat_idx[:num_train]  # + flat_idx_cf[:num_train_cf]
        # train_index.append(flat_idx_cf[:num_train_cf])
        dev_index = flat_idx[num_train:]
        splits.append((train_index, dev_index))
    print("Train Size", len(splits[0][0]), "Dev Size", len(splits[0][1]))
    return splits


def prediction_direction_of_feat(X, Y, F1):
    clf = LogisticRegression(C=1000.0, max_iter=10000).fit(X, Y)
    feat_imp = clf.coef_.flatten()
    train_pred = clf.predict(X)
    train_acc = np.sum(train_pred == Y) / train_pred.size
    print("Linear ACC:", train_acc)
    return feat_imp


def quantify_colum(x, k=2, method="val"):
    if method == "val":
        interval = np.arange(k) / k
        vals = (np.max(x) - np.min(x)) * interval + np.min(x)
    if method == "percent":
        q = 100 * np.arange(k) / k
        vals = np.percentile(x, q)

    new_x = np.zeros_like(x)

    for i, v in enumerate(vals):
        new_x[x >= v] = i / k
    return new_x


def quantify_dataset(X, vocab, k, type):
    # print(X.shape)
    for i in range(X.shape[1]):
        fname = vocab.get_word(i)
        if "BASELINE" in fname:
            continue
        if "BOW" in fname:
            continue
        # print(fname)
        X[:, i] = quantify_colum(X[:, i], k=k, method=type)
    return X, vocab


def main():
    args = _parse_args()
    # print(args)
    np.random.seed(args.seed)
    # data = load_bin(f'./src/data/{args.dataset}/calib_data_{args.method}_{args.model_name}.bin')
    data = load_bin(
        f"./src/data/squad_adversarial/calibration_data_sc_attn_llama_context_rel.bin"
    )

    def take(n, iterable):
        """Return the first n items of the iterable as a list."""
        return list(islice(iterable, n))

    # data = take(1000, data.items())
    data = proc_input_data(args, dict(data))
    print("Total samples: ", len(data))

    X, Y, F1, vocab = make_np_dataset(args, data)
    if args.quant_size > 0:
        X, vocab = quantify_dataset(X, vocab, args.quant_size, args.quant_type)

    if args.dataset == "squad":
        # filter cf baseids
        cf_baseids = [sample for sample in data if "-cf-" in sample]
        cf_baseids_filtered = [
            sample.split("-")[0] for sample in cf_baseids if "turk" in sample
        ]
        print(len(cf_baseids_filtered))
        data_without_cf = [sample for sample in data if sample not in cf_baseids]
        print("Number of counterfactuals: ", len(cf_baseids))
        baseids = [
            k.split("-")[0]
            for k in data_without_cf
            if args.dataset == "squad" and "turk" in k
        ]

    else:
        # filter cf baseids
        # cf_baseids = [sample for sample in data if "-cf-" in sample]
        # print("Number of counterfactuals: ", len(cf_baseids))
        cf_baseids_filtered = [
            sample.split("-")[0] for sample in data if "-cf-" in sample
        ]
        print("Number of filtered counterfactuals: ", len(cf_baseids_filtered))
        data_without_cf = [sample for sample in data if "-cf-" not in sample]
        # print(data_without_cf)
        print("Number of original samples: ", len(data_without_cf))
        baseids = [k for k in data_without_cf]

    predefined_splits = gen_predefined_train_test_splits(
        baseids, args.n_run, args.train_size
    )
    agg_results = []
    for train_test_split in tqdm(
        predefined_splits, total=len(predefined_splits), desc="Running Exp..."
    ):
        if args.include_cf:
            indices = []
            train_set = train_test_split[0]
            test_set = train_test_split[1]
            for i, b in enumerate(train_set):
                if baseids[b] in cf_baseids_filtered:
                    indices.append(baseids[b])

            cf_indices = []
            for i, b in enumerate(data):
                # print(b)
                if "-cf-" in b and b.split("-")[0] in indices:
                    cf_indices.append(i)
            # print(cf_indices)
            train_set.extend(cf_indices)

            np.random.shuffle(train_set)
            augmented_train_set = (train_set, test_set)
            print(len(train_set), len(test_set))
        else:
            augmented_train_set = train_test_split
        # print(len(augmented_train_set[0]), len(augmented_train_set[1]))
        agg_results.append(one_pass_exp(args, X, Y, F1, vocab, augmented_train_set))

    # print(agg_results)
    agg_base_acc = np.array([x[0] for x in agg_results])
    agg_train_acc = np.array([x[1] for x in agg_results])
    agg_dev_acc = np.array([x[2] for x in agg_results])
    agg_auc = np.array([x[3] for x in agg_results])
    agg_f1_curve = np.array([x[4] for x in agg_results]).mean(axis=0)
    agg_macro_ce = np.array([x[6] for x in agg_results]).mean()
    print("AVG MACRO CE {:.3f}".format(agg_macro_ce))

    print("AVG MAJORITY ACC {:.3f}".format(agg_base_acc.mean()))
    print("AVG TRAIN ACC {:.3f}".format(agg_train_acc.mean()))
    print(
        "AVG DEV ACC: {:.3f} +/- {:.3f}, AUC: {:.3f}, MAX: {:.3f}, MIN: {:.3f}".format(
            agg_dev_acc.mean(),
            agg_dev_acc.std(),
            agg_auc.mean(),
            agg_dev_acc.max(),
            agg_dev_acc.min(),
        )
    )

    # print numbers for copy paste
    exp_numbers = [agg_base_acc.mean(), agg_dev_acc.mean(), agg_auc.mean()] + [
        agg_macro_ce
    ]  # + agg_f1_curve.tolist()
    print(",".join(["{:.3f}".format(x) for x in exp_numbers]))
    exp_numbers = exp_numbers[1:]
    print(",".join(["{:.1f}".format(a * 100) for a in exp_numbers]))
    if agg_results[0][5] is not None and args.show_imp:
        agg_feat_imp = np.array([x[5] for x in agg_results])
        agg_feat_imp = np.mean(agg_feat_imp, axis=0)
        imp_idx = np.argsort(-np.abs(agg_feat_imp))
        for rank, i in enumerate(imp_idx[:100]):
            print(rank, vocab.get_word(i), agg_feat_imp[i])


if __name__ == "__main__":
    main()

src/calibration/baseline/modelling.py
UKPLab/CATfOOD
import string
from typing import List, Tuple
from collections import OrderedDict

import torch
import spacy
from spacy.tokens import Doc
from tqdm import tqdm

from transformers import (
    AutoTokenizer,
    # RobertaAdapterModel,
)

from transformers.models.gpt2.tokenization_gpt2 import bytes_to_unicode
from src.calibration.baseline import dataloader
import utils

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"
# BASE_PATH = "/home/sachdeva/projects/ukp/exp_calibration/"


class NER:
    def __init__(self, tokenizer):
        self.tokenizer = tokenizer
        self.cls_token = self.tokenizer.cls_token
        self.eos_token = self.tokenizer.eos_token

    def encode(
        self, inputs: list = None, add_special_tokens: bool = True, return_tensors=None
    ):
        """
        Encode inputs using the model tokenizer
        Args:
            inputs: question, context pair as a list
            add_special_tokens: where to add CLS, SEP tokens
            return_tensors: whether to return tensors
        Return:
            tokenized inputs
        """
        return self.tokenizer(
            inputs,
            add_special_tokens=add_special_tokens,
            return_tensors=return_tensors,
            padding=True,
            truncation=True,
            max_length=512,
        )

    def decode(self, input_ids: torch.Tensor, skip_special_tokens: bool) -> List[str]:
        """
        Decode received input_ids into a list of word tokens.
        Args:
            input_ids (torch.Tensor): Input ids representing
            word tokens for a sentence/document.
        """
        return self.tokenizer.convert_ids_to_tokens(
            input_ids[0], skip_special_tokens=skip_special_tokens
        )

    def _bpe_decode(self, tokens: List[str],) -> Tuple[List[str], List]:

        byte_encoder = bytes_to_unicode()
        byte_decoder = {v: k for k, v in byte_encoder.items()}
        decoded_each_tok = [
            bytearray([byte_decoder[c] for c in t]).decode(
                encoding="utf-8", errors="replace"
            )
            for t in tokens
        ]

        end_points = []
        force_break = False
        for idx, token in enumerate(decoded_each_tok):
            # special token, punctuation, alphanumeric
            if (
                token in self.tokenizer.all_special_tokens
                or token in string.punctuation
                or not any([x.isalnum() for x in token.lstrip()])
                or token.lstrip == "'s"
            ):
                end_points.append(idx)
                force_break = True
                continue

            if force_break:
                end_points.append(idx)
                force_break = False
                continue

            if token[0] == " ":
                tokens[idx] = token[:]
                end_points.append(idx)

        end_points.append(len(tokens))

        segments = []
        for i in range(1, len(end_points)):
            if end_points[i - 1] == end_points[i]:
                continue
            segments.append((end_points[i - 1], end_points[i]))

        filtered_tokens = []
        for s0, s1 in segments:
            filtered_tokens.append("".join(decoded_each_tok[s0:s1]))

        return filtered_tokens, segments

    def process_input(self, request):
        inputs = [[request["question"], request["context"]]]
        encoded_inputs = self.encode(
            inputs, add_special_tokens=True, return_tensors="pt"
        )

        decoded_text = self.decode(
            encoded_inputs["input_ids"], skip_special_tokens=False
        )
        filtered_tokens, segments = self._bpe_decode(decoded_text)
        return filtered_tokens, segments

    def assign_entity(self, tokens, nlp):
        words = [x.lstrip() for x in tokens]
        spaces = [
            False if i == len(tokens) - 1 else tokens[i + 1][0] == " "
            for i in range(len(tokens))
        ]

        valid_idx = [i for i, w in enumerate(words) if len(w)]
        words = [words[i] for i in valid_idx]
        spaces = [spaces[i] for i in valid_idx]
        doc = Doc(nlp.vocab, words=words, spaces=spaces)
        processed_tokens = nlp(doc)

        ent_info = [("", "NULL", "NULL")] * len(tokens)
        for i, token in zip(valid_idx, processed_tokens):
            ent_info[i] = (token.text, token.ent_iob_, token.ent_type_)
        return ent_info

    def tag_instance(self, request, nlp):
        words, segments = self.process_input(request)
        # tokens = list(filter(None, tokens))

        context_start = words.index(self.tokenizer.eos_token)
        question_tokens = words[1:context_start]
        context_tokens = words[context_start + 2 : -1]
        question_ent_info = self.assign_entity(question_tokens, nlp)
        context_ent_info = self.assign_entity(context_tokens, nlp)
        ent_info = (
            [(self.cls_token, "SOS", "SOS")]
            + question_ent_info
            + [(self.eos_token, "EOS", "EOS"), (self.eos_token, "EOS", "EOS")]
            + context_ent_info
            + [(self.eos_token, "EOS", "EOS")]
        )

        assert len(ent_info) == len(words), "entities and words not equal"
        instance_info = {"words": words, "segments": segments, "ents": ent_info}
        # print(instance_info)
        return instance_info


if __name__ == "__main__":
    # model_path = "./src/checkpoints/roberta_squad"
    model_path = "./roberta-squad"
    # model = RobertaModel.from_pretrained(model_path)
    tokenizer = AutoTokenizer.from_pretrained("roberta-base")

    # data = dataloader.PreprocessData("squad_adversarial",
    #                                  "AddSent",
    #                                  save_data=False,
    #                                  save_path="../../")
    # outputs = list()
    # ner_ = NER(tokenizer=tokenizer)
    # nlp = spacy.load("en_core_web_sm")
    # c = 0
    # processed_instances = OrderedDict()
    # # process OOD data
    # for ex in tqdm(data.processed_val_set()):
    #     # ex = remove_white_space(ex)
    #     try:
    #         # if ex["id"] == "56e1239acd28a01900c67641":
    #         #     print(ex)
    #         ent_info = ner_.tag_instance(
    #                         request={
    #                                     "id": ex["id"],
    #                                     "question": ex["question"],
    #                                     "context": ex["context"],
    #                                  },
    #                         nlp=nlp
    #                     )
    #         processed_instances[ex["id"]] = ent_info
    #         c += 1
    #     except Exception as e:
    #         print(f"Unable to get entities: {e}")
    #         print(ex)

    # # process counterfactuals
    # cf_path = BASE_PATH + "src/data/squad_adversarial/rag_counterfactuals_turk0_last.jsonl"
    # c = 0
    # with jsonlines.open(cf_path) as reader:
    #     for ex in tqdm(reader):
    #         try:
    #             ent_info = ner_.tag_instance(
    #                 request={
    #                     "id": ex["id"],
    #                     "question": ex["question"],
    #                     "context": ex["context"],
    #                 },
    #                 nlp=nlp
    #             )
    #             processed_instances[ex["id"].replace("_", "-cf-")] = ent_info
    #             c += 1
    #         except Exception as e:
    #             print(f"Unable to get cf attributions: {e}")
    #             print(ex)
    #     print(f"Processed {c} instances of counterfactual data")

    # utils.dump_to_bin(processed_instances,
    #                   BASE_PATH + "src/data/squad_adv_new/ent_info.bin")

    ### Trivia QA
    def remove_white_space(example):
        example["question_text"] = " ".join(example["question_text"].split())
        example["context_text"] = " ".join(example["context_text"].split())
        return example

    data = dataloader.get_dev_examples(BASE_PATH + "src/data", "dev_hotpot.json")
    # data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/NewsQA.jsonl")
    # data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/BioASQ-dev.jsonl")
    # data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/NaturalQuestionsShort.jsonl")

    ner_ = NER(tokenizer=tokenizer)
    nlp = spacy.load("en_core_web_sm")
    c = 0
    processed_instances = OrderedDict()
    # process OOD data
    for ex in tqdm(data):
        ex = remove_white_space(ex)
        try:
            # if ex["id"] == "56e1239acd28a01900c67641":
            #     print(ex)
            ent_info = ner_.tag_instance(
                request={
                    "id": ex["qas_id"],
                    "question": ex["question_text"],
                    "context": ex["context_text"],
                },
                nlp=nlp,
            )
            processed_instances[ex["qas_id"]] = ent_info
            c += 1
        except Exception as e:
            print(f"Unable to get entities: {e}")
            print(ex)

    print(f"Processed {c} instances of original data")

    # # process counterfactuals
    # cf_path = BASE_PATH + "src/data/natural_questions/rag_counterfactuals_turk0.jsonl"
    # c = 0
    # with jsonlines.open(cf_path) as reader:
    #     for ex in tqdm(reader):
    #         try:
    #             ent_info = ner_.tag_instance(
    #                 request={
    #                     "id": ex["id"],
    #                     "question": ex["question"],
    #                     "context": ex["context"],
    #                 },
    #                 nlp=nlp
    #             )
    #             processed_instances[ex["id"].replace("_", "-cf-")] = ent_info
    #             c += 1
    #         except Exception as e:
    #             print(f"Unable to get cf attributions: {e}")
    #             print(ex)
    #     print(f"Processed {c} instances of counterfactual data")
    #
    utils.dump_to_bin(
        processed_instances, BASE_PATH + "src/data/hotpot_qa/ent_info.bin"
    )

    print(f"Saved instances: {c}")

src/calibration/baseline/ner.py
UKPLab/CATfOOD
import string
from typing import List, Tuple
from collections import OrderedDict

import torch
import spacy
from spacy.tokens import Doc
from tqdm import tqdm

from transformers import (
    AutoTokenizer,
    # RobertaAdapterModel,
)

from transformers.models.gpt2.tokenization_gpt2 import bytes_to_unicode
from src.calibration.baseline import dataloader
import utils

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


class POSTagger:
    def __init__(self, tokenizer):
        self.tokenizer = tokenizer
        self.cls_token = self.tokenizer.cls_token
        self.eos_token = self.tokenizer.eos_token

    def encode(
        self, inputs: list = None, add_special_tokens: bool = True, return_tensors=None
    ):
        """
        Encode inputs using the model tokenizer
        Args:
            inputs: question, context pair as a list
            add_special_tokens: where to add CLS, SEP tokens
            return_tensors: whether to return tensors
        Return:
            tokenized inputs
        """
        return self.tokenizer(
            inputs,
            add_special_tokens=add_special_tokens,
            return_tensors=return_tensors,
            padding=True,
            truncation=True,
            max_length=512,
        )

    def decode(self, input_ids: torch.Tensor, skip_special_tokens: bool) -> List[str]:
        """
        Decode received input_ids into a list of word tokens.
        Args:
            input_ids (torch.Tensor): Input ids representing
            word tokens for a sentence/document.
        """
        return self.tokenizer.convert_ids_to_tokens(
            input_ids[0], skip_special_tokens=skip_special_tokens
        )

    def _bpe_decode(self, tokens: List[str],) -> Tuple[List[str], List]:

        byte_encoder = bytes_to_unicode()
        byte_decoder = {v: k for k, v in byte_encoder.items()}
        decoded_each_tok = [
            bytearray([byte_decoder[c] for c in t]).decode(
                encoding="utf-8", errors="replace"
            )
            for t in tokens
        ]

        end_points = []
        force_break = False
        for idx, token in enumerate(decoded_each_tok):
            # special token, punctuation, alphanumeric
            if (
                token in self.tokenizer.all_special_tokens
                or token in string.punctuation
                or not any([x.isalnum() for x in token.lstrip()])
                or token.lstrip == "'s"
            ):
                end_points.append(idx)
                force_break = True
                continue

            if force_break:
                end_points.append(idx)
                force_break = False
                continue

            if token[0] == " ":
                tokens[idx] = token[:]
                end_points.append(idx)

        end_points.append(len(tokens))

        segments = []
        for i in range(1, len(end_points)):
            if end_points[i - 1] == end_points[i]:
                continue
            segments.append((end_points[i - 1], end_points[i]))

        filtered_tokens = []
        for s0, s1 in segments:
            filtered_tokens.append("".join(decoded_each_tok[s0:s1]))
        # print(filtered_tokens)
        # print(segments)

        return filtered_tokens, segments

    def process_input(self, request):
        inputs = [[request["question"], request["context"]]]
        encoded_inputs = self.encode(
            inputs, add_special_tokens=True, return_tensors="pt"
        )

        decoded_text = self.decode(
            encoded_inputs["input_ids"], skip_special_tokens=False
        )
        # print(decoded_text)
        filtered_tokens, segments = self._bpe_decode(decoded_text)
        return filtered_tokens, segments

    def assign_pos_tags(self, tokens, nlp):
        words = [x.lstrip() for x in tokens]
        spaces = [
            False if i == len(tokens) - 1 else tokens[i + 1][0] == " "
            for i in range(len(tokens))
        ]

        valid_idx = [i for i, w in enumerate(words) if len(w)]
        words = [words[i] for i in valid_idx]
        spaces = [spaces[i] for i in valid_idx]
        doc = Doc(nlp.vocab, words=words, spaces=spaces)
        processed_tokens = nlp(doc)

        tag_info = [("", "NULL", "NULL")] * len(tokens)
        for i, proc_tok in zip(valid_idx, processed_tokens):
            tag_info[i] = (proc_tok.text, proc_tok.pos_, proc_tok.tag_)

        return tag_info

    def tag_instance(self, request, nlp):
        words, segments = self.process_input(request)
        # tokens = list(filter(None, tokens))

        context_start = words.index(self.tokenizer.eos_token)
        question_tokens = words[1:context_start]
        context_tokens = words[context_start + 2 : -1]
        question_tag_info = self.assign_pos_tags(question_tokens, nlp)
        context_tag_info = self.assign_pos_tags(context_tokens, nlp)
        tag_info = (
            [(self.cls_token, "SOS", "SOS")]
            + question_tag_info
            + [(self.eos_token, "EOS", "EOS"), (self.eos_token, "EOS", "EOS")]
            + context_tag_info
            + [(self.eos_token, "EOS", "EOS")]
        )

        assert len(tag_info) == len(words), "tags and words not equal"
        # print(tag_info)
        instance_info = {"words": words, "segments": segments, "tags": tag_info}
        # print(instance_info)
        return instance_info


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        description="Passing arguments for model, tokenizer, and dataset."
    )
    parser.add_argument(
        "--dataset", type=str, required=True, help="Specify the dataset to use."
    )

    args = parser.parse_args()

    if args.dataset == "squad":
        loader = dataloader.PreprocessData(
            "squad", "plain_text", save_data=False, save_path="../../../../../"
        )
        data = loader.processed_val_set()
    elif args.dataset == "squad_adversarial":
        loader = dataloader.PreprocessData(
            "squad_adversarial", "AddSent", save_data=False, save_path="../../../../../"
        )
        data = loader.processed_val_set()
    elif args.dataset == "trivia_qa":
        data = dataloader.get_dev_examples("./src/data", "dev_trivia.json")
    elif args.dataset == "hotpot_qa":
        data = dataloader.get_dev_examples("./src/data", "dev_hotpot.json")
    elif args.dataset == "news_qa":
        data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/NewsQA.jsonl")
    elif args.dataset == "bioasq":
        data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/BioASQ-dev.jsonl")
    elif args.dataset == "natural_questions":
        data = dataloader.get_dev_samples_mrqa(
            BASE_PATH + "src/data/NaturalQuestionsShort.jsonl"
        )
    else:
        raise ValueError("Dataset not supported.")

    tokenizer = AutoTokenizer.from_pretrained("roberta-base")
    outputs = list()
    tagger = POSTagger(tokenizer=tokenizer)
    nlp = spacy.load("en_core_web_sm")
    c = 0
    processed_instances = OrderedDict()

    if args.dataset == "squad_adversarial":
        for ex in tqdm(data.processed_val_set()):
            try:
                tag_info = tagger.tag_instance(
                    request={
                        "id": ex["id"],
                        "question": ex["question"],
                        "context": ex["context"],
                    },
                    nlp=nlp,
                )
                processed_instances[ex["id"]] = tag_info
                # print(tag_info)
                c += 1
            except Exception as e:
                print(f"Unable to get tags: {e}")
                print(ex)
    elif args.dataset in [
        "trivia_qa",
        "hotpot_qa",
        "news_qa",
        "natural_questions",
        "bioasq",
    ]:

        def remove_white_space(example):
            example["question_text"] = " ".join(example["question_text"].split())
            example["context_text"] = " ".join(example["context_text"].split())
            return example

        for ex in tqdm(data):
            ex = remove_white_space(ex)
            # print(ex)
            # break
            try:
                tag_info = tagger.tag_instance(
                    request={
                        "id": ex["qas_id"],
                        "question": ex["question_text"],
                        "context": ex["context_text"],
                    },
                    nlp=nlp,
                )
                processed_instances[ex["qas_id"]] = tag_info
                # print(tag_info)
                c += 1
            except Exception as e:
                print(f"Unable to get tags: {e}")
                print(ex)

    utils.dump_to_bin(
        processed_instances, BASE_PATH + f"src/data/{args.dataset}/pos_info.bin"
    )
    print(f"Saved instances: {c}")

src/calibration/baseline/tagger.py
UKPLab/CATfOOD
import os
import pickle
import json
from collections import OrderedDict
import torch
import pandas as pd

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"
# BASE_PATH = "/home/sachdeva/projects/ukp/exp_calibration//"


def dump_to_bin(obj, fname):
    with open(fname, "wb") as f:
        pickle.dump(obj, f)


def load_bin(fname):
    with open(fname, "rb") as f:
        return pickle.load(f)


def read_json(fname):
    with open(fname, encoding="utf-8") as f:
        return json.load(f)


def dump_json(obj, fname, indent=None):
    with open(fname, "w", encoding="utf-8") as f:
        return json.dump(obj, f, indent=indent)


def build_file_dict():
    # prefix = 'squad_sample-addsent_roberta-base'
    prefix = BASE_PATH + "interpretations/shap/trivia_dev_roberta"
    fnames = os.listdir(prefix)
    qa_ids = [x.split("-", 1)[1].split(".")[0] for x in fnames]
    fullnames = [os.path.join(prefix, x) for x in fnames]
    output = dict(zip(qa_ids, fullnames))
    # print(output)
    processed_instances = OrderedDict()
    for id in qa_ids:
        processed_instances[id] = torch.load(output[id])
        break
    print(processed_instances)
    # import torch
    # attr = torch.load()
    # return dict(zip(qa_ids, fullnames))


def lemmatize_pos_tag(x):
    """
    limit the amount of POS tags to reduce the number of features
    :param x:
    :return:
    """
    tok, pos, tag = x
    if tag == "NNS":
        tag = "NN"
    if tag == "NNPS":
        tag = "NNP"
    if tag.startswith("JJ"):
        tag = "JJ"
    if tag.startswith("RB"):
        tag = "RB"
    if tag.startswith("W"):
        tag = "W"
    if tag.startswith("PRP"):
        tag = "PRP"
    if tag.startswith("VB"):
        tag = "VB"
    if pos == "PUNCT":  # or tag in [":"]:
        tag = "PUNCT"
    # if tag == '$':
    #     tag = 'SYM'
    # if tag == "``" or tag == "''":
    #     tag = "AUX"
    return tok, pos, tag


# trivia_tag_roberta_base_info


def convert_predictions_to_json_format():
    pred_df = (
        pd.read_json(
            BASE_PATH
            + "src/data/squad/nbest_predictions_roberta_squad_easy_0.75_top20.json"
        )
        .T.rename_axis("id")
        .reset_index()
    )
    print(pred_df.head())
    answers = [f"answer_{i}" for i in range(20)]
    pred_df.columns = ["id"] + answers

    from src.calibration.baseline import dataloader
    from tqdm import tqdm

    data = dataloader.PreprocessData(
        "squad_adversarial", "AddSent", save_data=False, save_path="../../../../"
    )
    val_set = data.processed_val_set()
    gold_text = [(sample["id"], sample["answers"]["text"]) for sample in tqdm(val_set)]
    # gold_text = [(sample["id"], sample["answers"]) for sample in tqdm(trivia_data)]
    # print(gold_text)
    gold_df = pd.DataFrame(gold_text, columns=["id", "gold_answers"])
    data = pd.merge(pred_df, gold_df, on="id")
    # rename columns
    data = data.rename(columns={"gold_answers": "gold_text"})
    print(data.head())
    print(data.columns)
    print(data["answer_0"][0])


if __name__ == "__main__":
    # build_file_dict()
    # fname = "src/data/squad_adversarial/calib_cf_ig.bin"
    # x = load_bin(fname)
    # for i in x:
    #     print(i)
    # print(x[-1])
    # unq = []
    # p_unq = []
    # c = 0
    # for q_id in x:

    #     tags = x[q_id]
    #     words, tags_for_tok = tags['words'], tags['tags']
    #     if q_id == "56f879bdaef23719006260e2":
    #         print(tags)
    #     # print(tags_for_tok)
    #     tags_for_tok = [lemmatize_pos_tag(x) for x in tags_for_tok]# if x[2] not in ["``", "''", "SOS", "EOS"]]
    #     if q_id == "56f879bdaef23719006260e2":
    #         print(tags)
    #         print(tags_for_tok)
    #     p = [t[1] for t in tags_for_tok]
    #     t = [t[2] for t in tags_for_tok]
    #     for v in t:
    #         if v not in unq:
    #             unq.append(v)
    #     for v in p:
    #         if v not in p_unq:
    #             p_unq.append(v)
    #     # c += 1
    #     # if c == 1000:
    #     #     break
    # print(len(unq))
    # print(unq)
    #
    # print(len(p_unq))
    # print(p_unq)
    # base_path = "/home/sachdeva/projects/ukp/exp_calibration"
    # file = base_path + "/src/data/squad_adv_new/calibration_data.bin"
    # data = load_bin(file)
    # print(data)
    convert_predictions_to_json_format()

src/calibration/baseline/utils.py
UKPLab/CATfOOD
import torch
from tqdm import tqdm
from collections import OrderedDict
import numpy as np
import traceback
from src.calibration.baseline import dataloader, utils
from sklearn.decomposition import PCA


BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        description="Passing arguments for model, tokenizer, and dataset."
    )
    parser.add_argument(
        "--model_name", type=str, required=True, help="Specify the model to use."
    )
    parser.add_argument(
        "--dataset", type=str, required=True, help="Specify the dataset to use."
    )
    args = parser.parse_args()

    if args.dataset == "squad":
        loader = dataloader.PreprocessData(
            "squad", "plain_text", save_data=False, save_path="../../../../"
        )
        data = loader.processed_val_set()
    elif args.dataset == "squad_adversarial":
        loader = dataloader.PreprocessData(
            "squad_adversarial", "AddSent", save_data=False, save_path="../../../../"
        )
        data = loader.processed_val_set()
    elif args.dataset == "trivia_qa":
        data = dataloader.get_dev_examples("./src/data", "dev_trivia.json")
    elif args.dataset == "hotpot_qa":
        data = dataloader.get_dev_examples("./src/data", "dev_hotpot.json")
    elif args.dataset == "news_qa":
        data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/NewsQA.jsonl")
    elif args.dataset == "bioasq":
        data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/BioASQ-dev.jsonl")
    elif args.dataset == "natural_questions":
        data = dataloader.get_dev_samples_mrqa(
            BASE_PATH + "src/data/NaturalQuestionsShort.jsonl"
        )
    else:
        raise ValueError("Dataset not supported.")

    attributions = utils.load_bin(
        f"{BASE_PATH}src/data/{args.dataset}/dense_repr_info_{args.model_name}.bin"
    )
    num_samples = len(data)
    all_data = []
    processed_instances = OrderedDict()

    if args.dataset == "squad_adversarial":
        for ex in tqdm(data):
            try:
                attr = np.array(attributions[ex["id"]])
                pad_width = ((0, 0), (0, 512 - attr.shape[1]), (0, 0))
                padded_arr = np.pad(
                    attr, pad_width=pad_width, mode="constant", constant_values=0
                )
                all_data.append(padded_arr.squeeze())
            except Exception:
                print(ex)
                print(f"Unable to get representations: {traceback.format_exc()}")
    elif args.dataset in [
        "trivia_qa",
        "hotpot_qa",
        "news_qa",
        "natural_questions",
        "bioasq",
    ]:

        for ex in tqdm(data):
            try:
                attr = np.array(attributions[ex["qas_id"]])
                pad_width = ((0, 0), (0, 512 - attr.shape[1]), (0, 0))
                padded_arr = np.pad(
                    attr, pad_width=pad_width, mode="constant", constant_values=0
                )
                all_data.append(padded_arr.squeeze())
            except Exception:
                print(ex)
                print(f"Unable to get attributions: {traceback.format_exc()}")

    hidden_states = np.stack(all_data, axis=0)
    # print(hidden_states.shape)
    hidden_states_flattened = np.reshape(hidden_states, (-1, 768))
    # Instantiate PCA object with desired number of components
    num_components = 10
    pca = PCA(n_components=num_components)
    pca.fit(hidden_states_flattened)
    reduced_hidden_states = pca.transform(hidden_states_flattened)
    reduced_hidden_states = np.reshape(
        reduced_hidden_states, (num_samples, 512, num_components)
    )
    # print(reduced_hidden_states.shape)

    for i, ex in enumerate(tqdm(data)):
        states = np.expand_dims(reduced_hidden_states[i, :, :], axis=0)
        if args.dataset == "squad_adversarial":
            processed_instances[ex["id"]] = torch.tensor(states)
        else:
            processed_instances[ex["qas_id"]] = torch.tensor(states)
    utils.dump_to_bin(
        processed_instances,
        BASE_PATH
        + f"src/data/{args.dataset}/dense_repr_pca_10_info_{args.model_name}.bin",
    )

src/calibration/dense_features/dim_reduction.py
UKPLab/CATfOOD
from transformers import (
    AutoTokenizer,
    # BertAdapterModel,
    # RobertaAdapterModel,
    RobertaForQuestionAnswering,
)
import torch
from tqdm import tqdm
from collections import OrderedDict
import traceback
from src.calibration.baseline import dataloader, utils

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


class DenseRepresentations:
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer

    def extract_representations(self, inputs):

        input_ids = self.tokenizer.encode(
            inputs,
            add_special_tokens=True,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=512,
        )
        # Get the start and end logits
        with torch.no_grad():
            outputs = self.model(input_ids, output_hidden_states=True)
            hidden_states = outputs.hidden_states[0]

        return hidden_states


if __name__ == "__main__":
    # extract_representations()

    import argparse

    parser = argparse.ArgumentParser(
        description="Passing arguments for model, tokenizer, and dataset."
    )

    parser.add_argument(
        "--model_name",
        default="roberta-squad-flan-ul2-context-rel-noise-seed-42",
        type=str,
        required=False,
        help="Specify the model to use.",
    )
    parser.add_argument(
        "--tokenizer",
        default="roberta-base",
        type=str,
        required=False,
        help="Specify the tokenizer to use.",
    )
    parser.add_argument(
        "--dataset", type=str, required=True, help="Specify the dataset to use."
    )

    args = parser.parse_args()

    model = RobertaForQuestionAnswering.from_pretrained(BASE_PATH + args.model_name)
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)

    if args.dataset == "squad":
        loader = dataloader.PreprocessData(
            "squad", "plain_text", save_data=False, save_path="../../../../"
        )
        data = loader.processed_val_set()
    elif args.dataset == "squad_adversarial":
        loader = dataloader.PreprocessData(
            "squad_adversarial", "AddSent", save_data=False, save_path="../../../../"
        )
        data = loader.processed_val_set()
    elif args.dataset == "trivia_qa":
        data = dataloader.get_dev_examples("./src/data", "dev_trivia.json")
    elif args.dataset == "hotpot_qa":
        data = dataloader.get_dev_examples("./src/data", "dev_hotpot.json")
    elif args.dataset == "news_qa":
        data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/NewsQA.jsonl")
    elif args.dataset == "bioasq":
        data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/BioASQ-dev.jsonl")
    elif args.dataset == "natural_questions":
        data = dataloader.get_dev_samples_mrqa(
            BASE_PATH + "src/data/NaturalQuestionsShort.jsonl"
        )
    else:
        raise ValueError("Dataset not supported.")

    outputs = list()
    c = 0
    representation = DenseRepresentations(model=model, tokenizer=tokenizer)
    processed_instances = OrderedDict()

    if args.dataset == "squad_adversarial":
        for ex in tqdm(data):
            try:
                # if ex["id"] == "56f879bdaef23719006260e2":
                states = representation.extract_representations(
                    [ex["question"], ex["context"]]
                )
                processed_instances[ex["id"]] = states
                c += 1
            except Exception:
                print(ex)
                print(f"Unable to get representations: {traceback.format_exc()}")
    elif args.dataset in [
        "trivia_qa",
        "hotpot_qa",
        "news_qa",
        "natural_questions",
        "bioasq",
    ]:

        def remove_white_space(example):
            example["question_text"] = " ".join(example["question_text"].split())
            example["context_text"] = " ".join(example["context_text"].split())
            return example

        for ex in tqdm(data):
            ex = remove_white_space(ex)
            try:
                states = representation.extract_representations(
                    [ex["question_text"], ex["context_text"]]
                )
                processed_instances[ex["qas_id"]] = states
                c += 1
            except Exception:
                print(ex)
                print(f"Unable to get attributions: {traceback.format_exc()}")

    print(f"Processed {c} instances of original data")
    utils.dump_to_bin(
        processed_instances,
        BASE_PATH + f"src/data/{args.dataset}/dense_repr_info_{args.model_name}.bin",
    )
    print(f"Saved instances: {c}")

src/calibration/dense_features/explanation_features.py
UKPLab/CATfOOD
from typing import List
from tqdm import tqdm
from collections import OrderedDict
import traceback

from transformers import (
    AutoTokenizer,
    # BertAdapterModel,
    # RobertaAdapterModel,
    RobertaForQuestionAnswering,
    PreTrainedModel,
    PreTrainedTokenizer,
    logging,
)

from src.calibration.explainers.base_explainer import BaseExplainer
from src.calibration.baseline import dataloader, utils

logging.set_verbosity_error()

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


class AttnAttribution(BaseExplainer):
    def __init__(self, model: PreTrainedModel, tokenizer: PreTrainedTokenizer):
        super().__init__(model=model, tokenizer=tokenizer)

    def interpret(self, inputs: List[List]):
        # get predicted answer
        model_kwargs = {"output_attentions": True}
        outputs, _ = self._predict(inputs, **model_kwargs)
        attn = outputs["attentions"][-1]
        weights = attn[:, :, 0, :].mean(1)
        attributions = weights.cpu().detach().numpy()[0]
        outputs = {"attributions": attributions}

        return outputs


if __name__ == "__main__":
    # model_path = BASE_PATH+"roberta-squad-flan-ul2-v1-temp-0.7"
    # model = RobertaForQuestionAnswering.from_pretrained(model_path)
    # tokenizer = AutoTokenizer.from_pretrained("roberta-base")
    import argparse

    parser = argparse.ArgumentParser(
        description="Passing arguments for model, tokenizer, and dataset."
    )

    parser.add_argument(
        "--model_name",
        default="",
        type=str,
        required=False,
        help="Specify the model to use.",
    )
    parser.add_argument(
        "--tokenizer",
        default="roberta-base",
        type=str,
        required=False,
        help="Specify the tokenizer to use.",
    )
    parser.add_argument(
        "--dataset", type=str, required=True, help="Specify the dataset to use."
    )

    args = parser.parse_args()

    model = RobertaForQuestionAnswering.from_pretrained(BASE_PATH + args.model_name)
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)

    if args.dataset == "squad":
        loader = dataloader.PreprocessData(
            "squad", "plain_text", save_data=False, save_path="../../../../../"
        )
        data = loader.processed_val_set()
    elif args.dataset == "squad_adversarial":
        loader = dataloader.PreprocessData(
            "squad_adversarial", "AddSent", save_data=False, save_path="../../../../../"
        )
        data = loader.processed_val_set()
    elif args.dataset == "trivia_qa":
        data = dataloader.get_dev_examples("./src/data", "dev_trivia.json")
    elif args.dataset == "hotpot_qa":
        data = dataloader.get_dev_examples("./src/data", "dev_hotpot.json")
    elif args.dataset == "news_qa":
        data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/NewsQA.jsonl")
    elif args.dataset == "bioasq":
        data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/BioASQ-dev.jsonl")
    elif args.dataset == "natural_questions":
        data = dataloader.get_dev_samples_mrqa(
            BASE_PATH + "src/data/NaturalQuestionsShort.jsonl"
        )
    else:
        raise ValueError("Dataset not supported.")

    outputs = list()
    grads = AttnAttribution(model=model, tokenizer=tokenizer)
    c = 0
    processed_instances = OrderedDict()

    if args.dataset == "squad_adversarial":
        for ex in tqdm(data):
            try:
                # if ex["id"] == "56f879bdaef23719006260e2":
                scores = grads.interpret([[ex["question"], ex["context"]]])
                processed_instances[ex["id"]] = scores
                c += 1
            except Exception:
                print(ex)
                print(f"Unable to get attributions: {traceback.format_exc()}")
    elif args.dataset in [
        "trivia_qa",
        "hotpot_qa",
        "news_qa",
        "natural_questions",
        "bioasq",
    ]:

        def remove_white_space(example):
            example["question_text"] = " ".join(example["question_text"].split())
            example["context_text"] = " ".join(example["context_text"].split())
            return example

        # data = dataloader.get_dev_examples(BASE_PATH+"src/data", "dev_hotpot.json")
        # data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/NewsQA.jsonl")
        # data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/BioASQ-dev.jsonl")
        # data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/NaturalQuestionsShort.jsonl")
        for ex in tqdm(data):
            ex = remove_white_space(ex)
            try:
                # if ex["id"] == "56f879bdaef23719006260e2":
                scores = grads.interpret([[ex["question_text"], ex["context_text"]]])
                processed_instances[ex["qas_id"]] = scores
                c += 1
            except Exception:
                print(ex)
                print(f"Unable to get attributions: {traceback.format_exc()}")

    print(f"Processed {c} instances of original data")
    utils.dump_to_bin(
        processed_instances,
        BASE_PATH + f"src/data/{args.dataset}/attn_info_{args.model_name}.bin",
    )
    print(f"Saved instances: {c}")

    # # process counterfactuals
    # cf_path = BASE_PATH + "src/data/squad_adversarial/rag_counterfactuals_turk0_last.jsonl"
    # c = 0
    # with jsonlines.open(cf_path) as reader:
    #     for ex in tqdm(reader):
    #         try:
    #             scores = grads.interpret([[ex["question"], ex["context"]]])
    #             processed_instances[ex["id"].replace("_", "-cf-")] = scores
    #             c += 1
    #         except Exception as e:
    #             print(f"Unable to get cf attributions: {e}")
    #             print(ex)
    #
    #     print(f"Processed {c} instances of counterfactual data")

    # print(processed_instances)
    # print(ast.literal_eval(processed_instances["56be4db0acb8001400a502ee"]))
    # print(len(ast.literal_eval(processed_instances["56be4db0acb8001400a502ee"])["attributions"][0]))

    ### Trivia QA

    # print(f"Processed {c} instances of original data")

    # # process counterfactuals
    # cf_path = BASE_PATH + "src/data/natural_questions/rag_counterfactuals_turk0.jsonl"
    # c = 0
    # with jsonlines.open(cf_path) as reader:
    #     for ex in tqdm(reader):
    #         try:
    #             scores = grads.interpret([[ex["question"], ex["context"]]])
    #             processed_instances[ex["id"].replace("_", "-cf-")] = scores
    #             c += 1
    #         except Exception as e:
    #             print(f"Unable to get cf attributions: {e}")
    #             print(ex)
    #
    #     print(f"Processed {c} instances of counterfactual data")
    #
    # utils.dump_to_bin(processed_instances,
    #                   BASE_PATH + "src/data/trivia_qa/attn_info_flan_ul2_cleaned.bin")
    #
    # print(f"Saved instances: {c}")

src/calibration/explainers/attentions/attention.py
UKPLab/CATfOOD
import math
import numpy as np
from tqdm import tqdm
import traceback
from collections import OrderedDict
from transformers import (
    AutoTokenizer,
    # BertAdapterModel,
    RobertaForQuestionAnswering,
    # RobertaAdapterModel,
    PreTrainedModel,
    PreTrainedTokenizer,
    logging,
)

import torch
from torch import backends
from torch.nn import Module, ModuleList
from typing import List

from src.calibration.explainers.base_explainer import BaseExplainer
from src.calibration.baseline import dataloader, utils

logging.set_verbosity_error()

torch.manual_seed(4)
torch.cuda.manual_seed(4)
np.random.seed(4)

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


class ScaledAttention(BaseExplainer):
    def __init__(self, model: PreTrainedModel, tokenizer: PreTrainedTokenizer):
        super().__init__(model=model, tokenizer=tokenizer)

    def get_model_attentions(self) -> Module or ModuleList:
        """
        Get the model attention layer
        :return:
        """
        model_prefix = self.model.base_model_prefix
        model_base = getattr(self.model, model_prefix)
        model_enc = getattr(model_base, "encoder")
        # get attn weights from last layer
        attentions = model_enc.layer[-1].attention
        return attentions

    def _register_forward_hooks(self, attentions_list: List):
        """
        Register the model attentions during the forward pass
        :param attentions_list:
        :return:
        """

        def forward_hook(module, inputs, output):
            attentions_list.append(
                output[1][:, :, 0, :].mean(1).squeeze(0).clone().detach()
            )

        handles = []
        attn_layer = self.get_model_attentions()
        handles.append(attn_layer.register_forward_hook(forward_hook))
        return handles

    def _register_attention_gradient_hooks(self, attn_grads: List):
        """
        Register the model gradients during the backward pass
        :param embedding_grads:
        :return:
        """

        def hook_layers(module, grad_in, grad_out):
            grads = grad_out[0]
            attn_grads.append(grads)

        hooks = []
        attentions = self.get_model_attentions()
        hooks.append(attentions.register_full_backward_hook(hook_layers))
        return hooks

    def get_gradients(self, inputs, answer_start, answer_end):
        """
        Compute model gradients
        :param inputs: list of question and context
        :param answer_start: answer span start
        :param answer_end: answer span end
        :return: dict of model gradients
        """
        attn_gradients: List[torch.Tensor] = []
        # print(answer_start, answer_end)

        original_param_name_to_requires_grad_dict = {}
        for param_name, param in self.model.named_parameters():
            original_param_name_to_requires_grad_dict[param_name] = param.requires_grad
            param.requires_grad = True

        hooks: List = self._register_attention_gradient_hooks(attn_gradients)
        with backends.cudnn.flags(enabled=False):
            encoded_inputs = self.encode(inputs, return_tensors="pt")
            encoded_inputs.to(self.device)
            outputs = self.model(
                **encoded_inputs,
                start_positions=answer_start,
                end_positions=answer_end,
                output_attentions=True,
            )
            loss = outputs.loss
            # Zero gradients.
            # NOTE: this is actually more efficient than calling `self._model.zero_grad()`
            # because it avoids a read op when the gradients are first updated below.
            for p in self.model.parameters():
                p.grad = None
            loss.backward()

        for hook in hooks:
            hook.remove()

        grad_dict = dict()
        for idx, grad in enumerate(attn_gradients):
            key = "grad_input_" + str(idx + 1)
            grad_dict[key] = grad.detach().cpu().numpy()

        # restore the original requires_grad values of the parameters
        for param_name, param in self.model.named_parameters():
            param.requires_grad = original_param_name_to_requires_grad_dict[param_name]
        # print(grad_dict)
        return grad_dict

    def interpret(self, inputs: List[List]):
        # get predicted answer
        outputs, _ = self._predict(inputs)
        start_idx = torch.argmax(outputs[0])
        end_idx = torch.argmax(outputs[1])

        answer_start = torch.tensor([start_idx], device=self.device)
        answer_end = torch.tensor([end_idx], device=self.device)

        attentions_list: List[torch.Tensor] = []
        # Hook used for saving embeddings
        handles: List = self._register_forward_hooks(attentions_list)
        try:
            grads = self.get_gradients(inputs, answer_start, answer_end)
        finally:
            for handle in handles:
                handle.remove()

        # Gradients come back in the reverse order that they were sent into the network
        attentions_list.reverse()
        attentions_list = [attn.cpu().detach().numpy() for attn in attentions_list]

        instances_with_grads = dict()
        for key, grad in grads.items():
            input_idx = int(key[-1]) - 1
            attn_grad = np.sum(grad[0] * attentions_list[input_idx][0], axis=1)
            norm = np.linalg.norm(attn_grad, ord=1)
            normalized_grad = [math.fabs(e) / norm for e in attn_grad]
            grads[key] = normalized_grad

        instances_with_grads["instance_" + str(1)] = grads

        outputs = {"attributions": instances_with_grads["instance_1"]["grad_input_1"]}
        return outputs


if __name__ == "__main__":
    # load model
    import argparse

    parser = argparse.ArgumentParser(
        description="Passing arguments for model, tokenizer, and dataset."
    )

    parser.add_argument(
        "--model_name",
        default="",
        type=str,
        required=False,
        help="Specify the model to use.",
    )
    parser.add_argument(
        "--tokenizer",
        default="roberta-base",
        type=str,
        required=False,
        help="Specify the tokenizer to use.",
    )
    parser.add_argument(
        "--dataset", type=str, required=True, help="Specify the dataset to use."
    )

    args = parser.parse_args()

    model = RobertaForQuestionAnswering.from_pretrained(BASE_PATH + args.model_name)
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)

    if args.dataset == "squad":
        loader = dataloader.PreprocessData(
            "squad", "plain_text", save_data=False, save_path="../../../../../"
        )
        data = loader.processed_val_set()
    elif args.dataset == "squad_adversarial":
        loader = dataloader.PreprocessData(
            "squad_adversarial", "AddSent", save_data=False, save_path="../../../../../"
        )
        data = loader.processed_val_set()
    elif args.dataset == "trivia_qa":
        data = dataloader.get_dev_examples("./src/data", "dev_trivia.json")
    elif args.dataset == "hotpot_qa":
        data = dataloader.get_dev_examples("./src/data", "dev_hotpot.json")
    elif args.dataset == "news_qa":
        data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/NewsQA.jsonl")
    elif args.dataset == "bioasq":
        data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/BioASQ-dev.jsonl")
    elif args.dataset == "natural_questions":
        data = dataloader.get_dev_samples_mrqa(
            BASE_PATH + "src/data/NaturalQuestionsShort.jsonl"
        )
    else:
        raise ValueError("Dataset not supported.")

    outputs = list()
    grads = ScaledAttention(model=model, tokenizer=tokenizer)
    c = 0
    processed_instances = OrderedDict()
    if args.dataset == "squad_adversarial":
        for ex in tqdm(data):
            try:
                # if ex["id"] == "56f879bdaef23719006260e2":
                scores = grads.interpret([[ex["question"], ex["context"]]])
                processed_instances[ex["id"]] = scores
                c += 1
            except Exception:
                print(ex)
                print(f"Unable to get attributions: {traceback.format_exc()}")
    elif args.dataset in [
        "trivia_qa",
        "hotpot_qa",
        "news_qa",
        "natural_questions",
        "bioasq",
    ]:

        def remove_white_space(example):
            example["question_text"] = " ".join(example["question_text"].split())
            example["context_text"] = " ".join(example["context_text"].split())
            return example

        # data = dataloader.get_dev_examples(BASE_PATH+"src/data", "dev_hotpot.json")
        # data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/NewsQA.jsonl")
        # data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/BioASQ-dev.jsonl")
        # data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/NaturalQuestionsShort.jsonl")

        for ex in tqdm(data):
            ex = remove_white_space(ex)
            try:
                # if ex["id"] == "56f879bdaef23719006260e2":
                scores = grads.interpret([[ex["question_text"], ex["context_text"]]])
                processed_instances[ex["qas_id"]] = scores
                c += 1
            except Exception:
                print(ex)
                print(f"Unable to get attributions: {traceback.format_exc()}")

    print(f"Processed {c} instances of original data")
    utils.dump_to_bin(
        processed_instances,
        BASE_PATH + f"src/data/{args.dataset}/sc_attn_info_{args.model_name}.bin",
    )
    print(f"Saved instances: {c}")

src/calibration/explainers/attentions/scaled_attention.py
UKPLab/CATfOOD
import string
import numpy as np
from abc import ABC
from typing import List, Tuple, Dict, Optional

import torch
from torch import backends
from torch.nn import Module, ModuleList
from transformers import PreTrainedModel, PreTrainedTokenizer
from transformers.models.gpt2.tokenization_gpt2 import bytes_to_unicode


class BaseExplainer(ABC):
    def __init__(
        self, model: PreTrainedModel, tokenizer: PreTrainedTokenizer,
    ):
        self.model = model
        self.tokenizer = tokenizer

        if self.model.config.model_type == "gpt2":
            self.ref_token_id = self.tokenizer.eos_token_id
        else:
            self.ref_token_id = self.tokenizer.pad_token_id

        self.sep_token_id = (
            self.tokenizer.sep_token_id
            if self.tokenizer.sep_token_id is not None
            else self.tokenizer.eos_token_id
        )
        self.cls_token_id = (
            self.tokenizer.cls_token_id
            if self.tokenizer.cls_token_id is not None
            else self.tokenizer.bos_token_id
        )

        self.model_prefix = model.base_model_prefix

        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)

        self.word_embeddings = self.model.get_input_embeddings()
        self.position_embeddings = None
        self.token_type_embeddings = None

    def _ensure_tensor_on_device(self, **inputs):
        """
        Ensure PyTorch tensors are on the specified device.

        Args:
            inputs (keyword arguments that should be :obj:`torch.Tensor`): The
                tensors to place on :obj:`self.device`.

        Return:
            :obj:`Dict[str, torch.Tensor]`: The same as :obj:`inputs`
                but on the proper device.
        """
        return {name: tensor.to(self.model.device) for name, tensor in inputs.items()}

    def get_model_embeddings(
        self, embedding_type: str = "word_embeddings"
    ) -> Module or ModuleList:
        """
        Get the model embedding layer
        Args:
            embedding_type: can be one of word_embeddings,
                token_type_embeddings or position_embeddings
        """
        embeddings = Module or ModuleList
        model_prefix = self.model.base_model_prefix
        model_base = getattr(self.model, model_prefix)
        model_embeddings = getattr(model_base, "embeddings")
        if embedding_type == "word_embeddings":
            embeddings = model_embeddings.word_embeddings
        elif embedding_type == "token_type_embeddings":
            embeddings = model_embeddings.token_type_embeddings
        elif embedding_type == "position_embeddings":
            embeddings = model_embeddings.position_embeddings
        return embeddings

    def _register_hooks(self, embeddings_list: List, alpha: int):
        """
        Register the model embeddings during the forward pass
        Args:
            embeddings_list: list to store embeddings during forward pass
        """

        def forward_hook(module, inputs, output):
            if alpha == 0:
                embeddings_list.append(output.squeeze(0).clone().detach())

        handles = []
        embedding_layer = self.get_model_embeddings()
        handles.append(embedding_layer.register_forward_hook(forward_hook))
        return handles

    def _register_embedding_gradient_hooks(self, embedding_grads: List):
        """
        Register the model gradients during the backward pass
        Args:
            embedding_grads: list to store the gradients
        """

        def hook_layers(module, grad_in, grad_out):
            grads = grad_out[0]
            embedding_grads.append(grads)

        hooks = []
        embedding_layer = self.get_model_embeddings()
        hooks.append(embedding_layer.register_full_backward_hook(hook_layers))
        return hooks

    def get_gradients(self, inputs, answer_start, answer_end):
        """
        Compute model gradients
        Args:
            inputs: list of question and context
            answer_start: answer span start
            answer_end: answer span end
        Return:
            dict of model gradients
        """
        embedding_gradients: List[torch.Tensor] = []
        # print(answer_start, answer_end)

        original_param_name_to_requires_grad_dict = {}
        for param_name, param in self.model.named_parameters():
            original_param_name_to_requires_grad_dict[param_name] = param.requires_grad
            param.requires_grad = True

        hooks: List = self._register_embedding_gradient_hooks(embedding_gradients)
        with backends.cudnn.flags(enabled=False):
            encoded_inputs = self.encode(inputs, return_tensors="pt")
            encoded_inputs.to(self.device)
            outputs = self.model(
                **encoded_inputs,
                start_positions=answer_start.to(self.device),
                end_positions=answer_end.to(self.device),
            )
            loss = outputs.loss

            # Zero gradients.
            # NOTE: this is actually more efficient than calling `self._model.zero_grad()`
            # because it avoids a read op when the gradients are first updated below.
            for p in self.model.parameters():
                p.grad = None
            loss.backward()

        for hook in hooks:
            hook.remove()

        grad_dict = dict()
        for idx, grad in enumerate(embedding_gradients):
            key = "grad_input_" + str(idx + 1)
            grad_dict[key] = grad.detach().cpu().numpy()

        # restore the original requires_grad values of the parameters
        for param_name, param in self.model.named_parameters():
            param.requires_grad = original_param_name_to_requires_grad_dict[param_name]
        # print(grad_dict)
        return grad_dict

    def encode(
        self, inputs: list = None, add_special_tokens: bool = True, return_tensors=None
    ):
        """
        Encode inputs using the model tokenizer
        Args:
            inputs: question, context pair as a list
            add_special_tokens: where to add CLS, SEP tokens
            return_offsets_mapping: mapping of sub-words to tokens
            return_tensors: whether to return tensors
        Return:
            tokenized inputs
        """
        return self.tokenizer(
            inputs,
            add_special_tokens=add_special_tokens,
            return_tensors=return_tensors,
            padding=True,
            truncation=True,
            max_length=512,
        )

    def decode(self, input_ids: torch.Tensor, skip_special_tokens: bool) -> List[str]:
        """
        Decode received input_ids into a list of word tokens.
        Args:
            input_ids (torch.Tensor): Input ids representing
            word tokens for a sentence/document.
        """
        return self.tokenizer.convert_ids_to_tokens(
            input_ids[0], skip_special_tokens=skip_special_tokens
        )

    def _predict(self, inputs, **model_kwargs) -> tuple:
        """
        Inference on the input.
        Args:
            inputs: list of question, context input
        Returns:
             The model outputs and optionally the input features
        """
        encoded_inputs = self.encode(
            inputs, add_special_tokens=True, return_tensors="pt"
        )
        encoded_inputs.to(self.device)
        self.decoded_text = self.decode(
            encoded_inputs["input_ids"], skip_special_tokens=False
        )
        self.words_mapping = encoded_inputs.word_ids()

        all_predictions = list()
        self.model.to(self.device)
        predictions = self.model(**encoded_inputs, **model_kwargs)
        # print(predictions)
        all_predictions.append(predictions)
        keys = all_predictions[0].keys()
        final_prediction = {}
        for key in keys:
            if isinstance(all_predictions[0][key], tuple):
                tuple_of_lists = list(
                    zip(
                        *[
                            [
                                torch.stack(p).to(self.device)
                                if isinstance(p, tuple)
                                else p.to(self.device)
                                for p in tpl[key]
                            ]
                            for tpl in all_predictions
                        ]
                    )
                )
                final_prediction[key] = tuple(torch.cat(l) for l in tuple_of_lists)
            else:
                final_prediction[key] = torch.cat(
                    [p[key].to(self.device) for p in all_predictions]
                )

        return predictions, encoded_inputs

    def question_answering(self, request):
        """
        Span-based question answering for a given question and context.
        We expect the input to use the (question, context) format for the text pairs.
        Args:
          request: the prediction request

        """

        def decode(
            start_: np.ndarray,
            end_: np.ndarray,
            topk: int,
            max_answer_len: int,
            undesired_tokens_: np.ndarray,
        ) -> Tuple:
            """
            Take the output of any :obj:`ModelForQuestionAnswering` and will generate probabilities
            for each span to be the actual answer.

            In addition, it filters out some unwanted/impossible cases like answer len being greater
            than max_answer_len or answer end position being before the starting position. The method
            supports output the k-best answer through the topk argument.

            Args:
                start_ (:obj:`np.ndarray`): Individual start probabilities for each token.
                end (:obj:`np.ndarray`): Individual end_ probabilities for each token.
                topk (:obj:`int`): Indicates how many possible answer span(s) to extract from the model output.
                max_answer_len (:obj:`int`): Maximum size of the answer to extract from the model's output.
                undesired_tokens_ (:obj:`np.ndarray`): Mask determining tokens that can be part of the answer
            """
            # Ensure we have batch axis
            if start_.ndim == 1:
                start_ = start_[None]

            if end_.ndim == 1:
                end_ = end_[None]

            # Compute the score of each tuple(start_, end_) to be the real answer
            outer = np.matmul(np.expand_dims(start_, -1), np.expand_dims(end_, 1))

            # Remove candidate with end_ < start_ and end_ - start_ > max_answer_len
            candidates = np.tril(np.triu(outer), max_answer_len - 1)

            #  Inspired by Chen & al. (https://github.com/facebookresearch/DrQA)
            scores_flat = candidates.flatten()
            if topk == 1:
                idx_sort = [np.argmax(scores_flat)]
            elif len(scores_flat) < topk:
                idx_sort = np.argsort(-scores_flat)
            else:
                idx = np.argpartition(-scores_flat, topk)[0:topk]
                idx_sort = idx[np.argsort(-scores_flat[idx])]

            starts_, ends_ = np.unravel_index(idx_sort, candidates.shape)[1:]
            desired_spans = np.isin(starts_, undesired_tokens_.nonzero()) & np.isin(
                ends_, undesired_tokens_.nonzero()
            )
            starts_ = starts_[desired_spans]
            ends_ = ends_[desired_spans]
            scores_ = candidates[0, starts_, ends_]

            return starts_, ends_, scores_

        predictions, features = self._predict(request)

        task_outputs = {"answers": []}
        for idx, (start, end, (_, context)) in enumerate(
            zip(predictions["start_logits"], predictions["end_logits"], request)
        ):
            start = start.cpu().detach().numpy()
            end = end.cpu().detach().numpy()
            # Ensure padded tokens & question tokens cannot belong to the set of candidate answers.
            question_tokens = np.abs(
                np.array([s != 1 for s in features.sequence_ids(idx)]) - 1
            )
            # Unmask CLS token for 'no answer'
            question_tokens[0] = 1
            undesired_tokens = question_tokens & features["attention_mask"][idx].numpy()

            # Generate mask
            undesired_tokens_mask = undesired_tokens == 0.0

            # Make sure non-context indexes in the tensor cannot contribute to the softmax
            start = np.where(undesired_tokens_mask, -10000.0, start)
            end = np.where(undesired_tokens_mask, -10000.0, end)

            start = np.exp(
                start - np.log(np.sum(np.exp(start), axis=-1, keepdims=True))
            )
            end = np.exp(end - np.log(np.sum(np.exp(end), axis=-1, keepdims=True)))

            # Get score for 'no answer' then mask for decoding step (CLS token
            no_answer_score = (start[0] * end[0]).item()
            start[0] = end[0] = 0.0

            starts, ends, scores = decode(start, end, 1, 128, undesired_tokens)
            enc = features[idx]
            answers = [
                {
                    "score": score.item(),
                    "start": enc.word_to_chars(enc.token_to_word(s), sequence_index=1)[
                        0
                    ],
                    "end": enc.word_to_chars(enc.token_to_word(e), sequence_index=1)[1],
                    "answer": context[
                        enc.word_to_chars(enc.token_to_word(s), sequence_index=1)[
                            0
                        ] : enc.word_to_chars(enc.token_to_word(e), sequence_index=1)[1]
                    ],
                }
                for s, e, score in zip(starts, ends, scores)
            ]
            answers.append(
                {"score": no_answer_score, "start": 0, "end": 0, "answer": ""}
            )
            answers = sorted(answers, key=lambda x: x["score"], reverse=True)[:1]
            task_outputs["answers"].append(answers)
        return task_outputs

    def handle_edge_cases(self, tokens):
        context_start = tokens.index(self.tokenizer.eos_token)
        word_map = []
        # print(context_start)
        # force_break = False
        for i, t in enumerate(tokens):
            # special token
            if t in self.tokenizer.all_special_tokens:
                word_map.append(None)
                continue
            if i < context_start and (t == "'s" or t == "'t" or t == "th"):
                word_map.append(word_map[i - 1])
            elif i < context_start:
                if len(word_map) == 1:
                    word_map.append(0)
                else:
                    word_map.append(word_map[i - 1] + 1)
            elif i > context_start and (t == "'s" or t == "'t" or t == "th"):
                word_map.append(word_map[i - 1])
            elif i > context_start:
                if word_map[i - 1] is None:
                    word_map.append(0)
                else:
                    word_map.append(word_map[i - 1] + 1)
        return word_map

    def _bpe_decode(
        self,
        tokens: List[str],
        # attributions: List
    ) -> List:

        byte_encoder = bytes_to_unicode()
        byte_decoder = {v: k for k, v in byte_encoder.items()}
        decoded_each_tok = [
            bytearray([byte_decoder[c] for c in t]).decode(
                encoding="utf-8", errors="replace"
            )
            for t in tokens
        ]

        end_points = []
        force_break = False
        for idx, token in enumerate(decoded_each_tok):
            # special token, punctuation, alphanumeric
            if (
                token in self.tokenizer.all_special_tokens
                or token in string.punctuation
                or not any([x.isalnum() for x in token.lstrip()])
                or token.lstrip == "'s"
            ):
                end_points.append(idx)
                force_break = True
                continue

            if force_break:
                end_points.append(idx)
                force_break = False
                continue

            if token[0] == " ":
                tokens[idx] = token[:]
                end_points.append(idx)

        end_points.append(len(tokens))

        segments = []
        for i in range(1, len(end_points)):
            if end_points[i - 1] == end_points[i]:
                continue
            segments.append((end_points[i - 1], end_points[i]))

        # filtered_tokens, scores = [], []
        # for s0, s1 in segments:
        #     filtered_tokens.append(''.join(decoded_each_tok[s0:s1]))
        #     # scores.append(np.sum(attributions[s0:s1], axis=0))
        # filtered_tokens = [token.lstrip() for token in filtered_tokens]
        # attribution_score = np.stack(scores, axis=0)

        return segments  # attribution_score

    def _wordpiece_decode(
        self, tokens: List[str], attributions: List
    ) -> Tuple[List[str], np.array]:

        decoded_each_tok = tokens
        word_map = self.words_mapping
        chars_to_handle = ["s", "t", "ve", "re", "m", "n't"]

        context_start = tokens.index(self.tokenizer.sep_token)
        for idx, token in enumerate(decoded_each_tok[:-1]):
            if (
                token not in self.tokenizer.all_special_tokens
                and token == "'"
                and decoded_each_tok[idx + 1] in chars_to_handle
                and idx < context_start
            ):
                word_map[idx] = word_map[idx - 1]
                word_map[idx + 1] = word_map[idx - 1]
                word_map[idx + 2 : context_start] = [
                    w - 2 for w in word_map[idx + 2 : context_start] if w
                ]
                continue
            if (
                token not in self.tokenizer.all_special_tokens
                and token == "'"
                and decoded_each_tok[idx + 1] in chars_to_handle
                and idx > context_start
            ):
                word_map[idx] = word_map[idx - 1]
                word_map[idx + 1] = word_map[idx - 1]
                word_map[idx + 2 : -1] = [w - 2 for w in word_map[idx + 2 : -1] if w]
                continue

        filtered_tokens = [decoded_each_tok[0]]
        for idx, (word_idx, word) in enumerate(zip(word_map, decoded_each_tok[1:])):
            if word_idx == word_map[idx + 1] and not word == self.tokenizer.sep_token:
                filtered_tokens[-1] = f'{filtered_tokens[-1]}{word.replace("##", "")}'
            else:
                filtered_tokens.append(word)

        attribution_score = [attributions[0]]
        for idx, (word_idx, score) in enumerate(zip(word_map, attributions[1:])):
            if word_idx == word_map[idx + 1] and word_idx is not None:
                attribution_score[-1] = attribution_score[-1] + score
            else:
                attribution_score.append(score)

        return filtered_tokens, np.array(attribution_score)

    def process_outputs(self, attributions: List, top_k: int, mode: str) -> List[Dict]:
        """
        post-process the word attributions to merge the sub-words tokens
        to words
        Args:
            attributions: word importance scores
            top_k: number of top word attributions
            mode: whether to show attribution in question, context or both
        Returns:
            dict of processed words along with their scores
        """

        filtered_tokens: list = []
        importance: np.array = np.array([])
        sep_tokens: int = 0
        dec_text = self.decoded_text
        if self.model.config.model_type in ["roberta", "bart"]:
            filtered_tokens, importance = self._bpe_decode(dec_text, attributions)
            sep_tokens = 2
        elif self.model.config.model_type == "bert":
            filtered_tokens, importance = self._wordpiece_decode(dec_text, attributions)
            sep_tokens = 1

        normed_imp = [np.round(float(i) / sum(importance), 3) for i in importance]
        result = [(w, a) for w, a in zip(filtered_tokens, normed_imp) if w != ""]
        assert len(filtered_tokens) == len(normed_imp)
        # outputs = {"attributions": result}
        context_start = filtered_tokens.index(self.tokenizer.sep_token)
        # account for cls token in result
        question = [
            (idx, v[0], v[1])
            for idx, v in enumerate(result[1:])
            if idx < context_start - 1
        ]

        context = [
            (idx - len(question) - sep_tokens, v[0], v[1])
            for idx, v in enumerate(result[1:])
            if idx > context_start - 1 and v[0] != self.tokenizer.sep_token
        ]

        outputs, outputs_question, outputs_context = [], [], []
        if mode == "question" or mode == "all":
            outputs_question = [
                (i, k.lower(), v)
                for i, k, v in sorted(question, key=lambda item: item[2], reverse=True)[
                    :top_k
                ]
            ]
        if mode == "context" or mode == "all":
            outputs_context = [
                (i, k.lower(), v)
                for i, k, v in sorted(context, key=lambda item: item[2], reverse=True)[
                    :top_k
                ]
            ]

        outputs = [{"question": outputs_question, "context": outputs_context}]
        return outputs

src/calibration/explainers/base_explainer.py
UKPLab/CATfOOD
import math
import numpy as np
from tqdm import tqdm
from typing import List, Dict, Tuple, Any
from collections import OrderedDict

from transformers import (
    AutoTokenizer,
    # BertAdapterModel,
    # RobertaAdapterModel,
    RobertaForQuestionAnswering,
    PreTrainedModel,
    PreTrainedTokenizer,
    logging,
)

from src.calibration.explainers.base_explainer import BaseExplainer
from src.calibration.baseline import dataloader, utils

import traceback

import torch

logging.set_verbosity_error()

torch.manual_seed(4)
torch.cuda.manual_seed(4)
np.random.seed(4)

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


class IntegratedGradients(BaseExplainer):
    def __init__(self, model: PreTrainedModel, tokenizer: PreTrainedTokenizer):
        super().__init__(model=model, tokenizer=tokenizer)

    def _register_hooks(self, embeddings_list: List, alpha: int):
        def forward_hook(module, inputs, output):
            # Save the input for later use. Only do so on first call.
            if alpha == 0:
                embeddings_list.append(output.squeeze(0).clone().detach())
            # Scale the embedding by alpha
            output.mul_(alpha)

        handles = []
        embedding_layer = self.get_model_embeddings()
        handles.append(embedding_layer.register_forward_hook(forward_hook))
        return handles

    def _integrate_gradients(
        self, inputs: List[List]
    ) -> Tuple[Dict[str, np.ndarray], torch.Tensor, torch.Tensor]:
        """
        Returns:
             integrated gradients for the given [`Instance`]
        """
        ig_grads: Dict[str, Any] = {}

        # List of Embedding inputs
        embeddings_list: List[torch.Tensor] = []

        # answer prediction
        outputs, _ = self._predict(inputs)
        start_idx = torch.argmax(outputs[0])
        end_idx = torch.argmax(outputs[1])
        answer_start = torch.tensor([start_idx])
        answer_end = torch.tensor([end_idx])

        # Use 10 terms in the summation approximation of the integral in integrated grad
        steps = 10
        # Exclude the endpoint because we do a left point integral approximation
        for alpha in np.linspace(0, 1.0, num=steps, endpoint=False):
            handles = []
            # Hook for modifying embedding value
            handles = self._register_hooks(embeddings_list, alpha)

            try:
                grads = self.get_gradients(inputs, answer_start, answer_end)
            finally:
                for handle in handles:
                    handle.remove()

            # Running sum of gradients
            if ig_grads == {}:
                ig_grads = grads
            else:
                for key in grads.keys():
                    ig_grads[key] += grads[key]

        # Average of each gradient term
        for key in ig_grads.keys():
            ig_grads[key] /= steps

        # Gradients come back in the reverse order that they were sent into the network
        embeddings_list.reverse()
        embeddings_list = [
            embedding.cpu().detach().numpy() for embedding in embeddings_list
        ]
        # Element-wise multiply average gradient by the input
        for idx, input_embedding in enumerate(embeddings_list):
            # print(idx, input_embedding)
            key = "grad_input_" + str(idx + 1)
            ig_grads[key] *= input_embedding

        return ig_grads, answer_start, answer_end

    def interpret(self, inputs: List[List]):
        # run integrated grad
        grads, answer_start, answer_end = self._integrate_gradients(inputs)
        # normalize results
        instances_with_grads = dict()
        for key, grad in grads.items():
            # The [0] here is undo-ing the batching that happens in get_gradients.
            embedding_grad = np.sum(grad[0], axis=1)
            norm = np.linalg.norm(embedding_grad, ord=1)
            normalized_grad = np.array([math.fabs(e) / norm for e in embedding_grad])
            grads[key] = normalized_grad

        instances_with_grads["instance_" + str(1)] = grads
        outputs = {"attributions": instances_with_grads["instance_1"]["grad_input_1"]}
        # print(len(instances_with_grads["instance_1"]["grad_input_1"]))

        return outputs


if __name__ == "__main__":
    # model_path = BASE_PATH + "roberta-squad-flan-ul2-v1-temp-0.7"
    # model = RobertaForQuestionAnswering.from_pretrained(model_path)
    # tokenizer = AutoTokenizer.from_pretrained("roberta-base")

    import argparse

    parser = argparse.ArgumentParser(
        description="Passing arguments for model, tokenizer, and dataset."
    )

    parser.add_argument(
        "--model_name",
        default="",
        type=str,
        required=False,
        help="Specify the model to use.",
    )
    parser.add_argument(
        "--tokenizer",
        default="roberta-base",
        type=str,
        required=False,
        help="Specify the tokenizer to use.",
    )
    parser.add_argument(
        "--dataset", type=str, required=True, help="Specify the dataset to use."
    )

    args = parser.parse_args()

    model = RobertaForQuestionAnswering.from_pretrained(BASE_PATH + args.model_name)
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)

    if args.dataset == "squad":
        loader = dataloader.PreprocessData(
            "squad", "plain_text", save_data=False, save_path="../../../../../"
        )
        data = loader.processed_val_set()
    elif args.dataset == "squad_adversarial":
        loader = dataloader.PreprocessData(
            "squad_adversarial", "AddSent", save_data=False, save_path="../../../../../"
        )
        data = loader.processed_val_set()
    elif args.dataset == "trivia_qa":
        data = dataloader.get_dev_examples("./src/data", "dev_trivia.json")
    elif args.dataset == "hotpot_qa":
        data = dataloader.get_dev_examples("./src/data", "dev_hotpot.json")
    elif args.dataset == "news_qa":
        data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/NewsQA.jsonl")
    elif args.dataset == "bioasq":
        data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/BioASQ-dev.jsonl")
    elif args.dataset == "natural_questions":
        data = dataloader.get_dev_samples_mrqa(
            BASE_PATH + "src/data/NaturalQuestionsShort.jsonl"
        )
    else:
        raise ValueError("Dataset not supported.")

    outputs = list()

    grads = IntegratedGradients(model=model, tokenizer=tokenizer)
    c = 0
    processed_instances = OrderedDict()

    if args.dataset == "squad_adversarial":
        for ex in tqdm(data):
            try:
                # if ex["id"] == "56f879bdaef23719006260e2":
                scores = grads.interpret([[ex["question"], ex["context"]]])
                processed_instances[ex["id"]] = scores
                c += 1
            except Exception:
                print(ex)
                print(f"Unable to get attributions: {traceback.format_exc()}")
    elif args.dataset in [
        "trivia_qa",
        "hotpot_qa",
        "news_qa",
        "natural_questions",
        "bioasq",
    ]:

        def remove_white_space(example):
            example["question_text"] = " ".join(example["question_text"].split())
            example["context_text"] = " ".join(example["context_text"].split())
            return example

        # data = dataloader.get_dev_examples(BASE_PATH+"src/data", "dev_hotpot.json")
        # data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/NewsQA.jsonl")
        # data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/BioASQ-dev.jsonl")
        # data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/NaturalQuestionsShort.jsonl")
        for ex in tqdm(data):
            ex = remove_white_space(ex)
            try:
                # if ex["id"] == "56f879bdaef23719006260e2":
                scores = grads.interpret([[ex["question_text"], ex["context_text"]]])
                processed_instances[ex["qas_id"]] = scores
                c += 1
            except Exception:
                print(ex)
                print(f"Unable to get attributions: {traceback.format_exc()}")

    print(f"Processed {c} instances of original data")
    utils.dump_to_bin(
        processed_instances,
        BASE_PATH + f"src/data/{args.dataset}/ig_info_{args.model_name}.bin",
    )
    print(f"Saved instances: {c}")

src/calibration/explainers/gradients/integrated_grads.py
UKPLab/CATfOOD
import math
from tqdm import tqdm
from collections import OrderedDict
import traceback
import numpy as np
from typing import List
import torch
from transformers import (
    AutoTokenizer,
    # BertAdapterModel,
    # RobertaAdapterModel,
    # AutoModelWithHeads,
    RobertaForQuestionAnswering,
    PreTrainedModel,
    PreTrainedTokenizer,
    logging,
)
from src.calibration.explainers.base_explainer import BaseExplainer
from src.calibration.baseline import dataloader, utils

logging.set_verbosity_error()

torch.manual_seed(4)
torch.cuda.manual_seed(4)
np.random.seed(4)

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


class SimpleGradients(BaseExplainer):
    """
    class for the implementation of simple gradients' explanation method
    """

    def __init__(self, model: PreTrainedModel, tokenizer: PreTrainedTokenizer):
        super().__init__(model=model, tokenizer=tokenizer)

    def interpret(
        self, inputs: List[List],
    ):
        """
        gets the word attributions
        """
        # get predicted answer
        outputs, _ = self._predict(inputs)
        start_idx = torch.argmax(outputs[0])
        end_idx = torch.argmax(outputs[1])

        answer_start = torch.tensor([start_idx])
        answer_end = torch.tensor([end_idx])

        embeddings_list: List[torch.Tensor] = []
        # Hook used for saving embeddings
        handles: List = self._register_hooks(embeddings_list, alpha=0)
        try:
            grads = self.get_gradients(inputs, answer_start, answer_end)
        finally:
            for handle in handles:
                handle.remove()

        # Gradients come back in the reverse order that they were sent into the network
        embeddings_list.reverse()
        embeddings_list = [
            embedding.cpu().detach().numpy() for embedding in embeddings_list
        ]
        # token_offsets.reverse()
        # embeddings_list = self._aggregate_token_embeddings(embeddings_list, token_offsets)
        instances_with_grads = dict()
        for key, grad in grads.items():
            # Get number at the end of every gradient key (they look like grad_input_[int],
            # we're getting this [int] part and subtracting 1 for zero-based indexing).
            # This is then used as an index into the reversed input array to match up the
            # gradient and its respective embedding.
            input_idx = int(key[-1]) - 1
            # The [0] here is undo-ing the batching that happens in get_gradients.
            emb_grad = np.sum(grad[0] * embeddings_list[input_idx][0], axis=1)
            norm = np.linalg.norm(emb_grad, ord=1)
            normalized_grad = [math.fabs(e) / norm for e in emb_grad]
            grads[key] = normalized_grad

        instances_with_grads["instance_" + str(1)] = grads

        outputs = {"attributions": instances_with_grads["instance_1"]["grad_input_1"]}

        return outputs


if __name__ == "__main__":

    import argparse

    parser = argparse.ArgumentParser(
        description="Passing arguments for model, tokenizer, and dataset."
    )
    parser.add_argument(
        "--model_name",
        default="",
        type=str,
        required=False,
        help="Specify the model to use.",
    )
    parser.add_argument(
        "--tokenizer",
        default="roberta-base",
        type=str,
        required=False,
        help="Specify the tokenizer to use.",
    )
    parser.add_argument(
        "--dataset", type=str, required=True, help="Specify the dataset to use."
    )

    args = parser.parse_args()
    model = RobertaForQuestionAnswering.from_pretrained(BASE_PATH + args.model_name)
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)

    if args.dataset == "squad":
        loader = dataloader.PreprocessData(
            "squad", "plain_text", save_data=False, save_path="../../../../../"
        )
        data = loader.processed_val_set()
    elif args.dataset == "squad_adversarial":
        loader = dataloader.PreprocessData(
            "squad_adversarial", "AddSent", save_data=False, save_path="../../../../../"
        )
        data = loader.processed_val_set()
    elif args.dataset == "trivia_qa":
        data = dataloader.get_dev_examples("./src/data", "dev_trivia.json")
    elif args.dataset == "hotpot_qa":
        data = dataloader.get_dev_examples("./src/data", "dev_hotpot.json")
    elif args.dataset == "news_qa":
        data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/NewsQA.jsonl")
    elif args.dataset == "bioasq":
        data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/BioASQ-dev.jsonl")
    elif args.dataset == "natural_questions":
        data = dataloader.get_dev_samples_mrqa(
            BASE_PATH + "src/data/NaturalQuestionsShort.jsonl"
        )
    else:
        raise ValueError("Dataset not supported.")

    outputs = list()

    grads = SimpleGradients(model=model, tokenizer=tokenizer)
    c = 0
    processed_instances = OrderedDict()
    if args.dataset == "squad_adversarial":
        for ex in tqdm(data):
            try:
                # if ex["id"] == "56f879bdaef23719006260e2":
                scores = grads.interpret([[ex["question"], ex["context"]]])
                processed_instances[ex["id"]] = scores
                c += 1
            except Exception:
                print(ex)
                print(f"Unable to get attributions: {traceback.format_exc()}")
    elif args.dataset in [
        "trivia_qa",
        "hotpot_qa",
        "news_qa",
        "natural_questions",
        "bioasq",
    ]:

        def remove_white_space(example):
            example["question_text"] = " ".join(example["question_text"].split())
            example["context_text"] = " ".join(example["context_text"].split())
            return example

        # data = dataloader.get_dev_examples(BASE_PATH+"src/data", "dev_hotpot.json")
        # data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/NewsQA.jsonl")
        # data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/BioASQ-dev.jsonl")
        # data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/NaturalQuestionsShort.jsonl")

        for ex in tqdm(data):
            ex = remove_white_space(ex)
            try:
                # if ex["id"] == "56f879bdaef23719006260e2":
                scores = grads.interpret([[ex["question_text"], ex["context_text"]]])
                processed_instances[ex["qas_id"]] = scores
                c += 1
            except Exception:
                print(ex)
                print(f"Unable to get attributions: {traceback.format_exc()}")

    print(f"Processed {c} instances of original data")
    utils.dump_to_bin(
        processed_instances,
        BASE_PATH + f"src/data/{args.dataset}/simple_grads_info_{args.model_name}.bin",
    )
    print(f"Saved instances: {c}")

src/calibration/explainers/gradients/simple_grads.py
UKPLab/CATfOOD
import json
import numpy as np
from typing import List, Dict
import random
from tqdm import tqdm
from collections import OrderedDict
import traceback

import torch
from transformers import (
    AutoTokenizer,
    # BertAdapterModel,
    # RobertaAdapterModel,
    RobertaForQuestionAnswering,
    PreTrainedModel,
    PreTrainedTokenizer,
    logging,
)

from src.calibration.explainers.base_explainer import BaseExplainer
from src.calibration.baseline import dataloader, utils

logging.set_verbosity_error()


class RandomAttributions(BaseExplainer):
    def __init__(self, model: PreTrainedModel, tokenizer: PreTrainedTokenizer):
        super().__init__(model=model, tokenizer=tokenizer)

    def interpret(self, inputs: List[List], output: str = "processed"):
        # get predicted answer
        outputs, _ = self._predict(inputs)
        start_idx = torch.argmax(outputs[0])
        end_idx = torch.argmax(outputs[1])

        answer_start = torch.tensor([start_idx])
        answer_end = torch.tensor([end_idx])

        attributions = [random.uniform(0, 1) for i in range(outputs[0].size()[1])]

        if output == "raw":
            return (
                self.encode(inputs, add_special_tokens=True, return_tensors="pt"),
                attributions,
                answer_start,
                answer_end,
            )

        outputs = self.process_outputs(attributions=attributions)

        return json.dumps(outputs, indent=4)

    def process_outputs(self, attributions: List, top_k=10, mode="all") -> Dict:
        """
        post-process the word attributions to merge the sub-words tokens
        to words
        Args:
            attributions: word importance scores
            top_k: number of top word attributions
            mode: whether to show attribution in question, context or both
        Returns:
            dict of processed words along with their scores
        """

        segments: list = []
        importance: np.array = np.array([])
        dec_text = self.decoded_text
        if self.model.config.model_type in ["roberta", "bart"]:
            segments = self._bpe_decode(dec_text)  # , attributions)
        elif self.model.config.model_type == "bert":
            filtered_tokens, importance = self._wordpiece_decode(dec_text, attributions)

        # normed_imp = [np.round(float(i) / sum(attributions), 3)
        #               for i in attributions]
        # result = [(w, a) for w, a in zip(filtered_tokens, normed_imp)]
        # if w != '']
        # assert len(segments) == len(attributions)
        # print(len(result))
        outputs = {"attributions": attributions, "segments": segments}
        return outputs


if __name__ == "__main__":
    # base_model = "bert-base-uncased"
    # adapter_model = "AdapterHub/bert-base-uncased-pf-squad_v2"
    # model = BertAdapterModel.from_pretrained(base_model)
    # tokenizer = AutoTokenizer.from_pretrained(base_model)
    # adapter_name = model.load_adapter(adapter_model, source="hf")
    # model.active_adapters = adapter_name

    model_path = "./src/checkpoints/roberta_squad"
    model = RobertaForQuestionAnswering.from_pretrained(model_path)
    tokenizer = AutoTokenizer.from_pretrained(model_path)

    data = dataloader.PreprocessData(
        "squad_adversarial", "AddSent", save_data=False, save_path="../../../../"
    )
    # data = dataloader.get_dev_examples("./src/calibration/data", "dev_trivia.json")
    outputs = list()

    # def remove_white_space(example):
    #     example["question_text"] = ' '.join(example["question_text"].split())
    #     example["context_text"] = ' '.join(example["context_text"].split())
    #     return example

    grads = RandomAttributions(model=model, tokenizer=tokenizer)
    c = 0
    processed_instances = OrderedDict()
    for ex in tqdm(data.processed_val_set()):
        # ex = remove_white_space(ex)
        # print(ex)
        # ex["context"] = "Life's good when we are stress free."
        try:
            # if ex["id"] == "56f879bdaef23719006260e2":
            scores = grads.interpret([[ex["question"], ex["context"]]])
            processed_instances[ex["id"]] = scores
            c += 1
        except Exception:
            print(ex)
            print(f"Unable to get attributions: {traceback.format_exc()}")

        # else:
        #     continue
        # if c == 1:
        # break
    # import ast
    # print(ast.literal_eval(processed_instances["56f879bdaef23719006260e2"]))
    utils.dump_to_bin(
        processed_instances, "./src/results/squad_adv_random_roberta_base_info.bin"
    )
    print(f"Saved instances: {c}")

src/calibration/explainers/random_attribution.py
UKPLab/CATfOOD
import os

from transformers import MODEL_FOR_QUESTION_ANSWERING_MAPPING
from transformers import AutoConfig, AutoTokenizer

from dataclasses import dataclass

MODEL_CONFIG_CLASSES = list(MODEL_FOR_QUESTION_ANSWERING_MAPPING.keys())
MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
# BASE_PATH="/home/sachdeva/projects/ukp/exp_calibration"
BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration"

# roberta-squad-llama-13b-v2-temp-0.7-seed-1
# roberta-squad-flan-ul2-v1-temp-0.7
# roberta-squad-gpt-jt-v2-temp-0.7-seed-0
# roberta-squad-gpt-neox-context-rel-seed-42
# roberta-squad-flan-ul2-context-rel-noise-seed-42


@dataclass
class InterpConfig:
    model_type: str = os.getenv("MODEL_TYPE", "roberta")
    model_name_or_path: str = f"{BASE_PATH}/{os.getenv('MODEL_NAME_OR_PATH', 'roberta-squad')}"
    tokenizer_name_or_path: str = "roberta-base"
    dataset: str = os.getenv("DATASET", "trivia")
    dataset_config: str = "AddSent"
    output_dir: str = BASE_PATH + "/src/data/squad"
    train_file: str = None
    predict_file: str = BASE_PATH + "/src/data/"
    eval_batch_size: int = 1
    config_name: str = None
    max_seq_length: int = 384
    max_query_len: int = 64
    max_answer_len: int = 30
    do_lower_case: bool = True
    n_best_size: int = 20
    cache_dir: str = None
    seed: int = 42
    local_rank: int = -1
    threads: int = -1
    first_n_samples: int = 13000
    do_vis: bool = False
    no_cuda: bool = False
    n_gpu: int = 1
    percent: float = 0.3
    percent_augment: float = 0.15
    # interp_dir: str = f"{os.getenv('INTERP_DIR', 'exp_llama2_flan_ul2')}/shap/trivia/dev/roberta"
    # visual_dir: str = f"{os.getenv('INTERP_DIR', 'exp_llama2_flan_ul2')}/shap/trivia/dev/visual"
    interp_dir: str = "exp_roberta_base/shap/nq/dev/roberta"
    visual_dir: str = "exp_roberta_base/shap/nq/dev/visual"


def load_config_and_tokenizer(args):
    if args.dataset in ["simple", "synth", "comp"]:
        tokenizer = SimBertTokenizer()
        if args.dataset == "simple":
            config = SimBertConfig()
        elif args.dataset == "comp":
            config = CompBertConfig()
        else:
            config = SynBertConfig()
    elif args.dataset in [
        "hpqa",
        "squad",
        "bioasq",
        "newsqa",
        "natq",
        "trivia",
        "hotpot",
        "squad_adversarial",
    ]:
        config = AutoConfig.from_pretrained(
            args.config_name if args.config_name else args.model_name_or_path,
            cache_dir=args.cache_dir if args.cache_dir else None,
        )
        tokenizer = AutoTokenizer.from_pretrained(
            args.tokenizer_name_or_path
            if args.tokenizer_name_or_path
            else args.model_name_or_path,
            do_lower_case=args.do_lower_case,
            cache_dir=args.cache_dir if args.cache_dir else None,
        )
    else:
        raise RuntimeError("Dataset not supported")

    return config, tokenizer

src/calibration/explainers/shap/common/config.py
UKPLab/CATfOOD
from collections import OrderedDict


class IndexedFeature:
    def __init__(self):
        self.data = OrderedDict()

    def add(self, k, v=1.0):
        if k in self.data:
            self.data[k] = self.data[k] + v
        else:
            self.data[k] = v

    def add_if_absent(self, k, v=1.0):
        if k not in self.data:
            self.data[k] = v

    def add_set(self, other):
        for k, v in other.data.items():
            self.data[k] = v

    def __getitem__(self, k):
        return self.data.get(k, 0.0)

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        return str(self.data)

    def __len__(self):
        return len(self.data)

    def add_prefix(self, prefix):
        new_data = OrderedDict([(prefix + k, v) for (k, v) in self.data.items()])
        self.data = new_data


class FeatureVocab:
    def __init__(self):
        self.feat_to_id = {}
        self.id_to_feat = {}

    def __getitem__(self, word):
        return self.feat_to_id.get(word, -1)

    def __contains__(self, word):
        return word in self.feat_to_id

    def __len__(self):
        return len(self.feat_to_id)

    def size(self):
        return len(self.feat_to_id)

    def get_word(self, wid):
        return self.id_to_feat[wid]

    def get_names(self):
        return [self.id_to_feat[i] for i in range(len(self))]

    def add(self, word):
        if word not in self:
            wid = self.feat_to_id[word] = len(self)
            self.id_to_feat[wid] = word
            return wid

src/calibration/explainers/shap/common/index_feature.py
UKPLab/CATfOOD
# from transformers.data.datasets.squad import *
# from transformers.data.processors.squad import *
# from multiprocessing import Pool, cpu_count

import json
import os
from functools import partial
from multiprocessing import Pool, cpu_count

import numpy as np
from tqdm import tqdm

from transformers.file_utils import is_torch_available
from transformers.tokenization_utils_base import TruncationStrategy
from transformers.utils import logging
from transformers.data.processors.utils import DataProcessor
from data.dataset_utils import get_prefix_tokens

logger = logging.get_logger(__name__)

MULTI_SEP_TOKENS_TOKENIZERS_SET = {"roberta", "camembert", "bart"}

if is_torch_available():
    import torch
    from torch.utils.data import TensorDataset


def _improve_answer_span(
    doc_tokens, input_start, input_end, tokenizer, orig_answer_text
):
    """Returns tokenized answer spans that better match the annotated answer."""

    candidate_cleaned_forms = []
    candidate_cleaned_forms.append(" ".join(tokenizer.tokenize(orig_answer_text)))
    if tokenizer.__class__.__name__ in [
        "RobertaTokenizer",
        "BartTokenizer",
    ]:
        candidate_cleaned_forms.append(
            " ".join(tokenizer.tokenize(orig_answer_text, add_prefix_space=True))
        )

    for new_start in range(input_start, input_end + 1):
        for new_end in range(input_end, new_start - 1, -1):
            text_span = " ".join(doc_tokens[new_start : (new_end + 1)])
            if text_span in candidate_cleaned_forms:
                return (new_start, new_end)

    return (input_start, input_end)


def _check_is_max_context(doc_spans, cur_span_index, position):
    """Check if this is the 'max context' doc span for the token."""
    best_score = None
    best_span_index = None
    for (span_index, doc_span) in enumerate(doc_spans):
        end = doc_span.start + doc_span.length - 1
        if position < doc_span.start:
            continue
        if position > end:
            continue
        num_left_context = position - doc_span.start
        num_right_context = end - position
        score = min(num_left_context, num_right_context) + 0.01 * doc_span.length
        if best_score is None or score > best_score:
            best_score = score
            best_span_index = span_index

    return cur_span_index == best_span_index


def _new_check_is_max_context(doc_spans, cur_span_index, position):
    """Check if this is the 'max context' doc span for the token."""
    # if len(doc_spans) == 1:
    # return True
    best_score = None
    best_span_index = None
    for (span_index, doc_span) in enumerate(doc_spans):
        end = doc_span["start"] + doc_span["length"] - 1
        if position < doc_span["start"]:
            continue
        if position > end:
            continue
        num_left_context = position - doc_span["start"]
        num_right_context = end - position
        score = min(num_left_context, num_right_context) + 0.01 * doc_span["length"]
        if best_score is None or score > best_score:
            best_score = score
            best_span_index = span_index

    return cur_span_index == best_span_index


def _is_whitespace(c):
    if c == " " or c == "\t" or c == "\r" or c == "\n" or ord(c) == 0x202F:
        return True
    return False


def custom_squad_convert_example_to_features_init(tokenizer_for_convert):
    global tokenizer
    tokenizer = tokenizer_for_convert


def custom_squad_convert_examples_to_features(
    examples,
    tokenizer,
    max_seq_length,
    max_query_length,
    is_training,
    padding_strategy="max_length",
    return_dataset=False,
    threads=1,
    tqdm_enabled=True,
    dataset="hpqa,",
):
    """
    Converts a list of examples into a list of features that can be directly given as input to a model.
    It is model-dependant and takes advantage of many of the tokenizer's features to create the model's inputs.

    Args:
        examples: list of :class:`~transformers.data.processors.squad.SquadExample`
        tokenizer: an instance of a child of :class:`~transformers.PreTrainedTokenizer`
        max_seq_length: The maximum sequence length of the inputs.
        max_query_length: The maximum length of the query.
        is_training: whether to create features for model evaluation or model training.
        padding_strategy: Default to "max_length". Which padding strategy to use
        return_dataset: Default False. Either 'pt' or 'tf'.
            if 'pt': returns a torch.data.TensorDataset,
            if 'tf': returns a tf.data.Dataset
        threads: multiple processing threadsa-smi


    Returns:
        list of :class:`~transformers.data.processors.squad.SquadFeatures`

    Example::

        processor = SquadV2Processor()
        examples = processor.get_dev_examples(data_dir)

        features = squad_convert_examples_to_features(
            examples=examples,
            tokenizer=tokenizer,
            max_seq_length=args.max_seq_length,
            max_query_length=args.max_query_length,
            is_training=not evaluate,
        )
    """

    # Defining helper methods
    features = []
    threads = min(threads, cpu_count())
    with Pool(
        threads,
        initializer=custom_squad_convert_example_to_features_init,
        initargs=(tokenizer,),
    ) as p:
        annotate_ = partial(
            custom_squad_convert_example_to_features,
            max_seq_length=max_seq_length,
            max_query_length=max_query_length,
            padding_strategy=padding_strategy,
            is_training=is_training,
            dataset=dataset,
        )
        features = list(
            tqdm(
                p.imap(annotate_, examples, chunksize=32),
                total=len(examples),
                desc="convert squad examples to features",
                disable=not tqdm_enabled,
            )
        )
        # print(features)
    new_features = []
    unique_id = 1000000000
    example_index = 0
    for example_features in tqdm(
        features,
        total=len(features),
        desc="add example index and unique id",
        disable=not tqdm_enabled,
    ):
        if not example_features:
            continue
        for example_feature in example_features:
            example_feature.example_index = example_index
            example_feature.unique_id = unique_id
            new_features.append(example_feature)
            unique_id += 1
        example_index += 1
    features = new_features
    del new_features
    if return_dataset == "pt":
        if not is_torch_available():
            raise RuntimeError("PyTorch must be installed to return a PyTorch dataset.")

        # Convert to Tensors and build dataset
        all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)
        all_attention_masks = torch.tensor(
            [f.attention_mask for f in features], dtype=torch.long
        )
        all_token_type_ids = torch.tensor(
            [f.token_type_ids for f in features], dtype=torch.long
        )
        all_cls_index = torch.tensor([f.cls_index for f in features], dtype=torch.long)
        all_p_mask = torch.tensor([f.p_mask for f in features], dtype=torch.float)
        all_is_impossible = torch.tensor(
            [f.is_impossible for f in features], dtype=torch.float
        )

        if not is_training:
            all_feature_index = torch.arange(all_input_ids.size(0), dtype=torch.long)
            dataset = TensorDataset(
                all_input_ids,
                all_attention_masks,
                all_token_type_ids,
                all_feature_index,
                all_cls_index,
                all_p_mask,
            )
        else:
            all_start_positions = torch.tensor(
                [f.start_position for f in features], dtype=torch.long
            )
            all_end_positions = torch.tensor(
                [f.end_position for f in features], dtype=torch.long
            )
            dataset = TensorDataset(
                all_input_ids,
                all_attention_masks,
                all_token_type_ids,
                all_start_positions,
                all_end_positions,
                all_cls_index,
                all_p_mask,
                all_is_impossible,
            )

        return features, dataset
    else:
        return features


def custom_squad_convert_example_to_features(
    example, max_seq_length, max_query_length, padding_strategy, is_training, dataset
):
    tok_to_orig_index = []
    orig_to_tok_index = []
    all_doc_tokens = []

    prefix_tokens = get_prefix_tokens(dataset, tokenizer)

    augmented_doc_tokens = prefix_tokens + example.doc_tokens

    ex_start_position = example.start_position + len(prefix_tokens)
    ex_end_position = example.end_position + len(prefix_tokens)
    # transform context by prepending athings
    # insert tokens for special purpose

    for (i, token) in enumerate(augmented_doc_tokens):
        orig_to_tok_index.append(len(all_doc_tokens))

        if i < len(prefix_tokens):
            tok_to_orig_index.append(i)
            all_doc_tokens.append(token)

        else:
            if tokenizer.__class__.__name__ in [
                "RobertaTokenizer",
                "BartTokenizer",
            ]:
                sub_tokens = tokenizer.tokenize(token, add_prefix_space=True)
            else:
                sub_tokens = tokenizer.tokenize(token)

            for sub_token in sub_tokens:
                tok_to_orig_index.append(i)
                all_doc_tokens.append(sub_token)

    if is_training:
        if example.is_yesno:
            tok_start_position = prefix_tokens.index(example.answer_text)
            tok_end_position = tok_start_position
        elif not example.is_impossible:
            tok_start_position = orig_to_tok_index[ex_start_position]
            if ex_end_position < len(augmented_doc_tokens) - 1:
                tok_end_position = orig_to_tok_index[ex_end_position + 1] - 1
            else:
                tok_end_position = len(all_doc_tokens) - 1

            (tok_start_position, tok_end_position) = _improve_answer_span(
                all_doc_tokens,
                tok_start_position,
                tok_end_position,
                tokenizer,
                example.answer_text,
            )

    truncated_query = tokenizer.encode(
        example.question_text,
        add_special_tokens=False,
        truncation=True,
        max_length=max_query_length,
    )

    # Tokenizers who insert 2 SEP tokens in-between <context> & <question> need to have special handling
    # in the way they compute mask of added tokens.
    tokenizer_type = type(tokenizer).__name__.replace("Tokenizer", "").lower()
    sequence_added_tokens = (
        tokenizer.model_max_length - tokenizer.max_len_single_sentence + 1
        if tokenizer_type in MULTI_SEP_TOKENS_TOKENIZERS_SET
        else tokenizer.model_max_length - tokenizer.max_len_single_sentence
    )
    sequence_pair_added_tokens = (
        tokenizer.model_max_length - tokenizer.max_len_sentences_pair
    )

    span_doc_tokens = all_doc_tokens

    # Define the side we want to truncate / pad and the text/pair sorting
    if tokenizer.padding_side == "right":
        texts = truncated_query
        pairs = span_doc_tokens
        truncation = TruncationStrategy.ONLY_SECOND.value
    else:
        texts = span_doc_tokens
        pairs = truncated_query
        truncation = TruncationStrategy.ONLY_FIRST.value
    # print(texts)
    encoded_dict = tokenizer.encode_plus(  # TODO(thom) update this logic
        example.question_text,
        example.context_text,
        truncation=truncation,
        padding=padding_strategy,
        max_length=max_seq_length,
        return_overflowing_tokens=False,
        return_token_type_ids=True,
    )

    paragraph_len = min(
        len(all_doc_tokens),
        max_seq_length - len(truncated_query) - sequence_pair_added_tokens,
    )

    if tokenizer.pad_token_id in encoded_dict["input_ids"]:
        if tokenizer.padding_side == "right":
            non_padded_ids = encoded_dict["input_ids"][
                : encoded_dict["input_ids"].index(tokenizer.pad_token_id)
            ]
        else:
            last_padding_id_position = (
                len(encoded_dict["input_ids"])
                - 1
                - encoded_dict["input_ids"][::-1].index(tokenizer.pad_token_id)
            )
            non_padded_ids = encoded_dict["input_ids"][last_padding_id_position + 1 :]

    else:
        non_padded_ids = encoded_dict["input_ids"]

    tokens = tokenizer.convert_ids_to_tokens(non_padded_ids)

    token_to_orig_map = {}
    for i in range(paragraph_len):
        index = (
            len(truncated_query) + sequence_added_tokens + i
            if tokenizer.padding_side == "right"
            else i
        )
        token_to_orig_map[index] = tok_to_orig_index[i]

    encoded_dict["paragraph_len"] = paragraph_len
    encoded_dict["tokens"] = tokens
    encoded_dict["token_to_orig_map"] = token_to_orig_map
    encoded_dict["truncated_query_with_special_tokens_length"] = (
        len(truncated_query) + sequence_added_tokens
    )
    encoded_dict["token_is_max_context"] = {}
    encoded_dict["start"] = 0
    encoded_dict["length"] = paragraph_len

    span = encoded_dict

    # while end

    doc_span_index = 0
    for j in range(span["paragraph_len"]):
        is_max_context = _new_check_is_max_context([span], doc_span_index, j)
        index = (
            j
            if tokenizer.padding_side == "left"
            else span["truncated_query_with_special_tokens_length"] + j
        )
        span["token_is_max_context"][index] = is_max_context

    # for span in spans:
    # Identify the position of the CLS token
    cls_index = span["input_ids"].index(tokenizer.cls_token_id)

    # p_mask: mask with 1 for token than cannot be in the answer (0 for token which can be in an answer)
    # Original TF implem also keep the classification token (set to 0)
    p_mask = np.ones_like(span["token_type_ids"])
    if tokenizer.padding_side == "right":
        p_mask[len(truncated_query) + sequence_added_tokens :] = 0
    else:
        p_mask[
            -len(span["tokens"]) : -(len(truncated_query) + sequence_added_tokens)
        ] = 0

    pad_token_indices = np.where(span["input_ids"] == tokenizer.pad_token_id)
    special_token_indices = np.asarray(
        tokenizer.get_special_tokens_mask(
            span["input_ids"], already_has_special_tokens=True
        )
    ).nonzero()

    p_mask[pad_token_indices] = 1
    p_mask[special_token_indices] = 1

    # Set the cls index to 0: the CLS index can be used for impossible answers
    p_mask[cls_index] = 0

    span_is_impossible = example.is_impossible
    start_position = 0
    end_position = 0
    if is_training and not span_is_impossible:
        # For training, if our document chunk does not contain an annotation
        # we throw it out, since there is nothing to predict.
        doc_start = span["start"]
        doc_end = span["start"] + span["length"] - 1
        out_of_span = False

        if not (tok_start_position >= doc_start and tok_end_position <= doc_end):
            out_of_span = True

        if out_of_span:
            start_position = cls_index
            end_position = cls_index
            span_is_impossible = True
        else:
            if tokenizer.padding_side == "left":
                doc_offset = 0
            else:
                doc_offset = len(truncated_query) + sequence_added_tokens

            start_position = tok_start_position - doc_start + doc_offset
            end_position = tok_end_position - doc_start + doc_offset
    features = []
    features.append(
        SquadFeatures(
            span["input_ids"],
            span["attention_mask"],
            span["token_type_ids"],
            cls_index,
            p_mask.tolist(),
            example_index=0,  # Can not set unique_id and example_index here. They will be set after multiple processing.
            unique_id=0,
            paragraph_len=span["paragraph_len"],
            token_is_max_context=span["token_is_max_context"],
            tokens=span["tokens"],
            token_to_orig_map=span["token_to_orig_map"],
            start_position=start_position,
            end_position=end_position,
            is_impossible=span_is_impossible,
            qas_id=example.qas_id,
        )
    )
    return features


class SquadProcessor(DataProcessor):
    """
    Processor for the SQuAD data set.
    Overriden by SquadV1Processor and SquadV2Processor, used by the version 1.1 and version 2.0 of SQuAD, respectively.
    """

    def get_train_examples(self, data_dir, filename=None):
        """
        Returns the training examples from the data directory.
        Args:
            data_dir: Directory containing the data files used for training and evaluating.
            filename: None by default, specify this if the training file has a different name than the original one
                which is `train-v1.1.json` and `train-v2.0.json` for squad versions 1.1 and 2.0 respectively.
        """
        if data_dir is None:
            data_dir = ""

        with open(os.path.join(data_dir, filename), "r", encoding="utf-8") as reader:
            input_data = json.load(reader)["data"]
        return self._create_examples(input_data, "train")

    def get_dev_examples(self, data_dir, filename=None):
        """
        Returns the evaluation example from the data directory.
        Args:
            data_dir: Directory containing the data files used for training and evaluating.
            filename: None by default, specify this if the evaluation file has a different name than the original one
                which is `train-v1.1.json` and `train-v2.0.json` for squad versions 1.1 and 2.0 respectively.
        """
        if data_dir is None:
            data_dir = ""

        with open(os.path.join(data_dir, filename), "r", encoding="utf-8") as reader:
            input_data = json.load(reader)["data"]
        return self._create_examples(input_data, "dev")

    def _create_examples(self, input_data, set_type, tqdm_enabled=True):
        is_training = set_type == "train"
        examples = []
        for entry in tqdm(input_data, disable=not tqdm_enabled):
            title = entry["title"]
            for paragraph in entry["paragraphs"]:
                context_text = paragraph["context"]

                for qa in paragraph["qas"]:
                    qas_id = qa["id"]
                    question_text = qa["question"]
                    start_position_character = None
                    answer_text = None
                    answers = []
                    question_type = qa.get("question_type", "none")
                    is_yesno = qa.get("is_yesno", False)
                    is_impossible = qa.get("is_impossible", False)
                    if not is_impossible:
                        if is_training:
                            answer = qa["answers"][0]
                            answer_text = answer["text"]
                            start_position_character = answer["answer_start"]
                        else:
                            answers = qa["answers"]
                    example = SquadExample(
                        qas_id=qas_id,
                        question_text=question_text,
                        context_text=context_text,
                        answer_text=answer_text,
                        start_position_character=start_position_character,
                        title=title,
                        is_yesno=is_yesno,
                        is_impossible=is_impossible,
                        question_type=question_type,
                        answers=answers,
                    )
                    examples.append(example)
        return examples


class SquadExample:
    """
    A single training/test example for the Squad dataset, as loaded from disk.
    Args:
        qas_id: The example's unique identifier
        question_text: The question string
        context_text: The context string
        answer_text: The answer string
        start_position_character: The character position of the start of the answer
        title: The title of the example
        answers: None by default, this is used during evaluation. Holds answers as well as their start positions.
        is_impossible: False by default, set to True if the example has no possible answer.
    """

    def __init__(
        self,
        qas_id,
        question_text,
        context_text,
        answer_text,
        start_position_character,
        title,
        answers=[],
        question_type=None,
        is_yesno=False,
        is_impossible=False,
    ):
        self.qas_id = qas_id
        self.question_text = question_text
        self.context_text = context_text
        self.answer_text = answer_text
        self.title = title
        self.is_impossible = is_impossible
        self.is_yesno = is_yesno
        self.question_type = question_type
        self.answers = answers

        self.start_position, self.end_position = 0, 0

        doc_tokens = []
        char_to_word_offset = []
        prev_is_whitespace = True

        # Split on whitespace so that different tokens may be attributed to their original position.
        for c in self.context_text:
            if _is_whitespace(c):
                prev_is_whitespace = True
            else:
                if prev_is_whitespace:
                    doc_tokens.append(c)
                else:
                    doc_tokens[-1] += c
                prev_is_whitespace = False
            char_to_word_offset.append(len(doc_tokens) - 1)

        self.doc_tokens = doc_tokens
        self.char_to_word_offset = char_to_word_offset

        # Start and end positions only has a value during evaluation.
        if self.is_yesno:
            self.start_position = -1
            self.end_position = -1
        else:
            if start_position_character is not None and not is_impossible:
                self.start_position = char_to_word_offset[start_position_character]
                self.end_position = char_to_word_offset[
                    min(
                        start_position_character + len(answer_text) - 1,
                        len(char_to_word_offset) - 1,
                    )
                ]


class SquadFeatures:
    """
    Single squad example features to be fed to a model.
    Those features are model-specific and can be crafted from :class:`~transformers.data.processors.squad.SquadExample`
    using the :method:`~transformers.data.processors.squad.squad_convert_examples_to_features` method.
    Args:
        input_ids: Indices of input sequence tokens in the vocabulary.
        attention_mask: Mask to avoid performing attention on padding token indices.
        token_type_ids: Segment token indices to indicate first and second portions of the inputs.
        cls_index: the index of the CLS token.
        p_mask: Mask identifying tokens that can be answers vs. tokens that cannot.
            Mask with 1 for tokens than cannot be in the answer and 0 for token that can be in an answer
        example_index: the index of the example
        unique_id: The unique Feature identifier
        paragraph_len: The length of the context
        token_is_max_context: List of booleans identifying which tokens have their maximum context in this feature object.
            If a token does not have their maximum context in this feature object, it means that another feature object
            has more information related to that token and should be prioritized over this feature for that token.
        tokens: list of tokens corresponding to the input ids
        token_to_orig_map: mapping between the tokens and the original text, needed in order to identify the answer.
        start_position: start of the answer token index
        end_position: end of the answer token index
    """

    def __init__(
        self,
        input_ids,
        attention_mask,
        token_type_ids,
        cls_index,
        p_mask,
        example_index,
        unique_id,
        paragraph_len,
        token_is_max_context,
        tokens,
        token_to_orig_map,
        start_position,
        end_position,
        is_impossible,
        qas_id: str = None,
    ):
        self.input_ids = input_ids
        self.attention_mask = attention_mask
        self.token_type_ids = token_type_ids
        self.cls_index = cls_index
        self.p_mask = p_mask

        self.example_index = example_index
        self.unique_id = unique_id
        self.paragraph_len = paragraph_len
        self.token_is_max_context = token_is_max_context
        self.tokens = tokens
        self.token_to_orig_map = token_to_orig_map

        self.start_position = start_position
        self.end_position = end_position
        self.is_impossible = is_impossible
        self.qas_id = qas_id


class SquadResult:
    """
    Constructs a SquadResult which can be used to evaluate a model's output on the SQuAD dataset.
    Args:
        unique_id: The unique identifier corresponding to that example.
        start_logits: The logits corresponding to the start of the answer
        end_logits: The logits corresponding to the end of the answer
    """

    def __init__(
        self,
        unique_id,
        start_logits,
        end_logits,
        start_top_index=None,
        end_top_index=None,
        cls_logits=None,
    ):
        self.start_logits = start_logits
        self.end_logits = end_logits
        self.unique_id = unique_id

        if start_top_index:
            self.start_top_index = start_top_index
            self.end_top_index = end_top_index
            self.cls_logits = cls_logits

src/calibration/explainers/shap/data/custom_squad_feature.py
UKPLab/CATfOOD
from transformers import RobertaTokenizer
import string
from typing import List
from transformers.models.gpt2.tokenization_gpt2 import bytes_to_unicode


def get_prefix_tokens(dataset, tokenizer):
    if dataset == "hpqa":
        return ["yes", "no", "unk", tokenizer.sep_token]
    elif dataset in ["squad", "bioasq", "newsqa", "natq", "trivia", "hotpot"]:
        return []
    elif dataset == "synth":
        return []
    elif dataset == "simple":
        return []
    elif dataset == "comp":
        return []
    else:
        raise RuntimeError("invalid dataset")


def _merge_roberta_tokens_into_words(tokenizer, feature):
    tokens = feature[0]

    decoded_each_tok = [
        bytearray([tokenizer.byte_decoder[c] for c in t]).decode(
            "utf-8", errors=tokenizer.errors
        )
        for t in tokens
    ]

    token_to_orig_map = feature.token_to_orig_map

    end_points = []
    context_start = tokens.index(tokenizer.eos_token)
    force_break = False
    for i, t in enumerate(decoded_each_tok):
        # special token
        if t in tokenizer.all_special_tokens:
            end_points.append(i)
            force_break = True
            continue

        if t in string.punctuation:
            end_points.append(i)
            force_break = True
            continue

        if force_break:
            end_points.append(i)
            force_break = False
            continue

        # if in question segment
        if i <= context_start:
            if t[0] == " ":
                decoded_each_tok[i] = t[1:]
                end_points.append(i)
        else:
            if token_to_orig_map[i] != token_to_orig_map[i - 1]:
                end_points.append(i)
    end_points.append(len(decoded_each_tok))

    # if in context segment
    segments = []
    for i in range(1, len(end_points)):
        if end_points[i - 1] == end_points[i]:
            continue
        segments.append((end_points[i - 1], end_points[i]))

    merged_tokens = []
    for s0, s1 in segments:
        merged_tokens.append("".join(decoded_each_tok[s0:s1]))

    return merged_tokens, segments


def _bpe_decode(
    tokenizer,
    tokens: List[str],
    # attributions: List
):

    byte_encoder = bytes_to_unicode()
    byte_decoder = {v: k for k, v in byte_encoder.items()}
    decoded_each_tok = [
        bytearray([byte_decoder[c] for c in t]).decode(
            encoding="utf-8", errors="replace"
        )
        for t in tokens
    ]

    end_points = []
    force_break = False
    for idx, token in enumerate(decoded_each_tok):
        # special token, punctuation, alphanumeric
        if (
            token in tokenizer.all_special_tokens
            or token in string.punctuation
            or not any([x.isalnum() for x in token.lstrip()])
            or token.lstrip == "'s"
        ):
            end_points.append(idx)
            force_break = True
            continue

        if force_break:
            end_points.append(idx)
            force_break = False
            continue

        if token[0] == " ":
            tokens[idx] = token[:]
            end_points.append(idx)

    end_points.append(len(tokens))

    segments = []
    for i in range(1, len(end_points)):
        if end_points[i - 1] == end_points[i]:
            continue
        segments.append((end_points[i - 1], end_points[i]))

    merged_tokens = []
    for s0, s1 in segments:
        merged_tokens.append("".join(decoded_each_tok[s0:s1]))

    return merged_tokens, segments


def merge_tokens_into_words(tokenizer, feature):
    if isinstance(tokenizer, RobertaTokenizer):
        return _bpe_decode(tokenizer, feature)

src/calibration/explainers/shap/data/dataset_utils.py
UKPLab/CATfOOD
""" Very heavily inspired by the official evaluation script for SQuAD version 2.0 which was
modified by XLNet authors to update `find_best_threshold` scripts for SQuAD V2.0

In addition to basic functionality, we also compute additional statistics and
plot precision-recall curves if an additional na_prob.json file is provided.
This file is expected to map question ID's to the model's predicted probability
that a question is unanswerable.
"""


import collections
import json
import math
import re
import string

from transformers.models.bert.tokenization_bert import BasicTokenizer

from transformers.utils import logging


logger = logging.get_logger(__name__)


def normalize_answer(s):
    """Lower text and remove punctuation, articles and extra whitespace."""

    def remove_articles(text):
        regex = re.compile(r"\b(a|an|the)\b", re.UNICODE)
        return re.sub(regex, " ", text)

    def white_space_fix(text):
        return " ".join(text.split())

    def remove_punc(text):
        exclude = set(string.punctuation)
        return "".join(ch for ch in text if ch not in exclude)

    def lower(text):
        return text.lower()

    return white_space_fix(remove_articles(remove_punc(lower(s))))


def get_tokens(s):
    if not s:
        return []
    return normalize_answer(s).split()


def compute_exact(a_gold, a_pred):
    return int(normalize_answer(a_gold) == normalize_answer(a_pred))


def compute_f1(a_gold, a_pred):
    gold_toks = get_tokens(a_gold)
    pred_toks = get_tokens(a_pred)
    common = collections.Counter(gold_toks) & collections.Counter(pred_toks)
    num_same = sum(common.values())
    if len(gold_toks) == 0 or len(pred_toks) == 0:
        # If either is no-answer, then F1 is 1 if they agree, 0 otherwise
        return int(gold_toks == pred_toks)
    if num_same == 0:
        return 0
    precision = 1.0 * num_same / len(pred_toks)
    recall = 1.0 * num_same / len(gold_toks)
    f1 = (2 * precision * recall) / (precision + recall)
    return f1


def get_raw_scores(examples, preds):
    """
    Computes the exact and f1 scores from the examples and the model predictions
    """
    exact_scores = {}
    f1_scores = {}

    for idx in range(len(examples["id"])):
        qas_id = examples["id"][idx][0]
        # print(examples["answers"][idx])
        answers = examples["answers"][idx][0]["text"]
        print(answers)
        gold_answers = [answer for answer in answers if normalize_answer(answer)]
        print(gold_answers)

        if not gold_answers:
            # For unanswerable questions, only correct answer is empty string
            gold_answers = [""]

        print(preds)

        if qas_id not in preds["id"]:
            print("Missing prediction for %s" % qas_id)
            continue

        prediction = preds["answers"][idx][0]["answer"]
        exact_scores[qas_id] = max(compute_exact(a, prediction) for a in gold_answers)
        f1_scores[qas_id] = max(compute_f1(a, prediction) for a in gold_answers)

    return exact_scores, f1_scores


def apply_no_ans_threshold(scores, na_probs, qid_to_has_ans, na_prob_thresh):
    new_scores = {}
    for qid, s in scores.items():
        pred_na = int(na_probs["answers"]) > na_prob_thresh
        if pred_na:
            new_scores[qid] = float(not qid_to_has_ans[qid])
        else:
            new_scores[qid] = s
    return new_scores


def make_eval_dict(exact_scores, f1_scores, qid_list=None):
    if not qid_list:
        total = len(exact_scores)
        return collections.OrderedDict(
            [
                ("exact", 100.0 * sum(exact_scores.values()) / total),
                ("f1", 100.0 * sum(f1_scores.values()) / total),
                ("total", total),
            ]
        )
    else:
        total = len(qid_list)
        return collections.OrderedDict(
            [
                ("exact", 100.0 * sum(exact_scores[k] for k in qid_list) / total),
                ("f1", 100.0 * sum(f1_scores[k] for k in qid_list) / total),
                ("total", total),
            ]
        )


def merge_eval(main_eval, new_eval, prefix):
    for k in new_eval:
        main_eval["%s_%s" % (prefix, k)] = new_eval[k]


def find_best_thresh_v2(preds, scores, na_probs, qid_to_has_ans):
    num_no_ans = sum(1 for k in qid_to_has_ans if not qid_to_has_ans[k])
    cur_score = num_no_ans
    best_score = cur_score
    best_thresh = 0.0
    qid_list = sorted(na_probs, key=lambda k: na_probs[k])
    for i, qid in enumerate(qid_list):
        if qid not in scores:
            continue
        if qid_to_has_ans[qid]:
            diff = scores[qid]
        else:
            if preds[qid]:
                diff = -1
            else:
                diff = 0
        cur_score += diff
        if cur_score > best_score:
            best_score = cur_score
            best_thresh = na_probs[qid]

    has_ans_score, has_ans_cnt = 0, 0
    for qid in qid_list:
        if not qid_to_has_ans[qid]:
            continue
        has_ans_cnt += 1

        if qid not in scores:
            continue
        has_ans_score += scores[qid]

    return (
        100.0 * best_score / len(scores),
        best_thresh,
        1.0 * has_ans_score / has_ans_cnt,
    )


def find_all_best_thresh_v2(
    main_eval, preds, exact_raw, f1_raw, na_probs, qid_to_has_ans
):
    best_exact, exact_thresh, has_ans_exact = find_best_thresh_v2(
        preds, exact_raw, na_probs, qid_to_has_ans
    )
    best_f1, f1_thresh, has_ans_f1 = find_best_thresh_v2(
        preds, f1_raw, na_probs, qid_to_has_ans
    )
    main_eval["best_exact"] = best_exact
    main_eval["best_exact_thresh"] = exact_thresh
    main_eval["best_f1"] = best_f1
    main_eval["best_f1_thresh"] = f1_thresh
    main_eval["has_ans_exact"] = has_ans_exact
    main_eval["has_ans_f1"] = has_ans_f1


def find_best_thresh(preds, scores, na_probs, qid_to_has_ans):
    num_no_ans = sum(1 for k in qid_to_has_ans if not qid_to_has_ans[k])
    cur_score = num_no_ans
    best_score = cur_score
    best_thresh = 0.0
    qid_list = sorted(na_probs, key=lambda k: na_probs[k])
    for _, qid in enumerate(qid_list):
        if qid not in scores:
            continue
        if qid_to_has_ans[qid]:
            diff = scores[qid]
        else:
            if preds[qid]:
                diff = -1
            else:
                diff = 0
        cur_score += diff
        if cur_score > best_score:
            best_score = cur_score
            best_thresh = na_probs[qid]
    return 100.0 * best_score / len(scores), best_thresh


def find_all_best_thresh(main_eval, preds, exact_raw, f1_raw, na_probs, qid_to_has_ans):
    best_exact, exact_thresh = find_best_thresh(
        preds, exact_raw, na_probs, qid_to_has_ans
    )
    best_f1, f1_thresh = find_best_thresh(preds, f1_raw, na_probs, qid_to_has_ans)

    main_eval["best_exact"] = best_exact
    main_eval["best_exact_thresh"] = exact_thresh
    main_eval["best_f1"] = best_f1
    main_eval["best_f1_thresh"] = f1_thresh


def squad_evaluate(
    examples, preds, no_answer_probs=None, no_answer_probability_threshold=1.0
):
    qas_id_to_has_answer = {
        examples["id"][0][idx]: bool(examples["answers"][idx][0]["text"])
        for idx in range(len(examples["id"]))
    }

    has_answer_qids = [
        qas_id for qas_id, has_answer in qas_id_to_has_answer.items() if has_answer
    ]
    no_answer_qids = [
        qas_id for qas_id, has_answer in qas_id_to_has_answer.items() if not has_answer
    ]

    if no_answer_probs is None:
        no_answer_probs = {k: 0.0 for k in preds if k != "id"}

    exact, f1 = get_raw_scores(examples, preds)

    exact_threshold = apply_no_ans_threshold(
        exact, no_answer_probs, qas_id_to_has_answer, no_answer_probability_threshold
    )
    f1_threshold = apply_no_ans_threshold(
        f1, no_answer_probs, qas_id_to_has_answer, no_answer_probability_threshold
    )

    evaluation = make_eval_dict(exact_threshold, f1_threshold)

    if has_answer_qids:
        has_ans_eval = make_eval_dict(
            exact_threshold, f1_threshold, qid_list=has_answer_qids
        )
        merge_eval(evaluation, has_ans_eval, "HasAns")

    if no_answer_qids:
        no_ans_eval = make_eval_dict(
            exact_threshold, f1_threshold, qid_list=no_answer_qids
        )
        merge_eval(evaluation, no_ans_eval, "NoAns")

    if no_answer_probs:
        find_all_best_thresh(
            evaluation, preds, exact, f1, no_answer_probs, qas_id_to_has_answer
        )

    return evaluation


def hotpot_evaluate(
    examples, preds, no_answer_probs=None, no_answer_probability_threshold=1.0
):
    qas_id_to_has_answer = {
        example["id"]: bool(example["answers"]) for example in examples
    }
    has_answer_qids = [
        qas_id for qas_id, has_answer in qas_id_to_has_answer.items() if has_answer
    ]
    print("has_answer_qids", len(has_answer_qids))

    exact, f1 = get_raw_scores(examples, preds)

    evaluation = make_eval_dict(exact, f1)

    return evaluation


def get_final_text(pred_text, orig_text, do_lower_case, verbose_logging=False):
    """Project the tokenized prediction back to the original text."""

    # When we created the data, we kept track of the alignment between original
    # (whitespace tokenized) tokens and our WordPiece tokenized tokens. So
    # now `orig_text` contains the span of our original text corresponding to the
    # span that we predicted.
    #
    # However, `orig_text` may contain extra characters that we don't want in
    # our prediction.
    #
    # For example, let's say:
    #   pred_text = steve smith
    #   orig_text = Steve Smith's
    #
    # We don't want to return `orig_text` because it contains the extra "'s".
    #
    # We don't want to return `pred_text` because it's already been normalized
    # (the SQuAD eval script also does punctuation stripping/lower casing but
    # our tokenizer does additional normalization like stripping accent
    # characters).
    #
    # What we really want to return is "Steve Smith".
    #
    # Therefore, we have to apply a semi-complicated alignment heuristic between
    # `pred_text` and `orig_text` to get a character-to-character alignment. This
    # can fail in certain cases in which case we just return `orig_text`.

    def _strip_spaces(text):
        ns_chars = []
        ns_to_s_map = collections.OrderedDict()
        for (i, c) in enumerate(text):
            if c == " ":
                continue
            ns_to_s_map[len(ns_chars)] = i
            ns_chars.append(c)
        ns_text = "".join(ns_chars)
        return (ns_text, ns_to_s_map)

    # We first tokenize `orig_text`, strip whitespace from the result
    # and `pred_text`, and check if they are the same length. If they are
    # NOT the same length, the heuristic has failed. If they are the same
    # length, we assume the characters are one-to-one aligned.
    tokenizer = BasicTokenizer(do_lower_case=do_lower_case)

    tok_text = " ".join(tokenizer.tokenize(orig_text))

    start_position = tok_text.find(pred_text)
    if start_position == -1:
        if verbose_logging:
            logger.info("Unable to find text: '%s' in '%s'" % (pred_text, orig_text))
        return orig_text
    end_position = start_position + len(pred_text) - 1

    (orig_ns_text, orig_ns_to_s_map) = _strip_spaces(orig_text)
    (tok_ns_text, tok_ns_to_s_map) = _strip_spaces(tok_text)

    if len(orig_ns_text) != len(tok_ns_text):
        if verbose_logging:
            logger.info(
                "Length not equal after stripping spaces: '%s' vs '%s'",
                orig_ns_text,
                tok_ns_text,
            )
        return orig_text

    # We then project the characters in `pred_text` back to `orig_text` using
    # the character-to-character alignment.
    tok_s_to_ns_map = {}
    for (i, tok_index) in tok_ns_to_s_map.items():
        tok_s_to_ns_map[tok_index] = i

    orig_start_position = None
    if start_position in tok_s_to_ns_map:
        ns_start_position = tok_s_to_ns_map[start_position]
        if ns_start_position in orig_ns_to_s_map:
            orig_start_position = orig_ns_to_s_map[ns_start_position]

    if orig_start_position is None:
        if verbose_logging:
            logger.info("Couldn't map start position")
        return orig_text

    orig_end_position = None
    if end_position in tok_s_to_ns_map:
        ns_end_position = tok_s_to_ns_map[end_position]
        if ns_end_position in orig_ns_to_s_map:
            orig_end_position = orig_ns_to_s_map[ns_end_position]

    if orig_end_position is None:
        if verbose_logging:
            logger.info("Couldn't map end position")
        return orig_text

    output_text = orig_text[orig_start_position : (orig_end_position + 1)]
    return output_text


def _get_best_indexes(logits, n_best_size):
    """Get the n-best logits from a list."""
    index_and_score = sorted(enumerate(logits), key=lambda x: x[1], reverse=True)

    best_indexes = []
    for i in range(len(index_and_score)):
        if i >= n_best_size:
            break
        best_indexes.append(index_and_score[i][0])
    return best_indexes


def _compute_softmax(scores):
    """Compute softmax probability over raw logits."""
    if not scores:
        return []

    max_score = None
    for score in scores:
        if max_score is None or score > max_score:
            max_score = score

    exp_scores = []
    total_sum = 0.0
    for score in scores:
        x = math.exp(score - max_score)
        exp_scores.append(x)
        total_sum += x

    probs = []
    for score in exp_scores:
        probs.append(score / total_sum)
    return probs


def compute_predictions_logits(
    all_examples,
    all_features,
    all_results,
    n_best_size,
    max_answer_length,
    do_lower_case,
    output_prediction_file,
    output_nbest_file,
    output_null_log_odds_file,
    verbose_logging,
    version_2_with_negative,
    null_score_diff_threshold,
    tokenizer,
    dataset="hpqa",
):
    """Write final predictions to the json file and log-odds of null if needed."""
    # if output_prediction_file:
    #     logger.info(f"Writing predictions to: {output_prediction_file}")
    # if output_nbest_file:
    #     logger.info(f"Writing nbest to: {output_nbest_file}")
    # if output_null_log_odds_file and version_2_with_negative:
    #     logger.info(f"Writing null_log_odds to: {output_null_log_odds_file}")

    example_index_to_features = collections.defaultdict(list)
    for feature in all_features:
        example_index_to_features[feature.example_index].append(feature)

    unique_id_to_result = {}
    for result in all_results:
        unique_id_to_result[result.unique_id] = result

    _PrelimPrediction = collections.namedtuple(  # pylint: disable=invalid-name
        "PrelimPrediction",
        ["feature_index", "start_index", "end_index", "start_logit", "end_logit"],
    )

    all_predictions = collections.OrderedDict()
    all_nbest_json = collections.OrderedDict()
    scores_diff_json = collections.OrderedDict()

    for (example_index, example) in enumerate(all_examples):
        features = example_index_to_features[example_index]

        prelim_predictions = []
        # keep track of the minimum score of null start+end of position 0
        score_null = 1000000  # large and positive
        min_null_feature_index = 0  # the paragraph slice with min null score
        null_start_logit = 0  # the start logit at the slice with min null score
        null_end_logit = 0  # the end logit at the slice with min null score
        for (feature_index, feature) in enumerate(features):
            result = unique_id_to_result[feature.unique_id]
            start_indexes = _get_best_indexes(result.start_logits, n_best_size)
            end_indexes = _get_best_indexes(result.end_logits, n_best_size)
            # if we could have irrelevant answers, get the min score of irrelevant
            if version_2_with_negative:
                feature_null_score = result.start_logits[0] + result.end_logits[0]
                if feature_null_score < score_null:
                    score_null = feature_null_score
                    min_null_feature_index = feature_index
                    null_start_logit = result.start_logits[0]
                    null_end_logit = result.end_logits[0]
            for start_index in start_indexes:
                for end_index in end_indexes:
                    # We could hypothetically create invalid predictions, e.g., predict
                    # that the start of the span is in the question. We throw out all
                    # invalid predictions.
                    if start_index >= len(feature.tokens):
                        continue
                    if end_index >= len(feature.tokens):
                        continue
                    if start_index not in feature.token_to_orig_map:
                        continue
                    if end_index not in feature.token_to_orig_map:
                        continue
                    if not feature.token_is_max_context.get(start_index, False):
                        continue
                    if end_index < start_index:
                        continue
                    length = end_index - start_index + 1
                    if length > max_answer_length:
                        continue
                    prelim_predictions.append(
                        _PrelimPrediction(
                            feature_index=feature_index,
                            start_index=start_index,
                            end_index=end_index,
                            start_logit=result.start_logits[start_index],
                            end_logit=result.end_logits[end_index],
                        )
                    )
        if version_2_with_negative:
            prelim_predictions.append(
                _PrelimPrediction(
                    feature_index=min_null_feature_index,
                    start_index=0,
                    end_index=0,
                    start_logit=null_start_logit,
                    end_logit=null_end_logit,
                )
            )
        prelim_predictions = sorted(
            prelim_predictions,
            key=lambda x: (x.start_logit + x.end_logit),
            reverse=True,
        )

        _NbestPrediction = collections.namedtuple(  # pylint: disable=invalid-name
            "NbestPrediction", ["text", "start_logit", "end_logit"]
        )

        seen_predictions = {}
        nbest = []

        prefix_tokens = (
            ["yes", "no", "unk", tokenizer.sep_token] if dataset == "hpqa" else []
        )
        ex_doc_tokens = prefix_tokens + example.doc_tokens
        for pred in prelim_predictions:
            if len(nbest) >= n_best_size:
                break
            feature = features[pred.feature_index]
            if pred.start_index > 0:  # this is a non-null prediction
                tok_tokens = feature.tokens[pred.start_index : (pred.end_index + 1)]
                orig_doc_start = feature.token_to_orig_map[pred.start_index]
                orig_doc_end = feature.token_to_orig_map[pred.end_index]
                orig_tokens = ex_doc_tokens[orig_doc_start : (orig_doc_end + 1)]

                tok_text = tokenizer.convert_tokens_to_string(tok_tokens)

                # tok_text = " ".join(tok_tokens)
                #
                # # De-tokenize WordPieces that have been split off.
                # tok_text = tok_text.replace(" ##", "")
                # tok_text = tok_text.replace("##", "")

                # Clean whitespace
                tok_text = tok_text.strip()
                tok_text = " ".join(tok_text.split())
                orig_text = " ".join(orig_tokens)

                final_text = get_final_text(
                    tok_text, orig_text, do_lower_case, verbose_logging
                )
                if final_text in seen_predictions:
                    continue

                seen_predictions[final_text] = True
            else:
                final_text = ""
                seen_predictions[final_text] = True

            nbest.append(
                _NbestPrediction(
                    text=final_text,
                    start_logit=pred.start_logit,
                    end_logit=pred.end_logit,
                )
            )
        # if we didn't include the empty option in the n-best, include it
        if version_2_with_negative:
            if "" not in seen_predictions:
                nbest.append(
                    _NbestPrediction(
                        text="", start_logit=null_start_logit, end_logit=null_end_logit
                    )
                )

            # In very rare edge cases we could only have single null prediction.
            # So we just create a nonce prediction in this case to avoid failure.
            if len(nbest) == 1:
                nbest.insert(
                    0, _NbestPrediction(text="empty", start_logit=0.0, end_logit=0.0)
                )

        # In very rare edge cases we could have no valid predictions. So we
        # just create a nonce prediction in this case to avoid failure.
        if not nbest:
            nbest.append(_NbestPrediction(text="empty", start_logit=0.0, end_logit=0.0))

        assert len(nbest) >= 1, "No valid predictions"

        total_scores = []
        best_non_null_entry = None
        for entry in nbest:
            total_scores.append(entry.start_logit + entry.end_logit)
            if not best_non_null_entry:
                if entry.text:
                    best_non_null_entry = entry

        probs = _compute_softmax(total_scores)

        nbest_json = []
        for (i, entry) in enumerate(nbest):
            output = collections.OrderedDict()
            output["text"] = entry.text
            output["probability"] = probs[i]
            output["start_logit"] = entry.start_logit
            output["end_logit"] = entry.end_logit
            nbest_json.append(output)

        assert len(nbest_json) >= 1, "No valid predictions"

        if not version_2_with_negative:
            all_predictions[example.qas_id] = nbest_json[0]["text"]
        else:
            # predict "" iff the null score - the score of best non-null > threshold
            score_diff = (
                score_null
                - best_non_null_entry.start_logit
                - (best_non_null_entry.end_logit)
            )
            scores_diff_json[example.qas_id] = score_diff
            if score_diff > null_score_diff_threshold:
                all_predictions[example.qas_id] = ""
            else:
                all_predictions[example.qas_id] = best_non_null_entry.text
        all_nbest_json[example.qas_id] = nbest_json

    if output_prediction_file:
        with open(output_prediction_file, "w") as writer:
            writer.write(json.dumps(all_predictions, indent=4) + "\n")

    if output_nbest_file:
        with open(output_nbest_file, "w") as writer:
            writer.write(json.dumps(all_nbest_json, indent=4) + "\n")

    if output_null_log_odds_file and version_2_with_negative:
        with open(output_null_log_odds_file, "w") as writer:
            writer.write(json.dumps(scores_diff_json, indent=4) + "\n")

    return all_predictions

src/calibration/explainers/shap/data/qa_metrics.py
UKPLab/CATfOOD
import itertools
from pprint import pprint
import os
import json
import jsonlines
import re
from tqdm import tqdm

from typing import List

from collections import OrderedDict

from datasets import (
    load_dataset,
    concatenate_datasets,
    Features,
    Value,
    Sequence,
    Dataset,
)


class PreprocessData:
    def __init__(
        self,
        dataset_name: str,
        dataset_config: str,
        cf_path: str = None,
        save_data: bool = False,
        save_path: str = "",
    ):
        """
        load datasets
        """
        dataset = load_dataset(dataset_name, dataset_config)
        if dataset_name == "squad":
            self.train_set = dataset["train"]
        self.val_set = dataset["validation"]
        self.save_data = save_data
        self.save_path = save_path
        self.cf_path = cf_path
        self.count = 0

    def get_data_info(self):
        """
        data statistics
        """
        train_samples = self.train_set.shape[0]
        val_samples = self.val_set.shape[0]
        column_names = self.train_set.column_names
        train_ex = self.train_set.select([0])[0]
        val_ex = self.val_set.select([0])[0]
        info = {
            "train_samples": train_samples,
            "val_samples": val_samples,
            "column_names": column_names,
            "train_ex": train_ex,
            "val_ex": val_ex,
        }
        return OrderedDict(info)

    def remove_trailing_space(self, example):
        # for questions in squad that have extra starting end space
        example["question"] = example["question"].strip()
        return example

    def remove_white_space(self, example):
        example["question"] = " ".join(example["question"].split())
        example["context"] = " ".join(example["context"].split())
        return example

    def processed_train_val_set(self):
        # filter unanswerable questions
        # print(self.train_set.filter(lambda x: len(x["answers"]["text"]) != 1))
        # self.train_set = self.train_set.select(range(1000))
        # self.val_set = self.val_set.select(range(10))
        if self.cf_path is not None:
            self.train_set = self._add_counterfactuals()
        answerable_train_set = self.train_set.filter(
            lambda x: len(x["answers"]["text"]) != 0
        )
        answerable_train_set = answerable_train_set.map(self.remove_trailing_space)
        answerable_train_set = answerable_train_set.map(self.remove_white_space)

        answerable_val_set = self.val_set.filter(
            lambda x: len(x["answers"]["text"]) != 0
        )
        answerable_val_set = answerable_val_set.map(self.remove_trailing_space)
        answerable_val_set = answerable_val_set.map(self.remove_white_space)
        return answerable_train_set, answerable_val_set

    def processed_val_set(self):
        # filter unanswerable questions
        # print(self.train_set.filter(lambda x: len(x["answers"]["text"]) != 1))
        answerable_val_set = self.val_set.filter(
            lambda x: len(x["answers"]["text"]) != 0
        )
        answerable_val_set = answerable_val_set.map(self.remove_trailing_space)
        answerable_val_set = answerable_val_set.map(self.remove_white_space)
        return answerable_val_set

    def sanitize_ptb_tokenized_string(self, text: str) -> str:
        """
        Sanitizes string that was tokenized using PTBTokenizer
        """
        tokens = text.split(" ")
        if len(tokens) == 0:
            return text

        # Replace quotation marks and parentheses
        token_map = {
            "``": '"',
            "''": '"',
            "-lrb-": "(",
            "-rrb-": ")",
            "-lsb-": "[",
            "-rsb-": "]",
            "-lcb-": "{",
            "-rcb-": "}",
            "<s>": "",
            "</s>": "",
        }

        # Merge punctuation with previous tokens
        punct_forward = {"`", "$", "#"}
        punct_backward = {".", ",", "!", "?", ":", ";", "%", "'"}

        # Exact matches that get merged forward or backward
        em_forward = {"(", "[", "{"}
        em_backward = {"n't", "na", ")", "]", "}"}

        new_tokens: List[str] = []

        merge_fwd = False
        for i, orig_token in enumerate(tokens):
            tokens[i] = (
                token_map[orig_token.lower()]
                if orig_token.lower() in token_map
                else orig_token
            )
            new_token = tokens[i].lower()

            # merge_fwd was set by previous token, so it should be prepended to current token
            if merge_fwd:
                tokens[i] = tokens[i - 1] + tokens[i]

            if len(tokens[i]) == 0:
                continue

            # Special cases for `` and '', those tells us if " is the start or end of a quotation.
            # Also always merge tokens starting with ' backward and don't merge back if we just merged forward
            merge_bckwd = not merge_fwd and (
                orig_token == "''"
                or new_token in em_backward
                or new_token.startswith("'")
                or all(c in punct_backward for c in new_token)
            )
            merge_fwd = (
                orig_token == "``"
                or new_token in em_forward
                or all(c in punct_forward for c in new_token)
            )

            if merge_bckwd and new_tokens:
                new_tokens[-1] += tokens[i]
            elif not new_tokens or not merge_fwd or i == len(tokens) - 1:
                new_tokens.append(tokens[i])

        return " ".join(new_tokens)

    def _convert_answer_column(self, example):
        """
        convert answer column to squad format
        """

        context = example["context"].strip()
        # remove extra white space
        context = re.sub(r"\s+", " ", context)
        example["context"] = context

        answer = example["answer"]["text"]
        # handle spaces between special characters
        special_chars = ["$", "@"]
        if any(char in answer for char in special_chars):
            for char in special_chars:
                answer = answer.replace(f"{char} ", char)
        # remove spaces between special characters (quotations)
        answer = re.sub(r'"\s*([^"]*?)\s*"', r'"\1"', answer)
        answer = re.sub(r" - ", "-", answer)

        # answer = answer.replace("$ ", "$")
        answer_start = str(context.lower()).find(str(answer.lower()))
        if answer_start == -1:
            self.count += 1
            answer_start = str(context.lower()).find(
                self.sanitize_ptb_tokenized_string(answer).lower()
            )
            if answer_start != -1:
                answer = str(self.sanitize_ptb_tokenized_string(answer))
                self.count -= 1
            else:
                print("=====================================")
                print(context)
                print(answer)
                print("------")
                print(context.lower())
                print(self.sanitize_ptb_tokenized_string(answer))
                print(
                    context.lower().find(
                        self.sanitize_ptb_tokenized_string(answer).lower()
                    )
                )
                # print(context.lower().index(self.sanitize_ptb_tokenized_string(answer).lower()))
                print("=====================================")
                answer = ""

        example["answers"] = {"answer_start": [answer_start], "text": [answer]}
        return example

    def _add_counterfactuals(self):
        """
        add counterfactuals to data
        """
        counterfactuals_dataset = load_dataset("json", data_files=self.cf_path)
        # print(len(counterfactuals_dataset[]))
        new_title_column = [""] * len(counterfactuals_dataset["train"])
        counterfactuals_dataset = counterfactuals_dataset["train"].add_column(
            "title", new_title_column
        )
        # rename answer column
        # counterfactuals_dataset = counterfactuals_dataset.rename_column("answer", "answers")
        counterfactuals_dataset = counterfactuals_dataset.map(
            self._convert_answer_column
        )
        print("Count:", self.count)
        counterfactuals_dataset = counterfactuals_dataset.remove_columns("answer")
        # remove samples with no answer
        counterfactuals_dataset = counterfactuals_dataset.filter(
            lambda x: x["answers"]["answer_start"][0] != -1
        )
        counterfactuals_dataset = counterfactuals_dataset.cast(
            Features(
                {
                    "id": Value("string"),
                    "title": Value("string"),
                    "context": Value("string"),
                    "question": Value("string"),
                    "answers": Sequence(
                        {"answer_start": Value("int32"), "text": Value("string"),}
                    ),
                }
            )
        )
        self.train_set = self.train_set.cast(
            Features(
                {
                    "id": Value("string"),
                    "title": Value("string"),
                    "context": Value("string"),
                    "question": Value("string"),
                    "answers": Sequence(
                        {"answer_start": Value("int32"), "text": Value("string"),}
                    ),
                }
            )
        )

        # ex = counterfactuals_dataset[0]
        # print("CF: ", ex)
        # print(ex["context"].strip().lower())
        # print(ex["context"].strip().lower().find(ex["answers"]["text"][0]))
        # ex = self.train_set[0]
        # print("Train: ", ex)
        # print(ex["context"].find(ex["answers"]["text"][0]))
        # add counterfactuals to example data
        merged_dataset = concatenate_datasets([self.train_set, counterfactuals_dataset])
        return merged_dataset


def get_dev_examples(data_dir, file_name):
    """
    Data loader for dev set of trivia qa and hotpot qa
    """

    with open(os.path.join(data_dir, file_name), "r", encoding="utf-8") as reader:
        input_data = json.load(reader)["data"]
    sum = 0
    for i in input_data:
        sum += len(i["paragraphs"])
    print(sum)

    examples = []
    for entry in tqdm(input_data):
        title = entry["title"]
        for paragraph in entry["paragraphs"]:
            context_text = paragraph["context"]

            for qa in paragraph["qas"]:
                qas_id = qa["id"]
                question_text = qa["question"]
                start_position_character = None
                answer_text = None
                answers = []
                question_type = qa.get("question_type", "none")
                is_yesno = qa.get("is_yesno", False)
                is_impossible = qa.get("is_impossible", False)
                if not is_impossible:
                    answer = qa["answers"][0]
                    answer_text = answer["text"]
                    start_position_character = answer["answer_start"]
                else:
                    answers = qa["answers"]
                example = {
                    "qas_id": qas_id,
                    "question": question_text,
                    "context": context_text,
                    "answers": [answer_text],
                    "start_position_character": start_position_character,
                    # "answers": answers,
                }
                examples.append(example)
    return examples


def get_dev_examples_hf():
    """
    Data loader for dev set of trivia qa and hotpot qa
    """

    with open(
        os.path.join(
            "/storage/ukp/work/sachdeva/research_projects/exp_calibration/src/data/dev_trivia.json"
        ),
        "r",
        encoding="utf-8",
    ) as reader:
        input_data = json.load(reader)["data"]
    sum = 0
    for i in input_data:
        sum += len(i["paragraphs"])
    print(sum)

    examples = []
    for entry in tqdm(input_data):
        title = entry["title"]
        for paragraph in entry["paragraphs"]:
            context_text = paragraph["context"]

            for qa in paragraph["qas"]:
                qas_id = qa["id"]
                question_text = qa["question"]
                start_position_character = None
                answer_text = None
                answers = []
                question_type = qa.get("question_type", "none")
                is_yesno = qa.get("is_yesno", False)
                is_impossible = qa.get("is_impossible", False)
                if not is_impossible:
                    answer = qa["answers"][0]
                    answer_text = answer["text"]
                    start_position_character = answer["answer_start"]
                else:
                    answers = qa["answers"]
                yield {
                    "id": qas_id,
                    "question": question_text,
                    "context": context_text,
                    "answers": [answer_text],
                    "start_position_character": start_position_character,
                    # "answers": answers,
                }


def get_dev_samples_mrqa(data_path):
    examples = []
    with jsonlines.open(data_path) as reader:
        for example in tqdm(reader):
            # skip first header row
            if "header" in example:
                continue
            context_text = example["context"]
            questions = example["qas"]
            for question in questions:
                id = question["qid"]
                question_text = question["question"]
                answers = question["detected_answers"][0]
                answer_text = answers["text"]
                char_spans = answers["char_spans"]
                token_spans = answers["token_spans"]
                example = {
                    "qas_id": id,
                    "question_text": question_text,
                    "context_text": context_text,
                    "answer_text": answer_text,
                    "char_spans": char_spans,
                    "token_spans": token_spans,
                }
                examples.append(example)
        return examples


if __name__ == "__main__":
    # save_dir = "../data/squad/train_data/"
    # BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"
    BASE_PATH = "//"
    # dataloader = PreprocessData(
    #     "squad", "plain_text",
    #     # cf_path=None,
    #     cf_path=BASE_PATH+"src/data/squad/counterfactuals_with_answers_filtered_final_18k.jsonl",
    #     save_data=False,
    #     save_path="../../"
    # )
    # dataset = dataloader._add_counterfactuals()
    # print(len(dataset))
    # print(dataset[10:20])

    # ex = get_dev_examples(BASE_PATH+"src/data/", "dev_trivia.json")
    # print(len(ex))
    # print(ex[0])
    def group_batch(batch):
        return {k: [v] for k, v in batch.items()}

    dataset = Dataset.from_generator(get_dev_examples_hf)
    dataset = dataset.select(range(610, len(dataset)))
    # print(dataset[0:10])
    print(len(dataset))
    eval_dataloader = dataset.map(group_batch, batched=True, batch_size=2)
    # print(len(dataset))
    # print(dataset[0:10])
    #
    # print(len(eval_dataloader))
    # print(eval_dataloader[0])
    # import re
    # x = "from 1216 to 1227 nearly £10,000 was spent on the tower of london; in this period, only the work at windsor castle cost more (£15,000). most of the work was focused on the palatial buildings of the innermost ward. the tradition of whitewashing the white tower (from which it derives its name) began in 1240. beginning around 1238, the castle was expanded to the east, north, and north-west. the work lasted through the reign of henry iii and into that of edward i, interrupted occasionally by civil unrest. new creations included a new defensive perimeter, studded with towers, while on the west, north, and east sides, where the wall was not defended by the river, a defensive ditch was dug. the eastern extension took the castle beyond the bounds of the old roman settlement, marked by the city wall which had been incorporated into the castle's defences. the tower had long been a symbol of oppression, despised by londoners, and henry's building programme was unpopular. so when the gatehouse collapsed in 1240, the locals celebrated the setback. the expansion caused disruption locally and £166 was paid to st katherine's hospital and the prior of holy trinity in compensation."
    # # x = str.unicode(x, "UTF-8")
    # # x = x.replace("&nbsp;", " ")
    # x = re.sub(r'\s+', ' ', x)
    # # x = " ".join(x.split(" "))
    # print(x)
    # print(x.find("henry iii"))

    # example = dataset[0]
    # print(example)
    #
    # print(example["context"].find(example["answers"]["text"][0]))

    # print(dataset[-1])

    # print(len(train), len(test))
    # print(train[0])
    # test = dataloader._add_counterfactuals()
    # for ex in set:
    #     print(ex)
    #     print([ex["question"]])
    #     print([ex["context"]])
    #     break

    # dev_examples = get_dev_samples_mrqa("data/NewsQA.jsonl")

    # context = 'Super Bowl 50 was an American football game to determine the champion of the National Football League (NFL) for the 2015 season. The American Football Conference (AFC) champion Denver Broncos defeated the National Football Conference (NFC) champion Carolina Panthers 24–10 to earn their third Super Bowl title. The game was played on February 7, 2016, at Levi\'s Stadium in the San Francisco Bay Area at Santa Clara, California. As this was the 50th Super Bowl, the league emphasized the "golden anniversary" with various gold-themed initiatives, as well as temporarily suspending the tradition of naming each Super Bowl game with Roman numerals (under which the game would have been known as "Super Bowl L"), so that the logo could prominently feature the Arabic numerals 50.'
    # answer = 'Denver Broncos'
    # question = 'Which NFL team represented the AFC at Super Bowl 50?'
    # # get answer start position
    # answer_start = context.find(answer)
    # print(answer_start)

src/calibration/explainers/shap/dataloader.py
UKPLab/CATfOOD
import numpy as np
import os
import json
import collections
from tqdm import tqdm
from typing import List, Tuple

import torch
from transformers import (
    PreTrainedModel,
    PreTrainedTokenizer,
    AutoTokenizer,
    RobertaForQuestionAnswering,
    # BertAdapterModel,
    # RobertaAdapterModel
)
import dataloader

# BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"
BASE_PATH = "//"


class Inference:
    def __init__(
        self, model: PreTrainedModel, tokenizer: PreTrainedTokenizer,
    ):
        self.model = model
        self.tokenizer = tokenizer

        if self.model.config.model_type == "gpt2":
            self.ref_token_id = self.tokenizer.eos_token_id
        else:
            self.ref_token_id = self.tokenizer.pad_token_id

        self.sep_token_id = (
            self.tokenizer.sep_token_id
            if self.tokenizer.sep_token_id is not None
            else self.tokenizer.eos_token_id
        )
        self.cls_token_id = (
            self.tokenizer.cls_token_id
            if self.tokenizer.cls_token_id is not None
            else self.tokenizer.bos_token_id
        )

        self.model_prefix = model.base_model_prefix

        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)

    def _ensure_tensor_on_device(self, **inputs):
        """
        Ensure PyTorch tensors are on the specified device.

        Args:
            inputs (keyword arguments that should be :obj:`torch.Tensor`): The tensors to place on :obj:`self.device`.

        Return:
            :obj:`Dict[str, torch.Tensor]`: The same as :obj:`inputs` but on the proper device.
        """
        return {name: tensor.to(self.model.device) for name, tensor in inputs.items()}

    def encode(
        self, inputs: list = None, add_special_tokens: bool = True, return_tensors=None
    ):
        """
        Encode inputs using the model tokenizer
        Args:
            inputs: question, context pair as a list
            add_special_tokens: where to add CLS, SEP tokens
            return_tensors: whether to return tensors
        Return:
            tokenized inputs
        """
        return self.tokenizer(
            inputs,
            add_special_tokens=add_special_tokens,
            return_tensors=return_tensors,
            padding=True,
            truncation=True,
            max_length=512,
        )

    def decode(self, input_ids: torch.Tensor, skip_special_tokens: bool) -> List[str]:
        """
        Decode received input_ids into a list of word tokens.
        Args:
            input_ids (torch.Tensor): Input ids representing
            word tokens for a sentence/document.
        """
        return self.tokenizer.convert_ids_to_tokens(
            input_ids[0], skip_special_tokens=skip_special_tokens
        )

    def _predict(self, request, batch_size: int = 1) -> tuple:
        """
        Inference on the input.
        Args:
            inputs: list of question, context input
        Returns:
             The model outputs and optionally the input features
        """
        all_predictions = []
        # self.model.to("cuda" if torch.cuda.is_available() else "cpu")

        features = self.encode(request, add_special_tokens=True, return_tensors="pt",)
        print("tokens", features.tokens())
        print("input_ids", features)

        for start_idx in range(0, len(request), batch_size):
            with torch.no_grad():
                input_features = {
                    k: features[k][start_idx : start_idx + batch_size]
                    for k in features.keys()
                }
                input_features = self._ensure_tensor_on_device(**input_features)
                predictions = self.model(**input_features)
                all_predictions.append(predictions)
        keys = all_predictions[0].keys()
        final_prediction = {}
        for key in keys:
            if isinstance(all_predictions[0][key], tuple):
                tuple_of_lists = list(
                    zip(
                        *[
                            [
                                torch.stack(p).to(self.device)
                                if isinstance(p, tuple)
                                else p.to(self.device)
                                for p in tpl[key]
                            ]
                            for tpl in all_predictions
                        ]
                    )
                )
                final_prediction[key] = tuple(torch.cat(l) for l in tuple_of_lists)
            else:
                final_prediction[key] = torch.cat(
                    [p[key].to(self.device) for p in all_predictions]
                )
        # print(final_prediction)
        # print(features)
        return final_prediction, features

    def question_answering(self, request, top_k: int = 1):
        """
        Span-based question answering for a given question and context.
        We expect the input to use the (question, context) format for the text pairs.
        Args:
          request: the prediction request
          top_k: num preds

        """

        def decode(
            start_: np.ndarray,
            end_: np.ndarray,
            topk: int,
            max_answer_len: int,
            undesired_tokens_: np.ndarray,
        ) -> Tuple:
            """
            Take the output of any :obj:`ModelForQuestionAnswering` and will generate probabilities
            for each span to be the actual answer.

            In addition, it filters out some unwanted/impossible cases like answer len being greater
            than max_answer_len or answer end position being before the starting position. The method
            supports output the k-best answer through the topk argument.

            Args:
                start_ (:obj:`np.ndarray`): Individual start probabilities for each token.
                end (:obj:`np.ndarray`): Individual end_ probabilities for each token.
                topk (:obj:`int`): Indicates how many possible answer span(s) to extract from the model output.
                max_answer_len (:obj:`int`): Maximum size of the answer to extract from the model's output.
                undesired_tokens_ (:obj:`np.ndarray`): Mask determining tokens that can be part of the answer
            """
            # Ensure we have batch axis
            if start_.ndim == 1:
                start_ = start_[None]

            if end_.ndim == 1:
                end_ = end_[None]

            # Compute the score of each tuple(start_, end_) to be the real answer
            outer = np.matmul(np.expand_dims(start_, -1), np.expand_dims(end_, 1))

            # Remove candidate with end_ < start_ and end_ - start_ > max_answer_len
            candidates = np.tril(np.triu(outer), max_answer_len - 1)

            #  Inspired by Chen & al. (https://github.com/facebookresearch/DrQA)
            scores_flat = candidates.flatten()
            if topk == 1:
                idx_sort = [np.argmax(scores_flat)]
            elif len(scores_flat) < topk:
                idx_sort = np.argsort(-scores_flat)
            else:
                idx = np.argpartition(-scores_flat, topk)[0:topk]
                idx_sort = idx[np.argsort(-scores_flat[idx])]

            starts_, ends_ = np.unravel_index(idx_sort, candidates.shape)[1:]
            desired_spans = np.isin(starts_, undesired_tokens_.nonzero()) & np.isin(
                ends_, undesired_tokens_.nonzero()
            )
            starts_ = starts_[desired_spans]
            ends_ = ends_[desired_spans]
            scores_ = candidates[0, starts_, ends_]

            return starts_, ends_, scores_

        predictions, features = self._predict(request)
        start_idx = torch.argmax(predictions["start_logits"])
        end_idx = torch.argmax(predictions["end_logits"])
        answer_start = torch.tensor([start_idx])
        answer_end = torch.tensor([end_idx])

        # print(features)
        task_outputs = {"answers": []}
        for idx, (start, end, (_, context)) in enumerate(
            zip(predictions["start_logits"], predictions["end_logits"], request)
        ):
            start = start.cpu().detach().numpy()
            end = end.cpu().detach().numpy()
            # Ensure padded tokens & question tokens cannot belong to the set of candidate answers.
            question_tokens = np.abs(
                np.array([s != 1 for s in features.sequence_ids(idx)]) - 1
            )
            # Unmask CLS token for 'no answer'
            question_tokens[0] = 1
            undesired_tokens = question_tokens & features["attention_mask"][idx].numpy()

            # Generate mask
            undesired_tokens_mask = undesired_tokens == 0.0

            # Make sure non-context indexes in the tensor cannot contribute to the softmax
            start = np.where(undesired_tokens_mask, -10000.0, start)
            end = np.where(undesired_tokens_mask, -10000.0, end)

            start = np.exp(
                start - np.log(np.sum(np.exp(start), axis=-1, keepdims=True))
            )
            end = np.exp(end - np.log(np.sum(np.exp(end), axis=-1, keepdims=True)))

            # Get score for 'no answer' then mask for decoding step (CLS token
            no_answer_score = (start[0] * end[0]).item()
            start[0] = end[0] = 0.0

            starts, ends, scores = decode(start, end, top_k, 128, undesired_tokens)
            enc = features[idx]
            answers = [
                {
                    "score": score.item(),
                    "start": enc.word_to_chars(enc.token_to_word(s), sequence_index=1)[
                        0
                    ],
                    "end": enc.word_to_chars(enc.token_to_word(e), sequence_index=1)[1],
                    "answer": context[
                        enc.word_to_chars(enc.token_to_word(s), sequence_index=1)[
                            0
                        ] : enc.word_to_chars(enc.token_to_word(e), sequence_index=1)[1]
                    ],
                }
                for s, e, score in zip(starts, ends, scores)
            ]

            answers.append(
                {"score": no_answer_score, "start": 0, "end": 0, "answer": ""}
            )
            answers = sorted(answers, key=lambda x: x["score"], reverse=True)[:top_k]
            task_outputs["answers"].append(answers)
            task_outputs["answer_start"] = answer_start.cpu().detach().numpy()
            task_outputs["answer_end"] = answer_end.cpu().detach().numpy()
        return task_outputs


if __name__ == "__main__":
    # base_model = "roberta-base"
    # adapter_model = "AdapterHub/roberta-base-pf-squad"
    # model = RobertaAdapterModel.from_pretrained(base_model)
    # tokenizer = AutoTokenizer.from_pretrained(base_model)
    # adapter_name = model.load_adapter(adapter_model, source="hf")
    # model.active_adapters = adapter_name

    model_path = BASE_PATH + "src/checkpoints/roberta-squad-with-easy-0.75"
    # model_path = BASE_PATH+"roberta-squad-with-cfs-15k"
    model = RobertaForQuestionAnswering.from_pretrained(model_path)
    tokenizer = AutoTokenizer.from_pretrained("roberta-base")

    inf = Inference(model=model, tokenizer=tokenizer)

    # data = dataloader.PreprocessData("squad_adversarial", "AddSent", save_data=False, save_path="../../")

    # batch_data = [for ex in data.processed_val_set()]
    # def group_batch(batch):
    #     return {k: [v] for k, v in batch.items()}
    # data = data.processed_val_set().map(group_batch, batched=True, batch_size=2)
    # c=0
    # for ex in data:
    #     print(ex)
    #     c+=1
    #     if c==2:
    #         break
    # print(inf.predict(ex))

    # print(data)

    data = dataloader.get_dev_examples("./src/data", "dev_trivia.json")
    outputs = list()

    def remove_white_space(example):
        example["question_text"] = " ".join(example["question_text"].split())
        example["context_text"] = " ".join(example["context_text"].split())
        return example

    all_nbest_json = collections.OrderedDict()
    _NbestPrediction = collections.namedtuple(  # pylint: disable=invalid-name
        "NbestPrediction", ["text", "score", "answer_start", "answer_end"]
    )

    seen_predictions = {}
    nbest = []

    c = 0
    for ex in tqdm(data):
        # for ex in tqdm(data.processed_val_set()):
        # print(ex)
        ex = remove_white_space(ex)
        try:
            prediction = inf.question_answering(
                request=[[ex["question_text"], ex["context_text"],]], top_k=20
            )
            # break
            # print(prediction)
            for i, pred in enumerate(prediction["answers"][0]):
                # print(pred)
                nbest.append(
                    _NbestPrediction(
                        text=pred["answer"],
                        score=pred["score"],
                        answer_start=pred["start"],
                        answer_end=pred["end"],
                    )
                )

            nbest_json = []
            for (i, entry) in enumerate(nbest):
                output = collections.OrderedDict()
                output["text"] = entry.text
                output["probability"] = entry.score
                output["answer_start"] = entry.answer_start
                output["answer_end"] = entry.answer_end
                nbest_json.append(output)

            nbest = []
            all_nbest_json[ex["id"]] = nbest_json
            # print(all_nbest_json)
            c += 1
        except Exception as e:
            print(f"Unable to get prediction: {e}")
    #     # if c == 10:
    #     #     break
    #
    path = f"{BASE_PATH}src/data/trivia_qa/"
    output_nbest_file = os.path.join(
        path, "nbest_predictions_roberta_easy_cfs_0.75.json"
    )
    with open(output_nbest_file, "w") as writer:
        writer.write(json.dumps(all_nbest_json, indent=4) + "\n")

    # print(outputs)
    #
    # cf_path = BASE_PATH + "src/data/squad_adversarial/rag_counterfactuals_turk0_last.jsonl"
    # with jsonlines.open(cf_path) as reader:
    #     for ex in tqdm(reader):
    #         try:
    #             prediction = inf.question_answering(
    #                 request=[[
    #                     ex["question"],
    #                     ex["context"],
    #                 ]],
    #                 top_k=20
    #             )
    #             result = {
    #                 "id": ex["id"].replace("_", "-cf-"),
    #                 "question": ex["question"],
    #                 "context": ex["context"],
    #                 "gold_text": [ex["answer"]],
    #                 "pred_text": prediction["answers"],
    #                 "answer_start": prediction["answer_start"],
    #                 "answer_end": prediction["answer_end"]
    #             }
    #             outputs.append(result)
    #             c+=1
    #         except Exception as e:
    #             print(f"Unable to get cf prediction: {e}")
    #             print(ex)
    #
    #     print(f"Processed {c} instances of counterfactual data")

    # df = pd.DataFrame(outputs)
    # df.to_csv(f"{BASE_PATH}src/data/squad_adv_new/outputs.csv")

    # Trivia QA
    # data = dataloader.get_dev_examples(BASE_PATH+"src/data", "dev_hotpot.json")
    # data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/NewsQA.jsonl")
    # data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/BioASQ-dev.jsonl")
    # data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/NaturalQuestionsShort.jsonl")
    # def remove_white_space(example):
    #     example["question_text"] = ' '.join(example["question_text"].split())
    #     example["context_text"] = ' '.join(example["context_text"].split())
    #     return example
    #
    # outputs = list()
    # c = 0
    # for ex in tqdm(data):
    #     # print(ex)
    #     ex = remove_white_space(ex)
    #     try:
    #         prediction = inf.question_answering(
    #             request=[[
    #                 ex["question_text"],
    #                 ex["context_text"],
    #             ]],
    #             top_k=20
    #         )
    #         result = {
    #             "id": ex["qas_id"],
    #             "question": ex["question_text"],
    #             "context": ex["context_text"],
    #             "gold_text": ex["answer_text"],
    #             "pred_text": prediction["answers"],
    #             "answer_start": prediction["answer_start"],
    #             "answer_end": prediction["answer_end"]
    #         }
    #         outputs.append(result)
    #         c += 1
    #     except Exception as e:
    #         print(f"Unable to get prediction: {e}")
    #
    # cf_path = BASE_PATH + "src/data/natural_questions/rag_counterfactuals_turk0.jsonl"
    # with jsonlines.open(cf_path) as reader:
    #     for ex in tqdm(reader):
    #         try:
    #             prediction = inf.question_answering(
    #                 request=[[
    #                     ex["question"],
    #                     ex["context"],
    #                 ]],
    #                 top_k=20
    #             )
    #             result = {
    #                 "id": ex["id"].replace("_", "-cf-"),
    #                 "question": ex["question"],
    #                 "context": ex["context"],
    #                 "gold_text": [ex["answer"]],
    #                 "pred_text": prediction["answers"],
    #                 "answer_start": prediction["answer_start"],
    #                 "answer_end": prediction["answer_end"]
    #             }
    #             outputs.append(result)
    #             c += 1
    #         except Exception as e:
    #             print(f"Unable to get cf prediction: {e}")
    #             print(ex)
    #
    #     print(f"Processed {c} instances of counterfactual data")
    #
    # df = pd.DataFrame(outputs)
    # df.to_csv(f"{BASE_PATH}src/data/natural_questions/outputs_roberta.csv")

src/calibration/explainers/shap/inference.py
UKPLab/CATfOOD
from transformers import RobertaForQuestionAnswering
from transformers.models.roberta.modeling_roberta import (
    RobertaPreTrainedModel,
    RobertaPooler,
    RobertaModel,
    RobertaEncoder,
    RobertaAttention,
    RobertaSelfAttention,
    RobertaEmbeddings,
    RobertaAttention,
    RobertaSelfOutput,
    RobertaIntermediate,
    RobertaOutput,
    RobertaLayer,
    QuestionAnsweringModelOutput,
)

import numpy as np
from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
from torch.nn import CrossEntropyLoss
from tqdm import tqdm

from transformers.modeling_utils import (
    PreTrainedModel,
    apply_chunking_to_forward,
    find_pruneable_heads_and_indices,
    prune_linear_layer,
)
import torch
from torch import nn
import torch.nn.functional as F
import math

# from int_grad.ig_models import IGRobertaLayer

# Copied from transformers.modeling_bert.BertSelfAttention with Bert->Roberta
class ProbeRobertaSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(
            config, "embedding_size"
        ):
            raise ValueError(
                "The hidden size (%d) is not a multiple of the number of attention "
                "heads (%d)" % (config.hidden_size, config.num_attention_heads)
            )

        self.num_attention_heads = config.num_attention_heads
        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
        self.all_head_size = self.num_attention_heads * self.attention_head_size

        self.query = nn.Linear(config.hidden_size, self.all_head_size)
        self.key = nn.Linear(config.hidden_size, self.all_head_size)
        self.value = nn.Linear(config.hidden_size, self.all_head_size)

        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)

    def transpose_for_scores(self, x):
        new_x_shape = x.size()[:-1] + (
            self.num_attention_heads,
            self.attention_head_size,
        )
        x = x.view(*new_x_shape)
        return x.permute(0, 2, 1, 3)

    # possible probation
    # disable a layer, i.e., zeroing out all the links
    # mask a layer, i.e., mask some attention
    # specify attention
    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        output_attentions=False,
        input_attentions=None,
        link_mask=None,
    ):
        mixed_query_layer = self.query(hidden_states)

        # If this is instantiated as a cross-attention module, the keys
        # and values come from an encoder; the attention mask needs to be
        # such that the encoder's padding tokens are not attended to.
        if encoder_hidden_states is not None:
            mixed_key_layer = self.key(encoder_hidden_states)
            mixed_value_layer = self.value(encoder_hidden_states)
            attention_mask = encoder_attention_mask
        else:
            mixed_key_layer = self.key(hidden_states)
            mixed_value_layer = self.value(hidden_states)

        query_layer = self.transpose_for_scores(mixed_query_layer)
        key_layer = self.transpose_for_scores(mixed_key_layer)
        value_layer = self.transpose_for_scores(mixed_value_layer)

        if input_attentions is None:
            # Take the dot product between "query" and "key" to get the raw attention scores.
            attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
            attention_scores = attention_scores / math.sqrt(self.attention_head_size)
            if attention_mask is not None:
                # Apply the attention mask is (precomputed for all layers in RobertaModel forward() function)
                attention_scores = attention_scores + attention_mask

            # Normalize the attention scores to probabilities.
            attention_probs = nn.Softmax(dim=-1)(attention_scores)
        else:
            attention_probs = input_attentions

        if link_mask is not None:
            attention_probs = attention_probs * link_mask
        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        context_layer = torch.matmul(attention_probs, value_layer)

        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(*new_context_layer_shape)

        outputs = (
            (context_layer, attention_probs) if output_attentions else (context_layer,)
        )
        return outputs


# Copied from transformers.modeling_bert.BertAttention with Bert->Roberta
class ProbeRobertaAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.self = ProbeRobertaSelfAttention(config)
        self.output = RobertaSelfOutput(config)
        self.pruned_heads = set()

    def prune_heads(self, heads):
        if len(heads) == 0:
            return
        heads, index = find_pruneable_heads_and_indices(
            heads,
            self.self.num_attention_heads,
            self.self.attention_head_size,
            self.pruned_heads,
        )

        # Prune linear layers
        self.self.query = prune_linear_layer(self.self.query, index)
        self.self.key = prune_linear_layer(self.self.key, index)
        self.self.value = prune_linear_layer(self.self.value, index)
        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)

        # Update hyper params and store pruned heads
        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
        self.self.all_head_size = (
            self.self.attention_head_size * self.self.num_attention_heads
        )
        self.pruned_heads = self.pruned_heads.union(heads)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        output_attentions=False,
        input_attentions=None,
        link_mask=None,
    ):
        self_outputs = self.self(
            hidden_states,
            attention_mask,
            head_mask,
            encoder_hidden_states,
            encoder_attention_mask,
            output_attentions,
            input_attentions,
            link_mask,
        )
        attention_output = self.output(self_outputs[0], hidden_states)
        outputs = (attention_output,) + self_outputs[
            1:
        ]  # add attentions if we output them
        return outputs


# Copied from transformers.modeling_bert.BertLayer with Bert->Roberta
class ProbeRobertaLayer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.chunk_size_feed_forward = config.chunk_size_feed_forward
        self.seq_len_dim = 1
        self.attention = ProbeRobertaAttention(config)
        self.is_decoder = config.is_decoder
        self.add_cross_attention = config.add_cross_attention
        if self.add_cross_attention:
            assert (
                self.is_decoder
            ), f"{self} should be used as a decoder model if cross attention is added"
            self.crossattention = RobertaAttention(config)
        self.intermediate = RobertaIntermediate(config)
        self.output = RobertaOutput(config)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        output_attentions=False,
        input_attentions=None,
        link_mask=None,
    ):
        self_attention_outputs = self.attention(
            hidden_states,
            attention_mask,
            head_mask,
            output_attentions=output_attentions,
            input_attentions=input_attentions,
            link_mask=link_mask,
        )
        attention_output = self_attention_outputs[0]
        outputs = self_attention_outputs[
            1:
        ]  # add self attentions if we output attention weights

        if self.is_decoder and encoder_hidden_states is not None:
            assert hasattr(
                self, "crossattention"
            ), f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers by setting `config.add_cross_attention=True`"
            cross_attention_outputs = self.crossattention(
                attention_output,
                attention_mask,
                head_mask,
                encoder_hidden_states,
                encoder_attention_mask,
                output_attentions,
            )
            attention_output = cross_attention_outputs[0]
            outputs = (
                outputs + cross_attention_outputs[1:]
            )  # add cross attentions if we output attention weights

        layer_output = apply_chunking_to_forward(
            self.feed_forward_chunk,
            self.chunk_size_feed_forward,
            self.seq_len_dim,
            attention_output,
        )
        outputs = (layer_output,) + outputs
        return outputs

    def feed_forward_chunk(self, attention_output):
        intermediate_output = self.intermediate(attention_output)
        layer_output = self.output(intermediate_output, attention_output)
        return layer_output


class ProbeRobertaEncoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.layer = nn.ModuleList(
            [ProbeRobertaLayer(config) for _ in range(config.num_hidden_layers)]
        )

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        output_attentions=False,
        output_hidden_states=False,
        return_dict=False,
        input_attentions=None,
        link_mask=None,
    ):
        all_hidden_states = () if output_hidden_states else None
        all_attentions = () if output_attentions else None

        for i, layer_module in enumerate(self.layer):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            layer_head_mask = head_mask[i] if head_mask is not None else None

            layer_input_attentions = (
                input_attentions[i] if input_attentions is not None else None
            )
            layer_link_mask = link_mask[i] if link_mask is not None else None
            if getattr(self.config, "gradient_checkpointing", False):

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        return module(
                            *inputs, output_attentions, layer_input_attentions
                        )

                    return custom_forward

                layer_outputs = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(layer_module),
                    hidden_states,
                    attention_mask,
                    layer_head_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                )
            else:
                layer_outputs = layer_module(
                    hidden_states,
                    attention_mask,
                    layer_head_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    output_attentions,
                    layer_input_attentions,
                    layer_link_mask,
                )
            hidden_states = layer_outputs[0]
            if output_attentions:
                all_attentions = all_attentions + (layer_outputs[1],)

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return tuple(
                v
                for v in [hidden_states, all_hidden_states, all_attentions]
                if v is not None
            )
        return BaseModelOutput(
            last_hidden_state=hidden_states,
            hidden_states=all_hidden_states,
            attentions=all_attentions,
        )


class ProbeRobertaModel(RobertaPreTrainedModel):

    authorized_missing_keys = [r"position_ids"]

    # Copied from transformers.modeling_bert.BertModel.__init__ with Bert->Roberta
    def __init__(self, config, add_pooling_layer=True):
        super().__init__(config)
        self.config = config

        self.embeddings = RobertaEmbeddings(config)
        self.encoder = ProbeRobertaEncoder(config)

        self.pooler = RobertaPooler(config) if add_pooling_layer else None

        self.init_weights()

    def get_input_embeddings(self):
        return self.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.embeddings.word_embeddings = value

    def _prune_heads(self, heads_to_prune):

        for layer, heads in heads_to_prune.items():
            self.encoder.layer[layer].attention.prune_heads(heads)

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
        input_attentions=None,
        link_mask=None,
    ):
        r"""
        encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention
            if the model is configured as a decoder.
        encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
            Mask to avoid performing attention on the padding token indices of the encoder input. This mask
            is used in the cross-attention if the model is configured as a decoder.
            Mask values selected in ``[0, 1]``:
            ``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens.
        """
        output_attentions = (
            output_attentions
            if output_attentions is not None
            else self.config.output_attentions
        )
        output_hidden_states = (
            output_hidden_states
            if output_hidden_states is not None
            else self.config.output_hidden_states
        )
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError(
                "You cannot specify both input_ids and inputs_embeds at the same time"
            )
        elif input_ids is not None:
            input_shape = input_ids.size()
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        device = input_ids.device if input_ids is not None else inputs_embeds.device

        if attention_mask is None:
            attention_mask = torch.ones(input_shape, device=device)
        if token_type_ids is None:
            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)

        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
        # ourselves in which case we just need to make it broadcastable to all heads.
        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(
            attention_mask, input_shape, device
        )

        # If a 2D or 3D attention mask is provided for the cross-attention
        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
        if self.config.is_decoder and encoder_hidden_states is not None:
            (
                encoder_batch_size,
                encoder_sequence_length,
                _,
            ) = encoder_hidden_states.size()
            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
            if encoder_attention_mask is None:
                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
            encoder_extended_attention_mask = self.invert_attention_mask(
                encoder_attention_mask
            )
        else:
            encoder_extended_attention_mask = None

        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x n_heads x N x N
        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)

        embedding_output = self.embeddings(
            input_ids=input_ids,
            position_ids=position_ids,
            token_type_ids=token_type_ids,
            inputs_embeds=inputs_embeds,
        )
        encoder_outputs = self.encoder(
            embedding_output,
            attention_mask=extended_attention_mask,
            head_mask=head_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_extended_attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            input_attentions=input_attentions,
            link_mask=link_mask,
        )
        sequence_output = encoder_outputs[0]
        pooled_output = (
            self.pooler(sequence_output) if self.pooler is not None else None
        )

        if not return_dict:
            return (sequence_output, pooled_output) + encoder_outputs[1:]

        return BaseModelOutputWithPooling(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
        )


class ProbeRobertaForQuestionAnswering(RobertaPreTrainedModel):
    authorized_unexpected_keys = [r"pooler"]
    authorized_missing_keys = [r"position_ids"]

    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels

        self.roberta = ProbeRobertaModel(config, add_pooling_layer=False)
        self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)

        self.init_weights()
        self.num_hidden_layers = self.config.num_hidden_layers
        self.num_attention_heads = self.config.num_attention_heads

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        start_positions=None,
        end_positions=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
        input_attentions=None,
        link_mask=None,
    ):
        """
        start_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
            Labels for position (index) of the start of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (:obj:`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.
        end_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
            Labels for position (index) of the end of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (:obj:`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.
        """
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )
        outputs = self.roberta(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            input_attentions=input_attentions,
            link_mask=link_mask,
        )
        sequence_output = outputs[0]

        logits = self.qa_outputs(sequence_output)
        start_logits, end_logits = logits.split(1, dim=-1)
        start_logits = start_logits.squeeze(-1)
        end_logits = end_logits.squeeze(-1)

        total_loss = None
        if start_positions is not None and end_positions is not None:
            # If we are on multi-GPU, split add a dimension
            if len(start_positions.size()) > 1:
                start_positions = start_positions.squeeze(-1)
            if len(end_positions.size()) > 1:
                end_positions = end_positions.squeeze(-1)
            # sometimes the start/end positions are outside our model inputs, we ignore these terms
            ignored_index = start_logits.size(1)
            start_positions.clamp_(0, ignored_index)
            end_positions.clamp_(0, ignored_index)

            loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
            start_loss = loss_fct(start_logits, start_positions)
            end_loss = loss_fct(end_logits, end_positions)
            total_loss = (start_loss + end_loss) / 2

        if not return_dict:
            output = (start_logits, end_logits) + outputs[2:]
            return ((total_loss,) + output) if total_loss is not None else output

        # print("start_logits", start_logits)
        # print("end_logits", end_logits)

        return QuestionAnsweringModelOutput(
            loss=total_loss,
            start_logits=start_logits,
            end_logits=end_logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    # @staticmethod
    def probs_of_span(self, start_logits, end_logits, start_indexes, end_indexes):
        start_logits = F.softmax(start_logits, dim=1)
        end_logits = F.softmax(end_logits, dim=1)
        selected_start_logits = torch.gather(
            start_logits,
            1,
            torch.tensor(start_indexes).to(self.roberta.device).view(-1, 1),
        )
        selected_end_logits = torch.gather(
            end_logits, 1, torch.tensor(end_indexes).to(self.roberta.device).view(-1, 1)
        )
        sum_logits = selected_start_logits * selected_end_logits

        return sum_logits

    # batch size has to be 1
    def probe_forward(
        self,
        input_ids=None,
        answers=None,
        final_start_logits=None,  # for comparison
        final_end_logits=None,  # for comparison
        active_layers=None,
        link_mask=None,
        return_kl=True,
        **kwargs,
    ):

        if active_layers is None:
            input_attentions = None
        else:
            zero_input_attention = torch.zeros(
                [
                    input_ids.size(0),
                    self.config.num_attention_heads,
                    input_ids.size(1),
                    input_ids.size(1),
                ],
                device=input_ids.device,
            )
            input_attentions = [
                None if active_layers[i] else zero_input_attention
                for i in range(self.num_hidden_layers)
            ]

        outputs = self.forward(
            input_ids=input_ids,
            input_attentions=input_attentions,
            link_mask=link_mask,
            **kwargs,
        )
        batch_start_logits, batch_end_logits = outputs.start_logits, outputs.end_logits
        start_indexes, end_indexes = (
            np.array(answers["answer_start"]),
            np.array(answers["answer_end"]),
        )
        batch_probs = self.probs_of_span(
            batch_start_logits, batch_end_logits, start_indexes, end_indexes
        )
        if not return_kl:
            return batch_probs.item()

        kl_loss = F.kl_div(
            F.log_softmax(batch_start_logits, dim=1),
            F.softmax(final_start_logits, dim=1),
            reduction="batchmean",
        ) + F.kl_div(
            F.log_softmax(batch_end_logits, dim=1),
            F.softmax(final_end_logits, dim=1),
            reduction="batchmean",
        )
        return batch_probs.item(), kl_loss.item()

    def restricted_forward(
        self, input_ids=None, active_layers=None, link_mask=None, **kwargs,
    ):
        if active_layers is None:
            input_attentions = None
        else:
            zero_input_attention = torch.zeros(
                [
                    input_ids.size(0),
                    self.config.num_attention_heads,
                    input_ids.size(1),
                    input_ids.size(1),
                ],
                device=input_ids.device,
            )
            input_attentions = [
                None if active_layers[i] else zero_input_attention
                for i in range(self.num_hidden_layers)
            ]
        return self.forward(
            input_ids=input_ids,
            input_attentions=input_attentions,
            link_mask=link_mask,
            **kwargs,
        )

    def layer_attribute(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
        active_layers=None,
        link_mask=None,
        input_attentions=None,
        start_indexes=None,
        end_indexes=None,
        final_start_logits=None,  # for comparison
        final_end_logits=None,  # for comparison
        num_steps=300,
    ):
        # for parallel
        # input size B * L
        # attention mask B * L
        # N_Layer * B * N_HEAD * L * L

        # # sanity check
        # N_Layer * B * N_HEAD * L * L
        final_probs = self.probs_of_span(
            final_end_logits, final_end_logits, start_indexes, end_indexes
        )

        attributions = []

        if active_layers is None:
            active_layers = [1 for _ in range(self.num_hidden_layers)]
        zero_input_attention = torch.zeros(
            [
                input_ids.size(0),
                self.config.num_attention_heads,
                input_ids.size(1),
                input_ids.size(1),
            ],
            device=input_ids.device,
        )

        for layer_i in range(self.num_hidden_layers):

            if not active_layers[layer_i]:
                attributions.append(zero_input_attention)
                continue

            layer_acc_attribution = 0.0
            layer_input_attention_backbone = [
                None if active_layers[i] else zero_input_attention
                for i in range(self.num_hidden_layers)
            ]
            layer_attention = input_attentions[layer_i]

            for step_i in tqdm(range(num_steps)):
                # compose input
                step_attentions = layer_attention * step_i / num_steps
                step_attentions.requires_grad_(True)
                layer_input_attention_backbone[layer_i] = step_attentions

                outputs = self.roberta(
                    input_ids,
                    attention_mask=attention_mask,
                    token_type_ids=token_type_ids,
                    position_ids=position_ids,
                    head_mask=head_mask,
                    inputs_embeds=inputs_embeds,
                    output_attentions=output_attentions,
                    output_hidden_states=output_hidden_states,
                    input_attentions=layer_input_attention_backbone,
                    link_mask=link_mask,
                )

                sequence_output = outputs[0]

                logits = self.qa_outputs(sequence_output)
                start_logits, end_logits = logits.split(1, dim=-1)

                start_logits = start_logits.squeeze(-1)
                end_logits = end_logits.squeeze(-1)

                span_probs = self.probs_of_span(
                    start_logits, end_logits, start_indexes, end_indexes
                )
                step_loss = torch.sum(span_probs)
                step_loss.backward()

                step_attribution = (
                    step_attentions.grad.detach() * layer_attention / num_steps
                )
                layer_acc_attribution += step_attribution

                if step_i == 0:
                    baseline_probs = span_probs
            # print(layer_acc_attribution.size())
            sum_attribution = torch.sum(layer_acc_attribution, (3, 2, 1))
            sanity_probs = baseline_probs + sum_attribution
            # print(layer_i, final_probs.item(), sanity_probs.item(), baseline_probs.item(),  sum_attribution.item())
            attributions.append(layer_acc_attribution)
        attributions = torch.stack(attributions)
        # print(attributions.size())
        return attributions

src/calibration/explainers/shap/probe/probe_models.py
UKPLab/CATfOOD
import numpy as np
import torch


def stats_of_layer_attribution(attribution, feature):
    tokens = feature.tokens

    # pool
    attribution = attribution.numpy()
    attribution = np.sum(attribution, axis=0)

    num_tokens = attribution.shape[0]
    attribution_abs = np.abs(attribution)
    attribution_abs = attribution_abs.flatten()
    attribution_abs = attribution_abs / np.sum(attribution_abs)
    sorted_idx = np.argsort(-attribution_abs)
    attribution_abs_sorted = attribution_abs[sorted_idx]

    top_thresholds = [0.005, 0.01, 0.02]

    for threshold in top_thresholds:
        topk = int(threshold * attribution_abs_sorted.size)
        percent_pertained = np.sum(attribution_abs_sorted[:topk])
        print(
            "% Threshould: {:.1f}, Retained {:.1f}".format(
                threshold * 100, percent_pertained * 100
            )
        )

        selected_index = sorted_idx[:topk]
        attn_src = np.floor_divide(selected_index, num_tokens)
        attn_dst = np.remainder(selected_index, num_tokens)
        attn_src = np.sort(np.unique(attn_src))
        attn_dst = np.sort(np.unique(attn_dst))

        print("src involved", attn_src.size, [tokens[i] for i in attn_src])
        print("dst involved", attn_dst.size, [tokens[i] for i in attn_dst])

        # print('src involved', attn_src.size, 'dst involved', attn_dst.size)


def get_link_mask_by_thresholds(attribution, top_thresholds, pad_to_len=0):
    num_heads = attribution.shape[0]
    num_tokens = attribution.shape[1]

    attribution = attribution.numpy()
    attribution = np.sum(attribution, axis=0)
    attribution[attribution < 0] = 0
    attribution_abs = np.abs(attribution)
    attribution_abs = attribution_abs.flatten()
    sorted_idx = np.argsort(-attribution_abs)

    return_link_masks = []
    np_bool_type = np.dtype(bool)
    for threshold in top_thresholds:
        topk = int(threshold * attribution_abs.size)
        selected_index = sorted_idx[:topk]

        link_mask = np.zeros_like(sorted_idx, dtype=np_bool_type)
        link_mask[selected_index] = True
        link_mask = np.reshape(link_mask, (num_tokens, num_tokens))

        # link_mask[np.arange(num_tokens),np.arange(num_tokens)] = True
        # print(link_mask.shape, topk, np.sum(link_mask))
        if pad_to_len > 0:
            padded_mask = np.ones((pad_to_len, pad_to_len), dtype=np_bool_type)
            padded_mask[:num_tokens, :num_tokens] = link_mask
        else:
            padded_mask = link_mask

        padded_mask = torch.BoolTensor(padded_mask)
        padded_mask = padded_mask.expand(1, num_heads, -1, -1)
        return_link_masks.append(padded_mask)

    return return_link_masks


def get_link_mask_by_token_thresholds(attribution, top_thresholds, pad_to_len=0):
    num_heads = attribution.shape[0]
    num_tokens = attribution.shape[1]

    attribution = attribution.numpy()
    attribution = np.sum(attribution, axis=0)
    attribution[attribution < 0] = 0

    token_values = np.sum(attribution, axis=0) + np.sum(attribution, axis=1)
    sorted_idx = np.argsort(-token_values)

    return_link_masks = []
    np_bool_type = np.dtype(bool)
    for threshold in top_thresholds:
        topk = int(threshold * token_values.size)
        selected_index = sorted_idx[:topk]

        mask_row = np.zeros_like(sorted_idx, dtype=np_bool_type)
        mask_row[selected_index] = True
        link_mask = np.zeros([num_tokens, num_tokens], dtype=np_bool_type)
        link_mask[selected_index] = mask_row
        # link_mask[np.arange(num_tokens),np.arange(num_tokens)] = True
        # print(link_mask.shape, topk, np.sum(link_mask))
        if pad_to_len > 0:
            padded_mask = np.ones((pad_to_len, pad_to_len), dtype=np_bool_type)
            padded_mask[:num_tokens, :num_tokens] = link_mask
        else:
            padded_mask = link_mask

        padded_mask = torch.BoolTensor(padded_mask)
        padded_mask = padded_mask.expand(1, num_heads, -1, -1)
        return_link_masks.append(padded_mask)

    return return_link_masks

src/calibration/explainers/shap/probe/probe_utils.py
UKPLab/CATfOOD
"""
File to run the lime interpretability
"""

import logging

logging.basicConfig(level="ERROR")
import os
import shutil
import random
import timeit
import collections
import copy

import numpy as np
import torch
from typing import List, Tuple
from tqdm import tqdm

from datasets import Dataset
from common.config import InterpConfig, load_config_and_tokenizer

import dataloader

from transformers import (
    MODEL_FOR_QUESTION_ANSWERING_MAPPING,
    PreTrainedTokenizer,
    PreTrainedModel,
)
from transformers.models.roberta.modeling_roberta import (
    create_position_ids_from_input_ids,
)

from data.qa_metrics import squad_evaluate

from probe.probe_models import ProbeRobertaForQuestionAnswering
from vis_tools.vis_utils import visualize_token_attributions


logger = logging.getLogger(__name__)

MODEL_CONFIG_CLASSES = list(MODEL_FOR_QUESTION_ANSWERING_MAPPING.keys())
MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)

import itertools
from functools import partial
from utils import run_lime_attribution


def _mkdir_f(prefix):
    if os.path.exists(prefix):
        shutil.rmtree(prefix)
    os.makedirs(prefix)


def set_seed(args):
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)


def ig_analyze(args, tokenizer):
    filenames = os.listdir(args.interp_dir)
    # print(filenames)
    _mkdir_f(args.visual_dir)
    for fname in tqdm(filenames, desc="Visualizing"):
        interp_info = torch.load(os.path.join(args.interp_dir, fname))
        # datset_stats.append(stats_of_ig_interpretation(tokenizer, interp_info))
        visualize_token_attributions(args, tokenizer, interp_info, fname)


class LimeLM:
    def __init__(
        self, args, model: PreTrainedModel, tokenizer: PreTrainedTokenizer,
    ):
        self.args = args
        self.model = model
        self.tokenizer = tokenizer

        if self.model.config.model_type == "gpt2":
            self.ref_token_id = self.tokenizer.eos_token_id
        else:
            self.ref_token_id = self.tokenizer.pad_token_id

        self.sep_token_id = (
            self.tokenizer.sep_token_id
            if self.tokenizer.sep_token_id is not None
            else self.tokenizer.eos_token_id
        )
        self.cls_token_id = (
            self.tokenizer.cls_token_id
            if self.tokenizer.cls_token_id is not None
            else self.tokenizer.bos_token_id
        )

        self.model_prefix = self.model.base_model_prefix
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)

    def _ensure_tensor_on_device(self, **inputs):
        """
        Ensure PyTorch tensors are on the specified device.

        Args:
            inputs (keyword arguments that should be :obj:`torch.Tensor`): The tensors to place on :obj:`self.device`.

        Return:
            :obj:`Dict[str, torch.Tensor]`: The same as :obj:`inputs` but on the proper device.
        """
        return {name: tensor.to(self.model.device) for name, tensor in inputs.items()}

    def encode(
        self, inputs: list = None, add_special_tokens: bool = True, return_tensors=None
    ):
        """
        Encode inputs using the model tokenizer
        Args:
            inputs: question, context pair as a list
            add_special_tokens: where to add CLS, SEP tokens
            return_tensors: whether to return tensors
        Return:
            tokenized inputs
        """
        return self.tokenizer(
            inputs,
            add_special_tokens=add_special_tokens,
            return_tensors=return_tensors,
            padding=True,
            truncation=True,
            max_length=512,
        )

    def decode(self, input_ids: torch.Tensor, skip_special_tokens: bool) -> List[str]:
        """
        Decode received input_ids into a list of word tokens.
        Args:
            input_ids (torch.Tensor): Input ids representing
            word tokens for a sentence/document.
        """
        return self.tokenizer.convert_ids_to_tokens(
            input_ids[0], skip_special_tokens=skip_special_tokens
        )

    def _predict(self, request, batch_size: int = 1) -> tuple:
        """
        Inference on the input.
        Args:
            inputs: list of question, context input
        Returns:
             The model outputs and optionally the input features
        """
        all_predictions = []
        # self.model.to("cuda" if torch.cuda.is_available() else "cpu")

        features = self.encode(request, add_special_tokens=True, return_tensors="pt",)

        for start_idx in range(0, len(request), batch_size):
            with torch.no_grad():
                input_features = {
                    k: features[k][start_idx : start_idx + batch_size]
                    for k in features.keys()
                }
                input_features = self._ensure_tensor_on_device(**input_features)
                predictions = self.model(**input_features)
                all_predictions.append(predictions)
        keys = all_predictions[0].keys()
        final_prediction = {}
        for key in keys:
            if isinstance(all_predictions[0][key], tuple):
                tuple_of_lists = list(
                    zip(
                        *[
                            [
                                torch.stack(p).to(self.device)
                                if isinstance(p, tuple)
                                else p.to(self.device)
                                for p in tpl[key]
                            ]
                            for tpl in all_predictions
                        ]
                    )
                )
                final_prediction[key] = tuple(torch.cat(l) for l in tuple_of_lists)
            else:
                final_prediction[key] = torch.cat(
                    [p[key].to(self.device) for p in all_predictions]
                )
        return final_prediction, features

    def question_answering(self, request, top_k: int = 1):
        """
        Span-based question answering for a given question and context.
        We expect the input to use the (question, context) format for the text pairs.
        Args:
          request: the prediction request
          top_k: num preds

        """

        def decode(
            start_: np.ndarray,
            end_: np.ndarray,
            topk: int,
            max_answer_len: int,
            undesired_tokens_: np.ndarray,
        ) -> Tuple:
            """
            Take the output of any :obj:`ModelForQuestionAnswering` and will generate probabilities
            for each span to be the actual answer.

            In addition, it filters out some unwanted/impossible cases like answer len being greater
            than max_answer_len or answer end position being before the starting position. The method
            supports output the k-best answer through the topk argument.

            Args:
                start_ (:obj:`np.ndarray`): Individual start probabilities for each token.
                end (:obj:`np.ndarray`): Individual end_ probabilities for each token.
                topk (:obj:`int`): Indicates how many possible answer span(s) to extract from the model output.
                max_answer_len (:obj:`int`): Maximum size of the answer to extract from the model's output.
                undesired_tokens_ (:obj:`np.ndarray`): Mask determining tokens that can be part of the answer
            """
            # Ensure we have batch axis
            if start_.ndim == 1:
                start_ = start_[None]

            if end_.ndim == 1:
                end_ = end_[None]

            # Compute the score of each tuple(start_, end_) to be the real answer
            outer = np.matmul(np.expand_dims(start_, -1), np.expand_dims(end_, 1))

            # Remove candidate with end_ < start_ and end_ - start_ > max_answer_len
            candidates = np.tril(np.triu(outer), max_answer_len - 1)

            #  Inspired by Chen & al. (https://github.com/facebookresearch/DrQA)
            scores_flat = candidates.flatten()
            if topk == 1:
                idx_sort = [np.argmax(scores_flat)]
            elif len(scores_flat) < topk:
                idx_sort = np.argsort(-scores_flat)
            else:
                idx = np.argpartition(-scores_flat, topk)[0:topk]
                idx_sort = idx[np.argsort(-scores_flat[idx])]

            starts_, ends_ = np.unravel_index(idx_sort, candidates.shape)[1:]
            desired_spans = np.isin(starts_, undesired_tokens_.nonzero()) & np.isin(
                ends_, undesired_tokens_.nonzero()
            )
            starts_ = starts_[desired_spans]
            ends_ = ends_[desired_spans]
            scores_ = candidates[0, starts_, ends_]

            return starts_, ends_, scores_

        predictions, features = self._predict(request)

        task_outputs = {"answers": [], "answer_start": [], "answer_end": []}
        for idx, (start, end, (_, context)) in enumerate(
            zip(predictions["start_logits"], predictions["end_logits"], request)
        ):

            start_idx = torch.argmax(start)
            end_idx = torch.argmax(end)
            answer_start = torch.tensor([start_idx])
            answer_end = torch.tensor([end_idx])

            start = start.cpu().detach().numpy()
            end = end.cpu().detach().numpy()
            # Ensure padded tokens & question tokens cannot belong to the set of candidate answers.
            question_tokens = np.abs(
                np.array([s != 1 for s in features.sequence_ids(idx)]) - 1
            )
            # Unmask CLS token for 'no answer'
            question_tokens[0] = 1
            undesired_tokens = question_tokens & features["attention_mask"][idx].numpy()

            # Generate mask
            undesired_tokens_mask = undesired_tokens == 0.0

            # Make sure non-context indexes in the tensor cannot contribute to the softmax
            start = np.where(undesired_tokens_mask, -10000.0, start)
            end = np.where(undesired_tokens_mask, -10000.0, end)

            start = np.exp(
                start - np.log(np.sum(np.exp(start), axis=-1, keepdims=True))
            )
            end = np.exp(end - np.log(np.sum(np.exp(end), axis=-1, keepdims=True)))

            # Get score for 'no answer' then mask for decoding step (CLS token
            no_answer_score = (start[0] * end[0]).item()
            start[0] = end[0] = 0.0

            starts, ends, scores = decode(start, end, top_k, 128, undesired_tokens)
            enc = features[idx]
            answers = [
                {
                    "score": score.item(),
                    "start": enc.word_to_chars(enc.token_to_word(s), sequence_index=1)[
                        0
                    ],
                    "end": enc.word_to_chars(enc.token_to_word(e), sequence_index=1)[1],
                    "answer": context[
                        enc.word_to_chars(enc.token_to_word(s), sequence_index=1)[
                            0
                        ] : enc.word_to_chars(enc.token_to_word(e), sequence_index=1)[1]
                    ],
                }
                for s, e, score in zip(starts, ends, scores)
            ]

            answers.append(
                {"score": no_answer_score, "start": 0, "end": 0, "answer": ""}
            )
            answers = sorted(answers, key=lambda x: x["score"], reverse=True)[:top_k]
            task_outputs["answers"].append(answers)
            task_outputs["answer_start"].extend(answer_start.cpu().detach().numpy())
            task_outputs["answer_end"].extend(answer_end.cpu().detach().numpy())
        return predictions, features, task_outputs

    def predict_with_mask(
        self, active_mask, tokenizer, model, base_inputs, answers, full_input_ids
    ):
        input_ids = tokenizer.mask_token_id * torch.ones_like(full_input_ids)
        input_ids[0, active_mask == 1] = full_input_ids[0, active_mask == 1]
        prob = model.probe_forward(**base_inputs, input_ids=input_ids, answers=answers)
        return prob

    def run_lime(self, features, answers):

        # make a copy of the input features
        input_features = copy.deepcopy(features)
        input_features = input_features.to(self.device)
        tokens = input_features.tokens()
        input_features["return_kl"] = False

        full_input_ids = input_features.pop("input_ids")
        full_position_ids = create_position_ids_from_input_ids(
            full_input_ids, self.tokenizer.pad_token_id
        ).to(full_input_ids.device)

        # fix position id
        input_features["position_ids"] = full_position_ids
        # fix cls ? maybe
        score_fn = partial(
            self.predict_with_mask,
            tokenizer=self.tokenizer,
            model=self.model,
            base_inputs=input_features,
            answers=answers,
            full_input_ids=full_input_ids,
        )

        np_attribution = run_lime_attribution(len(tokens), score_fn).reshape((1, -1))
        return torch.from_numpy(np_attribution)

    def lime_interp(self, prefix=""):
        args = self.args
        if not os.path.exists(args.interp_dir):
            os.makedirs(args.interp_dir)

        # fix the model
        self.model.requires_grad_(False)

        def group_batch(batch):
            return {k: [v] for k, v in batch.items()}

        # data = dataloader.PreprocessData(
        #     args.dataset,
        #     args.dataset_config,
        #     save_data=False,
        #     save_path="../../"
        # )
        # dataset = data.processed_val_set()
        # trivia_qa
        dataset = Dataset.from_generator(dataloader.get_dev_examples_hf)
        dataset = dataset.select(range(535, 2500))

        # for shortcuts
        # dataloader = token_in_context.Shortcut(
        #     args.dataset,
        #     args.dataset_config,
        #     args.percent,
        #     args.percent_augment
        # )
        # _, dataset = dataloader.create_synthetic_set()

        eval_dataloader = dataset.map(
            group_batch, batched=True, batch_size=args.eval_batch_size
        )

        if not os.path.exists(args.output_dir) and args.local_rank in [-1, 0]:
            os.makedirs(args.output_dir)

        # Eval!
        logger.info("***** Running evaluation {} *****".format(prefix))
        logger.info("  Num examples = %d", len(dataset))
        logger.info("  Batch size = %d", args.eval_batch_size)

        all_predictions = []
        start_time = timeit.default_timer()

        for idx, batch in tqdm(
            enumerate(eval_dataloader),
            desc="Interpreting",
            total=min(len(dataset), args.first_n_samples),
        ):
            if idx == args.first_n_samples:
                break
            # print(batch)
            # if batch["id"][0] == "6b67d06dbf4b4e2eae802b3e341ce8d4":

            predictions, features, answers = self.question_answering(
                request=[
                    [ques, cxt]
                    for ques, cxt in zip(batch["question"], batch["context"])
                ],
            )

            # input_features = features
            with torch.no_grad():
                importances = self.run_lime(features, answers)

            answers["id"] = batch["id"]
            # print(answers)
            batch_predictions = collections.OrderedDict(answers)
            # print("batch_predictions", batch_predictions)

            self.dump_lime_info(args, batch, features, answers, importances)
            # lots of info, dump to files immediately
            all_predictions.append(batch_predictions)
            # break

        evalTime = timeit.default_timer() - start_time
        logger.info(
            "  Evaluation done in total %f secs (%f sec per example)",
            evalTime,
            evalTime / len(dataset),
        )

        # Compute the F1 and exact scores.
        def merge_predictions(dicts):
            return dict(itertools.chain(*[list(x.items()) for x in dicts]))

        all_predictions = merge_predictions(all_predictions)
        results = squad_evaluate(
            eval_dataloader[: len(all_predictions["id"])], all_predictions
        )
        return results

    def dump_lime_info(self, args, examples, features, predictions, attributions):
        """
        Save the LIME attributions to a file.
        """
        for idx in range(len(examples["id"])):

            actual_len = len(features.tokens(idx))
            attribution = attributions[idx][:actual_len].clone().detach()
            example = [
                examples["question"][idx],
                examples["context"][idx],
                examples["answers"][idx],
            ]

            filename = os.path.join(args.interp_dir, f'{examples["id"][idx]}.bin')
            prelim_result = [
                predictions["answer_start"][idx],
                predictions["answer_end"][idx],
            ]
            prediction = predictions["answers"][idx]
            feature = [
                features.tokens(idx),
                features.input_ids[idx].numpy(),
                features.attention_mask[idx].numpy(),
            ]
            torch.save(
                {
                    "example": example,
                    "feature": feature,
                    "prediction": prediction,
                    "prelim_result": prelim_result,
                    "attribution": attribution,
                },
                filename,
            )


def main():
    args = InterpConfig()

    # Setup CUDA, GPU & distributed training
    if args.local_rank == -1 or args.no_cuda:
        device = torch.device(
            "cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu"
        )
        args.n_gpu = 0 if args.no_cuda else torch.cuda.device_count()
    else:  # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
        torch.cuda.set_device(args.local_rank)
        device = torch.device("cuda", args.local_rank)
        torch.distributed.init_process_group(backend="nccl")
        args.n_gpu = 1
    args.device = device

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN,
    )
    logger.warning(
        "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s",
        args.local_rank,
        device,
        args.n_gpu,
        bool(args.local_rank != -1),
    )

    # Load pretrained model and tokenizer
    if args.local_rank not in [-1, 0]:
        # Make sure only the first process in distributed training will download model & vocab
        torch.distributed.barrier()

    args.model_type = args.model_type.lower()
    config, tokenizer = load_config_and_tokenizer(args)
    # Set seed
    set_seed(args)

    if args.do_vis:
        ig_analyze(args, tokenizer)
    else:
        # Evaluation - we can ask to evaluate all the checkpoints (sub-directories) in a directory
        logger.info("Loading checkpoint %s for evaluation", args.model_name_or_path)
        checkpoint = args.model_name_or_path
        logger.info("Evaluate the following checkpoints: %s", checkpoint)

        # Reload the model
        model = ProbeRobertaForQuestionAnswering.from_pretrained(checkpoint)
        model.to(args.device)

        interpret = LimeLM(args, model, tokenizer)
        # Evaluate
        result = interpret.lime_interp()
        logger.info("Results: {}".format(result))

        return result


if __name__ == "__main__":
    main()

src/calibration/explainers/shap/run_lime.py
UKPLab/CATfOOD
"""
File to run SHAP interpretation
"""

import logging
import os
import shutil
import random
import timeit
import collections
import copy

import numpy as np
import torch
from typing import List, Tuple
from tqdm import tqdm
import pandas as pd

from datasets import Dataset
from common.config import InterpConfig, load_config_and_tokenizer

import dataloader

from transformers import MODEL_FOR_QUESTION_ANSWERING_MAPPING

from transformers.models.roberta.modeling_roberta import (
    create_position_ids_from_input_ids,
)
from transformers import PreTrainedTokenizer, PreTrainedModel
from data.qa_metrics import squad_evaluate

from probe.probe_models import ProbeRobertaForQuestionAnswering
from vis_tools.vis_utils import visualize_token_attributions

logger = logging.getLogger(__name__)

MODEL_CONFIG_CLASSES = list(MODEL_FOR_QUESTION_ANSWERING_MAPPING.keys())
MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"
# BASE_PATH = "/home/sachdeva/projects/ukp/exp_calibration/"

import itertools
from functools import partial
from utils import run_shap_attribution


def _mkdir_f(prefix):
    if os.path.exists(prefix):
        shutil.rmtree(prefix)
    os.makedirs(prefix)


def set_seed(args):
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)


def ig_analyze(args, tokenizer):
    filenames = os.listdir(args.interp_dir)
    # print(filenames)
    filenames = filenames[:2]
    _mkdir_f(args.visual_dir)
    for fname in tqdm(filenames, desc="Visualizing"):
        interp_info = torch.load(os.path.join(args.interp_dir, fname))
        # datset_stats.append(stats_of_ig_interpretation(tokenizer, interp_info))
        visualize_token_attributions(args, tokenizer, interp_info, fname)


class ShapLM:
    def __init__(
        self, args, model: PreTrainedModel, tokenizer: PreTrainedTokenizer,
    ):
        self.args = args
        self.model = model
        self.tokenizer = tokenizer

        if self.model.config.model_type == "gpt2":
            self.ref_token_id = self.tokenizer.eos_token_id
        else:
            self.ref_token_id = self.tokenizer.pad_token_id

        self.sep_token_id = (
            self.tokenizer.sep_token_id
            if self.tokenizer.sep_token_id is not None
            else self.tokenizer.eos_token_id
        )
        self.cls_token_id = (
            self.tokenizer.cls_token_id
            if self.tokenizer.cls_token_id is not None
            else self.tokenizer.bos_token_id
        )

        self.model_prefix = self.model.base_model_prefix
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)

    def _ensure_tensor_on_device(self, **inputs):
        """
        Ensure PyTorch tensors are on the specified device.

        Args:
            inputs (keyword arguments that should be :obj:`torch.Tensor`): The tensors to place on :obj:`self.device`.

        Return:
            :obj:`Dict[str, torch.Tensor]`: The same as :obj:`inputs` but on the proper device.
        """
        return {name: tensor.to(self.model.device) for name, tensor in inputs.items()}

    def encode(
        self, inputs: list = None, add_special_tokens: bool = True, return_tensors=None
    ):
        """
        Encode inputs using the model tokenizer
        Args:
            inputs: question, context pair as a list
            add_special_tokens: where to add CLS, SEP tokens
            return_tensors: whether to return tensors
        Return:
            tokenized inputs
        """
        return self.tokenizer(
            inputs,
            add_special_tokens=add_special_tokens,
            return_tensors=return_tensors,
            padding=True,
            truncation=True,
            max_length=512,
        )

    def decode(self, input_ids: torch.Tensor, skip_special_tokens: bool) -> List[str]:
        """
        Decode received input_ids into a list of word tokens.
        Args:
            input_ids (torch.Tensor): Input ids representing
            word tokens for a sentence/document.
        """
        return self.tokenizer.convert_ids_to_tokens(
            input_ids[0], skip_special_tokens=skip_special_tokens
        )

    def _predict(self, request, batch_size: int = 1) -> tuple:
        """
        Inference on the input.
        Args:
            inputs: list of question, context input
        Returns:
             The model outputs and optionally the input features
        """
        all_predictions = []
        # self.model.to("cuda" if torch.cuda.is_available() else "cpu")

        features = self.encode(request, add_special_tokens=True, return_tensors="pt",)

        for start_idx in range(0, len(request), batch_size):
            with torch.no_grad():
                input_features = {
                    k: features[k][start_idx : start_idx + batch_size]
                    for k in features.keys()
                }
                input_features = self._ensure_tensor_on_device(**input_features)
                predictions = self.model(**input_features)
                all_predictions.append(predictions)
        keys = all_predictions[0].keys()
        final_prediction = {}
        for key in keys:
            if isinstance(all_predictions[0][key], tuple):
                tuple_of_lists = list(
                    zip(
                        *[
                            [
                                torch.stack(p).to(self.device)
                                if isinstance(p, tuple)
                                else p.to(self.device)
                                for p in tpl[key]
                            ]
                            for tpl in all_predictions
                        ]
                    )
                )
                final_prediction[key] = tuple(torch.cat(l) for l in tuple_of_lists)
            else:
                final_prediction[key] = torch.cat(
                    [p[key].to(self.device) for p in all_predictions]
                )
        return final_prediction, features

    def question_answering(self, request, top_k: int = 1):
        """
        Span-based question answering for a given question and context.
        We expect the input to use the (question, context) format for the text pairs.
        Args:
          request: the prediction request
          top_k: num preds

        """

        def decode(
            start_: np.ndarray,
            end_: np.ndarray,
            topk: int,
            max_answer_len: int,
            undesired_tokens_: np.ndarray,
        ) -> Tuple:
            """
            Take the output of any :obj:`ModelForQuestionAnswering` and will generate probabilities
            for each span to be the actual answer.

            In addition, it filters out some unwanted/impossible cases like answer len being greater
            than max_answer_len or answer end position being before the starting position. The method
            supports output the k-best answer through the topk argument.

            Args:
                start_ (:obj:`np.ndarray`): Individual start probabilities for each token.
                end (:obj:`np.ndarray`): Individual end_ probabilities for each token.
                topk (:obj:`int`): Indicates how many possible answer span(s) to extract from the model output.
                max_answer_len (:obj:`int`): Maximum size of the answer to extract from the model's output.
                undesired_tokens_ (:obj:`np.ndarray`): Mask determining tokens that can be part of the answer
            """
            # Ensure we have batch axis
            if start_.ndim == 1:
                start_ = start_[None]

            if end_.ndim == 1:
                end_ = end_[None]

            # Compute the score of each tuple(start_, end_) to be the real answer
            outer = np.matmul(np.expand_dims(start_, -1), np.expand_dims(end_, 1))

            # Remove candidate with end_ < start_ and end_ - start_ > max_answer_len
            candidates = np.tril(np.triu(outer), max_answer_len - 1)

            #  Inspired by Chen & al. (https://github.com/facebookresearch/DrQA)
            scores_flat = candidates.flatten()
            if topk == 1:
                idx_sort = [np.argmax(scores_flat)]
            elif len(scores_flat) < topk:
                idx_sort = np.argsort(-scores_flat)
            else:
                idx = np.argpartition(-scores_flat, topk)[0:topk]
                idx_sort = idx[np.argsort(-scores_flat[idx])]

            starts_, ends_ = np.unravel_index(idx_sort, candidates.shape)[1:]
            desired_spans = np.isin(starts_, undesired_tokens_.nonzero()) & np.isin(
                ends_, undesired_tokens_.nonzero()
            )
            starts_ = starts_[desired_spans]
            ends_ = ends_[desired_spans]
            scores_ = candidates[0, starts_, ends_]

            return starts_, ends_, scores_

        predictions, features = self._predict(request)

        task_outputs = {"answers": [], "answer_start": [], "answer_end": []}
        for idx, (start, end, (_, context)) in enumerate(
            zip(predictions["start_logits"], predictions["end_logits"], request)
        ):

            start_idx = torch.argmax(start)
            end_idx = torch.argmax(end)
            answer_start = torch.tensor([start_idx])
            answer_end = torch.tensor([end_idx])

            start = start.cpu().detach().numpy()
            end = end.cpu().detach().numpy()
            # Ensure padded tokens & question tokens cannot belong to the set of candidate answers.
            question_tokens = np.abs(
                np.array([s != 1 for s in features.sequence_ids(idx)]) - 1
            )
            # Unmask CLS token for 'no answer'
            question_tokens[0] = 1
            undesired_tokens = question_tokens & features["attention_mask"][idx].numpy()

            # Generate mask
            undesired_tokens_mask = undesired_tokens == 0.0

            # Make sure non-context indexes in the tensor cannot contribute to the softmax
            start = np.where(undesired_tokens_mask, -10000.0, start)
            end = np.where(undesired_tokens_mask, -10000.0, end)

            start = np.exp(
                start - np.log(np.sum(np.exp(start), axis=-1, keepdims=True))
            )
            end = np.exp(end - np.log(np.sum(np.exp(end), axis=-1, keepdims=True)))

            # Get score for 'no answer' then mask for decoding step (CLS token
            no_answer_score = (start[0] * end[0]).item()
            start[0] = end[0] = 0.0

            starts, ends, scores = decode(start, end, top_k, 128, undesired_tokens)
            enc = features[idx]
            answers = [
                {
                    "score": score.item(),
                    "start": enc.word_to_chars(enc.token_to_word(s), sequence_index=1)[
                        0
                    ],
                    "end": enc.word_to_chars(enc.token_to_word(e), sequence_index=1)[1],
                    "answer": context[
                        enc.word_to_chars(enc.token_to_word(s), sequence_index=1)[
                            0
                        ] : enc.word_to_chars(enc.token_to_word(e), sequence_index=1)[1]
                    ],
                }
                for s, e, score in zip(starts, ends, scores)
            ]

            answers.append(
                {"score": no_answer_score, "start": 0, "end": 0, "answer": ""}
            )
            answers = sorted(answers, key=lambda x: x["score"], reverse=True)[:top_k]
            task_outputs["answers"].append(answers)
            task_outputs["answer_start"].extend(answer_start.cpu().detach().numpy())
            task_outputs["answer_end"].extend(answer_end.cpu().detach().numpy())
        return predictions, features, task_outputs

    def predict_with_mask(
        self, active_mask, tokenizer, model, base_inputs, answers, full_input_ids
    ):
        input_ids = tokenizer.mask_token_id * torch.ones_like(full_input_ids)
        input_ids[0, active_mask == 1] = full_input_ids[0, active_mask == 1]
        prob = model.probe_forward(**base_inputs, input_ids=input_ids, answers=answers)
        return prob

    def run_shap(self, features, answers):

        # make a copy of the input features
        input_features = copy.deepcopy(features)
        input_features = input_features.to(self.device)
        tokens = input_features.tokens()
        input_features["return_kl"] = False

        full_input_ids = input_features.pop("input_ids")
        full_position_ids = create_position_ids_from_input_ids(
            full_input_ids, self.tokenizer.pad_token_id
        ).to(full_input_ids.device)

        # fix position id
        input_features["position_ids"] = full_position_ids
        # fix cls ? maybe
        score_fn = partial(
            self.predict_with_mask,
            tokenizer=self.tokenizer,
            model=self.model,
            base_inputs=input_features,
            answers=answers,
            full_input_ids=full_input_ids,
        )

        np_attribution = run_shap_attribution(self.args, len(tokens), score_fn).reshape(
            (1, -1)
        )
        return torch.from_numpy(np_attribution)

    def shap_interp(self, prefix=""):
        args = self.args
        if not os.path.exists(args.interp_dir):
            os.makedirs(args.interp_dir)

        # fix the model
        self.model.requires_grad_(False)

        def group_batch(batch):
            return {k: [v] for k, v in batch.items()}

        # data = dataloader.PreprocessData(
        #     args.dataset,
        #     args.dataset_config,
        #     save_data=False,
        #     save_path="../../"
        # )
        # dataset = data.processed_val_set()
        # dataset = dataset.select(range(1445, 3560))

        # trivia_qa
        # dataset = Dataset.from_generator(dataloader.get_dev_examples_hf)
        # dataset = dataset.select(range(0, 3000))
        # dataset = dataset.filter(lambda example: example['id'] == "")

        # for shortcuts
        # dataloader = token_in_context.Shortcut(
        #     args.dataset,
        #     args.dataset_config,
        #     args.percent,
        #     args.percent_augment
        # )
        # _, dataset = dataloader.create_synthetic_set()

        # dataset = Dataset.from_generator(dataloader.get_dev_examples_hf)
        # dataset = dataset.select(range(0, 618))

        # MRQA datasets
        ##############################################################################################
        data = dataloader.get_dev_samples_mrqa(
            BASE_PATH + "src/data/NaturalQuestionsShort.jsonl"
        )
        data = data[11700:11850]

        # Define a dictionary to map old keys to new keys
        key_mapping = {
            "qas_id": "id",
            "question_text": "question",
            "context_text": "context",
            "answer_text": "answers",
            # Add more key mappings...
        }

        # Rename specified keys in each sample
        # Rename specified keys and wrap values in lists in each sample
        renamed_samples = []
        for sample in data:
            renamed_sample = {
                new_key: [sample[old_key]]
                if old_key == "answer_text"
                else sample[old_key]
                for old_key, new_key in key_mapping.items()
            }
            renamed_samples.append(renamed_sample)

        df = pd.DataFrame(renamed_samples)
        # Create a pseudo-dataset with the list
        dataset = Dataset.from_pandas(df)

        ################################################################################################

        eval_dataloader = dataset.map(
            group_batch, batched=True, batch_size=args.eval_batch_size
        )

        if not os.path.exists(args.output_dir) and args.local_rank in [-1, 0]:
            os.makedirs(args.output_dir)

        # Eval!
        logger.info("***** Running evaluation {} *****".format(prefix))
        logger.info("  Num examples = %d", len(dataset))
        logger.info("  Batch size = %d", args.eval_batch_size)

        all_predictions = []
        start_time = timeit.default_timer()

        for idx, batch in tqdm(
            enumerate(eval_dataloader),
            desc="Interpreting",
            total=min(len(dataset), args.first_n_samples),
        ):
            if idx == args.first_n_samples:
                break
            # changes here for dataset
            predictions, features, answers = self.question_answering(
                request=[
                    [ques, cxt]
                    for ques, cxt in zip(batch["question"], batch["context"])
                ],
            )

            # input_features = features
            with torch.no_grad():
                importances = self.run_shap(features, answers)

            answers["id"] = batch["id"]
            batch_predictions = collections.OrderedDict(answers)
            # print("batch_predictions", batch_predictions)

            self.dump_shap_info(args, batch, features, answers, importances)
            # lots of info, dump to files immediately
            all_predictions.append(batch_predictions)
            # break

        evalTime = timeit.default_timer() - start_time
        logger.info(
            "  Evaluation done in total %f secs (%f sec per example)",
            evalTime,
            evalTime / len(dataset),
        )

        # Compute the F1 and exact scores.
        def merge_predictions(dicts):
            return dict(itertools.chain(*[list(x.items()) for x in dicts]))

        all_predictions = merge_predictions(all_predictions)
        results = squad_evaluate(
            eval_dataloader[: len(all_predictions["id"])], all_predictions
        )
        return results

    def dump_shap_info(self, args, examples, features, predictions, attributions):
        """
        Save the SHAP attributions to a file.
        """
        for idx in range(len(examples["id"])):

            actual_len = len(features.tokens(idx))
            attribution = attributions[idx][:actual_len].clone().detach()
            example = [
                examples["question"][idx],
                examples["context"][idx],
                examples["answers"][idx],
            ]

            filename = os.path.join(args.interp_dir, f'{examples["id"][idx]}.bin')
            prelim_result = [
                predictions["answer_start"][idx],
                predictions["answer_end"][idx],
            ]
            prediction = predictions["answers"][idx]
            feature = [
                features.tokens(idx),
                features.input_ids[idx].numpy(),
                features.attention_mask[idx].numpy(),
            ]
            torch.save(
                {
                    "example": example,
                    "feature": feature,
                    "prediction": prediction,
                    "prelim_result": prelim_result,
                    "attribution": attribution,
                },
                filename,
            )


def main():
    args = InterpConfig()

    # Setup CUDA, GPU & distributed training
    if args.local_rank == -1 or args.no_cuda:
        device = torch.device(
            "cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu"
        )
        args.n_gpu = 0 if args.no_cuda else torch.cuda.device_count()
    else:  # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
        torch.cuda.set_device(args.local_rank)
        device = torch.device("cuda", args.local_rank)
        torch.distributed.init_process_group(backend="nccl")
        args.n_gpu = 1
    args.device = device

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN,
    )
    logger.warning(
        "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s",
        args.local_rank,
        device,
        args.n_gpu,
        bool(args.local_rank != -1),
    )

    # Load pretrained model and tokenizer
    if args.local_rank not in [-1, 0]:
        # Make sure only the first process in distributed training will download model & vocab
        torch.distributed.barrier()

    args.model_type = args.model_type.lower()
    config, tokenizer = load_config_and_tokenizer(args)
    # Set seed
    set_seed(args)

    if args.do_vis:
        ig_analyze(args, tokenizer)
    else:
        # Evaluation - we can ask to evaluate all the checkpoints (sub-directories) in a directory
        logger.info("Loading checkpoint %s for evaluation", args.model_name_or_path)
        checkpoint = args.model_name_or_path
        logger.info("Evaluate the following checkpoints: %s", checkpoint)

        # Reload the model
        model = ProbeRobertaForQuestionAnswering.from_pretrained(checkpoint)
        model.to(args.device)

        interpret = ShapLM(args, model, tokenizer)
        # Evaluate
        result = interpret.shap_interp()
        logger.info("Results: {}".format(result))

        return result


if __name__ == "__main__":
    main()

src/calibration/explainers/shap/run_shap.py
UKPLab/CATfOOD
import numpy as np
import scipy as sp
import sklearn
from sklearn.linear_model import Ridge, LinearRegression
from scipy.special import binom
from itertools import combinations


############################### LIME ###############################

## modified from original lime implementation
def lime_kernel(d, kernel_width=25):
    return np.sqrt(np.exp(-(d ** 2) / kernel_width ** 2))


def lime_feature_selection(data, labels, weights, num_features, method):
    clf = Ridge(alpha=0.01, fit_intercept=True)
    clf.fit(data, labels, sample_weight=weights)

    coef = clf.coef_
    if sp.sparse.issparse(data):
        coef = sp.sparse.csr_matrix(clf.coef_)
        weighted_data = coef.multiply(data[0])
        # Note: most efficient to slice the data before reversing
        sdata = len(weighted_data.data)
        argsort_data = np.abs(weighted_data.data).argsort()
        # Edge case where data is more sparse than requested number of feature importances
        # In that case, we just pad with zero-valued features
        if sdata < num_features:
            nnz_indexes = argsort_data[::-1]
            indices = weighted_data.indices[nnz_indexes]
            num_to_pad = num_features - sdata
            indices = np.concatenate(
                (indices, np.zeros(num_to_pad, dtype=indices.dtype))
            )
            indices_set = set(indices)
            pad_counter = 0
            for i in range(data.shape[1]):
                if i not in indices_set:
                    indices[pad_counter + sdata] = i
                    pad_counter += 1
                    if pad_counter >= num_to_pad:
                        break
        else:
            nnz_indexes = argsort_data[sdata - num_features : sdata][::-1]
            indices = weighted_data.indices[nnz_indexes]
        return indices
    else:
        weighted_data = coef * data[0]
        feature_weights = sorted(
            zip(range(data.shape[1]), weighted_data),
            key=lambda x: np.abs(x[1]),
            reverse=True,
        )
        return np.array([x[0] for x in feature_weights[:num_features]])


def lime_explain_instance_with_data(
    neighborhood_data, neighborhood_labels, distances, feature_selection="auto"
):
    # print("nbr labels", neighborhood_labels)
    # print("nbr data", neighborhood_data)
    weights = lime_kernel(distances)
    labels_column = neighborhood_labels
    # used_features = self.feature_selection(neighborhood_data, labels_column, weights, num_features, feature_selection)
    model_regressor = Ridge(alpha=1, fit_intercept=True)
    easy_model = model_regressor
    easy_model.fit(neighborhood_data, labels_column, sample_weight=weights)
    return easy_model.coef_


def lime_feat_labels_distances(
    doc_size, classifier_fn, num_samples=5000, distance_metric="cosine"
):
    def distance_fn(x):
        return (
            sklearn.metrics.pairwise.pairwise_distances(
                x, x[0], metric=distance_metric
            ).ravel()
            * 100
        )

    sample = np.random.randint(1, doc_size + 1, num_samples - 1)
    data = np.ones((num_samples, doc_size))
    data[0] = np.ones(doc_size, dtype=np.int64)
    features_range = range(doc_size)
    for i, size in enumerate(sample, start=1):
        # random perturbation of data
        inactive = np.random.choice(features_range, size, replace=False)
        data[i, inactive] = 0
    # print("data:", data)
    labels = np.array([classifier_fn(d) for d in data])
    # print("labels", labels)
    distances = distance_fn(sp.sparse.csr_matrix(data))
    # print(data, labels, distances)
    return data, labels, distances


def run_lime_attribution(doc_size, classifier_fn):
    data, labels, distances = lime_feat_labels_distances(doc_size, classifier_fn)
    return lime_explain_instance_with_data(data, labels, distances)


############################### SHAP ###############################

# hyper-parameter inherited from original shap implementation
def shap_feat_label_weights(doc_size, classifier_fn, verbose=False):
    # not working for small seq for now, needs more complex way
    num_sample = 2 * doc_size + 2 ** 11
    if num_sample > (2 ** doc_size - 2):
        num_sample = 2 ** doc_size - 2

    if verbose:
        print("Doc size", doc_size, "Num sample", num_sample)

    num_included = 0
    subset_size = 1
    data = []
    kernel_weights = []

    features_idx = np.arange(doc_size)
    while True:
        num_left = num_sample - num_included
        num_sample_of_size = int(binom(doc_size, subset_size))
        if not (2 * subset_size == doc_size):
            num_sample_of_size *= 2
        # not able to contain all of this size
        if num_left < num_sample_of_size:
            if verbose:
                print(
                    "for size",
                    subset_size,
                    "needed",
                    num_sample_of_size,
                    "remaining",
                    num_left,
                )
            break
        if verbose:
            print("including everything of", subset_size)
        weight_of_size = (
            (doc_size - 1)
            / (subset_size * (doc_size - subset_size))
            / binom(doc_size, subset_size)
        )
        # add all combination of this size
        for inds in combinations(features_idx, subset_size):
            pos_mask = np.ones(doc_size, dtype=np.int64)
            pos_mask[np.array(inds, dtype=np.int64)] = 0
            data.append(pos_mask)
            kernel_weights.append(weight_of_size)
            if not (2 * subset_size == doc_size):
                neg_mask = 1 - pos_mask
                data.append(neg_mask)
                kernel_weights.append(weight_of_size)
        subset_size += 1
        num_included += num_sample_of_size

    num_fixed_data = num_included
    unincluded_size = subset_size
    num_left = num_sample - num_included
    if num_left > 0:
        remaining_possible_size = list(
            range(unincluded_size, doc_size + 1 - unincluded_size)
        )
        remaining_possible_weight = np.array(
            [(doc_size - 1) / (s * (doc_size - s)) for s in remaining_possible_size]
        )
        remaining_sample_weight = remaining_possible_weight.sum() / num_left
        if verbose:
            print("already included", num_fixed_data)
            print("remaining size to sample from", remaining_possible_size)
            print("with the following weight", remaining_possible_weight)

        # normalize sampling rate to 1
        remaining_possible_weight = (
            remaining_possible_weight / remaining_possible_weight.sum()
        )
        size_sample = np.random.choice(
            remaining_possible_size, int(num_left / 2), p=remaining_possible_weight
        )

        for size in size_sample:
            selected = np.random.choice(features_idx, size, replace=False)
            pos_mask = np.ones(doc_size, dtype=np.int64)
            pos_mask[np.array(selected, dtype=np.int64)] = 0
            neg_mask = 1 - pos_mask
            data.append(pos_mask)
            data.append(neg_mask)
            kernel_weights.append(remaining_sample_weight)
            kernel_weights.append(remaining_sample_weight)

        if verbose:
            print("sum of fixed weights", sum(kernel_weights[:num_fixed_data]))
            print(
                "theoretical sum of fixed weights",
                2
                * sum(
                    [
                        (doc_size - 1) / (s * (doc_size - s))
                        for s in range(1, unincluded_size)
                    ]
                ),
            )
            print(
                "sum of randomly sampled weights", sum(kernel_weights[num_fixed_data:])
            )
            print(
                "theoretical sum of randomly sampled weights",
                sum(
                    [
                        (doc_size - 1) / (s * (doc_size - s))
                        for s in remaining_possible_size
                    ]
                ),
            )

    data = np.stack(data)
    labels = np.array([classifier_fn(d) for d in data])
    kernel_weights = np.array(kernel_weights)
    return data, labels, kernel_weights


def shap_explain_instance_with_data(data, labels, weights):
    model_regressor = LinearRegression(fit_intercept=False)
    model_regressor.fit(data, labels, sample_weight=weights)
    return model_regressor.coef_


def run_shap_attribution(args, doc_size, classifier_fn):
    data, labels, weights = shap_feat_label_weights(doc_size, classifier_fn)
    # print(data.shape, labels.shape, weights.shape)
    return shap_explain_instance_with_data(data, labels, weights)


if __name__ == "__main__":
    dummy_fn = lambda x: np.sum(x)
    # run_shap_attribution(None, 4, dummy_fn)
    # run_shap_attribution(None, 10, dummy_fn)
    # run_shap_attribution(None, 20, dummy_fn)
    val = run_shap_attribution(None, 100, dummy_fn)
    print(val)

src/calibration/explainers/shap/utils.py
UKPLab/CATfOOD
from PIL import ImageFont, ImageDraw, Image
from colour import Color
import numpy as np
import math
from itertools import chain
import os
from os.path import join
import shutil
import string
from .vis_attention import Element, TextElement, merge_unimportant_blocks

GLOBAL_W_PADDING = 200
GLOBAL_H_PADDING = 80
CONTEXT_W = 800
FONT_SIZE = 10
FONT = "arial"
TEXT_HEIGHT = 15
TEXT_W_PADDING = 4
BG_COLOR = Color("white")
POS_COLOR = Color("blue")
NEG_COLOR = Color("red")
MERGE_THRESHOLD = 0.001
PLOT_THRESHOLD = 0.001


def positive_color(val):
    val = min(val, 1.0)
    val = max(val, 0.0)
    new_c = []
    for a, b in zip(POS_COLOR.rgb, BG_COLOR.rgb):
        c = val * a + (1 - val) * b
        new_c.append(c)
    return tuple(int(x * 255) for x in new_c)


def negative_color(val):
    val = min(val, 1.0)
    val = max(val, 0.0)
    new_c = []
    for a, b in zip(NEG_COLOR.rgb, BG_COLOR.rgb):
        c = val * a + (1 - val) * b
        new_c.append(c)
    return tuple(int(x * 255) for x in new_c)


class WeightedTextElement(Element):
    def __init__(self, token, weight, index, intense, x=0, y=0, w=0, h=0):
        super().__init__(x, y, w, h)
        self.token = token
        weight_str = "%.3f" % weight
        if weight >= 0:
            self.weight = weight_str[1:]
        else:
            self.weight = weight_str[2:]
        self.index = index
        self.token_box = TextElement(self.token, index)
        self.weight_box = TextElement(self.weight, index)
        self.intense = intense

    def box_size(self, draw, font):
        self.token = self.token.encode("latin-1", "replace").decode("latin-1")
        tw, th = draw.textsize(self.token, font=font)
        ww, wh = draw.textsize(self.weight, font=font)
        self.token_box.w = tw
        self.weight_box.w = ww
        self.w = max(tw, ww)
        self.h = 2 * TEXT_HEIGHT
        return self.w, self.h

    def set_position(self, x, y):
        self.x = x
        self.y = y
        self.token_box.x = x + (self.w - self.token_box.w) / 2
        self.token_box.y = y
        self.weight_box.x = x + (self.w - self.weight_box.w) / 2
        self.weight_box.y = y + TEXT_HEIGHT

    def render(self, draw, font):
        if self.intense > 0:
            intense = self.intense
            rescaled_intense = 1.0 * intense + 0.2 * (1 - intense)
            c = positive_color(rescaled_intense)
        else:
            intense = -self.intense
            rescaled_intense = 1.0 * intense + 0.2 * (1 - intense)
            c = negative_color(rescaled_intense)
        # black = (0,0,0)
        # draw.text((self.token_box.x, self.token_box.y), self.token_box.text, font=font, fill=(0, 0, 0))

        self.token_box.text = self.token_box.text.encode("latin-1", "replace").decode(
            "latin-1"
        )
        draw.text(
            (self.token_box.x, self.token_box.y), self.token_box.text, font=font, fill=c
        )
        draw.text(
            (self.weight_box.x, self.weight_box.y),
            self.weight_box.text,
            font=font,
            fill=c,
        )


class TokenGraph:
    def __init__(self, tokens, connection, additional_info=None):
        self.tokens = tokens

        sum_attr, max_attr, min_attr = (
            np.sum(connection),
            np.max(connection),
            np.min(connection),
        )
        max_connection_abs_val = np.max(np.abs(connection))
        tokens, connection = merge_unimportant_blocks(
            tokens, connection, threshold=(max_connection_abs_val * MERGE_THRESHOLD)
        )
        self.connection = connection
        # aggregated
        gather_weight = np.sum(connection, axis=1)
        dispatch_weight = np.sum(connection, axis=0)
        agg_weight = (gather_weight + dispatch_weight) / 2
        # draw figure
        self.aggregated_boxes = []
        max_agg_val = np.max(np.abs(agg_weight))
        for i, (tok, w) in enumerate(zip(tokens, agg_weight)):
            box = WeightedTextElement(tok, w, i, w / max_agg_val)
            self.aggregated_boxes.append(box)

        self.gather_boxes = []
        max_gather_val = np.max(np.abs(gather_weight))
        for i, (tok, w) in enumerate(zip(tokens, gather_weight)):
            box = WeightedTextElement(tok, w, i, w / max_gather_val)
            self.gather_boxes.append(box)

        self.dispatch_boxes = []
        max_dispatch_val = np.max(np.abs(dispatch_weight))
        for i, (tok, w) in enumerate(zip(tokens, dispatch_weight)):
            box = WeightedTextElement(tok, w, i, w / max_dispatch_val)
            self.dispatch_boxes.append(box)

        self.title_text = (
            "Aggregated   -   Gather (X attends to)   -   Dispatach (X is attended)"
        )
        self.add_info = additional_info

        self.w = 0
        self.h = 0

    def get_global_size(self):
        w = GLOBAL_W_PADDING + CONTEXT_W + GLOBAL_W_PADDING
        # estimation
        h = GLOBAL_H_PADDING + 1500 + GLOBAL_H_PADDING
        return w, h

    def _arrange(self, boxes, draw, font, start_x, start_y):
        line_h = TEXT_HEIGHT * 2
        line = 0
        cursor = start_x
        for i, box in enumerate(boxes):
            bw, bh = box.box_size(draw, font)

            # new line
            if cursor + bw > start_x + CONTEXT_W:
                line += 1
                cursor = start_x
                bx = cursor
            else:
                bx = cursor
            cursor += bw + TEXT_W_PADDING
            by = line * TEXT_HEIGHT * 2 + start_y
            box.set_position(bx, by)

        ending_h = start_y + (line + 1) * line_h
        return ending_h

    def arrange(self, draw, font):
        endding = GLOBAL_H_PADDING
        endding = self._arrange(
            self.aggregated_boxes, draw, font, GLOBAL_W_PADDING, endding
        )
        endding += GLOBAL_H_PADDING
        endding = self._arrange(
            self.gather_boxes, draw, font, GLOBAL_W_PADDING, endding
        )
        endding += GLOBAL_H_PADDING
        endding = self._arrange(
            self.dispatch_boxes, draw, font, GLOBAL_W_PADDING, endding
        )

    def render(self, draw, font):
        draw.text((GLOBAL_W_PADDING, 30), self.title_text, font=font, fill=(0, 0, 0))
        if self.add_info:
            draw.text(
                (GLOBAL_W_PADDING, 30 + TEXT_HEIGHT),
                self.add_info,
                font=font,
                fill=(0, 0, 0),
            )

        for box in self.aggregated_boxes:
            box.render(draw, font)

        for box in self.gather_boxes:
            box.render(draw, font)

        for box in self.dispatch_boxes:
            box.render(draw, font)


def visualize_tok_attribution(filename, tokens, connection, interp_info):
    # connection should be n_toks * n_toks
    assert (len(tokens), len(tokens)) == connection.shape

    if interp_info["example"].answer_text:
        gt_text = interp_info["example"].answer_text
    elif len(interp_info["example"].answers):
        gt_text = interp_info["example"].answers[0]["text"]
    else:
        gt_text = "NULL"

    p_text = interp_info["prediction"]
    additional_info = "P: {}        G: {}".format(p_text, gt_text)
    graph = TokenGraph(tokens, connection, additional_info=additional_info)
    global_w, global_h = graph.get_global_size()

    font = ImageFont.truetype("%s.ttf" % (FONT), FONT_SIZE)

    image = Image.new("RGB", (global_w, global_h), (255, 255, 255))
    draw = ImageDraw.Draw(image)

    graph.arrange(draw, font)
    graph.render(draw, font)

    image.save(filename)


def merge_token_attribution_by_segments(attributions, segments):
    new_val = []
    for a, b in segments:
        new_val.append(np.sum(attributions[a:b]))
    attention = np.array(new_val)
    return attention

src/calibration/explainers/shap/vis_tools/vis_token.py
UKPLab/CATfOOD
# from PIL import ImageFont, ImageDraw, Image
# from colour import Color
import numpy as np
import math
from itertools import chain
from typing import List

import os
from os.path import join
import shutil
import string
from transformers import RobertaTokenizer
from transformers.models.gpt2.tokenization_gpt2 import bytes_to_unicode

from .vis_attention import visualize_connection, merge_attention_by_segments
from .vis_token import visualize_tok_attribution, merge_token_attribution_by_segments
from .vis_vanilla_token import visualize_vanilla_tok_attribution

# from src.rag.shap.data.dataset_utils import merge_tokens_into_words


def _mkdir_f(prefix):
    if os.path.exists(prefix):
        shutil.rmtree(prefix)
    os.makedirs(prefix)


def visualize_token_attributions(args, tokenizer, interp_info, fname):
    assert args.visual_dir is not None
    # prefix
    prefix = join(args.visual_dir, fname.split(".")[0])
    # attribution
    # N Layer * N Head
    attribution = interp_info["attribution"]
    attribution_val = attribution.numpy()
    tokens = list(interp_info["feature"][0])
    words, segments = _bpe_decode(tokenizer, tokens)

    # plot aggregated
    # along layers
    aggregated_attribution = merge_token_attribution_by_segments(
        attribution_val, segments
    )
    visualize_vanilla_tok_attribution(
        prefix + ".jpg", words, aggregated_attribution, interp_info
    )


def _bpe_decode(
    tokenizer,
    tokens: List[str],
    # attributions: List
):

    byte_encoder = bytes_to_unicode()
    byte_decoder = {v: k for k, v in byte_encoder.items()}
    decoded_each_tok = [
        bytearray([byte_decoder[c] for c in t]).decode(
            encoding="utf-8", errors="replace"
        )
        for t in tokens
    ]

    end_points = []
    force_break = False
    for idx, token in enumerate(decoded_each_tok):
        # special token, punctuation, alphanumeric
        if (
            token in tokenizer.all_special_tokens
            or token in string.punctuation
            or not any([x.isalnum() for x in token.lstrip()])
            or token.lstrip == "'s"
        ):
            end_points.append(idx)
            force_break = True
            continue

        if force_break:
            end_points.append(idx)
            force_break = False
            continue

        if token[0] == " ":
            tokens[idx] = token[:]
            end_points.append(idx)

    end_points.append(len(tokens))
    segments = []
    for i in range(1, len(end_points)):
        if end_points[i - 1] == end_points[i]:
            continue
        segments.append((end_points[i - 1], end_points[i]))

    merged_tokens = []
    for s0, s1 in segments:
        merged_tokens.append("".join(decoded_each_tok[s0:s1]))

    return merged_tokens, segments


def visualize_attributions(args, tokenizer, interp_info, do_head=False, do_layer=True):
    assert args.visual_dir is not None
    feature = interp_info["feature"]

    # prefix
    prefix = join(args.visual_dir, f"{feature.example_index}-{feature.qas_id}")
    _mkdir_f(prefix)

    # attribution
    # N Layer * N Head
    attribution = interp_info["attribution"]
    # attribution = attribution.res
    n_layers, n_heads, n_tokens, _ = tuple(attribution.size())

    attribution_val = attribution.numpy()
    attribution_diff = np.sum(attribution_val)
    attribution_val = attribution_val / attribution_diff

    words, segments = merge_tokens_into_words(tokenizer, interp_info["feature"])

    # plot aggregated
    # along layers
    aggregated_attribution = np.sum(attribution_val, axis=0)
    aggregated_attribution = np.sum(aggregated_attribution, axis=0)
    aggregated_attribution = merge_attention_by_segments(
        aggregated_attribution, segments
    )
    visualize_connection(
        join(prefix, "aggregated.jpg"), words, aggregated_attribution, interp_info
    )
    visualize_tok_attribution(
        join(prefix, "token_attribution.jpg"),
        words,
        aggregated_attribution,
        interp_info,
    )
    if do_head:
        aggregated_by_head = np.sum(attribution_val, axis=0)
        for i_head in range(aggregated_by_head.shape[0]):
            aggregated_head_i = aggregated_by_head[i_head]
            aggregated_head_i = merge_attention_by_segments(aggregated_head_i, segments)
            visualize_connection(
                join(prefix, f"head-{i_head}.jpg"),
                words,
                aggregated_head_i,
                interp_info,
            )
    if do_layer:
        aggregated_by_layer = np.sum(attribution_val, axis=1)
        for i_layer in range(aggregated_by_layer.shape[0]):
            aggregated_layer_i = aggregated_by_layer[i_layer]
            aggregated_layer_i = merge_attention_by_segments(
                aggregated_layer_i, segments
            )
            visualize_connection(
                join(prefix, f"layer-{i_layer}.jpg"),
                words,
                aggregated_layer_i,
                interp_info,
            )


def visualize_layer_attributions(args, tokenizer, interp_info, do_layer=True):
    assert args.visual_dir is not None
    feature = interp_info["feature"]

    # prefix
    prefix = join(args.visual_dir, f"{feature.example_index}-{feature.qas_id}")
    _mkdir_f(prefix)

    # attribution
    # N Layer * N Head
    attribution = interp_info["attribution"]
    attention = interp_info["attention"]
    prelim_result = interp_info["prelim_result"]
    # attribution = attribution.res
    n_layers, n_heads, n_tokens, _ = tuple(attribution.size())

    attribution_val = attribution.numpy()

    active_layers = interp_info["active_layers"]
    words, segments = merge_tokens_into_words(tokenizer, interp_info["feature"])

    aggregated_attribution = np.sum(attribution_val, axis=0)
    aggregated_attribution = np.sum(aggregated_attribution, axis=0)
    aggregated_attribution = merge_attention_by_segments(
        aggregated_attribution, segments
    )
    visualize_connection(
        join(prefix, "aggregated.jpg"), words, aggregated_attribution, interp_info
    )
    visualize_tok_attribution(
        join(prefix, "token_attribution.jpg"),
        words,
        aggregated_attribution,
        interp_info,
    )

    if do_layer:
        aggregated_by_layer = np.sum(attribution_val, axis=1)
        for i_layer in range(aggregated_by_layer.shape[0]):
            if not active_layers[i_layer]:
                continue
            aggregated_layer_i = aggregated_by_layer[i_layer]
            aggregated_layer_i = merge_attention_by_segments(
                aggregated_layer_i, segments
            )
            visualize_connection(
                join(prefix, f"layer-{i_layer}.jpg"),
                words,
                aggregated_layer_i,
                interp_info,
            )


def visualize_pruned_layer_attributions(args, tokenizer, interp_info, do_layer=True):
    assert args.visual_dir is not None
    feature = interp_info["feature"]

    # prefix
    prefix = join(args.visual_dir, f"{feature.example_index}-{feature.qas_id}")
    _mkdir_f(prefix)

    # attribution
    # N Layer * N Head
    attribution = interp_info["attribution"]
    attention = interp_info["attention"]
    prelim_result = interp_info["prelim_result"]
    # attribution = attribution.res
    n_layers, n_heads, n_tokens, _ = tuple(attribution.size())

    attribution_val = attribution.numpy()

    active_layers = interp_info["active_layers"]
    words, segments = merge_tokens_into_words(tokenizer, interp_info["feature"])
    if do_layer:
        aggregated_by_layer = np.sum(attribution_val, axis=1)
        for i_layer in range(aggregated_by_layer.shape[0]):
            if not active_layers[i_layer]:
                continue
            aggregated_layer_i = aggregated_by_layer[i_layer]
            aggregated_layer_i = merge_attention_by_segments(
                aggregated_layer_i, segments
            )
            visualize_connection(
                join(prefix, f"layer-{i_layer}.jpg"),
                words,
                aggregated_layer_i,
                interp_info,
                vis_negative=False,
            )

src/calibration/explainers/shap/vis_tools/vis_utils.py
UKPLab/CATfOOD
from PIL import ImageFont, ImageDraw, Image
from colour import Color
import numpy as np
import math
from itertools import chain
import os
from os.path import join
import shutil
import string
from .vis_token import WeightedTextElement, positive_color, negative_color

GLOBAL_W_PADDING = 200
GLOBAL_H_PADDING = 80
CONTEXT_W = 800
FONT_SIZE = 10
FONT = "arial"
TEXT_HEIGHT = 15
TEXT_W_PADDING = 4


class VanillaTokenGraph:
    def __init__(self, tokens, attributions, additional_info=None):
        self.tokens = tokens

        sum_attr, max_attr, min_attr = (
            np.sum(attributions),
            np.max(attributions),
            np.min(attributions),
        )
        max_connection_abs_val = np.max(np.abs(attributions))
        # self.connection = connection
        # aggregated
        agg_weight = attributions
        # draw figure
        self.aggregated_boxes = []
        max_agg_val = np.max(np.abs(agg_weight))
        for i, (tok, w) in enumerate(zip(tokens, agg_weight)):
            box = WeightedTextElement(tok, w, i, w / max_agg_val)
            self.aggregated_boxes.append(box)

        self.title_text = "sum: {:.3f}, max_attr: {:.3f}, min_attr: {:.3f}".format(
            sum_attr, max_attr, min_attr
        )
        self.add_info = additional_info

        self.w = 0
        self.h = 0

    def get_global_size(self):
        w = GLOBAL_W_PADDING + CONTEXT_W + GLOBAL_W_PADDING
        # estimation
        h = GLOBAL_H_PADDING + 300 + GLOBAL_H_PADDING
        return w, h

    def _arrange(self, boxes, draw, font, start_x, start_y):
        line_h = TEXT_HEIGHT * 2
        line = 0
        cursor = start_x
        for i, box in enumerate(boxes):
            bw, bh = box.box_size(draw, font)

            # new line
            if cursor + bw > start_x + CONTEXT_W:
                line += 1
                cursor = start_x
                bx = cursor
            else:
                bx = cursor
            cursor += bw + TEXT_W_PADDING
            by = line * TEXT_HEIGHT * 2 + start_y
            box.set_position(bx, by)

        ending_h = start_y + (line + 1) * line_h
        return ending_h

    def arrange(self, draw, font):
        endding = GLOBAL_H_PADDING
        endding = self._arrange(
            self.aggregated_boxes, draw, font, GLOBAL_W_PADDING, endding
        )

    def render(self, draw, font):
        draw.text((GLOBAL_W_PADDING, 30), self.title_text, font=font, fill=(0, 0, 0))
        if self.add_info:
            draw.text(
                (GLOBAL_W_PADDING, 30 + TEXT_HEIGHT),
                self.add_info,
                font=font,
                fill=(0, 0, 0),
            )

        for box in self.aggregated_boxes:
            box.render(draw, font)


def visualize_vanilla_tok_attribution(filename, tokens, attributions, interp_info):
    # connection should be n_toks * n_toks
    assert len(tokens) == attributions.size
    # print(interp_info)
    # if interp_info['example'].answer_text:
    #     gt_text = interp_info['example'].answer_text
    # elif len(interp_info['example'].answers):
    #     gt_text = interp_info['example'].answers[0]['text']
    # else:
    #     gt_text = "NULL"
    gt_text = interp_info["example"][2]["text"]
    print(gt_text)

    p_text = interp_info["prediction"][0]["answer"]
    print(p_text)
    additional_info = "P: {}        G: {}".format(p_text, gt_text)
    graph = VanillaTokenGraph(tokens, attributions, additional_info=additional_info)
    global_w, global_h = graph.get_global_size()

    # font = ImageFont.truetype("%s.ttf"%(FONT), FONT_SIZE)
    font = ImageFont.load_default()

    image = Image.new("RGB", (global_w, global_h), (255, 255, 255))
    draw = ImageDraw.Draw(image)

    graph.arrange(draw, font)
    graph.render(draw, font)

    image.save(filename)

src/calibration/explainers/shap/vis_tools/vis_vanilla_token.py
UKPLab/CATfOOD
import matplotlib
import matplotlib.pyplot as plt
from matplotlib import rc
import seaborn as sns
import pandas as pd
from typing import List, Dict

sns.set_theme()
# sns.set_style("white")
sns.set_context("talk")
plt.style.use("seaborn-deep")
rc("font", **{"family": "serif", "serif": ["Computer Modern"]})
sns.set_context("paper")
# from matplotlib import rc
# rc('font', family='serif')
plt.rcParams["text.color"] = "black"
# rc('text', usetex=True)

# matplotlib.use("pgf")
# matplotlib.rcParams.update({
#     "pgf.texsystem": "pdflatex",
#     'font.family': 'serif',
#     'text.usetex': True,
#     'pgf.rcfonts': False,
# })

matplotlib.rcParams.update({"text.usetex": True})
matplotlib.rc("text.latex", preamble=r"\usepackage{xcolor}")

# print([k for k in plt.rcParams.keys() if k.startswith('lines')])
import csv

import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import plotly.io as pio

pio.kaleido.scope.mathjax = None

plt.tight_layout()


def save_results_conf(path):
    data = {"model": [], "dataset": [], "acc": [], "auc": [], "mce": []}

    # Add data to the dictionary
    models = ["Base", "RGF", "LLaMA", "GPT-NeoxT", "Flan-UL2"]
    datasets = ["squad_adversarial", "trivia", "hotpot"]
    acc_values = [
        [0.642, 0.644, 0.656, 0.660, 0.666],
        [0.601, 0.659, 0.669, 0.678, 0.665],
        [0.611, 0.635, 0.643, 0.630, 0.651],
    ]
    auc_values = [
        [0.685, 0.709, 0.716, 0.713, 0.749],
        [0.583, 0.688, 0.748, 0.745, 0.732],
        [0.740, 0.793, 0.800, 0.800, 0.816],
    ]

    mce_values = [
        [0.474, 0.465, 0.467, 0.473, 0.454],
        [0.539, 0.549, 0.507, 0.507, 0.520],
        [0.502, 0.495, 0.487, 0.482, 0.488],
    ]

    for model in models:
        for i, dataset in enumerate(datasets):
            data["model"].append(model)
            data["dataset"].append(dataset)
            data["acc"].extend(acc_values[i])
            data["auc"].extend(auc_values[i])
            data["mce"].extend(mce_values[i])

    # Write data to CSV file
    with open(path, "w", newline="") as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(["model", "dataset", "acc", "auc", "mce"])
        for i in range(len(data["model"])):
            writer.writerow(
                [
                    data["model"][i],
                    data["dataset"][i],
                    data["acc"][i],
                    data["auc"][i],
                    data["mce"][i],
                ]
            )

    print("CSV file saved successfully.")


def load_data(method, dataset):
    acc, auc, mce = [], [], []
    if method == "conf" and dataset == "squad_adversarial":
        # conf squad
        acc = [0.642, 0.644, 0.656, 0.660, 0.666]
        auc = [0.685, 0.709, 0.716, 0.713, 0.749]
        mce = [0.474, 0.465, 0.467, 0.473, 0.454]
    elif method == "conf" and dataset == "trivia_qa":
        # conf trivia
        acc = [0.601, 0.659, 0.669, 0.678, 0.665]
        auc = [0.583, 0.688, 0.748, 0.745, 0.732]
        mce = [0.539, 0.549, 0.507, 0.507, 0.520]
    elif method == "conf" and dataset == "hotpot_qa":
        # conf hotpot
        acc = [0.611, 0.635, 0.643, 0.630, 0.651]
        auc = [0.740, 0.793, 0.800, 0.800, 0.816]
        mce = [0.502, 0.495, 0.487, 0.482, 0.488]
    elif method == "shap" and dataset == "squad_adversarial":
        acc = [0.750, 0.736, 0.734, 0.775, 0.745, 0.782, 0.752, 0.791]
        auc = [0.843, 0.840, 0.840, 0.871, 0.847, 0.875, 0.854, 0.884]
        mce = [0.471, 0.469, 0.468, 0.461, 0.470, 0.461, 0.468, 0.460]
    elif method == "shap" and dataset == "trivia_qa":
        acc = [0.720, 0.707, 0.707, 0.707, 0.703, 0.702, 0.707, 0.708]
        auc = [0.718, 0.723, 0.761, 0.764, 0.755, 0.756, 0.753, 0.753]
        mce = [0.545, 0.527, 0.505, 0.506, 0.508, 0.509, 0.515, 0.517]
    elif method == "shap" and dataset == "hotpot_qa":
        acc = [0.633, 0.646, 0.653, 0.655, 0.627, 0.630, 0.653, 0.654]
        auc = [0.755, 0.797, 0.800, 0.801, 0.796, 0.797, 0.811, 0.812]
        mce = [0.504, 0.490, 0.493, 0.494, 0.493, 0.493, 0.489, 0.490]
    elif method == "sc_attn" and dataset == "squad_adversarial":
        acc = [0.657, 0.661, 0.654, 0.716, 0.670, 0.729, 0.672, 0.732]
        auc = [0.781, 0.794, 0.781, 0.845, 0.790, 0.849, 0.791, 0.858]
        mce = [0.476, 0.469, 0.475, 0.461, 0.475, 0.459, 0.479, 0.461]
    elif method == "sc_attn" and dataset == "trivia_qa":
        acc = [0.735, 0.713, 0.705, 0.705, 0.699, 0.701, 0.707, 0.708]
        auc = [0.719, 0.722, 0.763, 0.766, 0.756, 0.758, 0.752, 0.756]
        mce = [0.558, 0.539, 0.510, 0.511, 0.511, 0.511, 0.520, 0.521]
    elif method == "sc_attn" and dataset == "hotpot_qa":
        acc = [0.631, 0.641, 0.640, 0.642, 0.625, 0.626, 0.642, 0.643]
        auc = [0.745, 0.789, 0.789, 0.791, 0.787, 0.788, 0.800, 0.803]
        mce = [0.509, 0.492, 0.493, 0.494, 0.492, 0.494, 0.491, 0.491]
    elif method == "ig" and dataset == "squad_adversarial":
        acc = [0.655, 0.653, 0.665, 0.724, 0.666, 0.730, 0.687, 0.746]
        auc = [0.779, 0.779, 0.791, 0.845, 0.792, 0.847, 0.810, 0.858]
        mce = [0.476, 0.475, 0.474, 0.464, 0.474, 0.462, 0.475, 0.464]
    elif method == "ig" and dataset == "trivia_qa":
        acc = [0.729, 0.710, 0.705, 0.706, 0.702, 0.702, 0.706, 0.708]
        auc = [0.722, 0.724, 0.760, 0.763, 0.752, 0.756, 0.749, 0.753]
        mce = [0.554, 0.535, 0.508, 0.510, 0.511, 0.511, 0.521, 0.524]
    elif method == "ig" and dataset == "hotpot_qa":
        acc = [0.629, 0.636, 0.641, 0.644, 0.620, 0.619, 0.642, 0.643]
        auc = [0.746, 0.784, 0.789, 0.791, 0.785, 0.783, 0.799, 0.804]
        mce = [0.507, 0.493, 0.493, 0.494, 0.493, 0.495, 0.492, 0.492]

    if method == "conf" and dataset == "natural_questions":
        acc = [0.686, 0.71, 0.727, 0.72, 0.706]
        auc = [0.701, 0.78, 0.805, 0.813, 0.80]
        mce = [0.531, 0.516, 0.49, 0.476, 0.494]
    elif method == "conf" and dataset == "news_qa":
        acc = [0.68, 0.678, 0.699, 0.703, 0.687]
        auc = [0.713, 0.744, 0.751, 0.751, 0.767]
        mce = [0.535, 0.536, 0.530, 0.534, 0.527]
    elif method == "conf" and dataset == "bioasq":
        acc = [0.633, 0.67, 0.668, 0.651, 0.675]
        auc = [0.721, 0.754, 0.795, 0.802, 0.77]
        mce = [0.499, 0.507, 0.477, 0.465, 0.494]
    elif method == "shap" and dataset == "natural_questions":
        acc = [0.769, 0.747, 0.738, 0.738, 0.73, 0.731, 0.739, 0.741]
        auc = [0.779, 0.80, 0.803, 0.805, 0.81, 0.812, 0.814, 0.815]
        mce = [0.520, 0.491, 0.484, 0.486, 0.478, 0.480, 0.493, 0.494]
    elif method == "shap" and dataset == "news_qa":
        acc = [0.709, 0.703, 0.70, 0.70, 0.698, 0.698, 0.711, 0.712]
        auc = [0.752, 0.763, 0.753, 0.753, 0.74, 0.74, 0.774, 0.775]
        mce = [0.517, 0.515, 0.517, 0.516, 0.518, 0.517, 0.514, 0.514]
    elif method == "shap" and dataset == "bioasq":
        acc = [0.699, 0.713, 0.721, 0.72, 0.71, 0.697, 0.732, 0.721]
        auc = [0.768, 0.784, 0.822, 0.822, 0.824, 0.823, 0.81, 0.81]
        mce = [0.504, 0.496, 0.493, 0.493, 0.485, 0.485, 0.497, 0.497]
    elif method == "sc_attn" and dataset == "natural_questions":
        acc = [0.766, 0.741, 0.74, 0.741, 0.727, 0.729, 0.745, 0.744]
        auc = [0.777, 0.797, 0.804, 0.805, 0.809, 0.81, 0.813, 0.814]
        mce = [0.526, 0.496, 0.489, 0.49, 0.481, 0.482, 0.495, 0.496]
    elif method == "sc_attn" and dataset == "news_qa":
        acc = [0.720, 0.707, 0.704, 0.705, 0.703, 0.702, 0.713, 0.716]
        auc = [0.762, 0.77, 0.756, 0.758, 0.745, 0.748, 0.777, 0.781]
        mce = [0.538, 0.533, 0.528, 0.529, 0.526, 0.526, 0.527, 0.528]
    elif method == "sc_attn" and dataset == "bioasq":
        acc = [0.717, 0.723, 0.727, 0.724, 0.729, 0.72, 0.737, 0.733]
        auc = [0.781, 0.787, 0.825, 0.823, 0.834, 0.83, 0.810, 0.813]
        mce = [0.509, 0.502, 0.494, 0.495, 0.486, 0.487, 0.502, 0.501]
    elif method == "ig" and dataset == "natural_questions":
        acc = [0.767, 0.745, 0.74, 0.74, 0.73, 0.730, 0.742, 0.743]
        auc = [0.777, 0.797, 0.804, 0.805, 0.807, 0.808, 0.813, 0.814]
        mce = [0.526, 0.497, 0.489, 0.490, 0.482, 0.483, 0.496, 0.497]
    elif method == "ig" and dataset == "news_qa":
        acc = [0.720, 0.706, 0.706, 0.705, 0.70, 0.702, 0.714, 0.717]
        auc = [0.762, 0.767, 0.758, 0.76, 0.743, 0.746, 0.779, 0.783]
        mce = [0.536, 0.528, 0.526, 0.527, 0.525, 0.524, 0.528, 0.529]
    elif method == "ig" and dataset == "bioasq":
        acc = [0.704, 0.713, 0.716, 0.72, 0.716, 0.703, 0.731, 0.73]
        auc = [0.77, 0.777, 0.819, 0.822, 0.824, 0.825, 0.807, 0.809]
        mce = [0.507, 0.502, 0.495, 0.495, 0.488, 0.487, 0.502, 0.501]
    return acc, auc, mce


def calculate_calibration_gain(
    method="conf", dataset=["squad_adversarial"], metric="auc"
):
    all_scores = []
    for set in dataset:
        acc, auc, mce = load_data(method, set)
        if metric == "acc":
            score = [round(element * 100, 2) for element in acc]
        elif metric == "auc":
            score = [round(element * 100, 2) for element in auc]

        result = []
        for s in score[1:]:
            result.append(round(s - score[0], 2))
        # print(result)
        all_scores.append(result)

    print(all_scores)
    # Calculate the sums for each position
    sums = [sum(col) for col in zip(*all_scores)]

    # Calculate the averages
    averages = [round(s / len(all_scores), 2) for s in sums]
    print(averages)


def visualize_calibration(method="conf", dataset="squad_adversarial"):

    acc, auc, mce = load_data(method, dataset)
    mce_flip = [round(1 - m, 3) for m in mce]

    if method == "conf":
        header = ["Base", "RGF", "LLaMA", "GPT-NeoxT", "Flan-UL2"]
        colors = ["#1f77b4", "#ff7f0e", "#d62728", "#9467bd", "#2ca02c"]
        # marker_sym = ['D', 's', '*', '^', 'o']
        dashes = ["solid", "dashed", (0, (3, 2, 1, 2)), "dotted", "dashdot"]

    else:
        header = [
            "Base",
            "RGF",
            "LLaMA",
            "LLaMA + F",
            "GPT-NeoxT",
            "GPT-NeoxT + F",
            "Flan-UL2",
            "Flan-UL2 + F",
        ]
        colors = [
            "#1f77b4",
            "#ff7f0e",
            "#d62728",
            "#8c564b",
            "#9467bd",
            "#e377c2",
            "#2ca02c",
            "#008080",
        ]
        dashes = [
            "solid",
            "dashed",
            (0, (3, 2, 1, 2)),  # small line dashdot
            (0, (5, 1)),
            "dotted",
            (0, (3, 1, 1, 1)),  # longer dashed
            "dashdot",
            (0, (4, 3, 1, 2, 1, 3)),  # double dot line
        ]
        # marker_sym = ['D', 's', '*', 'P', '^', 'X', 'o', 'H']

    data = pd.DataFrame({"Accuracy": acc, "AUC": auc, "1-MCE": mce_flip})
    data = data.transpose()
    data.columns = header
    # print(data.head())

    figsize = (3.5, 2.5)  # Adjust the width and height values as desired
    plt.figure(figsize=figsize)

    plt.ylim(0.4, 0.90)
    # plt.yticks([0.4, 0.5, 0.6, 0.7, 0.8, 0.9])

    ax = sns.lineplot(data=data, marker="o", palette=colors)
    # Set marker patterns for each model
    # Set marker patterns for each line
    from itertools import cycle

    marker_labels = header
    # markers = cycle(marker_sym)
    # for line in ax.lines:
    #     line.set_marker(next(markers))

    # Set line dash styles
    line_styles = cycle(dashes)
    color_palette = cycle(colors)
    for line in ax.lines:
        line.set_linestyle(next(line_styles))
        line.set_color(next(color_palette))

    # Hide the legend box
    legend = plt.legend()
    legend.remove()

    # legend_handles = []
    # for i, label in enumerate(marker_labels):
    #     color = ax.lines[i].get_color()
    #     handle, = plt.plot([], [],# marker=marker_sym[i],
    #                        color='black', label=marker_labels[i], #markersize=4,
    #                        # linestyle='None',
    #                        markeredgecolor='None', markerfacecolor=color)
    #     legend_handles.append(handle)
    #
    # # Show the legend
    # legend = plt.legend(handles=legend_handles, loc='lower left', prop={'size': 8})
    # for text in legend.get_texts():
    #     text.set_fontsize(8)  # Set the desired font size
    # plt.setp(legend.get_frame(), alpha=0.2)  # Optional: Set the background opacity
    # legend.get_frame().set_edgecolor('black')
    # plt.setp(legend.get_frame(), width=0.2)  # Set the width of the legend box

    # Set tick color and text color to black
    plt.tick_params(colors="black")
    plt.rc("xtick", color="black")
    plt.rc("ytick", color="black")

    # Set the font size of x-tick labels and y-tick labels

    # plt.xlabel("metrics", fontsize=12, color="r")
    # plt.ylabel("scores", fontsize=12, color="r")
    # plt.show()

    plt.tight_layout()
    # Hide the x-tick labels
    # plt.tick_params(axis='x', which='both', bottom=False, top=False, labelbottom=False)

    # Show both x and y gridlines
    plt.grid(True, which="both")
    # plt.xticks(fontsize=12)
    plt.yticks(fontsize=12)
    plt.gca().xaxis.set_tick_params(labelbottom=False)

    plt.savefig(f"{method}_{dataset}_1.pdf")


def plot_bar_chart(
    methods: List,
    datasets: List,
    subplot_titles: List,
    show_legend: bool,
    save_path: str,
):
    """
    plot the calibration results for models
    """

    metrics = ["Accuracy", "AUC", "1-MCE"]
    complete_data: List = []
    # Define consistent colors for models
    # model_colors = {
    #     "Base": "#1f77b4",
    #     "RGF": "#ff7f0e",
    #     "LLaMA": "#2ca02c",
    #     "LLaMA + F<sub>d": "#8c564b",
    #     "GPT-NeoxT": "#d62728",
    #     "GPT-NeoxT + F<sub>d": "#e377c2",
    #     "Flan-UL2": "#9467bd",
    #     "Flan-UL2 + F<sub>d": "#17becf"
    # }

    # model_colors = {
    #     "Base": 'rgb(251,180,174)',
    #     "RGF": 'rgb(179,205,227)',
    #     "LLaMA": 'rgb(204,235,197)',
    #     "LLaMA + F<sub>d": 'rgb(222,203,228)',
    #     "GPT-NeoxT": 'rgb(254,217,166)',
    #     "GPT-NeoxT + F<sub>d": 'rgb(255,255,204)',
    #     "Flan-UL2": 'rgb(229,216,189)',
    #     "Flan-UL2 + F<sub>d": 'rgb(253,218,236)'
    # }

    # model_colors = {
    #     "Base": 'rgb(102, 197, 204)',
    #     "RGF":  'rgb(246, 207, 113)',
    #     "LLaMA":  'rgb(248, 156, 116)',
    #     "LLaMA + F<sub>d": 'rgb(220, 176, 242)',
    #     "GPT-NeoxT":  'rgb(135, 197, 95)',
    #     "GPT-NeoxT + F<sub>d":  'rgb(158, 185, 243)',
    #     "Flan-UL2":  'rgb(254, 136, 177)',
    #     "Flan-UL2 + F<sub>d":  'rgb(201, 219, 116)'
    # }
    # nein
    # model_colors = {
    #     "Base": 'rgb(102,194,165)',
    #     "RGF": 'rgb(252,141,98)',
    #     "LLaMA":  'rgb(141,160,203)',
    #     "LLaMA + F<sub>d":  'rgb(231,138,195)',
    #     "GPT-NeoxT": 'rgb(166,216,84)',
    #     "GPT-NeoxT + F<sub>d": 'rgb(255,217,47)',
    #     "Flan-UL2": 'rgb(229,196,148)',
    #     "Flan-UL2 + F<sub>d": 'rgb(179,179,179)'
    # }

    # model_colors = {
    #     "Base": 'rgb(141,211,199)',
    #     "RGF": 'rgb(255,255,179)',
    #     "LLaMA":  'rgb(190,186,218)',
    #     "LLaMA + F<sub>d": 'rgb(251,128,114)',
    #     "GPT-NeoxT":   'rgb(128,177,211)',
    #     "GPT-NeoxT + F<sub>d":  'rgb(253,180,98)',
    #     "Flan-UL2": 'rgb(179,222,105)',
    #     "Flan-UL2 + F<sub>d":  'rgb(252,205,229)'
    # }

    model_colors = {
        "Base": "#636EFA",
        "RGF": "#EF553B",
        "LLaMA": "#00CC96",
        "LLaMA + F<sub>d": "#B6E880",
        "GPT-NeoxT": "#FFA15A",
        "GPT-NeoxT + F<sub>d": "#FECB52",
        "Flan-UL2": "#FF6692",
        "Flan-UL2 + F<sub>d": "#FF97FF",
    }

    for i in range(len(methods)):
        data: Dict = {}
        data["Metric"] = metrics
        if methods[i] == "conf":
            models = ["Base", "RGF", "LLaMA", "GPT-NeoxT", "Flan-UL2"]
        else:
            models = [
                "Base",
                "RGF",
                "LLaMA",
                "LLaMA + F<sub>d",
                "GPT-NeoxT",
                "GPT-NeoxT + F<sub>d",
                "Flan-UL2",
                "Flan-UL2 + F<sub>d",
            ]
        acc, auc, mce = load_data(methods[i], datasets[i])
        mce_flip = [round(1 - m, 3) for m in mce]
        df = pd.DataFrame({"Accuracy": acc, "AUC": auc, "1-MCE": mce_flip})
        for idx in df.index:
            data[models[idx]] = df.loc[idx].tolist()
        # print(data)
        complete_data.append(data)

    # Create subplots with n row and m columns
    num_cols = 3
    num_rows = len(methods) // num_cols
    fig = make_subplots(
        rows=num_rows,
        cols=num_cols,
        shared_yaxes=False,
        subplot_titles=subplot_titles,  #  ["NQ", "News QA", "BioASQ"]
        horizontal_spacing=0.05,
        vertical_spacing=0.075,
    )

    for i, _data in enumerate(complete_data):
        if show_legend:
            if i + 1 in [1, 4]:
                showlegend = True
            else:
                showlegend = False
        else:
            showlegend = False
        subplot = go.Figure()
        if methods[i] == "conf":
            models = ["Base", "RGF", "LLaMA", "GPT-NeoxT", "Flan-UL2"]
            legend_models = models
        else:
            models = [
                "Base",
                "RGF",
                "LLaMA + F<sub>d",
                "GPT-NeoxT + F<sub>d",
                "Flan-UL2 + F<sub>d",
            ]
            legend_models = [
                "LLaMA + F<sub>d",
                "GPT-NeoxT + F<sub>d",
                "Flan-UL2 + F<sub>d",
            ]
        for j, model in enumerate(models):
            subplot.add_trace(
                go.Bar(
                    x=_data["Metric"],
                    y=_data[model],
                    showlegend=showlegend and model in legend_models,
                    name=model,
                    marker_color=model_colors[model],
                )
            )
        # subplot.update_traces(marker=dict(line=dict(color='black', width=1)))  # Adjust color and width as needed
        # Add subplot to the main figure
        for trace in subplot.data:
            row = i // 3 + 1
            col = (i + 1) % 3 if (i + 1) % 3 != 0 else 3
            fig.add_trace(trace, row=row, col=col)

    # Apply layout to each subplot
    fig.update_layout(
        font=dict(family="Times New Roman", size=12, color="black"),
        plot_bgcolor="white",  # Set plot background color
        # showlegend=True,
        legend=dict(bgcolor="white", bordercolor="black", borderwidth=1),
        xaxis=dict(
            title_font=dict(size=16, color="black"),
            ticks="outside",
            mirror=True,
            showline=True,
            linewidth=1.5,
            linecolor="black",
        ),  # Optional: Move x-axis ticks outside
        yaxis=dict(
            title_font=dict(size=16, color="black"),
            ticks="outside",
            mirror=True,
            showline=True,
            linewidth=1.5,
            linecolor="black",
            range=[0, 1],
        ),  # Optional: Move y-axis ticks outside
        xaxis2=dict(
            title_font=dict(size=16, color="black"),
            ticks="outside",
            mirror=True,
            showline=True,
            linewidth=1.5,
            linecolor="black",
        ),  # Optional: Move x-axis ticks outside
        yaxis2=dict(
            title_font=dict(size=16, color="black"),
            ticks="outside",
            mirror=True,
            showline=True,
            linewidth=1.5,
            linecolor="black",
            range=[0, 1],
        ),  # Optional: Move y-axis ticks outside
        xaxis3=dict(
            title_font=dict(size=16, color="black"),
            ticks="outside",
            mirror=True,
            showline=True,
            linewidth=1.5,
            linecolor="black",
        ),  # Optional: Move x-axis ticks outside
        yaxis3=dict(
            title_font=dict(size=16, color="black"),
            ticks="outside",
            mirror=True,
            showline=True,
            linewidth=1.5,
            linecolor="black",
            range=[0, 1],
        ),  # Optional: Move y-axis ticks outside
        xaxis4=dict(
            title_font=dict(size=16, color="black"),
            ticks="outside",
            mirror=True,
            showline=True,
            linewidth=1.5,
            linecolor="black",
        ),  # Optional: Move x-axis ticks outside
        yaxis4=dict(
            title_font=dict(size=16, color="black"),
            ticks="outside",
            mirror=True,
            showline=True,
            linewidth=1.5,
            linecolor="black",
            range=[0, 1],
        ),  # Optional: Move y-axis ticks outside
        xaxis5=dict(
            title_font=dict(size=16, color="black"),
            ticks="outside",
            mirror=True,
            showline=True,
            linewidth=1.5,
            linecolor="black",
        ),  # Optional: Move x-axis ticks outside
        yaxis5=dict(
            title_font=dict(size=16, color="black"),
            ticks="outside",
            mirror=True,
            showline=True,
            linewidth=1.5,
            linecolor="black",
            range=[0, 1],
        ),  # Optional: Move y-axis ticks outside
        xaxis6=dict(
            title_font=dict(size=16, color="black"),
            ticks="outside",
            mirror=True,
            showline=True,
            linewidth=1.5,
            linecolor="black",
        ),  # Optional: Move x-axis ticks outside
        yaxis6=dict(
            title_font=dict(size=16, color="black"),
            ticks="outside",
            mirror=True,
            showline=True,
            linewidth=1.5,
            linecolor="black",
            range=[0, 1],
        ),  # Optional: Move y-axis ticks outside
        xaxis7=dict(
            title_font=dict(size=16, color="black"),
            ticks="outside",
            mirror=True,
            showline=True,
            linewidth=1.5,
            linecolor="black",
        ),  # Optional: Move x-axis ticks outside
        yaxis7=dict(
            title_font=dict(size=16, color="black"),
            ticks="outside",
            mirror=True,
            showline=True,
            linewidth=1.5,
            linecolor="black",
            range=[0, 1],
        ),  # Optional: Move y-axis ticks outside
        xaxis8=dict(
            title_font=dict(size=16, color="black"),
            ticks="outside",
            mirror=True,
            showline=True,
            linewidth=1.5,
            linecolor="black",
        ),  # Optional: Move x-axis ticks outside
        yaxis8=dict(
            title_font=dict(size=16, color="black"),
            ticks="outside",
            mirror=True,
            showline=True,
            linewidth=1.5,
            linecolor="black",
            range=[0, 1],
        ),  # Optional: Move y-axis ticks outside
        xaxis9=dict(
            title_font=dict(size=16, color="black"),
            ticks="outside",
            mirror=True,
            showline=True,
            linewidth=1.5,
            linecolor="black",
        ),  # Optional: Move x-axis ticks outside
        yaxis9=dict(
            title_font=dict(size=16, color="black"),
            ticks="outside",
            mirror=True,
            showline=True,
            linewidth=1.5,
            linecolor="black",
            range=[0, 1],
        ),  # Optional: Move y-axis ticks outside
        xaxis10=dict(
            title_font=dict(size=16, color="black"),
            ticks="outside",
            mirror=True,
            showline=True,
            linewidth=1.5,
            linecolor="black",
        ),  # Optional: Move x-axis ticks outside
        yaxis10=dict(
            title_font=dict(size=16, color="black"),
            ticks="outside",
            mirror=True,
            showline=True,
            linewidth=1.5,
            linecolor="black",
            range=[0, 1],
        ),  # Optional: Move y-axis ticks outside
        xaxis11=dict(
            title_font=dict(size=16, color="black"),
            ticks="outside",
            mirror=True,
            showline=True,
            linewidth=1.5,
            linecolor="black",
        ),  # Optional: Move x-axis ticks outside
        yaxis11=dict(
            title_font=dict(size=16, color="black"),
            ticks="outside",
            mirror=True,
            showline=True,
            linewidth=1.5,
            linecolor="black",
            range=[0, 1],
        ),  # Optional: Move y-axis ticks outside
        xaxis12=dict(
            title_font=dict(size=16, color="black"),
            ticks="outside",
            mirror=True,
            showline=True,
            linewidth=1.5,
            linecolor="black",
        ),  # Optional: Move x-axis ticks outside
        yaxis12=dict(
            title_font=dict(size=16, color="black"),
            ticks="outside",
            mirror=True,
            showline=True,
            linewidth=1.5,
            linecolor="black",
            range=[0, 1],
        ),  # Optional: Move y-axis ticks outside
        yaxis_title="CONF<br>Scores",
        yaxis4_title="SHAP<br>Scores",
        yaxis7_title="SC. ATTN.<br>Scores",
        yaxis10_title="IG<br>Scores",
    )
    # Move legend above subplot
    fig.update_layout(
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.05,
            xanchor="center",
            x=0.5,
            font=dict(size=16),
        )
    )
    additional_y_ticks = [0, 0.25, 0.5, 0.75, 1]
    fig.update_yaxes(tickvals=additional_y_ticks)
    # fig.show()
    fig.update_layout(
        width=1200, height=700, template="ggplot2", margin=dict(t=10,),
    )
    pio.write_image(fig, save_path)


if __name__ == "__main__":
    # visualize_calibration(method="sc_attn", dataset="hotpot")
    # save_results_conf("./src/calibration/visualize/data/calib_results_exp.csv")
    plot_bar_chart(
        methods=["conf"] * 3 + ["shap"] * 3 + ["sc_attn"] * 3 + ["ig"] * 3,
        datasets=["natural_questions", "news_qa", "bioasq"] * 4,
        subplot_titles=["NQ", "News QA", "BioASQ"],
        # datasets=["squad_adversarial", "trivia_qa", "hotpot_qa"] * 4,
        # subplot_titles =  ['SQuAD Adversarial', 'Trivia QA', 'Hotpot QA'],
        show_legend=True,
        save_path="calibration_plots_5.pdf",
    )
    # print(px.colors.qualitative.Plotly)
    # calculate_calibration_gain(
    #     method="ig",
    #     dataset=["squad_adversarial", "trivia_qa", "hotpot_qa", "natural_questions", "news_qa", "bioasq"],
    #     metric="auc"
    # )

src/calibration/visualize/calib_plots.py
UKPLab/CATfOOD
from nltk.translate.bleu_score import sentence_bleu
from zss import simple_distance, Node
import numpy as np
import edit_distance
from munch import Munch
import itertools
import random

random.seed(42)

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


def normalized_levenshtein_distance(s1, s2):
    def levenshtein_distance(s1, s2):
        if len(s1) < len(s2):
            s1, s2 = s2, s1

        distances = range(len(s1) + 1)
        for i2, c2 in enumerate(s2):
            distances_ = [i2 + 1]
            for i1, c1 in enumerate(s1):
                if c1 == c2:
                    distances_.append(distances[i1])
                else:
                    distances_.append(
                        1 + min((distances[i1], distances[i1 + 1], distances_[-1]))
                    )
            distances = distances_
        return distances[-1]

    distance = levenshtein_distance(s1, s2)
    max_length = max(len(s1), len(s2))
    return distance / max_length


def compute_lev_distance(doc1, doc2):
    sm = edit_distance.SequenceMatcher(a=[t for t in doc1], b=[t for t in doc2])
    return 1 - sm.ratio()


def add_node_to_root(root):
    if not root or not list(root.children):
        return None
    curr_node = Node(root.text.lower())
    for c in root.children:
        cnode = add_node_to_root(c)
        if cnode:
            curr_node.addkid(cnode)
    return curr_node


def get_nodes(doc):
    span = list(doc.sents)[0]
    return add_node_to_root(span.root)


def compute_tree_edit_distance(doc1, doc2):
    try:
        dist = simple_distance(get_nodes(doc1), get_nodes(doc2))
        return dist  # / len(doc1)
    except:
        return 0


def compute_closeness(docs, base_doc, sentence_similarity=None):
    sem_dist, tree_dist, edit_dist = [], [], []
    for doc in docs:
        if sentence_similarity:
            sem_dist.append(sentence_similarity(base_doc, doc))
        tree_dist.append(compute_tree_edit_distance(base_doc, doc))
        edit_dist.append(compute_lev_distance(base_doc, doc))
    return Munch(
        sem_dist=np.mean(sem_dist),
        tree_dist=np.mean(tree_dist),
        edit_dist=np.mean(edit_dist),
    )


if __name__ == "__main__":
    # rag_counterfactuals_complete_noise_min_filtered_final_dedup_1
    # counterfactual_data_llama_13b_v1_qg_pipeline_all_data_cleaned.jsonl
    # flan_ul2_collated_data_with_answers_processed.jsonl
    # flan-t5-xxl-v3_collated_data_with_answers_processed.jsonl
    # counterfactual_data_alpaca_13b_v2_qg_pipeline_all_data_cleaned.jsonl
    # counterfactual_data_gpt_jt_v2_qg_pipeline_all_data_cleaned.jsonl
    # counterfactual_data_gpt_neox_20b_v2_qg_pipeline_all_data_cleaned.jsonl

    from datasets import load_dataset
    from tqdm import tqdm
    import jsonlines
    import statistics

    # load squad data
    dataset = load_dataset("squad", "plain_text")
    train_data = dataset["train"]
    squad_data = [
        sample
        for sample in tqdm(
            train_data, total=len(train_data), desc="Loading SQuAD data ... "
        )
    ]

    # for sample in squad_data:
    #     if sample["id"] == "570c4f18fed7b91900d45893":
    #         print(sample)

    levenshtein_dist = []
    c = 0
    with jsonlines.open(
        f"{BASE_PATH}src/data/squad/t5_squad_counterfactuals/rag_counterfactuals_complete_noise_min_filtered_final_dedup_1.jsonl"
    ) as reader:
        for example in tqdm(reader):
            id = example["id"].split("_")[0]
            cf_question = example["question"]
            orig_example = [sample for sample in squad_data if sample["id"] == id][0]
            orig_question = orig_example["question"]
            # print("original: ", orig_question)
            # print("cf: ", cf_question)
            levenshtein_dist.append(
                normalized_levenshtein_distance(orig_question, cf_question)
            )
            # print(levenshtein_dist)
        #
        # c += 1
        # # break
        # if c == 20:
        #     break
        dist_score = statistics.mean(levenshtein_dist)
    print(dist_score)

    # x = levenshtein_distance("cap", "cap")
    # print(x)

src/cf_generation/analysis/closeness.py
UKPLab/CATfOOD
import numpy as np
from scipy import stats
from typing import List


def calculate_correlation(x: List, y: List):
    arr_a = np.array(x)
    arr_b = np.array(y)

    res = stats.spearmanr(arr_a, arr_b)
    corr, p_value = res
    return corr, p_value


if __name__ == "__main__":
    # sbert_sim  = [0.55, 0.48, 0.50, 0.50, 0.45, 0.41, 0.41]
    # semantic_eq = [0.52, 0.46, 0.51, 0.55, 0.46, 0.41, 0.40]
    # leven_dist = [0.61, 0.67, 0.65, 0.67, 0.68, 0.71, 0.71]
    # self_bleu = [0.30, 0.26, 0.28, 0.27, 0.24, 0.19, 0.19]

    sbert_sim = [0.55, 0.50, 0.45, 0.41]
    semantic_eq = [0.52, 0.51, 0.46, 0.40]
    leven_dist = [0.61, 0.65, 0.68, 0.71]
    self_bleu = [0.30, 0.28, 0.24, 0.19]

    ood_gain = [1.24, 3.02, 3.45, 1.27, 3.61, 2.36, 3.48]
    calib_conf_gain = [5.42, 7.87, 8.02, 8.18]
    calib_shap_gain = [1.53, 2.73, 2.62, 3.35]
    calib_sc_attn_gain = [1.57, 2.55, 2.6, 2.97]
    calib_ig_gain = [1.2, 2.75, 2.45, 3.35]

    calib_shap_gain_f = [1.53, 3.35, 3.13, 3.9]
    calib_sc_attn_gain_f = [1.57, 3.72, 3.63, 4.33]
    calib_ig_gain_f = [1.2, 3.83, 3.48, 4.42]

    bleu_corr = -calculate_correlation(self_bleu, calib_ig_gain_f)[0]
    leven_corr = calculate_correlation(leven_dist, calib_ig_gain_f)[0]  # label flipping
    sbert_corr = -calculate_correlation(sbert_sim, calib_ig_gain_f)[0]
    semantic_corr = -calculate_correlation(semantic_eq, calib_ig_gain_f)[0]

    avg_corr = (bleu_corr + leven_corr + sbert_corr + semantic_corr) / 4
    print(round(avg_corr, 2))

    print("self bleu: ", bleu_corr)
    print("levensh dist: ", leven_corr)
    print("sbert sim: ", sbert_corr)
    print("semantic eq: ", semantic_corr)

src/cf_generation/analysis/correlation.py
UKPLab/CATfOOD
# -*- coding: utf-8 -*-
# @Date    : 2022-11-05 11:48:38
# @Author  : Chenghao Mou (mouchenghao@gmail.com)
from __future__ import annotations

import argparse
import os
import random
import subprocess
from collections import deque
from pathlib import Path
from pathlib import PosixPath
from typing import Deque
from typing import Generator
from typing import List
from typing import Literal
from typing import Sequence
from typing import Tuple

import datasets
from datasets import load_dataset

from text_dedup import logger
from text_dedup.utils import add_io_args
from text_dedup.utils import add_meta_args
from text_dedup.utils import add_sa_args
from text_dedup.utils.timer import Timer

random.seed(42)
datasets.logging.set_verbosity_error()


def merge_intervals(
    intervals: List[slice],
    merge_strategy: Literal["longest", "overlapping"] = "longest",
) -> List[slice]:
    """
    Merge overlapping intervals.
    Parameters
    ----------
    intervals : List[slice]
        List of intervals
    merge_strategy : Literal["longest", "overlapping"]
        Strategy to merge intervals, by default "longest"
        "overlapping": merge overlapping intervals
        "longest": only ignore duplicate substrings, this is useful because
        when [2, 4] and [3, 5] are duplicates, [2, 5] might be not
    Returns
    -------
    List[slice]
        List of merged intervals
    Examples
    --------
    >>> merge_intervals([slice(0, 10, None), slice(1, 11, None), slice(2, 12, None), slice(3, 13, None),
    ... slice(4, 14, None), slice(5, 15, None), slice(6, 16, None), slice(7, 21, None)], merge_strategy='overlapping')
    [slice(0, 21, None)]
    >>> merge_intervals([slice(0, 10, None), slice(1, 11, None), slice(2, 12, None), slice(3, 13, None),
    ... slice(4, 14, None), slice(5, 15, None), slice(6, 16, None), slice(7, 21, None)],
    ... merge_strategy='longest') #doctest: +ELLIPSIS
    [slice(0, 10, None), slice(1, 11, None), slice(2, 12, None), ... slice(7, 21, None)]
    >>> merge_intervals([slice(0, 2), slice(2, 4), slice(4, 5)], 'overlapping')
    [slice(0, 5, None)]
    >>> merge_intervals([slice(0, 4), slice(2, 4), slice(4, 5)], 'longest')
    [slice(0, 4, None), slice(4, 5, None)]
    >>> merge_intervals([slice(0, 10, None), slice(0, 10, None), slice(0, 10, None), slice(0, 10, None), slice(0, 10, None)])
    [slice(0, 10, None)]
    """
    if len(intervals) == 0:
        return []

    q = deque(
        sorted(
            list(
                map(
                    lambda s: slice(s[0], s[1]), {(s.start, s.stop) for s in intervals},
                ),
            ),
            key=lambda x: (x.start, -x.stop),
        )
    )

    merged: List[slice] = []

    while q:
        current = q.popleft()

        if not merged:
            merged.append(current)
            continue

        prev = merged[-1]
        if merge_strategy == "overlapping":
            if prev.stop >= current.start:
                merged[-1] = slice(prev.start, max(prev.stop, current.stop))
            else:
                merged.append(current)
        elif merge_strategy == "longest":
            if current.stop <= prev.stop:  # ignore substrings
                continue
            else:
                merged.append(current)

    return merged


def restore(
    boundaries: Sequence[slice], segments: str | Path | Sequence[slice],
) -> Generator:
    """
    Restore the duplicate slices from seg_file to their original document boundaries.
    Parameters
    ----------
    boundaries : List[slice]
        List of slices document boundary offsets.
    segments: Union[str, List[slice]]
        Path to the segmented file with duplicate offsets or a list of duplicate slices.
    Yields
    ------
    int, slice
        index and offset
    Examples
    --------
    >>> list(restore(
    ... [slice(0, 10, None),
    ... slice(10, 20, None)],
    ... [slice(0, 5, None),
    ... slice(5, 10, None),
    ... slice(5, 15, None),
    ... slice(5, 19, None)]))
    [(0, slice(0, 5, None)), (0, slice(5, 10, None)), (1, slice(0, 5, None)), (1, slice(0, 9, None))]
    """
    indices: Deque[slice] = deque([])

    if isinstance(segments, (str, PosixPath)):
        with open(segments) as f:
            for line in f:
                try:
                    left, right = line.strip().split(" ", 1)
                    if left.isdigit() and right.isdigit():
                        indices.append(slice(int(left), int(right)))
                except ValueError:
                    continue
    elif isinstance(segments, list):
        indices = deque(segments)

    for i, s in enumerate(boundaries):
        while indices:
            curr_slice = indices.popleft()
            while curr_slice.stop <= s.start and indices:
                curr_slice = indices.popleft()

            x, y = curr_slice.start, curr_slice.stop
            if y <= s.start:
                break

            if x >= s.stop:
                indices.appendleft(slice(x, y))
                break

            if s.start <= x < s.stop <= y:
                yield i, slice(x - s.start, s.stop - s.start)
                if y > s.stop:
                    indices.appendleft(slice(s.stop, y))
                break
            elif s.start <= x < y <= s.stop:
                yield i, slice(x - s.start, y - s.start)
                continue
            elif x < s.start < y <= s.stop:
                yield i, slice(0, y - s.start)
                continue
            elif x < s.start < s.stop <= y:
                yield i, slice(0, s.stop - s.start)
                if y > s.stop:
                    indices.appendleft(slice(s.stop, y))
                break


def restore_and_merge(
    boundaries: Sequence[slice],
    segments: str | Path | Sequence[slice],
    k: int,
    merge_strategy: Literal["longest", "overlapping"] = "longest",
) -> Tuple[List[List[slice]], int]:
    """
    Restore the duplicate slices from seg_file to their original document boundaries and merge them.
    Parameters
    ----------
    boundaries : List[slice]
        List of slices document boundary offsets.
    segments: Union[str, List[slice]]
        Path to the segmented file with duplicate offsets or a list of duplicate slices.
    k : int
        The minimum duplicate substring byte length.
    merge_strategy : Literal["longest", "overlapping"], optional
        The merge strategy to use, by default "longest"
    Returns
    -------
    Tuple[List[List[slice]], int]
        List of merged slices for each document and the duplicate size.
    Examples
    --------
    >>> restore_and_merge(
    ... [slice(0, 10, None), slice(10, 20, None)],
    ... [slice(0, 5, None), slice(5, 10, None), slice(12, 19, None)],
    ... 5, 'longest')
    ([[slice(0, 5, None), slice(5, 10, None)], [slice(2, 9, None)]], 17)
    >>> restore_and_merge(
    ... [slice(0, 10, None), slice(10, 20, None)],
    ... [slice(0, 5, None), slice(5, 10, None), slice(12, 19, None)],
    ... 5, 'overlapping')
    ([[slice(0, 10, None)], [slice(2, 9, None)]], 17)
    """
    duplicate_size = 0
    results: List[List[slice]] = [[] for _ in boundaries]
    for idx, s in restore(boundaries, segments):
        if s.stop - s.start >= k:
            results[int(idx)].append(s)
    for i in range(len(results)):
        results[i] = merge_intervals(results[i], merge_strategy)
        duplicate_size += sum([s.stop - s.start for s in results[i]])
    return results, duplicate_size


def __run_command(cmd: str, cwd: str):
    p = subprocess.Popen(cmd, shell=True, cwd=cwd,)
    code = p.wait()
    if code != 0:
        raise RuntimeError(f"Command {cmd} failed with code {code}. CWD: {cwd}")


def clean_up(text: str, slices: List[slice]) -> str:
    """
    Remove duplicate substrings from the text.
    Parameters
    ----------
    text : str
        Text to remove duplicate substrings from.
    slices : List[slice]
        List of slices to remove.
    Returns
    -------
    str
        Text with duplicate substrings removed.
    Examples
    --------
    >>> clean_up("This is a test.", [slice(0, 4, None), slice(5, 7, None)])
    '  a test.'
    """
    byte_array = bytearray(text, "utf-8")
    result = bytearray()
    start = 0
    for s in slices:
        result.extend(byte_array[start : s.start])
        start = s.stop
    result.extend(byte_array[start:])

    return result.decode("utf-8", errors="ignore")


if __name__ == "__main__":

    parser = argparse.ArgumentParser(
        prog="text-dedup.suffixarray",
        description="Deduplicate text using Suffix Array Deduplication",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser = add_io_args(parser)
    parser = add_meta_args(parser)
    parser = add_sa_args(parser)

    args = parser.parse_args()
    # args.path = "squad"
    # args.google_repo_path = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/dedup_data"
    # args.name = "squad_dedup"
    # args.output = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/output/suffix_array/squad_dedup"
    # args.split = "train"
    # args.column = "question"
    # args.cache_dir = "./cache"

    assert args.path is not None, "Please specify `path` for `load_dataset`."

    (Path(args.google_repo_path) / "output").mkdir(exist_ok=True, parents=True)
    (Path(args.google_repo_path) / "tmp").mkdir(exist_ok=True, parents=True)
    temp_text = "output/temp_text.txt"
    temp_output = "output/temp_output.txt"
    timer = Timer()

    with timer("Total"):
        with timer("Loading"):
            ds = load_dataset(
                path=args.path,
                name=args.name,
                data_dir=args.data_dir,
                data_files=args.data_files,
                split=args.split,
                revision=args.revision,
                cache_dir=args.cache_dir,
                use_auth_token=args.use_auth_token,
            )

        with timer("Preprocessing"):
            offsets: List[slice] = []
            start = 0
            with open(Path(args.google_repo_path) / temp_text, "wb") as f:
                for doc in ds:
                    doc_bytes = doc[args.column].encode("utf-8")
                    end = start + len(doc_bytes)
                    offsets.append(slice(start, end))
                    start = end
                    f.write(doc_bytes)

        with timer("SuffixArray"):
            __run_command(
                f"python scripts/make_suffix_array.py {temp_text}",
                args.google_repo_path,
            )

        with timer("SelfSimilar"):
            __run_command(
                f"cargo run self-similar --data-file {temp_text}"
                f" --length-threshold {args.k} --cache-dir {args.cache_dir} --num-threads {os.cpu_count()}",
                args.google_repo_path,
            )
            __run_command(
                f"cargo run collect --data-file {temp_text}"
                f" --length-threshold {args.k} --cache-dir {args.cache_dir} >"
                f" {temp_output}",
                args.google_repo_path,
            )

        with timer("Restore"):
            duplicate_slices, duplicate_size = restore_and_merge(
                offsets,
                Path(args.google_repo_path) / temp_output,
                args.k,
                args.strategy,
            )

        with timer("Deduplicate"):
            ds = ds.map(
                lambda content, idx: {
                    args.column: clean_up(content, duplicate_slices[idx]),
                },
                with_indices=True,
                input_columns=[args.column],
                desc="Deduplicating",
            ).filter(
                lambda content: len(content) > 0,
                input_columns=[args.column],
                desc="Filtering empty documents",
            )

        with timer("Saving"):
            ds.save_to_disk(args.output)

    PAD = 30
    for k, v in timer.elapsed_times.items():
        logger.info(f"{k:<{PAD}}: {v:.2f} seconds")

    logger.info(f"{'Before':<{PAD}}: {start} bytes ({len(offsets)})")
    logger.info(f"{'After':<{PAD}}: {start - duplicate_size} bytes ({len(ds)})")

src/cf_generation/analysis/dedup.py
UKPLab/CATfOOD
import os.path
from typing import List
import numpy as np
from munch import Munch
import itertools

from nltk.translate.bleu_score import sentence_bleu
from sentence_transformers import SentenceTransformer, util
import scipy

from datasets import load_dataset
from tqdm import tqdm
import jsonlines
import statistics
import random

random.seed(42)

from src.cf_generation.llm_generation.utils import save_to_disk


BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


class Diversity:
    def __init__(self, data_path, metric, save_path):
        # load squad data
        dataset = load_dataset("squad", "plain_text")
        train_data = dataset["train"]
        self.squad_data = [
            sample
            for sample in tqdm(
                train_data, total=len(train_data), desc="Loading SQuAD data ... "
            )
        ]
        self.scores: List = []
        self.data_path = data_path
        self.metric = metric
        self.save_path = save_path

    def calculate_self_bleu(self, docs, base_doc, kwargs=None):
        # it should just be augments around one example.
        scores = []
        if len(docs) == 0:
            return Munch(bleu4=1)
        included = []
        data_points = []
        for doc in docs:
            included.append(doc)
            data_points.append([d for d in doc])
        # print("Included: ", included)
        # print("data points: ", data_points)
        included.append(base_doc)
        # print("Included: ", included)
        data_points.append([d for d in base_doc])
        # print("data points: ", data_points)

        points = list(itertools.combinations(range(len(included)), 2))
        # print(points)
        for i, j in points:
            scores.append(sentence_bleu([data_points[i]], data_points[j]))
        # print(Munch(bleu4=np.mean(scores)))
        return Munch(bleu4=np.mean(scores))

    def normalized_levenshtein_distance(self, sentence1, sentence2):
        """
        calculate levenshtein distance between sentences
        """

        def levenshtein_distance(s1, s2):
            if len(s1) < len(s2):
                s1, s2 = s2, s1

            distances = range(len(s1) + 1)
            for i2, c2 in enumerate(s2):
                distances_ = [i2 + 1]
                for i1, c1 in enumerate(s1):
                    if c1 == c2:
                        distances_.append(distances[i1])
                    else:
                        distances_.append(
                            1 + min((distances[i1], distances[i1 + 1], distances_[-1]))
                        )
                distances = distances_
            return distances[-1]

        distance = levenshtein_distance(sentence1, sentence2)
        max_length = max(len(sentence1), len(sentence2))
        return distance / max_length

    def calculate_semantic_similarity(self, sentence1, sentence2):
        model = SentenceTransformer("paraphrase-MiniLM-L6-v2")
        # Encode the sentences into embeddings
        embeddings1 = model.encode(sentence1, convert_to_tensor=True)
        embeddings2 = model.encode(sentence2, convert_to_tensor=True)

        # Calculate the cosine similarity between the embeddings
        similarity_score = util.pytorch_cos_sim(embeddings1, embeddings2)

        # Convert the similarity score to a Python float
        similarity_score = similarity_score.item()
        return similarity_score

    def calculate_semantic_similarity_batched(self, sentences1, sentences2):
        model = SentenceTransformer("paraphrase-MiniLM-L6-v2")
        batch_size = 64  # Adjust the batch size based on your available GPU memory
        num_batches = len(sentences1) // batch_size + 1
        print(num_batches)

        similarity_scores = []
        for i in tqdm(range(num_batches)):
            batch_sentences1 = sentences1[i * batch_size : (i + 1) * batch_size]
            batch_sentences2 = sentences2[i * batch_size : (i + 1) * batch_size]
            if batch_sentences1:
                embeddings1 = model.encode(batch_sentences1, convert_to_tensor=True)
                embeddings2 = model.encode(batch_sentences2, convert_to_tensor=True)

                similarity_batch = util.pytorch_cos_sim(embeddings1, embeddings2)
                similarity_scores.extend(similarity_batch.cpu().numpy().diagonal())
            # break

        return similarity_scores

    def compute_scores(self):
        results = []
        with jsonlines.open(
            os.path.join(f"{BASE_PATH}src/data/squad/", self.data_path)
        ) as reader:
            for example in tqdm(reader):
                id = example["id"].split("_")[0]
                cf_question = example["question"]
                orig_example = [
                    sample for sample in self.squad_data if sample["id"] == id
                ][0]
                orig_question = orig_example["question"]
                # print(cf_question, orig_question)
                if self.metric == "self_bleu":
                    sim_score = self.calculate_self_bleu([orig_question], cf_question)[
                        "bleu4"
                    ]
                elif self.metric == "sbert_sim":
                    sim_score = self.calculate_semantic_similarity(
                        orig_question, cf_question
                    )
                self.scores.append(sim_score)
                results.append(
                    {
                        "id": id,
                        "sentence1": orig_question,
                        "sentence2": cf_question,
                        "similarity_score": float(sim_score),
                    }
                )
            # average scores
            avg_score = statistics.mean(self.scores)
            if self.save_path:
                with jsonlines.open(
                    os.path.join(f"{BASE_PATH}src/data/squad/", self.save_path),
                    mode="w",
                ) as writer:
                    for result in results:
                        writer.write(result)
        return avg_score

    def compute_scores_batched(self):
        results = []
        c = 0
        # Create a dictionary to map "id" to "question" sentences
        squad_data_map = {
            sample["id"]: sample["question"] for sample in self.squad_data
        }

        with jsonlines.open(
            os.path.join(f"{BASE_PATH}src/data/squad/", self.data_path)
        ) as reader:
            sentences1 = []
            sentences2 = []
            idx = []
            for example in tqdm(reader):
                # Process the data points and add them to the batches
                sentences1.append(example["question"])
                orig_example_id = example["id"].split("_")[0]
                idx.append(orig_example_id)
                sentences2.append(
                    squad_data_map.get(orig_example_id, "")
                )  # Use get() to handle cases where "id" not found

        # Calculate semantic similarity in batches
        similarity_scores = self.calculate_semantic_similarity_batched(
            sentences1, sentences2
        )
        for i, (id, sentence1, sentence2, sim_score) in enumerate(
            zip(idx, sentences1, sentences2, similarity_scores)
        ):
            cf_question = sentence1
            orig_question = sentence2
            results.append(
                {
                    "id": id,
                    "sentence1": orig_question,
                    "sentence2": cf_question,
                    "similarity_score": float(sim_score),
                }
            )

        if self.save_path:
            with jsonlines.open(
                os.path.join(f"{BASE_PATH}src/data/squad/", self.save_path), mode="w"
            ) as writer:
                for result in results:
                    writer.write(result)
        # print(similarity_scores)
        return statistics.mean(similarity_scores)

    def compute_upper_bound(self):
        results = []
        c = 0
        # Create a list for random instances from SQuAD
        num_random_instances = len(
            self.squad_data
        )  # Adjust the number of random instances as needed
        random_instances = [
            random.choice(self.squad_data) for _ in range(num_random_instances)
        ]
        sentences1, sentences2 = [], []
        for i in tqdm(range(len(self.squad_data))):
            sentences1.append(self.squad_data[i]["question"])
            sentences2.append(random_instances[i]["question"])

        similarity_scores = []
        # Calculate semantic similarity in batches
        if self.metric == "sbert_sim":
            similarity_scores = self.calculate_semantic_similarity_batched(
                sentences1, sentences2
            )
        elif self.metric == "self_bleu":
            for i in tqdm(range(len(sentences1))):
                similarity_scores.append(
                    self.calculate_self_bleu([sentences1[i]], sentences2[i])["bleu4"]
                )
        elif self.metric == "levenshtein":
            for i in tqdm(range(len(sentences1))):
                similarity_scores.append(
                    self.normalized_levenshtein_distance([sentences1[i]], sentences2[i])
                )
        for i, (sentence1, sentence2, sim_score) in enumerate(
            zip(sentences1, sentences2, similarity_scores)
        ):
            cf_question = sentence1
            orig_question = sentence2
            results.append(
                {
                    "sentence1": orig_question,
                    "sentence2": cf_question,
                    "similarity_score": float(sim_score),
                }
            )

        if self.save_path:
            with jsonlines.open(
                os.path.join(f"{BASE_PATH}src/data/squad/", self.save_path), mode="w"
            ) as writer:
                for result in results:
                    writer.write(result)
        # print(similarity_scores)
        return statistics.mean(similarity_scores)


if __name__ == "__main__":
    diverse = Diversity(
        data_path="llama_collated_data_with_answers_processed_context_relevance.jsonl",
        metric="levenshtein",
        save_path="",  # "flan_ul2_sbert_sim_batched.jsonl"
    )
    score = diverse.compute_upper_bound()
    print(score)

    # self_blue = []
    # examples_sim = []
    # examples_div, examples_div1, examples_div2, examples_div3 = [], [], [], []
    # c = 0

    # bleu_score = calculate_self_bleu([orig_question], cf_question)["bleu4"]
    # self_blue.append(bleu_score)
    # print(self_blue)

    # c+=1
    # # if c == 10:
    # #     break
    # # if 0.1 <= bleu_score < 0.2:
    # #     examples_div.append(example)
    # if 0 <= bleu_score < 0.15:
    #     examples_div1.append(example)
    # elif 0.15 <= bleu_score < 0.3:
    #     examples_div2.append(example)
    # elif 0.3 <= bleu_score < 0.45:
    #     examples_div3.append(example)
    # else:
    #     examples_div.append(example)

    # div_score = statistics.mean(self_blue)
    # save_to_disk(examples_sim, "counterfactual_data_flan_ul2_qg_pipeline_all_data_cleaned_sim_0.45.jsonl")
    # save_to_disk(examples_div, "counterfactual_data_flan_ul2_qg_pipeline_all_data_cleaned_div_0.1_0.2.jsonl")
    # save_to_disk(examples_div1, "counterfactual_data_flan_ul2_qg_pipeline_all_data_cleaned_div_0_0.15.jsonl")
    # save_to_disk(examples_div2, "counterfactual_data_flan_ul2_qg_pipeline_all_data_cleaned_div_0.15_0.3.jsonl")
    # save_to_disk(examples_div3, "counterfactual_data_flan_ul2_qg_pipeline_all_data_cleaned_div_0.3_0.45.jsonl")

    # print(div_score)

src/cf_generation/analysis/diversity.py
UKPLab/CATfOOD
import math
import numpy as np
from munch import Munch
from transformers import GPT2LMHeadModel, GPT2TokenizerFast

import torch
from copy import deepcopy


#########################################################################
### compute perplexity
#########################################################################


def _add_special_tokens(text, tokenizer):
    return tokenizer.bos_token + text + tokenizer.eos_token


def _tokens_log_prob(texts, model, tokenizer, batch_size=128, is_cuda=True):
    outputs = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i : i + batch_size]
        outputs.extend(
            _tokens_log_prob_for_batch(batch, model, tokenizer, is_cuda=is_cuda)
        )
    return outputs


def _tokens_log_prob_for_batch(texts, model, tokenizer, is_cuda=True):
    device = "cuda" if is_cuda else "cpu"
    outputs = []
    texts = [_add_special_tokens(text, tokenizer) for text in deepcopy(texts)]
    print(texts)
    # encoding = tokenizer.batch_encode_plus(texts, return_tensors='pt')
    encoding = tokenizer.batch_encode_plus(
        texts, return_tensors="pt", truncation=True, padding=True
    )
    with torch.no_grad():
        ids = encoding["input_ids"].to(device)
        attention_mask = encoding["attention_mask"].to(device)
        # nopad_mask = ids != tokenizer.pad_token_id
        nopad_mask = ids != tokenizer.pad_token_id
        outputs = model(ids, attention_mask=attention_mask)
        logits = outputs.logits

    for sent_index in range(len(texts)):
        sent_nopad_mask = nopad_mask[sent_index]
        sent_tokens = [
            tok
            for i, tok in enumerate(encoding.tokens(sent_index))
            if sent_nopad_mask[i] and i != 0
        ]
        sent_ids = ids[sent_index, sent_nopad_mask][1:]
        sent_logits = logits[sent_index, sent_nopad_mask][:-1, :]
        sent_logits[:, tokenizer.pad_token_id] = float("-inf")
        sent_ids_scores = sent_logits.gather(1, sent_ids.unsqueeze(1)).squeeze(1)
        sent_log_probs = sent_ids_scores - sent_logits.logsumexp(1)
        # sent_log_probs = cast(torch.DoubleTensor, sent_log_probs)
        # sent_ids = cast(torch.LongTensor, sent_ids)

        output = (sent_log_probs.cpu().numpy(), sent_ids.cpu().numpy(), sent_tokens)
        outputs.append(output)
    return outputs


def load_perplex_scorer(model_id="gpt2"):
    model = GPT2LMHeadModel.from_pretrained(model_id)
    tokenizer = GPT2TokenizerFast.from_pretrained(
        model_id, use_fast=True, add_special_tokens=False
    )
    device = "cuda" if torch.cuda.is_available() else "cpu"
    tokenizer.add_special_tokens({"additional_special_tokens": ["<|pad|>"]})
    tokenizer.pad_token = "<|pad|>"
    model.resize_token_embeddings(len(tokenizer))
    model.eval()
    model.to(device)
    return Munch(model=model, tokenizer=tokenizer)


def reduce_perplex_prob(log_probs, log=False, reduce="prod"):
    tlen = log_probs.shape[0]
    if reduce == "prod":
        score = log_probs.sum()
    elif reduce == "mean":
        score = log_probs.logsumexp(0) - math.log(tlen)
    elif reduce == "gmean":
        score = log_probs.mean(0)
    elif reduce == "hmean":
        score = log_probs.neg().logsumexp(0).neg() + math.log(tlen)
    else:
        raise ValueError("Unrecognized scoring strategy: %s" % reduce)
    if not log:
        score = score.exp()
    return score.item()


def normalize_score(log_score, slen, alpha=0.8):
    # Elephant in the Room: An Evaluation Framework for Assessing Adversarial Examples in NLP
    return log_score / math.pow((5 + slen) / 6, alpha)


def compute_sent_perplexity(
    sentences, perplex_scorer, log=True, reduce="prod", is_normalize=False, is_cuda=True
):
    """Compute the sentence perplexity. For filtering.
    Args:
        sentences ([type]): [description]
        perplex_scorer ([type]): [description]
        log (bool, optional): [description]. Defaults to True.
        reduce (str, optional): [description]. Defaults to "prod".
        is_normalize (bool, optional): [description]. Defaults to False.
    Returns:
        [type]: [description]
    """
    scores = []
    model, tokenizer = perplex_scorer.model, perplex_scorer.tokenizer
    outputs = _tokens_log_prob(sentences, model, tokenizer, is_cuda=is_cuda)
    for sent_log_prob, sent_ids, sent_tokens in outputs:
        score = reduce_perplex_prob(sent_log_prob, reduce=reduce, log=log)
        if is_normalize:
            score = normalize_score(score, len(sent_tokens))
        scores.append(score)
    return scores


def filter_by_sent_perplexity(sentences, perplex_scorer, thred=20, is_cuda=True):
    scores = compute_sent_perplexity(
        sentences,
        perplex_scorer,
        log=True,
        reduce="prod",
        is_normalize=False,
        is_cuda=is_cuda,
    )
    idxes = np.where(np.array(scores) <= thred)[0]
    filtered = [sentences[i] for i in idxes]


def compute_phrase_perplexity(
    sentence_phrase_tuples,
    perplex_scorer,
    log=True,
    reduce="prod",
    is_normalize=False,
    is_cuda=True,
):
    scores = []
    sentence_phrase_tuples = (
        sentence_phrase_tuples
        if type(sentence_phrase_tuples) != tuple
        else [sentence_phrase_tuples]
    )
    if len(sentence_phrase_tuples) == 0:
        return scores
    model, tokenizer = perplex_scorer.model, perplex_scorer.tokenizer
    outputs = _tokens_log_prob(
        [s[0] for s in sentence_phrase_tuples], model, tokenizer, is_cuda=is_cuda
    )
    for idx, (sentence, phrase) in enumerate(sentence_phrase_tuples):
        log_probs_all = outputs[idx][0]
        full_len = len(outputs[idx][1]) - 1
        if phrase:
            prefix_len = len(tokenizer(sentence.split(phrase)[0].strip())["input_ids"])
        else:
            prefix_len = 0
        phrase_len = len(tokenizer(phrase)["input_ids"])
        prefix_idx, phrase_idx = [0, prefix_len], [prefix_len, prefix_len + phrase_len]

        log_probs = log_probs_all[phrase_idx[0] : phrase_idx[1]]
        # print(sentence.split(phrase)[0].strip(), perplex_scorer.tokenizer(sentence.split(phrase)[0].strip()))
        # print(sentence, phrase, phrase_idx)
        full_sent_score = reduce_perplex_prob(log_probs_all, log=log, reduce=reduce)
        phrase_score = reduce_perplex_prob(log_probs, log=log, reduce=reduce)
        if is_normalize:
            full_sent_score = normalize_score(full_sent_score, full_len)
            phrase_score = normalize_score(phrase_score, phrase_len)
        scores.append((full_sent_score, phrase_score))
    return scores


def compute_delta_perplexity(
    edit_ops, perplex_scorer, is_normalize=False, is_cuda=True
):
    """This is to compute the perplexity
    Args:
        edit_ops ([type]): [description]
        perplex_scorer ([type]): [description]
        is_normalize (bool, optional): [description]. Defaults to False.
    Returns:
        [type]: [description]
    """
    tuples = []
    # print(metadata.primary.acore.doc.text)
    # print(metadata.primary.bcore.doc.text)
    edit_ops = [o for o in edit_ops if o.op != "equal"]
    for op in edit_ops:
        aphrase, bphrase = (
            (op.fromz_full, op.toz_full)
            if op.op == "insert" or op.op == "delete"
            else (op.fromz_core, op.toz_core)
        )
        asent, bsent = aphrase.doc, bphrase.doc
        tuples += [(asent.text, aphrase.text), (bsent.text, bphrase.text)]
    # print(tuples)
    scores = compute_phrase_perplexity(
        tuples, perplex_scorer, is_normalize=is_normalize, is_cuda=is_cuda
    )
    # print(scores)
    paired_scores = []
    for i in range(len(edit_ops)):
        # because of negative, it's i - i+1; lower the better.
        # print(scores[2*i])
        # print(scores[2*i+1])
        paired_scores.append(
            Munch(
                pr_sent=scores[2 * i][0] - scores[2 * i + 1][0],
                pr_phrase=scores[2 * i][1] - scores[2 * i + 1][1],
            )
        )
    paired_scores = sorted(
        paired_scores, key=lambda x: (max(x.pr_sent, x.pr_phrase)), reverse=True
    )  # use the most ungrammar part as the
    return paired_scores[0]


def get_token_probabilities(perplex_scorer, text):
    """
    Get the probabilities of each token in the given text using GPT-2 model.
    """
    model, tokenizer = perplex_scorer.model, perplex_scorer.tokenizer
    input_ids = tokenizer.encode(text, return_tensors="pt")
    with torch.no_grad():
        outputs = model(input_ids)
        logits = outputs.logits

    probs = torch.softmax(logits, dim=-1)
    token_probs = probs[0, :-1].gather(2, input_ids[:, 1:].unsqueeze(-1)).squeeze()

    return token_probs.tolist()


def calculate_perplexity(probabilities):
    """
    Calculate perplexity given a list of probabilities.

    # Example probabilities for each token in the generated answer
    token_probabilities = [0.9, 0.8, 0.95, 0.7, 0.85]
    """
    log_probs = np.log2(probabilities)
    avg_log_prob = np.mean(log_probs)
    perplexity = 2 ** (-avg_log_prob)

    return perplexity


def generate_text(prompt, perplex_scorer, max_length=100):
    """
    Generate text using the GPT-2 model given a prompt.
    """
    model, tokenizer = perplex_scorer.model, perplex_scorer.tokenizer
    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    with torch.no_grad():
        output = model.generate(
            input_ids, max_length=max_length, num_return_sequences=1, do_sample=True
        )
    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
    return generated_text


if __name__ == "__main__":
    model_id = "gpt2"
    perplex_scorer = load_perplex_scorer(model_id)

    prompt = (
        "Q: Sydney walked past a homeless woman asking for change but did not have any money they "
        "could give to her. Sydney felt bad afterwards. How would you describe Sydney?  Sympathetic, "
        "Like a person who was unable to help, Incredulous? A:"
    )

    # log_probs = _tokens_log_prob([prompt], model, tokenizer, batch_size=1, is_cuda=True)
    # print(log_probs)

    # Extract the answer by removing the prompt
    # answer = generated_output.replace(prompt, "").strip()

    # sent_perplexity = compute_sent_perplexity(
    #     [prompt], perplex_scorer, log=True, reduce="prod", is_normalize=False, is_cuda=False
    # )
    # print(sent_perplexity)

    print(generate_text(prompt, perplex_scorer, max_length=100))

src/cf_generation/analysis/perplexity.py
UKPLab/CATfOOD
from datasets import load_dataset
from collections import Counter
from tqdm import tqdm

import plotly.graph_objs as go
import plotly.io as pio

# BASE_PATH = "/home/sachdeva/projects/ukp/exp_calibration/"
BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


def topk_questions(questions, k=5):
    """
    Counts the occurrences of the first word in each question and returns the top k most common questions.

    Args:
        questions (list): A list of questions to count.
        k (int): The number of top questions to return. Default is 5.

    Returns:
        list: A list of the top k most common questions based on the first word.
    """
    tokens_to_remove = ["[", "'", '"', "]", "/"]
    # Create a translation table that maps each unwanted token to None
    translator = str.maketrans({token: None for token in tokens_to_remove})
    questions = [ques.translate(translator).strip() for ques in questions]
    # remove empty questions
    questions = [ques for ques in questions if ques]
    print(len(questions))
    try:
        first_words = [q.split()[0].lower() for q in questions]
    except:
        print(questions)
        raise
    counter = Counter(first_words)
    return [(q, count) for q, count in counter.most_common(k)]


def plot_question_types(questions, k=5, file_name=None):
    """
    Creates a bar chart based on the types of questions in the input list.

    Args:
        questions (list): A list of questions to analyze.
        k (int): The number of top questions to display. Default is 5.
        file_name (str): The name of the file to save the plotly figure as..
    """
    # Get the top k most common question types and their counts
    topk = topk_questions(questions, k=k)
    labels, counts = zip(*topk)

    # Calculate the total number of questions
    total = len(questions)

    # Calculate the percentage of each question type
    percentages = [count / total * 100 for count in counts]

    # Create the bar chart
    fig = go.Figure([go.Bar(x=labels, y=percentages)])

    # Set the chart title and axis labels
    # fig.update_layout(title='Distribution of Question Types', xaxis_title='Question Type', yaxis_title='Count')

    # Adjust the font size based on the chart height
    height = 300
    width = 400

    # Set the height of the figure
    fig.update_layout(height=height)

    font_size = int(14 * (fig.layout.height / height))
    fig.update_layout(font=dict(size=font_size))

    # Adjust the margin to remove the side margins
    fig.update_layout(margin=dict(l=5, r=5, t=30, b=60))
    # Add a white background to the chart
    fig.update_layout(plot_bgcolor="white", paper_bgcolor="white")

    # Set the showline and line attributes for the xaxis and yaxis objects to create a border around the chart
    fig.update_xaxes(showline=True, linecolor="gray", linewidth=1, mirror=True)
    # showgrid=True, gridcolor='lightgray', gridwidth=1)
    fig.update_yaxes(
        showline=True,
        linecolor="gray",
        linewidth=1,
        mirror=True,
        showgrid=True,
        gridcolor="lightgray",
        gridwidth=1,
    )

    # Save the chart as an SVG file
    pio.write_image(fig, file_name, format="svg", width=width, height=height, scale=3)

    # Display the chart
    # fig.show()


if __name__ == "__main__":
    data_type = "counterfactuals"
    if data_type == "squad":
        # load squad data
        dataset = load_dataset("squad", "plain_text")
        train_data = dataset["train"]
        data = [
            sample
            for sample in tqdm(
                train_data, total=len(train_data), desc="Loading SQuAD data ... "
            )
        ]
    elif data_type == "counterfactuals":
        # read in counterfactual data from jsonl file
        import jsonlines

        with jsonlines.open(
            f"{BASE_PATH}src/data/squad/flan_ul2_collated_data_with_answers_processed.jsonl"
        ) as reader:
            data = [
                sample
                for sample in tqdm(reader, desc="Loading Counterfactual data ... ")
            ]
    questions = [sample["question"] for sample in data]
    # print(len(questions))
    top_questions = topk_questions(questions, k=20)
    print(top_questions)

    # plot_question_types(questions, k=10, file_name="question_types_cf.svg")

src/cf_generation/analysis/question_type.py
UKPLab/CATfOOD
import os
import jsonlines
from tqdm import tqdm
import torch
import numpy as np
import random
from torch.utils.data import Dataset, DataLoader

import evaluate
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"

seed = 42
random.seed(seed)
np.random.seed(seed)
# Fix torch random seed
torch.manual_seed(seed)


class MyDataset(Dataset):
    def __init__(self, data, tokenizer, max_seq_length):
        self.data = data
        self.tokenizer = tokenizer
        self.max_seq_length = max_seq_length

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        id, qa_1, qa_2 = self.data[idx]
        inputs = qa_1 + " [SEP] " + qa_2
        reverse_inputs = qa_2 + " [SEP] " + qa_1

        encoded_input = self.tokenizer.encode_plus(
            inputs,
            padding="max_length",
            max_length=self.max_seq_length,
            return_tensors="pt",
            truncation=True,
        )
        encoded_reverse_input = self.tokenizer.encode_plus(
            reverse_inputs,
            padding="max_length",
            max_length=self.max_seq_length,
            return_tensors="pt",
            truncation=True,
        )

        return (
            encoded_input["input_ids"].squeeze(),
            encoded_reverse_input["input_ids"].squeeze(),
            id,
            qa_1,
            qa_2,
        )


class SemanticSimilarity:
    def __init__(self, data_path, save_path, batch_size):
        self.data_path = data_path
        self.save_path = save_path
        self.batch_size = batch_size
        dataset = load_dataset("squad", "plain_text")
        train_data = dataset["train"]
        self.squad_data = [
            sample
            for sample in tqdm(
                train_data, total=len(train_data), desc="Loading SQuAD data ... "
            )
        ]

        self.device = "cuda" if torch.cuda.is_available() else "cpu"

    def _load_data(self, squad_data_map):
        with jsonlines.open(
            os.path.join(f"{BASE_PATH}src/data/squad/", self.data_path)
        ) as reader:
            data = list(reader)  # Load all data into memory at once
        sentences1 = [example["question"] for example in data]
        sentences2 = [
            squad_data_map.get(example["id"].split("_")[0], "") for example in data
        ]
        idx = [example["id"].split("_")[0] for example in data]

        # num_random_instances = len(self.squad_data)
        # data = [random.choice(self.squad_data) for _ in range(num_random_instances)]
        # sentences1 = [example["question"] for example in data]
        # sentences2 = [sample["question"] for sample in self.squad_data]
        # idx = [example["id"].split("_")[0] for example in self.squad_data]
        return idx, sentences1, sentences2

    def compute_entailment(self):
        model_name = "microsoft/deberta-large-mnli"
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForSequenceClassification.from_pretrained(model_name).to(
            self.device
        )

        squad_data_map = {
            sample["id"]: sample["question"] for sample in self.squad_data
        }
        idx, sentences1, sentences2 = self._load_data(squad_data_map)

        data = list(zip(idx, sentences1, sentences2))
        max_seq_length = 128  # You can adjust this value based on your needs
        dataset = MyDataset(data, tokenizer, max_seq_length)
        dataloader = DataLoader(dataset, batch_size=self.batch_size, shuffle=False)

        semantic_set_ids = {}
        deberta_predictions = []
        results = []
        forward_entailment_probs = []
        reverse_entailment_probs = []
        for batch in tqdm(dataloader):
            padded_inputs, padded_reverse_inputs, ids, qa_1, qa_2 = batch
            # print(padded_inputs)
            # print(len(padded_inputs))
            padded_inputs = padded_inputs.to(
                self.device
            )  # Ensure tensor type and conversion
            padded_reverse_inputs = padded_reverse_inputs.to(
                self.device
            )  # Ensure tensor type and conversion

            prediction = model(input_ids=padded_inputs)["logits"]
            reverse_prediction = model(input_ids=padded_reverse_inputs)["logits"]

            # Class 0: Contradiction
            # Class 1: Neutral
            # Class 2: Entailment
            predicted_labels = torch.argmax(prediction, dim=1)
            # print(predicted_labels)
            reverse_predicted_labels = torch.argmax(reverse_prediction, dim=1)
            # print(reverse_predicted_labels)

            deberta_predictions_batch = (predicted_labels != 0) & (
                reverse_predicted_labels != 0
            )
            deberta_predictions.extend(deberta_predictions_batch.tolist())
            # print(deberta_predictions)

            forward_entailment_probs_batch = torch.softmax(prediction, dim=1)[
                :, 2
            ].tolist()  # Probability of entailment
            reverse_entailment_probs_batch = torch.softmax(reverse_prediction, dim=1)[
                :, 2
            ].tolist()
            # print(forward_entailment_probs_batch)

            for i in range(len(qa_1)):
                forward_deberta_pred = predicted_labels[i] != 0  # Class 2 is Entailment
                reverse_deberta_pred = (
                    reverse_predicted_labels[i] != 0
                )  # Class 2 is Entailment

                forward_entailment_prob = forward_entailment_probs_batch[i]
                reverse_entailment_prob = reverse_entailment_probs_batch[i]

                # Include the example in results only if DeBERTa prediction is True for both directions
                if forward_deberta_pred and reverse_deberta_pred:
                    result = {
                        "id": ids[i],  # Assuming you want to use qa_2 as the ID
                        "sentence1": qa_1[i],
                        "sentence2": qa_2[i],
                        "forward_entailment_prob": forward_entailment_prob,
                        "reverse_entailment_prob": reverse_entailment_prob,
                    }
                    results.append(result)
                    forward_entailment_probs.append(forward_entailment_prob)
                    reverse_entailment_probs.append(reverse_entailment_prob)
                # print(forward_entailment_probs)
                # print(reverse_entailment_probs)
            # break

        # Calculate and print the average entailment probability for True samples in both directions
        if forward_entailment_probs and reverse_entailment_probs:
            avg_forward_entailment_prob = sum(forward_entailment_probs) / len(
                forward_entailment_probs
            )
            avg_reverse_entailment_prob = sum(reverse_entailment_probs) / len(
                reverse_entailment_probs
            )
            print(
                f"Average Forward Entailment Probability for True samples: {avg_forward_entailment_prob:.4f}"
            )
            print(
                f"Average Reverse Entailment Probability for True samples: {avg_reverse_entailment_prob:.4f}"
            )

            # Calculate and print the overall average of both directions
            overall_avg_entailment_prob = (
                avg_forward_entailment_prob + avg_reverse_entailment_prob
            ) / 2
            print(
                f"Overall Average Entailment Probability: {overall_avg_entailment_prob:.4f}"
            )
        else:
            print("No DeBERTa True predictions found for both directions.")

        # save results
        if self.save_path:
            with jsonlines.open(self.save_path, mode="w") as writer:
                writer.write_all(results)

    def compute_syntactic_similarity(self):
        pass


if __name__ == "__main__":
    sim = SemanticSimilarity(
        data_path="t5_squad_counterfactuals/rag_counterfactuals_complete_noise_min_filtered_final_dedup_1.jsonl",
        save_path="",  # "semantic_uncertainty_rag_ent_neu.jsonl",
        batch_size=32,
    )
    sim.compute_entailment()

src/cf_generation/analysis/semantic_uncertainty.py
UKPLab/CATfOOD
"""
context_text <hl>answer<hl> context_text
target: question_text
"""

import os

import numpy as np
from tqdm import tqdm

import jsonlines

DOC_STRIDE = 320
MAX_LENGTH = 640
SEED = 42


def _add_eos_examples(example):
    pass


def get_strided_contexts_and_ans(
    example, tokenizer, doc_stride=256, max_length=512, assertion=True
):
    # overlap will be of doc_stride - q_len

    context = example["context"]
    answer = example["answer"]
    answer_text = answer["text"]
    start_token = answer["start_token"]
    end_token = answer["end_token"]
    hl_token = "<hl>"
    # sep_token = "<sep>"

    # if example["answer"]["category"] == 3:
    #     print(answer)

    # input = f"generate question: {context[:start_token]} {hl_token} {answer} {hl_token} {context[end_token + 1:]}"

    # process yes/no questions
    if example["answer"]["category"] in [3, 4]:
        input = f"generate question: {context} {hl_token} {answer_text[0]} {hl_token}"
    else:
        # for short/long questions
        input = f"generate question: {context[:start_token]} {hl_token} {answer_text} {hl_token} {context[end_token + 1:]}"

    label = example["question"]

    # print(input)
    # print(label)
    input_ids = tokenizer(input)["input_ids"]
    label_ids = tokenizer(label, max_length=256, padding=True, truncation=True)
    label_ids = np.array(label_ids["input_ids"])

    # print(tokenizer.pad_token_id)
    # print(label_ids)
    label_ids[label_ids == tokenizer.pad_token_id] = -100
    # print(label_ids)
    #
    # # if (isinstance(label_ids, int)):
    # print(label)

    # print(input_ids)
    # print(tokenizer(context)["input_ids"])

    cmd_indices = input_ids[:3]
    cmd_len = len(cmd_indices)
    end_token_id = input_ids[-1]  # </s>
    # print(input_ids)
    # print(end_token_id)
    # print(cmd_indices)

    # print(len(input_ids))

    # input_ids = tokenizer(example["question"]["text"], out["context"]).input_ids
    # q_len = input_ids.index(tokenizer.sep_token_id) + 1
    #
    # return yes/no
    if answer["category"] in [3, 4]:  # category is an integer
        inputs = []
        category = []
        doc_start_indices = range(cmd_len, len(input_ids), max_length - doc_stride)
        for i in doc_start_indices:
            end_index = i + max_length - cmd_len
            slice = input_ids[i:end_index]
            inputs.append(cmd_indices + slice)
            if inputs[-1] != end_token:
                inputs.append(end_token)
            category.append(answer["category"])
            if slice[-1] == tokenizer.sep_token_id:
                break

        # for input in inputs:
        #     print(input)
        # print({
        #     "example_id": example["id"],
        #     "input_ids": inputs,
        #     "labels": [[label_ids]] * len(category),
        #     "answers": {
        #         "start_token": [-100] * len(category),
        #         "end_token": [-100] * len(category),
        #         "category": category,
        #     },
        # })

        return {
            "example_id": example["id"],
            "input_ids": inputs,
            "labels": [label_ids.tolist()] * len(category),
            "answers": {
                "start_token": [-100] * len(category),
                "end_token": [-100] * len(category),
                "category": category,
            },
        }

    # print(answer)
    splitted_context = example["context"].split(" ")
    complete_end_token = splitted_context[answer["end_token"]]
    # print(complete_end_token)
    answer["start_token"] = len(
        tokenizer(
            " ".join(splitted_context[: answer["start_token"]]),
            add_special_tokens=True,
        ).input_ids
    )
    answer["end_token"] = len(
        tokenizer(
            " ".join(splitted_context[: answer["end_token"]]), add_special_tokens=True
        ).input_ids
    )
    # print(answer)

    answer["start_token"] += cmd_len + 1  # include highlight token
    answer["end_token"] += cmd_len + 1  # include highlight token

    # fixing end token
    num_sub_tokens = len(
        tokenizer(complete_end_token, add_special_tokens=False).input_ids
    )
    if num_sub_tokens > 1:
        answer["end_token"] += num_sub_tokens - 1

    old = input_ids[
        answer["start_token"] : answer["end_token"] + 1
    ]  # right & left are inclusive
    start_token = answer["start_token"]
    end_token = answer["end_token"]

    # print(start_token, end_token)

    if assertion:
        """This won't match exactly because of extra gaps => visually inspect everything"""
        new = tokenizer.decode(old)
        if answer_text != new:
            print("ISSUE IN TOKENIZATION")
            print("OLD:", answer_text)
            print("NEW:", new, end="\n\n")

    if len(input_ids) <= max_length:
        # print(label)
        # print(input)
        # print(input_ids)
        # print(label_ids)
        # print("here")
        return {
            "example_id": example["id"],
            "input_ids": [input_ids],
            "labels": [label_ids.tolist()],
            "answers": {
                "start_token": [answer["start_token"]],
                "end_token": [answer["end_token"]],
                "category": [answer["category"]],
            },
        }

    doc_start_indices = range(cmd_len, len(input_ids), max_length - doc_stride)
    # print(doc_start_indices)

    inputs = []
    answers_start_token = []
    answers_end_token = []
    answers_category = []  # null, yes, no, long, short
    for i in doc_start_indices:
        end_index = i + max_length - cmd_len - 1
        slice = input_ids[i:end_index]
        if slice[-1] != end_token_id:
            slice.append(end_token_id)
        inputs.append(cmd_indices + slice)
        assert len(inputs[-1]) <= max_length, "Issue in truncating length"

        if start_token >= i and end_token <= end_index - 1:
            start_token = start_token
            end_token = end_token
            answers_category.append(answer["category"])  # ["short"] -> "short"
        else:
            start_token = -100
            end_token = -100
            answers_category.append("null")
        # print(start_token, end_token)
        new = inputs[-1][start_token : end_token + 1]

        answers_start_token.append(start_token)
        answers_end_token.append(end_token)
        if assertion:
            """checking if above code is working as expected for all the samples"""
            if new != old and new != [tokenizer.cls_token_id]:
                print("ISSUE in strided for ID:", example["id"])
                print("New:", tokenizer.decode(new))
                print("Old:", tokenizer.decode(old), end="\n\n")
        # if slice[-1] == end_token_id:
        #     break
    # print(tokenizer.sep_token_id)
    # print(answer_text)
    # print(answers_start_token)
    # print(answers_end_token)
    # print(inputs[0])
    # print(inputs[1])
    # print(inputs)
    # print(len(inputs))
    # print(len(inputs[0]))

    # for input in inputs:
    #     print(input)
    return {
        "example_id": example["id"],
        "input_ids": inputs,
        "labels": [label_ids.tolist()] * len(inputs),
        "answers": {
            "start_token": answers_start_token,
            "end_token": answers_end_token,
            "category": answers_category,
        },
    }


def prepare_inputs(example, tokenizer, doc_stride=256, max_length=512, assertion=False):
    example = get_strided_contexts_and_ans(
        example,
        tokenizer,
        doc_stride=doc_stride,
        max_length=max_length,
        assertion=assertion,
    )

    return example


def save_to_disk(hf_data, file_name):
    with jsonlines.open(file_name, "a") as writer:
        for example in tqdm(hf_data, total=len(hf_data), desc="Saving samples ... "):
            start = example["answers"]["start_token"]
            end = example["answers"]["end_token"]
            cat = example["answers"]["category"]
            # print(start, end , cat)
            for input_ids, labels, start, end, cat in zip(
                example["input_ids"], example["labels"], start, end, cat
            ):
                if start == -100 and end == -100:
                    continue  # remove unanswerable questions

                writer.write(
                    {
                        "id": example["example_id"],
                        "input_ids": input_ids,
                        "labels": labels,
                        "answers": {
                            "start_token": start,
                            "end_token": end,
                            "category": cat,
                        },
                    },
                )


if __name__ == "__main__":
    """Running area"""
    # from datasets import load_dataset
    from transformers import T5Tokenizer

    tokenizer = T5Tokenizer.from_pretrained("t5-base")
    tokenizer.add_special_tokens({"additional_special_tokens": ["<hl>"]})
    # model.resize_token_embeddings(len(tokenizer))
    c = 0
    data = []
    cache_file_name = "data/nq-train-tokenized-short-qg"
    with jsonlines.open("data/nq-train-short-only.jsonl") as file:
        for line in tqdm(file.iter()):
            # print("--------------------------")
            c += 1
            # if line["answer"]["category"] == 3:
            # print(line)
            example = prepare_inputs(
                line, tokenizer, max_length=MAX_LENGTH, doc_stride=DOC_STRIDE
            )
            # print(example)
            # if c == 100:
            #     break
            data.append(example)
            if c % 1000 == 0:
                save_to_disk(data, file_name=cache_file_name + ".jsonl")
                data = []
            if c % 5000 == 0:
                print(f"Samples processed: {c}")
    print(c)

    # with jsonlines.open("data/nq-tokenized-qg.jsonl", "r") as file:
    #     for line in file.iter():
    #         print(line)

src/cf_generation/baseline_generation/batch_nq.py
UKPLab/CATfOOD
"""
format:
question <sep_token> context

sep_token: ">>"

target: answer_text
"""

import os

import numpy as np
from tqdm import tqdm

import jsonlines

DOC_STRIDE = 320
MAX_LENGTH = 640
SEED = 42


def get_strided_contexts_and_ans(
    example, tokenizer, doc_stride=256, max_length=512, assertion=True
):
    # overlap will be of doc_stride - q_len

    question = example["question"]
    context = example["context"]
    answer = example["answer"]
    title = example["title"]
    answer_text = answer["text"]
    start_token = answer["start_token"]
    end_token = answer["end_token"]
    sep_token = ">>"

    input = f"{question} {sep_token} {context}"
    label = answer["text"]

    input_ids = tokenizer(input)["input_ids"]
    label_ids = tokenizer(label, max_length=256, padding=True, truncation=True)
    label_ids = np.array(label_ids["input_ids"])

    sep_token_id = tokenizer.convert_tokens_to_ids(sep_token)
    q_len = input_ids.index(sep_token_id) + 1
    q_indices = input_ids[:q_len]
    end_token_id = input_ids[-1]  # </s>
    # print(input_ids)

    # 1. return yes/no
    if answer["category"] in [3, 4]:  # category is an integer
        inputs = []
        category = []
        doc_start_indices = range(q_len, len(input_ids), max_length - doc_stride)
        for i in doc_start_indices:
            end_index = i + max_length - q_len
            slice = input_ids[i:end_index]
            inputs.append(q_indices + slice)
            if inputs[-1] != end_token:
                inputs.append(end_token)
            category.append(answer["category"])
            # if slice[-1] == tokenizer.sep_token_id:
            #     break

        return {
            "example_id": example["id"],
            "input_ids": inputs,
            "labels": [label_ids.tolist()] * len(category),
            "answers": {
                "start_token": [-100] * len(category),
                "end_token": [-100] * len(category),
                "category": category,
            },
        }

    # print(answer)
    splitted_context = example["context"].split(" ")
    complete_end_token = splitted_context[answer["end_token"]]
    # print(complete_end_token)
    answer["start_token"] = len(
        tokenizer(
            " ".join(splitted_context[: answer["start_token"]]),
            add_special_tokens=True,
        ).input_ids
    )
    answer["end_token"] = len(
        tokenizer(
            " ".join(splitted_context[: answer["end_token"]]), add_special_tokens=True
        ).input_ids
    )
    # print(answer)

    answer["start_token"] += q_len - 1  # include sep
    answer["end_token"] += q_len - 1

    # fixing end token
    num_sub_tokens = len(
        tokenizer(complete_end_token, add_special_tokens=False).input_ids
    )
    if num_sub_tokens > 1:
        answer["end_token"] += num_sub_tokens - 1

    old = input_ids[
        answer["start_token"] : answer["end_token"] + 1
    ]  # right & left are inclusive
    start_token = answer["start_token"]
    end_token = answer["end_token"]

    # print(start_token, end_token)

    if assertion:
        """This won't match exactly because of extra gaps => visually inspect everything"""
        new = tokenizer.decode(old)
        if answer_text != new:
            print("ISSUE IN TOKENIZATION")
            print("OLD:", answer_text)
            print("NEW:", new, end="\n\n")

    # 2. when input length is less than our max_length
    if len(input_ids) <= max_length:
        return {
            "example_id": example["id"],
            "input_ids": [input_ids],
            "labels": [label_ids.tolist()],
            "answers": {
                "start_token": [answer["start_token"]],
                "end_token": [answer["end_token"]],
                "category": [answer["category"]],
            },
        }

    doc_start_indices = range(q_len, len(input_ids), max_length - doc_stride)
    # print(doc_start_indices)

    # 3. for strided qa samples
    inputs = []
    answers_start_token = []
    answers_end_token = []
    answers_category = []  # null, yes, no, long, short
    for i in doc_start_indices:
        end_index = i + max_length - q_len - 1
        slice = input_ids[i:end_index]
        if slice[-1] != end_token_id:
            slice.append(end_token_id)
        inputs.append(q_indices + slice)
        assert len(inputs[-1]) <= max_length, "Issue in truncating length"

        if start_token >= i and end_token <= end_index - 1:
            start_token = start_token
            end_token = end_token
            answers_category.append(answer["category"])  # ["short"] -> "short"
        else:
            start_token = -100
            end_token = -100
            answers_category.append("null")
        # print(start_token, end_token)
        new = inputs[-1][start_token : end_token + 1]

        answers_start_token.append(start_token)
        answers_end_token.append(end_token)
        if assertion:
            """checking if above code is working as expected for all the samples"""
            if new != old and new != [tokenizer.cls_token_id]:
                print("ISSUE in strided for ID:", example["id"])
                print("New:", tokenizer.decode(new))
                print("Old:", tokenizer.decode(old), end="\n\n")

    return {
        "example_id": example["id"],
        "input_ids": inputs,
        "labels": [label_ids.tolist()] * len(inputs),
        "answers": {
            "start_token": answers_start_token,
            "end_token": answers_end_token,
            "category": answers_category,
        },
    }


def prepare_inputs(example, tokenizer, doc_stride=256, max_length=512, assertion=True):
    example = get_strided_contexts_and_ans(
        example,
        tokenizer,
        doc_stride=doc_stride,
        max_length=max_length,
        assertion=assertion,
    )

    return example


def save_to_disk(hf_data, file_name):
    with jsonlines.open(file_name, "a") as writer:
        for example in tqdm(hf_data, total=len(hf_data), desc="Saving samples ... "):
            start = example["answers"]["start_token"]
            end = example["answers"]["end_token"]
            cat = example["answers"]["category"]
            # print(start, end , cat)
            for input_ids, labels, start, end, cat in zip(
                example["input_ids"], example["labels"], start, end, cat
            ):
                if start == -100 and end == -100:
                    continue  # remove unanswerable questions

                writer.write(
                    {
                        "id": example["example_id"],
                        "input_ids": input_ids,
                        "labels": labels,
                        "answers": {
                            "start_token": start,
                            "end_token": end,
                            "category": cat,
                        },
                    },
                )


if __name__ == "__main__":
    """Running area"""
    # from datasets import load_dataset
    from transformers import T5Tokenizer

    tokenizer = T5Tokenizer.from_pretrained("t5-base")
    tokenizer.add_special_tokens({"additional_special_tokens": [">>"]})
    # print(tokenizer.convert_tokens_to_ids(">>"))
    # model.resize_token_embeddings(len(tokenizer))
    c = 0
    data = []
    cache_file_name = "data/nq-train-tokenized-short-qa"
    with jsonlines.open("data/nq-train-short-only.jsonl") as file:
        for line in tqdm(file.iter()):
            # print("--------------------------")
            c += 1
            # if line["answer"]["category"] == 3:
            # print(line)
            example = prepare_inputs(
                line,
                tokenizer,
                max_length=MAX_LENGTH,
                doc_stride=DOC_STRIDE,
                assertion=False,
            )
            # print(example)
            # if c == 10:
            #     break
            data.append(example)
            if c % 1000 == 0:
                save_to_disk(data, file_name=cache_file_name + ".jsonl")
                data = []
            if c % 5000 == 0:
                print(f"Samples processed: {c}")
    print(c)

src/cf_generation/baseline_generation/batch_nq_qa.py
UKPLab/CATfOOD
import os

import numpy as np
from tqdm import tqdm

import jsonlines

DOC_STRIDE = 320
MAX_LENGTH = 640
SEED = 42


def _add_eos_examples(example):
    pass


def get_strided_contexts_and_ans(
    example, tokenizer, doc_stride=256, max_length=512, assertion=True
):
    # overlap will be of doc_stride - q_len

    context = example["context"]
    answer = example["answer"]
    title = example["title"]
    answer_text = answer["text"]
    start_token = answer["start_token"]
    end_token = answer["end_token"]
    hl_token = "<<"
    sep_token = ">>"
    # sep_token = "<sep>"

    # if example["answer"]["category"] == 3:
    #     print(answer)

    # input = f"generate question: {context[:start_token]} {hl_token} {answer} {hl_token} {context[end_token + 1:]}"

    # process yes/no questions
    if example["answer"]["category"] in [3, 4]:
        input = f"{title} {sep_token} {context} {hl_token} answer = {answer_text[0]} {sep_token}"
    else:
        # for short/long questions
        input = f"{title} {sep_token} {context[:start_token]} {hl_token} answer = {answer_text} {sep_token} {context[end_token + 1:]}"

    label = example["question"]

    # print(input)
    # print(label)
    input_ids = tokenizer(input)["input_ids"]
    label_ids = tokenizer(label)
    label_ids = np.array(label_ids["input_ids"])
    # print(tokenizer.pad_token_id)
    # print(label_ids)
    label_ids[label_ids == tokenizer.pad_token_id] = -100
    # print(label_ids)

    # print(input_ids)
    # print(tokenizer(context)["input_ids"])
    sep_token_id = tokenizer.convert_tokens_to_ids(sep_token)
    title_len = input_ids.index(sep_token_id) + 1
    title_indices = input_ids[:title_len]
    end_token_id = input_ids[-1]  # </s>
    # print(input_ids)
    #
    # return yes/no
    if answer["category"] in [3, 4]:  # category is an integer
        inputs = []
        category = []
        doc_start_indices = range(title_len, len(input_ids), max_length - doc_stride)
        for i in doc_start_indices:
            end_index = i + max_length - title_len
            slice = input_ids[i:end_index]
            inputs.append(title_indices + slice)
            if inputs[-1] != end_token:
                inputs.append(end_token)
            category.append(answer["category"])
            if slice[-1] == tokenizer.sep_token_id:
                break

        # for input in inputs:
        #     print(input)
        # print({
        #     "example_id": example["id"],
        #     "input_ids": inputs,
        #     "labels": [[label_ids]] * len(category),
        #     "answers": {
        #         "start_token": [-100] * len(category),
        #         "end_token": [-100] * len(category),
        #         "category": category,
        #     },
        # })

        return {
            "example_id": example["id"],
            "input_ids": inputs,
            "labels": [label_ids.tolist()] * len(category),
            "answers": {
                "start_token": [-100] * len(category),
                "end_token": [-100] * len(category),
                "category": category,
            },
        }

    # print(answer)
    splitted_context = example["context"].split(" ")
    complete_end_token = splitted_context[answer["end_token"]]
    # print(complete_end_token)
    answer["start_token"] = len(
        tokenizer(
            " ".join(splitted_context[: answer["start_token"]]),
            add_special_tokens=True,
        ).input_ids
    )
    answer["end_token"] = len(
        tokenizer(
            " ".join(splitted_context[: answer["end_token"]]), add_special_tokens=True
        ).input_ids
    )
    # print(answer)

    answer["start_token"] += title_len + 1  # include highlight token
    answer["end_token"] += title_len + 1  # include highlight token

    # fixing end token
    num_sub_tokens = len(
        tokenizer(complete_end_token, add_special_tokens=False).input_ids
    )
    if num_sub_tokens > 1:
        answer["end_token"] += num_sub_tokens - 1

    old = input_ids[
        answer["start_token"] : answer["end_token"] + 1
    ]  # right & left are inclusive
    start_token = answer["start_token"]
    end_token = answer["end_token"]

    # print(start_token, end_token)

    if assertion:
        """This won't match exactly because of extra gaps => visually inspect everything"""
        new = tokenizer.decode(old)
        if answer_text != new:
            print("ISSUE IN TOKENIZATION")
            print("OLD:", answer_text)
            print("NEW:", new, end="\n\n")

    if len(input_ids) <= max_length:
        return {
            "example_id": example["id"],
            "input_ids": [input_ids],
            "labels": label_ids.tolist(),
            "answers": {
                "start_token": [answer["start_token"]],
                "end_token": [answer["end_token"]],
                "category": [answer["category"]],
            },
        }

    doc_start_indices = range(title_len, len(input_ids), max_length - doc_stride)
    # print(doc_start_indices)

    inputs = []
    answers_start_token = []
    answers_end_token = []
    answers_category = []  # null, yes, no, long, short
    for i in doc_start_indices:
        end_index = i + max_length - title_len - 1
        slice = input_ids[i:end_index]
        if slice[-1] != end_token_id:
            slice.append(end_token_id)
        inputs.append(title_indices + slice)
        assert len(inputs[-1]) <= max_length, "Issue in truncating length"

        if start_token >= i and end_token <= end_index - 1:
            start_token = start_token
            end_token = end_token
            answers_category.append(answer["category"])  # ["short"] -> "short"
        else:
            start_token = -100
            end_token = -100
            answers_category.append("null")
        # print(start_token, end_token)
        new = inputs[-1][start_token : end_token + 1]

        answers_start_token.append(start_token)
        answers_end_token.append(end_token)
        if assertion:
            """checking if above code is working as expected for all the samples"""
            if new != old and new != [tokenizer.cls_token_id]:
                print("ISSUE in strided for ID:", example["id"])
                print("New:", tokenizer.decode(new))
                print("Old:", tokenizer.decode(old), end="\n\n")
        # if slice[-1] == end_token_id:
        #     break
    # print(tokenizer.sep_token_id)
    # print(answer_text)
    # print(answers_start_token)
    # print(answers_end_token)
    # print(inputs[0])
    # print(inputs[1])
    # print(inputs)
    # print(len(inputs))
    # print(len(inputs[0]))

    # for input in inputs:
    #     print(input)
    return {
        "example_id": example["id"],
        "input_ids": inputs,
        "labels": [label_ids.tolist()] * len(inputs),
        "answers": {
            "start_token": answers_start_token,
            "end_token": answers_end_token,
            "category": answers_category,
        },
    }


def prepare_inputs(example, tokenizer, doc_stride=256, max_length=512, assertion=False):
    example = get_strided_contexts_and_ans(
        example,
        tokenizer,
        doc_stride=doc_stride,
        max_length=max_length,
        assertion=assertion,
    )

    return example


def save_to_disk(hf_data, file_name):
    with jsonlines.open(file_name, "a") as writer:
        for example in tqdm(hf_data, total=len(hf_data), desc="Saving samples ... "):
            start = example["answers"]["start_token"]
            end = example["answers"]["end_token"]
            cat = example["answers"]["category"]
            # print(start, end , cat)
            for input_ids, labels, start, end, cat in zip(
                example["input_ids"], example["labels"], start, end, cat
            ):
                if start == -100 and end == -100:
                    continue  # remove unanswerable questions
                # if cat == 0:
                #     continue
                writer.write(
                    {
                        "id": example["example_id"],
                        "input_ids": input_ids,
                        "labels": labels,
                        "answers": {
                            "start_token": start,
                            "end_token": end,
                            "category": cat,
                        },
                    },
                )


if __name__ == "__main__":
    """Running area"""
    # from datasets import load_dataset
    from transformers import T5Tokenizer

    tokenizer = T5Tokenizer.from_pretrained("t5-base")
    tokenizer.add_special_tokens({"additional_special_tokens": ["<<", ">>"]})
    print(tokenizer.convert_tokens_to_ids(">>"))
    # model.resize_token_embeddings(len(tokenizer))
    # c = 0
    # data = []
    # cache_file_name = "data/nq-tokenized-short-qg"
    # with jsonlines.open("data/nq-train-short-only.jsonl") as file:
    #     for line in tqdm(file.iter()):
    #         # print("--------------------------")
    #         c += 1
    #         # if line["answer"]["category"] == 3:
    #         # print(line)
    #         example = prepare_inputs(line, tokenizer, max_length=MAX_LENGTH, doc_stride=DOC_STRIDE)
    # print(example)
    # if c == 100:
    #     break
    #         data.append(example)
    #         if c % 1000 == 0:
    #             save_to_disk(data, file_name=cache_file_name + ".jsonl")
    #             data = []
    #         if c % 5000 == 0:
    #             print(f"Samples processed: {c}")
    # print(c)

    # with jsonlines.open("data/nq-tokenized-qg.jsonl", "r") as file:
    #     for line in file.iter():
    #         print(line)

src/cf_generation/baseline_generation/batch_nq_short.py
UKPLab/CATfOOD
import os

import numpy as np
from tqdm import tqdm

from datasets import load_dataset
from transformers import T5Tokenizer

import jsonlines

DOC_STRIDE = 320
MAX_LENGTH = 640
SEED = 42

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


def _add_eos_examples(example):
    pass


def get_strided_contexts_and_ans(
    example, tokenizer, doc_stride=256, max_length=512, assertion=True
):
    # overlap will be of doc_stride - q_len

    context = example["context"]
    answer = example["answers"]
    answer_text = answer["text"][0]
    start_token = answer["answer_start"][0]
    end_token = start_token + len(answer_text)
    hl_token = "<hl>"
    # sep_token = "<sep>"

    # for short/long questions
    input = f"generate question: {context[:start_token]} {hl_token} {answer_text} {hl_token} {context[end_token + 1:]}"
    label = example["question"]
    input_ids = tokenizer(input)["input_ids"]
    label_ids = tokenizer(label, max_length=256, padding=True, truncation=True)
    label_ids = np.array(label_ids["input_ids"])

    label_ids[label_ids == tokenizer.pad_token_id] = -100

    cmd_indices = input_ids[:3]
    cmd_len = len(cmd_indices)
    end_token_id = input_ids[-1]  # </s>

    splitted_context = example["context"].split(" ")
    complete_end_token = splitted_context[end_token]
    # print(complete_end_token)
    answer["start_token"] = len(
        tokenizer(
            " ".join(splitted_context[:start_token]), add_special_tokens=True,
        ).input_ids
    )
    answer["end_token"] = len(
        tokenizer(
            " ".join(splitted_context[:end_token]), add_special_tokens=True
        ).input_ids
    )
    # print(answer)

    answer["start_token"] += cmd_len + 1  # include highlight token
    answer["end_token"] += cmd_len + 1  # include highlight token

    # fixing end token
    num_sub_tokens = len(
        tokenizer(complete_end_token, add_special_tokens=False).input_ids
    )
    if num_sub_tokens > 1:
        answer["end_token"] += num_sub_tokens - 1

    old = input_ids[
        answer["start_token"] : answer["end_token"] + 1
    ]  # right & left are inclusive
    start_token = answer["start_token"]
    end_token = answer["end_token"]

    # print(start_token, end_token)

    if assertion:
        """This won't match exactly because of extra gaps => visually inspect everything"""
        new = tokenizer.decode(old)
        if answer_text != new:
            print("ISSUE IN TOKENIZATION")
            print("OLD:", answer_text)
            print("NEW:", new, end="\n\n")

    if len(input_ids) <= max_length:
        return {
            "example_id": example["id"],
            "input_ids": [input_ids],
            "labels": [label_ids.tolist()],
            "answers": {
                "start_token": [answer["start_token"]],
                "end_token": [answer["end_token"]],
            },
        }

    doc_start_indices = range(cmd_len, len(input_ids), max_length - doc_stride)
    # print(doc_start_indices)

    inputs = []
    answers_start_token = []
    answers_end_token = []
    for i in doc_start_indices:
        end_index = i + max_length - cmd_len - 1
        slice = input_ids[i:end_index]
        if slice[-1] != end_token_id:
            slice.append(end_token_id)
        inputs.append(cmd_indices + slice)
        assert len(inputs[-1]) <= max_length, "Issue in truncating length"

        if start_token >= i and end_token <= end_index - 1:
            start_token = start_token
            end_token = end_token
        else:
            start_token = -100
            end_token = -100
        # print(start_token, end_token)
        new = inputs[-1][start_token : end_token + 1]

        answers_start_token.append(start_token)
        answers_end_token.append(end_token)
        if assertion:
            """checking if above code is working as expected for all the samples"""
            if new != old and new != [tokenizer.cls_token_id]:
                print("ISSUE in strided for ID:", example["id"])
                print("New:", tokenizer.decode(new))
                print("Old:", tokenizer.decode(old), end="\n\n")

    return {
        "example_id": example["id"],
        "input_ids": inputs,
        "labels": [label_ids.tolist()] * len(inputs),
        "answers": {
            "start_token": answers_start_token,
            "end_token": answers_end_token,
        },
    }


def prepare_inputs(example, tokenizer, doc_stride=256, max_length=512, assertion=False):
    example = get_strided_contexts_and_ans(
        example,
        tokenizer,
        doc_stride=doc_stride,
        max_length=max_length,
        assertion=assertion,
    )

    return example


def save_to_disk(hf_data, file_name):
    with jsonlines.open(file_name, "a") as writer:
        for example in tqdm(hf_data, total=len(hf_data), desc="Saving samples ... "):
            start = example["answers"]["start_token"]
            end = example["answers"]["end_token"]
            for input_ids, labels, start, end, cat in zip(
                example["input_ids"], example["labels"], start, end, cat
            ):
                if start == -100 and end == -100:
                    continue  # remove unanswerable questions

                writer.write(
                    {
                        "id": example["example_id"],
                        "input_ids": input_ids,
                        "labels": labels,
                        "answers": {"start_token": start, "end_token": end,},
                    },
                )


if __name__ == "__main__":
    """Running area"""

    tokenizer = T5Tokenizer.from_pretrained("t5-base")
    tokenizer.add_special_tokens({"additional_special_tokens": ["<hl>"]})
    # model.resize_token_embeddings(len(tokenizer))
    c = 0
    data = []
    cache_file_name = os.path.join(BASE_PATH, "src/data/squad/squad_qg_tokenized")

    # # load squad data
    dataset = load_dataset("squad", "plain_text")
    train_data = dataset["train"]
    squad_data = [
        sample
        for sample in tqdm(
            train_data, total=len(train_data), desc="Loading SQuAD data ... "
        )
    ]
    #
    for ex in tqdm(train_data):
        example = prepare_inputs(
            ex, tokenizer, max_length=MAX_LENGTH, doc_stride=DOC_STRIDE
        )
        data.append(example)
        if c % 1000 == 0:
            save_to_disk(data, file_name=cache_file_name + ".jsonl")
            data = []
        if c % 5000 == 0:
            print(f"Samples processed: {c}")
    print(c)

src/cf_generation/baseline_generation/batch_squad.py
UKPLab/CATfOOD
"""
Utilities for data handling.
"""
import json
import logging
import os
import pandas as pd
import shutil
import re
import random
from tqdm import tqdm

from typing import Dict

logger = logging.getLogger(__name__)


def read_data(file_path: str, task_name: str, guid_as_int: bool = False):
    """
  Reads task-specific datasets from corresponding GLUE-style TSV files.
  """
    logger.warning(
        "Data reading only works when data is in TSV format, "
        " and last column as classification label."
    )

    # `guid_index`: should be 2 for SNLI, 0 for MNLI and None for any random tsv file.
    if task_name == "MNLI":
        return read_glue_tsv(file_path, guid_index=0, guid_as_int=guid_as_int)
    elif task_name == "SNLI":
        return read_glue_tsv(file_path, guid_index=2, guid_as_int=guid_as_int)
    elif task_name == "WINOGRANDE":
        return read_glue_tsv(file_path, guid_index=0)
    elif task_name == "QNLI":
        return read_glue_tsv(file_path, guid_index=0)
    else:
        raise NotImplementedError(f"Reader for {task_name} not implemented.")


def convert_tsv_entries_to_dataframe(tsv_dict: Dict, header: str) -> pd.DataFrame:
    """
  Converts entries from TSV file to Pandas DataFrame for faster processing.
  """
    header_fields = header.strip().split("\t")
    data = {header: [] for header in header_fields}

    for line in tsv_dict.values():
        fields = line.strip().split("\t")
        assert len(header_fields) == len(fields)
        for field, header in zip(fields, header_fields):
            data[header].append(field)

    df = pd.DataFrame(data, columns=header_fields)
    return df


def copy_dev_test(
    task_name: str, from_dir: os.path, to_dir: os.path, extension: str = ".tsv"
):
    """
  Copies development and test sets (for data selection experiments) from `from_dir` to `to_dir`.
  """
    if task_name == "MNLI":
        dev_filename = "dev_matched.tsv"
        test_filename = "dev_mismatched.tsv"
    elif task_name in ["SNLI", "QNLI", "WINOGRANDE"]:
        dev_filename = f"dev{extension}"
        test_filename = f"test{extension}"
    else:
        raise NotImplementedError(f"Logic for {task_name} not implemented.")

    dev_path = os.path.join(from_dir, dev_filename)
    if os.path.exists(dev_path):
        shutil.copyfile(dev_path, os.path.join(to_dir, dev_filename))
    else:
        raise ValueError(f"No file found at {dev_path}")

    test_path = os.path.join(from_dir, test_filename)
    if os.path.exists(test_path):
        shutil.copyfile(test_path, os.path.join(to_dir, test_filename))
    else:
        raise ValueError(f"No file found at {test_path}")


def read_jsonl(file_path: str, key: str = "pairID"):
    """
  Reads JSONL file to recover mapping between one particular key field
  in the line and the result of the line as a JSON dict.
  If no key is provided, return a list of JSON dicts.
  """
    df = pd.read_json(file_path, lines=True)
    records = df.to_dict("records")
    logger.info(f"Read {len(records)} JSON records from {file_path}.")

    if key:
        assert key in df.columns
        return {record[key]: record for record in records}
    return


def read_glue_tsv(
    file_path: str, guid_index: int, label_index: int = -1, guid_as_int: bool = False
):
    """
  Reads TSV files for GLUE-style text classification tasks.
  Returns:
    - a mapping between the example ID and the entire line as a string.
    - the header of the TSV file.
  """
    tsv_dict = {}

    i = -1
    with open(file_path, "r") as tsv_file:
        for line in tqdm.tqdm([line for line in tsv_file]):
            i += 1
            if i == 0:
                header = line.strip()
                field_names = line.strip().split("\t")
                continue

            fields = line.strip().split("\t")
            label = fields[label_index]
            if len(fields) > len(field_names):
                # SNLI / MNLI fields sometimes contain multiple annotator labels.
                # Ignore all except the gold label.
                reformatted_fields = fields[: len(field_names) - 1] + [label]
                assert len(reformatted_fields) == len(field_names)
                reformatted_line = "\t".join(reformatted_fields)
            else:
                reformatted_line = line.strip()

            if label == "-" or label == "":
                logger.info(f"Skippping line: {line}")
                continue

            if guid_index is None:
                guid = i
            else:
                guid = fields[guid_index]  # PairID.
            if guid in tsv_dict:
                logger.info(f"Found clash in IDs ... skipping example {guid}.")
                continue
            tsv_dict[guid] = reformatted_line.strip()

    logger.info(
        f"Read {len(tsv_dict)} valid examples, with unique IDS, out of {i} from {file_path}"
    )
    if guid_as_int:
        tsv_numeric = {
            int(convert_string_to_unique_number(k)): v for k, v in tsv_dict.items()
        }
        return tsv_numeric, header
    return tsv_dict, header


def convert_string_to_unique_number(string: str) -> int:
    """
  Hack to convert SNLI ID into a unique integer ID, for tensorizing.
  """
    id_map = {"e": "0", "c": "1", "n": "2"}

    # SNLI-specific hacks.
    if string.startswith("vg_len"):
        code = "555"
    elif string.startswith("vg_verb"):
        code = "444"
    else:
        code = "000"

    try:
        number = int(code + re.sub(r"\D", "", string) + id_map.get(string[-1], "3"))
    except:
        number = random.randint(10000, 99999)
        logger.info(f"Cannot find ID for {string}, using random number {number}.")
    return number

src/cf_generation/baseline_generation/cartography/data_utils.py
UKPLab/CATfOOD
import itertools
from pprint import pprint
import os
import json
import jsonlines
import re
from tqdm import tqdm

from typing import List

from collections import OrderedDict

from datasets import load_dataset, concatenate_datasets, Features, Value, Sequence


class PreprocessData:
    def __init__(
        self,
        dataset_name: str,
        dataset_config: str,
        cf_path: str = None,
        save_data: bool = False,
        save_path: str = "",
    ):
        """
        load datasets
        """
        dataset = load_dataset(dataset_name, dataset_config)
        if dataset_name == "squad":
            self.train_set = dataset["train"]
        self.val_set = dataset["validation"]
        self.save_data = save_data
        self.save_path = save_path
        self.cf_path = cf_path
        self.count = 0

    def get_data_info(self):
        """
        data statistics
        """
        train_samples = self.train_set.shape[0]
        val_samples = self.val_set.shape[0]
        column_names = self.train_set.column_names
        train_ex = self.train_set.select([0])[0]
        val_ex = self.val_set.select([0])[0]
        info = {
            "train_samples": train_samples,
            "val_samples": val_samples,
            "column_names": column_names,
            "train_ex": train_ex,
            "val_ex": val_ex,
        }
        return OrderedDict(info)

    def remove_trailing_space(self, example):
        # for questions in squad that have extra starting end space
        example["question"] = example["question"].strip()
        return example

    def remove_white_space(self, example):
        example["question"] = " ".join(example["question"].split())
        example["context"] = " ".join(example["context"].split())
        return example

    def processed_train_val_set(self):
        # filter unanswerable questions
        # print(self.train_set.filter(lambda x: len(x["answers"]["text"]) != 1))
        # self.train_set = self.train_set.select(range(1000))
        # self.val_set = self.val_set.select(range(10))
        if self.cf_path is not None:
            self.train_set = self._add_counterfactuals()
        answerable_train_set = self.train_set.filter(
            lambda x: len(x["answers"]["text"]) != 0
        )
        answerable_train_set = answerable_train_set.map(self.remove_trailing_space)
        answerable_train_set = answerable_train_set.map(self.remove_white_space)

        answerable_val_set = self.val_set.filter(
            lambda x: len(x["answers"]["text"]) != 0
        )
        answerable_val_set = answerable_val_set.map(self.remove_trailing_space)
        answerable_val_set = answerable_val_set.map(self.remove_white_space)
        return answerable_train_set, answerable_val_set

    def processed_val_set(self):
        # filter unanswerable questions
        # print(self.train_set.filter(lambda x: len(x["answers"]["text"]) != 1))
        answerable_val_set = self.val_set.filter(
            lambda x: len(x["answers"]["text"]) != 0
        )
        answerable_val_set = answerable_val_set.map(self.remove_trailing_space)
        answerable_val_set = answerable_val_set.map(self.remove_white_space)
        return answerable_val_set

    def processed_counterfactuals(self):
        """
        process counterfactuals
        """
        cf_data = self._get_counterfactuals()
        answerable_train_set = cf_data.filter(lambda x: len(x["answers"]["text"]) != 0)
        answerable_train_set = answerable_train_set.map(self.remove_trailing_space)
        answerable_train_set = answerable_train_set.map(self.remove_white_space)

        answerable_val_set = self.val_set.filter(
            lambda x: len(x["answers"]["text"]) != 0
        )
        answerable_val_set = answerable_val_set.map(self.remove_trailing_space)
        answerable_val_set = answerable_val_set.map(self.remove_white_space)
        return answerable_train_set, answerable_val_set

    def sanitize_ptb_tokenized_string(self, text: str) -> str:
        """
        Sanitizes string that was tokenized using PTBTokenizer
        """
        tokens = text.split(" ")
        if len(tokens) == 0:
            return text

        # Replace quotation marks and parentheses
        token_map = {
            "``": '"',
            "''": '"',
            "-lrb-": "(",
            "-rrb-": ")",
            "-lsb-": "[",
            "-rsb-": "]",
            "-lcb-": "{",
            "-rcb-": "}",
            "<s>": "",
            "</s>": "",
        }

        # Merge punctuation with previous tokens
        punct_forward = {"`", "$", "#"}
        punct_backward = {".", ",", "!", "?", ":", ";", "%", "'"}

        # Exact matches that get merged forward or backward
        em_forward = {"(", "[", "{"}
        em_backward = {"n't", "na", ")", "]", "}"}

        new_tokens: List[str] = []

        merge_fwd = False
        for i, orig_token in enumerate(tokens):
            tokens[i] = (
                token_map[orig_token.lower()]
                if orig_token.lower() in token_map
                else orig_token
            )
            new_token = tokens[i].lower()

            # merge_fwd was set by previous token, so it should be prepended to current token
            if merge_fwd:
                tokens[i] = tokens[i - 1] + tokens[i]

            if len(tokens[i]) == 0:
                continue

            # Special cases for `` and '', those tells us if " is the start or end of a quotation.
            # Also always merge tokens starting with ' backward and don't merge back if we just merged forward
            merge_bckwd = not merge_fwd and (
                orig_token == "''"
                or new_token in em_backward
                or new_token.startswith("'")
                or all(c in punct_backward for c in new_token)
            )
            merge_fwd = (
                orig_token == "``"
                or new_token in em_forward
                or all(c in punct_forward for c in new_token)
            )

            if merge_bckwd and new_tokens:
                new_tokens[-1] += tokens[i]
            elif not new_tokens or not merge_fwd or i == len(tokens) - 1:
                new_tokens.append(tokens[i])

        return " ".join(new_tokens)

    def _convert_answer_column(self, example):
        """
        convert answer column to squad format
        """

        context = example["context"].strip()
        # remove extra white space
        context = re.sub(r"\s+", " ", context)
        example["context"] = context

        answer = example["answers"]["text"]
        # handle spaces between special characters
        special_chars = ["$", "@"]
        if any(char in answer for char in special_chars):
            for char in special_chars:
                answer = answer.replace(f"{char} ", char)
        # remove spaces between special characters (quotations)
        answer = re.sub(r'"\s*([^"]*?)\s*"', r'"\1"', answer)
        answer = re.sub(r" - ", "-", answer)

        # answer = answer.replace("$ ", "$")
        answer_start = str(context.lower()).find(str(answer.lower()))
        if answer_start == -1:
            self.count += 1
            answer_start = str(context.lower()).find(
                self.sanitize_ptb_tokenized_string(answer).lower()
            )
            if answer_start != -1:
                answer = str(self.sanitize_ptb_tokenized_string(answer))
                self.count -= 1
            else:
                answer = ""

        example["answers"] = {"answer_start": [answer_start], "text": [answer]}
        return example

    def _get_counterfactuals(self):
        """
        Load and process counterfactuals
        """
        counterfactuals_dataset = load_dataset("json", data_files=self.cf_path)
        # counterfactuals_dataset = counterfactuals_dataset.remove_columns("similarity")
        # print(len(counterfactuals_dataset[]))
        new_title_column = [""] * len(counterfactuals_dataset["train"])
        counterfactuals_dataset = counterfactuals_dataset["train"].add_column(
            "title", new_title_column
        )
        # rename answer column
        # counterfactuals_dataset = counterfactuals_dataset.rename_column("answer", "answers")
        counterfactuals_dataset = counterfactuals_dataset.map(
            self._convert_answer_column
        )
        # print(counterfactuals_dataset[:2])
        # counterfactuals_dataset = counterfactuals_dataset.remove_columns("answer")
        # remove samples with no answer
        counterfactuals_dataset = counterfactuals_dataset.filter(
            lambda x: x["answers"]["answer_start"][0] != -1
        )
        counterfactuals_dataset = counterfactuals_dataset.cast(
            Features(
                {
                    "id": Value("string"),
                    "title": Value("string"),
                    "context": Value("string"),
                    "question": Value("string"),
                    "answers": Sequence(
                        {
                            "answer_start": Value("int32"),
                            # "answer_end": Value("int32"),
                            "text": Value("string"),
                        }
                    ),
                }
            )
        )
        return counterfactuals_dataset

    def _add_counterfactuals(self):
        """
        add counterfactuals to data
        """
        counterfactuals_dataset = self._get_counterfactuals()
        self.train_set = self.train_set.cast(
            Features(
                {
                    "id": Value("string"),
                    "title": Value("string"),
                    "context": Value("string"),
                    "question": Value("string"),
                    "answers": Sequence(
                        {"answer_start": Value("int32"), "text": Value("string"),}
                    ),
                }
            )
        )

        merged_dataset = concatenate_datasets([self.train_set, counterfactuals_dataset])
        return merged_dataset


if __name__ == "__main__":
    save_dir = "../data/squad/train_data/"
    BASE_PATH = "//"
    dataloader = PreprocessData(
        "squad",
        "plain_text",
        cf_path=BASE_PATH
        + "src/data/squad/squad_counterfactuals_noise_min_filtered_final_2.jsonl",
        save_data=False,
        save_path=save_dir,
    )
    dataset = dataloader._get_counterfactuals()
    print(len(dataset))

src/cf_generation/baseline_generation/cartography/dataloader.py
UKPLab/CATfOOD
"""
Filtering and dataset mapping methods based on training dynamics.
By default, this module reads training dynamics from a given trained model and
computes the metrics---confidence, variability, correctness,
as well as baseline metrics of forgetfulness and threshold closeness
for each instance in the training data.
If specified, data maps can be plotted with respect to confidence and variability.
Moreover, datasets can be filtered with respect any of the other metrics.
"""
import argparse
import json
import logging
import matplotlib.pyplot as plt
import numpy as np
import os
import pandas as pd
import seaborn as sns
import torch
from tqdm import tqdm
import jsonlines

from collections import defaultdict
from typing import List

from src.rag.cartography.selection_utils import read_training_dynamics
from src.cartography.dataloader import PreprocessData


logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", level=logging.INFO
)
logger = logging.getLogger(__name__)


def compute_forgetfulness(correctness_trend: List[float]) -> int:
    """
    Given a epoch-wise trend of train predictions, compute frequency with which
    an example is forgotten, i.e. predicted incorrectly _after_ being predicted correctly.
    Based on: https://arxiv.org/abs/1812.05159
    """
    if not any(correctness_trend):  # Example is never predicted correctly, or learnt!
        return 1000
    learnt = False  # Predicted correctly in the current epoch.
    times_forgotten = 0
    for is_correct in correctness_trend:
        if (not learnt and not is_correct) or (learnt and is_correct):
            # nothing changed.
            continue
        elif learnt and not is_correct:
            # Forgot after learning at some point!
            learnt = False
            times_forgotten += 1
        elif not learnt and is_correct:
            # Learnt!
            learnt = True
    return times_forgotten


def compute_correctness(trend: List[float]) -> float:
    """
    Aggregate #times an example is predicted correctly during all training epochs.
    """
    return sum(trend)


def compute_train_dy_metrics(training_dynamics, args):
    """
    Given the training dynamics (logits for each training instance across epochs), compute metrics
    based on it, for data map coordinates.
    Computed metrics are: confidence, variability, correctness, forgetfulness, threshold_closeness---
    the last two being baselines from prior work
    (Example Forgetting: https://arxiv.org/abs/1812.05159 and
    Active Bias: https://arxiv.org/abs/1704.07433 respectively).
    Returns:
    - DataFrame with these metrics.
    - DataFrame with more typical training evaluation metrics, such as accuracy / loss.
    """
    confidence_ = {}
    variability_ = {}
    threshold_closeness_ = {}
    correctness_ = {}
    forgetfulness_ = {}

    # Functions to be applied to the data.
    variability_func = lambda conf: np.std(conf)
    if (
        args.include_ci
    ):  # Based on prior work on active bias (https://arxiv.org/abs/1704.07433)
        variability_func = lambda conf: np.sqrt(
            np.var(conf) + np.var(conf) * np.var(conf) / (len(conf) - 1)
        )
    threshold_closeness_func = lambda conf: conf * (1 - conf)

    num_tot_epochs = len(list(training_dynamics.values())[0]["start_logits"])
    if args.burn_out < num_tot_epochs:
        logger.info(
            f"Computing training dynamics. Burning out at {args.burn_out} of {num_tot_epochs}. "
        )
    else:
        logger.info(f"Computing training dynamics across {num_tot_epochs} epochs")
        logger.info(
            "Metrics computed: confidence, variability, correctness, forgetfulness, threshold_closeness"
        )

    start_logits = {i: [] for i in range(num_tot_epochs)}
    end_logits = {i: [] for i in range(num_tot_epochs)}
    start_targets = {i: [] for i in range(num_tot_epochs)}
    end_targets = {i: [] for i in range(num_tot_epochs)}
    training_accuracy = defaultdict(float)

    try:
        for guid in tqdm(training_dynamics):
            correctness_trend = []
            true_start_probs_trend = []
            true_end_probs_trend = []

            record = training_dynamics[guid]
            if len(record["start_logits"]) > num_tot_epochs:
                continue

            for i, (start_epoch_logits, end_epoch_logits) in enumerate(
                zip(record["start_logits"], record["end_logits"])
            ):
                start_probs = torch.nn.functional.softmax(
                    torch.Tensor(start_epoch_logits), dim=-1
                )
                # print(start_probs)
                end_probs = torch.nn.functional.softmax(
                    torch.Tensor(end_epoch_logits), dim=-1
                )
                # print(end_probs)
                true_start_prob = float(start_probs[record["start_gold"]])
                true_end_prob = float(end_probs[record["end_gold"]])
                true_start_probs_trend.append(true_start_prob)
                true_end_probs_trend.append(true_end_prob)

                start_prediction = np.argmax(start_epoch_logits)
                end_prediction = np.argmax(end_epoch_logits)
                # print(start_prediction, end_prediction)
                # print(record["start_gold"], record["end_gold"])
                is_correct = (start_prediction == record["start_gold"]).item() & (
                    end_prediction == record["end_gold"]
                ).item()
                correctness_trend.append(is_correct)
                # break

                training_accuracy[i] += is_correct
                start_logits[i].append(start_epoch_logits)
                end_logits[i].append(end_epoch_logits)
                start_targets[i].append(record["start_gold"])
                end_targets[i].append(record["end_gold"])

            if args.burn_out < num_tot_epochs:
                correctness_trend = correctness_trend[: args.burn_out]
                true_start_probs_trend = true_start_probs_trend[: args.burn_out]
                true_end_probs_trend = true_end_probs_trend[: args.burn_out]

            correctness_[guid] = compute_correctness(correctness_trend)
            confidence_[guid] = (
                np.sum(np.mean(true_start_probs_trend) + np.mean(true_end_probs_trend))
                / 2
            )
            variability_[guid] = (
                np.sum(
                    variability_func(true_start_probs_trend)
                    + variability_func(true_end_probs_trend)
                )
                / 2
            )
            forgetfulness_[guid] = compute_forgetfulness(correctness_trend)
            threshold_closeness_[guid] = threshold_closeness_func(confidence_[guid])

    except Exception as e:
        logger.info(e)
        logger.info("Error in computing training dynamics metrics")

    # print(training_accuracy)
    # print("Correctness: ", correctness_)
    # print("Conf: ", confidence_)
    # print("Var: ", variability_)
    # print("forget: ", forgetfulness_)
    # print("thresh: ", threshold_closeness_)

    # Should not affect ranking, so ignoring.
    epsilon_var = np.mean(list(variability_.values()))

    column_names = [
        "guid",
        "index",
        "threshold_closeness",
        "confidence",
        "variability",
        "correctness",
        "forgetfulness",
    ]
    df = pd.DataFrame(
        [
            [
                guid,
                i,
                threshold_closeness_[guid],
                confidence_[guid],
                variability_[guid],
                correctness_[guid],
                forgetfulness_[guid],
            ]
            for i, guid in enumerate(correctness_)
        ],
        columns=column_names,
    )

    # compute loss
    def compute_loss(idx: int):
        loss = torch.nn.CrossEntropyLoss()
        combined_loss = (
            loss(
                torch.tensor(start_logits[idx]), torch.LongTensor(start_targets[idx])
            ).item()
            + loss(
                torch.tensor(end_logits[idx]), torch.LongTensor(end_targets[idx])
            ).item()
        )
        combined_loss /= 2
        combined_loss /= len(training_dynamics)
        return combined_loss

    df_train = pd.DataFrame(
        [
            [i, compute_loss(i), training_accuracy[i] / len(training_dynamics)]
            for i in range(num_tot_epochs)
        ],
        columns=["epoch", "loss", "train_acc"],
    )
    return df, df_train


def consider_ascending_order(filtering_metric: str) -> bool:
    """
    Determine if the metric values' sorting order to get the most `valuable` examples for training.
    """
    if filtering_metric == "variability":
        return False
    elif filtering_metric == "confidence":
        return True
    elif filtering_metric == "threshold_closeness":
        return False
    elif filtering_metric == "forgetfulness":
        return False
    elif filtering_metric == "correctness":
        return True
    else:
        raise NotImplementedError(
            f"Filtering based on {filtering_metric} not implemented!"
        )


def write_filtered_data(args, train_dy_metrics):
    """
    Filter data based on the given metric, and write it in TSV format to train GLUE-style classifier.
    """
    # First save the args for filtering, to keep track of which model was used for filtering.
    argparse_dict = vars(args)
    with open(
        os.path.join(args.filtering_output_dir, f"filtering_configs.json"), "w"
    ) as outfile:
        outfile.write(json.dumps(argparse_dict, indent=4, sort_keys=True) + "\n")

    # Determine whether to sort data in ascending order or not, based on the metric.
    is_ascending = consider_ascending_order(args.metric)
    if args.worst:
        is_ascending = not is_ascending

    # Sort by selection.
    sorted_scores = train_dy_metrics.sort_values(
        by=[args.metric], ascending=is_ascending
    )

    print("Sorted scores:", sorted_scores)
    print("Train dynamics:", train_dy_metrics)

    # original_train_file = os.path.join(os.path.join(args.data_dir, args.task_name), f"train.tsv")
    # train_numeric, header = read_data(original_train_file, task_name=args.task_name, guid_as_int=True)

    dataloader = PreprocessData(
        args.dataset_name,
        args.dataset_config,
        cf_path=os.path.join(
            args.data_dir,
            "rag_counterfactuals_complete_noise_min_filtered_final_dedup_1.jsonl",
        ),
        save_data=False,
        save_path="",
    )
    train_data, _ = dataloader.processed_counterfactuals()
    for fraction in [0.01, 0.05, 0.10, 0.1667, 0.25, 0.3319, 0.50, 0.75]:
        outdir = os.path.join(
            args.filtering_output_dir,
            f"cartography_{args.metric}_{fraction:.2f}/{args.task_name}",
        )
        if not os.path.exists(outdir):
            os.makedirs(outdir)

        num_samples = int(fraction * len(train_data))
        filtered_idx = []

        selected = sorted_scores.head(n=num_samples + 1)
        # print("Selected:", selected)
        if args.both_ends:
            hardest = sorted_scores.head(n=int(num_samples * 0.7))
            easiest = sorted_scores.tail(n=num_samples - hardest.shape[0])
            selected = pd.concat([hardest, easiest])
            fm = args.metric
            logger.info(
                f"Selecting both ends: {fm} = "
                f"({hardest.head(1)[fm].values[0]:3f}: {hardest.tail(1)[fm].values[0]:3f}) "
                f"& ({easiest.head(1)[fm].values[0]:3f}: {easiest.tail(1)[fm].values[0]:3f})"
            )

        selection_iterator = tqdm(range(len(selected)))
        for idx in selection_iterator:
            selection_iterator.set_description(
                f"{args.metric} = {selected.iloc[idx][args.metric]:.4f}"
            )

            selected_id = selected.iloc[idx]["guid"]
            filtered_idx.append(selected_id)
            # print("filtered ids: ", filtered_idx)
        with jsonlines.open(os.path.join(outdir, f"train.jsonl"), "w") as writer:
            for example in tqdm(
                train_data, total=len(train_data), desc="Saving samples ... "
            ):
                # print(example["id"])
                if example["id"] in filtered_idx:
                    writer.write(
                        {
                            "id": example["id"],
                            "title": example["title"],
                            "question": example["question"],
                            "context": example["context"],
                            "answers": example["answers"],
                        }
                    )
    # logger.info(f"Wrote {num_samples} samples to {outdir}.")


def plot_data_map(
    dataframe: pd.DataFrame,
    plot_dir: os.path,
    hue_metric: str = "correct.",
    title: str = "",
    model: str = "RoBERTa",
    show_hist: bool = False,
    max_instances_to_plot=55000,
):
    # Set style.
    sns.set(style="whitegrid", font_scale=1.6, font="Georgia", context="paper")
    logger.info(f"Plotting figure for {title} using the {model} model ...")

    # Subsample data to plot, so the plot is not too busy.
    dataframe = dataframe.sample(
        n=max_instances_to_plot
        if dataframe.shape[0] > max_instances_to_plot
        else len(dataframe)
    )

    # Normalize correctness to a value between 0 and 1.
    dataframe = dataframe.assign(
        corr_frac=lambda d: d.correctness / d.correctness.max()
    )
    dataframe["correct."] = [f"{x:.1f}" for x in dataframe["corr_frac"]]

    main_metric = "variability"
    other_metric = "confidence"

    hue = hue_metric
    num_hues = len(dataframe[hue].unique().tolist())
    style = hue_metric if num_hues < 8 else None

    if not show_hist:
        fig, ax0 = plt.subplots(1, 1, figsize=(8, 6))
    else:
        fig = plt.figure(figsize=(14, 10),)
        gs = fig.add_gridspec(3, 2, width_ratios=[5, 1])
        ax0 = fig.add_subplot(gs[:, 0])

    # Make the scatterplot.
    # Choose a palette.
    pal = sns.diverging_palette(260, 15, n=num_hues, sep=10, center="dark")

    plot = sns.scatterplot(
        x=main_metric,
        y=other_metric,
        ax=ax0,
        data=dataframe,
        hue=hue,
        palette=pal,
        style=style,
        s=30,
    )

    # Annotate Regions.
    bb = lambda c: dict(boxstyle="round,pad=0.3", ec=c, lw=2, fc="white")
    func_annotate = lambda text, xyc, bbc: ax0.annotate(
        text,
        xy=xyc,
        xycoords="axes fraction",
        fontsize=15,
        color="black",
        va="center",
        ha="center",
        rotation=350,
        bbox=bb(bbc),
    )
    an1 = func_annotate("ambiguous", xyc=(0.9, 0.5), bbc="black")
    an2 = func_annotate("easy-to-learn", xyc=(0.27, 0.85), bbc="r")
    an3 = func_annotate("hard-to-learn", xyc=(0.35, 0.25), bbc="b")

    if not show_hist:
        plot.legend(ncol=1, bbox_to_anchor=[0.175, 0.5], loc="right")
    else:
        plot.legend(fancybox=True, shadow=True, ncol=1)
    plot.set_xlabel("variability")
    plot.set_ylabel("confidence")

    if show_hist:
        plot.set_title(f"{title}-{model} Data Map", fontsize=17)

        # Make the histograms.
        ax1 = fig.add_subplot(gs[0, 1])
        ax2 = fig.add_subplot(gs[1, 1])
        ax3 = fig.add_subplot(gs[2, 1])

        plott0 = dataframe.hist(column=["confidence"], ax=ax1, color="#622a87")
        plott0[0].set_title("")
        plott0[0].set_xlabel("confidence")
        plott0[0].set_ylabel("density")

        plott1 = dataframe.hist(column=["variability"], ax=ax2, color="teal")
        plott1[0].set_title("")
        plott1[0].set_xlabel("variability")
        plott1[0].set_ylabel("density")

        plot2 = sns.countplot(x="correct.", data=dataframe, ax=ax3, color="#86bf91")
        ax3.xaxis.grid(True)  # Show the vertical gridlines

        plot2.set_title("")
        plot2.set_xlabel("correctness")
        plot2.set_ylabel("density")

    fig.tight_layout()
    filename = (
        f"{plot_dir}/{title}_{model}.pdf"
        if show_hist
        else f"figures/compact_{title}_{model}.pdf"
    )
    fig.savefig(filename, dpi=300)
    logger.info(f"Plot saved to {filename}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--filter",
        action="store_true",
        help="Whether to filter data subsets based on specified `metric`.",
    )
    parser.add_argument(
        "--plot",
        action="store_true",
        help="Whether to plot data maps and save as `pdf`.",
    )
    parser.add_argument(
        "--model_dir",
        "-o",
        required=True,
        type=os.path.abspath,
        help="Directory where model training dynamics stats reside.",
    )
    parser.add_argument(
        "--data_dir",
        "-d",
        default="./src/data/squad/",
        type=os.path.abspath,
        help="Directory where data for task resides.",
    )
    parser.add_argument(
        "--plots_dir",
        default="./cartography/",
        type=os.path.abspath,
        help="Directory where plots are to be saved.",
    )
    parser.add_argument(
        "--task_name",
        "-t",
        default="SQUAD",
        choices=("SNLI", "MNLI", "QNLI", "WINOGRANDE", "SQUAD"),
        help="Which task are we plotting or filtering for.",
    )
    parser.add_argument(
        "--metric",
        choices=(
            "threshold_closeness",
            "confidence",
            "variability",
            "correctness",
            "forgetfulness",
        ),
        help="Metric to filter data by.",
    )
    parser.add_argument(
        "--include_ci",
        action="store_true",
        help="Compute the confidence interval for variability.",
    )
    parser.add_argument(
        "--filtering_output_dir",
        "-f",
        default="./roberta-squad-t5-squad-cfs-amb/",
        type=os.path.abspath,
        help="Output directory where filtered datasets are to be written.",
    )
    parser.add_argument(
        "--worst",
        action="store_true",
        help="Select from the opposite end of the spectrum acc. to metric,"
        "for baselines",
    )
    parser.add_argument(
        "--both_ends",
        action="store_true",
        help="Select from both ends of the spectrum acc. to metric,",
    )
    parser.add_argument(
        "--burn_out",
        type=int,
        default=100,
        help="# Epochs for which to compute train dynamics.",
    )
    parser.add_argument(
        "--model", default="RoBERTa", help="Model for which data map is being plotted"
    )
    parser.add_argument(
        "--dataset_name",
        default="squad",
        help="The dataset name for which the data map is being plotted",
    )
    parser.add_argument(
        "--dataset_config",
        default="plain_text",
        help="dataset config as per HF datasets",
    )

    args = parser.parse_args()

    training_dynamics = read_training_dynamics(
        args.model_dir,
        strip_last=True if args.task_name in ["QNLI"] else False,
        burn_out=args.burn_out if args.burn_out < 100 else None,
    )
    total_epochs = len(list(training_dynamics.values())[0]["start_logits"])
    print(f"Total epochs: {total_epochs}")

    if args.burn_out > total_epochs:
        args.burn_out = total_epochs
    logger.info(f"Total epochs found: {args.burn_out}")

    # compute_train_dy_metrics(training_dynamics, args)
    train_dy_metrics, df_train = compute_train_dy_metrics(training_dynamics, args)
    print(train_dy_metrics)
    print(df_train.head())

    burn_out_str = f"_{args.burn_out}" if args.burn_out > total_epochs else ""
    train_dy_filename = os.path.join(args.model_dir, f"td_metrics{burn_out_str}.jsonl")
    train_dy_metrics.to_json(train_dy_filename, orient="records", lines=True)
    logger.info(f"Metrics based on Training Dynamics written to {train_dy_filename}")

    if args.filter:
        assert args.filtering_output_dir
        if not os.path.exists(args.filtering_output_dir):
            os.makedirs(args.filtering_output_dir)
        assert args.metric
        write_filtered_data(args, train_dy_metrics)

    if args.plot:
        assert args.plots_dir
        if not os.path.exists(args.plots_dir):
            os.makedirs(args.plots_dir)
        plot_data_map(
            train_dy_metrics,
            args.plots_dir,
            title="Squad_CF",
            show_hist=True,
            model=args.model,
        )

src/cf_generation/baseline_generation/cartography/dynamics_filtering.py
UKPLab/CATfOOD
"""
Train roberta model with counterfactuals
"""

import os
import logging
import json
from typing import List, Tuple
import collections
import pandas as pd
import numpy as np

import torch
from transformers import (
    AutoTokenizer,
    AutoConfig,
    AutoModelForQuestionAnswering,
    TrainingArguments,
    Trainer,
)

import wandb
from src.cartography.dataloader import PreprocessData

logger = logging.getLogger(__name__)

# hyperparameters
LEARNING_RATE = 1.0e-5
WARMUP_RATIO = 0.06
MAX_EPOCHS = 5
FP16 = False
MODEL_NAME = "roberta-base"
BATCH_SIZE = 16

RESUME_TRAINING = None

OUTPUT_DIR = "roberta-squad-t5-squad-cfs-cartography"
CHECKPOINT = None

DO_TRAIN = True
DO_EVAL = True


# set to None when using all data
MAX_TRAIN_SAMPLES = None
MAX_VAL_SAMPLES = None

WANDB_RUN_NAME = OUTPUT_DIR

wandb.init(
    config={
        "lr": LEARNING_RATE,
        "warmup_ratio": WARMUP_RATIO,
        "epochs": MAX_EPOCHS,
        "model": MODEL_NAME,
        "batch_size": BATCH_SIZE,
        "output_dir": OUTPUT_DIR,
    }
)
wandb.run.name = WANDB_RUN_NAME
BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


class CustomTrainer(Trainer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.ep_train_ids = None
        self.ep_start_logits = None
        self.ep_end_logits = None
        self.ep_gold_start = None
        self.ep_gold_end = None

    def training_step(self, model: torch.nn.Module, inputs: dict) -> torch.Tensor:
        model.train()
        idx = inputs.pop("idx")

        inputs = {k: v.to(self.args.device) for k, v in inputs.items()}
        outputs = model(**inputs)
        loss = outputs[0]

        if self.args.n_gpu > 1:
            loss = loss.mean()

        if self.args.gradient_accumulation_steps > 1 and not self.deepspeed:
            # deepspeed handles loss scaling by gradient_accumulation_steps in its `backward`
            loss = loss / self.args.gradient_accumulation_steps

        if self.do_grad_scaling:
            self.scaler.scale(loss).backward()
        else:
            loss.backward()

        epoch = int(self.state.epoch)
        start_positions = inputs["start_positions"].detach().cpu().numpy()
        end_positions = inputs["end_positions"].detach().cpu().numpy()
        start_logits = outputs[1].detach().cpu().numpy()
        end_logits = outputs[2].detach().cpu().numpy()

        # store the logits
        if self.ep_train_ids is None:
            self.ep_train_ids = np.array(idx)
            self.ep_start_logits = start_logits
            self.ep_end_logits = end_logits
            self.ep_gold_start = start_positions
            self.ep_gold_end = end_positions
        else:
            self.ep_train_ids = np.concatenate(
                (self.ep_train_ids, np.array(idx)), axis=0
            )
            self.ep_start_logits = np.concatenate(
                (self.ep_start_logits, start_logits), axis=0
            )
            self.ep_end_logits = np.concatenate(
                (self.ep_end_logits, end_logits), axis=0
            )
            self.ep_gold_start = np.concatenate(
                (self.ep_gold_start, start_positions), axis=0
            )
            self.ep_gold_end = np.concatenate((self.ep_gold_end, end_positions), axis=0)

        # log the metrics after every epoch
        if len(self.ep_train_ids) == len(self.train_dataset):
            self.log_training_dynamics(
                self.args.output_dir,
                epoch,
                list(self.ep_train_ids),
                list(self.ep_start_logits),
                list(self.ep_end_logits),
                list(self.ep_gold_start),
                list(self.ep_gold_end),
            )
            self.ep_train_ids = None
            self.ep_start_logits = None
            self.ep_end_logits = None
            self.ep_gold_start = None
            self.ep_gold_end = None

        return loss.detach()

    @staticmethod
    def log_training_dynamics(
        output_dir: os.path,
        epoch: int,
        train_ids: List[int],
        start_logits: List[List[float]],
        end_logits: List[List[float]],
        start_golds: List[int],
        end_golds: List[int],
    ):
        """
        Save training dynamics (logits) from given epoch as records of a `.jsonl` file.
        """
        td_df = pd.DataFrame(
            {
                "guid": train_ids,
                f"start_logits_epoch_{epoch}": start_logits,
                f"end_logits_epoch_{epoch}": end_logits,
                "start_gold": start_golds,
                "end_gold": end_golds,
            }
        )

        logging_dir = os.path.join(output_dir, f"training_dynamics")
        # Create directory for logging training dynamics, if it doesn't already exist.
        if not os.path.exists(logging_dir):
            os.makedirs(logging_dir)
        epoch_file_name = os.path.join(logging_dir, f"dynamics_epoch_{epoch}.jsonl")
        td_df.to_json(epoch_file_name, lines=True, orient="records")
        logger.info(f"Training Dynamics logged to {epoch_file_name}")


class RobertaSquad:
    def __init__(
        self,
        dataset_name: str,
        dataset_config: str,
        cf_path: str,
        max_src_len: int,
        stride: int,
        do_train: bool,
        do_eval: bool,
        model_name: str,
        cache_dir: str = None,
        max_train_samples: int = None,
        max_val_samples: int = None,
        save_results: bool = False,
    ):
        self.dataset_name = dataset_name
        self.dataset_config = dataset_config
        self.cf_path = cf_path
        self.max_src_len = max_src_len
        self.stride = stride
        self.do_train = do_train
        self.do_eval = do_eval
        self.model_name = model_name
        self.cache_dir = cache_dir
        self.max_train_samples = max_train_samples
        self.max_val_samples = max_val_samples
        self.padding = "max_length"
        self.save_results = save_results
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        # load resources
        self._prepare_data()

    def _ensure_tensor_on_device(self, **inputs):
        """
        Ensure PyTorch tensors are on the specified device.
        Args:
            inputs (keyword arguments that should be :obj:`torch.Tensor`):
                The tensors to place on :obj:`self.device`.
        Return:
            :obj:`Dict[str, torch.Tensor]`: The same as :obj:`inputs` but on the proper device.
        """
        return {name: tensor.to(self.model.device) for name, tensor in inputs.items()}

    def _prepare_data(self):

        """
        Load model, tokenizer and data
        """

        config = AutoConfig.from_pretrained(self.model_name, cache_dir=self.cache_dir,)
        self.tokenizer = AutoTokenizer.from_pretrained(
            self.model_name, cache_dir=self.cache_dir, use_fast=True,
        )
        self.model = AutoModelForQuestionAnswering.from_pretrained(
            self.model_name, config=config, cache_dir=self.cache_dir,
        )
        dataloader = PreprocessData(
            self.dataset_name,
            self.dataset_config,
            cf_path=BASE_PATH + self.cf_path,
            save_data=False,
            save_path="",
        )
        if self.do_train and not self.do_eval:
            self.train_set, _ = dataloader.processed_counterfactuals()
        elif self.do_train and self.do_eval:
            self.train_set, self.val_set = dataloader.processed_counterfactuals()

    def _preprocess_training_examples(self, examples):
        """
        Preprocess training examples
        """

        inputs = self.tokenizer(
            examples["question"],
            examples["context"],
            max_length=self.max_src_len,
            truncation="only_second",
            stride=self.stride,
            return_overflowing_tokens=True,
            return_offsets_mapping=True,
            padding="max_length",
        )

        offset_mapping = inputs.pop("offset_mapping")
        sample_map = inputs.pop("overflow_to_sample_mapping")
        answers = examples["answers"]
        batch_ids = examples["id"]

        processed_batch_ids = []
        start_positions = []
        end_positions = []

        for i, offset in enumerate(offset_mapping):
            sample_idx = sample_map[i]
            processed_batch_ids.append(batch_ids[sample_idx])
            answer = answers[sample_idx]
            start_char = answer["answer_start"][0]
            end_char = answer["answer_start"][0] + len(answer["text"][0])
            sequence_ids = inputs.sequence_ids(i)

            # Find the start and end of the context
            idx = 0
            while sequence_ids[idx] != 1:
                idx += 1
            context_start = idx
            while sequence_ids[idx] == 1:
                idx += 1
            context_end = idx - 1

            # If the answer is not fully inside the context, label is (0, 0)
            if (
                offset[context_start][0] > start_char
                or offset[context_end][1] < end_char
            ):
                start_positions.append(0)
                end_positions.append(0)
            else:
                # Otherwise it's the start and end token positions
                idx = context_start
                while idx <= context_end and offset[idx][0] <= start_char:
                    idx += 1
                start_positions.append(idx - 1)

                idx = context_end
                while idx >= context_start and offset[idx][1] >= end_char:
                    idx -= 1
                end_positions.append(idx + 1)

        inputs["start_positions"] = start_positions
        inputs["end_positions"] = end_positions
        inputs["id"] = processed_batch_ids

        return inputs

    def data_collator(self, features):
        idx = [x["id"] for x in features]
        input_ids = [x["input_ids"] for x in features]
        input_ids = torch.tensor(input_ids, dtype=torch.long)
        attention_mask = [x["attention_mask"] for x in features]
        attention_mask = torch.tensor(attention_mask, dtype=torch.long)
        start_positions = [x["start_positions"] for x in features]
        start_positions = torch.tensor(start_positions, dtype=torch.long)
        end_positions = [x["end_positions"] for x in features]
        end_positions = torch.tensor(end_positions, dtype=torch.long)

        return {
            "idx": idx,
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "start_positions": start_positions,
            "end_positions": end_positions,
        }

    def train(self):
        if self.do_train:
            logger.info("*** Train ***")

            if self.max_train_samples is not None:
                # We will select sample from whole data if argument is specified
                max_train_samples = min(len(self.train_set), self.max_train_samples)
                self.train_set = self.train_set.select(range(max_train_samples))

            train_dataset = self.train_set.map(
                self._preprocess_training_examples,
                batched=True,
                remove_columns=self.train_set.column_names,
            )
            print(len(self.train_set), len(train_dataset))

        if self.do_eval:
            if self.max_val_samples is not None:
                # We will select sample from whole data if argument is specified
                max_val_samples = min(len(self.val_set), self.max_val_samples)
                self.val_set = self.val_set.select(range(max_val_samples))

            validation_dataset = self.val_set.map(
                self._preprocess_training_examples,
                batched=True,
                remove_columns=self.val_set.column_names,
            )
            print(len(self.val_set), len(validation_dataset))

        args = TrainingArguments(
            output_dir=OUTPUT_DIR,
            overwrite_output_dir=False,
            do_train=self.do_train,
            do_eval=self.do_eval,
            # evaluation_strategy="epoch",
            save_strategy="epoch",
            # eval_steps=4000,
            per_device_train_batch_size=BATCH_SIZE,
            per_device_eval_batch_size=BATCH_SIZE,
            learning_rate=LEARNING_RATE,
            warmup_ratio=WARMUP_RATIO,
            num_train_epochs=MAX_EPOCHS,
            logging_steps=10,
            save_total_limit=2,
            # load_best_model_at_end=True,
            run_name=WANDB_RUN_NAME,
            disable_tqdm=False,
            report_to=["wandb"],
            remove_unused_columns=False,
            fp16=FP16,
            # label_names=[
            #     "eval_loss",
            # ],  # it's important to log eval_loss
        )
        print("Batch Size", args.train_batch_size)
        print("Parallel Mode", args.parallel_mode)

        trainer = CustomTrainer(
            model=self.model,
            args=args,
            data_collator=self.data_collator,
            train_dataset=train_dataset if self.do_train else None,
            eval_dataset=validation_dataset if self.do_eval else None,
        )
        try:
            if self.do_train:
                checkpoint = None
                if RESUME_TRAINING is not None:
                    checkpoint = RESUME_TRAINING
                trainer.train(resume_from_checkpoint=checkpoint)
                trainer.save_model()
        except KeyboardInterrupt:
            trainer.save_model("interrupted-squad")

        wandb.finish()


if __name__ == "__main__":
    trainer = RobertaSquad(
        dataset_name="squad",
        dataset_config="plain_text",
        cf_path="src/data/squad/t5_squad_counterfactuals/rag_counterfactuals_complete_noise_min_filtered_final_dedup_1.jsonl",
        max_src_len=384,
        stride=128,
        do_train=DO_TRAIN,
        do_eval=DO_EVAL,
        model_name=MODEL_NAME,
        max_train_samples=MAX_TRAIN_SAMPLES,
        max_val_samples=MAX_VAL_SAMPLES,
        save_results=True,
    )
    trainer.train()

src/cf_generation/baseline_generation/cartography/run_train.py
UKPLab/CATfOOD
import json
import logging
import numpy as np
import os
import pandas as pd
import tqdm

from typing import List

logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", level=logging.INFO
)
logger = logging.getLogger(__name__)


def log_training_dynamics(
    output_dir: os.path,
    epoch: int,
    train_ids: List[int],
    train_logits: List[List[float]],
    train_golds: List[int],
):
    """
  Save training dynamics (logits) from given epoch as records of a `.jsonl` file.
  """
    td_df = pd.DataFrame(
        {"guid": train_ids, f"logits_epoch_{epoch}": train_logits, "gold": train_golds}
    )

    logging_dir = os.path.join(output_dir, f"training_dynamics")
    # Create directory for logging training dynamics, if it doesn't already exist.
    if not os.path.exists(logging_dir):
        os.makedirs(logging_dir)
    epoch_file_name = os.path.join(logging_dir, f"dynamics_epoch_{epoch}.jsonl")
    td_df.to_json(epoch_file_name, lines=True, orient="records")
    logger.info(f"Training Dynamics logged to {epoch_file_name}")


def read_training_dynamics(
    model_dir: os.path,
    strip_last: bool = False,
    id_field: str = "guid",
    burn_out: int = None,
):
    """
  Given path to logged training dynamics, merge stats across epochs.
  Returns:
  - Dict between ID of a train instances and its gold label, and the list of logits across epochs.
  """
    train_dynamics = {}

    td_dir = os.path.join(model_dir, "training_dynamics")
    num_epochs = len(
        [f for f in os.listdir(td_dir) if os.path.isfile(os.path.join(td_dir, f))]
    )
    if burn_out:
        num_epochs = burn_out

    logger.info(f"Reading {num_epochs} files from {td_dir} ...")
    for epoch_num in tqdm.tqdm(range(num_epochs)):
        epoch_file = os.path.join(td_dir, f"dynamics_epoch_{epoch_num}.jsonl")
        assert os.path.exists(epoch_file)

        with open(epoch_file, "r") as infile:
            for line in infile:
                record = json.loads(line.strip())
                guid = record[id_field] if not strip_last else record[id_field][:-1]

                if guid not in train_dynamics:
                    assert epoch_num == 0
                    train_dynamics[guid] = {
                        "start_gold": record["start_gold"],
                        "end_gold": record["end_gold"],
                        "start_logits": [],
                        "end_logits": [],
                    }
                train_dynamics[guid]["start_logits"].append(
                    record[f"start_logits_epoch_{epoch_num}"]
                )
                train_dynamics[guid]["end_logits"].append(
                    record[f"end_logits_epoch_{epoch_num}"]
                )

    logger.info(f"Read training dynamics for {len(train_dynamics)} train instances.")
    return train_dynamics

src/cf_generation/baseline_generation/cartography/selection_utils.py
UKPLab/CATfOOD
import jsonlines
import string
from tqdm import tqdm
import nltk

nltk.download("stopwords")
from nltk.corpus import stopwords


BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


class CounterfactualDataset:
    def __init__(self, data_path, save_path):
        self.data_path = data_path
        self.save_path = save_path

    def _load_data(self):
        """
        load datasets
        """
        examples = []
        c = 0
        with jsonlines.open(self.data_path) as reader:
            for ex in tqdm(reader):
                examples.append(ex)
                c += 1
                # if c==15000:
                #     break

        return examples

    def _clean_text(self, text):
        """
        clean dataset
        """

        def remove_punctuations(txt, punct=string.punctuation):
            return "".join([c for c in txt if c not in punct])

        def remove_stopwords(txt, sw=list(stopwords.words("english"))):
            return " ".join([w for w in txt.split() if w.lower() not in sw])

        text = text.replace("\n", " ").replace("\r", " ").replace("'", "")
        text = remove_punctuations(text)
        text = remove_stopwords(text)
        return text.lower()

    def _noise_filtering(self, examples):
        """
        Filter out noisy examples
        """
        filtered_examples = []
        for example in tqdm(examples):
            predicted_answer = self._clean_text(example["predicted_answer"]["text"])
            alternate_answers = [
                self._clean_text(ex) for ex in example["alternate_answers"]
            ]
            # check if predicted answer is in alternate answers
            if predicted_answer not in alternate_answers:
                continue
            # if the predicted answer is in alternate answers,
            # check if it is same as the majority of the alternate answers
            if predicted_answer == max(
                set(alternate_answers), key=alternate_answers.count
            ):
                filtered_examples.append(example)

        return filtered_examples

    def _minimal_filtering(self, examples):
        """
        Filter out minimal examples
        """
        filtered_examples = []
        processed_ids = []
        c = 0
        # print(examples)
        for ex in tqdm(examples):
            id = ex["id"].split("_")[0]
            # actual_answer = self._clean_text(ex["answers"]["text"][0])
            if id not in processed_ids:
                # find all examples with the same id
                examples_with_same_id = [
                    example
                    for example in examples
                    if example["id"].split("_")[0] == id
                    and self._clean_text(example["predicted_answer"]["text"])
                    != self._clean_text(example["answers"]["text"][0])
                ]
                # print("Ex same", examples_with_same_id)
                if not examples_with_same_id:
                    continue
                # get similarity scores and example ids
                similarity_scores = [
                    (example["similarity"], example["id"])
                    for example in examples_with_same_id
                ]
                # sort examples by similarity scores in ascending order
                # sorted_examples = sorted(similarity_scores, key=lambda x: x[0])
                # get last example id
                # last_example_id = sorted_examples[-1][1]
                # top_k_ids = [ex[1] for ex in sorted_examples[:3]]
                # get the index of the example with the lowest similarity score
                try:
                    min_index = min(similarity_scores, key=lambda x: x[0])[1]
                    # print(min_index)
                    # get the example with the lowest similarity score i.e. the minimally different example
                    min_example = [
                        example
                        for example in examples_with_same_id
                        if example["id"] in min_index
                    ]
                    # print(min_example)
                    filtered_examples.extend(min_example)
                    # add the id to the processed ids
                    processed_ids.append(id)
                    c += 1
                except:
                    processed_ids.append(id)
                    continue
        print(c)
        return filtered_examples

    def get_counterfactuals(self):
        """
        Get counterfactuals
        """
        examples = self._load_data()
        # counterfactuals = self._noise_filtering(examples)
        counterfactuals = self._minimal_filtering(examples)

        # save dataset
        self._save_dataset(counterfactuals)
        return counterfactuals

    def _save_dataset(self, counterfactuals):
        """
        Save dataset
        """
        with jsonlines.open(self.save_path, mode="w") as writer:
            for cf in tqdm(counterfactuals):
                writer.write(
                    {
                        "id": cf["id"],
                        "question": cf["predicted_question"],
                        "context": cf["retrieved_context"],
                        "answers": cf["predicted_answer"],
                        # "similarity": cf["similarity"]
                    },
                )


if __name__ == "__main__":
    data_path = (
        BASE_PATH
        + "src/data/squad/t5_squad_counterfactuals/rag_counterfactuals_complete_noise_filtered_final.jsonl"
    )
    save_path = (
        BASE_PATH
        + "src/data/squad/t5_squad_counterfactuals/rag_counterfactuals_complete_noise_min_filtered_final_2.jsonl"
    )
    dataset = CounterfactualDataset(data_path, save_path)
    counterfactuals = dataset.get_counterfactuals()

    # print(counterfactuals)
    # set 1: according to paper [nf + mf]
    # set 2: only minimal without checking pred == actual
    # if very less in set1 , them nf + mf w/o answer check in min.

src/cf_generation/baseline_generation/cf_dataset.py
UKPLab/CATfOOD
import logging
from argparse import ArgumentParser
from transformers import RealmConfig
from transformers.models.realm.modeling_realm import logger

from model import get_openqa_tf_finetuned, get_openqa_tf_pretrained

logger.setLevel(logging.INFO)


def get_arg_parser():
    parser = ArgumentParser()

    # ./data/enwiki-20181220/blocks.tfr
    parser.add_argument(
        "--block_records_path", type=str, required=True, help="Block records path."
    )
    # ./data/cc_news_pretrained/embedder/encoded/encoded.ckpt
    parser.add_argument(
        "--block_emb_path", type=str, required=True, help="Block embeddings path."
    )

    pretrained_group = parser.add_argument_group("pretrained conversion")

    pretrained_group.add_argument(
        "--embedder_path",
        type=str,
        default=r"./data/cc_news_pretrained/embedder/variables/variables",
        help="Pretrained embedder path.",
    )
    pretrained_group.add_argument(
        "--bert_path",
        type=str,
        default=r"./data/cc_news_pretrained/bert/variables/variables",
        help="Pretrained bert path.",
    )

    finetuned_group = parser.add_argument_group("finetuned conversion")

    finetuned_group.add_argument(
        "--checkpoint_path",
        type=str,
        default=r"./data/orqa_nq_model_from_realm/export/best_default/checkpoint/model.ckpt-300000",
        help="Finetuned checkpoint path.",
    )

    parser.add_argument(
        "--output_path",
        type=str,
        default=r"./converted_model/",
        help="Converted checkpoint path.",
    )
    parser.add_argument(
        "--from_pretrained",
        action="store_true",
        help="Whether to convert from a pretrained checkpoint or a fine-tuned checkpoint.",
    )

    return parser


def main(args):
    config = RealmConfig()

    if args.from_pretrained:
        model = get_openqa_tf_pretrained(args, config)
    else:
        model = get_openqa_tf_finetuned(args, config)

    model.save_pretrained(args.output_path)
    model.retriever.save_pretrained(args.output_path)


if __name__ == "__main__":
    parser = get_arg_parser()
    args = parser.parse_args()
    main(args)

src/cf_generation/baseline_generation/convert.py
UKPLab/CATfOOD
import jsonlines
from tqdm import tqdm

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


def remove_duplicates():
    examples = []
    with jsonlines.open(
        BASE_PATH + "src/data/squad/squad_counterfactuals_noise_min_filtered.jsonl"
    ) as reader:
        for example in tqdm(reader):
            # get all samples for a particular id
            # idx = example["id"]
            # if "turk0" not in idx:
            #     continue
            examples.append(example)
    filtered = []
    for example in tqdm(examples, total=len(examples), desc="Removing duplicates ... "):
        question = example["question"]
        # remove duplicate questions from examples
        if question not in [e["question"] for e in filtered]:
            filtered.append(example)
    save_to_disk(
        filtered,
        BASE_PATH
        + "src/data/squad/squad_counterfactuals_noise_min_filtered_final.jsonl",
    )
    # for sample in examples:
    #     original_idx = sample["id"].split("_")[0]
    #     original_question = example["question"]


def remove_duplicate_counterfactuals():
    """
    Keep only one example per id
    """
    examples = []
    with jsonlines.open(
        BASE_PATH
        + "src/data/squad/t5_squad_counterfactuals/rag_counterfactuals_complete_noise_min_filtered_final_1.jsonl"
    ) as reader:
        for example in tqdm(reader):
            examples.append(example)
    filtered = []
    for example in tqdm(
        examples, total=len(examples), desc="Removing duplicate cfs ... "
    ):
        idx = example["id"].split("_")[0]
        # remove duplicate questions from examples
        if idx not in [e["id"].split("_")[0] for e in filtered]:
            filtered.append(example)
    save_to_disk(
        filtered,
        BASE_PATH
        + "src/data/squad/t5_squad_counterfactuals/rag_counterfactuals_complete_noise_min_filtered_final_dedup_1.jsonl",
    )


def save_to_disk(data, file_name):
    with jsonlines.open(file_name, "a") as writer:
        for example in tqdm(data, total=len(data), desc="Saving samples ... "):
            writer.write(example)


if __name__ == "__main__":
    remove_duplicate_counterfactuals()

src/cf_generation/baseline_generation/duplicates.py
UKPLab/CATfOOD
from collections import Counter
import re
import string
import pandas as pd

import src.calibration.baseline.dataloader as dataloader


# BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"
BASE_PATH = "//"


def normalize_answer(s):
    """
    Lower text and remove punctuation, articles and extra whitespace.
    """

    def remove_articles(text):
        return re.sub(r"\b(a|an|the)\b", " ", text)

    def white_space_fix(text):
        return " ".join(text.split())

    def remove_punc(text):
        exclude = set(string.punctuation)
        return "".join(ch for ch in text if ch not in exclude)

    def lower(text):
        return text.lower()

    return white_space_fix(remove_articles(remove_punc(lower(s))))


def f1_score(prediction, ground_truth):
    """
    Compute f1 score of prediction
    :param prediction:
    :param ground_truth:
    :return:
    """
    prediction_tokens = normalize_answer(prediction).split()
    ground_truth_tokens = normalize_answer(ground_truth).split()
    common = Counter(prediction_tokens) & Counter(ground_truth_tokens)
    num_same = sum(common.values())
    if num_same == 0:
        return 0
    precision = 1.0 * num_same / len(prediction_tokens)
    recall = 1.0 * num_same / len(ground_truth_tokens)
    f1 = (2 * precision * recall) / (precision + recall)
    return f1


def exact_match_score(prediction, ground_truth):
    """
    Exact match for Squad evaluation
    :param prediction:
    :param ground_truth:
    :return:
    """
    # print("-------------------")
    # print(normalize_answer(prediction), normalize_answer(ground_truth))
    if normalize_answer(prediction) == normalize_answer(ground_truth):
        return 1
    else:
        return 0


def metric_max_over_ground_truths(metric_fn, prediction, ground_truths):
    scores_for_ground_truths = []
    for ground_truth in ground_truths:
        score = metric_fn(prediction, ground_truth)
        scores_for_ground_truths.append(score)
    # print(scores_for_ground_truths)
    return max(scores_for_ground_truths)


def get_score(metric, pred_text, gold_text):
    # print(pred_text, gold_text)
    # print(type(pred_text[0][0]["answer"]))
    if not isinstance(pred_text, str):
        pred_text = pred_text[0][0]["answer"]
    # print(gold_text, pred_text)
    if isinstance(gold_text, str):
        gold_text = [gold_text]
    if metric == "exact_match":
        score = metric_max_over_ground_truths(exact_match_score, pred_text, gold_text)
    elif metric == "f1":
        score = metric_max_over_ground_truths(f1_score, pred_text, gold_text)
    return score


if __name__ == "__main__":
    pred_df = (
        pd.read_json(
            BASE_PATH + "src/data/squad_adv_new/nbest_predictions_roberta_top20.json"
        )
        .T.rename_axis("id")
        .reset_index()
    )
    answers = [f"answer_{i}" for i in range(20)]
    pred_df.columns = ["id"] + answers
    squad_data = dataloader.PreprocessData(
        "squad_adversarial", "AddSent", save_data=False, save_path="../../../"
    )
    # trivia_data = Dataset.from_generator(dataloader.get_dev_examples_hf)
    # print(trivia_data)
    gold_text = [
        (sample["id"], sample["answers"]["text"])
        for sample in squad_data.processed_val_set()
    ]
    # gold_text = [(sample["id"], sample["answers"]) for sample in tqdm(trivia_data)]
    # print(gold_text)
    gold_df = pd.DataFrame(gold_text, columns=["id", "gold_answers"])
    data = pd.merge(pred_df, gold_df, on="id")
    print(data.head())
    metrics = ["exact_match", "f1"]
    for metric in metrics:
        scores = []
        for i in range(len(data)):
            scores.append(
                get_score(metric, data["answer_0"][i]["text"], data["gold_answers"][i])
            )
        print(metric, sum(scores) / len(scores) * 100)

src/cf_generation/baseline_generation/evaluate.py
UKPLab/CATfOOD
import os

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

BASE_PATH = os.getenv("PYTHONPATH", "/home/sachdeva/projects/exp_calibration/")

if __name__ == "__main__":
    model_path = BASE_PATH + "t5-base-nq-short-qa-10ep/checkpoint-107990"
    tokenizer = AutoTokenizer.from_pretrained("t5-base")
    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
    model.eval()
    sep_token = ">>"
    question = "What is the name of the stadium where the Super Bowl was held?"
    context = "The Super bowl was held at the Levi's Stadium in the New York City."
    prepared_input = f"{question} {sep_token} {context}"
    features = tokenizer(
        prepared_input,
        max_length=128,
        padding="max_length",
        truncation=True,
        return_tensors="pt",
    )
    # print(features)
    outputs = model.generate(**features, max_length=128, num_beams=2)
    dec_preds = tokenizer.decode(outputs[0], skip_special_tokens=True)
    print("The prediction is: ", dec_preds)

src/cf_generation/baseline_generation/evaluation/eval_qa.py
UKPLab/CATfOOD
import os

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

BASE_PATH = os.getenv("PYTHONPATH", "/home/sachdeva/projects/exp_calibration/")

if __name__ == "__main__":
    model_path = BASE_PATH + "t5-base-nq-short-qg-10ep/checkpoint-109150"
    tokenizer = AutoTokenizer.from_pretrained("t5-base")
    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
    model.eval()
    hl_token = "<hl>"
    context = f"The Super bowl was held at the Levi's Stadium in the {hl_token} New York {hl_token}."
    prepared_input = f"generate question: {context}"
    features = tokenizer(
        prepared_input,
        max_length=128,
        padding="max_length",
        truncation=True,
        return_tensors="pt",
    )
    # print(features)
    outputs = model.generate(**features, max_length=128, num_beams=2)
    dec_preds = tokenizer.decode(outputs[0], skip_special_tokens=True)
    print("The prediction is: ", dec_preds)

src/cf_generation/baseline_generation/evaluation/eval_qg.py
UKPLab/CATfOOD
import os
import torch
import numpy as np
from typing import List
from transformers import (
    RealmRetriever,
    RealmTokenizer,
    RealmConfig,
    AutoModelForSeq2SeqLM,
    AutoTokenizer,
)
from src.rag.modelling_realm import RealmForOpenQA

device = "cuda" if torch.cuda.is_available() else "cpu"

retriever = RealmRetriever.from_pretrained("google/realm-orqa-nq-openqa")
tokenizer = RealmTokenizer.from_pretrained("google/realm-orqa-nq-openqa")
tokenizer.add_tokens("<hl>")
# tokenizer.add_special_tokens(
#         {"additional_special_tokens": ["<hl>"]}
#     )
hl_token_id = tokenizer.convert_tokens_to_ids(["<hl>"])[0]
config = RealmConfig.from_pretrained("google/realm-orqa-nq-openqa")
# print(config)
config.reader_beam_size = 5
model = RealmForOpenQA.from_pretrained(
    "google/realm-orqa-nq-openqa", retriever=retriever, config=config
)
# model.resize_token_embeddings(len(tokenizer))
model.to(device)

question = "Where did the super bowl 50 take place?"
question_ids = tokenizer([question], return_tensors="pt").to(device)
# answer_ids = tokenizer(
#     ["alan mathison turing"],
#     add_special_tokens=False,
#     return_token_type_ids=False,
#     return_attention_mask=False,
# ).input_ids

reader_output = model(**question_ids, return_dict=True)
# print(reader_output)
print(reader_output.input_ids)
print(hl_token_id)
start_positions = reader_output.reader_output.start_pos.cpu().numpy()
end_positions = reader_output.reader_output.end_pos.cpu().numpy()
# add hl_token_id at tensor index start position and end position
highlight_input_ids = reader_output.input_ids.cpu().numpy()
inputs: List = []
for idx, (start, end) in enumerate(zip(start_positions, end_positions)):
    start = int(start)
    end = int(end)
    # get sep idx
    sep_idx = np.where(highlight_input_ids[idx, :start] == tokenizer.sep_token_id)[0][
        -1
    ]
    inputs.append(
        np.concatenate(
            [
                highlight_input_ids[idx, sep_idx:start],
                [hl_token_id],
                highlight_input_ids[idx, start : end + 1],
                [hl_token_id],
                highlight_input_ids[idx, end + 1 :],
            ],
            axis=0,
        )
    )


predicted_answer = tokenizer.batch_decode(
    reader_output.predicted_answer_ids, skip_special_tokens=True
)
contexts = tokenizer.batch_decode(
    torch.tensor(np.array(inputs), device=device),
    skip_special_tokens=True,
    clean_up_tokenization_spaces=True,
)
# loss = reader_output.loss
print("==============================")
print(contexts)
print("==============================")
# print(loss)

retrieved_blocks = np.take(
    retriever.block_records,
    indices=reader_output.retrieved_block_ids.detach().cpu().numpy(),
    axis=0,
)
print(retrieved_blocks)

# print(retriever.block_records[0])
# print(retriever.block_records[101])

# text_pair = []
# for retrieved_block in retrieved_blocks:
#     text_pair.append(retriever.block_records[retrieved_block])
#
# for text in text_pair:
#     if "alan mathison turing" in text.decode("utf-8"):
#         print('====================')
#         print("found")
#         print(text)
#         break

# context = "The Super Bowl is currently played in early February (the game originally took place in early " \
#           "to mid-January), culminating a season that generally begins in September of the previous " \
#           "calendar year. For example, Super Bowl 50, which was played on February 7, 2016, determined " \
#           "the league champion for the 2015 NFL season. The years shown below refer to the season, not " \
#           "the date that the Super Bowl was actually played. The Pittsburgh Steelers (1974, 1975, 1978, " \
#           "1979, 2005, 2008) are first with six Vince Lombardi Trophies. The Dallas Cowboys (1971, 1977, " \
#           "1992, 1993, 1995), the San Francisco 49ers (1981, 1984, 1988, 1989, 1994), and the New " \
#           "England Patriots (2001, 2003, 2004, 2014, 2016) are tied for second with 5 each. The Green " \
#           "Bay Packers (1966, 1967, 1996, 2010) and the New York Giants (1986, 1990, 2007, 2011) are " \
#           "tied for fifth with four each. The Oakland Raiders (1976, 1980, 1983), the Washington " \
#           "Redskins (1982, 1987, 1991) and the Denver Broncos (1997, 1998, 2015) are tied for " \
#           "seventh with three each."

# x = torch.tensor([[8, 2, 3],
#                   [4, 5, 6]])
#
# print(torch.max(x, dim=0).values)
# print(torch.argmax(torch.max(x, dim=0).values))

# for answer, block in zip(predicted_answer, retrieved_blocks):
#     pass
#
BASE_PATH = os.getenv("PYTHONPATH", "/home/sachdeva/projects/exp_calibration/")
model_path = BASE_PATH + "t5-large-squad-qg-seed-42"
tokenizer = AutoTokenizer.from_pretrained("t5-large")
model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
model.to(device)
model.eval()
hl_token = "<hl>"
prefix = "generate question: "

prepared_inputs = [prefix + context for context in contexts]
features = tokenizer(
    prepared_inputs,
    max_length=640,
    padding="max_length",
    truncation=True,
    return_tensors="pt",
).to(device)
# print(features)
outputs = model.generate(**features, max_length=128, num_beams=15, early_stopping=True)
dec_preds = tokenizer.batch_decode(outputs, skip_special_tokens=True)
print("The predicted questions are: ", dec_preds)
print("The predicted answers are: ", predicted_answer)

src/cf_generation/baseline_generation/hf_realm.py
UKPLab/CATfOOD
"""
Minimality filtering of questions with the counterfactuals
The closest question is the one with the lowest similarity to the original question

We select the closest question to the original question where the
answer is different.

Now this value might seem relatively arbitrary to you, it’s hard to determine if
this value reflects that the content is plagiarized or not. The larger the value
is the less likely it is to be considered plagiarized based on our understanding
of Levenshtein distance. However, it’s difficult to determine the threshold of
what distance is not large enough.
"""


from functools import lru_cache
import jsonlines
from tqdm import tqdm
import string
import Levenshtein as lev

import nltk
from nltk.corpus import stopwords


BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"
nltk.download(
    "stopwords", download_dir="/ukp-storage-1/sachdeva/miniconda3/envs/llm/nltk_data"
)


def remove_punctuations(txt, punct=string.punctuation):
    """
    This function will remove punctuations from the input text
    """
    return "".join([c for c in txt if c not in punct])


def remove_stopwords(txt, sw=list(stopwords.words("english"))):
    """
    This function will remove the stopwords from the input txt
    """
    return " ".join([w for w in txt.split() if w.lower() not in sw])


def clean_text(txt):
    """
    This function will clean the text being passed by removing specific line feed characters
    like '\n', '\r', and '\'
    """

    txt = txt.replace("\n", " ").replace("\r", " ").replace("'", "")
    txt = remove_punctuations(txt)
    txt = remove_stopwords(txt)
    return txt.lower()


def similarity(a, b):
    return lev.distance(clean_text(a), clean_text(b))


def save_to_disk(data, file_name):
    with jsonlines.open(file_name, "a") as writer:
        for example in tqdm(data, total=len(data), desc="Saving samples ... "):
            writer.write(example)


if __name__ == "__main__":
    examples = []
    with jsonlines.open(
        BASE_PATH
        + "src/data/squad/t5_squad_counterfactuals/rag_counterfactuals_complete.jsonl"
    ) as reader:
        for example in tqdm(reader):
            # get all samples for a particular id
            idx = example["id"].split("_")[0]
            original_question = example["question"]
            predicted_question = example["predicted_question"]
            example["similarity"] = similarity(original_question, predicted_question)
            examples.append(example)
    save_to_disk(
        examples,
        BASE_PATH
        + "src/data/squad/t5_squad_counterfactuals/rag_counterfactuals_complete_mf.jsonl",
    )

src/cf_generation/baseline_generation/minimality_filtering.py
UKPLab/CATfOOD
import numpy as np
import torch
from transformers import (
    RealmConfig,
    RealmReader,
    RealmRetriever,
    RealmScorer,
    RealmForOpenQA,
    RealmTokenizerFast,
    load_tf_weights_in_realm,
)
from transformers.models.realm.retrieval_realm import convert_tfrecord_to_np


def add_additional_documents(openqa, additional_documents_path):
    documents = np.load(additional_documents_path, allow_pickle=True)
    total_documents = documents.shape[0]

    retriever = openqa.retriever
    tokenizer = openqa.retriever.tokenizer

    # docs
    retriever.block_records = np.concatenate(
        (retriever.block_records, documents), axis=0
    )

    # embeds
    documents = [doc.decode() for doc in documents]
    inputs = tokenizer(documents, padding=True, truncation=True, return_tensors="pt")

    with torch.no_grad():
        projected_score = openqa.embedder(**inputs, return_dict=True).projected_score
        openqa.block_emb = torch.cat((openqa.block_emb, projected_score), dim=0)

    openqa.config.num_block_records += total_documents


def get_openqa_tf_finetuned(args, config=None):
    if config is None:
        config = RealmConfig(hidden_act="gelu_new")

    tokenizer = RealmTokenizerFast.from_pretrained(
        "google/realm-cc-news-pretrained-embedder", do_lower_case=True
    )

    block_records = convert_tfrecord_to_np(
        args.block_records_path, config.num_block_records
    )
    retriever = RealmRetriever(block_records, tokenizer)

    openqa = RealmForOpenQA(config, retriever)

    openqa = load_tf_weights_in_realm(openqa, config, args.checkpoint_path,)

    openqa = load_tf_weights_in_realm(openqa, config, args.block_emb_path,)

    openqa.eval()

    return openqa


def get_openqa_tf_pretrained(args, config=None):
    if config is None:
        config = RealmConfig(hidden_act="gelu_new")

    tokenizer = RealmTokenizerFast.from_pretrained(
        "google/realm-cc-news-pretrained-embedder", do_lower_case=True
    )

    block_records = convert_tfrecord_to_np(
        args.block_records_path, config.num_block_records
    )
    retriever = RealmRetriever(block_records, tokenizer)

    openqa = RealmForOpenQA(config, retriever)

    openqa = load_tf_weights_in_realm(openqa, config, args.bert_path,)

    openqa = load_tf_weights_in_realm(openqa, config, args.embedder_path,)

    openqa = load_tf_weights_in_realm(openqa, config, args.block_emb_path,)

    openqa.eval()

    return openqa


def get_openqa(args, config=None):
    if config is None:
        config = RealmConfig(hidden_act="gelu_new")

    retriever = RealmRetriever.from_pretrained(args.checkpoint_pretrained_name)

    openqa = RealmForOpenQA.from_pretrained(
        args.checkpoint_pretrained_name, retriever=retriever, config=config,
    )
    openqa.eval()

    return openqa


def get_scorer_reader_tokenizer_tf(args, config=None):
    if config is None:
        config = RealmConfig(hidden_act="gelu_new")
    scorer = RealmScorer(config, args.block_records_path)

    # Load retriever weights
    scorer = load_tf_weights_in_realm(scorer, config, args.retriever_path,)

    # Load block_emb weights
    scorer = load_tf_weights_in_realm(scorer, config, args.block_emb_path,)
    scorer.eval()

    reader = RealmReader.from_pretrained(
        args.checkpoint_path, config=config, from_tf=True,
    )
    reader.eval()

    tokenizer = RealmTokenizerFast.from_pretrained(
        "google/realm-cc-news-pretrained-embedder", do_lower_case=True
    )

    return scorer, reader, tokenizer


def get_scorer_reader_tokenizer_pt_pretrained(args, config=None):
    if config is None:
        config = RealmConfig(hidden_act="gelu_new")
    scorer = RealmScorer.from_pretrained(
        args.retriever_pretrained_name, args.block_records_path, config=config
    )

    # Load block_emb weights
    scorer = load_tf_weights_in_realm(scorer, config, args.block_emb_path,)
    scorer.eval()

    reader = RealmReader.from_pretrained(args.checkpoint_pretrained_name, config=config)
    reader.eval()

    tokenizer = RealmTokenizerFast.from_pretrained(
        "google/realm-cc-news-pretrained-embedder", do_lower_case=True
    )

    return scorer, reader, tokenizer


def get_scorer_reader_tokenizer_pt_finetuned(args, config=None):
    if config is None:
        config = RealmConfig(hidden_act="gelu_new")
    scorer = RealmScorer.from_pretrained(
        args.retriever_pretrained_name, args.block_records_path, config=config
    )
    scorer.eval()

    reader = RealmReader.from_pretrained(args.checkpoint_pretrained_name, config=config)
    reader.eval()

    tokenizer = RealmTokenizerFast.from_pretrained(
        "google/realm-cc-news-pretrained-embedder", do_lower_case=True
    )

    return scorer, reader, tokenizer


def get_scorer_reader_tokenizer(args, config=None):
    if config is None:
        config = RealmConfig(hidden_act="gelu_new")

    scorer = RealmScorer(config, args.block_records_path)
    reader = RealmReader(config)
    tokenizer = RealmTokenizerFast.from_pretrained(
        "google/realm-cc-news-pretrained-embedder", do_lower_case=True
    )

    return scorer, reader, tokenizer

src/cf_generation/baseline_generation/model.py
UKPLab/CATfOOD
# coding=utf-8
# Copyright 2022 The REALM authors and The HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
""" PyTorch REALM model."""

import math
import os
from dataclasses import dataclass
from typing import Optional, Tuple, Union, List

import torch
from packaging import version
from torch import nn
from torch.nn import CrossEntropyLoss

from transformers.activations import ACT2FN
from transformers.modeling_outputs import (
    BaseModelOutputWithPastAndCrossAttentions,
    BaseModelOutputWithPoolingAndCrossAttentions,
    MaskedLMOutput,
    ModelOutput,
)
from transformers.modeling_utils import PreTrainedModel
from transformers.pytorch_utils import (
    apply_chunking_to_forward,
    find_pruneable_heads_and_indices,
    prune_linear_layer,
)
from transformers.utils import (
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from transformers.models.realm.configuration_realm import RealmConfig


logger = logging.get_logger(__name__)
_EMBEDDER_CHECKPOINT_FOR_DOC = "google/realm-cc-news-pretrained-embedder"
_ENCODER_CHECKPOINT_FOR_DOC = "google/realm-cc-news-pretrained-encoder"
_SCORER_CHECKPOINT_FOR_DOC = "google/realm-cc-news-pretrained-scorer"
_CONFIG_FOR_DOC = "RealmConfig"
_TOKENIZER_FOR_DOC = "RealmTokenizer"

REALM_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "google/realm-cc-news-pretrained-embedder",
    "google/realm-cc-news-pretrained-encoder",
    "google/realm-cc-news-pretrained-scorer",
    "google/realm-cc-news-pretrained-openqa",
    "google/realm-orqa-nq-openqa",
    "google/realm-orqa-nq-reader",
    "google/realm-orqa-wq-openqa",
    "google/realm-orqa-wq-reader",
    # See all REALM models at https://huggingface.co/models?filter=realm
]


def load_tf_weights_in_realm(model, config, tf_checkpoint_path):
    """Load tf checkpoints in a pytorch model."""
    try:
        import re

        import numpy as np
        import tensorflow as tf
    except ImportError:
        logger.error(
            "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "
            "https://www.tensorflow.org/install/ for installation instructions."
        )
        raise
    tf_path = os.path.abspath(tf_checkpoint_path)
    logger.info(f"Converting TensorFlow checkpoint from {tf_path}")
    # Load weights from TF model
    init_vars = tf.train.list_variables(tf_path)
    names = []
    arrays = []

    for name, shape in init_vars:
        logger.info(f"Loading TF weight {name} with shape {shape}")
        array = tf.train.load_variable(tf_path, name)
        names.append(name)
        arrays.append(array)

    for name, array in zip(names, arrays):
        if isinstance(model, RealmReader) and "reader" not in name:
            logger.info(
                f"Skipping {name} as it is not {model.__class__.__name__}'s parameter"
            )
            continue

        # For pretrained openqa reader
        if (name.startswith("bert") or name.startswith("cls")) and isinstance(
            model, RealmForOpenQA
        ):
            name = name.replace("bert/", "reader/realm/")
            name = name.replace("cls/", "reader/cls/")

        # For pretrained encoder
        if (name.startswith("bert") or name.startswith("cls")) and isinstance(
            model, RealmKnowledgeAugEncoder
        ):
            name = name.replace("bert/", "realm/")

        # For finetuned reader
        if name.startswith("reader"):
            reader_prefix = "" if isinstance(model, RealmReader) else "reader/"
            name = name.replace("reader/module/bert/", f"{reader_prefix}realm/")
            name = name.replace("reader/module/cls/", f"{reader_prefix}cls/")
            name = name.replace(
                "reader/dense/", f"{reader_prefix}qa_outputs/dense_intermediate/"
            )
            name = name.replace(
                "reader/dense_1/", f"{reader_prefix}qa_outputs/dense_output/"
            )
            name = name.replace(
                "reader/layer_normalization",
                f"{reader_prefix}qa_outputs/layer_normalization",
            )

        # For embedder and scorer
        if name.startswith("module/module/module/"):  # finetuned
            embedder_prefix = "" if isinstance(model, RealmEmbedder) else "embedder/"
            name = name.replace(
                "module/module/module/module/bert/", f"{embedder_prefix}realm/"
            )
            name = name.replace(
                "module/module/module/LayerNorm/", f"{embedder_prefix}cls/LayerNorm/"
            )
            name = name.replace(
                "module/module/module/dense/", f"{embedder_prefix}cls/dense/"
            )
            name = name.replace(
                "module/module/module/module/cls/predictions/",
                f"{embedder_prefix}cls/predictions/",
            )
            name = name.replace(
                "module/module/module/bert/", f"{embedder_prefix}realm/"
            )
            name = name.replace(
                "module/module/module/cls/predictions/",
                f"{embedder_prefix}cls/predictions/",
            )
        elif name.startswith("module/module/"):  # pretrained
            embedder_prefix = "" if isinstance(model, RealmEmbedder) else "embedder/"
            name = name.replace(
                "module/module/LayerNorm/", f"{embedder_prefix}cls/LayerNorm/"
            )
            name = name.replace("module/module/dense/", f"{embedder_prefix}cls/dense/")

        name = name.split("/")
        # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v
        # which are not required for using pretrained model
        if any(
            n
            in [
                "adam_v",
                "adam_m",
                "AdamWeightDecayOptimizer",
                "AdamWeightDecayOptimizer_1",
                "global_step",
            ]
            for n in name
        ):
            logger.info(f"Skipping {'/'.join(name)}")
            continue
        pointer = model
        for m_name in name:
            if re.fullmatch(r"[A-Za-z]+_\d+", m_name):
                scope_names = re.split(r"_(\d+)", m_name)
            else:
                scope_names = [m_name]
            if scope_names[0] == "kernel" or scope_names[0] == "gamma":
                pointer = getattr(pointer, "weight")
            elif scope_names[0] == "output_bias" or scope_names[0] == "beta":
                pointer = getattr(pointer, "bias")
            else:
                try:
                    pointer = getattr(pointer, scope_names[0])
                except AttributeError:
                    logger.info(f"Skipping {'/'.join(name)}")
                    continue
            if len(scope_names) >= 2:
                num = int(scope_names[1])
                pointer = pointer[num]
        if m_name[-11:] == "_embeddings":
            pointer = getattr(pointer, "weight")
        elif m_name == "kernel":
            array = np.transpose(array)
        try:
            assert (
                pointer.shape == array.shape
            ), f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched"
        except AssertionError as e:
            e.args += (pointer.shape, array.shape)
            raise
        logger.info(f"Initialize PyTorch weight {name}")
        pointer.data = torch.from_numpy(array)
    return model


# Copied from transformers.models.bert.modeling_bert.BertEmbeddings with Bert->Realm
class RealmEmbeddings(nn.Module):
    """Construct the embeddings from word, position and token_type embeddings."""

    def __init__(self, config):
        super().__init__()
        self.word_embeddings = nn.Embedding(
            config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id
        )
        self.position_embeddings = nn.Embedding(
            config.max_position_embeddings, config.hidden_size
        )
        self.token_type_embeddings = nn.Embedding(
            config.type_vocab_size, config.hidden_size
        )

        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
        # any TensorFlow checkpoint file
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
        self.position_embedding_type = getattr(
            config, "position_embedding_type", "absolute"
        )
        self.register_buffer(
            "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1))
        )
        if version.parse(torch.__version__) > version.parse("1.6.0"):
            self.register_buffer(
                "token_type_ids",
                torch.zeros(self.position_ids.size(), dtype=torch.long),
                persistent=False,
            )

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        token_type_ids: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        past_key_values_length: int = 0,
    ) -> torch.Tensor:
        if input_ids is not None:
            input_shape = input_ids.size()
        else:
            input_shape = inputs_embeds.size()[:-1]

        seq_length = input_shape[1]

        if position_ids is None:
            position_ids = self.position_ids[
                :, past_key_values_length : seq_length + past_key_values_length
            ]

        # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs
        # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves
        # issue #5664
        if token_type_ids is None:
            if hasattr(self, "token_type_ids"):
                buffered_token_type_ids = self.token_type_ids[:, :seq_length]
                buffered_token_type_ids_expanded = buffered_token_type_ids.expand(
                    input_shape[0], seq_length
                )
                token_type_ids = buffered_token_type_ids_expanded
            else:
                token_type_ids = torch.zeros(
                    input_shape, dtype=torch.long, device=self.position_ids.device
                )

        if inputs_embeds is None:
            inputs_embeds = self.word_embeddings(input_ids)
        token_type_embeddings = self.token_type_embeddings(token_type_ids)

        embeddings = inputs_embeds + token_type_embeddings
        if self.position_embedding_type == "absolute":
            position_embeddings = self.position_embeddings(position_ids)
            embeddings += position_embeddings
        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)
        return embeddings


# Copied from transformers.models.bert.modeling_bert.BertSelfAttention with Bert->Realm
class RealmSelfAttention(nn.Module):
    def __init__(self, config, position_embedding_type=None):
        super().__init__()
        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(
            config, "embedding_size"
        ):
            raise ValueError(
                f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
                f"heads ({config.num_attention_heads})"
            )

        self.num_attention_heads = config.num_attention_heads
        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
        self.all_head_size = self.num_attention_heads * self.attention_head_size

        self.query = nn.Linear(config.hidden_size, self.all_head_size)
        self.key = nn.Linear(config.hidden_size, self.all_head_size)
        self.value = nn.Linear(config.hidden_size, self.all_head_size)

        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
        self.position_embedding_type = position_embedding_type or getattr(
            config, "position_embedding_type", "absolute"
        )
        if (
            self.position_embedding_type == "relative_key"
            or self.position_embedding_type == "relative_key_query"
        ):
            self.max_position_embeddings = config.max_position_embeddings
            self.distance_embedding = nn.Embedding(
                2 * config.max_position_embeddings - 1, self.attention_head_size
            )

        self.is_decoder = config.is_decoder

    def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
        new_x_shape = x.size()[:-1] + (
            self.num_attention_heads,
            self.attention_head_size,
        )
        x = x.view(new_x_shape)
        return x.permute(0, 2, 1, 3)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[torch.Tensor]:
        mixed_query_layer = self.query(hidden_states)

        # If this is instantiated as a cross-attention module, the keys
        # and values come from an encoder; the attention mask needs to be
        # such that the encoder's padding tokens are not attended to.
        is_cross_attention = encoder_hidden_states is not None

        if is_cross_attention and past_key_value is not None:
            # reuse k,v, cross_attentions
            key_layer = past_key_value[0]
            value_layer = past_key_value[1]
            attention_mask = encoder_attention_mask
        elif is_cross_attention:
            key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
            value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
            attention_mask = encoder_attention_mask
        elif past_key_value is not None:
            key_layer = self.transpose_for_scores(self.key(hidden_states))
            value_layer = self.transpose_for_scores(self.value(hidden_states))
            key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
            value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
        else:
            key_layer = self.transpose_for_scores(self.key(hidden_states))
            value_layer = self.transpose_for_scores(self.value(hidden_states))

        query_layer = self.transpose_for_scores(mixed_query_layer)

        if self.is_decoder:
            # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
            # Further calls to cross_attention layer can then reuse all cross-attention
            # key/value_states (first "if" case)
            # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
            # all previous decoder key/value_states. Further calls to uni-directional self-attention
            # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
            # if encoder bi-directional self-attention `past_key_value` is always `None`
            past_key_value = (key_layer, value_layer)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))

        if (
            self.position_embedding_type == "relative_key"
            or self.position_embedding_type == "relative_key_query"
        ):
            seq_length = hidden_states.size()[1]
            position_ids_l = torch.arange(
                seq_length, dtype=torch.long, device=hidden_states.device
            ).view(-1, 1)
            position_ids_r = torch.arange(
                seq_length, dtype=torch.long, device=hidden_states.device
            ).view(1, -1)
            distance = position_ids_l - position_ids_r
            positional_embedding = self.distance_embedding(
                distance + self.max_position_embeddings - 1
            )
            positional_embedding = positional_embedding.to(
                dtype=query_layer.dtype
            )  # fp16 compatibility

            if self.position_embedding_type == "relative_key":
                relative_position_scores = torch.einsum(
                    "bhld,lrd->bhlr", query_layer, positional_embedding
                )
                attention_scores = attention_scores + relative_position_scores
            elif self.position_embedding_type == "relative_key_query":
                relative_position_scores_query = torch.einsum(
                    "bhld,lrd->bhlr", query_layer, positional_embedding
                )
                relative_position_scores_key = torch.einsum(
                    "bhrd,lrd->bhlr", key_layer, positional_embedding
                )
                attention_scores = (
                    attention_scores
                    + relative_position_scores_query
                    + relative_position_scores_key
                )

        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
        if attention_mask is not None:
            # Apply the attention mask is (precomputed for all layers in RealmModel forward() function)
            attention_scores = attention_scores + attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = nn.functional.softmax(attention_scores, dim=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        context_layer = torch.matmul(attention_probs, value_layer)

        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(new_context_layer_shape)

        outputs = (
            (context_layer, attention_probs) if output_attentions else (context_layer,)
        )

        if self.is_decoder:
            outputs = outputs + (past_key_value,)
        return outputs


# Copied from transformers.models.bert.modeling_bert.BertSelfOutput with Bert->Realm
class RealmSelfOutput(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(
        self, hidden_states: torch.Tensor, input_tensor: torch.Tensor
    ) -> torch.Tensor:
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        return hidden_states


# Copied from transformers.models.bert.modeling_bert.BertAttention with Bert->Realm
class RealmAttention(nn.Module):
    def __init__(self, config, position_embedding_type=None):
        super().__init__()
        self.self = RealmSelfAttention(
            config, position_embedding_type=position_embedding_type
        )
        self.output = RealmSelfOutput(config)
        self.pruned_heads = set()

    def prune_heads(self, heads):
        if len(heads) == 0:
            return
        heads, index = find_pruneable_heads_and_indices(
            heads,
            self.self.num_attention_heads,
            self.self.attention_head_size,
            self.pruned_heads,
        )

        # Prune linear layers
        self.self.query = prune_linear_layer(self.self.query, index)
        self.self.key = prune_linear_layer(self.self.key, index)
        self.self.value = prune_linear_layer(self.self.value, index)
        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)

        # Update hyper params and store pruned heads
        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
        self.self.all_head_size = (
            self.self.attention_head_size * self.self.num_attention_heads
        )
        self.pruned_heads = self.pruned_heads.union(heads)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[torch.Tensor]:
        self_outputs = self.self(
            hidden_states,
            attention_mask,
            head_mask,
            encoder_hidden_states,
            encoder_attention_mask,
            past_key_value,
            output_attentions,
        )
        attention_output = self.output(self_outputs[0], hidden_states)
        outputs = (attention_output,) + self_outputs[
            1:
        ]  # add attentions if we output them
        return outputs


# Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert->Realm
class RealmIntermediate(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
        if isinstance(config.hidden_act, str):
            self.intermediate_act_fn = ACT2FN[config.hidden_act]
        else:
            self.intermediate_act_fn = config.hidden_act

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.dense(hidden_states)
        hidden_states = self.intermediate_act_fn(hidden_states)
        return hidden_states


# Copied from transformers.models.bert.modeling_bert.BertOutput with Bert->Realm
class RealmOutput(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(
        self, hidden_states: torch.Tensor, input_tensor: torch.Tensor
    ) -> torch.Tensor:
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        return hidden_states


# Copied from transformers.models.bert.modeling_bert.BertLayer with Bert->Realm
class RealmLayer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.chunk_size_feed_forward = config.chunk_size_feed_forward
        self.seq_len_dim = 1
        self.attention = RealmAttention(config)
        self.is_decoder = config.is_decoder
        self.add_cross_attention = config.add_cross_attention
        if self.add_cross_attention:
            if not self.is_decoder:
                raise ValueError(
                    f"{self} should be used as a decoder model if cross attention is added"
                )
            self.crossattention = RealmAttention(
                config, position_embedding_type="absolute"
            )
        self.intermediate = RealmIntermediate(config)
        self.output = RealmOutput(config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[torch.Tensor]:
        # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
        self_attn_past_key_value = (
            past_key_value[:2] if past_key_value is not None else None
        )
        self_attention_outputs = self.attention(
            hidden_states,
            attention_mask,
            head_mask,
            output_attentions=output_attentions,
            past_key_value=self_attn_past_key_value,
        )
        attention_output = self_attention_outputs[0]

        # if decoder, the last output is tuple of self-attn cache
        if self.is_decoder:
            outputs = self_attention_outputs[1:-1]
            present_key_value = self_attention_outputs[-1]
        else:
            outputs = self_attention_outputs[
                1:
            ]  # add self attentions if we output attention weights

        cross_attn_present_key_value = None
        if self.is_decoder and encoder_hidden_states is not None:
            if not hasattr(self, "crossattention"):
                raise ValueError(
                    f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers"
                    " by setting `config.add_cross_attention=True`"
                )

            # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple
            cross_attn_past_key_value = (
                past_key_value[-2:] if past_key_value is not None else None
            )
            cross_attention_outputs = self.crossattention(
                attention_output,
                attention_mask,
                head_mask,
                encoder_hidden_states,
                encoder_attention_mask,
                cross_attn_past_key_value,
                output_attentions,
            )
            attention_output = cross_attention_outputs[0]
            outputs = (
                outputs + cross_attention_outputs[1:-1]
            )  # add cross attentions if we output attention weights

            # add cross-attn cache to positions 3,4 of present_key_value tuple
            cross_attn_present_key_value = cross_attention_outputs[-1]
            present_key_value = present_key_value + cross_attn_present_key_value

        layer_output = apply_chunking_to_forward(
            self.feed_forward_chunk,
            self.chunk_size_feed_forward,
            self.seq_len_dim,
            attention_output,
        )
        outputs = (layer_output,) + outputs

        # if decoder, return the attn key/values as the last output
        if self.is_decoder:
            outputs = outputs + (present_key_value,)

        return outputs

    def feed_forward_chunk(self, attention_output):
        intermediate_output = self.intermediate(attention_output)
        layer_output = self.output(intermediate_output, attention_output)
        return layer_output


# Copied from transformers.models.bert.modeling_bert.BertEncoder with Bert->Realm
class RealmEncoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.layer = nn.ModuleList(
            [RealmLayer(config) for _ in range(config.num_hidden_layers)]
        )
        self.gradient_checkpointing = False

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = False,
        output_hidden_states: Optional[bool] = False,
        return_dict: Optional[bool] = True,
    ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]:
        all_hidden_states = () if output_hidden_states else None
        all_self_attentions = () if output_attentions else None
        all_cross_attentions = (
            () if output_attentions and self.config.add_cross_attention else None
        )

        next_decoder_cache = () if use_cache else None
        for i, layer_module in enumerate(self.layer):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            layer_head_mask = head_mask[i] if head_mask is not None else None
            past_key_value = past_key_values[i] if past_key_values is not None else None

            if self.gradient_checkpointing and self.training:

                if use_cache:
                    logger.warning(
                        "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
                    )
                    use_cache = False

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        return module(*inputs, past_key_value, output_attentions)

                    return custom_forward

                layer_outputs = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(layer_module),
                    hidden_states,
                    attention_mask,
                    layer_head_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                )
            else:
                layer_outputs = layer_module(
                    hidden_states,
                    attention_mask,
                    layer_head_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    past_key_value,
                    output_attentions,
                )

            hidden_states = layer_outputs[0]
            if use_cache:
                next_decoder_cache += (layer_outputs[-1],)
            if output_attentions:
                all_self_attentions = all_self_attentions + (layer_outputs[1],)
                if self.config.add_cross_attention:
                    all_cross_attentions = all_cross_attentions + (layer_outputs[2],)

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return tuple(
                v
                for v in [
                    hidden_states,
                    next_decoder_cache,
                    all_hidden_states,
                    all_self_attentions,
                    all_cross_attentions,
                ]
                if v is not None
            )
        return BaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,
            past_key_values=next_decoder_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
            cross_attentions=all_cross_attentions,
        )


# Copied from transformers.models.bert.modeling_bert.BertPooler with Bert->Realm
class RealmPooler(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.activation = nn.Tanh()

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        # We "pool" the model by simply taking the hidden state corresponding
        # to the first token.
        first_token_tensor = hidden_states[:, 0]
        pooled_output = self.dense(first_token_tensor)
        pooled_output = self.activation(pooled_output)
        return pooled_output


@dataclass
class RealmEmbedderOutput(ModelOutput):
    """
    Outputs of [`RealmEmbedder`] models.

    Args:
        projected_score (`torch.FloatTensor` of shape `(batch_size, config.retriever_proj_size)`):

            Projected score.
        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
            Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
            shape `(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
            sequence_length)`.

            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
            heads.
    """

    projected_score: torch.FloatTensor = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None


@dataclass
class RealmScorerOutput(ModelOutput):
    """
    Outputs of [`RealmScorer`] models.

    Args:
        relevance_score (`torch.FloatTensor` of shape `(batch_size, config.num_candidates)`):
            The relevance score of document candidates (before softmax).
        query_score (`torch.FloatTensor` of shape `(batch_size, config.retriever_proj_size)`):
            Query score derived from the query embedder.
        candidate_score (`torch.FloatTensor` of shape `(batch_size, config.num_candidates, config.retriever_proj_size)`):
            Candidate score derived from the embedder.
    """

    relevance_score: torch.FloatTensor = None
    query_score: torch.FloatTensor = None
    candidate_score: torch.FloatTensor = None


@dataclass
class RealmReaderOutput(ModelOutput):
    """
    Outputs of [`RealmReader`] models.

    Args:
        loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `start_positions`, `end_positions`, `has_answers` are provided):
            Total loss.
        retriever_loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `start_positions`, `end_positions`, `has_answers` are provided):
            Retriever loss.
        reader_loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `start_positions`, `end_positions`, `has_answers` are provided):
            Reader loss.
        retriever_correct (`torch.BoolTensor` of shape `(config.searcher_beam_size,)`, *optional*):
            Whether or not an evidence block contains answer.
        reader_correct (`torch.BoolTensor` of shape `(config.reader_beam_size, num_candidates)`, *optional*):
            Whether or not a span candidate contains answer.
        block_idx (`torch.LongTensor` of shape `()`):
            The index of the retrieved evidence block in which the predicted answer is most likely.
        candidate (`torch.LongTensor` of shape `()`):
            The index of the retrieved span candidates in which the predicted answer is most likely.
        start_pos (`torch.IntTensor` of shape `()`):
            Predicted answer starting position in *RealmReader*'s inputs.
        end_pos: (`torch.IntTensor` of shape `()`):
            Predicted answer ending position in *RealmReader*'s inputs.
        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
            Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
            shape `(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
            sequence_length)`.

            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
            heads.
    """

    loss: torch.FloatTensor = None
    retriever_loss: torch.FloatTensor = None
    reader_loss: torch.FloatTensor = None
    retriever_correct: torch.BoolTensor = None
    reader_correct: torch.BoolTensor = None
    block_idx: torch.LongTensor = None
    candidate: torch.LongTensor = None
    start_pos: torch.int32 = None
    end_pos: torch.int32 = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None


@dataclass
class RealmForOpenQAOutput(ModelOutput):
    """

    Outputs of [`RealmForOpenQA`] models.

    Args:
        reader_output (`dict`):
            Reader output.
        predicted_answer_ids (`torch.LongTensor` of shape `(answer_sequence_length)`):
            Predicted answer ids.
    """

    reader_output: dict = None
    retrieved_block_ids: torch.LongTensor = None
    input_ids: torch.LongTensor = None
    predicted_answer_ids: torch.LongTensor = None


class RealmPredictionHeadTransform(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        if isinstance(config.hidden_act, str):
            self.transform_act_fn = ACT2FN[config.hidden_act]
        else:
            self.transform_act_fn = config.hidden_act
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)

    def forward(self, hidden_states):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.transform_act_fn(hidden_states)
        hidden_states = self.LayerNorm(hidden_states)
        return hidden_states


class RealmLMPredictionHead(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.transform = RealmPredictionHeadTransform(config)

        # The output weights are the same as the input embeddings, but there is
        # an output-only bias for each token.
        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        self.bias = nn.Parameter(torch.zeros(config.vocab_size))

        # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
        self.decoder.bias = self.bias

    def forward(self, hidden_states):
        hidden_states = self.transform(hidden_states)
        hidden_states = self.decoder(hidden_states)
        return hidden_states


class RealmOnlyMLMHead(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.predictions = RealmLMPredictionHead(config)

    def forward(self, sequence_output):
        prediction_scores = self.predictions(sequence_output)
        return prediction_scores


class RealmScorerProjection(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.predictions = RealmLMPredictionHead(config)
        self.dense = nn.Linear(config.hidden_size, config.retriever_proj_size)
        self.LayerNorm = nn.LayerNorm(
            config.retriever_proj_size, eps=config.layer_norm_eps
        )

    def forward(self, hidden_states):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.LayerNorm(hidden_states)
        return hidden_states


class RealmReaderProjection(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.dense_intermediate = nn.Linear(
            config.hidden_size, config.span_hidden_size * 2
        )
        self.dense_output = nn.Linear(config.span_hidden_size, 1)
        self.layer_normalization = nn.LayerNorm(
            config.span_hidden_size, eps=config.reader_layer_norm_eps
        )
        self.relu = nn.ReLU()

    def forward(self, hidden_states, block_mask):
        def span_candidates(masks):
            """
            Generate span candidates.

            Args:
                masks: <bool> [num_retrievals, max_sequence_len]

            Returns:
                starts: <int32> [num_spans] ends: <int32> [num_spans] span_masks: <int32> [num_retrievals, num_spans]
                whether spans locate in evidence block.
            """
            _, max_sequence_len = masks.shape

            def _spans_given_width(width):
                current_starts = torch.arange(
                    max_sequence_len - width + 1, device=masks.device
                )
                current_ends = torch.arange(
                    width - 1, max_sequence_len, device=masks.device
                )
                return current_starts, current_ends

            starts, ends = zip(
                *(_spans_given_width(w + 1) for w in range(self.config.max_span_width))
            )

            # [num_spans]
            starts = torch.cat(starts, 0)
            ends = torch.cat(ends, 0)

            # [num_retrievals, num_spans]
            start_masks = torch.index_select(masks, dim=-1, index=starts)
            end_masks = torch.index_select(masks, dim=-1, index=ends)
            span_masks = start_masks * end_masks

            return starts, ends, span_masks

        def mask_to_score(mask):
            return (1.0 - mask.type(torch.float32)) * -10000.0

        # [reader_beam_size, max_sequence_len, span_hidden_size * 2]
        hidden_states = self.dense_intermediate(hidden_states)
        # [reader_beam_size, max_sequence_len, span_hidden_size]
        start_projection, end_projection = hidden_states.chunk(2, dim=-1)

        candidate_starts, candidate_ends, candidate_mask = span_candidates(block_mask)

        candidate_start_projections = torch.index_select(
            start_projection, dim=1, index=candidate_starts
        )
        candidate_end_projections = torch.index_select(
            end_projection, dim=1, index=candidate_ends
        )
        candidate_hidden = candidate_start_projections + candidate_end_projections

        # [reader_beam_size, num_candidates, span_hidden_size]
        candidate_hidden = self.relu(candidate_hidden)
        # [reader_beam_size, num_candidates, span_hidden_size]
        candidate_hidden = self.layer_normalization(candidate_hidden)
        # [reader_beam_size, num_candidates]
        reader_logits = self.dense_output(candidate_hidden).squeeze(-1)
        # [reader_beam_size, num_candidates]
        reader_logits += mask_to_score(candidate_mask)

        return reader_logits, candidate_starts, candidate_ends


REALM_START_DOCSTRING = r"""
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`RealmConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""

REALM_INPUTS_DOCSTRING = r"""
    Args:
        input_ids (`torch.LongTensor` of shape `({0})`):
            Indices of input sequence tokens in the vocabulary.

            Indices can be obtained using [`RealmTokenizer`]. See [`PreTrainedTokenizer.encode`] and
            [`PreTrainedTokenizer.__call__`] for details.

            [What are input IDs?](../glossary#input-ids)
        attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:

            - 1 for tokens that are **not masked**,
            - 0 for tokens that are **masked**.

            [What are attention masks?](../glossary#attention-mask)
        token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*):
            Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
            1]`:

            - 0 corresponds to a *sentence A* token,
            - 1 corresponds to a *sentence B* token.

            [What are token type IDs?](../glossary#token-type-ids)
        position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
            config.max_position_embeddings - 1]`.

            [What are position IDs?](../glossary#position-ids)
        head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
            Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:

            - 1 indicates the head is **not masked**,
            - 0 indicates the head is **masked**.

        inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):
            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
            is useful if you want more control over how to convert *input_ids* indices into associated vectors than the
            model's internal embedding lookup matrix.
        output_attentions (`bool`, *optional*):
            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
            tensors for more detail.
        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
            more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
"""


class RealmPreTrainedModel(PreTrainedModel):
    """
    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
    models.
    """

    config_class = RealmConfig
    load_tf_weights = load_tf_weights_in_realm
    base_model_prefix = "realm"
    _keys_to_ignore_on_load_missing = [r"position_ids"]

    def _init_weights(self, module):
        """Initialize the weights"""
        if isinstance(module, nn.Linear):
            # Slightly different from the TF version which uses truncated_normal for initialization
            # cf https://github.com/pytorch/pytorch/pull/5617
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)

    def _flatten_inputs(self, *inputs):
        """Flatten inputs' shape to (-1, input_shape[-1])"""
        flattened_inputs = []
        for tensor in inputs:
            if tensor is None:
                flattened_inputs.append(None)
            else:
                input_shape = tensor.shape
                if len(input_shape) > 2:
                    tensor = tensor.view((-1, input_shape[-1]))
                flattened_inputs.append(tensor)
        return flattened_inputs


class RealmBertModel(RealmPreTrainedModel):
    """
    Same as the original BertModel but remove docstrings.
    """

    def __init__(self, config, add_pooling_layer=True):
        super().__init__(config)
        self.config = config

        self.embeddings = RealmEmbeddings(config)
        self.encoder = RealmEncoder(config)

        self.pooler = RealmPooler(config) if add_pooling_layer else None

        # Weights initialization is mostly managed by other Realm models,
        # but we also have them initialized here to keep a consistency.
        self.post_init()

    def get_input_embeddings(self):
        return self.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.embeddings.word_embeddings = value

    def _prune_heads(self, heads_to_prune):
        """
        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
        class PreTrainedModel
        """
        for layer, heads in heads_to_prune.items():
            self.encoder.layer[layer].attention.prune_heads(heads)

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        past_key_values=None,
        use_cache=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        output_attentions = (
            output_attentions
            if output_attentions is not None
            else self.config.output_attentions
        )
        output_hidden_states = (
            output_hidden_states
            if output_hidden_states is not None
            else self.config.output_hidden_states
        )
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        if self.config.is_decoder:
            use_cache = use_cache if use_cache is not None else self.config.use_cache
        else:
            use_cache = False

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError(
                "You cannot specify both input_ids and inputs_embeds at the same time"
            )
        elif input_ids is not None:
            input_shape = input_ids.size()
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        batch_size, seq_length = input_shape
        device = input_ids.device if input_ids is not None else inputs_embeds.device

        # past_key_values_length
        past_key_values_length = (
            past_key_values[0][0].shape[2] if past_key_values is not None else 0
        )

        if attention_mask is None:
            attention_mask = torch.ones(
                ((batch_size, seq_length + past_key_values_length)), device=device
            )

        if token_type_ids is None:
            if hasattr(self.embeddings, "token_type_ids"):
                buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length]
                buffered_token_type_ids_expanded = buffered_token_type_ids.expand(
                    batch_size, seq_length
                )
                token_type_ids = buffered_token_type_ids_expanded
            else:
                token_type_ids = torch.zeros(
                    input_shape, dtype=torch.long, device=device
                )

        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
        # ourselves in which case we just need to make it broadcastable to all heads.
        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(
            attention_mask, input_shape
        )

        # If a 2D or 3D attention mask is provided for the cross-attention
        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
        if self.config.is_decoder and encoder_hidden_states is not None:
            (
                encoder_batch_size,
                encoder_sequence_length,
                _,
            ) = encoder_hidden_states.size()
            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
            if encoder_attention_mask is None:
                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
            encoder_extended_attention_mask = self.invert_attention_mask(
                encoder_attention_mask
            )
        else:
            encoder_extended_attention_mask = None

        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x n_heads x N x N
        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)

        embedding_output = self.embeddings(
            input_ids=input_ids,
            position_ids=position_ids,
            token_type_ids=token_type_ids,
            inputs_embeds=inputs_embeds,
            past_key_values_length=past_key_values_length,
        )
        encoder_outputs = self.encoder(
            embedding_output,
            attention_mask=extended_attention_mask,
            head_mask=head_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_extended_attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        sequence_output = encoder_outputs[0]
        pooled_output = (
            self.pooler(sequence_output) if self.pooler is not None else None
        )

        if not return_dict:
            return (sequence_output, pooled_output) + encoder_outputs[1:]

        return BaseModelOutputWithPoolingAndCrossAttentions(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
            past_key_values=encoder_outputs.past_key_values,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
            cross_attentions=encoder_outputs.cross_attentions,
        )


@add_start_docstrings(
    "The embedder of REALM outputting projected score that will be used to calculate relevance score.",
    REALM_START_DOCSTRING,
)
class RealmEmbedder(RealmPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)

        self.realm = RealmBertModel(self.config)
        self.cls = RealmScorerProjection(self.config)
        self.post_init()

    def get_input_embeddings(self):
        return self.realm.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.realm.embeddings.word_embeddings = value

    @add_start_docstrings_to_model_forward(
        REALM_INPUTS_DOCSTRING.format("batch_size, sequence_length")
    )
    @replace_return_docstrings(
        output_type=RealmEmbedderOutput, config_class=_CONFIG_FOR_DOC
    )
    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        r"""
        Returns:

        Example:

        ```python
        >>> from transformers import RealmTokenizer, RealmEmbedder
        >>> import torch

        >>> tokenizer = RealmTokenizer.from_pretrained("google/realm-cc-news-pretrained-embedder")
        >>> model = RealmEmbedder.from_pretrained("google/realm-cc-news-pretrained-embedder")

        >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
        >>> outputs = model(**inputs)

        >>> projected_score = outputs.projected_score
        ```
        """

        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        realm_outputs = self.realm(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        # [batch_size, hidden_size]
        pooler_output = realm_outputs[1]
        # [batch_size, retriever_proj_size]
        projected_score = self.cls(pooler_output)

        if not return_dict:
            return (projected_score,) + realm_outputs[2:4]
        else:
            return RealmEmbedderOutput(
                projected_score=projected_score,
                hidden_states=realm_outputs.hidden_states,
                attentions=realm_outputs.attentions,
            )


@add_start_docstrings(
    "The scorer of REALM outputting relevance scores representing the score of document candidates (before softmax).",
    REALM_START_DOCSTRING,
)
class RealmScorer(RealmPreTrainedModel):
    r"""
    Args:
        query_embedder ([`RealmEmbedder`]):
            Embedder for input sequences. If not specified, it will use the same embedder as candidate sequences.
    """

    def __init__(self, config, query_embedder=None):
        super().__init__(config)

        self.embedder = RealmEmbedder(self.config)

        self.query_embedder = (
            query_embedder if query_embedder is not None else self.embedder
        )

        self.post_init()

    @add_start_docstrings_to_model_forward(
        REALM_INPUTS_DOCSTRING.format("batch_size, sequence_length")
    )
    @replace_return_docstrings(
        output_type=RealmScorerOutput, config_class=_CONFIG_FOR_DOC
    )
    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        candidate_input_ids=None,
        candidate_attention_mask=None,
        candidate_token_type_ids=None,
        candidate_inputs_embeds=None,
        head_mask=None,
        inputs_embeds=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        r"""
        candidate_input_ids (`torch.LongTensor` of shape `(batch_size, num_candidates, sequence_length)`):
            Indices of candidate input sequence tokens in the vocabulary.

            Indices can be obtained using [`RealmTokenizer`]. See [`PreTrainedTokenizer.encode`] and
            [`PreTrainedTokenizer.__call__`] for details.

            [What are input IDs?](../glossary#input-ids)
        candidate_attention_mask (`torch.FloatTensor` of shape `(batch_size, num_candidates, sequence_length)`, *optional*):
            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:

            - 1 for tokens that are **not masked**,
            - 0 for tokens that are **masked**.

            [What are attention masks?](../glossary#attention-mask)
        candidate_token_type_ids (`torch.LongTensor` of shape `(batch_size, num_candidates, sequence_length)`, *optional*):
            Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
            1]`:

            - 0 corresponds to a *sentence A* token,
            - 1 corresponds to a *sentence B* token.

            [What are token type IDs?](../glossary#token-type-ids)
        candidate_inputs_embeds (`torch.FloatTensor` of shape `(batch_size * num_candidates, sequence_length, hidden_size)`, *optional*):
            Optionally, instead of passing `candidate_input_ids` you can choose to directly pass an embedded
            representation. This is useful if you want more control over how to convert *candidate_input_ids* indices
            into associated vectors than the model's internal embedding lookup matrix.

        Returns:

        Example:

        ```python
        >>> import torch
        >>> from transformers import RealmTokenizer, RealmScorer

        >>> tokenizer = RealmTokenizer.from_pretrained("google/realm-cc-news-pretrained-scorer")
        >>> model = RealmScorer.from_pretrained("google/realm-cc-news-pretrained-scorer", num_candidates=2)

        >>> # batch_size = 2, num_candidates = 2
        >>> input_texts = ["How are you?", "What is the item in the picture?"]
        >>> candidates_texts = [["Hello world!", "Nice to meet you!"], ["A cute cat.", "An adorable dog."]]

        >>> inputs = tokenizer(input_texts, return_tensors="pt")
        >>> candidates_inputs = tokenizer.batch_encode_candidates(candidates_texts, max_length=10, return_tensors="pt")

        >>> outputs = model(
        ...     **inputs,
        ...     candidate_input_ids=candidates_inputs.input_ids,
        ...     candidate_attention_mask=candidates_inputs.attention_mask,
        ...     candidate_token_type_ids=candidates_inputs.token_type_ids,
        ... )
        >>> relevance_score = outputs.relevance_score
        ```"""

        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        if input_ids is None and inputs_embeds is None:
            raise ValueError("You have to specify either input_ids or input_embeds.")

        if candidate_input_ids is None and candidate_inputs_embeds is None:
            raise ValueError(
                "You have to specify either candidate_input_ids or candidate_inputs_embeds."
            )

        query_outputs = self.query_embedder(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        # [batch_size * num_candidates, candidate_seq_len]
        (
            flattened_input_ids,
            flattened_attention_mask,
            flattened_token_type_ids,
        ) = self._flatten_inputs(
            candidate_input_ids, candidate_attention_mask, candidate_token_type_ids
        )

        candidate_outputs = self.embedder(
            flattened_input_ids,
            attention_mask=flattened_attention_mask,
            token_type_ids=flattened_token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=candidate_inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        # [batch_size, retriever_proj_size]
        query_score = query_outputs[0]
        # [batch_size * num_candidates, retriever_proj_size]
        candidate_score = candidate_outputs[0]
        # [batch_size, num_candidates, retriever_proj_size]
        candidate_score = candidate_score.view(
            -1, self.config.num_candidates, self.config.retriever_proj_size
        )
        # [batch_size, num_candidates]
        relevance_score = torch.einsum("BD,BND->BN", query_score, candidate_score)

        if not return_dict:
            return relevance_score, query_score, candidate_score

        return RealmScorerOutput(
            relevance_score=relevance_score,
            query_score=query_score,
            candidate_score=candidate_score,
        )


@add_start_docstrings(
    "The knowledge-augmented encoder of REALM outputting masked language model logits and marginal log-likelihood"
    " loss.",
    REALM_START_DOCSTRING,
)
class RealmKnowledgeAugEncoder(RealmPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.realm = RealmBertModel(self.config)
        self.cls = RealmOnlyMLMHead(self.config)
        self.post_init()

    def get_input_embeddings(self):
        return self.realm.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.realm.embeddings.word_embeddings = value

    def get_output_embeddings(self):
        return self.cls.predictions.decoder

    def set_output_embeddings(self, new_embeddings):
        self.cls.predictions.decoder = new_embeddings

    @add_start_docstrings_to_model_forward(
        REALM_INPUTS_DOCSTRING.format("batch_size, num_candidates, sequence_length")
    )
    @replace_return_docstrings(output_type=MaskedLMOutput, config_class=_CONFIG_FOR_DOC)
    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        relevance_score=None,
        labels=None,
        mlm_mask=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        r"""
        relevance_score (`torch.FloatTensor` of shape `(batch_size, num_candidates)`, *optional*):
            Relevance score derived from RealmScorer, must be specified if you want to compute the masked language
            modeling loss.

        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
            config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
            loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`

        mlm_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Mask to avoid calculating joint loss on certain positions. If not specified, the loss will not be masked.
            Mask values selected in `[0, 1]`:

            - 1 for tokens that are **not masked**,
            - 0 for tokens that are **masked**.

        Returns:

        Example:

        ```python
        >>> import torch
        >>> from transformers import RealmTokenizer, RealmKnowledgeAugEncoder

        >>> tokenizer = RealmTokenizer.from_pretrained("google/realm-cc-news-pretrained-encoder")
        >>> model = RealmKnowledgeAugEncoder.from_pretrained(
        ...     "google/realm-cc-news-pretrained-encoder", num_candidates=2
        ... )

        >>> # batch_size = 2, num_candidates = 2
        >>> text = [["Hello world!", "Nice to meet you!"], ["The cute cat.", "The adorable dog."]]

        >>> inputs = tokenizer.batch_encode_candidates(text, max_length=10, return_tensors="pt")
        >>> outputs = model(**inputs)
        >>> logits = outputs.logits
        ```"""

        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        (
            flattened_input_ids,
            flattened_attention_mask,
            flattened_token_type_ids,
        ) = self._flatten_inputs(input_ids, attention_mask, token_type_ids)

        joint_outputs = self.realm(
            flattened_input_ids,
            attention_mask=flattened_attention_mask,
            token_type_ids=flattened_token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        # [batch_size * num_candidates, joint_seq_len, hidden_size]
        joint_output = joint_outputs[0]
        # [batch_size * num_candidates, joint_seq_len, vocab_size]
        prediction_scores = self.cls(joint_output)
        # [batch_size, num_candidates]
        candidate_score = relevance_score

        masked_lm_loss = None
        if labels is not None:
            if candidate_score is None:
                raise ValueError(
                    "You have to specify `relevance_score` when `labels` is specified in order to compute loss."
                )

            batch_size, seq_length = labels.size()

            if mlm_mask is None:
                mlm_mask = torch.ones_like(labels, dtype=torch.float32)
            else:
                mlm_mask = mlm_mask.type(torch.float32)

            # Compute marginal log-likelihood
            loss_fct = CrossEntropyLoss(reduction="none")  # -100 index = padding token

            # [batch_size * num_candidates * joint_seq_len, vocab_size]
            mlm_logits = prediction_scores.view(-1, self.config.vocab_size)
            # [batch_size * num_candidates * joint_seq_len]
            mlm_targets = labels.tile(1, self.config.num_candidates).view(-1)
            # [batch_size, num_candidates, joint_seq_len]
            masked_lm_log_prob = -loss_fct(mlm_logits, mlm_targets).view(
                batch_size, self.config.num_candidates, seq_length
            )
            # [batch_size, num_candidates, 1]
            candidate_log_prob = candidate_score.log_softmax(-1).unsqueeze(-1)
            # [batch_size, num_candidates, joint_seq_len]
            joint_gold_log_prob = candidate_log_prob + masked_lm_log_prob
            # [batch_size, joint_seq_len]
            marginal_gold_log_probs = joint_gold_log_prob.logsumexp(1)
            # []
            masked_lm_loss = -torch.nansum(
                torch.sum(marginal_gold_log_probs * mlm_mask) / torch.sum(mlm_mask)
            )

        if not return_dict:
            output = (prediction_scores,) + joint_outputs[2:4]
            return (
                ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
            )

        return MaskedLMOutput(
            loss=masked_lm_loss,
            logits=prediction_scores,
            hidden_states=joint_outputs.hidden_states,
            attentions=joint_outputs.attentions,
        )


@add_start_docstrings("The reader of REALM.", REALM_START_DOCSTRING)
class RealmReader(RealmPreTrainedModel):

    _keys_to_ignore_on_load_unexpected = [r"pooler", "cls"]

    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels

        self.realm = RealmBertModel(config)
        self.cls = RealmOnlyMLMHead(config)
        self.qa_outputs = RealmReaderProjection(config)

        self.post_init()

    @add_start_docstrings_to_model_forward(
        REALM_INPUTS_DOCSTRING.format("reader_beam_size, sequence_length")
    )
    @replace_return_docstrings(
        output_type=RealmReaderOutput, config_class=_CONFIG_FOR_DOC
    )
    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        relevance_score=None,
        block_mask=None,
        start_positions=None,
        end_positions=None,
        has_answers=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        r"""
        relevance_score (`torch.FloatTensor` of shape `(searcher_beam_size,)`, *optional*):
            Relevance score, which must be specified if you want to compute the logits and marginal log loss.
        block_mask (`torch.BoolTensor` of shape `(searcher_beam_size, sequence_length)`, *optional*):
            The mask of the evidence block, which must be specified if you want to compute the logits and marginal log
            loss.
        start_positions (`torch.LongTensor` of shape `(searcher_beam_size,)`, *optional*):
            Labels for position (index) of the start of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
            are not taken into account for computing the loss.
        end_positions (`torch.LongTensor` of shape `(searcher_beam_size,)`, *optional*):
            Labels for position (index) of the end of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
            are not taken into account for computing the loss.
        has_answers (`torch.BoolTensor` of shape `(searcher_beam_size,)`, *optional*):
            Whether or not the evidence block has answer(s).

        Returns:
        """
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        if relevance_score is None:
            raise ValueError(
                "You have to specify `relevance_score` to calculate logits and loss."
            )
        if block_mask is None:
            raise ValueError(
                "You have to specify `block_mask` to separate question block and evidence block."
            )
        if token_type_ids.size(1) < self.config.max_span_width:
            raise ValueError(
                "The input sequence length must be greater than or equal to config.max_span_width."
            )

        outputs = self.realm(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        # [reader_beam_size, joint_seq_len, hidden_size]
        # print("outputs: ", outputs)
        sequence_output = outputs[0]
        # [reader_beam_size, num_candidates], [num_candidates], [num_candidates]
        reader_logits, candidate_starts, candidate_ends = self.qa_outputs(
            sequence_output, block_mask[0 : self.config.reader_beam_size]
        )
        # [searcher_beam_size, 1]
        retriever_logits = torch.unsqueeze(
            relevance_score[0 : self.config.reader_beam_size], -1
        )
        # [reader_beam_size, num_candidates]
        reader_logits += retriever_logits
        # []
        starts: torch.LongTensor = torch.tensor([], device=reader_logits.device)
        ends: torch.LongTensor = torch.tensor([], device=reader_logits.device)

        for logits in reader_logits:
            logits = torch.unsqueeze(logits, 0)
            predicted_block_index = torch.argmax(torch.max(logits, dim=1).values)
            predicted_candidate = torch.argmax(torch.max(logits, dim=0).values)
            predicted_start = torch.index_select(
                candidate_starts, dim=0, index=predicted_candidate
            )
            # [1]
            predicted_end = torch.index_select(
                candidate_ends, dim=0, index=predicted_candidate
            )
            starts = torch.cat((starts, predicted_start), dim=0)
            ends = torch.cat((ends, predicted_end), dim=0)

        # get the context to retrieve the answer
        # predicted_block_index = torch.argmax(torch.max(reader_logits, dim=1).values)
        # # []
        # predicted_candidate = torch.argmax(torch.max(reader_logits, dim=0).values)
        # # [1]
        # predicted_start = torch.index_select(candidate_starts, dim=0, index=predicted_candidate)
        # # [1]
        # predicted_end = torch.index_select(candidate_ends, dim=0, index=predicted_candidate)

        total_loss = None
        retriever_loss = None
        reader_loss = None
        retriever_correct = None
        reader_correct = None
        if (
            start_positions is not None
            and end_positions is not None
            and has_answers is not None
        ):

            def compute_correct_candidates(
                candidate_starts, candidate_ends, gold_starts, gold_ends
            ):
                """Compute correct span."""
                # [reader_beam_size, num_answers, num_candidates]
                is_gold_start = torch.eq(
                    torch.unsqueeze(torch.unsqueeze(candidate_starts, 0), 0),
                    torch.unsqueeze(gold_starts, -1),
                )
                is_gold_end = torch.eq(
                    torch.unsqueeze(torch.unsqueeze(candidate_ends, 0), 0),
                    torch.unsqueeze(gold_ends, -1),
                )

                # [reader_beam_size, num_candidates]
                return torch.any(torch.logical_and(is_gold_start, is_gold_end), 1)

            def marginal_log_loss(logits, is_correct):
                """Loss based on the negative marginal log-likelihood."""

                def mask_to_score(mask):
                    return (1.0 - mask.type(torch.float32)) * -10000.0

                # []
                log_numerator = torch.logsumexp(
                    logits + mask_to_score(is_correct), dim=-1
                )
                log_denominator = torch.logsumexp(logits, dim=-1)
                return log_denominator - log_numerator

            # sometimes the start/end positions are outside our model inputs, we ignore these terms
            # `-1` is reserved for no answer.
            ignored_index = sequence_output.size(1)
            start_positions = start_positions.clamp(-1, ignored_index)
            end_positions = end_positions.clamp(-1, ignored_index)

            retriever_correct = has_answers
            any_retriever_correct = torch.any(retriever_correct)

            reader_correct = compute_correct_candidates(
                candidate_starts=candidate_starts,
                candidate_ends=candidate_ends,
                gold_starts=start_positions[0 : self.config.reader_beam_size],
                gold_ends=end_positions[0 : self.config.reader_beam_size],
            )
            any_reader_correct = torch.any(reader_correct)

            retriever_loss = marginal_log_loss(relevance_score, retriever_correct)
            reader_loss = marginal_log_loss(
                reader_logits.view(-1), reader_correct.view(-1)
            )
            retriever_loss *= any_retriever_correct.type(torch.float32)
            reader_loss *= any_reader_correct.type(torch.float32)

            total_loss = (retriever_loss + reader_loss).mean()

        if not return_dict:
            output = (
                predicted_block_index,
                predicted_candidate,
                starts,
                ends,
            ) + outputs[2:]
            return (
                (
                    (
                        total_loss,
                        retriever_loss,
                        reader_loss,
                        retriever_correct,
                        reader_correct,
                    )
                    + output
                )
                if total_loss is not None
                else output
            )

        return RealmReaderOutput(
            loss=total_loss,
            retriever_loss=retriever_loss,
            reader_loss=reader_loss,
            retriever_correct=retriever_correct,
            reader_correct=reader_correct,
            block_idx=predicted_block_index,
            candidate=predicted_candidate,
            start_pos=starts,
            end_pos=ends,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


REALM_FOR_OPEN_QA_DOCSTRING = r"""
    Args:
        input_ids (`torch.LongTensor` of shape `({0})`):
            Indices of input sequence tokens in the vocabulary.

            Indices can be obtained using [`RealmTokenizer`]. See [`PreTrainedTokenizer.encode`] and
            [`PreTrainedTokenizer.__call__`] for details.

            [What are input IDs?](../glossary#input-ids)
        attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:

            - 1 for tokens that are **not masked**,
            - 0 for tokens that are **masked**.

            [What are attention masks?](../glossary#attention-mask)
        token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*):
            Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
            1]`:

            - 0 corresponds to a *sentence A* token,
            - 1 corresponds to a *sentence B* token (should not be used in this model by design).

            [What are token type IDs?](../glossary#token-type-ids)
        answer_ids (`list` of shape `(num_answers, answer_length)`, *optional*):
            Answer ids for computing the marginal log-likelihood loss. Indices should be in `[-1, 0, ...,
            config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-1` are ignored (masked), the
            loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
"""


@add_start_docstrings(
    "`RealmForOpenQA` for end-to-end open domain question answering.",
    REALM_START_DOCSTRING,
)
class RealmForOpenQA(RealmPreTrainedModel):
    def __init__(self, config, retriever=None):
        super().__init__(config)
        self.embedder = RealmEmbedder(config)
        self.reader = RealmReader(config)
        self.register_buffer(
            "block_emb",
            torch.zeros(()).new_empty(
                size=(config.num_block_records, config.retriever_proj_size),
                dtype=torch.float32,
                device=torch.device("cpu"),
            ),
        )
        self.retriever = retriever

        self.post_init()

    @property
    def searcher_beam_size(self):
        if self.training:
            return self.config.searcher_beam_size
        return self.config.reader_beam_size

    def block_embedding_to(self, device):
        """Send `self.block_emb` to a specific device.

        Args:
            device (`str` or `torch.device`):
                The device to which `self.block_emb` will be sent.
        """

        self.block_emb = self.block_emb.to(device)

    @add_start_docstrings_to_model_forward(
        REALM_FOR_OPEN_QA_DOCSTRING.format("1, sequence_length")
    )
    @replace_return_docstrings(
        output_type=RealmForOpenQAOutput, config_class=_CONFIG_FOR_DOC
    )
    def forward(
        self,
        input_ids,
        attention_mask=None,
        token_type_ids=None,
        answer_ids=None,
        return_dict=None,
    ):
        r"""
        Returns:

        Example:

        ```python
        >>> import torch
        >>> from transformers import RealmForOpenQA, RealmRetriever, RealmTokenizer

        >>> retriever = RealmRetriever.from_pretrained("google/realm-orqa-nq-openqa")
        >>> tokenizer = RealmTokenizer.from_pretrained("google/realm-orqa-nq-openqa")
        >>> model = RealmForOpenQA.from_pretrained("google/realm-orqa-nq-openqa", retriever=retriever)

        >>> question = "Who is the pioneer in modern computer science?"
        >>> question_ids = tokenizer([question], return_tensors="pt")
        >>> answer_ids = tokenizer(
        ...     ["alan mathison turing"],
        ...     add_special_tokens=False,
        ...     return_token_type_ids=False,
        ...     return_attention_mask=False,
        ... ).input_ids

        >>> reader_output, predicted_answer_ids = model(**question_ids, answer_ids=answer_ids, return_dict=False)
        >>> predicted_answer = tokenizer.decode(predicted_answer_ids)
        >>> loss = reader_output.loss
        ```"""

        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        if input_ids is not None and input_ids.shape[0] != 1:
            raise ValueError("The batch_size of the inputs must be 1.")

        question_outputs = self.embedder(
            input_ids=input_ids,
            token_type_ids=token_type_ids,
            attention_mask=attention_mask,
            return_dict=True,
        )
        # [1, projection_size]
        question_projection = question_outputs[0]

        # CPU computation starts.
        # [1, block_emb_size]
        batch_scores = torch.einsum(
            "BD,QD->QB", self.block_emb, question_projection.to(self.block_emb.device)
        )
        # [1, searcher_beam_size]
        _, retrieved_block_ids = torch.topk(
            batch_scores, k=self.searcher_beam_size, dim=-1
        )

        # [searcher_beam_size]
        retrieved_block_ids = retrieved_block_ids.squeeze()
        # [searcher_beam_size, projection_size]
        retrieved_block_emb = torch.index_select(
            self.block_emb, dim=0, index=retrieved_block_ids
        )
        # CPU computation ends.
        # Retrieve possible answers
        has_answers, start_pos, end_pos, concat_inputs = self.retriever(
            retrieved_block_ids.cpu(),
            input_ids,
            answer_ids,
            max_length=self.config.reader_seq_len,
        )

        concat_inputs = concat_inputs.to(self.reader.device)
        block_mask = concat_inputs.special_tokens_mask.type(torch.bool).to(
            device=self.reader.device
        )
        block_mask.logical_not_().logical_and_(
            concat_inputs.token_type_ids.type(torch.bool)
        )

        if has_answers is not None:
            has_answers = torch.tensor(
                has_answers, dtype=torch.bool, device=self.reader.device
            )
            start_pos = torch.tensor(
                start_pos, dtype=torch.long, device=self.reader.device
            )
            end_pos = torch.tensor(end_pos, dtype=torch.long, device=self.reader.device)

        # [searcher_beam_size]
        retrieved_logits = torch.einsum(
            "D,BD->B",
            question_projection.squeeze(),
            retrieved_block_emb.to(self.reader.device),
        )

        reader_output = self.reader(
            input_ids=concat_inputs.input_ids[0 : self.config.reader_beam_size],
            attention_mask=concat_inputs.attention_mask[
                0 : self.config.reader_beam_size
            ],
            token_type_ids=concat_inputs.token_type_ids[
                0 : self.config.reader_beam_size
            ],
            relevance_score=retrieved_logits,
            block_mask=block_mask,
            has_answers=has_answers,
            start_positions=start_pos,
            end_positions=end_pos,
            return_dict=True,
        )

        #     tokenizer = self.config.tokenizer.add_special_tokens(
        #     {"additional_special_tokens": ["<hl>"]}
        # )

        batch_answer_ids: list = []
        sequence_length = 128
        for idx in range(self.config.reader_beam_size):
            predicted_block = concat_inputs.input_ids[idx]
            start = reader_output.start_pos.detach().cpu().numpy()[idx]
            end = reader_output.end_pos.detach().cpu().numpy()[idx]
            predicted_answer_ids = predicted_block[int(start) : int(end) + 1]
            pad = torch.zeros(
                sequence_length - len(predicted_answer_ids),
                dtype=torch.long,
                device=self.reader.device,
            )
            batch_answer_ids.append(torch.cat((predicted_answer_ids, pad), dim=0))
            # print(predicted_block)
        batch_answer_ids = torch.stack(batch_answer_ids)

        if not return_dict:
            return reader_output, batch_answer_ids

        return RealmForOpenQAOutput(
            reader_output=reader_output,
            retrieved_block_ids=retrieved_block_ids,
            input_ids=concat_inputs.input_ids,
            predicted_answer_ids=batch_answer_ids,
        )

src/cf_generation/baseline_generation/modelling_realm.py
UKPLab/CATfOOD
"""
Use an ensemble of 6 RC models using different random seeds
to filter noisy counterfactuals

we keep any generated (q′, c′, a′) triples where at least 5 of the 6
models agree on the answer
"""

from tqdm import tqdm
import jsonlines
import torch
from transformers import T5ForConditionalGeneration, T5Tokenizer

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"

NUM_BEAMS = 4
BATCH_SIZE = 32

# dev mode
TEST = True


class NoiseFiltering:
    def __init__(self, model_path, tokenizer):

        self.sep_token = ">>"
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model_path = model_path
        self.tokenizer = tokenizer
        self.tokenizer.add_special_tokens(
            {"additional_special_tokens": [self.sep_token]}
        )
        self.model = T5ForConditionalGeneration.from_pretrained(self.model_path).to(
            self.device
        )
        self.model.eval()

        self.batch_size = BATCH_SIZE

    def _prepare_inputs(self, examples):
        """
        Prepare the inputs for the model
        :param examples: the examples to prepare
        :return: the prepared inputs
        """

        # prepare the inputs
        def generate_input(example):
            return " ".join(
                [
                    example["predicted_question"],
                    self.sep_token,
                    example["retrieved_context"],
                ]
            )

        inputs = [generate_input(example) for example in examples]
        return inputs

    def filter(self, examples):
        """
        Filter the counterfactuals
        :return: the filtered counterfactuals
        """

        # tokenize the inputs
        inputs = self._prepare_inputs(examples)
        features = self.tokenizer(
            inputs,
            max_length=512,
            padding="max_length",
            truncation=True,
            return_tensors="pt",
        ).to(self.device)

        with torch.no_grad():
            outputs = self.model.generate(
                **features, max_length=128, num_beams=NUM_BEAMS, early_stopping=True
            )
            predictions = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
        return predictions

    def predict(self, examples):
        """
        Predict the answers for the counterfactuals
        """
        print(len(examples))
        model_outputs = []
        for example in tqdm(range(0, len(examples), self.batch_size)):
            data = examples[example : example + self.batch_size]
            predictions = self.filter(data)
            model_outputs.extend(predictions)
        return model_outputs


def save_to_disk(data, file_name):
    with jsonlines.open(file_name, "a") as writer:
        for example in tqdm(data, total=len(data), desc="Saving samples ... "):
            writer.write(example)


if __name__ == "__main__":
    paths = ["t5-large-squad-qa-seed-42"]
    examples = []
    with jsonlines.open(
        BASE_PATH
        + "src/data/squad/t5_squad_counterfactuals/rag_counterfactuals_complete.jsonl"
    ) as reader:
        for example in reader:
            example["alternate_answers"] = []
            examples.append(example)
    # sample first n% of examples
    # num_samples = int(0.20 * len(examples))
    # print("Number of samples: ", num_samples)
    # examples = examples[: num_samples]
    print(len(examples))

    c = 0  # counter
    for path in paths:
        model_path = BASE_PATH + path
        tokenizer = T5Tokenizer.from_pretrained(model_path)
        nf = NoiseFiltering(model_path=model_path, tokenizer=tokenizer)
        predictions = nf.predict(examples)
        for i, ex in enumerate(examples):
            ex["alternate_answers"].append(predictions[i])
            c += 1
        # if c % 1000 == 0:
        #     save_to_disk(examples, BASE_PATH + "src/data/squad/rag_predictions_noise_filtered_10p.jsonl")
        # save examples to disk
        save_to_disk(
            examples,
            BASE_PATH
            + f"src/data/squad/t5_squad_counterfactuals/rag_counterfactuals_complete_nf_{path}_1.jsonl",
        )

src/cf_generation/baseline_generation/noise_filtering.py
UKPLab/CATfOOD
import os
import re
import numpy as np
from tqdm import tqdm

import jsonlines

DOC_STRIDE = 256
MAX_LENGTH = 512
SEED = 42
PROCESS_TRAIN = os.environ.pop("PROCESS_TRAIN", "false")
CATEGORY_MAPPING = {"null": 0, "short": 1, "long": 2, "yes": 3, "no": 4}


def _get_single_answer(example):
    def choose_first(answer, is_long_answer=False):
        assert isinstance(answer, list)
        if len(answer) == 1:
            answer = answer[0]
            return {k: [answer[k]] for k in answer} if is_long_answer else answer
        for a in answer:
            if is_long_answer:
                a = {k: [a[k]] for k in a}
            if len(a["start_token"]) > 0:
                break
        return a

    answer = {"id": example["example_id"]}
    annotation = example["annotations"]
    print(annotation)
    yes_no_answer = annotation[0]["yes_no_answer"]
    print(yes_no_answer)
    if yes_no_answer != "NONE":
        answer["category"] = ["yes"] if 1 in yes_no_answer else ["no"]
        answer["start_token"] = answer["end_token"] = []
        answer["start_byte"] = answer["end_byte"] = []
        answer["text"] = ["<cls>"]
    else:
        answer["category"] = ["short"]
        out = choose_first(annotation[0]["short_answers"])
        out["text"] = []
        out["start_token"] = out["start_token"]
        out["end_token"] = out["end_token"]
        print(out)
        if out["start_token"] == 0:
            # answer will be long if short is not available
            answer["category"] = ["long"]
            out = choose_first(annotation[0]["long_answer"], is_long_answer=True)
            out["text"] = []
        print(out)
        answer.update(out)

    # disregard some samples
    if answer["start_token"] == 0 or answer["start_token"] == answer["end_token"]:
        answer["remove_it"] = True
    else:
        answer["remove_it"] = False
    print(answer)
    # cols = ["start_token", "end_token", "text"]
    # if not all([isinstance(answer[k], list) for k in cols]):
    #     raise ValueError("Issue in ID", example["example_id"])

    return answer


def get_context_and_ans(example, assertion=False):
    """Gives new context after removing <html> & new answer tokens as per new context"""
    answer = _get_single_answer(example)
    # bytes are of no use
    # del answer["start_byte"]
    # del answer["end_byte"]

    # handle yes_no answers explicitly
    if answer["category"][0] in ["yes", "no"]:  # category is list with one element
        doc = example["document"]["tokens"]
        context = []
        for i in range(len(doc["token"])):
            if not doc["is_html"][i]:
                context.append(doc["token"][i])
        return {
            "context": " ".join(context),
            "answer": {
                "start_token": -100,  # ignore index in cross-entropy
                "end_token": -100,  # ignore index in cross-entropy
                "category": answer["category"],
                "span": answer["category"],  # extra
            },
        }

    # later, help in removing all no answers
    if answer["start_token"] == [-1]:
        return {
            "context": "None",
            "answer": {
                "start_token": -1,
                "end_token": -1,
                "category": "null",
                "span": "None",  # extra
            },
        }

    # handling normal samples

    cols = ["start_token", "end_token"]
    print("ans", answer)
    answer.update({k: answer[k] for k in cols})  # e.g. [10] == 10

    # doc = example["document"]["tokens"]
    start_token = answer["start_token"]
    end_token = answer["end_token"]
    #
    # context = []
    # for i in range(len(doc["token"])):
    #     if not doc["is_html"][i]:
    #         context.append(doc["token"][i])
    #     else:
    #         if answer["start_token"] > i:
    #             start_token -= 1
    #         if answer["end_token"] > i:
    #             end_token -= 1
    # new = " ".join(context[start_token:end_token])
    new = " ".join(example["document_text"].split(" ")[start_token:end_token])
    print(new)

    CLEANR = re.compile("<.*?>")

    def cleanhtml(raw_html):
        cleantext = re.sub(CLEANR, "", raw_html)
        return cleantext

    # checking above code
    # if assertion:
    #     """checking if above code is working as expected for all the samples"""
    #     is_html = doc["is_html"][answer["start_token"] : answer["end_token"]]
    #     old = doc["token"][answer["start_token"] : answer["end_token"]]
    #     old = " ".join([old[i] for i in range(len(old)) if not is_html[i]])
    #     if new != old:
    #         print("ID:", example["id"])
    #         print("New:", new, end="\n")
    #         print("Old:", old, end="\n\n")
    context = cleanhtml(example["document_text"])
    context = " ".join(context.split())
    print(context)

    return {
        "context": context,
        "answer": {
            "start_token": start_token,
            "end_token": end_token - 1,  # this makes it inclusive
            "category": answer["category"],  # either long or short
            "span": new,  # extra
        },
    }


def get_strided_contexts_and_ans(
    example, tokenizer, doc_stride=256, max_length=512, assertion=True
):
    # overlap will be of doc_stride - q_len

    out = get_context_and_ans(example, assertion=assertion)
    print(out)
    answer = out["answer"]

    # later, removing these samples
    if answer["start_token"] == -1:
        return {
            "example_id": example["id"],
            "input_ids": [[-1]],
            "labels": {"start_token": [-1], "end_token": [-1], "category": ["null"],},
        }

    input_ids = tokenizer(example["question_text"], out["context"]).input_ids
    print(input_ids)
    print(tokenizer.eos_token_id)
    q_len = input_ids.index(tokenizer.eos_token_id) + 1

    # return yes/no
    if answer["category"][0] in ["yes", "no"]:  # category is list with one element
        inputs = []
        category = []
        q_indices = input_ids[:q_len]
        doc_start_indices = range(q_len, len(input_ids), max_length - doc_stride)
        for i in doc_start_indices:
            end_index = i + max_length - q_len
            slice = input_ids[i:end_index]
            inputs.append(q_indices + slice)
            category.append(answer["category"][0])
            if slice[-1] == tokenizer.sep_token_id:
                break

        return {
            "example_id": example["id"],
            "input_ids": inputs,
            "labels": {
                "start_token": [-100] * len(category),
                "end_token": [-100] * len(category),
                "category": category,
            },
        }

    splitted_context = out["context"].split()
    complete_end_token = splitted_context[answer["end_token"]]
    answer["start_token"] = len(
        tokenizer(
            " ".join(splitted_context[: answer["start_token"]]),
            add_special_tokens=False,
        ).input_ids
    )
    answer["end_token"] = len(
        tokenizer(
            " ".join(splitted_context[: answer["end_token"]]), add_special_tokens=False
        ).input_ids
    )

    answer["start_token"] += q_len
    answer["end_token"] += q_len

    # fixing end token
    num_sub_tokens = len(
        tokenizer(complete_end_token, add_special_tokens=False).input_ids
    )
    if num_sub_tokens > 1:
        answer["end_token"] += num_sub_tokens - 1

    old = input_ids[
        answer["start_token"] : answer["end_token"] + 1
    ]  # right & left are inclusive
    start_token = answer["start_token"]
    end_token = answer["end_token"]

    if assertion:
        """This won't match exactly because of extra gaps => visaully inspect everything"""
        new = tokenizer.decode(old)
        if answer["span"] != new:
            print("ISSUE IN TOKENIZATION")
            print("OLD:", answer["span"])
            print("NEW:", new, end="\n\n")

    if len(input_ids) <= max_length:
        return {
            "example_id": example["id"],
            "input_ids": [input_ids],
            "labels": {
                "start_token": [answer["start_token"]],
                "end_token": [answer["end_token"]],
                "category": answer["category"],
            },
        }

    q_indices = input_ids[:q_len]
    doc_start_indices = range(q_len, len(input_ids), max_length - doc_stride)

    inputs = []
    answers_start_token = []
    answers_end_token = []
    answers_category = []  # null, yes, no, long, short
    for i in doc_start_indices:
        end_index = i + max_length - q_len
        slice = input_ids[i:end_index]
        inputs.append(q_indices + slice)
        assert len(inputs[-1]) <= max_length, "Issue in truncating length"

        if start_token >= i and end_token <= end_index - 1:
            start_token = start_token - i + q_len
            end_token = end_token - i + q_len
            answers_category.append(answer["category"][0])  # ["short"] -> "short"
        else:
            start_token = -100
            end_token = -100
            answers_category.append("null")
        new = inputs[-1][start_token : end_token + 1]

        answers_start_token.append(start_token)
        answers_end_token.append(end_token)
        if assertion:
            """checking if above code is working as expected for all the samples"""
            if new != old and new != [tokenizer.cls_token_id]:
                print("ISSUE in strided for ID:", example["id"])
                print("New:", tokenizer.decode(new))
                print("Old:", tokenizer.decode(old), end="\n\n")
        if slice[-1] == tokenizer.sep_token_id:
            break
    print(inputs)
    print(len(inputs))
    print(len(inputs[0]))
    print(
        {
            "start_token": answers_start_token,
            "end_token": answers_end_token,
            "category": answers_category,
        }
    )

    return {
        "example_id": example["example_id"],
        "input_ids": inputs,
        "labels": {
            "start_token": answers_start_token,
            "end_token": answers_end_token,
            "category": answers_category,
        },
    }


def prepare_inputs(example, tokenizer, doc_stride=256, max_length=512, assertion=False):
    example = get_strided_contexts_and_ans(
        example,
        tokenizer,
        doc_stride=doc_stride,
        max_length=max_length,
        assertion=assertion,
    )

    return example


def save_to_disk(hf_data, file_name):
    with jsonlines.open(file_name, "a") as writer:
        for example in tqdm(hf_data, total=len(hf_data), desc="Saving samples ... "):
            labels = example["labels"]
            for ids, start, end, cat in zip(
                example["input_ids"],
                labels["start_token"],
                labels["end_token"],
                labels["category"],
            ):
                if start == -1 and end == -1:
                    continue  # leave waste samples with no answer
                if cat == "null" and np.random.rand() < 0.6:
                    continue  # removing 50 % samples
                writer.write(
                    {
                        "input_ids": ids,
                        "start_token": start,
                        "end_token": end,
                        "category": CATEGORY_MAPPING[cat],
                    }
                )


if __name__ == "__main__":
    """Running area"""
    from datasets import load_dataset
    from transformers import T5Tokenizer

    # data = load_dataset("natural_questions")
    tokenizer = T5Tokenizer.from_pretrained("t5-base")
    # # print(data["validation"][0])
    #
    # data = data["train" if PROCESS_TRAIN == "true" else "validation"]
    #
    # cache_file_name = "data/nq-training" if PROCESS_TRAIN == "true" else "data/nq-validation"
    # fn_kwargs = dict(
    #     tokenizer=tokenizer,
    #     doc_stride=DOC_STRIDE,
    #     max_length=MAX_LENGTH,
    #     assertion=False,
    # )
    # data = data.map(
    #     prepare_inputs, fn_kwargs=fn_kwargs, cache_file_name=cache_file_name
    # )
    # data = data.remove_columns(["annotations", "document", "id", "question"])
    # print(data)
    #
    # np.random.seed(SEED)
    # save_to_disk(data, file_name=cache_file_name + ".jsonl")
    import jsonlines

    with jsonlines.open("data/simplified-nq-train.jsonl") as file:
        for line in file.iter():
            print(line)
            prepare_inputs(line, tokenizer)
            break

src/cf_generation/baseline_generation/nq_dataset.py
UKPLab/CATfOOD
import json
import os

import numpy as np
from tqdm import tqdm

SEED = 42
PROCESS_TRAIN = os.environ.pop("PROCESS_TRAIN", "false")
CATEGORY_MAPPING = {"null": 0, "short": 1, "long": 2, "yes": 3, "no": 4}


def _get_single_answer(example):
    def choose_first(answer, is_long_answer=False):
        # print(answer)
        if is_long_answer:
            assert isinstance(answer, dict)
            answer = [answer]
        else:
            assert isinstance(answer, list)
        if len(answer) == 1:
            answer = answer[0]
            return {k: [answer[k]][0] for k in answer} if is_long_answer else answer
        for a in answer:
            if is_long_answer:
                a = {k: [a[k]] for k in a}
            if len(answer) > 1:
                break
        return a

    # print(example)
    answer = {"id": example["example_id"]}
    annotation = example["annotations"]
    # print(annotation)
    if isinstance(annotation, list):
        annotation = annotation[0]
    yes_no_answer = annotation["yes_no_answer"]
    if yes_no_answer != "NONE":
        answer["category"] = ["yes"] if yes_no_answer == "YES" else ["no"]
        answer["start_token"] = answer["end_token"] = []
        answer["start_byte"] = answer["end_byte"] = []
        answer["text"] = ["<cls>"]
    elif annotation["short_answers"]:
        answer["category"] = ["short"]
        # print(annotation["short_answers"])
        out = choose_first(annotation["short_answers"])
        out["text"] = []
        answer.update(out)
        # print(out["start_token"])
    else:
        # answer will be long if short is not available
        answer["category"] = ["long"]
        out = choose_first(annotation["long_answer"], is_long_answer=True)
        out["text"] = []
        answer.update(out)
    # print(answer)

    # disregard some samples
    if answer["start_token"] == answer["end_token"]:
        answer["remove_it"] = True
    else:
        answer["remove_it"] = False

    # cols = ["start_token", "end_token", "start_byte", "end_byte", "text"]
    # if not all([isinstance(answer[k], list or int) for k in cols]):
    #     raise ValueError("Issue in ID", example["example_id"])

    return answer


def get_context_and_ans(example, assertion=False):
    """Gives new context after removing <html> & new answer tokens as per new context"""
    answer = _get_single_answer(example)
    # bytes are of no use
    del answer["start_byte"]
    del answer["end_byte"]

    # handle yes_no answers explicitly
    if answer["category"][0] in ["yes", "no"]:  # category is list with one element
        doc = example["document_tokens"]
        context = []
        for i in range(len(doc)):
            if not doc[i]["html_token"]:
                context.append(doc[i]["token"])
        return {
            "context": " ".join(context),
            "answer": {
                "start_token": -100,  # ignore index in cross-entropy
                "end_token": -100,  # ignore index in cross-entropy
                "category": answer["category"],
                "span": answer["category"],  # extra
            },
        }

    # later, help in removing all no answers
    if answer["start_token"] == [-1]:
        return {
            "context": "None",
            "answer": {
                "start_token": -1,
                "end_token": -1,
                "category": "null",
                "span": "None",  # extra
            },
        }

    # handling normal samples
    # cols = ["start_token", "end_token"]
    # answer.update(
    #     {k: answer[k][0] if len(answer[k]) > 0 else answer[k] for k in cols}
    # )  # e.g. [10] == 10

    doc = example["document_tokens"]
    # if len(doc) > 1:
    #     doc = doc[0]
    start_token = answer["start_token"]
    end_token = answer["end_token"]

    context = []
    for i in range(len(doc)):
        if not doc[i]["html_token"]:
            context.append(doc[i]["token"])
        else:
            if answer["start_token"] > i:
                start_token -= 1
            if answer["end_token"] > i:
                end_token -= 1
    new = " ".join(context[start_token:end_token])
    # print(new)

    # checking above code
    if assertion:
        """checking if above code is working as expected for all the samples"""
        is_html = doc["html_token"][answer["start_token"] : answer["end_token"]]
        old = doc["token"][answer["start_token"] : answer["end_token"]]
        old = " ".join([old[i] for i in range(len(old)) if not is_html[i]])
        if new != old:
            print("ID:", example["id"])
            print("New:", new, end="\n")
            print("Old:", old, end="\n\n")

    return {
        "context": " ".join(context),
        "answer": {
            "start_token": start_token,
            "end_token": end_token - 1,  # this makes it inclusive
            "category": answer["category"],  # either long or short
            "span": new,  # extra
        },
    }


def get_contexts_and_ans(example, assertion=True):
    # overlap will be of doc_stride - q_len
    # print(example)
    out = get_context_and_ans(example, assertion=assertion)
    out["question"] = example["question_text"]
    out["example_id"] = example["example_id"]
    out["document_title"] = example["document_title"]
    answer = out["answer"]

    # later, removing these samples
    if answer["start_token"] == -1:
        return {
            "example_id": example["example_id"],
            "title": example["document_title"],
            "question": example["question_text"],
            "context": "",
            "answer": {
                "start_token": -1,
                "end_token": -1,
                "category": ["null"],
                "span": "",
            },
        }

    return out


def prepare_inputs(example, assertion=False):
    example = get_contexts_and_ans(example, assertion=assertion,)

    return example


def save_to_disk(hf_data, file_name):
    with jsonlines.open(file_name, "a") as writer:
        for example in tqdm(hf_data, total=len(hf_data), desc="Saving samples ... "):
            answer = example["answer"]
            if answer["start_token"] == -1 and answer["end_token"] == -1:
                continue  # leave waste samples with no answer
            if answer["category"][0] == "null" and np.random.rand() < 0.6:
                continue  # removing 50 % samples
            writer.write(
                {
                    "id": example["example_id"],
                    "title": example["document_title"],
                    "question": example["question"] + "?",
                    "context": example["context"],
                    "answer": {
                        "start_token": answer["start_token"],
                        "end_token": answer["end_token"],
                        "category": CATEGORY_MAPPING[answer["category"][0]],
                        "text": answer["span"],
                    },
                }
            )


if __name__ == "__main__":
    """Running area"""
    from pathlib import Path
    import jsonlines
    import os
    import gzip

    path = Path("/home/sachdeva/Downloads/v1.0/dev")
    files = os.listdir(path)
    dataset_files = [path / str(f) for f in files]
    # dataset_file = dataset_files[0]

    c = 0
    data = []
    cache_file_name = "data/nq-dev-short-only"
    for data_file in tqdm(dataset_files):
        with gzip.open(data_file, "rb") as f:
            for line in f:
                # print(json.loads(line)["annotations"])
                # print(len(json.loads(line)["annotations"]))
                if json.loads(line)["annotations"][0]["short_answers"]:
                    #     print(json.loads(line)["annotations"][0]["short_answers"])
                    #     print(json.loads(line)["annotations"][0]["long_answer"])
                    # print("--------------------------")
                    # c += 1
                    # if c == 100:
                    #     break
                    example = prepare_inputs(json.loads(line))
                    data.append(example)
                    c += 1
                if c % 1000 == 0:
                    save_to_disk(data, file_name=cache_file_name + ".jsonl")
                    data = []
                if c % 5000 == 0:
                    print(f"Samples processed: {c}")
            # if c == 20:
            #     break
            # break

    # print(example)
    # save_to_disk(data, file_name=cache_file_name + ".jsonl")
    # with jsonlines.open("/home/sachdeva/Downloads/v1.0/train/data.jsonl") as file:
    #     for line in file.iter():
    #         # print(line)
    #         # print("--------------------------")
    #         c += 1
    #         # prepare_inputs(line, tokenizer)
    #         # if c == 10:
    #         #     break
    # print(c)

src/cf_generation/baseline_generation/nq_short_only.py
UKPLab/CATfOOD
import os
import torch
import jsonlines
from tqdm import tqdm
import numpy as np
from typing import List
from collections import OrderedDict
from transformers import (
    RealmRetriever,
    RealmTokenizer,
    RealmConfig,
    AutoModelForSeq2SeqLM,
    AutoTokenizer,
)

from src.cf_generation.baseline_generation.modelling_realm import RealmForOpenQA
from src.calibration.baseline import dataloader

BASE_PATH = os.getenv("PYTHONPATH", "/home/sachdeva/projects/exp_calibration/")
NUM_BEAMS = 15


class CounterfactualGeneration:
    def __init__(self, beam_size: int = 5, prefix: str = None):

        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.hl_token = "<hl>"
        self.prefix = prefix

        # --------------------------------
        # load retriever model
        # --------------------------------
        self.retriever = RealmRetriever.from_pretrained("google/realm-orqa-nq-openqa")
        self.realm_tokenizer = RealmTokenizer.from_pretrained(
            "google/realm-orqa-nq-openqa"
        )
        self.realm_tokenizer.add_tokens(self.hl_token)

        self.hl_token_id = self.realm_tokenizer.convert_tokens_to_ids(["<hl>"])[0]

        realm_config = RealmConfig.from_pretrained("google/realm-orqa-nq-openqa")
        # print(config)
        realm_config.reader_beam_size = beam_size
        self.realm_model = RealmForOpenQA.from_pretrained(
            "google/realm-orqa-nq-openqa",
            retriever=self.retriever,
            config=realm_config,
            torch_dtype=torch.bfloat16,
        )
        self.realm_model.to(self.device)

        # --------------------------------
        # load question generation model
        # --------------------------------
        model_path = BASE_PATH + "t5-3b-squad-qg-seed-42"
        self.qg_tokenizer = AutoTokenizer.from_pretrained(model_path)
        # self.qg_tokenizer.add_special_tokens(
        #     {"additional_special_tokens": [self.hl_token]}
        # )
        self.qg_model = AutoModelForSeq2SeqLM.from_pretrained(
            model_path, torch_dtype=torch.bfloat16
        )
        self.qg_model.to(self.device)
        self.qg_model.eval()

    def _retrieve_context_and_answers(self, question: str = None):
        question_ids = self.realm_tokenizer([question], return_tensors="pt").to(
            self.device
        )
        reader_output = self.realm_model(**question_ids, return_dict=True)
        start_positions = reader_output.reader_output.start_pos.cpu().numpy()
        end_positions = reader_output.reader_output.end_pos.cpu().numpy()
        # add hl_token_id at tensor index start position and end position
        highlight_input_ids = reader_output.input_ids.cpu().numpy()
        inputs: List = []
        for idx, (start, end) in enumerate(zip(start_positions, end_positions)):
            start = int(start)
            end = int(end)
            # get sep idx to retrieve context
            sep_idx = np.where(
                highlight_input_ids[idx, :start] == self.realm_tokenizer.sep_token_id
            )[0][-1]
            inputs.append(
                np.concatenate(
                    [
                        highlight_input_ids[idx, sep_idx:start],
                        [self.hl_token_id],
                        highlight_input_ids[idx, start : end + 1],
                        [self.hl_token_id],
                        highlight_input_ids[idx, end + 1 :],
                    ],
                    axis=0,
                )
            )

        predicted_answers = self.realm_tokenizer.batch_decode(
            reader_output.predicted_answer_ids, skip_special_tokens=True
        )
        contexts = self.realm_tokenizer.batch_decode(
            torch.tensor(np.array(inputs), device=self.device),
            skip_special_tokens=True,
            clean_up_tokenization_spaces=True,
        )

        retrieved_blocks = np.take(
            self.retriever.block_records,
            indices=reader_output.retrieved_block_ids.detach().cpu().numpy(),
            axis=0,
        )

        return (
            contexts,
            predicted_answers,
            retrieved_blocks,
            start_positions,
            end_positions,
        )

    def generate_question(self, question: str = None):
        (
            contexts,
            predicted_answers,
            retrieved_blocks,
            start_pos,
            end_pos,
        ) = self._retrieve_context_and_answers(question)
        prepared_inputs = [self.prefix + context for context in contexts]
        features = self.qg_tokenizer(
            prepared_inputs,
            max_length=512,
            padding="max_length",
            truncation=True,
            return_tensors="pt",
        ).to(self.device)
        # print(features)
        outputs = self.qg_model.generate(
            **features, max_length=128, num_beams=NUM_BEAMS, early_stopping=True
        )
        predicted_questions = self.qg_tokenizer.batch_decode(
            outputs, skip_special_tokens=True
        )
        return (
            predicted_questions,
            predicted_answers,
            retrieved_blocks,
            start_pos,
            end_pos,
        )


def save_to_disk(data, file_name):
    with jsonlines.open(file_name, "a") as writer:
        for example in tqdm(data, total=len(data), desc="Saving samples ... "):
            # print(example)
            for idx, (question, context, answer, start_pos, end_pos) in enumerate(
                zip(
                    example["predicted_questions"],
                    example["retrieved_contexts"],
                    example["predicted_answers"],
                    example["start_positions"],
                    example["end_positions"],
                )
            ):
                # print(question)
                # print(question, context, answer, start_pos, end_pos)
                writer.write(
                    {
                        "id": example["id"] + "_" + str(idx),
                        "question": example["question"],
                        "context": example["context"],
                        "answers": example["answers"],
                        "predicted_question": question,
                        "retrieved_context": context.decode("utf-8"),
                        # "original_retrieved_context": context,
                        "predicted_answer": {
                            "text": answer,
                            "answer_start": int(start_pos),
                            "answer_end": int(end_pos),
                        },
                    },
                )


if __name__ == "__main__":
    data = dataloader.PreprocessData(
        "squad", "plain_text", save_data=False, save_path="../../../"
    )

    prefix = "generate question: "
    cg = CounterfactualGeneration(beam_size=15, prefix=prefix)

    c = 0
    cache_file_name = (
        BASE_PATH
        + "src/data/squad/rag_counterfactuals_with_answers_t5_3b_squad_qg_06082023"
    )
    processed_instances = OrderedDict()
    outputs = []
    train_data, val_data = data.processed_train_val_set()
    chunk_size = 1000
    train_len = len(train_data)
    num_chunks = int(train_len / chunk_size) + 1
    for ex in tqdm(train_data):
        # ex = remove_white_space(ex)
        try:
            c += 1
            # if c<=87000:
            #     continue

            (
                predicted_questions,
                predicted_answers,
                retrieved_blocks,
                start_pos,
                end_pos,
            ) = cg.generate_question(ex["question"])
            # processed_instances[ex["id"]] = data
            outputs.append(
                {
                    "id": ex["id"],
                    "question": ex["question"],
                    "context": ex["context"],
                    "answers": ex["answers"],
                    "predicted_questions": predicted_questions,
                    "retrieved_contexts": retrieved_blocks,
                    "predicted_answers": predicted_answers,
                    "start_positions": list(start_pos),
                    "end_positions": list(end_pos),
                }
            )

            if c % 1000 == 0:
                save_to_disk(outputs, file_name=cache_file_name + ".jsonl")
                outputs = []
                num_chunks -= 1
            # save last chunk
            # if num_chunks == 1 and c == train_len:
            #     save_to_disk(outputs, file_name=cache_file_name + ".jsonl")
            #     outputs = []

        except Exception as e:
            print(f"Unable to get counterfactuals: {e}")
            print(ex)

    save_to_disk(outputs, file_name=cache_file_name + ".jsonl")
    # if c == 1:
    #     break
    # print(processed_instances)

    # data = dataloader.get_dev_examples(BASE_PATH + "src/data", "dev_hotpot.json")
    # data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/SearchQA-dev.jsonl")
    #
    # def remove_white_space(example):
    #     example["question_text"] = ' '.join(example["question_text"].split())
    #     example["context_text"] = ' '.join(example["context_text"].split())
    #     return example
    #
    # c = 0
    # cache_file_name = BASE_PATH + "src/data/search_qa/rag_counterfactuals"
    # processed_instances = OrderedDict()
    # outputs = []
    # for ex in tqdm(data):
    #     ex = remove_white_space(ex)
    #     try:
    #         # if ex["id"] == "56e1239acd28a01900c67641":
    #         #     print(ex)
    #         predicted_questions, predicted_answers, retrieved_blocks = cg.generate_question(ex["question_text"])
    #         # processed_instances[ex["id"]] = data
    #         outputs.append({
    #             "id": ex["qas_id"],
    #             "question": ex["question_text"],
    #             "context": ex["context_text"],
    #             "answers": ex["answer_text"],
    #             "predicted_questions": predicted_questions,
    #             "retrieved_contexts": retrieved_blocks,
    #             "predicted_answers": predicted_answers,
    #         })
    #         # print(tag_info)
    #         c += 1
    #         if c % 1000 == 0:
    #             save_to_disk(outputs, file_name=cache_file_name + ".jsonl")
    #             outputs = []
    #     except Exception as e:
    #         print(f"Unable to get counterfactuals: {e}")
    #         print(ex)

src/cf_generation/baseline_generation/pipeline.py
UKPLab/CATfOOD
import json
import os

import numpy as np
from tqdm import tqdm

SEED = 42
PROCESS_TRAIN = os.environ.pop("PROCESS_TRAIN", "false")
CATEGORY_MAPPING = {"null": 0, "short": 1, "long": 2, "yes": 3, "no": 4}


def _get_single_answer(example):
    def choose_first(answer, is_long_answer=False):
        # print(answer)
        if is_long_answer:
            assert isinstance(answer, dict)
            answer = [answer]
        else:
            assert isinstance(answer, list)
        if len(answer) == 1:
            answer = answer[0]
            return {k: [answer[k]][0] for k in answer} if is_long_answer else answer
        for a in answer:
            if is_long_answer:
                a = {k: [a[k]] for k in a}
            if len(answer) > 1:
                break
        return a

    # print(example)
    answer = {"id": example["example_id"]}
    annotation = example["annotations"]
    # print(annotation)
    if isinstance(annotation, list):
        annotation = annotation[0]
    yes_no_answer = annotation["yes_no_answer"]
    if yes_no_answer != "NONE":
        answer["category"] = ["yes"] if yes_no_answer == "YES" else ["no"]
        answer["start_token"] = answer["end_token"] = []
        answer["start_byte"] = answer["end_byte"] = []
        answer["text"] = ["<cls>"]
    elif annotation["short_answers"]:
        answer["category"] = ["short"]
        # print(annotation["short_answers"])
        out = choose_first(annotation["short_answers"])
        out["text"] = []
        answer.update(out)
        # print(out["start_token"])
    else:
        # answer will be long if short is not available
        answer["category"] = ["long"]
        out = choose_first(annotation["long_answer"], is_long_answer=True)
        out["text"] = []
        answer.update(out)
    # print(answer)

    # disregard some samples
    if answer["start_token"] == answer["end_token"]:
        answer["remove_it"] = True
    else:
        answer["remove_it"] = False

    # cols = ["start_token", "end_token", "start_byte", "end_byte", "text"]
    # if not all([isinstance(answer[k], list or int) for k in cols]):
    #     raise ValueError("Issue in ID", example["example_id"])

    return answer


def get_context_and_ans(example, assertion=False):
    """Gives new context after removing <html> & new answer tokens as per new context"""
    answer = _get_single_answer(example)
    # bytes are of no use
    del answer["start_byte"]
    del answer["end_byte"]

    # handle yes_no answers explicitly
    if answer["category"][0] in ["yes", "no"]:  # category is list with one element
        doc = example["document_tokens"]
        context = []
        for i in range(len(doc)):
            if not doc[i]["html_token"]:
                context.append(doc[i]["token"])
        return {
            "context": " ".join(context),
            "answer": {
                "start_token": -100,  # ignore index in cross-entropy
                "end_token": -100,  # ignore index in cross-entropy
                "category": answer["category"],
                "span": answer["category"],  # extra
            },
        }

    # later, help in removing all no answers
    if answer["start_token"] == [-1]:
        return {
            "context": "None",
            "answer": {
                "start_token": -1,
                "end_token": -1,
                "category": "null",
                "span": "None",  # extra
            },
        }

    # handling normal samples
    # cols = ["start_token", "end_token"]
    # answer.update(
    #     {k: answer[k][0] if len(answer[k]) > 0 else answer[k] for k in cols}
    # )  # e.g. [10] == 10

    doc = example["document_tokens"]
    # if len(doc) > 1:
    #     doc = doc[0]
    start_token = answer["start_token"]
    end_token = answer["end_token"]

    context = []
    for i in range(len(doc)):
        if not doc[i]["html_token"]:
            context.append(doc[i]["token"])
        else:
            if answer["start_token"] > i:
                start_token -= 1
            if answer["end_token"] > i:
                end_token -= 1
    new = " ".join(context[start_token:end_token])
    # print(new)

    # checking above code
    if assertion:
        """checking if above code is working as expected for all the samples"""
        is_html = doc["html_token"][answer["start_token"] : answer["end_token"]]
        old = doc["token"][answer["start_token"] : answer["end_token"]]
        old = " ".join([old[i] for i in range(len(old)) if not is_html[i]])
        if new != old:
            print("ID:", example["id"])
            print("New:", new, end="\n")
            print("Old:", old, end="\n\n")

    return {
        "context": " ".join(context),
        "answer": {
            "start_token": start_token,
            "end_token": end_token - 1,  # this makes it inclusive
            "category": answer["category"],  # either long or short
            "span": new,  # extra
        },
    }


def get_contexts_and_ans(example, assertion=True):
    # overlap will be of doc_stride - q_len

    out = get_context_and_ans(example, assertion=assertion)
    out["question"] = example["question_text"]
    out["example_id"] = example["example_id"]
    answer = out["answer"]

    # later, removing these samples
    if answer["start_token"] == -1:
        return {
            "example_id": example["example_id"],
            "question": example["question_text"],
            "context": "",
            "answer": {
                "start_token": -1,
                "end_token": -1,
                "category": ["null"],
                "span": "",
            },
        }

    return out


def prepare_inputs(example, assertion=False):
    example = get_contexts_and_ans(example, assertion=assertion,)

    return example


def save_to_disk(hf_data, file_name):
    with jsonlines.open(file_name, "a") as writer:
        for example in tqdm(hf_data, total=len(hf_data), desc="Saving samples ... "):
            answer = example["answer"]
            if answer["start_token"] == -1 and answer["end_token"] == -1:
                continue  # leave waste samples with no answer
            if answer["category"][0] == "null" and np.random.rand() < 0.6:
                continue  # removing 50 % samples
            writer.write(
                {
                    "id": example["example_id"],
                    "question": example["question"] + "?",
                    "context": example["context"],
                    "answer": {
                        "start_token": answer["start_token"],
                        "end_token": answer["end_token"],
                        "category": CATEGORY_MAPPING[answer["category"][0]],
                        "text": answer["span"],
                    },
                }
            )


if __name__ == "__main__":
    """Running area"""
    from pathlib import Path
    import jsonlines
    import os
    import gzip

    path = Path("/home/sachdeva/Downloads/v1.0/train")
    files = os.listdir(path)
    dataset_files = [path / str(f) for f in files]
    # dataset_file = dataset_files[0]

    c = 0
    data = []
    cache_file_name = "data/nq-train"
    for data_file in tqdm(dataset_files):
        with gzip.open(data_file, "rb") as f:
            for line in f:
                example = prepare_inputs(json.loads(line))
                data.append(example)
                c += 1
                if c % 1000 == 0:
                    save_to_disk(data, file_name=cache_file_name + ".jsonl")
                    data = []
                if c % 5000 == 0:
                    print(f"Samples processed: {c}")
            # if c == 20:
            #     break

    # print(example)
    # save_to_disk(data, file_name=cache_file_name + ".jsonl")
    # with jsonlines.open("/home/sachdeva/Downloads/v1.0/train/data.jsonl") as file:
    #     for line in file.iter():
    #         # print(line)
    #         # print("--------------------------")
    #         c += 1
    #         # prepare_inputs(line, tokenizer)
    #         # if c == 10:
    #         #     break
    # print(c)

src/cf_generation/baseline_generation/preprocess_nq.py
UKPLab/CATfOOD
from argparse import ArgumentParser

import torch

from model import get_openqa, add_additional_documents

from transformers.models.realm.modeling_realm import logger
from transformers.utils import logging

logger.setLevel(logging.INFO)
torch.set_printoptions(precision=8)


def get_arg_parser():
    parser = ArgumentParser()

    parser.add_argument("--question", type=str, required=True, help="Input question.")
    parser.add_argument(
        "--checkpoint_pretrained_name",
        type=str,
        default=r"converted_model",
        help="Checkpoint name or path.",
    )
    parser.add_argument(
        "--additional_documents_path",
        type=str,
        default=None,
        help="Additional document entries for retrieval. " "Must be .npy format.",
    )

    return parser


def main(args):
    openqa = get_openqa(args)
    tokenizer = openqa.retriever.tokenizer

    if args.additional_documents_path is not None:
        add_additional_documents(openqa, args.additional_documents_path)

    question_ids = tokenizer(args.question, return_tensors="pt").input_ids

    with torch.no_grad():
        outputs = openqa(input_ids=question_ids, return_dict=True,)

    predicted_answer = tokenizer.decode(outputs.predicted_answer_ids)

    print(f"Question: {args.question}\nAnswer: {predicted_answer}")

    return predicted_answer


if __name__ == "__main__":
    parser = get_arg_parser()
    args = parser.parse_args()
    main(args)

src/cf_generation/baseline_generation/realm.py
UKPLab/CATfOOD
"""
Fine-tuning the library models for question answering.
"""
# You can also adapt this script on your own question answering task. Pointers for this are left as comments.

import logging
import os
import sys
import pprint
from collections import OrderedDict
from dataclasses import dataclass, field
from typing import Optional

from datasets import load_dataset, load_metric

import transformers
from trainer_qa import QuestionAnsweringTrainer
from transformers import (
    AutoConfig,
    AutoModelForQuestionAnswering,
    AutoTokenizer,
    DataCollatorWithPadding,
    EvalPrediction,
    HfArgumentParser,
    PreTrainedTokenizerFast,
    TrainingArguments,
    default_data_collator,
    set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
from knockknock import slack_sender

from utils_qa import postprocess_qa_predictions
from training_args_seq2seq import Seq2SeqTrainingArguments

logger = logging.getLogger(__name__)


@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
    """

    model_name_or_path: str = field(
        metadata={
            "help": "Path to pretrained model or model identifier from huggingface.co/models"
        }
    )
    config_name: Optional[str] = field(
        default=None,
        metadata={
            "help": "Pretrained config name or path if not the same as model_name"
        },
    )
    tokenizer_name: Optional[str] = field(
        default=None,
        metadata={
            "help": "Pretrained tokenizer name or path if not the same as model_name"
        },
    )
    cache_dir: Optional[str] = field(
        default=None,
        metadata={
            "help": "Path to directory to store the pretrained models downloaded from huggingface.co"
        },
    )
    model_revision: str = field(
        default="main",
        metadata={
            "help": "The specific model version to use (can be a branch name, tag name or commit id)."
        },
    )
    use_auth_token: bool = field(
        default=False,
        metadata={
            "help": "Will use the token generated when running `transformers-cli login` (necessary to use this script "
            "with private models)."
        },
    )


@dataclass
class DataTrainingArguments:
    """
    Arguments pertaining to what data we are going to input our model for training and eval.
    """

    dataset_name: Optional[str] = field(
        default=None,
        metadata={"help": "The name of the dataset to use (via the datasets library)."},
    )
    dataset_config_name: Optional[str] = field(
        default=None,
        metadata={
            "help": "The configuration name of the dataset to use (via the datasets library)."
        },
    )
    train_file: Optional[str] = field(
        default=None, metadata={"help": "The input training data file (a text file)."}
    )
    validation_file: Optional[str] = field(
        default=None,
        metadata={
            "help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."
        },
    )
    test_file: Optional[str] = field(
        default=None,
        metadata={
            "help": "An optional input test data file to evaluate the perplexity on (a text file)."
        },
    )
    overwrite_cache: bool = field(
        default=False,
        metadata={"help": "Overwrite the cached training and evaluation sets"},
    )
    preprocessing_num_workers: Optional[int] = field(
        default=None,
        metadata={"help": "The number of processes to use for the preprocessing."},
    )
    max_source_length: int = field(
        default=384,
        metadata={
            "help": "The maximum total input sequence length after tokenization. Sequences longer "
            "than this will be truncated, sequences shorter will be padded."
        },
    )
    pad_to_max_length: bool = field(
        default=True,
        metadata={
            "help": "Whether to pad all samples to `max_seq_length`. "
            "If False, will pad the samples dynamically when batching to the maximum length in the batch (which can "
            "be faster on GPU but will be slower on TPU)."
        },
    )
    max_train_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
            "value if set."
        },
    )
    max_eval_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
            "value if set."
        },
    )
    max_predict_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": "For debugging purposes or quicker training, truncate the number of prediction examples to this "
            "value if set."
        },
    )
    num_beams: Optional[int] = field(
        default=None,
        metadata={
            "help": "Number of beams to use for evaluation. This argument will be passed to ``model.generate``, "
            "which is used during ``evaluate`` and ``predict``."
        },
    )
    version_2_with_negative: bool = field(
        default=False,
        metadata={"help": "If true, some of the examples do not have an answer."},
    )
    null_score_diff_threshold: float = field(
        default=0.0,
        metadata={
            "help": "The threshold used to select the null answer: if the best answer has a score that is less than "
            "the score of the null answer minus this threshold, the null answer is selected for this example. "
            "Only useful when `version_2_with_negative=True`."
        },
    )
    doc_stride: int = field(
        default=128,
        metadata={
            "help": "When splitting up a long document into chunks, how much stride to take between chunks."
        },
    )
    n_best_size: int = field(
        default=20,
        metadata={
            "help": "The total number of n-best predictions to generate when looking for an answer."
        },
    )
    max_target_length: int = field(
        default=30,
        metadata={
            "help": "The maximum length of an answer that can be generated. This is needed because the start "
            "and end predictions are not conditioned on one another."
        },
    )
    val_max_target_length: Optional[int] = field(
        default=30,
        metadata={
            "help": "The maximum total sequence length for validation target text after tokenization. Sequences longer "
            "than this will be truncated, sequences shorter will be padded. Will default to `max_target_length`."
            "This argument is also used to override the ``max_length`` param of ``model.generate``, which is used "
            "during ``evaluate`` and ``predict``."
        },
    )


dummy_dir = "/home/ubuntu/sagemaker_test/"


def decorate_condition(webhook_url, channel, decorator):
    return decorator if webhook_url is not None and channel is not None else lambda x: x


# monitor training via huggingface knockknock
webhook_url = (
    os.environ["SM_HP_SLACK_WEBHOOK_URL"]
    if "SM_HP_SLACK_WEBHOOK_URL" in os.environ
    else None
)
channel = (
    os.environ["SM_HP_SLACK_CHANNEL"] if "SM_HP_SLACK_CHANNEL" in os.environ else None
)


@decorate_condition(
    webhook_url, channel, slack_sender(webhook_url=webhook_url, channel=channel)
)
def main():
    # See all possible arguments in src/transformers/training_args.py
    # or by passing the --help flag to this script.
    # We now keep distinct sets of args, for a cleaner separation of concerns.

    parser = HfArgumentParser(
        (ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments)
    )
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, training_args = parser.parse_json_file(
            json_file=os.path.abspath(sys.argv[1])
        )
    else:
        (
            model_args,
            data_args,
            sagemaker_args,
            training_args,
        ) = parser.parse_args_into_dataclasses()

    # checks for input data
    # TODO: support for .txt files
    if not data_args.dataset_name:
        raise ValueError("Need either a dataset name or a training/validation file.")
    # if not data_args.val_max_target_length:
    #     data_args.val_max_target_length = data_args.max_target_length

    # metrics logger initialization
    if "wandb" in training_args.report_to:
        import wandb

        wandb.init(project="sm-qa")

    # Detecting last checkpoint.
    last_checkpoint = None
    if (
        os.path.isdir(training_args.output_dir)
        and training_args.do_train
        and not training_args.overwrite_output_dir
    ):
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif (
            last_checkpoint is not None and training_args.resume_from_checkpoint is None
        ):
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )
    logger.setLevel(logging.INFO if training_args.should_log else logging.WARN)

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
    )

    if training_args.should_log:
        transformers.utils.logging.set_verbosity_info()
        transformers.utils.logging.enable_default_handler()
        transformers.utils.logging.enable_explicit_format()
    logger.info(f"Training/evaluation parameters {training_args}")

    # Set seed before initializing model.
    set_seed(training_args.seed)

    # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
    # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
    # (the dataset will be downloaded automatically from the datasets Hub).
    #
    # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
    # 'text' is found. You can easily tweak this behavior (see below).
    #
    # In distributed training, the load_dataset function guarantee that only one local process can concurrently
    # download the dataset.
    if data_args.dataset_name:
        # Downloading and loading a dataset from the hub.
        datasets = load_dataset(data_args.dataset_name, data_args.dataset_config_name)

    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
    # https://huggingface.co/docs/datasets/loading_datasets.html.

    # Load pretrained model and tokenizer
    #
    # Distributed training:
    # The .from_pretrained methods guarantee that only one local process can concurrently
    # download model & vocab.
    config = AutoConfig.from_pretrained(
        model_args.config_name
        if model_args.config_name
        else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.tokenizer_name
        if model_args.tokenizer_name
        else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        use_fast=True,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    )
    model = AutoModelForQuestionAnswering.from_pretrained(
        model_args.model_name_or_path,
        from_tf=bool(".ckpt" in model_args.model_name_or_path),
        config=config,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    )

    # Tokenizer check: this script requires a fast tokenizer.
    if not isinstance(tokenizer, PreTrainedTokenizerFast):
        raise ValueError(
            "This example script only works for models that have a fast tokenizer. Checkout the big table of models "
            "at https://huggingface.co/transformers/index.html#supported-frameworks to find the model types that meet this "
            "requirement"
        )

    # Preprocessing the datasets.
    # Preprocessing is slighlty different for training and evaluation.
    if training_args.do_train:
        column_names = datasets["train"].column_names
    elif training_args.do_eval:
        column_names = datasets["validation"].column_names
    else:
        column_names = datasets["test"].column_names
    question_column_name = "question" if "question" in column_names else column_names[0]
    context_column_name = "context" if "context" in column_names else column_names[1]
    answer_column_name = "answers" if "answers" in column_names else column_names[2]

    # Padding side determines if we do (question|context) or (context|question).
    pad_on_right = tokenizer.padding_side == "right"

    if data_args.max_source_length > tokenizer.model_max_length:
        logger.warning(
            f"The max_seq_length passed ({data_args.max_source_length}) is larger than the maximum length for the"
            f"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}."
        )
    max_seq_length = min(data_args.max_source_length, tokenizer.model_max_length)

    # Training preprocessing
    def prepare_train_features(examples):
        # Tokenize our examples with truncation and maybe padding, but keep the overflows using a stride. This results
        # in one example possible giving several features when a context is long, each of those features having a
        # context that overlaps a bit the context of the previous feature.
        tokenized_examples = tokenizer(
            examples[question_column_name if pad_on_right else context_column_name],
            examples[context_column_name if pad_on_right else question_column_name],
            truncation="only_second" if pad_on_right else "only_first",
            max_length=max_seq_length,
            stride=data_args.doc_stride,
            return_overflowing_tokens=True,  # in long contexts
            return_offsets_mapping=True,
            padding="max_length" if data_args.pad_to_max_length else False,
        )

        # Since one example might give us several features if it has a long context, we need a map from a feature to
        # its corresponding example. This key gives us just that.
        sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")
        # The offset mappings will give us a map from token to character position in the original context. This will
        # help us compute the start_positions and end_positions.
        offset_mapping = tokenized_examples.pop("offset_mapping")

        # Let's label those examples!
        tokenized_examples["start_positions"] = []
        tokenized_examples["end_positions"] = []

        for i, offsets in enumerate(offset_mapping):
            # We will label impossible answers with the index of the CLS token.
            input_ids = tokenized_examples["input_ids"][i]
            cls_index = input_ids.index(tokenizer.cls_token_id)

            # Grab the sequence corresponding to that example (to know what is the context and what is the question).
            sequence_ids = tokenized_examples.sequence_ids(i)

            # One example can give several spans, this is the index of the example containing this span of text.
            sample_index = sample_mapping[i]
            answers = examples[answer_column_name][sample_index]
            # If no answers are given, set the cls_index as answer.
            if len(answers["answer_start"]) == 0:
                tokenized_examples["start_positions"].append(cls_index)
                tokenized_examples["end_positions"].append(cls_index)
            else:
                # Start/end character index of the answer in the text.
                start_char = answers["answer_start"][0]
                end_char = start_char + len(answers["text"][0])

                # Start token index of the current span in the text.
                token_start_index = 0
                while sequence_ids[token_start_index] != (1 if pad_on_right else 0):
                    token_start_index += 1

                # End token index of the current span in the text.
                token_end_index = len(input_ids) - 1
                while sequence_ids[token_end_index] != (1 if pad_on_right else 0):
                    token_end_index -= 1

                # Detect if the answer is out of the span (in which case this feature is labeled with the CLS index).
                if not (
                    offsets[token_start_index][0] <= start_char
                    and offsets[token_end_index][1] >= end_char
                ):
                    tokenized_examples["start_positions"].append(cls_index)
                    tokenized_examples["end_positions"].append(cls_index)
                else:
                    # Otherwise move the token_start_index and token_end_index to the two ends of the answer.
                    # Note: we could go after the last offset if the answer is the last word (edge case).
                    while (
                        token_start_index < len(offsets)
                        and offsets[token_start_index][0] <= start_char
                    ):
                        token_start_index += 1
                    tokenized_examples["start_positions"].append(token_start_index - 1)
                    while offsets[token_end_index][1] >= end_char:
                        token_end_index -= 1
                    tokenized_examples["end_positions"].append(token_end_index + 1)

        return tokenized_examples

    if training_args.do_train:
        if "train" not in datasets:
            raise ValueError("--do_train requires a train dataset")
        train_dataset = datasets["train"]
        if data_args.max_train_samples is not None:
            # We will select sample from whole data if agument is specified
            train_dataset = train_dataset.select(range(data_args.max_train_samples))
        # Create train feature from dataset
        train_dataset = train_dataset.map(
            prepare_train_features,
            batched=True,
            num_proc=data_args.preprocessing_num_workers,
            remove_columns=column_names,
            load_from_cache_file=not data_args.overwrite_cache,
        )
        if data_args.max_train_samples is not None:
            # Number of samples might increase during Feature Creation, We select only specified max samples
            train_dataset = train_dataset.select(range(data_args.max_train_samples))

    # Validation preprocessing
    def prepare_validation_features(examples):
        # Tokenize our examples with truncation and maybe padding, but keep the overflows using a stride. This results
        # in one example possible giving several features when a context is long, each of those features having a
        # context that overlaps a bit the context of the previous feature.
        tokenized_examples = tokenizer(
            examples[question_column_name if pad_on_right else context_column_name],
            examples[context_column_name if pad_on_right else question_column_name],
            truncation="only_second" if pad_on_right else "only_first",
            max_length=max_seq_length,
            stride=data_args.doc_stride,
            return_overflowing_tokens=True,
            return_offsets_mapping=True,
            padding="max_length" if data_args.pad_to_max_length else False,
        )

        # Since one example might give us several features if it has a long context, we need a map from a feature to
        # its corresponding example. This key gives us just that.
        sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")

        # For evaluation, we will need to convert our predictions to substrings of the context, so we keep the
        # corresponding example_id and we will store the offset mappings.
        tokenized_examples["example_id"] = []

        for i in range(len(tokenized_examples["input_ids"])):
            # Grab the sequence corresponding to that example (to know what is the context and what is the question).
            sequence_ids = tokenized_examples.sequence_ids(i)
            context_index = 1 if pad_on_right else 0

            # One example can give several spans, this is the index of the example containing this span of text.
            sample_index = sample_mapping[i]
            tokenized_examples["example_id"].append(examples["id"][sample_index])

            # Set to None the offset_mapping that are not part of the context so it's easy to determine if a token
            # position is part of the context or not.
            tokenized_examples["offset_mapping"][i] = [
                (o if sequence_ids[k] == context_index else None)
                for k, o in enumerate(tokenized_examples["offset_mapping"][i])
            ]

        return tokenized_examples

    if training_args.do_eval:
        if "validation" not in datasets:
            raise ValueError("--do_eval requires a validation dataset")
        eval_examples = datasets["validation"]
        if data_args.max_eval_samples is not None:
            # We will select sample from whole data
            eval_examples = eval_examples.select(range(data_args.max_eval_samples))
        # Validation Feature Creation
        eval_dataset = eval_examples.map(
            prepare_validation_features,
            batched=True,
            num_proc=data_args.preprocessing_num_workers,
            remove_columns=column_names,
            load_from_cache_file=not data_args.overwrite_cache,
        )
        if data_args.max_eval_samples is not None:
            # During Feature creation dataset samples might increase, we will select required samples again
            eval_dataset = eval_dataset.select(range(data_args.max_eval_samples))

    if training_args.do_predict:
        if "test" not in datasets:
            raise ValueError("--do_predict requires a test dataset")
        predict_examples = datasets["test"]
        if data_args.max_predict_samples is not None:
            # We will select sample from whole data
            predict_examples = predict_examples.select(
                range(data_args.max_predict_samples)
            )
        # Predict Feature Creation
        predict_dataset = predict_examples.map(
            prepare_validation_features,
            batched=True,
            num_proc=data_args.preprocessing_num_workers,
            remove_columns=column_names,
            load_from_cache_file=not data_args.overwrite_cache,
        )
        if data_args.max_predict_samples is not None:
            # During Feature creation dataset samples might increase, we will select required samples again
            predict_dataset = predict_dataset.select(
                range(data_args.max_predict_samples)
            )

    # Data collator
    # We have already padded to max length if the corresponding flag is True, otherwise we need to pad in the data
    # collator.
    data_collator = (
        default_data_collator
        if data_args.pad_to_max_length
        else DataCollatorWithPadding(
            tokenizer, pad_to_multiple_of=8 if training_args.fp16 else None
        )
    )

    # Post-processing:
    def post_processing_function(examples, features, predictions, stage="eval"):
        # Post-processing: we match the start logits and end logits to answers in the original context.
        predictions = postprocess_qa_predictions(
            examples=examples,
            features=features,
            predictions=predictions,
            version_2_with_negative=data_args.version_2_with_negative,
            n_best_size=data_args.n_best_size,
            max_answer_length=data_args.max_target_length,
            null_score_diff_threshold=data_args.null_score_diff_threshold,
            output_dir=sagemaker_args.output_data_dir,
            is_world_process_zero=trainer.is_world_process_zero(),
            prefix=stage,
        )
        # Format the result to the format the metric expects.
        if data_args.version_2_with_negative:
            formatted_predictions = [
                {"id": k, "prediction_text": v, "no_answer_probability": 0.0}
                for k, v in predictions.items()
            ]
        else:
            formatted_predictions = [
                {"id": k, "prediction_text": v} for k, v in predictions.items()
            ]

        references = [
            {"id": ex["id"], "answers": ex[answer_column_name]} for ex in examples
        ]
        return EvalPrediction(predictions=formatted_predictions, label_ids=references)

    metric = load_metric("squad_v2" if data_args.version_2_with_negative else "squad")

    def compute_metrics(p: EvalPrediction):
        return metric.compute(predictions=p.predictions, references=p.label_ids)

    # Initialize our Trainer
    trainer = QuestionAnsweringTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset if training_args.do_train else None,
        eval_dataset=eval_dataset if training_args.do_eval else None,
        eval_examples=eval_examples if training_args.do_eval else None,
        tokenizer=tokenizer,
        data_collator=data_collator,
        post_process_function=post_processing_function,
        compute_metrics=compute_metrics,
    )

    # Training
    results = OrderedDict()
    if training_args.do_train:
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            logger.info("***** continue training *****")
            checkpoint = training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
            logger.info("***** continue training *****")
            checkpoint = last_checkpoint
        train_result = trainer.train(resume_from_checkpoint=checkpoint)

        trainer.save_model(
            sagemaker_args.model_dir
        )  # Saves the model and tokenizer to S3

        metrics = train_result.metrics
        max_train_samples = (
            data_args.max_train_samples
            if data_args.max_train_samples is not None
            else len(train_dataset)
        )
        metrics["train_samples"] = min(max_train_samples, len(train_dataset))
        results["train"] = {
            k: v for k, v in metrics.items() if not k.__contains__("mem")
        }

    # Evaluation
    if training_args.do_eval:
        logger.info("*** Evaluate ***")

        metrics = trainer.evaluate()
        max_eval_samples = (
            data_args.max_eval_samples
            if data_args.max_eval_samples is not None
            else len(eval_dataset)
        )
        metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset))
        results["eval"] = {
            k: v for k, v in metrics.items() if not k.__contains__("mem")
        }

    if training_args.do_predict:
        logger.info("*** Predict ***")

        predict_results = trainer.predict(
            predict_dataset,
            metric_key_prefix="predict",
            max_length=data_args.val_max_target_length,
            num_beams=data_args.num_beams,
        )
        metrics = predict_results.metrics
        max_predict_samples = (
            data_args.max_predict_samples
            if data_args.max_predict_samples is not None
            else len(predict_dataset)
        )
        metrics["predict_samples"] = min(max_predict_samples, len(predict_dataset))
        results["test"] = {
            k: v for k, v in metrics.items() if not k.__contains__("mem")
        }

        if trainer.is_world_process_zero():
            if training_args.predict_with_generate:
                predictions = tokenizer.batch_decode(
                    predict_results.predictions,
                    skip_special_tokens=True,
                    clean_up_tokenization_spaces=True,
                )
                predictions = [pred.strip() for pred in predictions]
                output_prediction_file = os.path.join(
                    sagemaker_args.output_data_dir, "generated_predictions.txt"
                )
                with open(output_prediction_file, "w") as writer:
                    writer.write("\n".join(predictions))

    # write results to file
    output_results_file = os.path.join(sagemaker_args.output_data_dir, "results.txt")
    with open(output_results_file, "w") as writer:
        if training_args.do_train:
            writer.write("*" * 20 + "TRAIN METRICS" + "*" * 20 + "\n")
            writer.write(
                "\n"
                + "".join(
                    ["{} = {}\n".format(k, v) for k, v in results["train"].items()]
                )
            )
        if training_args.do_eval:
            writer.write("\n" + "*" * 20 + "VALIDATION METRICS" + "*" * 20 + "\n")
            writer.write(
                "\n"
                + "".join(
                    ["{} = {}\n".format(k, v) for k, v in results["eval"].items()]
                )
            )
        if training_args.do_predict:
            writer.write("\n" + "*" * 20 + "TEST METRICS" + "*" * 20 + "\n")
            writer.write(
                "\n"
                + "".join(
                    ["{} = {}\n".format(k, v) for k, v in results["test"].items()]
                )
            )

    return pprint.pformat(results)


def _mp_fn(index):
    # For xla_spawn (TPUs)
    main()


if __name__ == "__main__":
    main()

src/cf_generation/baseline_generation/run_qa.py
UKPLab/CATfOOD
# coding=utf-8
# Copyright 2018 The Google AI Language Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
r"""Utilities to simplify the csachdevaical NQ data.
The csachdevaical NQ data contains the HTML of each Wikipedia page along with a
sequence of tokens on that page, each of which is indexed into the HTML.
Many users will not want to use the HTML at all, and this file provides
utilities to extract only the text into a new record of the form:
  {
    "example_id": 3902,
    "document_url": "http://wikipedia.org/en/strings"
    "question_text": "what is a string",
    "document_text": "<P> A string is a list of characters in order . </P>",
    "annotations": [{
      "long_answer": { "start_token": 0, "end_token": 12 },
      "short_answers": [{ "start_token": 5, "end_token": 8 }],
      "yes_no_answer": "NONE",
    }],
    "long_answer_candidates": [
      {"start_token": 0, "end_token": 12, "top_level": True}
    ]
  }
which leads to a much smaller training set (4.4Gb instead of 41Gb).
In this representation, the [start, end) indices are into the blank separated
sequence of tokens. So, answer spans can be extracted using the following
snippet:
  " ".join(example["document_text"].split(" ")[`start_token`:`end_token`]).
WARNING: Use `split(" ")` instead of `split()` to avoid complications from
  characters such as `\u180e` which may or may not be recognized as a whitespace
  character depending on your python version.
To avoid complications at test time, we do not provide a simplified version
of the development data, and there is no simplified version of the hidden test
set. If you rely on the simplified data, then you must call the
`simplify_nq_example` function below on every example that is passed in at test
time.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import re
import jsonlines


def get_nq_tokens(simplified_nq_example):
    """Returns list of blank separated tokens."""

    if "document_text" not in simplified_nq_example:
        raise ValueError(
            "`get_nq_tokens` should be called on a simplified NQ"
            "example that contains the `document_text` field."
        )

    return simplified_nq_example["document_text"].split(" ")


def simplify_nq_example(nq_example):
    r"""Returns dictionary with blank separated tokens in `document_text` field.
  Removes byte offsets from annotations, and removes `document_html` and
  `document_tokens` fields. All annotations in the ouput are represented as
  [start_token, end_token) offsets into the blank separated tokens in the
  `document_text` field.
  WARNING: Tokens are separated by a single blank character. Do not split on
    arbitrary whitespace since different implementations have different
    treatments of some unicode characters such as \u180e.
  Args:
    nq_example: Dictionary containing original NQ example fields.
  Returns:
    Dictionary containing `document_text` field, not containing
    `document_tokens` or `document_html`, and with all annotations represented
    as [`start_token`, `end_token`) offsets into the space separated sequence.
  """

    def _clean_token(token):
        """Returns token in which blanks are replaced with underscores.
    HTML table cell openers may contain blanks if they span multiple columns.
    There are also a very few unicode characters that are prepended with blanks.
    Args:
      token: Dictionary representation of token in original NQ format.
    Returns:
      String token.
    """
        return re.sub(u" ", "_", token["token"])

    text = " ".join([_clean_token(t) for t in nq_example["document_tokens"]])

    def _remove_html_byte_offsets(span):
        if "start_byte" in span:
            del span["start_byte"]

        if "end_byte" in span:
            del span["end_byte"]

        return span

    def _clean_annotation(annotation):
        annotation["long_answer"] = _remove_html_byte_offsets(annotation["long_answer"])
        annotation["short_answers"] = [
            _remove_html_byte_offsets(sa) for sa in annotation["short_answers"]
        ]
        return annotation

    simplified_nq_example = {
        "question_text": nq_example["question_text"],
        "example_id": nq_example["example_id"],
        "document_url": nq_example["document_url"],
        "document_text": text,
        "long_answer_candidates": [
            _remove_html_byte_offsets(c) for c in nq_example["long_answer_candidates"]
        ],
        "annotations": [_clean_annotation(a) for a in nq_example["annotations"]],
    }

    if len(get_nq_tokens(simplified_nq_example)) != len(nq_example["document_tokens"]):
        raise ValueError("Incorrect number of tokens.")

    return simplified_nq_example


if __name__ == "__main__":
    with jsonlines.open("data/nq_dev_set.jsonl") as file:
        for line in file.iter():
            print(line)
            print("-------------------")
            print(simplify_nq_example(line))
            print("-------------------")
            print(get_nq_tokens(simplify_nq_example(line)))
            break

src/cf_generation/baseline_generation/simplified_nq.py
UKPLab/CATfOOD
import os

import numpy as np
import torch
import wandb
from src.calibration.baseline.dataloader import PreprocessData

from transformers import (
    T5ForConditionalGeneration,
    T5Tokenizer,
    Trainer,
    TrainingArguments,
    T5Config,
)

# BASE_PATH = "/home/sachdeva/projects/exp_calibration/src/rag/"
# CACHE_DIR = "/home/sachdeva/.cache"
BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/src/rag/"
CACHE_DIR = "/storage/ukp/work/sachdeva/.cache"

os.environ["TRANSFORMERS_CACHE"] = CACHE_DIR
os.environ["WANDB_PROJECT"] = "t5-squad-qg"
# wandb.init(dir=os.getenv("WANDB_DIR", BASE_PATH))


class T5ForQGeneration(T5ForConditionalGeneration):
    """T5ForConditionalGeneration for QA"""

    def __init__(self, config):
        super().__init__(config)

    def forward(
        self, input_ids=None, attention_mask=None, labels=None,
    ):
        outputs = super().forward(
            input_ids=input_ids, attention_mask=attention_mask, labels=labels
        )
        loss = outputs[0]

        return {"loss": loss}


class QGen:
    def __init__(
        self,
        dataset_name: str,
        dataset_config: str,
        cf_path: str,
        max_src_len: int,
        stride: int,
        do_train: bool,
        do_eval: bool,
        model_name: str,
        cache_dir: str = None,
        max_train_samples: int = None,
        max_val_samples: int = None,
        save_results: bool = False,
    ):
        self.dataset_name = dataset_name
        self.dataset_config = dataset_config
        self.cf_path = cf_path
        self.max_src_len = max_src_len
        self.stride = stride
        self.do_train = do_train
        self.do_eval = do_eval
        self.model_name = model_name
        self.cache_dir = cache_dir
        self.max_train_samples = max_train_samples
        self.max_val_samples = max_val_samples
        self.padding = "max_length"
        self.save_results = save_results
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        # load resources
        self._prepare_data()

    def _preprocess_data(self, example):
        question = example["question"].strip()
        context = example["context"]
        answer = example["answers"]
        answer_text = answer["text"][0]
        start_token = answer["answer_start"][0]
        end_token = start_token + len(answer_text)
        sep_token = ">>"

        input = f"{question} {sep_token} {context}"
        label = answer_text

        input_ids = self.tokenizer(input, max_length=512, truncation=True)["input_ids"]
        label_ids = self.tokenizer(label, max_length=256, padding=True, truncation=True)
        label_ids = np.array(label_ids["input_ids"])

        # print(tokenizer.pad_token_id)
        # print(label_ids)
        label_ids[label_ids == self.tokenizer.pad_token_id] = -100
        example["input_ids"] = input_ids
        example["labels"] = label_ids
        return example

    def _prepare_data(self):

        """
        Load model, tokenizer and data
        """

        config = T5Config.from_pretrained(self.model_name, cache_dir=self.cache_dir,)
        self.tokenizer = T5Tokenizer.from_pretrained(
            self.model_name, cache_dir=self.cache_dir, use_fast=True,
        )
        self.tokenizer.add_special_tokens({"additional_special_tokens": [">>"]})

        self.model = T5ForQGeneration.from_pretrained(
            self.model_name, config=config, cache_dir=self.cache_dir,
        )
        self.model.resize_token_embeddings(len(self.tokenizer))
        dataloader = PreprocessData(
            self.dataset_name,
            self.dataset_config,
            cf_path=None,
            save_data=False,
            save_path="",
        )
        self.train_set, self.val_set = dataloader.processed_train_val_set()

    def collate_fn(self, features, pad_id=0, threshold=512):
        def pad_elems(ls, pad_id, maxlen):
            while len(ls) < maxlen:
                ls.append(pad_id)
            return ls

        maxlen = max([len(x["input_ids"]) for x in features])
        # avoid attention_type switching
        if maxlen < threshold:
            maxlen = threshold

        # dynamic padding
        input_ids = [
            pad_elems(x["input_ids"], pad_id=pad_id, maxlen=maxlen) for x in features
        ]
        input_ids = torch.tensor(input_ids, dtype=torch.long)
        # print(input_ids)
        # print(len(input_ids[0]))

        # padding mask...
        attention_mask = input_ids.clone()
        attention_mask[attention_mask != pad_id] = 1
        attention_mask[attention_mask == pad_id] = 0
        # replace padding token id's of the labels by -100, so it's ignored by the loss
        labels = [pad_elems(x["labels"], pad_id=-100, maxlen=256) for x in features]
        labels = torch.tensor(labels, dtype=torch.long)
        # print(labels)
        # print(len(labels[0]))

        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "labels": labels,
        }

    def train(self):

        args = TrainingArguments(
            output_dir=OUTPUT_DIR,
            overwrite_output_dir=False,
            do_train=DO_TRAIN,
            do_eval=DO_EVAL,
            evaluation_strategy="epoch",
            # eval_steps=4000,
            per_device_train_batch_size=BATCH_SIZE,
            per_device_eval_batch_size=BATCH_SIZE,
            learning_rate=LEARNING_RATE,
            warmup_steps=100,
            lr_scheduler_type="linear",
            num_train_epochs=MAX_EPOCHS,
            logging_steps=10,
            # save_steps=250,
            save_total_limit=2,
            save_strategy="epoch",
            load_best_model_at_end=False,
            run_name=WANDB_RUN_NAME,
            disable_tqdm=False,
            report_to=["wandb"],
            remove_unused_columns=False,
            fp16=FP16,
            seed=SEED,
            label_names=["labels"],  # it's important to log eval_loss
        )
        print("Batch Size", args.train_batch_size)
        print("Parallel Mode", args.parallel_mode)

        train_data = self.train_set.map(self._preprocess_data)
        val_data = self.val_set.map(self._preprocess_data)

        trainer = Trainer(
            model=self.model,
            args=args,
            data_collator=self.collate_fn,
            train_dataset=train_data,
            eval_dataset=val_data,
        )
        try:
            checkpoint = None
            if RESUME_TRAINING is not None:
                checkpoint = RESUME_TRAINING
            trainer.train(resume_from_checkpoint=checkpoint)
            trainer.save_model(OUTPUT_DIR)
            self.tokenizer.save_pretrained(save_directory=OUTPUT_DIR)
        except KeyboardInterrupt:
            trainer.save_model("interrupted-squad-qg")
        wandb.finish()


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        description="Passing arguments for model, tokenizer, and dataset."
    )
    parser.add_argument(
        "--model_name",
        default="roberta-base",
        type=str,
        required=False,
        help="Specify the model to use.",
    )
    parser.add_argument(
        "--tokenizer",
        default="roberta-base",
        type=str,
        required=False,
        help="Specify the tokenizer to use.",
    )
    parser.add_argument(
        "--output_dir", type=str, required=True, help="Specify the output directory."
    )
    parser.add_argument(
        "--seed",
        default=42,
        type=int,
        required=False,
        help="Specify the seed for reproducibility.",
    )
    parser.add_argument(
        "--cf_path",
        default=None,
        type=str,
        required=False,
        help="Specify the path to counterfactuals.",
    )
    parser.add_argument(
        "--epochs", default=3, type=int, required=False, help="Epochs to train."
    )
    parser.add_argument(
        "--batch_size", default=4, type=int, required=False, help="Train batch size."
    )
    parser.add_argument(
        "--lr", default=2.0e-5, type=float, required=False, help="Learning rate."
    )
    parser.add_argument(
        "--warmup_ratio", default=0.06, type=float, required=False, help="Warmup ratio."
    )

    args = parser.parse_args()

    # hyperparameters
    LEARNING_RATE = args.lr
    WARMUP_RATIO = args.warmup_ratio
    MAX_EPOCHS = args.epochs
    BATCH_SIZE = args.batch_size
    OUTPUT_DIR = args.output_dir
    SEED = args.seed

    CHECKPOINT = None
    RESUME_TRAINING = None

    DO_TRAIN = True
    DO_EVAL = True
    FP16 = False

    # set to None when using all data
    MAX_TRAIN_SAMPLES = None
    MAX_VAL_SAMPLES = None

    WANDB_RUN_NAME = OUTPUT_DIR

    wandb.init(
        config={
            "lr": LEARNING_RATE,
            "warmup_ratio": WARMUP_RATIO,
            "epochs": MAX_EPOCHS,
            "model": args.model_name,
            "batch_size": BATCH_SIZE,
            "output_dir": OUTPUT_DIR,
        }
    )
    wandb.run.name = WANDB_RUN_NAME

    trainer = QGen(
        dataset_name="squad",
        dataset_config="plain_text",
        cf_path=args.cf_path,
        # "src/data/squad/counterfactual_data_gpt_neox_20b_v2_qg_pipeline_all_data_cleaned.jsonl",
        max_src_len=384,
        stride=128,
        do_train=DO_TRAIN,
        do_eval=DO_EVAL,
        model_name=args.model_name,
        max_train_samples=MAX_TRAIN_SAMPLES,
        max_val_samples=MAX_VAL_SAMPLES,
        save_results=True,
    )
    trainer.train()

src/cf_generation/baseline_generation/train_gen_qa.py
UKPLab/CATfOOD
import os

import numpy as np
import torch
import torch.nn as nn
import wandb
from datasets import load_dataset

from transformers import (
    T5ForConditionalGeneration,
    T5Tokenizer,
    Trainer,
    TrainingArguments,
    DataCollatorForSeq2Seq,
)

# BASE_PATH = "/home/sachdeva/projects/exp_calibration/src/rag/"
# CACHE_DIR = "/home/sachdeva/.cache"
BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/src/rag/"
CACHE_DIR = "/storage/ukp/work/sachdeva/.cache"

OUTPUT_DIR = "t5-3b-nq-short-qa-seed-42"

TRAIN_ON_SMALL = os.environ.pop("TRAIN_ON_SMALL", "false")
RESUME_TRAINING = None

SEED = 42
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)

GROUP_BY_LENGTH = True
LEARNING_RATE = 2 * 1.0e-5
WARMUP_STEPS = 100
DROPOUT = 0.1
MAX_EPOCHS = 3
FP16 = False
SCHEDULER = "linear"
MODEL_NAME = "t5-3b"

BATCH_SIZE = 4

DO_TRAIN = True
DO_EVAL = True
DO_PREDICT = True

os.environ["TRANSFORMERS_CACHE"] = CACHE_DIR
os.environ["WANDB_PROJECT"] = "t5-natural-questions-qa"
# wandb.init(dir=os.getenv("WANDB_DIR", BASE_PATH))

WANDB_RUN_NAME = OUTPUT_DIR

wandb.init(
    config={
        "seed": SEED,
        "lr": LEARNING_RATE,
        "warmup_steps": WARMUP_STEPS,
        "dropout": 0.1,
        "epochs": MAX_EPOCHS,
        "scheduler": SCHEDULER,
        "model": MODEL_NAME,
        "batch_size": BATCH_SIZE,
        "output_dir": OUTPUT_DIR,
    }
)
wandb.run.name = WANDB_RUN_NAME


def collate_fn(features, pad_id=0, threshold=640):
    def pad_elems(ls, pad_id, maxlen):
        while len(ls) < maxlen:
            ls.append(pad_id)
        return ls

    maxlen = max([len(x["input_ids"]) for x in features])
    # avoid attention_type switching
    if maxlen < threshold:
        maxlen = threshold

    # dynamic padding
    input_ids = [pad_elems(x["input_ids"], pad_id, maxlen) for x in features]
    input_ids = torch.tensor(input_ids, dtype=torch.long)

    # padding mask
    attention_mask = input_ids.clone()
    attention_mask[attention_mask != pad_id] = 1
    attention_mask[attention_mask == pad_id] = 0

    # replace padding token id's of the labels by -100, so it's ignored by the loss
    labels = [pad_elems(x["labels"], pad_id=-100, maxlen=256) for x in features]
    labels = torch.tensor(labels, dtype=torch.long)

    return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}


class T5ForNaturalQuestions(T5ForConditionalGeneration):
    """T5 with CLS Head over the top for predicting category"""

    def __init__(self, config):
        super().__init__(config)
        self.cls = nn.Linear(config.hidden_size, 5)

    def forward(
        self, input_ids, attention_mask=None, labels=None,
    ):

        outputs = super().forward(
            input_ids=input_ids, attention_mask=attention_mask, labels=labels
        )
        loss = outputs[0]

        return {"loss": loss}


if __name__ == "__main__":
    # "nq-training.jsonl" & "nq-validation.jsonl" are obtained from running `prepare_nq.py`
    tr_dataset = load_dataset(
        "json", data_files=BASE_PATH + "data/nq-train-tokenized-short-qa.jsonl"
    )["train"]
    val_dataset = load_dataset(
        "json", data_files=BASE_PATH + "data/nq-dev-tokenized-short-qa.jsonl"
    )["train"]
    # print(tr_dataset[0])

    if TRAIN_ON_SMALL == "true":
        # this will run for ~12 hrs on 2 K80 GPU
        np.random.seed(SEED)
        indices = np.random.randint(0, 86389, size=8)
        tr_dataset = tr_dataset.select(indices)
        np.random.seed(SEED)
        indices = np.random.randint(0, 1500, size=2)
        val_dataset = val_dataset.select(indices)

    tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR)
    tokenizer.add_special_tokens({"additional_special_tokens": [">>"]})
    model = T5ForNaturalQuestions.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR)
    model.resize_token_embeddings(len(tokenizer))

    args = TrainingArguments(
        output_dir=OUTPUT_DIR,
        overwrite_output_dir=False,
        do_train=DO_TRAIN,
        do_eval=DO_EVAL,
        evaluation_strategy="epoch",
        # eval_steps=4000,
        per_device_train_batch_size=BATCH_SIZE,
        per_device_eval_batch_size=BATCH_SIZE,
        learning_rate=LEARNING_RATE,
        warmup_steps=WARMUP_STEPS,
        lr_scheduler_type=SCHEDULER,
        num_train_epochs=MAX_EPOCHS,
        logging_steps=100,
        # save_steps=250,
        save_total_limit=1,
        save_strategy="epoch",
        load_best_model_at_end=False,
        run_name=WANDB_RUN_NAME,
        disable_tqdm=False,
        report_to=["wandb"],
        remove_unused_columns=False,
        fp16=FP16,
        seed=SEED,
        label_names=["labels"],  # it's important to log eval_loss
    )
    print("Batch Size", args.train_batch_size)
    print("Parallel Mode", args.parallel_mode)

    trainer = Trainer(
        model=model,
        args=args,
        data_collator=collate_fn,
        train_dataset=tr_dataset,
        eval_dataset=val_dataset,
    )
    try:
        trainer.train(resume_from_checkpoint=RESUME_TRAINING)
        trainer.save_model(OUTPUT_DIR)
        tokenizer.save_pretrained(save_directory=OUTPUT_DIR)
    except KeyboardInterrupt:
        trainer.save_model("interrupted-natural-questions-qa")
    wandb.finish()

src/cf_generation/baseline_generation/train_nq_qa.py
UKPLab/CATfOOD
import os

import numpy as np
import torch
import torch.nn as nn
import wandb
from datasets import load_dataset

from transformers import (
    T5ForConditionalGeneration,
    T5Tokenizer,
    Trainer,
    TrainingArguments,
    DataCollatorForSeq2Seq,
)

# BASE_PATH = "/home/sachdeva/projects/exp_calibration/src/rag/"
# CACHE_DIR = "/home/sachdeva/.cache"
BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/src/rag/"
CACHE_DIR = "/storage/ukp/work/sachdeva/.cache"

OUTPUT_DIR = "t5-3b-nq-short-qg-seed-42"

TRAIN_ON_SMALL = os.environ.pop("TRAIN_ON_SMALL", "false")
RESUME_TRAINING = None

SEED = 42

LEARNING_RATE = 2 * 1.0e-5
WARMUP_STEPS = 100
DROPOUT = 0.1
MAX_EPOCHS = 3
FP16 = False
SCHEDULER = "linear"
MODEL_NAME = "t5-3b"

BATCH_SIZE = 4

DO_TRAIN = True
DO_EVAL = True
DO_PREDICT = True

os.environ["TRANSFORMERS_CACHE"] = CACHE_DIR
os.environ["WANDB_PROJECT"] = "t5-natural-questions-qg"
# wandb.init(dir=os.getenv("WANDB_DIR", BASE_PATH))

WANDB_RUN_NAME = OUTPUT_DIR

wandb.init(
    config={
        "seed": SEED,
        "lr": LEARNING_RATE,
        "warmup_steps": WARMUP_STEPS,
        "dropout": 0.1,
        "epochs": MAX_EPOCHS,
        "scheduler": SCHEDULER,
        "model": MODEL_NAME,
        "batch_size": BATCH_SIZE,
        "output_dir": OUTPUT_DIR,
    }
)
wandb.run.name = WANDB_RUN_NAME


def collate_fn(features, pad_id=0, threshold=640):
    def pad_elems(ls, pad_id, maxlen):
        while len(ls) < maxlen:
            ls.append(pad_id)
        return ls

    maxlen = max([len(x["input_ids"]) for x in features])
    # avoid attention_type switching
    if maxlen < threshold:
        maxlen = threshold

    # dynamic padding
    input_ids = [pad_elems(x["input_ids"], pad_id=pad_id, maxlen=640) for x in features]
    input_ids = torch.tensor(input_ids, dtype=torch.long)
    # print(input_ids)
    # print(len(input_ids[0]))

    # padding mask...
    attention_mask = input_ids.clone()
    attention_mask[attention_mask != pad_id] = 1
    attention_mask[attention_mask == pad_id] = 0
    # print(attention_mask)
    # print(len(attention_mask[0]))
    # print(torch.tensor(
    #         [x["labels"] for x in features], dtype=torch.long
    #     ))
    # replace padding token id's of the labels by -100, so it's ignored by the loss
    labels = [pad_elems(x["labels"], pad_id=-100, maxlen=256) for x in features]
    labels = torch.tensor(labels, dtype=torch.long)
    # print(labels)
    # print(len(labels[0]))

    return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}


class T5ForNaturalQuestions(T5ForConditionalGeneration):
    """T5ForConditionalGeneration for NQ QA"""

    def __init__(self, config):
        super().__init__(config)

    def forward(
        self, input_ids=None, attention_mask=None, labels=None,
    ):
        outputs = super().forward(
            input_ids=input_ids, attention_mask=attention_mask, labels=labels
        )
        loss = outputs[0]

        return {"loss": loss}


if __name__ == "__main__":
    # "nq-training.jsonl" & "nq-validation.jsonl" are obtained from running `prepare_nq.py`
    tr_dataset = load_dataset(
        "json",
        data_files=BASE_PATH + "data/nq-train-tokenized-short-qg.jsonl",
        cache_dir=CACHE_DIR,
    )["train"]
    val_dataset = load_dataset(
        "json",
        data_files=BASE_PATH + "data/nq-dev-tokenized-short-qg.jsonl",
        cache_dir=CACHE_DIR,
    )["train"]

    if TRAIN_ON_SMALL == "true":
        np.random.seed(SEED)
        indices = np.random.randint(0, 87000, size=100)
        tr_dataset = tr_dataset.select(indices)
        np.random.seed(SEED)
        indices = np.random.randint(0, 2000, size=5)
        val_dataset = val_dataset.select(indices)
    # print(tr_dataset, val_dataset)

    tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR)
    tokenizer.add_special_tokens({"additional_special_tokens": ["<hl>"]})
    model = T5ForNaturalQuestions.from_pretrained(
        MODEL_NAME, cache_dir=CACHE_DIR, torch_dtype=torch.float16, device_map="auto"
    )
    model.resize_token_embeddings(len(tokenizer))

    args = TrainingArguments(
        output_dir=OUTPUT_DIR,
        overwrite_output_dir=False,
        do_train=DO_TRAIN,
        do_eval=DO_EVAL,
        evaluation_strategy="epoch",
        # eval_steps=4000,
        per_device_train_batch_size=BATCH_SIZE,
        per_device_eval_batch_size=BATCH_SIZE,
        learning_rate=LEARNING_RATE,
        warmup_steps=WARMUP_STEPS,
        lr_scheduler_type=SCHEDULER,
        num_train_epochs=MAX_EPOCHS,
        logging_steps=10,
        # save_steps=250,
        save_total_limit=2,
        save_strategy="epoch",
        load_best_model_at_end=False,
        run_name=WANDB_RUN_NAME,
        disable_tqdm=False,
        report_to=["wandb"],
        remove_unused_columns=False,
        fp16=FP16,
        seed=SEED,
        label_names=["labels"],  # it's important to log eval_loss
    )
    print("Batch Size", args.train_batch_size)
    print("Parallel Mode", args.parallel_mode)

    trainer = Trainer(
        model=model,
        args=args,
        data_collator=collate_fn,
        train_dataset=tr_dataset,
        eval_dataset=val_dataset,
    )
    try:
        checkpoint = None
        if RESUME_TRAINING is not None:
            checkpoint = RESUME_TRAINING
        trainer.train(resume_from_checkpoint=checkpoint)
        trainer.save_model(OUTPUT_DIR)
        tokenizer.save_pretrained(save_directory=OUTPUT_DIR)
    except KeyboardInterrupt:
        trainer.save_model("interrupted-natural-questions")
    wandb.finish()

src/cf_generation/baseline_generation/train_nq_qg.py
UKPLab/CATfOOD
"""
Train roberta model with counterfactuals
"""

import logging
import os.path

import torch
from transformers import (
    AutoTokenizer,
    AutoConfig,
    AutoModelForQuestionAnswering,
    TrainingArguments,
    Trainer,
)

import wandb
from src.calibration.baseline.dataloader import PreprocessData

logger = logging.getLogger(__name__)

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


class RobertaSquad:
    def __init__(
        self,
        dataset_name: str,
        dataset_config: str,
        cf_path: str,
        max_src_len: int,
        stride: int,
        do_train: bool,
        do_eval: bool,
        model_name: str,
        cache_dir: str = None,
        max_train_samples: int = None,
        max_val_samples: int = None,
        save_results: bool = False,
    ):
        self.dataset_name = dataset_name
        self.dataset_config = dataset_config
        self.cf_path = cf_path
        self.max_src_len = max_src_len
        self.stride = stride
        self.do_train = do_train
        self.do_eval = do_eval
        self.model_name = model_name
        self.cache_dir = cache_dir
        self.max_train_samples = max_train_samples
        self.max_val_samples = max_val_samples
        self.padding = "max_length"
        self.save_results = save_results
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        # load resources
        self._prepare_data()

    def _ensure_tensor_on_device(self, **inputs):
        """
        Ensure PyTorch tensors are on the specified device.
        Args:
            inputs (keyword arguments that should be :obj:`torch.Tensor`):
                The tensors to place on :obj:`self.device`.
        Return:
            :obj:`Dict[str, torch.Tensor]`: The same as :obj:`inputs` but on the proper device.
        """
        return {name: tensor.to(self.model.device) for name, tensor in inputs.items()}

    def _prepare_data(self):

        """
        Load model, tokenizer and data
        """

        config = AutoConfig.from_pretrained(self.model_name, cache_dir=self.cache_dir,)
        self.tokenizer = AutoTokenizer.from_pretrained(
            self.model_name, cache_dir=self.cache_dir, use_fast=True,
        )
        self.model = AutoModelForQuestionAnswering.from_pretrained(
            self.model_name, config=config, cache_dir=self.cache_dir,
        )
        dataloader = PreprocessData(
            self.dataset_name,
            self.dataset_config,
            cf_path=os.path.join(BASE_PATH, self.cf_path),
            save_data=False,
            save_path="",
        )
        self.train_set, self.val_set = dataloader.processed_train_val_set()

    def _preprocess_training_examples(self, examples):
        """
        Preprocess training examples
        """
        inputs = self.tokenizer(
            examples["question"],
            examples["context"],
            max_length=self.max_src_len,
            truncation="only_second",
            stride=self.stride,
            return_overflowing_tokens=True,
            return_offsets_mapping=True,
            padding="max_length",
        )

        offset_mapping = inputs.pop("offset_mapping")
        sample_map = inputs.pop("overflow_to_sample_mapping")
        answers = examples["answers"]
        start_positions = []
        end_positions = []

        for i, offset in enumerate(offset_mapping):
            sample_idx = sample_map[i]
            answer = answers[sample_idx]
            start_char = answer["answer_start"][0]
            end_char = answer["answer_start"][0] + len(answer["text"][0])
            sequence_ids = inputs.sequence_ids(i)

            # Find the start and end of the context
            idx = 0
            while sequence_ids[idx] != 1:
                idx += 1
            context_start = idx
            while sequence_ids[idx] == 1:
                idx += 1
            context_end = idx - 1

            # If the answer is not fully inside the context, label is (0, 0)
            if (
                offset[context_start][0] > start_char
                or offset[context_end][1] < end_char
            ):
                start_positions.append(0)
                end_positions.append(0)
            else:
                # Otherwise it's the start and end token positions
                idx = context_start
                while idx <= context_end and offset[idx][0] <= start_char:
                    idx += 1
                start_positions.append(idx - 1)

                idx = context_end
                while idx >= context_start and offset[idx][1] >= end_char:
                    idx -= 1
                end_positions.append(idx + 1)

        inputs["start_positions"] = start_positions
        inputs["end_positions"] = end_positions
        return inputs

    def data_collator(self, features):
        input_ids = [x["input_ids"] for x in features]
        input_ids = torch.tensor(input_ids, dtype=torch.long)
        attention_mask = [x["attention_mask"] for x in features]
        attention_mask = torch.tensor(attention_mask, dtype=torch.long)
        start_positions = [x["start_positions"] for x in features]
        start_positions = torch.tensor(start_positions, dtype=torch.long)
        end_positions = [x["end_positions"] for x in features]
        end_positions = torch.tensor(end_positions, dtype=torch.long)

        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "start_positions": start_positions,
            "end_positions": end_positions,
        }

    def train(self):
        if self.do_train:
            logger.info("*** Train ***")

            if self.max_train_samples is not None:
                # We will select sample from whole data if argument is specified
                max_train_samples = min(len(self.train_set), self.max_train_samples)
                self.train_set = self.train_set.select(range(max_train_samples))

            train_dataset = self.train_set.map(
                self._preprocess_training_examples,
                batched=True,
                remove_columns=self.train_set.column_names,
            )
            print(len(self.train_set), len(train_dataset))

        if self.do_eval or self.do_predict:
            if self.max_val_samples is not None:
                # We will select sample from whole data if argument is specified
                max_val_samples = min(len(self.val_set), self.max_val_samples)
                self.val_set = self.val_set.select(range(max_val_samples))

            validation_dataset = self.val_set.map(
                self._preprocess_training_examples,
                batched=True,
                remove_columns=self.val_set.column_names,
            )
            print(len(self.val_set), len(validation_dataset))

        args = TrainingArguments(
            output_dir=OUTPUT_DIR,
            overwrite_output_dir=False,
            do_train=True,
            do_eval=True,
            evaluation_strategy="epoch",
            save_strategy="epoch",
            # eval_steps=4000,
            per_device_train_batch_size=BATCH_SIZE,
            per_device_eval_batch_size=BATCH_SIZE,
            learning_rate=LEARNING_RATE,
            warmup_ratio=WARMUP_RATIO,
            num_train_epochs=MAX_EPOCHS,
            logging_steps=10,
            save_total_limit=2,
            seed=SEED,
            load_best_model_at_end=True,
            run_name=WANDB_RUN_NAME,
            disable_tqdm=False,
            report_to=["wandb"],
            remove_unused_columns=False,
            fp16=FP16,
        )
        print("Batch Size", args.train_batch_size)
        print("Parallel Mode", args.parallel_mode)

        trainer = Trainer(
            model=self.model,
            args=args,
            data_collator=self.data_collator,
            train_dataset=train_dataset if self.do_train else None,
            eval_dataset=validation_dataset if self.do_eval else None,
        )
        try:
            if self.do_train:
                checkpoint = None
                if RESUME_TRAINING is not None:
                    checkpoint = RESUME_TRAINING
                trainer.train(resume_from_checkpoint=checkpoint)
                trainer.save_model()
        except KeyboardInterrupt:
            trainer.save_model("interrupted-squad")

        wandb.finish()


if __name__ == "__main__":

    import argparse

    parser = argparse.ArgumentParser(
        description="Passing arguments for model, tokenizer, and dataset."
    )
    parser.add_argument(
        "--model_name",
        default="roberta-base",
        type=str,
        required=False,
        help="Specify the model to use.",
    )
    parser.add_argument(
        "--tokenizer",
        default="roberta-base",
        type=str,
        required=False,
        help="Specify the tokenizer to use.",
    )
    parser.add_argument(
        "--output_dir", type=str, required=True, help="Specify the output directory."
    )
    parser.add_argument(
        "--seed",
        default=42,
        type=int,
        required=False,
        help="Specify the seed for reproducibility.",
    )
    parser.add_argument(
        "--cf_path",
        type=str,
        required=True,
        help="Specify the path to counterfactuals.",
    )
    parser.add_argument(
        "--epochs", default=5, type=int, required=False, help="Epochs to train."
    )
    parser.add_argument(
        "--batch_size", default=64, type=int, required=False, help="Train batch size."
    )
    parser.add_argument(
        "--lr", default=1.0e-5, type=float, required=False, help="Learning rate."
    )
    parser.add_argument(
        "--warmup_ratio", default=0.06, type=float, required=False, help="Warmup ratio."
    )

    args = parser.parse_args()

    # hyperparameters
    LEARNING_RATE = args.lr
    WARMUP_RATIO = args.warmup_ratio
    MAX_EPOCHS = args.epochs
    BATCH_SIZE = args.batch_size
    OUTPUT_DIR = args.output_dir
    SEED = args.seed

    CHECKPOINT = None
    RESUME_TRAINING = None

    DO_TRAIN = True
    DO_EVAL = True
    FP16 = False

    # set to None when using all data
    MAX_TRAIN_SAMPLES = None
    MAX_VAL_SAMPLES = None

    WANDB_RUN_NAME = OUTPUT_DIR

    wandb.init(
        config={
            "lr": LEARNING_RATE,
            "warmup_ratio": WARMUP_RATIO,
            "epochs": MAX_EPOCHS,
            "model": args.model_name,
            "batch_size": BATCH_SIZE,
            "output_dir": OUTPUT_DIR,
        }
    )
    wandb.run.name = WANDB_RUN_NAME

    trainer = RobertaSquad(
        dataset_name="squad",
        dataset_config="plain_text",
        cf_path=args.cf_path,  # "src/data/squad/counterfactual_data_gpt_neox_20b_v2_qg_pipeline_all_data_cleaned.jsonl",
        max_src_len=384,
        stride=128,
        do_train=DO_TRAIN,
        do_eval=DO_EVAL,
        model_name=args.model_name,
        max_train_samples=MAX_TRAIN_SAMPLES,
        max_val_samples=MAX_VAL_SAMPLES,
        save_results=True,
    )
    trainer.train()

src/cf_generation/baseline_generation/train_qa.py
UKPLab/CATfOOD
import os

import numpy as np
import torch
import wandb
from src.calibration.baseline.dataloader import PreprocessData

from transformers import (
    T5ForConditionalGeneration,
    T5Tokenizer,
    Trainer,
    TrainingArguments,
    T5Config,
)

# BASE_PATH = "/home/sachdeva/projects/exp_calibration/src/rag/"
# CACHE_DIR = "/home/sachdeva/.cache"
BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/src/rag/"
CACHE_DIR = "/storage/ukp/work/sachdeva/.cache"

os.environ["TRANSFORMERS_CACHE"] = CACHE_DIR
os.environ["WANDB_PROJECT"] = "t5-squad-qg"
# wandb.init(dir=os.getenv("WANDB_DIR", BASE_PATH))


class T5ForQGeneration(T5ForConditionalGeneration):
    """T5ForConditionalGeneration for QA"""

    def __init__(self, config):
        super().__init__(config)

    def forward(
        self, input_ids=None, attention_mask=None, labels=None,
    ):
        outputs = super().forward(
            input_ids=input_ids, attention_mask=attention_mask, labels=labels
        )
        loss = outputs[0]

        return {"loss": loss}


class QGen:
    def __init__(
        self,
        dataset_name: str,
        dataset_config: str,
        cf_path: str,
        max_src_len: int,
        stride: int,
        do_train: bool,
        do_eval: bool,
        model_name: str,
        cache_dir: str = None,
        max_train_samples: int = None,
        max_val_samples: int = None,
        save_results: bool = False,
    ):
        self.dataset_name = dataset_name
        self.dataset_config = dataset_config
        self.cf_path = cf_path
        self.max_src_len = max_src_len
        self.stride = stride
        self.do_train = do_train
        self.do_eval = do_eval
        self.model_name = model_name
        self.cache_dir = cache_dir
        self.max_train_samples = max_train_samples
        self.max_val_samples = max_val_samples
        self.padding = "max_length"
        self.save_results = save_results
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        # load resources
        self._prepare_data()

    def _preprocess_data(self, example):
        question = example["question"].strip()
        context = example["context"]
        answer = example["answers"]
        answer_text = answer["text"][0]
        start_token = answer["answer_start"][0]
        end_token = start_token + len(answer_text)
        hl_token = "<hl>"
        label = question
        input = f"generate question: {context[:start_token]} {hl_token} {answer_text} {hl_token} {context[end_token + 1:]}"

        input_ids = self.tokenizer(input, max_length=512, truncation=True)["input_ids"]
        label_ids = self.tokenizer(label, max_length=256, padding=True, truncation=True)
        label_ids = np.array(label_ids["input_ids"])

        # print(tokenizer.pad_token_id)
        # print(label_ids)
        label_ids[label_ids == self.tokenizer.pad_token_id] = -100
        example["input_ids"] = input_ids
        example["labels"] = label_ids
        return example

    def _prepare_data(self):

        """
        Load model, tokenizer and data
        """

        config = T5Config.from_pretrained(self.model_name, cache_dir=self.cache_dir,)
        self.tokenizer = T5Tokenizer.from_pretrained(
            self.model_name, cache_dir=self.cache_dir, use_fast=True,
        )
        self.tokenizer.add_special_tokens({"additional_special_tokens": ["<hl>"]})

        self.model = T5ForQGeneration.from_pretrained(
            self.model_name, config=config, cache_dir=self.cache_dir,
        )
        self.model.resize_token_embeddings(len(self.tokenizer))
        dataloader = PreprocessData(
            self.dataset_name,
            self.dataset_config,
            cf_path=None,
            save_data=False,
            save_path="",
        )
        self.train_set, self.val_set = dataloader.processed_train_val_set()

    def collate_fn(self, features, pad_id=0, threshold=512):
        def pad_elems(ls, pad_id, maxlen):
            while len(ls) < maxlen:
                ls.append(pad_id)
            return ls

        maxlen = max([len(x["input_ids"]) for x in features])
        # avoid attention_type switching
        if maxlen < threshold:
            maxlen = threshold

        # dynamic padding
        input_ids = [
            pad_elems(x["input_ids"], pad_id=pad_id, maxlen=maxlen) for x in features
        ]
        input_ids = torch.tensor(input_ids, dtype=torch.long)
        # print(input_ids)
        # print(len(input_ids[0]))

        # padding mask...
        attention_mask = input_ids.clone()
        attention_mask[attention_mask != pad_id] = 1
        attention_mask[attention_mask == pad_id] = 0
        # replace padding token id's of the labels by -100, so it's ignored by the loss
        labels = [pad_elems(x["labels"], pad_id=-100, maxlen=256) for x in features]
        labels = torch.tensor(labels, dtype=torch.long)
        # print(labels)
        # print(len(labels[0]))

        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "labels": labels,
        }

    def train(self):

        args = TrainingArguments(
            output_dir=OUTPUT_DIR,
            overwrite_output_dir=False,
            do_train=DO_TRAIN,
            do_eval=DO_EVAL,
            evaluation_strategy="epoch",
            # eval_steps=4000,
            per_device_train_batch_size=BATCH_SIZE,
            per_device_eval_batch_size=BATCH_SIZE,
            learning_rate=LEARNING_RATE,
            warmup_steps=100,
            lr_scheduler_type="linear",
            num_train_epochs=MAX_EPOCHS,
            logging_steps=10,
            # save_steps=250,
            save_total_limit=2,
            save_strategy="epoch",
            load_best_model_at_end=False,
            run_name=WANDB_RUN_NAME,
            disable_tqdm=False,
            report_to=["wandb"],
            remove_unused_columns=False,
            fp16=FP16,
            seed=SEED,
            label_names=["labels"],  # it's important to log eval_loss
        )
        print("Batch Size", args.train_batch_size)
        print("Parallel Mode", args.parallel_mode)

        train_data = self.train_set.map(self._preprocess_data)
        val_data = self.val_set.map(self._preprocess_data)
        # print(train_data[0])

        trainer = Trainer(
            model=self.model,
            args=args,
            data_collator=self.collate_fn,
            train_dataset=train_data,
            eval_dataset=val_data,
        )
        try:
            checkpoint = None
            if RESUME_TRAINING is not None:
                checkpoint = RESUME_TRAINING
            trainer.train(resume_from_checkpoint=checkpoint)
            trainer.save_model(OUTPUT_DIR)
            self.tokenizer.save_pretrained(save_directory=OUTPUT_DIR)
        except KeyboardInterrupt:
            trainer.save_model("interrupted-squad-qg")
        wandb.finish()


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        description="Passing arguments for model, tokenizer, and dataset."
    )
    parser.add_argument(
        "--model_name",
        default="roberta-base",
        type=str,
        required=False,
        help="Specify the model to use.",
    )
    parser.add_argument(
        "--tokenizer",
        default="roberta-base",
        type=str,
        required=False,
        help="Specify the tokenizer to use.",
    )
    parser.add_argument(
        "--output_dir", type=str, required=True, help="Specify the output directory."
    )
    parser.add_argument(
        "--seed",
        default=42,
        type=int,
        required=False,
        help="Specify the seed for reproducibility.",
    )
    parser.add_argument(
        "--cf_path",
        default=None,
        type=str,
        required=False,
        help="Specify the path to counterfactuals.",
    )
    parser.add_argument(
        "--epochs", default=3, type=int, required=False, help="Epochs to train."
    )
    parser.add_argument(
        "--batch_size", default=8, type=int, required=False, help="Train batch size."
    )
    parser.add_argument(
        "--lr", default=2.0e-5, type=float, required=False, help="Learning rate."
    )
    parser.add_argument(
        "--warmup_ratio", default=0.06, type=float, required=False, help="Warmup ratio."
    )

    args = parser.parse_args()

    # hyperparameters
    LEARNING_RATE = args.lr
    WARMUP_RATIO = args.warmup_ratio
    MAX_EPOCHS = args.epochs
    BATCH_SIZE = args.batch_size
    OUTPUT_DIR = args.output_dir
    SEED = args.seed

    CHECKPOINT = None
    RESUME_TRAINING = None

    DO_TRAIN = True
    DO_EVAL = True
    FP16 = False

    # set to None when using all data
    MAX_TRAIN_SAMPLES = None
    MAX_VAL_SAMPLES = None

    WANDB_RUN_NAME = OUTPUT_DIR

    wandb.init(
        config={
            "lr": LEARNING_RATE,
            "warmup_ratio": WARMUP_RATIO,
            "epochs": MAX_EPOCHS,
            "model": args.model_name,
            "batch_size": BATCH_SIZE,
            "output_dir": OUTPUT_DIR,
        }
    )
    wandb.run.name = WANDB_RUN_NAME

    trainer = QGen(
        dataset_name="squad",
        dataset_config="plain_text",
        cf_path=args.cf_path,
        # "src/data/squad/counterfactual_data_gpt_neox_20b_v2_qg_pipeline_all_data_cleaned.jsonl",
        max_src_len=384,
        stride=128,
        do_train=DO_TRAIN,
        do_eval=DO_EVAL,
        model_name=args.model_name,
        max_train_samples=MAX_TRAIN_SAMPLES,
        max_val_samples=MAX_VAL_SAMPLES,
        save_results=True,
    )
    trainer.train()

src/cf_generation/baseline_generation/train_qg.py
UKPLab/CATfOOD
import jsonlines
from tqdm import tqdm
import pickle


def merge_jsonl_files(files, out_file):
    with jsonlines.open(out_file, "w") as writer:
        for file in tqdm(files):
            with jsonlines.open(file) as reader:
                for obj in reader:
                    writer.write(obj)


def _load_binary_file(file):
    with open(file, "rb") as f:
        return pickle.load(f)


if __name__ == "__main__":
    # base_path = "/storage/ukp/work/sachdeva/research_projects/exp_calibration"
    # file_base_path = base_path + "/src/data/squad/rag_predictions_with_answers_noise_filtered_20p_chunk"
    # chunks = 5
    # files = [file_base_path + str(i+1) + ".jsonl" for i in range(chunks)]
    # out_file = base_path + "/src/data/squad/squad_counterfactuals_noise_filtered.jsonl"
    # merge_jsonl_files(files, out_file)

    base_path = "/"
    file = base_path + "/src/data/squad_adv_new/calibration_data_final.bin"
    data = _load_binary_file(file)
    print(data)

src/cf_generation/baseline_generation/utils.py
UKPLab/CATfOOD
import os
import sys

import fire
import gradio as gr
import torch
import transformers
from peft import PeftModel
from transformers import GenerationConfig, LlamaForCausalLM, LlamaTokenizer

from src.few_shot.prompter import Prompter

if torch.cuda.is_available():
    device = "cuda"
else:
    device = "cpu"

try:
    if torch.backends.mps.is_available():
        device = "mps"
except:  # noqa: E722
    pass


def main(
    load_8bit: bool = False,
    base_model: str = "",
    lora_weights: str = "tloen/alpaca-lora-7b",
    prompt_template: str = "",  # The prompt template to use, will default to alpaca.
    server_name: str = "0.0.0.0",  # Allows to listen on all interfaces by providing '0.
    share_gradio: bool = False,
):
    base_model = base_model or os.environ.get("BASE_MODEL", "")
    assert (
        base_model
    ), "Please specify a --base_model, e.g. --base_model='decapoda-research/llama-7b-hf'"

    prompter = Prompter(prompt_template)
    tokenizer = LlamaTokenizer.from_pretrained(base_model)
    if device == "cuda":
        model = LlamaForCausalLM.from_pretrained(
            base_model,
            load_in_8bit=load_8bit,
            torch_dtype=torch.float16,
            device_map="auto",
        )
        model = PeftModel.from_pretrained(
            model, lora_weights, torch_dtype=torch.float16,
        )
    elif device == "mps":
        model = LlamaForCausalLM.from_pretrained(
            base_model, device_map={"": device}, torch_dtype=torch.float16,
        )
        model = PeftModel.from_pretrained(
            model, lora_weights, device_map={"": device}, torch_dtype=torch.float16,
        )
    else:
        model = LlamaForCausalLM.from_pretrained(
            base_model, device_map={"": device}, low_cpu_mem_usage=True
        )
        model = PeftModel.from_pretrained(model, lora_weights, device_map={"": device},)

    # unwind broken decapoda-research config
    model.config.pad_token_id = tokenizer.pad_token_id = 0  # unk
    model.config.bos_token_id = 1
    model.config.eos_token_id = 2

    if not load_8bit:
        model.half()  # seems to fix bugs for some users.

    model.eval()
    if torch.__version__ >= "2" and sys.platform != "win32":
        model = torch.compile(model)

    def evaluate(
        instruction,
        input=None,
        temperature=0.1,
        top_p=0.75,
        top_k=40,
        num_beams=4,
        max_new_tokens=128,
        **kwargs,
    ):
        prompt = prompter.generate_prompt(instruction, input)
        inputs = tokenizer(prompt, return_tensors="pt")
        input_ids = inputs["input_ids"].to(device)
        generation_config = GenerationConfig(
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            num_beams=num_beams,
            **kwargs,
        )
        with torch.no_grad():
            generation_output = model.generate(
                input_ids=input_ids,
                generation_config=generation_config,
                return_dict_in_generate=True,
                output_scores=True,
                max_new_tokens=max_new_tokens,
            )
        s = generation_output.sequences[0]
        output = tokenizer.decode(s)
        return prompter.get_response(output)

    gr.Interface(
        fn=evaluate,
        inputs=[
            gr.components.Textbox(
                lines=2, label="Instruction", placeholder="Tell me about alpacas.",
            ),
            gr.components.Textbox(lines=2, label="Input", placeholder="none"),
            gr.components.Slider(minimum=0, maximum=1, value=0.1, label="Temperature"),
            gr.components.Slider(minimum=0, maximum=1, value=0.75, label="Top p"),
            gr.components.Slider(
                minimum=0, maximum=100, step=1, value=40, label="Top k"
            ),
            gr.components.Slider(minimum=1, maximum=4, step=1, value=4, label="Beams"),
            gr.components.Slider(
                minimum=1, maximum=2000, step=1, value=128, label="Max tokens"
            ),
        ],
        outputs=[gr.inputs.Textbox(lines=5, label="Output",)],
        title="🦙🌲 Alpaca-LoRA",
        description="Alpaca-LoRA is a 7B-parameter LLaMA model finetuned to follow instructions. It is trained on the [Stanford Alpaca](https://github.com/tatsu-lab/stanford_alpaca) dataset and makes use of the Huggingface LLaMA implementation. For more information, please visit [the project's website](https://github.com/tloen/alpaca-lora).",  # noqa: E501
    ).launch(server_name="0.0.0.0", share=share_gradio)
    # Old testing code follows.

    """
    # testing code for readme
    for instruction in [
        "Tell me about alpacas.",
        "Tell me about the president of Mexico in 2019.",
        "Tell me about the king of France in 2019.",
        "List all Canadian provinces in alphabetical order.",
        "Write a Python program that prints the first 10 Fibonacci numbers.",
        "Write a program that prints the numbers from 1 to 100. But for multiples of three print 'Fizz' instead of the number and for the multiples of five print 'Buzz'. For numbers which are multiples of both three and five print 'FizzBuzz'.",  # noqa: E501
        "Tell me five words that rhyme with 'shock'.",
        "Translate the sentence 'I have no mouth but I must scream' into Spanish.",
        "Count up from 1 to 500.",
    ]:
        print("Instruction:", instruction)
        print("Response:", evaluate(instruction))
        print()
    """


if __name__ == "__main__":
    fire.Fire(main)

src/cf_generation/llm_generation/alpaca.py
UKPLab/CATfOOD
import os
import jsonlines
from tqdm import tqdm
import torch
from peft import PeftModel

from typing import Optional

from datasets import load_dataset
from langchain import PromptTemplate, FewShotPromptTemplate
from transformers import AutoModelForCausalLM, T5Tokenizer, AutoTokenizer, AutoConfig
from transformers import GenerationConfig, LlamaForCausalLM, LlamaTokenizer

# from src.few_shot.utils import save_to_disk
# from src.few_shot.together import infer


BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


def save_to_disk(data, file_name):
    with jsonlines.open(file_name, "a") as writer:
        for example in tqdm(data, total=len(data), desc="Saving samples ... "):
            writer.write(example)


class LLMClient:
    def __init__(
        self,
        template,
        model_name_or_path: str = None,
        tokenizer_name_or_path: Optional[str] = None,
        data_path: str = None,
        threshold: float = 0.5,
        task: str = "qg",
        max_new_tokens: int = 50,
        temperature: float = 0.01,
        top_p: float = 1,
        top_k: int = 40,
        repetition_penalty: float = 1.0,
        save_results: bool = True,
        max_samples: int = None,
        stop: str = "\n",
    ):
        self.base_model = model_name_or_path
        self.tokenizer = tokenizer_name_or_path
        self.max_new_tokens = max_new_tokens
        self.temperature = temperature
        self.top_p = top_p
        self.top_k = top_k
        self.repetition_penalty = repetition_penalty
        self.data_path = data_path
        self.task = task
        self.threshold = threshold
        self.save_results = save_results
        self.max_samples = max_samples
        self.stop = stop.split(";")
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.lora_weights = f"{BASE_PATH}alpaca-cot-13b"

    def _load_model(self):
        tokenizer = LlamaTokenizer.from_pretrained(self.base_model)
        model = LlamaForCausalLM.from_pretrained(
            self.base_model,
            load_in_8bit=True,
            torch_dtype=torch.float16,
            device_map="auto",
        )
        model = PeftModel.from_pretrained(
            model, self.lora_weights, torch_dtype=torch.float16,
        )
        model.to(self.device)
        # unwind broken decapoda-research config
        model.config.pad_token_id = tokenizer.pad_token_id = 0  # unk
        model.config.bos_token_id = 1
        model.config.eos_token_id = 2
        return model, tokenizer

    def _create_zero_shot_prompt(self, context, question, answer):
        prompt = qg_template.format(context=context, question=question, answer=answer)
        return prompt

    def generate(self):
        c = 0
        skipped_instances = 0
        examples = []
        model, tokenizer = self._load_model()

        model_identifier = self.base_model.split("/")[-1]
        save_path = (
            BASE_PATH
            + f"src/data/squad/{model_identifier}_{self.task}_pipeline_temp_0.7"
        )
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        # load squad data
        dataset = load_dataset("squad", "plain_text")
        train_data = dataset["train"]
        squad_data = [
            sample
            for sample in tqdm(
                train_data, total=len(train_data), desc="Loading SQuAD data ... "
            )
        ]

        current_files = []
        file_names = []
        file_path = BASE_PATH + f"src/data/squad/{self.data_path}"
        file_names = [file_path]

        model.eval()
        for file_name in file_names:
            with jsonlines.open(file_name) as reader:
                for example in tqdm(reader):
                    try:
                        id = example["id"].split("_")[0]
                        context = example["context"]
                        orig_example = [
                            sample for sample in squad_data if sample["id"] == id
                        ][0]
                        # print(orig_example)
                        orig_context = orig_example["context"]
                        orig_question = orig_example["question"]
                        orig_answer = orig_example["answers"]
                        c += 1

                        if self.max_samples:
                            if c == self.max_samples:
                                break

                        prompt = self._create_zero_shot_prompt(
                            context=orig_context,
                            question=orig_question,
                            answer=orig_answer["text"][0],
                        )
                        # print(prompt)

                        inputs = tokenizer(prompt, return_tensors="pt").to(self.device)

                        outputs = model.generate(
                            **inputs,
                            max_new_tokens=self.max_new_tokens,
                            temperature=self.temperature,
                            top_p=self.top_p,
                            top_k=self.top_k,
                            repetition_penalty=self.repetition_penalty,
                            do_sample=True,
                            # num_return_sequences=1,
                            early_stopping=True,
                            pad_token_id=tokenizer.eos_token_id,
                        )
                        output = tokenizer.decode(outputs[0], skip_special_tokens=True)
                        # remove the context from the output
                        output = output[len(prompt) :]

                        result = {"id": example["id"], "context": context}

                        print("Context:", orig_context)
                        print("Question:", orig_question)
                        print("Answer:", orig_answer["text"][0])
                        print("Explanation:", output)
                        print("-" * 100)
                        #
                        # examples.append(result)
                        # if self.save_results:
                        #     if c % self.threshold == 0:
                        #         save_to_disk(
                        #             examples,
                        #             f"{save_path}/counterfactual_samples_{model_identifier}_{c}.jsonl"
                        #         )
                        #         examples = []
                    except Exception as e:
                        # print(outputs)
                        skipped_instances += 1
                        print(f"Skipped instance {c} due to error: {e}.")
                        continue

                    if c == 5:
                        break

            # save the remaining examples
            # if self.save_results:
            #     if examples:
            #         save_to_disk(
            #             examples,
            #             f"{save_path}/counterfactual_samples_{model_identifier}_{c}.jsonl"
            #         )
            break


if __name__ == "__main__":

    qg_template = """
As an answer explainer, your job is to give a rationale for the answer to the following question given the context it is derived from. 
The rationale should express a clear thought of reasoning that led to the answer.

Context: {context}

Question: {question}

Answer: {answer}

Rationale: Let's think step by step,

""".strip()

    model = "decapoda-research/llama-13b-hf"
    client = LLMClient(
        template=qg_template,
        model_name_or_path=model,
        tokenizer_name_or_path=model,
        task="zero-shot-cot",
        data_path="squad_counterfactuals_28_03.jsonl",
        threshold=1000,
        max_new_tokens=128,
        temperature=0.7,
        top_p=1,
        top_k=40,
        repetition_penalty=1.0,
        save_results=True,
        max_samples=None,
    )

    client.generate()

src/cf_generation/llm_generation/alpaca_cot.py
UKPLab/CATfOOD
import os
import sys

import fire
import gradio as gr
import torch
import transformers
from peft import PeftModel
from transformers import GenerationConfig, LlamaForCausalLM, LlamaTokenizer

from src.few_shot.prompter import Prompter

if torch.cuda.is_available():
    device = "cuda"
else:
    device = "cpu"

try:
    if torch.backends.mps.is_available():
        device = "mps"
except:  # noqa: E722
    pass


def main(
    load_8bit: bool = False,
    base_model: str = "",
    lora_weights: str = "tloen/alpaca-lora-7b",
    prompt_template: str = "",  # The prompt template to use, will default to alpaca.
    server_name: str = "0.0.0.0",  # Allows to listen on all interfaces by providing '0.
    share_gradio: bool = False,
):
    base_model = base_model or os.environ.get("BASE_MODEL", "")
    assert (
        base_model
    ), "Please specify a --base_model, e.g. --base_model='decapoda-research/llama-7b-hf'"

    prompter = Prompter(prompt_template)
    tokenizer = LlamaTokenizer.from_pretrained(base_model)
    if device == "cuda":
        print(
            "Loading model in 8bit mode" if load_8bit else "Loading model in 16bit mode"
        )
        model = LlamaForCausalLM.from_pretrained(
            "./hf_ckpt",
            load_in_8bit=False,
            torch_dtype=torch.float16,
            device_map="auto",
        )
        # model = PeftModel.from_pretrained(
        #     model,
        #     lora_weights,
        #     torch_dtype=torch.float16,
        # )
    elif device == "mps":
        model = LlamaForCausalLM.from_pretrained(
            base_model, device_map={"": device}, torch_dtype=torch.float16,
        )
        model = PeftModel.from_pretrained(
            model, lora_weights, device_map={"": device}, torch_dtype=torch.float16,
        )
    else:
        model = LlamaForCausalLM.from_pretrained(
            base_model, device_map={"": device}, low_cpu_mem_usage=True
        )
        model = PeftModel.from_pretrained(model, lora_weights, device_map={"": device},)

    # unwind broken decapoda-research config
    model.config.pad_token_id = tokenizer.pad_token_id = 0  # unk
    model.config.bos_token_id = 1
    model.config.eos_token_id = 2

    # model.config.pad_token_id = -1
    # model.config.bos_token_id = 0
    # model.config.eos_token_id = 2

    if not load_8bit:
        model.half()  # seems to fix bugs for some users.

    model.eval()
    if torch.__version__ >= "2" and sys.platform != "win32":
        model = torch.compile(model)

    def evaluate(
        instruction,
        input=None,
        temperature=0.1,
        top_p=0.75,
        top_k=40,
        num_beams=4,
        max_new_tokens=128,
        **kwargs,
    ):

        prompt = prompter.generate_prompt(instruction, input)
        inputs = tokenizer(prompt, return_tensors="pt")
        input_ids = inputs["input_ids"].to(device)
        generation_config = GenerationConfig(
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            num_beams=num_beams,
            **kwargs,
        )
        with torch.no_grad():
            generation_output = model.generate(
                input_ids=input_ids,
                generation_config=generation_config,
                return_dict_in_generate=True,
                output_scores=True,
                max_new_tokens=max_new_tokens,
            )
        # print(generation_output)
        s = generation_output.sequences[0]
        output = tokenizer.decode(s)
        return prompter.get_response(output)

    # gr.Interface(
    #     fn=evaluate,
    #     inputs=[
    #         gr.components.Textbox(
    #             lines=2,
    #             label="Instruction",
    #             placeholder="Tell me about alpacas.",
    #         ),
    #         gr.components.Textbox(lines=2, label="Input", placeholder="none"),
    #         gr.components.Slider(
    #             minimum=0, maximum=1, value=0.1, label="Temperature"
    #         ),
    #         gr.components.Slider(
    #             minimum=0, maximum=1, value=0.75, label="Top p"
    #         ),
    #         gr.components.Slider(
    #             minimum=0, maximum=100, step=1, value=40, label="Top k"
    #         ),
    #         gr.components.Slider(
    #             minimum=1, maximum=4, step=1, value=4, label="Beams"
    #         ),
    #         gr.components.Slider(
    #             minimum=1, maximum=2000, step=1, value=128, label="Max tokens"
    #         ),
    #     ],
    #     outputs=[
    #         gr.inputs.Textbox(
    #             lines=5,
    #             label="Output",
    #         )
    #     ],
    #     title="🦙🌲 Alpaca-LoRA",
    #     description="Alpaca-LoRA is a 7B-parameter LLaMA model finetuned to follow instructions. It is trained on the [Stanford Alpaca](https://github.com/tatsu-lab/stanford_alpaca) dataset and makes use of the Huggingface LLaMA implementation. For more information, please visit [the project's website](https://github.com/tloen/alpaca-lora).",  # noqa: E501
    # ).launch(server_name="0.0.0.0", share=share_gradio)
    # Old testing code follows.

    # testing code for readme
    for instruction in [
        "Tell me about alpacas.",
        "Tell me about the president of Mexico in 2019.",
        "Tell me about the king of France in 2019.",
        "List all Canadian provinces in alphabetical order.",
        "Write a Python program that prints the first 10 Fibonacci numbers.",
        "Write a program that prints the numbers from 1 to 100. But for multiples of three print 'Fizz' instead of the number and for the multiples of five print 'Buzz'. For numbers which are multiples of both three and five print 'FizzBuzz'.",  # noqa: E501
        "Tell me five words that rhyme with 'shock'.",
        "Translate the sentence 'I have no mouth but I must scream' into Spanish.",
        "Count up from 1 to 500.",
    ]:
        print("Instruction:", instruction)
        print("Response:", evaluate(instruction))
        print()


if __name__ == "__main__":
    fire.Fire(main)

src/cf_generation/llm_generation/alpaca_generate.py
UKPLab/CATfOOD
import os
import jsonlines
from tqdm import tqdm
import torch
from peft import PeftModel

from typing import Optional

from datasets import load_dataset
from langchain import PromptTemplate, FewShotPromptTemplate
from transformers import AutoModelForCausalLM, T5Tokenizer, AutoTokenizer, AutoConfig
from transformers import GenerationConfig, LlamaForCausalLM, LlamaTokenizer

# from src.few_shot.utils import save_to_disk
# from src.few_shot.together import infer

# Set the seed
seed = 0
torch.manual_seed(seed)


BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


def save_to_disk(data, file_name):
    with jsonlines.open(file_name, "a") as writer:
        for example in tqdm(data, total=len(data), desc="Saving samples ... "):
            writer.write(example)


class LLMClient:
    def __init__(
        self,
        template,
        model_name_or_path: str = None,
        tokenizer_name_or_path: Optional[str] = None,
        data_path: str = None,
        threshold: float = 0.5,
        task: str = "qg",
        max_new_tokens: int = 50,
        temperature: float = 0.01,
        top_p: float = 1,
        top_k: int = 40,
        repetition_penalty: float = 1.0,
        save_results: bool = True,
        max_samples: int = None,
        stop: str = "\n",
    ):
        self.template = template
        self.base_model = model_name_or_path
        self.tokenizer = tokenizer_name_or_path
        self.max_new_tokens = max_new_tokens
        self.temperature = temperature
        self.top_p = top_p
        self.top_k = top_k
        self.repetition_penalty = repetition_penalty
        self.data_path = data_path
        self.task = task
        self.threshold = threshold
        self.save_results = save_results
        self.max_samples = max_samples
        self.stop = stop.split(";")
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.lora_weights = f"{BASE_PATH}lora-alpaca-13b-10ep"

    def _load_model(self):
        tokenizer = LlamaTokenizer.from_pretrained(self.base_model)
        llama_model = LlamaForCausalLM.from_pretrained(
            self.base_model,
            # load_in_8bit=True,
            torch_dtype=torch.float16,
            device_map="auto",
        )
        model = PeftModel.from_pretrained(
            llama_model,
            self.lora_weights,
            torch_dtype=torch.float16,
            device_map="auto",
        )

        # if not load_8bit:
        model.half()

        # model.to(self.device)
        # unwind broken decapoda-research config
        model.config.pad_token_id = tokenizer.pad_token_id = 0  # unk
        model.config.bos_token_id = 1
        model.config.eos_token_id = 2
        return model, tokenizer

    def _create_few_shot_prompt(self, examples, context, question=None):
        example_template = """
### Input:
Context: {context}

### Response:
Question: {question}
Answer: {answer}
""".strip()

        suffix = """
### Input:
Context: {context}

### Response:
Question:
""".strip()
        example_prompt = PromptTemplate(
            input_variables=["context", "question", "answer"],
            template=example_template,
        )

        if self.task == "qg":
            few_shot_template = FewShotPromptTemplate(
                examples=examples,
                example_prompt=example_prompt,
                prefix=self.template,
                suffix=suffix,
                input_variables=["context"],
                example_separator="\n",
            )
            prompt = few_shot_template.format(context=context)
        # print("Prompt:", prompt)

        return prompt

    def generate(self):
        c = 0
        skipped_instances = 0
        examples = []
        model, tokenizer = self._load_model()

        model_identifier = self.base_model.split("/")[-1]
        model_identifier = "alpaca-lora-hf"
        save_path = (
            BASE_PATH
            + f"src/data/squad/t5_squad_counterfactuals/{model_identifier}_{self.task}_pipeline_remaining_temp_0.7"
        )
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        # load squad data
        dataset = load_dataset("squad", "plain_text")
        train_data = dataset["train"]
        squad_data = [
            sample
            for sample in tqdm(
                train_data, total=len(train_data), desc="Loading SQuAD data ... "
            )
        ]

        current_files = []
        file_names = []
        file_path = BASE_PATH + f"src/data/squad/{self.data_path}"
        if self.task == "qg":
            file_names = [file_path]
        elif self.task == "qa":
            all_files = os.listdir(file_path)
            file_names = [
                os.path.join(file_path, file)
                for file in all_files
                if file not in current_files
            ]

        model.eval()
        for file_name in file_names:
            with jsonlines.open(file_name) as reader:
                for example in tqdm(reader):
                    try:
                        id = example["id"].split("_")[0]
                        context = example["context"]
                        orig_example = [
                            sample for sample in squad_data if sample["id"] == id
                        ][0]
                        # print(orig_example)
                        orig_context = orig_example["context"]
                        orig_question = orig_example["question"]
                        orig_answer = orig_example["answers"]
                        c += 1

                        if c <= 5000:
                            continue

                        if self.max_samples:
                            if c == self.max_samples:
                                break
                        if self.task == "qa":
                            # preprocess the question
                            # remove any new line characters
                            question = example["question"]
                            question = question.replace("\n", "")
                            # check for truncated question words
                            # check for truncated question words
                            question_words = [
                                "what",
                                "who",
                                "when",
                                "where",
                                "why",
                                "how",
                                "which",
                                "whom",
                                "whose",
                            ]
                            for word in question_words:
                                if question.lower().startswith(word[1:]):
                                    # replace first word with the correct word
                                    question = word + question[len(word[1:]) :]
                                    # question = word + question[3:]
                                    break

                        prompt = self._create_few_shot_prompt(
                            [
                                {
                                    "context": orig_context,
                                    "question": orig_question,
                                    "answer": orig_answer["text"][0],
                                }
                            ],
                            context,
                            question=question if self.task == "qa" else None,
                        )
                        # print(prompt)

                        inputs = tokenizer(prompt, return_tensors="pt").to(self.device)

                        outputs = model.generate(
                            **inputs,
                            max_new_tokens=self.max_new_tokens,
                            temperature=self.temperature,
                            top_p=self.top_p,
                            top_k=self.top_k,
                            repetition_penalty=self.repetition_penalty,
                            do_sample=True,
                            # num_return_sequences=1,
                            early_stopping=True,
                            num_beams=1,
                            pad_token_id=tokenizer.eos_token_id,
                        )
                        output = tokenizer.decode(outputs[0], skip_special_tokens=True)
                        # remove the context from the output
                        output = output[len(prompt) :]

                        # output = infer(
                        #     model_name="GPT-JT-6B-v1",
                        #     prompt=prompt,
                        # )

                        result = {"id": example["id"], "context": context}

                        if self.task == "qg":
                            # add the question mark if it is not present
                            # if output[-1] != "?":
                            #     output = output + "?"
                            # question = output

                            # for stop_word in self.stop:
                            #     if stop_word != '' and stop_word in output:
                            #         question = output[:output.find(stop_word)]
                            #     # in some cases the stop word is at the beginning of the output
                            #     if question == "":
                            #         question = output[output.find(stop_word):]
                            # Split the text into lines
                            lines = output.split("\n")

                            # Extract the question and answer
                            question = lines[0].replace("Question: ", "")
                            question = question.strip()
                            answer = lines[1].replace("Answer: ", "")
                            answer = answer.replace("<bot>:", "")
                            answer = answer.strip()

                            result["question"] = question
                            result["answer"] = answer

                        elif self.task == "qa":
                            answer = output
                            # print("Output:", output)
                            for stop_word in self.stop:
                                if stop_word != "" and stop_word in output:
                                    answer = output[output.find(stop_word) :]
                                    # remove the first stop word from answer
                                    answer = answer.split(stop_word)
                                    answer = [a for a in answer if a != ""]
                                    if answer:
                                        answer = answer[0]
                                    else:
                                        answer = "I don't know."

                                # in some cases the stop word is at the beginning of the output
                            #     if answer == "":
                            #         answer = output[output.find(stop_word):]
                            #         answer = answer.replace(stop_word, "")
                            # if answer == "":
                            #     answer = answer.split(stop_word)
                            #     answer = [a for a in answer if a != ""][0]
                            skip_words = [
                                "A: ",
                                "A. ",
                                "Answer: ",
                                "Output: ",
                                "Label: ",
                            ]
                            for skip_word in skip_words:
                                if skip_word in answer:
                                    answer = answer.replace(skip_word, "")

                            # if answer.startswith("A: ") or answer.startswith("A. "):
                            #     answer = answer[3:]
                            # elif answer.startswith("Answer: "):
                            #     answer = answer.replace("Answer: ", "")
                            # question = example["question"]
                            result["question"] = question
                            result["answer"] = answer

                        # print("Context:", context)
                        # print("Question:", question)
                        # print("Answer:", answer)
                        # print("-" * 100)

                        examples.append(result)
                        if self.save_results:
                            if c % self.threshold == 0:
                                save_to_disk(
                                    examples,
                                    f"{save_path}/counterfactual_samples_{model_identifier}_{c}.jsonl",
                                )
                                examples = []
                    except Exception as e:
                        # print(outputs)
                        skipped_instances += 1
                        print(f"Skipped instance {c} due to error: {e}.")
                        continue

                    # if c==10:
                    #     break

            # save the remaining examples
            if self.save_results:
                if examples:
                    save_to_disk(
                        examples,
                        f"{save_path}/counterfactual_samples_{model_identifier}_{c}.jsonl",
                    )
            # break


if __name__ == "__main__":

    qg_template = """
As a question generator, your task is to create a concise and clear question that can be answered by an answer span within a given context. 
The context should be a piece of text, such as a news article or historical document, and the question should require understanding and analysis of the information presented in the context. 
Your generated question should focus on key details or events described in the context and should demonstrate your ability to identify important information. 
Additionally, please ensure that your question is specific enough to have a single correct answer within the given context. 
Please note that you may need to read through the provided context multiple times to fully understand its contents before generating an appropriate question. 
""".strip()

    alpaca_prompt = """
Below is an instruction that describes a task, paired with an input that provides further context. 
Write a response that appropriately completes the request.
### Instruction:
As a question generator, your task is to create a concise and clear question that can be answered by an answer span within the given context. 
The context should be a piece of text, such as a news article or historical document, and the question should require understanding and analysis of the information presented in the context. 
Your generated question should focus on key details or events described in the context and should demonstrate your ability to identify important information. 
Additionally, please ensure that your question is specific enough to have a single correct answer within the given context. 
Please note that you may need to read through the provided context multiple times to fully understand its contents before generating an appropriate question. 
""".strip()

    task = "qg"
    model = "decapoda-research/llama-13b-hf"
    client = LLMClient(
        template=alpaca_prompt,
        model_name_or_path=model,
        tokenizer_name_or_path=model,
        data_path="t5_squad_counterfactuals/rag_counterfactuals_complete_min_filtered_dedup.jsonl",
        threshold=1000,
        task=task,
        max_new_tokens=50,
        temperature=0.7,
        top_p=1,
        top_k=40,
        repetition_penalty=1.0,
        save_results=True,
        max_samples=None,
    )

    client.generate()

src/cf_generation/llm_generation/alpaca_prompter.py
UKPLab/CATfOOD
import inseq
import torch
import json
import os
from collections import Counter
import pandas as pd
from tqdm import tqdm
from itertools import islice
import collections

# "{q} The possibile answers are {options_list with a, b, c}, but the correct of a, b, c is"

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


def process_data():
    # read json file
    with open(f"{BASE_PATH}src/few_shot/big_bench_tasks/social_iqa.json", "r") as file:
        data = json.load(file)
        examples = data["examples"]
    questions = [sample["input"] for sample in examples]
    options = [list(sample["target_scores"].keys()) for sample in examples]
    answers = [list(sample["target_scores"].values()) for sample in examples]
    return questions, options, answers


def _merge_attributions(attributions):
    pass


def evaluate(file_name):
    correct = 0
    mapper = {"a": 0, "b": 1, "c": 2}

    with open(file_name, "r") as file:
        data = json.load(file)
    total_samples = len(data)
    for sample in tqdm(data):
        # print(sample)
        answer = sample["answer"][0]
        pred = sample["prediction"][0].strip()
        if len(pred) > 1:
            pred = pred[0]
        answer_key = mapper.get(pred, None)
        # print("answer key: ", answer_key)
        if answer_key is not None:
            if sample["options"][answer_key] == answer:
                correct += 1
        else:
            # print(sample["options"][answer_key])
            print(answer)
            print("-")
        # break

    accuracy = (correct / total_samples) * 100
    print(accuracy)


def get_attribution(model_name, attribution_type, max_samples=None, top_k=5):

    questions, options, answers = process_data()

    if max_samples:
        questions = questions[:max_samples]
        options = options[:max_samples]
        answers = answers[:max_samples]

    model = inseq.load_model(
        model_name, attribution_type, torch_dtype=torch.bfloat16, device_map="auto"
    )
    # load_in_8bit=True)

    feature_attributions = []
    topk_attributions = []
    outputs = []
    for ques, option, ans in tqdm(zip(questions, options, answers)):

        len_options = len(option)
        if len_options == 3:
            # prompt = f"Question: {ques} " \
            #          f"\n Options: a) {option[0]}, b) {option[1]}, c) {option[2]}" \
            #          "\n Answer: "
            prompt = f"{ques} The possible answers are a) {option[0]}, b) {option[1]}, c) {option[2]}, but the correct of a, b, c is"
        elif len_options == 2:
            # prompt = f"Question: {ques} " \
            #          f"\n Options: a) {option[0]}, b) {option[1]}" \
            #          "\n Answer: "
            prompt = f"{ques} The possible answers are a) {option[0]}, b) {option[1]}, but the correct of a, b is"

        attributions = model.attribute(prompt, generation_args={"max_new_tokens": 15},)
        # attributions.show()
        # print(attributions)
        # merged_attributions = inseq.FeatureAttributionOutput.merge_attributions([attributions])
        # print("Merged:", merged_attributions)
        aggr = attributions.sequence_attributions[0].aggregate()
        # print(aggr)
        # Creating a mapping of [src_token, tgt_token] -> attribution score
        # score_map = {}
        # for src_idx, src_tok in enumerate(aggr.source):
        #     for tgt_idx, tgt_tok in enumerate(aggr.target):
        #         score_map[(src_tok.token, tgt_tok.token)] = aggr.source_attributions[src_idx, tgt_idx].item()
        #
        # print(score_map)
        # print(attributions.get_scores_dicts())
        score_dict = aggr.get_scores_dicts()
        source_attributions = score_dict["source_attributions"]
        # feature_attributions.append(source_attributions)
        # df = pd.DataFrame(source_attributions)
        # print(df.head())

        averages = {}
        for token, token_dict in source_attributions.items():
            for key, value in token_dict.items():
                if key not in averages:
                    averages[key] = [value, 1]
                else:
                    averages[key][0] += value
                    averages[key][1] += 1

        averages_dict = {key: value[0] / value[1] for key, value in averages.items()}
        # print(averages_dict)

        x_before_options = {}
        for key, value in averages_dict.items():
            if key == "▁Options":
                break
            x_before_options[key] = value
        # print(x_before_options)
        sorted_averages = dict(
            sorted(x_before_options.items(), key=lambda x: x[1], reverse=True)
        )
        # print(sorted_averages)
        top_attributions = dict(islice(sorted_averages.items(), top_k))
        # topk_attributions.append(top_attributions)

        output = collections.OrderedDict()
        output["question"] = ques
        output["options"] = option
        output["answer"] = [option[ans.index(1)]]
        output["prediction"] = attributions.info["generated_texts"]
        output["attribution"] = top_attributions
        outputs.append(output)
        # print(output)

    save_file_name = model_name.split("/")[-1] + f"_{attribution_type}_social_iqa"
    output_nbest_file = os.path.join(
        BASE_PATH, f"predictions_{save_file_name}_prompt_2.json"
    )
    with open(output_nbest_file, "w") as writer:
        writer.write(json.dumps(outputs, indent=4) + "\n")

    return feature_attributions


if __name__ == "__main__":
    # source_attributions = get_attribution(
    #     model_name="google/flan-t5-large",
    #     attribution_type="attention",
    #     max_samples=None,
    #     top_k=5
    # )
    # print(source_attributions)
    # top_attrs = postprocess_attributions(source_attributions, top_k=10)
    # print(top_attrs)
    # process_data()

    ########### EVALUATE ############
    model_name = "flan-t5-large"
    evaluate(os.path.join(BASE_PATH, f"predictions_{model_name}_attention.json"))

    # x = [{'▁': {'▁Question': 0.0004824649658985436, ':': 0.09963615983724594, '▁Tracy': 0.023480601608753204,
    #             '▁didn': 0.0008805241086520255, "'": 0.0013725516619160771, 't':0.0015102395555004478,
    #             '▁go': 0.0005502908607013524, '▁home': 0.0014923501294106245, '▁that': 0.0009823078289628029,
    #             '▁evening': 0.0018404676811769605, '▁and': 0.00046898151049390435, '▁': 0.09211429953575134,
    #             're': 0.0004362099862191826, 'sisted': 0.0027353481855243444, '▁Riley': 0.007774821948260069,
    #             's': 0.0036468186881393194,'▁attacks': 0.004170892760157585, '.': 0.0030060261487960815,
    #             '▁What': 0.002392962109297514, '▁does': 0.0019515401218086481, '▁need': 0.013643212616443634,
    #             '▁to': 0.0005810720031149685, '▁do': 0.0037925769574940205, '▁before': 0.00945215579122305,
    #             '▁this': 0.0017502920236438513, '?': 0.002131211804226041, '▁Options': 0.013728966936469078,
    #             'a': 0.004350507166236639, ')': 0.0063089970499277115, '▁Make': 0.05659395083785057,
    #             '▁new': 0.0028916974551975727, '▁plan': 0.0028365186881273985, ',': 0.013565066270530224,
    #             'b': 0.008598132990300655, '▁Go': 0.0036058370023965836, '▁see': 0.0018858902622014284 }},
    # {'▁': {'▁Question': 0.0004824649658985436, ':': 0.09963615983724594, '▁Tracy': 0.023480601608753204,
    #        '▁didn': 0.0008805241086520255, "'": 0.0013725516619160771, 't': 0.0015102395555004478,
    #        '▁go': 0.0005502908607013524, '▁home': 0.0014923501294106245, '▁that': 0.0009823078289628029,
    #        '▁evening': 0.0018404676811769605, '▁and': 0.00046898151049390435, '▁': 0.09211429953575134,
    #        're': 0.0004362099862191826, 'sisted': 0.0027353481855243444, '▁Riley': 0.007774821948260069,
    #        's': 0.0036468186881393194, '▁attacks': 0.004170892760157585, '.': 0.0030060261487960815,
    #        '▁What': 0.002392962109297514, '▁does': 0.0019515401218086481, '▁need': 0.013643212616443634,
    #        '▁to': 0.0005810720031149685, '▁do': 0.0037925769574940205, '▁before': 0.00945215579122305,
    #        '▁this': 0.0017502920236438513, '?': 0.002131211804226041, '▁Options': 0.013728966936469078,
    #        'a': 0.004350507166236639, ')': 0.0063089970499277115, '▁Make': 0.05659395083785057,
    #        '▁new': 0.0028916974551975727, '▁plan': 0.0028365186881273985, ',': 0.013565066270530224,
    #        'b': 0.008598132990300655, '▁Go': 0.0036058370023965836, '▁see': 0.0018858902622014284}}
    # ]

    # attrs = [list(attr.values())[0] for attr in x]
    # # print(attrs)
    # tokens, attributions = list(attrs[0].keys()), list(attrs[0].values())
    # print(tokens)
    #
    # new_dict = {}
    #
    # for key, value in attrs[0].items():
    #     if key.startswith('_'):
    #         continue
    #     new_key = ''.join(filter(lambda x: x != '_', key))
    #     if new_key not in new_dict:
    #         new_dict[new_key] = value
    #     else:
    #         new_dict[new_key] += value
    #
    # print(new_dict)

    # filtered_tokens = [tokens[0]]
    # for word in tokens[1:]:
    #     print(word)
    #     if not word.startswith("▁"):
    #         print("here word:", word)
    #         filtered_tokens[-1] += f'{word}'
    #         # print("fil:", filtered_tokens)
    #     else:
    #         filtered_tokens.append(word)
    # print(filtered_tokens)
    #
    # attribution_score = [attributions[0]]
    # for idx, (word_idx, score) in enumerate(zip(word_map, attributions[1:])):
    #     if word_idx == word_map[idx + 1] and word_idx is not None:
    #         attribution_score[-1] = attribution_score[-1] + score
    #     else:
    #         attribution_score.append(score)

src/cf_generation/llm_generation/explainer.py
UKPLab/CATfOOD
import os

import torch
import transformers
from peft import PeftModel
from transformers import LlamaForCausalLM, LlamaTokenizer  # noqa: F402

BASE_MODEL = os.environ.get("BASE_MODEL", "llama-13b-hf")
assert (
    BASE_MODEL
), "Please specify a value for BASE_MODEL environment variable, e.g. `export BASE_MODEL=...`"  # noqa: E501

tokenizer = LlamaTokenizer.from_pretrained(BASE_MODEL)

base_model = LlamaForCausalLM.from_pretrained(
    BASE_MODEL, load_in_8bit=False, torch_dtype=torch.float16, device_map={"": "cpu"},
)

first_weight = base_model.model.layers[0].self_attn.q_proj.weight
first_weight_old = first_weight.clone()

lora_model = PeftModel.from_pretrained(
    base_model,
    "./lora-alpaca-13b-10ep",
    device_map={"": "cpu"},
    torch_dtype=torch.float16,
)

lora_weight = lora_model.base_model.model.model.layers[0].self_attn.q_proj.weight

assert torch.allclose(first_weight_old, first_weight)

# merge weights
for layer in lora_model.base_model.model.model.layers:
    layer.self_attn.q_proj.merge_weights = True
    layer.self_attn.v_proj.merge_weights = True

lora_model.train(False)

# did we do anything?
assert not torch.allclose(first_weight_old, first_weight)

lora_model_sd = lora_model.state_dict()
deloreanized_sd = {
    k.replace("base_model.model.", ""): v
    for k, v in lora_model_sd.items()
    if "lora" not in k
}

LlamaForCausalLM.save_pretrained(
    base_model, "./hf_ckpt", state_dict=deloreanized_sd, max_shard_size="400MB"
)

src/cf_generation/llm_generation/export_hf_ckpt.py
UKPLab/CATfOOD
import os

from transformers import AutoModelForCausalLM, AutoTokenizer


import torch
import jsonlines
from tqdm import tqdm
from datasets import load_dataset

from src.few_shot.utils import save_to_disk

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"
# BASE_PATH = "/home/sachdeva/projects/ukp/exp_calibration/"


if __name__ == "__main__":

    # load squad data
    dataset = load_dataset("squad", "plain_text")
    train_data = dataset["train"]
    squad_data = [
        sample
        for sample in tqdm(
            train_data, total=len(train_data), desc="Loading SQuAD data ... "
        )
    ]

    c = 0
    examples = []

    # prompt = "Generate a fluent and answerable question from the given context. Ensure that the answer " \
    #          "is a span in the context and is less than 10 words."

    # prompt = "Given the context below, generate a fluent question that can be answered from it."  # simple
    # prompt = "Using the context provided, produce a well-crafted question that can be answered from it."  # vague
    # prompt = "From the given context, create a clear and concise question that can be answered from it."  # vague
    # prompt = "Given the context below, create a specific question that can be answered using only information " \
    #          "directly stated in the passage."
    # prompt = "Using the given context, create a question that can be answered by selecting a relevant span from the text."
    # prompt = "From the provided context, create a question that can be accurately answered by extracting " \
    #          "information directly from the context."
    # prompt = "Construct a question that can be answered by selecting a key piece of information from the given context."
    prompt = (
        "Generate a fluent and answerable question from the given context. Ensure that the answer "
        "is a span in the context and is less than 10 words."
    )

    model_name = "facebook/opt-13b"
    model_identifier = model_name.split("/")[-1]
    save_path = BASE_PATH + f"src/data/squad/{model_identifier}_qg"
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    model = AutoModelForCausalLM.from_pretrained(
        model_name, torch_dtype=torch.float16
    ).cuda()
    tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)

    with jsonlines.open(
        BASE_PATH
        + "src/data/squad/squad_counterfactuals_noise_min_filtered_final_2.jsonl"
    ) as reader:
        for example in tqdm(reader):
            c += 1
            id = example["id"].split("_")[0]
            context = example["context"]
            orig_example = [sample for sample in squad_data if sample["id"] == id][0]
            # print(orig_example)
            orig_context = orig_example["context"]
            orig_question = orig_example["question"]
            orig_answer = orig_example["answers"]

            # input = f"{prompt} \nContext: {context} \nQuestion: "
            input = f"{prompt} \n\nContext: {context} \n\nQuestion: "
            #
            # input = f"{prompt} \nContext: {orig_context} \nQuestion: {orig_question} \nAnswer span: {orig_answer['text'][0]} " \
            #         f"\n{prompt} \nContext: {context} " \
            #         f"\nQuestion: \nAnswer span: "

            # print(input)
            inputs = tokenizer(input, return_tensors="pt").to("cuda")
            generated_ids = model.generate(
                inputs.input_ids,
                attention_mask=inputs.attention_mask,
                max_new_tokens=50,
            )
            outputs = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
            print(outputs)

            if c == 2:
                break
            # break
            # result = {
            #     "id": example["id"],
            #     "question": tokenizer.decode(outputs[0], skip_special_tokens=True),
            #     "context": context
            # }
            # print(result)
        #     examples.append(result)
        #     if c % 5000 == 0:
        #         save_to_disk(
        #             examples,
        #             f"{save_path}/counterfactual_questions_{model_identifier}_{c}.jsonl"
        #         )
        #         examples = []
        #
        # # save the remaining examples
        # if examples:
        #     save_to_disk(
        #         examples,
        #         f"{save_path}/counterfactual_questions_{model_identifier}_{c}.jsonl"
        #     )

src/cf_generation/llm_generation/fb_opt_qgen.py
UKPLab/CATfOOD
import os
import sys
from typing import List

import fire
import torch
import transformers
from datasets import load_dataset

"""
Unused imports:
import torch.nn as nn
import bitsandbytes as bnb
"""

from peft import (
    LoraConfig,
    get_peft_model,
    get_peft_model_state_dict,
    prepare_model_for_int8_training,
    set_peft_model_state_dict,
)
from transformers import LlamaForCausalLM, LlamaTokenizer

from src.few_shot.prompter import Prompter


def train(
    # model/data params
    base_model: str = "",  # the only required argument
    data_path: str = "yahma/alpaca-cleaned",
    output_dir: str = "./lora-alpaca",
    # training hyperparams
    batch_size: int = 128,
    micro_batch_size: int = 4,
    num_epochs: int = 3,
    learning_rate: float = 3e-4,
    cutoff_len: int = 256,
    val_set_size: int = 2000,
    # lora hyperparams
    lora_r: int = 8,
    lora_alpha: int = 16,
    lora_dropout: float = 0.05,
    lora_target_modules: List[str] = ["q_proj", "v_proj",],
    # llm hyperparams
    train_on_inputs: bool = True,  # if False, masks out inputs in loss
    group_by_length: bool = False,  # faster, but produces an odd training loss curve
    # wandb params
    wandb_project: str = "LLMs",
    wandb_run_name: str = "alpaca-13b",
    wandb_watch: str = "",  # options: false | gradients | all
    wandb_log_model: str = "",  # options: false | true
    resume_from_checkpoint: str = None,  # either training checkpoint or final adapter
    prompt_template_name: str = "alpaca",  # The prompt template to use, will default to alpaca.
):
    if int(os.environ.get("LOCAL_RANK", 0)) == 0:
        print(
            f"Training Alpaca-LoRA model with params:\n"
            f"base_model: {base_model}\n"
            f"data_path: {data_path}\n"
            f"output_dir: {output_dir}\n"
            f"batch_size: {batch_size}\n"
            f"micro_batch_size: {micro_batch_size}\n"
            f"num_epochs: {num_epochs}\n"
            f"learning_rate: {learning_rate}\n"
            f"cutoff_len: {cutoff_len}\n"
            f"val_set_size: {val_set_size}\n"
            f"lora_r: {lora_r}\n"
            f"lora_alpha: {lora_alpha}\n"
            f"lora_dropout: {lora_dropout}\n"
            f"lora_target_modules: {lora_target_modules}\n"
            f"train_on_inputs: {train_on_inputs}\n"
            f"group_by_length: {group_by_length}\n"
            f"wandb_project: {wandb_project}\n"
            f"wandb_run_name: {wandb_run_name}\n"
            f"wandb_watch: {wandb_watch}\n"
            f"wandb_log_model: {wandb_log_model}\n"
            f"resume_from_checkpoint: {resume_from_checkpoint or False}\n"
            f"prompt template: {prompt_template_name}\n"
        )
    assert (
        base_model
    ), "Please specify a --base_model, e.g. --base_model='decapoda-research/llama-7b-hf'"
    gradient_accumulation_steps = batch_size // micro_batch_size

    prompter = Prompter(prompt_template_name)

    device_map = "auto"
    world_size = int(os.environ.get("WORLD_SIZE", 1))
    ddp = world_size != 1
    if ddp:
        device_map = {"": int(os.environ.get("LOCAL_RANK") or 0)}
        gradient_accumulation_steps = gradient_accumulation_steps // world_size

    # Check if parameter passed or if set within environ
    use_wandb = len(wandb_project) > 0 or (
        "WANDB_PROJECT" in os.environ and len(os.environ["WANDB_PROJECT"]) > 0
    )
    # Only overwrite environ if wandb param passed
    if len(wandb_project) > 0:
        os.environ["WANDB_PROJECT"] = wandb_project
    if len(wandb_watch) > 0:
        os.environ["WANDB_WATCH"] = wandb_watch
    if len(wandb_log_model) > 0:
        os.environ["WANDB_LOG_MODEL"] = wandb_log_model

    model = LlamaForCausalLM.from_pretrained(
        base_model, load_in_8bit=True, torch_dtype=torch.float16, device_map=device_map,
    )

    tokenizer = LlamaTokenizer.from_pretrained(base_model)

    tokenizer.pad_token_id = 0  # unk. we want this to be different from the eos token
    tokenizer.padding_side = "left"  # Allow batched inference

    def tokenize(prompt, add_eos_token=True):
        # there's probably a way to do this with the tokenizer settings
        # but again, gotta move fast
        result = tokenizer(
            prompt,
            truncation=True,
            max_length=cutoff_len,
            padding=False,
            return_tensors=None,
        )
        if (
            result["input_ids"][-1] != tokenizer.eos_token_id
            and len(result["input_ids"]) < cutoff_len
            and add_eos_token
        ):
            result["input_ids"].append(tokenizer.eos_token_id)
            result["attention_mask"].append(1)

        result["labels"] = result["input_ids"].copy()

        return result

    def generate_and_tokenize_prompt(data_point):
        full_prompt = prompter.generate_prompt(
            data_point["instruction"], data_point["input"], data_point["output"],
        )
        tokenized_full_prompt = tokenize(full_prompt)
        if not train_on_inputs:
            user_prompt = prompter.generate_prompt(
                data_point["instruction"], data_point["input"]
            )
            tokenized_user_prompt = tokenize(user_prompt, add_eos_token=False)
            user_prompt_len = len(tokenized_user_prompt["input_ids"])

            tokenized_full_prompt["labels"] = (
                [-100] * user_prompt_len
                + tokenized_full_prompt["labels"][user_prompt_len:]
            )  # could be sped up, probably
        return tokenized_full_prompt

    model = prepare_model_for_int8_training(model)

    config = LoraConfig(
        r=lora_r,
        lora_alpha=lora_alpha,
        target_modules=lora_target_modules,
        lora_dropout=lora_dropout,
        bias="none",
        task_type="CAUSAL_LM",
    )
    model = get_peft_model(model, config)

    if data_path.endswith(".json") or data_path.endswith(".jsonl"):
        data = load_dataset("json", data_files=data_path)
    else:
        data = load_dataset(data_path)

    if resume_from_checkpoint:
        # Check the available weights and load them
        checkpoint_name = os.path.join(
            resume_from_checkpoint, "pytorch_model.bin"
        )  # Full checkpoint
        if not os.path.exists(checkpoint_name):
            checkpoint_name = os.path.join(
                resume_from_checkpoint, "adapter_model.bin"
            )  # only LoRA model - LoRA config above has to fit
            resume_from_checkpoint = False  # So the trainer won't try loading its state
        # The two files above have a different name depending on how they were saved, but are actually the same.
        if os.path.exists(checkpoint_name):
            print(f"Restarting from {checkpoint_name}")
            adapters_weights = torch.load(checkpoint_name)
            model = set_peft_model_state_dict(model, adapters_weights)
        else:
            print(f"Checkpoint {checkpoint_name} not found")

    model.print_trainable_parameters()  # Be more transparent about the % of trainable params.

    if val_set_size > 0:
        train_val = data["train"].train_test_split(
            test_size=val_set_size, shuffle=True, seed=42
        )
        train_data = train_val["train"].shuffle().map(generate_and_tokenize_prompt)
        val_data = train_val["test"].shuffle().map(generate_and_tokenize_prompt)
    else:
        train_data = data["train"].shuffle().map(generate_and_tokenize_prompt)
        val_data = None

    if not ddp and torch.cuda.device_count() > 1:
        # keeps Trainer from trying its own DataParallelism when more than 1 gpu is available
        model.is_parallelizable = True
        model.model_parallel = True

    trainer = transformers.Trainer(
        model=model,
        train_dataset=train_data,
        eval_dataset=val_data,
        args=transformers.TrainingArguments(
            per_device_train_batch_size=micro_batch_size,
            gradient_accumulation_steps=gradient_accumulation_steps,
            warmup_steps=100,
            num_train_epochs=num_epochs,
            learning_rate=learning_rate,
            fp16=True,
            logging_steps=10,
            optim="adamw_torch",
            evaluation_strategy="steps" if val_set_size > 0 else "no",
            save_strategy="steps",
            eval_steps=200 if val_set_size > 0 else None,
            save_steps=200,
            output_dir=output_dir,
            save_total_limit=3,
            load_best_model_at_end=True if val_set_size > 0 else False,
            ddp_find_unused_parameters=False if ddp else None,
            group_by_length=group_by_length,
            report_to="wandb" if use_wandb else None,
            run_name=wandb_run_name if use_wandb else None,
        ),
        data_collator=transformers.DataCollatorForSeq2Seq(
            tokenizer, pad_to_multiple_of=8, return_tensors="pt", padding=True
        ),
    )
    model.config.use_cache = False

    old_state_dict = model.state_dict
    model.state_dict = (
        lambda self, *_, **__: get_peft_model_state_dict(self, old_state_dict())
    ).__get__(model, type(model))

    if torch.__version__ >= "2" and sys.platform != "win32":
        model = torch.compile(model)

    trainer.train(resume_from_checkpoint=resume_from_checkpoint)

    model.save_pretrained(output_dir)

    print("\n If there's a warning about missing keys above, please disregard :)")


if __name__ == "__main__":
    fire.Fire(train)

src/cf_generation/llm_generation/finetune_alpaca.py
UKPLab/CATfOOD
from transformers import T5ForConditionalGeneration, T5Tokenizer, AutoTokenizer
import torch
import os
import jsonlines
import argparse
from tqdm import tqdm
from datasets import load_dataset

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"
# BASE_PATH = "/home/sachdeva/projects/ukp/exp_calibration/"


def save_to_disk(data, file_name):
    with jsonlines.open(file_name, "a") as writer:
        for example in tqdm(data, total=len(data), desc="Saving samples ... "):
            writer.write(example)


if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()

    # Set the seed
    torch.manual_seed(args.seed)

    # load squad data
    dataset = load_dataset("squad", "plain_text")
    train_data = dataset["train"]
    squad_data = [
        sample
        for sample in tqdm(
            train_data, total=len(train_data), desc="Loading SQuAD data ... "
        )
    ]

    current_files = []

    c = 0
    # prompt = "Given the following question and context, provide an answer " \
    #          "that best fits the question. Ensure that the answer " \
    #          "is a span in the context."
    prompt = (
        "Answer the question based on the context below. If the question cannot be answered "
        "using the information provided, then answer with 'I don't know'."
    )

    # to test
    # prompt = "Your task is to answer a question based on the given context. If the information provided " \
    #          "is insufficient to answer the question, please respond with 'I don't know.' Your response " \
    #          "should be clear and concise, providing only relevant information necessary to answer the question. " \
    #          "Please note that you should make every effort to provide an accurate and complete response based " \
    #          "on the available information."

    model_name = "google/flan-t5-xxl"
    model_identifier = model_name.split("/")[-1]

    save_path = (
        BASE_PATH
        + f"src/data/squad/noise_filter_{model_identifier}_qa_eval_seed_{args.seed}"
    )
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    model = T5ForConditionalGeneration.from_pretrained(
        model_name, torch_dtype=torch.bfloat16, device_map="auto"
    )
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    # all_files = os.listdir(BASE_PATH + f"src/data/squad/few_shot_{model_identifier}_qg_temp_0.7/")
    # files = [file for file in all_files if file not in current_files]
    # print(files)
    files = [
        "counterfactual_samples_Llama-2-13b-chat-hf_flan-t5-xxl_context_filtered_complete.jsonl"
    ]

    for file in files:
        c += 1
        # if c<=9:
        #     continue
        print(f"Processing file: {file}")
        examples = []
        with jsonlines.open(BASE_PATH + f"src/data/squad/" + file) as reader:
            # with jsonlines.open(BASE_PATH + f"src/data/squad/few_shot_{model_identifier}_qg_temp_0.7/" + file) as reader:
            for example in tqdm(reader):
                id = example["id"].split("_")[0]
                context = example["context"]
                question = example["question"]
                tokens_to_remove = ["[", "'", '"', "]"]
                # Create a translation table that maps each unwanted token to None
                translator = str.maketrans({token: None for token in tokens_to_remove})
                question = question.translate(translator).strip()

                orig_example = [sample for sample in squad_data if sample["id"] == id][
                    0
                ]

                orig_context = orig_example["context"]
                orig_question = orig_example["question"]
                orig_answer = orig_example["answers"]

                # input = f"{prompt} \nQuestion: {orig_question} \nContext: {orig_context} " \
                #         f"\nAnswer: {orig_answer['text'][0]}" \
                #         f"\n{prompt} \nQuestion: {question}  \nContext: {context} " \
                #         f"\nAnswer: "

                input = (
                    f"{prompt} \nContext: {orig_context} \nQuestion: {orig_question} "
                    f"\nAnswer: {orig_answer['text'][0]}"
                    f"\n{prompt} \nContext: {context} \nQuestion: {question} "
                    f"\nAnswer: "
                )

                inputs = tokenizer(input, return_tensors="pt").input_ids.to("cuda")
                outputs = model.generate(
                    inputs,
                    max_length=200,
                    temperature=0.7,
                    num_beams=1,
                    num_return_sequences=1,
                )
                result = {
                    "id": example["id"],
                    "question": question,
                    "context": context,
                    "answer": tokenizer.decode(outputs[0], skip_special_tokens=True),
                }
                # print(result)
                examples.append(result)

        save_to_disk(
            examples, f"{save_path}/counterfactual_samples_{model_identifier}_{c}.jsonl"
        )
        # c += 1

src/cf_generation/llm_generation/flan_ul2_qa.py
UKPLab/CATfOOD
from transformers import T5ForConditionalGeneration, T5Tokenizer, AutoTokenizer
import torch
import os
import jsonlines
from tqdm import tqdm
from datasets import load_dataset

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"
# BASE_PATH = "/home/sachdeva/projects/ukp/exp_calibration/"


def save_to_disk(data, file_name):
    with jsonlines.open(file_name, "a") as writer:
        for example in tqdm(data, total=len(data), desc="Saving samples ... "):
            writer.write(example)


if __name__ == "__main__":

    # load squad data
    dataset = load_dataset("squad", "plain_text")
    train_data = dataset["train"]
    squad_data = [
        sample
        for sample in tqdm(
            train_data, total=len(train_data), desc="Loading SQuAD data ... "
        )
    ]

    # print(len(squad_data))
    # print(squad_data[0])

    c = 0
    examples = []

    prompt = (
        "Generate a fluent and answerable question from the given context. Ensure that the answer "
        "is a span in the context and is less than 10 words."
    )

    # flan t5 xxl
    # prompt = "As a question generation model, your task is to generate a question from a given context " \
    #          "that can be answered by a specific span of text highlighted in the context. The question " \
    #          "should be clear and concise, using appropriate language and grammar. It should also focus on " \
    #          "extracting relevant information from the highlighted span. Please note that you should be able " \
    #          "to handle various types of contexts, including those with complex sentence structures or " \
    #          "technical jargon. Your response should provide a well-formed question that accurately captures " \
    #          "the meaning of the highlighted span and encourages critical thinking."

    model_name = "google/flan-t5-xxl"
    model_identifier = model_name.split("/")[-1]

    save_path = BASE_PATH + f"src/data/squad/few_shot_{model_identifier}_qg_temp_0.7"
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    model = T5ForConditionalGeneration.from_pretrained(
        model_name, torch_dtype=torch.bfloat16, device_map="auto"
    )
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    with jsonlines.open(
        BASE_PATH + "src/data/squad/squad_counterfactuals_28_03.jsonl"
    ) as reader:
        for example in tqdm(reader):
            c += 1
            id = example["id"].split("_")[0]
            context = example["context"]
            orig_example = [sample for sample in squad_data if sample["id"] == id][0]
            # print(orig_example)
            orig_context = orig_example["context"]
            orig_question = orig_example["question"]
            orig_answer = orig_example["answers"]

            input = (
                f"{prompt} \nContext: {orig_context} \nQuestion: {orig_question} \nAnswer span: {orig_answer['text'][0]} "
                f"\n{prompt} \nContext: {context} "
                f"\nQuestion: \nAnswer span: "
            )

            # print(input)
            inputs = tokenizer(input, return_tensors="pt").input_ids.to("cuda")
            outputs = model.generate(
                inputs,
                max_new_tokens=50,
                top_p=1,
                top_k=40,
                temperature=0.7,
                do_sample=True,
                num_return_sequences=1,
            )
            result = {
                "id": example["id"],
                "question": tokenizer.decode(outputs[0], skip_special_tokens=True),
                "context": context,
            }

            examples.append(result)
            if c % 5000 == 0:
                save_to_disk(
                    examples,
                    f"{save_path}/counterfactual_questions_{model_identifier}_{c}.jsonl",
                )
                examples = []

        # save the remaining examples
        if examples:
            save_to_disk(
                examples,
                f"{save_path}/counterfactual_questions_{model_identifier}_{c}.jsonl",
            )

src/cf_generation/llm_generation/flan_ul2_qgen.py
UKPLab/CATfOOD
from typing import List

import torch


class LLaMA:
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer

    def generate(
        self,
        prompts: List[str],
        max_gen_len: int,
        temperature: float = 0.8,
        top_p: float = 0.95,
    ) -> List[str]:
        bsz = len(prompts)
        params = self.model.num_parameters()
        assert bsz <= params.max_batch_size, (bsz, params.max_batch_size)

        prompt_tokens = [self.tokenizer.encode(x, bos=True, eos=False) for x in prompts]

        min_prompt_size = min([len(t) for t in prompt_tokens])
        max_prompt_size = max([len(t) for t in prompt_tokens])

        total_len = min(params.max_seq_len, max_gen_len + max_prompt_size)

        tokens = torch.full((bsz, total_len), self.tokenizer.pad_id).cuda().long()
        for k, t in enumerate(prompt_tokens):
            tokens[k, : len(t)] = torch.tensor(t).long()
        input_text_mask = tokens != self.tokenizer.pad_id
        start_pos = min_prompt_size
        prev_pos = 0
        for cur_pos in range(start_pos, total_len):
            logits = self.model.forward(tokens[:, prev_pos:cur_pos], prev_pos)
            if temperature > 0:
                probs = torch.softmax(logits / temperature, dim=-1)
                next_token = sample_top_p(probs, top_p)
            else:
                next_token = torch.argmax(logits, dim=-1)
            next_token = next_token.reshape(-1)
            # only replace token if prompt has already been generated
            next_token = torch.where(
                input_text_mask[:, cur_pos], tokens[:, cur_pos], next_token
            )
            tokens[:, cur_pos] = next_token
            prev_pos = cur_pos

        decoded = []
        for i, t in enumerate(tokens.tolist()):
            # cut to max gen len
            t = t[: len(prompt_tokens[i]) + max_gen_len]
            # cut to eos tok if any
            try:
                t = t[: t.index(self.tokenizer.eos_id)]
            except ValueError:
                pass
            decoded.append(self.tokenizer.decode(t))
        return decoded


def sample_top_p(probs, p):
    probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
    probs_sum = torch.cumsum(probs_sort, dim=-1)
    mask = probs_sum - probs_sort > p
    probs_sort[mask] = 0.0
    probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
    next_token = torch.multinomial(probs_sort, num_samples=1)
    next_token = torch.gather(probs_idx, -1, next_token)
    return next_token

src/cf_generation/llm_generation/generation.py
UKPLab/CATfOOD
import os
import torch
import jsonlines
from tqdm import tqdm
from datasets import load_dataset

from src.cf_generation.llm_generation.utils import save_to_disk

import openai

openai.api_key = os.getenv("OPENAI_API_KEY")

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


def get_response(prompt, max_tokens, model_name):
    """
    :param prompt:
    :param max_tokens:
    :param model_name:
    :return:

    Sample output
    generated_ans =
    {
          "id": "chatcmpl-7dKbsNzkrAjNumSBMc4IhwCZAEMlo",
          "object": "chat.completion",
          "created": 1689608372,
          "model": "gpt-4-0613",
          "choices": [
            {
              "index": 0,
              "message": {
                "role": "assistant",
                "content": "The phenomena of electric and magnetic fields being perpendicular to each other is a
                fundamental part of how electromagnetic waves work. \n\nLet's consider an analogy to understand
                this better. ...},
              "finish_reason": "stop"
            }
          ],
          "usage": {
            "prompt_tokens": 222,
            "completion_tokens": 247,
            "total_tokens": 469
          }
    }

    """
    response = openai.ChatCompletion.create(
        model=model_name,
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": prompt},
        ],
        max_tokens=max_tokens,
    )
    return response


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()

    # Set the seed
    torch.manual_seed(args.seed)

    # load squad data
    dataset = load_dataset("squad", "plain_text")
    train_data = dataset["train"]
    squad_data = [
        sample
        for sample in tqdm(
            train_data, total=len(train_data), desc="Loading SQuAD data ... "
        )
    ]

    model_name = "gpt-4-0314"
    save_path = (
        BASE_PATH + f"src/data/squad/{model_name}_qa_relevance_seed_{args.seed}/"
    )
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    file_path = os.path.join(
        BASE_PATH,
        f"src/data/squad/counterfactual_samples_flan_t5_xxl_context_irrelevance.jsonl",
    )
    files = [file_path]
    skipped = 0
    c = 0
    for file_name in files:
        examples = []
        with jsonlines.open(file_name) as reader:
            for example in tqdm(reader):
                try:
                    c += 1
                    # if c <= 100:
                    #     continue
                    id = example["id"].split("_")[0]
                    context = example["context"]
                    question = example["question"]
                    answer = example["answers"]["text"][0]

                    # print("Given ans:", example["answers"])

                    orig_example = [
                        sample for sample in squad_data if sample["id"] == id
                    ][0]

                    orig_context = orig_example["context"]
                    orig_question = orig_example["question"]
                    orig_answer = orig_example["answers"]

                    # input = GRADE_DOCS_PROMPT_FAST.format(query=question, result=context, answer=answer)
                    # print(input)

                    template = (
                        f"Given the question: \n"
                        f"{question} \n"
                        f"Decide if the following retrieved context is relevant to the {answer}: \n"
                        f"{context} \n"
                        "Answer in the following format: \n"
                        "'Context is relevant: True or False.' \n"
                    )
                    # print(template)
                    # break

                    output = get_response(
                        prompt=template, max_tokens=10, model_name=model_name
                    )

                    result = {
                        "id": example["id"],
                        "question": question,
                        "context": context,
                        "answer": answer,
                        "context_relevance": output,
                    }
                    # print(result)
                    # break
                    examples.append(result)
                    if c == 500:
                        break
                except Exception as e:
                    print("Skip")
        # # save the remaining examples
        if examples:
            save_to_disk(
                examples,
                f"{save_path}counterfactual_samples_{model_name}_flan_t5_xxl_{c}.jsonl",
            )
        # print(examples)

src/cf_generation/llm_generation/gpt_evaluator.py
UKPLab/CATfOOD
import os

from transformers import AutoModelForCausalLM, T5Tokenizer, AutoTokenizer, AutoConfig
import torch
import jsonlines
from tqdm import tqdm
from datasets import load_dataset

# from src.few_shot.utils import save_to_disk

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"
# BASE_PATH = "/home/sachdeva/projects/ukp/exp_calibration/"


def save_to_disk(data, file_name):
    with jsonlines.open(file_name, "a") as writer:
        for example in tqdm(data, total=len(data), desc="Saving samples ... "):
            writer.write(example)


if __name__ == "__main__":

    # load squad data
    dataset = load_dataset("squad", "plain_text")
    train_data = dataset["train"]
    squad_data = [
        sample
        for sample in tqdm(
            train_data, total=len(train_data), desc="Loading SQuAD data ... "
        )
    ]

    c = 0
    examples = []

    # prompt = "Generate a fluent and answerable question from the given context. Ensure that the answer " \
    #          "is a span in the context and is less than 10 words."
    # prompt = "Given the context below, generate a fluent question that can be answered from it."  # does not work
    # prompt = "Using the provided context, generate a coherent question that can " \
    #          "be answered based on the information within the context."   # okay to use

    # prompt = "Given the context below, generate a question that can be answered using information " \
    #          "directly stated in the context."   # does not work

    # prompt = "Using the given context, generate a question that requires selecting a short " \
    #          "and specific answer from it."

    prompt = (
        "You are a question generation model. Given the context below, please generate a question that "
        "can be answered based on the information given in the context."
    )  # use this

    # prompt = "As a question generation model, your task is to generate a relevant and informative question " \
    #          "based on the given context. The question should be of type what, where, which, why, who, how, etc., " \
    #          "and should be answerable based on the information provided in the context. Please ensure that the " \
    #          "generated question is clear and concise, providing specific details that allow for an accurate " \
    #          "response. Additionally, please take into account any relevant background or contextual information " \
    #          "when generating your question."

    # prompt = "Your task is to create a question that can only be answered by identifying a specific " \
    #          "span of text within a given context. The question you generate should require the reader " \
    #          "to identify key details or information contained within the specified span of text. This " \
    #          "could include facts, figures, names, dates, events, or other relevant information. Please make " \
    #          "sure that the question is challenging but not impossible to answer based on the given context. " \
    #          "It should also be clear and concise in order to avoid confusion."

    # prompt = "As a question generation model, your task is to generate a relevant and answerable question " \
    #          "based on the given context. Please carefully read the context provided and craft a question that " \
    #          "can be answered using information from the text. Your question should be clear and concise, with " \
    #          "specific focus on key details or concepts within the context. You should also ensure that your " \
    #          "question is open-ended enough to allow for various possible answers while still being grounded in " \
    #          "the information presented in the context. Please note that you should aim to create questions that " \
    #          "encourage critical thinking and deeper understanding of the topic at hand."

    # prompt = "As a question generation model, your task is to generate a relevant and answerable question " \
    #          "based on the given context. Please carefully read the context provided and craft a question that " \
    #          "can be answered using information from the text. Your question should be clear and concise, with " \
    #          "specific focus on key details or concepts within the context. You should also ensure that your " \
    #          "question should allow one possible answer while still being grounded in " \
    #          "the information presented in the context. Please note that you should aim to create questions that " \
    #          "encourage critical thinking and deeper understanding of the topic at hand."  # nope

    # prompt = "Your task is to create a concise and clear question that can be answered from the given context. " \
    #          "The question should be fluent and answerable, with an answer span of less than 10 words. Please " \
    #          "ensure that the question is relevant to the context provided and encourages critical thinking or " \
    #          "analysis. Your response should demonstrate your ability to distill key information into a focused " \
    #          "and effective question."  # nope

    # prompt = "Your task is to generate a question that can be answered by an answer span within a " \
    #          "given context. The context should be a piece of text, such as a news article or historical " \
    #          "document, and the question should require understanding and analysis of the information presented " \
    #          "in the context. The generated question should be clear and concise, focusing on key details or events " \
    #          "described in the context. It should also be specific enough to have a single correct answer that " \
    #          "can be found within the context."   # good prompt but unanswerable questions

    # prompt = "As a question generation model, your task is to generate a question from a given context that" \
    #          " can be answered by a specific span of text highlighted in the context. The question should be " \
    #          "clear and concise, using appropriate language and grammar. It should also focus on extracting " \
    #          "relevant information from the highlighted span. Please note that you should be able to handle " \
    #          "various types of contexts, including those with complex sentence structures or technical jargon. " \
    #          "Your response should provide a well-formed question that accurately captures the meaning of the " \
    #          "highlighted span and encourages critical thinking."  # okayish

    # prompt = "Using the given context, create a question that is grammatically correct and coherent. The answer " \
    #          "to the question should be a specific span of text found within the context. Please ensure that your " \
    #          "question is clear and concise, and that it accurately reflects the information provided in the context." # okayish

    # prompt = "Create a question based on the context provided below that elicits a fluent and " \
    #          "concise answer from within the given text. Please provide a clear and specific question " \
    #          "that can be answered directly by a single span of text within the passage, without " \
    #          "requiring additional information or interpretation. The question should also be relevant " \
    #          "to the content of the passage and demonstrate an understanding of its key points."
    # Note " \
    #          "that your question should encourage creativity in generating unique and interesting " \
    #          "questions while still maintaining focus on accuracy and relevance."   # nope

    # Context:\n some text \nQuestion: What is the question?
    # prompt = "Generate a fluent and answerable question from the given context. Ensure that the answer " \
    #          "is a span in the context and is less than 10 words."

    # prompt = "You are a question generation model. Given the context below, please generate an answerable question " \
    #     "from it."

    # prompt = "Your task is to generate a fluent and coherent question based on the given context, " \
    #          "which should strictly be answerable by a span within the context. The question should " \
    #          "take inspiration from extractive question answering datasets like SQuAD. Please ensure " \
    #          "that your question is clear and concise, with a focus on accuracy and relevance. Your " \
    #          "question should also demonstrate an understanding of the key details and concepts presented " \
    #          "in the context, while encouraging critical thinking and analysis. Please note that your " \
    #          "response should allow for various relevant and creative questions."  # nope

    model_name = "EleutherAI/gpt-neox-20b"
    model_identifier = model_name.split("/")[-1]
    save_path = BASE_PATH + f"src/data/squad/{model_identifier}_qg"
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    config = AutoConfig.from_pretrained(model_name, pad_token_id=0, eos_token_id=0)
    model = AutoModelForCausalLM.from_pretrained(
        model_name, torch_dtype=torch.bfloat16, device_map="auto", config=config
    )
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    model.eval()
    skipped_instances = 0
    with jsonlines.open(
        BASE_PATH
        + "src/data/squad/squad_counterfactuals_noise_min_filtered_final_2.jsonl"
    ) as reader:
        for example in tqdm(reader):
            try:
                c += 1
                if c <= 35000:
                    continue
                # print(c)
                id = example["id"].split("_")[0]
                context = example["context"]
                orig_example = [sample for sample in squad_data if sample["id"] == id][
                    0
                ]
                # print(orig_example)
                orig_context = orig_example["context"]
                orig_question = orig_example["question"]
                orig_answer = orig_example["answers"]
                # context = "The Eiffel Tower ( ; ) is a wrought-iron lattice tower on the Champ de Mars in Paris, France. " \
                #           "It is named after the engineer Gustave Eiffel, whose company designed and " \
                #           "built the tower. Constructed from 1887–1889 as the entrance to the 1889 World's Fair, it " \
                #           "was initially criticized by some of France's leading artists and intellectuals for " \
                #           "its design, but it has become a global cultural icon of France and one of the most " \
                #           "recognisable structures in the world. The Eiffel Tower is the most-visited paid monument " \
                #           "in the world; 6.91 million people ascended it in 2015. The tower is tall, about the same " \
                #           "height as an 81- building, and the tallest structure in Paris. Its base is square, measuring " \
                #           "on each side. During its construction, the Eiffel Tower surpassed the Washington Monument " \
                #           "to become the tallest man-made structure in the world, a title it held for 41 years until " \
                #           "the Chrysler Building in New York City was finished in 1930. Due to the addition of a " \
                #           "broadcasting aerial at the top of the tower in 1957, it is now taller than the Chrysler " \
                #           "Building by. Excluding transmitters, the Eiffel Tower is the second tallest free-standing " \
                #           "structure in France after the Millau Viaduct. The tower has three levels for visitors, " \
                #           "with restaurants on the first and second levels."

                # input = f"{prompt} \nContext: {context} \nQuestion: "
                # input = f"{prompt} \nContext: {orig_context} \nQuestion: {orig_question} \nAnswer: {orig_answer['text'][0]} " \
                #         f"\n### \n{prompt} \nContext: {context} " \
                #         f"\nQuestion: \nAnswer: \n###"

                input = (
                    f"{prompt} \n\nContext: {context} \n\nQuestion: "  # use this one
                )

                # input = f"{prompt}\n Context: {context}\n Question: "

                # print(input)
                inputs = tokenizer(input, return_tensors="pt").to("cuda")
                generated_ids = model.generate(
                    inputs.input_ids,
                    attention_mask=inputs.attention_mask,
                    max_new_tokens=50,
                    # temperature=0.5,
                )
                outputs = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
                # print(outputs)

                sentences = outputs.split("\nQuestion: ")[1].split("\n\n")
                # remove empty strings from the list
                sentences = [sentence for sentence in sentences if sentence]
                question = sentences[0]
                # remove ``` from the question
                question = question.replace("```", "")
                # remove new line characters
                question = question.replace("\n", "")
                # print("Context: ", context)
                # print("Original outputs: ", outputs)
                # print("Actual question: ", question)
                # print("-"*100)

                # if c == 20:
                #     break
                # break
                result = {"id": example["id"], "question": question, "context": context}
                # print(result)
                examples.append(result)
                if c % 5000 == 0:
                    save_to_disk(
                        examples,
                        f"{save_path}/counterfactual_questions_{model_identifier}_{c}.jsonl",
                    )
                    examples = []

            except Exception as e:
                print(outputs)
                skipped_instances += 1
                continue

        # save the remaining examples
        if examples:
            save_to_disk(
                examples,
                f"{save_path}/counterfactual_questions_{model_identifier}_{c}.jsonl",
            )

src/cf_generation/llm_generation/gpt_neox.py
UKPLab/CATfOOD
import os

from transformers import AutoModelForCausalLM, T5Tokenizer, AutoTokenizer, AutoConfig
import torch
import jsonlines
from tqdm import tqdm
from datasets import load_dataset

# from src.few_shot.utils import save_to_disk

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"
# BASE_PATH = "/home/sachdeva/projects/ukp/exp_calibration/"


def save_to_disk(data, file_name):
    with jsonlines.open(file_name, "a") as writer:
        for example in tqdm(data, total=len(data), desc="Saving samples ... "):
            writer.write(example)


if __name__ == "__main__":

    # load squad data
    dataset = load_dataset("squad", "plain_text")
    train_data = dataset["train"]
    squad_data = [
        sample
        for sample in tqdm(
            train_data, total=len(train_data), desc="Loading SQuAD data ... "
        )
    ]

    c = 0
    examples = []

    # prompt = "Generate a fluent and answerable question from the given context. Ensure that the answer " \
    #          "is a span in the context and is less than 10 words."
    # prompt = "Given the context below, generate a fluent question that can be answered from it."  # does not work
    # prompt = "Using the provided context, generate a coherent question that can " \
    #          "be answered based on the information within the context."   # okay to use

    # prompt = "Given the context below, generate a question that can be answered using information " \
    #          "directly stated in the context."   # does not work

    # prompt = "Using the given context, generate a question that requires selecting a short " \
    #          "and specific answer from it."

    # prompt = "You are a question generation model. Given the context below, please generate a question that " \
    #          "can be answered based on the information given in the context."  # use this

    # prompt = "As a question generation model, your task is to generate a relevant and informative question " \
    #          "based on the given context. The question should be of type what, where, which, why, who, how, etc., " \
    #          "and should be answerable based on the information provided in the context. Please ensure that the " \
    #          "generated question is clear and concise, providing specific details that allow for an accurate " \
    #          "response. Additionally, please take into account any relevant background or contextual information " \
    #          "when generating your question."

    # prompt = "Your task is to create a question that can only be answered by identifying a specific " \
    #          "span of text within a given context. The question you generate should require the reader " \
    #          "to identify key details or information contained within the specified span of text. This " \
    #          "could include facts, figures, names, dates, events, or other relevant information. Please make " \
    #          "sure that the question is challenging but not impossible to answer based on the given context. " \
    #          "It should also be clear and concise in order to avoid confusion."

    # prompt = "As a question generation model, your task is to generate a relevant and answerable question " \
    #          "based on the given context. Please carefully read the context provided and craft a question that " \
    #          "can be answered using information from the text. Your question should be clear and concise, with " \
    #          "specific focus on key details or concepts within the context. You should also ensure that your " \
    #          "question is open-ended enough to allow for various possible answers while still being grounded in " \
    #          "the information presented in the context. Please note that you should aim to create questions that " \
    #          "encourage critical thinking and deeper understanding of the topic at hand."

    # prompt = "As a question generation model, your task is to generate a relevant and answerable question " \
    #          "based on the given context. Please carefully read the context provided and craft a question that " \
    #          "can be answered using information from the text. Your question should be clear and concise, with " \
    #          "specific focus on key details or concepts within the context. You should also ensure that your " \
    #          "question should allow one possible answer while still being grounded in " \
    #          "the information presented in the context. Please note that you should aim to create questions that " \
    #          "encourage critical thinking and deeper understanding of the topic at hand."  # nope

    # prompt = "Your task is to create a concise and clear question that can be answered from the given context. " \
    #          "The question should be fluent and answerable, with an answer span of less than 10 words. Please " \
    #          "ensure that the question is relevant to the context provided and encourages critical thinking or " \
    #          "analysis. Your response should demonstrate your ability to distill key information into a focused " \
    #          "and effective question."  # nope

    # prompt = "Your task is to generate a question that can be answered by an answer span within a " \
    #          "given context. The context should be a piece of text, such as a news article or historical " \
    #          "document, and the question should require understanding and analysis of the information presented " \
    #          "in the context. The generated question should be clear and concise, focusing on key details or events " \
    #          "described in the context. It should also be specific enough to have a single correct answer that " \
    #          "can be found within the context."   # good prompt but unanswerable questions

    # prompt = "As a question generation model, your task is to generate a question from a given context that" \
    #          " can be answered by a specific span of text highlighted in the context. The question should be " \
    #          "clear and concise, using appropriate language and grammar. It should also focus on extracting " \
    #          "relevant information from the highlighted span. Please note that you should be able to handle " \
    #          "various types of contexts, including those with complex sentence structures or technical jargon. " \
    #          "Your response should provide a well-formed question that accurately captures the meaning of the " \
    #          "highlighted span and encourages critical thinking."  # okayish

    # prompt = "Using the given context, create a question that is grammatically correct and coherent. The answer " \
    #          "to the question should be a specific span of text found within the context. Please ensure that your " \
    #          "question is clear and concise, and that it accurately reflects the information provided in the context." # okayish

    # prompt = "Create a question based on the context provided below that elicits a fluent and " \
    #          "concise answer from within the given text. Please provide a clear and specific question " \
    #          "that can be answered directly by a single span of text within the passage, without " \
    #          "requiring additional information or interpretation. The question should also be relevant " \
    #          "to the content of the passage and demonstrate an understanding of its key points."
    # Note " \
    #          "that your question should encourage creativity in generating unique and interesting " \
    #          "questions while still maintaining focus on accuracy and relevance."   # nope

    # Context:\n some text \nQuestion: What is the question?
    # prompt = "Generate a fluent and answerable question from the given context. Ensure that the answer " \
    #          "is a span in the context and is less than 10 words."

    # prompt = "You are a question generation model. Given the context below, please generate an answerable question " \
    #     "from it."

    # prompt = "Your task is to generate a fluent and coherent question based on the given context, " \
    #          "which should strictly be answerable by a span within the context. The question should " \
    #          "take inspiration from extractive question answering datasets like SQuAD. Please ensure " \
    #          "that your question is clear and concise, with a focus on accuracy and relevance. Your " \
    #          "question should also demonstrate an understanding of the key details and concepts presented " \
    #          "in the context, while encouraging critical thinking and analysis. Please note that your " \
    #          "response should allow for various relevant and creative questions."  # nope

    prompt = "Given the context please answer the question."

    model_name = "EleutherAI/gpt-neox-20b"
    model_identifier = model_name.split("/")[-1]
    save_path = BASE_PATH + f"src/data/squad/{model_identifier}_qg"
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    config = AutoConfig.from_pretrained(model_name, pad_token_id=0, eos_token_id=0)
    model = AutoModelForCausalLM.from_pretrained(
        model_name, torch_dtype=torch.bfloat16, device_map="auto", config=config
    )
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    model.eval()
    skipped_instances = 0
    with jsonlines.open(
        BASE_PATH
        + "src/data/squad/counterfactual_samples_Llama-2-13b-chat-hf_gpt_neox_context_filtered_complete.jsonl"
    ) as reader:
        for example in tqdm(reader):
            try:
                c += 1
                if c == 20:
                    break
                # print(c)
                id = example["id"].split("_")[0]
                context = example["context"]
                question = example["question"]
                orig_example = [sample for sample in squad_data if sample["id"] == id][
                    0
                ]
                # print(orig_example)
                orig_context = orig_example["context"]
                orig_question = orig_example["question"]
                orig_answer = orig_example["answers"]

                input = f"{prompt} \nContext: {context} \nQuestion: {question} \nAnswer: "  # use this one

                # input = f"{prompt}\n Context: {context}\n Question: "

                # print(input)
                inputs = tokenizer(input, return_tensors="pt").to("cuda")
                generated_ids = model.generate(
                    inputs.input_ids,
                    attention_mask=inputs.attention_mask,
                    max_new_tokens=20,
                    # temperature=0.5,
                )
                outputs = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
                answer = outputs[len(input) :].strip().split("\n")[0]
                print("pred: ", answer)
                print("orig: ", example["answers"])
                print("-" * 100)

                # sentences = outputs.split("\nQuestion: ")[1].split("\n\n")
                # # remove empty strings from the list
                # sentences = [sentence for sentence in sentences if sentence]
                # question = sentences[0]
                # # remove ``` from the question
                # question = question.replace("```", "")
                # # remove new line characters
                # question = question.replace("\n", "")
                # print("Context: ", context)
                # print("Original outputs: ", outputs)
                # print("Actual question: ", question)
                # print("-"*100)

                # if c == 20:
                #     break
                # break
                # result = {
                #     "id": example["id"],
                #     "question": question,
                #     "context": context,
                #     "answer": answer
                # }
                # # print(result)
                # examples.append(result)
                # if c % 5000 == 0:
                #     save_to_disk(
                #         examples,
                #         f"{save_path}/counterfactual_questions_{model_identifier}_{c}.jsonl"
                #     )
                #     examples = []

            except Exception as e:
                # print(outputs)
                skipped_instances += 1
                continue

        # save the remaining examples
        # if examples:
        #     save_to_disk(
        #         examples,
        #         f"{save_path}/counterfactual_questions_{model_identifier}_{c}.jsonl"
        #     )

src/cf_generation/llm_generation/gpt_neox_qa.py
UKPLab/CATfOOD
import os

from transformers import AutoModelForCausalLM, T5Tokenizer, AutoTokenizer
import torch
import jsonlines
from tqdm import tqdm
from datasets import load_dataset

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"
# BASE_PATH = "/home/sachdeva/projects/ukp/exp_calibration/"


def save_to_disk(data, file_name):
    with jsonlines.open(file_name, "a") as writer:
        for example in tqdm(data, total=len(data), desc="Saving samples ... "):
            writer.write(example)


if __name__ == "__main__":

    # load squad data
    dataset = load_dataset("squad", "plain_text")
    train_data = dataset["train"]
    squad_data = [
        sample
        for sample in tqdm(
            train_data, total=len(train_data), desc="Loading SQuAD data ... "
        )
    ]

    c = 0
    examples = []

    prompt = (
        "Generate a fluent and answerable question from the given context. Ensure that the answer "
        "is a span in the context and is less than 10 words."
    )

    model_name = "EleutherAI/gpt-j-6B"
    model_identifier = model_name.split("/")[-1]
    save_path = BASE_PATH + f"src/data/squad/{model_identifier}_qg"
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    model = AutoModelForCausalLM.from_pretrained(
        model_name, torch_dtype=torch.bfloat16, device_map="auto"
    )
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    with jsonlines.open(
        BASE_PATH
        + "src/data/squad/squad_counterfactuals_noise_min_filtered_final_2.jsonl"
    ) as reader:
        for example in tqdm(reader):
            c += 1
            id = example["id"].split("_")[0]
            context = example["context"]
            orig_example = [sample for sample in squad_data if sample["id"] == id][0]
            # print(orig_example)
            orig_context = orig_example["context"]
            orig_question = orig_example["question"]
            orig_answer = orig_example["answers"]

            input = (
                f"{prompt} \nContext: {orig_context} \nQuestion: {orig_question} "
                f"\n{prompt} \nContext: {context} "
                f"\nQuestion: "
            )

            # print(input)
            inputs = tokenizer(input, return_tensors="pt").input_ids.to("cuda")
            outputs = model.generate(inputs, max_length=200)
            result = {
                "id": example["id"],
                "question": tokenizer.decode(outputs[0], skip_special_tokens=True),
                "context": context,
            }
            # print(result)
            examples.append(result)
            if c % 5000 == 0:
                save_to_disk(
                    examples,
                    f"{save_path}/counterfactual_questions_{model_identifier}_{c}.jsonl",
                )
                examples = []

        # save the remaining examples
        if examples:
            save_to_disk(
                examples,
                f"{save_path}/counterfactual_questions_{model_identifier}_{c}.jsonl",
            )

src/cf_generation/llm_generation/gptj_qgen.py
UKPLab/CATfOOD
import os
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

import jsonlines
from tqdm import tqdm
from datasets import load_dataset

from src.cf_generation.llm_generation.utils import save_to_disk

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"

# Set the seed
seed = 0
torch.manual_seed(seed)
device = "cuda" if torch.cuda.is_available() else "cpu"

model_name = "Llama-2-13b-chat-hf"

tokenizer = AutoTokenizer.from_pretrained(f"meta-llama/{model_name}")
model = AutoModelForCausalLM.from_pretrained(
    f"meta-llama/{model_name}"
)  # , torch_dtype=torch.bfloat16)
model.to(device)


def get_response(prompt):
    inputs = tokenizer(prompt, return_tensors="pt", return_token_type_ids=False).to(
        device
    )
    outputs = model.generate(**inputs, max_new_tokens=10,)
    output = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return output


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()

    # Set the seed
    torch.manual_seed(args.seed)

    # load squad data
    dataset = load_dataset("squad", "plain_text")
    train_data = dataset["train"]
    squad_data = [
        sample
        for sample in tqdm(
            train_data, total=len(train_data), desc="Loading SQuAD data ... "
        )
    ]
    test_model = "alpaca_13b"
    save_path = (
        BASE_PATH
        + f"src/data/squad/{model_name}_qa_relevance_{test_model}_seed_{args.seed}/"
    )
    # save_path = BASE_PATH + f"src/data/squad/"
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    file_path = os.path.join(
        BASE_PATH,
        # f"src/data/squad/gpt-4-0314_qa_relevance_seed_{args.seed}/counterfactual_samples_gpt-4-0314_llama_500.jsonl"
        f"src/data/squad/counterfactual_data_{test_model}_v2_qg_pipeline_all_data_cleaned.jsonl",
    )
    files = [file_path]
    skipped = 0
    c = 0
    for file_name in files:
        examples = []
        with jsonlines.open(file_name) as reader:
            for example in tqdm(reader):
                try:
                    c += 1

                    id = example["id"].split("_")[0]
                    context = example["context"]
                    question = example["question"]
                    answer = example["answers"]["text"][0]
                    orig_example = [
                        sample for sample in squad_data if sample["id"] == id
                    ][0]

                    orig_context = orig_example["context"]
                    orig_question = orig_example["question"]
                    orig_answer = orig_example["answers"]

                    # prompt = \
                    #     "You are a helpful assistant. Please think critically. "  \
                    #     f"Given the question: \n" \
                    #     f"{question} \n" \
                    #     f"Decide if the following retrieved context is relevant to answer this question: \n" \
                    #     f"{context} \n" \
                    #     "Answer in the following format: \n" \
                    #     "'Context is relevant: False or True.' \n" \
                    #     "Please answer False if you think that the question cannot be answered from the context. \n" \
                    #     "Assistant: Context is relevant: "

                    # prompt = \
                    #     f"Given the question: \n" \
                    #     f"{question} \n" \
                    #     f"Decide if the following retrieved context is relevant to answer the question: \n" \
                    #     f"{context} \n" \
                    #     "Your answer should be either True or False." \
                    #     "Assistant: "

                    # 2nd BEST TILL NOW (~75% accuracy)
                    ############################################################################################
                    # system_prompt = "You are a language/context evaluator. Your job is to identify whether " \
                    #                 "a given context is relevant for answering a given question."
                    #
                    # prompt = \
                    #     f"Given the question: \n" \
                    #     f"{question} \n" \
                    #     f"Decide if the following retrieved context is relevant to the {answer}: \n" \
                    #     f"{context} \n" \
                    #     "Your answer should be either True or False." \
                    #     "Assistant: "

                    #############################################################################################

                    # 2nd place
                    ############################################################################################
                    # system_prompt = "You are a language/context evaluator. Your job is to identify whether " \
                    #                 "a given context is relevant for answering a given question."
                    #
                    # prompt = \
                    #     f"Given the question: \n" \
                    #     f"{question} \n" \
                    #     f"Decide if the following retrieved context is relevant to the question: \n" \
                    #     f"{context} \n" \
                    #     "Your answer should be either True or False." \
                    #     "Assistant: "

                    #############################################################################################

                    # BEST TILL NOW(~81 % accuracy)
                    ############################################################################################
                    system_prompt = (
                        "You are a language/context evaluator. Your job is to identify whether "
                        "a given context is relevant for answering a given question."
                    )

                    task_prompt = (
                        f"Given the question: \n"
                        f"{question} \n"
                        f"Decide if the following retrieved context is relevant to the {answer}: \n"
                        f"{context} \n"
                        "Answer in the following format:\n True or False."
                        "Assistant: "
                    )

                    #############################################################################################

                    template = (
                        f"[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n"
                        f"{task_prompt}"
                        "[/INST]"
                    )

                    # print(template)
                    # break
                    # print("GPT-4 answer: ", example["context_relevance"]["choices"][0]["message"]["content"])
                    output = get_response(prompt=template)
                    output = output[len(template) :].strip().split("\n")[0]
                    # print("LLAMA answer: ", output)
                    # print("----"*10)

                    result = {
                        "id": example["id"],
                        "question": question,
                        "context": context,
                        "answer": answer,
                        "context_relevance": output,
                    }
                    # print(result)
                    # break
                    examples.append(result)
                    if c % 5000 == 0:
                        save_to_disk(
                            examples,
                            f"{save_path}counterfactual_samples_{model_name}_{test_model}_{c}.jsonl",
                        )
                        examples = []
                        print(f"----Saved {c} samples----")
                    # if c == 200:
                    #     break
                except Exception as e:
                    print("Skip")
        # save the remaining examples
        if examples:
            save_to_disk(
                examples,
                f"{save_path}counterfactual_samples_{model_name}_{test_model}_{c}.jsonl",
            )

src/cf_generation/llm_generation/llama2.py
UKPLab/CATfOOD
import os
import jsonlines
from tqdm import tqdm
import torch

from typing import Optional

from datasets import load_dataset
from langchain import PromptTemplate, FewShotPromptTemplate
from transformers import AutoModelForCausalLM, T5Tokenizer, AutoTokenizer, AutoConfig
from transformers import GenerationConfig, LlamaForCausalLM, LlamaTokenizer

# from src.few_shot.utils import save_to_disk
# from src.few_shot.together import infer

# Set the seed
seed = 0
torch.manual_seed(seed)

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


def save_to_disk(data, file_name):
    with jsonlines.open(file_name, "a") as writer:
        for example in tqdm(data, total=len(data), desc="Saving samples ... "):
            writer.write(example)


def get_llama(model):
    def skip(*args, **kwargs):
        pass

    torch.nn.init.kaiming_uniform_ = skip
    torch.nn.init.uniform_ = skip
    torch.nn.init.normal_ = skip

    model = LlamaForCausalLM.from_pretrained(model, torch_dtype="auto")
    model.seqlen = 2048
    return model


class LLMClient:
    def __init__(
        self,
        template,
        model_name_or_path: str = None,
        tokenizer_name_or_path: Optional[str] = None,
        data_path: str = None,
        threshold: float = 0.5,
        task: str = "qg",
        max_new_tokens: int = 50,
        temperature: float = 0.01,
        top_p: float = 1,
        top_k: int = 40,
        repetition_penalty: float = 1.0,
        save_results: bool = True,
        max_samples: int = None,
        stop: str = "\n",
    ):
        self.model = model_name_or_path
        self.tokenizer = tokenizer_name_or_path
        self.max_new_tokens = max_new_tokens
        self.temperature = temperature
        self.top_p = top_p
        self.top_k = top_k
        self.repetition_penalty = repetition_penalty
        self.data_path = data_path
        self.task = task
        self.threshold = threshold
        self.save_results = save_results
        self.max_samples = max_samples
        self.stop = stop.split(";")
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

    def _load_model(self):
        model = get_llama(self.model).to(self.device)
        tokenizer = LlamaTokenizer.from_pretrained(self.model)
        return model, tokenizer

    def _create_few_shot_prompt(self, examples, context, question=None):
        example_template = """
Context: {context}
Question: {question}
Answer: {answer}
""".strip()

        suffix = """
Context: {context}
Question: 
""".strip()
        example_prompt = PromptTemplate(
            input_variables=["context", "question", "answer"],
            template=example_template,
        )

        if self.task == "qg":
            few_shot_template = FewShotPromptTemplate(
                examples=examples,
                example_prompt=example_prompt,
                prefix=qg_template if self.task == "qg" else qa_template,
                suffix=suffix,
                input_variables=["context"],
                example_separator="\n",
            )
            prompt = few_shot_template.format(context=context)
        else:
            few_shot_template = FewShotPromptTemplate(
                examples=examples,
                example_prompt=example_prompt,
                prefix=qa_template,
                suffix=suffix,
                input_variables=["context", "question"],
                example_separator="\n",
            )
            prompt = few_shot_template.format(context=context, question=question)
        return prompt

    def generate(self):
        c = 0
        skipped_instances = 0
        examples = []
        model, tokenizer = self._load_model()
        model_identifier = self.model.split("/")[-1]
        save_path = (
            BASE_PATH
            + f"src/data/squad/t5_squad_counterfactuals/{model_identifier}_{self.task}_pipeline_remaining_temp_0.7"
        )
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        # load squad data
        dataset = load_dataset("squad", "plain_text")
        train_data = dataset["train"]
        squad_data = [
            sample
            for sample in tqdm(
                train_data, total=len(train_data), desc="Loading SQuAD data ... "
            )
        ]

        current_files = []
        file_names = []
        file_path = BASE_PATH + f"src/data/squad/{self.data_path}"
        if self.task == "qg":
            file_names = [file_path]
        elif self.task == "qa":
            all_files = os.listdir(file_path)
            file_names = [
                os.path.join(file_path, file)
                for file in all_files
                if file not in current_files
            ]

        model.eval()
        for file_name in file_names:
            with jsonlines.open(file_name) as reader:
                for example in tqdm(reader):
                    try:
                        id = example["id"].split("_")[0]
                        context = example["context"]
                        orig_example = [
                            sample for sample in squad_data if sample["id"] == id
                        ][0]
                        # print(orig_example)
                        orig_context = orig_example["context"]
                        orig_question = orig_example["question"]
                        orig_answer = orig_example["answers"]
                        c += 1
                        #
                        if c <= 28000:
                            continue

                        if self.max_samples:
                            if c == self.max_samples:
                                break
                        if self.task == "qa":
                            # preprocess the question
                            # remove any new line characters
                            question = example["question"]
                            question = question.replace("\n", "")
                            # check for truncated question words
                            # check for truncated question words
                            question_words = [
                                "what",
                                "who",
                                "when",
                                "where",
                                "why",
                                "how",
                                "which",
                                "whom",
                                "whose",
                            ]
                            for word in question_words:
                                if question.lower().startswith(word[1:]):
                                    # replace first word with the correct word
                                    question = word + question[len(word[1:]) :]
                                    # question = word + question[3:]
                                    break

                        prompt = self._create_few_shot_prompt(
                            [
                                {
                                    "context": orig_context,
                                    "question": orig_question,
                                    "answer": orig_answer["text"][0],
                                }
                            ],
                            context,
                            question=question if self.task == "qa" else None,
                        )
                        # print(prompt)
                        inputs = tokenizer(prompt, return_tensors="pt").to(self.device)
                        outputs = model.generate(
                            **inputs,
                            max_new_tokens=self.max_new_tokens,
                            temperature=self.temperature,
                            top_p=self.top_p,
                            top_k=self.top_k,
                            repetition_penalty=self.repetition_penalty,
                            do_sample=True,
                            # num_return_sequences=1,
                            early_stopping=True,
                            pad_token_id=tokenizer.eos_token_id,
                        )
                        output = tokenizer.decode(outputs[0], skip_special_tokens=True)
                        # remove the context from the output
                        output = output[len(prompt) :]

                        # output = infer(
                        #     model_name="GPT-JT-6B-v1",
                        #     prompt=prompt,
                        # )

                        result = {"id": example["id"], "context": context}

                        if self.task == "qg":
                            # add the question mark if it is not present
                            # if output[-1] != "?":
                            #     output = output + "?"
                            # question = output

                            # for stop_word in self.stop:
                            #     if stop_word != '' and stop_word in output:
                            #         question = output[:output.find(stop_word)]
                            #     # in some cases the stop word is at the beginning of the output
                            #     if question == "":
                            #         question = output[output.find(stop_word):]
                            # Split the text into lines
                            lines = output.split("\n")

                            # Extract the question and answer
                            question = lines[0].replace("Question: ", "")
                            question = question.strip()
                            answer = lines[1].replace("Answer: ", "")
                            answer = answer.replace("<bot>:", "")
                            answer = answer.strip()

                            result["question"] = question
                            result["answer"] = answer

                        elif self.task == "qa":
                            answer = output
                            # print("Output:", output)
                            for stop_word in self.stop:
                                if stop_word != "" and stop_word in output:
                                    answer = output[output.find(stop_word) :]
                                    # remove the first stop word from answer
                                    answer = answer.split(stop_word)
                                    answer = [a for a in answer if a != ""]
                                    if answer:
                                        answer = answer[0]
                                    else:
                                        answer = "I don't know."

                                # in some cases the stop word is at the beginning of the output
                            #     if answer == "":
                            #         answer = output[output.find(stop_word):]
                            #         answer = answer.replace(stop_word, "")
                            # if answer == "":
                            #     answer = answer.split(stop_word)
                            #     answer = [a for a in answer if a != ""][0]
                            skip_words = [
                                "A: ",
                                "A. ",
                                "Answer: ",
                                "Output: ",
                                "Label: ",
                            ]
                            for skip_word in skip_words:
                                if skip_word in answer:
                                    answer = answer.replace(skip_word, "")

                            # if answer.startswith("A: ") or answer.startswith("A. "):
                            #     answer = answer[3:]
                            # elif answer.startswith("Answer: "):
                            #     answer = answer.replace("Answer: ", "")
                            # question = example["question"]
                            result["question"] = question
                            result["answer"] = answer

                        # print("Context:", context)
                        # print("Question:", question)
                        # print("Answer:", answer)
                        # print("-" * 100)

                        examples.append(result)
                        if self.save_results:
                            if c % self.threshold == 0:
                                save_to_disk(
                                    examples,
                                    f"{save_path}/counterfactual_samples_{model_identifier}_{c}.jsonl",
                                )
                                examples = []
                    except Exception as e:
                        # print(outputs)
                        skipped_instances += 1
                        print(f"Skipped instance {c} due to error: {e}.")
                        continue

                    # if c==5:
                    #     break

            # save the remaining examples
            if self.save_results:
                if examples:
                    save_to_disk(
                        examples,
                        f"{save_path}/counterfactual_samples_{model_identifier}_{c}.jsonl",
                    )
            # break


if __name__ == "__main__":

    qg_template = """
As a question generator, your task is to create a concise and clear question that can be answered by an answer span within a given context. 
The context should be a piece of text, such as a news article or historical document, and the question should require understanding and analysis of the information presented in the context. 
Your generated question should focus on key details or events described in the context and should demonstrate your ability to identify important information. 
Additionally, please ensure that your question is specific enough to have a single correct answer within the given context. 
Please note that you may need to read through the provided context multiple times to fully understand its contents before generating an appropriate question. 
""".strip()

    chatgpt_qg_template = """
As a question generator, your task is to create a clear and concise question that can be answered by an answer span within a given context. 
The context should be a piece of text, such as a news article or historical document, and the question should require understanding and analysis of the information presented in the context. 
Your generated question should focus on key details or events described in the context, requiring readers to carefully read and analyze the provided text. 
Please ensure that your question is specific enough to have only one correct answer span within the given context. 
Please note that you should aim for clarity and concision while still maintaining accuracy and relevance to the provided text.
""".strip()

    cf_qa_template = """
Given the context and question, generate a question that negates the question.
""".strip()

    qa_template = """
As a language model, your task is to generate an answer that is strictly a span in the given context provided with a specific question.
Please provide a concise and relevant response that directly answers the question asked within the context.
Your answer should be accurate and reflect the information presented in the context.
Please note that you may need to use contextual clues and understanding to identify the most appropriate section of text to use as your answer span.


Context: {context}


Question: {question}


Answer: 
"""

    task = "qg"
    model = "llama-13b-hf"
    client = LLMClient(
        template=qg_template if task == "qg" else qa_template,
        model_name_or_path=model,
        tokenizer_name_or_path=model,
        data_path="t5_squad_counterfactuals/rag_counterfactuals_complete_min_filtered_dedup.jsonl",
        threshold=1000,
        task=task,
        max_new_tokens=50,
        temperature=0.7,
        top_p=1,
        top_k=40,
        repetition_penalty=1.0,
        save_results=True,
        max_samples=None,
    )

    client.generate()

src/cf_generation/llm_generation/llama_prompter.py
UKPLab/CATfOOD
import os
import torch
from transformers import GenerationConfig, LlamaForCausalLM, LlamaTokenizer
from transformers import T5ForConditionalGeneration, T5Tokenizer, AutoTokenizer
from langchain import PromptTemplate
import jsonlines
import re
from tqdm import tqdm
from datasets import load_dataset
from transformers import AutoModelForCausalLM, T5Tokenizer, AutoTokenizer, AutoConfig

from src.few_shot.utils import save_to_disk

# from src.few_shot.generation import LLaMA

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


# Set the seed
seed = 0
torch.manual_seed(seed)


def get_llama(model):
    def skip(*args, **kwargs):
        pass

    torch.nn.init.kaiming_uniform_ = skip
    torch.nn.init.uniform_ = skip
    torch.nn.init.normal_ = skip

    model = LlamaForCausalLM.from_pretrained(model, torch_dtype="auto")
    model.seqlen = 2048
    return model


from transformers import T5ForConditionalGeneration, T5Tokenizer, AutoTokenizer
import torch
import os
import jsonlines
import argparse
from tqdm import tqdm
from datasets import load_dataset

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"
# BASE_PATH = "/home/sachdeva/projects/ukp/exp_calibration/"


def save_to_disk(data, file_name):
    with jsonlines.open(file_name, "a") as writer:
        for example in tqdm(data, total=len(data), desc="Saving samples ... "):
            writer.write(example)


if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()

    # Set the seed
    torch.manual_seed(args.seed)

    # load squad data
    dataset = load_dataset("squad", "plain_text")
    train_data = dataset["train"]
    squad_data = [
        sample
        for sample in tqdm(
            train_data, total=len(train_data), desc="Loading SQuAD data ... "
        )
    ]

    current_files = []

    c = 0
    # prompt = "Given the following question and context, provide an answer " \
    #          "that best fits the question. Ensure that the answer " \
    #          "is a span in the context."
    prompt = (
        "Answer the question based on the context below. If the question cannot be answered "
        "using the information provided, then answer with 'I don't know'."
    )

    # to test
    # prompt = "Your task is to answer a question based on the given context. If the information provided " \
    #          "is insufficient to answer the question, please respond with 'I don't know.' Your response " \
    #          "should be clear and concise, providing only relevant information necessary to answer the question. " \
    #          "Please note that you should make every effort to provide an accurate and complete response based " \
    #          "on the available information."

    model_name = "llama-13b-hf"
    # model_name = "google/flan-ul2"
    model_identifier = model_name.split("/")[-1]

    save_path = (
        BASE_PATH
        + f"src/data/squad/few_shot_{model_identifier}_qa_eval_seed_{args.seed}"
    )
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    model = get_llama(model_name).to("cuda")
    tokenizer = LlamaTokenizer.from_pretrained(model_name)

    # all_files = os.listdir(BASE_PATH + f"src/data/squad/few_shot_{model_identifier}_qg_temp_0.7/")
    # files = [file for file in all_files if file not in current_files]
    # print(files)
    # files = ["llama_collated_data_with_answers_processed_context_irrelevance.jsonl"]
    files = ["counterfactual_data_llama_13b_v1_qg_pipeline_all_data_cleaned.jsonl"]

    for file in files:

        # if c<=9:
        #     continue
        print(f"Processing file: {file}")
        examples = []
        with jsonlines.open(BASE_PATH + f"src/data/squad/" + file) as reader:
            # with jsonlines.open(BASE_PATH + f"src/data/squad/few_shot_{model_identifier}_qg_temp_0.7/" + file) as reader:
            for example in tqdm(reader):
                c += 1
                id = example["id"].split("_")[0]
                context = example["context"]
                question = example["question"]
                tokens_to_remove = ["[", "'", '"', "]"]
                # Create a translation table that maps each unwanted token to None
                translator = str.maketrans({token: None for token in tokens_to_remove})
                question = question.translate(translator).strip()

                orig_example = [sample for sample in squad_data if sample["id"] == id][
                    0
                ]

                orig_context = orig_example["context"]
                orig_question = orig_example["question"]
                orig_answer = orig_example["answers"]

                input = (
                    f"{prompt} \nContext: {orig_context} \nQuestion: {orig_question} "
                    f"\nAnswer: {orig_answer['text'][0]}"
                    f"\n{prompt} \nContext: {context} \nQuestion: {question} "
                    f"\nAnswer: "
                )

                inputs = tokenizer(input, return_tensors="pt").input_ids.to("cuda")
                outputs = model.generate(
                    inputs, max_new_tokens=50, temperature=0.7, top_p=1, top_k=40,
                )
                output = tokenizer.decode(outputs[0], skip_special_tokens=True)
                # remove the context from the output
                output = output[len(input) :]
                answer = output.split("\n")[0]
                result = {
                    "id": example["id"],
                    "question": question,
                    "context": context,
                    "answer": answer,
                }
                print("question: ", question)
                print("context: ", context)
                print("answer: ", result["answer"])
                print("original answer: ", example["answers"])
                print("-" * 10)
                examples.append(result)

                if c == 10:
                    break

        # save_to_disk(examples, f"{save_path}/counterfactual_samples_{model_identifier}_{c}.jsonl")
        # c += 1

src/cf_generation/llm_generation/llama_qa.py
UKPLab/CATfOOD
import os
import torch
from transformers import LLaMAForCausalLM
from transformers.models.llama.tokenization_llama import LLaMATokenizer

import jsonlines
from tqdm import tqdm
from datasets import load_dataset

from src.few_shot.utils import save_to_disk

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


def get_llama(model):
    def skip(*args, **kwargs):
        pass

    torch.nn.init.kaiming_uniform_ = skip
    torch.nn.init.uniform_ = skip
    torch.nn.init.normal_ = skip

    model = LLaMAForCausalLM.from_pretrained(model, torch_dtype="auto")
    model.seqlen = 2048
    return model


if __name__ == "__main__":
    # load squad data
    dataset = load_dataset("squad", "plain_text")
    train_data = dataset["train"]
    squad_data = [
        sample
        for sample in tqdm(
            train_data, total=len(train_data), desc="Loading SQuAD data ... "
        )
    ]

    c = 0
    examples = []

    prompt = (
        "Generate a fluent and answerable question from the given context. Ensure that the answer "
        "is a span in the context and is less than 10 words."
    )
    # prompt = "Please generate a question that can be answered from the given context. Here is an example:" \

    device = torch.device("cuda:0")
    model_name = "llama-13b-hf"
    model = get_llama(model_name).to(device)
    tokenizer = LLaMATokenizer.from_pretrained(model_name)
    # print(model.num_parameters())

    model_identifier = model_name.split("/")[-1]
    save_path = BASE_PATH + f"src/data/squad/{model_identifier}_qg_1/"
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    """
    How are you doing today?

    I asked **LlaMA**: "How are you doing today?"
    And got this hilarious reply

    How are you doing today? Good to see you!
    I know, we haven't seen each other in awhile, but that's okay. Things are different now.
    I am older and wiser, and you are older and wiser. We have both grown so much in the past few years,
    and we have so much to share with each other. I really want us to be friends.
    I know, you're thinking that I'm just another girl who wants to get close to you for what you can
    give me. I do want to be close to you, but I don't want anything from you. I have my own life, my own
    career, my own friends, and my own family. I have everything I want and need. I'm not looking for you
    to make me happy, because I know that I can do that all by myself. I just want to have a good
    relationship with you, and I think that we can

    """
    skipped_instances = 0
    model.eval()
    with jsonlines.open(
        BASE_PATH
        + "src/data/squad/squad_counterfactuals_noise_min_filtered_final_2.jsonl"
    ) as reader:
        for example in tqdm(reader):
            try:
                c += 1
                if c <= 60000:
                    continue
                id = example["id"].split("_")[0]
                context = example["context"]
                orig_example = [sample for sample in squad_data if sample["id"] == id][
                    0
                ]
                # print(orig_example)
                orig_context = orig_example["context"]
                orig_question = orig_example["question"]
                orig_answer = orig_example["answers"]

                input = f"{prompt} \n\nC: {context} \n\nQ: "

                input_ids = tokenizer.encode(input, return_tensors="pt").to(device)
                with torch.no_grad():
                    generated_ids = model.generate(
                        input_ids,
                        # do_sample=True,
                        # max_length=1024,
                        max_new_tokens=50,
                        # top_p=0.95,
                        # temperature=0.8,
                    )
                outputs = tokenizer.decode([el.item() for el in generated_ids[0]])
                sentences = outputs.split("\nQ: ")[1].split("\n\n")
                # remove empty strings from the list
                sentences = [sentence for sentence in sentences if sentence]
                question = sentences[0]
                # remove ``` from the question
                question = question.replace("```", "")
                # remove new line characters
                question = question.replace("\n", "")
                # print("Actual question: ", question)
                # print("-" * 100)

                result = {"id": example["id"], "question": question, "context": context}

                examples.append(result)
                if c % 5000 == 0:
                    save_to_disk(
                        examples,
                        f"{save_path}counterfactual_questions_{model_identifier}_{c}.jsonl",
                    )
                    examples = []

            except Exception as e:
                print(outputs)
                skipped_instances += 1
                continue

        # save the remaining examples
        if examples:
            save_to_disk(
                examples,
                f"{save_path}counterfactual_questions_{model_identifier}_{c}.jsonl",
            )
        print(f"Skipped instances: {skipped_instances}")

src/cf_generation/llm_generation/llama_qgen.py
UKPLab/CATfOOD
import os
import torch
from transformers import GenerationConfig, LlamaForCausalLM, LlamaTokenizer
from transformers import T5ForConditionalGeneration, T5Tokenizer, AutoTokenizer
from langchain import PromptTemplate
import jsonlines
import re
from tqdm import tqdm
from datasets import load_dataset
from transformers import AutoModelForCausalLM, T5Tokenizer, AutoTokenizer, AutoConfig

from src.cf_generation.llm_generation.utils import save_to_disk


BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()

    # Set the seed
    torch.manual_seed(args.seed)

    # load squad data
    dataset = load_dataset("squad", "plain_text")
    train_data = dataset["train"]
    squad_data = [
        sample
        for sample in tqdm(
            train_data, total=len(train_data), desc="Loading SQuAD data ... "
        )
    ]

    template = """ 
        Given the question: \n
        {query}
        Decide if the following retrieved context is relevant to the {answer}: \n
        {result}
        Answer in the following format: \n
        "Context is relevant: True or False." \n """.strip()

    GRADE_DOCS_PROMPT_FAST = PromptTemplate(
        input_variables=["query", "result", "answer"], template=template
    )

    device = torch.device("cuda:0")
    model_name = "google/flan-ul2"

    model = T5ForConditionalGeneration.from_pretrained(
        model_name, torch_dtype=torch.bfloat16, device_map="auto"
    )
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    model_identifier = model_name.split("/")[-1]
    model_id = "alpaca"
    save_path = BASE_PATH + f"src/data/squad/{model_id}_qa_relevance_seed_{args.seed}/"
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    file_path = os.path.join(
        BASE_PATH,
        f"src/data/squad/counterfactual_data_alpaca_13b_v2_qg_pipeline_all_data_cleaned.jsonl",
    )
    files = [file_path]
    skipped = 0
    c = 0
    for file_name in files:
        examples = []
        with jsonlines.open(file_name) as reader:
            for example in tqdm(reader):
                try:
                    # c+=1
                    # if c <= 25000:
                    #     continue
                    id = example["id"].split("_")[0]
                    context = example["context"]
                    question = example["question"]
                    answer = example["answers"]["text"][0]

                    # print("Given ans:", example["answers"])

                    orig_example = [
                        sample for sample in squad_data if sample["id"] == id
                    ][0]

                    orig_context = orig_example["context"]
                    orig_question = orig_example["question"]
                    orig_answer = orig_example["answers"]

                    input = GRADE_DOCS_PROMPT_FAST.format(
                        query=question, result=context, answer=answer
                    )
                    # print(input)
                    input_ids = tokenizer.encode(input, return_tensors="pt").to(device)
                    # outputs = generator.generate([input], max_new_tokens=50, top_p=0.95)

                    with torch.no_grad():
                        outputs = model.generate(
                            input_ids,
                            max_new_tokens=10,
                            temperature=0,
                            top_p=1,
                            top_k=40,
                            repetition_penalty=1.0,
                            pad_token_id=tokenizer.eos_token_id,
                        )
                    output = tokenizer.decode(outputs[0], skip_special_tokens=True)

                    c += 1
                    # if c==50:
                    #     break

                    result = {
                        "id": example["id"],
                        "question": question,
                        "context": context,
                        "answer": answer,
                        "context_relevance": output,
                    }
                    # print(result)
                    # break
                    examples.append(result)
                    if c % 5000 == 0:
                        save_to_disk(
                            examples,
                            f"{save_path}counterfactual_samples_{model_id}_{c}.jsonl",
                        )
                        examples = []
                # break
                except Exception as e:
                    print("Skip")
        # save the remaining examples
        if examples:
            save_to_disk(
                examples, f"{save_path}counterfactual_samples_{model_id}_{c}.jsonl"
            )

src/cf_generation/llm_generation/llm_evaluator.py
UKPLab/CATfOOD
import jsonlines
from tqdm import tqdm
from datasets import load_dataset
import os


# BASE_PATH = "/home/sachdeva/projects/ukp/exp_calibration/"
BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


def postprocess_question(question):
    # question = question.strip()
    """Postprocess the questions to remove the question type and the question mark."""
    tokens_to_remove = ["[", "'", '"', "]"]
    # Create a translation table that maps each unwanted token to None
    translator = str.maketrans({token: None for token in tokens_to_remove})
    ques = question.translate(translator).strip()
    # check if the question words are less than 3
    if len(ques.split()) < 3:
        return None
    # question corrector
    # Define a dictionary to map misspelled words to correct words
    corrections = {
        "ho": "who",
        "hat": "what",
        "hich": "which",
        "ow": "how",
        "ich": "which",
        "here": "where",
        "hen": "when",
    }
    # if question starts with a misspelled word, replace it with the correct word
    # print("Question before correction: ", ques.split())0])
    if ques.split()[0].lower() in list(corrections.keys()):
        ques = ques.replace(ques.split()[0], corrections[ques.split()[0]])
        ques = ques.capitalize()
    return ques


def postprocess_answer(answer):
    """Postprocess the answer to remove the answer type."""
    if answer == " ":
        return None
    return answer


def remove_duplicate_examples(input_file, output_file):
    # Load SQuAD dataset
    squad_dataset = load_dataset("squad")

    # Collect all unique QA strings in SQuAD dataset
    squad_examples = set()
    for example in squad_dataset["train"]:
        qa_string = example["question"] + example["answers"]["text"][0]
        squad_examples.add(qa_string)

    seen_examples = set()

    with jsonlines.open(output_file, "w") as writer:
        with jsonlines.open(input_file) as reader:
            for example in tqdm(reader):
                qa_string = example["question"] + example["answers"]["text"][0]
                if qa_string not in seen_examples and qa_string not in squad_examples:
                    writer.write(example)
                    seen_examples.add(qa_string)


def postprocess_pipeline(data_path, save_path, model_name):
    """
    Postprocess the data generated by the pipeline
    """

    # step 1: collate all the data from the jsonl files
    from src.cf_generation.llm_generation.utils import collate_jsonl_files

    collate_jsonl_files(
        data_path, os.path.join(save_path, f"{model_name}_collated_data.jsonl")
    )
    print(
        "Collated data saved to: ",
        os.path.join(save_path, f"{model_name}_collated_data.jsonl"),
    )

    seen_examples = set()
    total = 0
    # filter out samples with same ids
    with jsonlines.open(
        os.path.join(save_path, f"{model_name}_collated_data_cleaned.jsonl"), "w"
    ) as writer:
        with jsonlines.open(
            os.path.join(save_path, f"{model_name}_collated_data.jsonl")
        ) as reader:
            for example in tqdm(reader):
                id = example["id"].split("_")[0]
                total += 1
                if id not in seen_examples:
                    writer.write(example)
                    seen_examples.add(id)
    print("Total samples before cleaning: ", total)
    print("Total samples after cleaning: ", len(seen_examples))

    # step 2: get the answer indices
    from src.cf_generation.llm_generation.utils import get_answer_start, save_to_disk

    data_type = "counterfactuals"
    if data_type == "squad":
        # load squad data
        dataset = load_dataset("squad", "plain_text")
        train_data = dataset["train"]
        data = [
            sample
            for sample in tqdm(
                train_data, total=len(train_data), desc="Loading SQuAD data ... "
            )
        ]
    elif data_type == "counterfactuals":
        # read in counterfactual data from jsonl file
        with jsonlines.open(
            os.path.join(save_path, f"{model_name}_collated_data_cleaned.jsonl")
        ) as reader:
            data = [
                sample
                for sample in tqdm(
                    reader, desc="Loading Counterfactual data for answer gen. ... "
                )
            ]

            # assert if data length is equal to the number of unique samples via the id
            assert len(data) == len(set([sample["id"] for sample in data]))

    no_answer_found = 0
    c = 0
    samples = []
    for example in tqdm(data):
        try:
            c += 1

            # print("question: ", example["question"])
            # print("context: ", example["context"])
            # print("answer: ", answer)
            # print("-" * 100)

            answer_start = get_answer_start(
                example["question"], example["context"], example["answer"]
            )
            if answer_start == -1:
                no_answer_found += 1
                continue
            samples.append(
                {
                    "id": example["id"],
                    "question": example["question"],
                    "context": example["context"],
                    "answers": {
                        "text": [example["answer"]],
                        "answer_start": [answer_start],
                        "answer_end": [answer_start + len(example["answer"])],
                    },
                },
            )
            # if c == 20:
            #     break
        except Exception as e:
            print("Error in example: ", example)

    print("No answer found: ", no_answer_found)
    save_to_disk(
        samples,
        os.path.join(save_path, f"{model_name}_collated_data_with_answers.jsonl"),
    )

    # step 3: postprocess data
    with jsonlines.open(
        os.path.join(save_path, f"{model_name}_collated_data_with_answers.jsonl")
    ) as reader:
        data = [
            sample
            for sample in tqdm(reader, desc="Loading data for postprocessing... ")
        ]

        # Load SQuAD dataset
        squad_dataset = load_dataset("squad")
        # Collect all unique QA strings in SQuAD dataset
        squad_examples = set()
        for example in squad_dataset["train"]:
            qa_string = example["question"] + example["answers"]["text"][0]
            squad_examples.add(qa_string)

        c = 0
        with jsonlines.open(
            os.path.join(
                save_path, f"{model_name}_collated_data_with_answers_processed.jsonl"
            ),
            "a",
        ) as writer:
            seen_examples = set()
            for example in tqdm(data, total=len(data), desc="Saving samples ... "):
                question = postprocess_question(example["question"])
                answer = postprocess_answer(example["answers"]["text"][0])
                if question is None or answer is None:
                    # print("Question is None: ", example["question"])
                    continue
                example["question"] = question
                example["answers"]["text"][0] = answer
                qa_string = question + answer
                # remove duplicate examples
                if qa_string not in seen_examples and qa_string not in squad_examples:
                    writer.write(example)
                    seen_examples.add(qa_string)
                    c += 1
                # else:
                # print("Duplicate example: ", example["question"], example["answers"]["text"][0])
                # writer.write(example)

        print("Final sample count: ", c)


if __name__ == "__main__":
    data_path = BASE_PATH + "src/data/squad/few_shot_flan-ul2_thinking/"
    save_path = BASE_PATH + "src/data/squad/"
    model_name = "flan-ul2_thinking"
    postprocess_pipeline(data_path, save_path, model_name)

    # input_file = f"{BASE_PATH}src/data/squad/counterfactual_data_gpt_neox_20b_v2_qg_pipeline_all_data.jsonl"
    # save_file = f"{BASE_PATH}src/data/squad/counterfactual_data_gpt_neox_20b_v2_qg_pipeline_all_data_cleaned.jsonl"

    # remaining_samples(complete_file, subset_file, save_path)

    # remove_duplicate_examples(input_file, save_file)

    # data_type = "counterfactuals"
    # if data_type == "squad":
    #     # load squad data
    #     dataset = load_dataset("squad", "plain_text")
    #     train_data = dataset["train"]
    #     data = [sample for sample in tqdm(train_data, total=len(train_data), desc="Loading SQuAD data ... ")]
    # elif data_type == "counterfactuals":
    #     # read in counterfactual data from jsonl file
    #     with jsonlines.open(f"{BASE_PATH}src/data/squad/counterfactual_data_GPT-NeoXT-Chat-Base-20B_qg_pipeline_remaining_final.jsonl") as reader:
    #         data = [sample for sample in tqdm(reader, desc="Loading Counterfactual data ... ")]
    #
    # # Load SQuAD dataset
    # squad_dataset = load_dataset("squad")
    #
    # # Collect all unique QA strings in SQuAD dataset
    # squad_examples = set()
    # for example in squad_dataset["train"]:
    #     qa_string = example["question"] + example["answers"]["text"][0]
    #     squad_examples.add(qa_string)
    #
    # c = 0
    # with jsonlines.open(f"{BASE_PATH}src/data/squad/counterfactual_data_GPT-NeoXT-Chat-Base-20B_qg_pipeline_remaining_final_cleaned.jsonl",
    #                     "a") as writer:
    #     seen_examples = set()
    #     for example in tqdm(data, total=len(data), desc="Saving samples ... "):
    #         question = postprocess_question(example["question"])
    #         answer = postprocess_answer(example["answers"]["text"][0])
    #         if question is None or answer is None:
    #             # print("Question is None: ", example["question"])
    #             continue
    #         example["question"] = question
    #         example["answers"]["text"][0] = answer
    #         qa_string = question + answer
    #         # remove duplicate examples
    #         if qa_string not in seen_examples and qa_string not in squad_examples:
    #             writer.write(example)
    #             seen_examples.add(qa_string)
    #             c += 1
    #         # else:
    #             # print("Duplicate example: ", example["question"], example["answers"]["text"][0])
    #         # writer.write(example)
    # print("Count: ", c)

src/cf_generation/llm_generation/postprocess.py
UKPLab/CATfOOD
import os
import jsonlines
from tqdm import tqdm
import torch

from typing import Optional

from datasets import load_dataset
from langchain import PromptTemplate, FewShotPromptTemplate
from transformers import AutoModelForCausalLM, T5Tokenizer, AutoTokenizer, AutoConfig

# from src.few_shot.utils import save_to_disk
from src.few_shot.together import infer

# Set the seed
seed = 0
torch.manual_seed(seed)


BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"


def save_to_disk(data, file_name):
    with jsonlines.open(file_name, "a") as writer:
        for example in tqdm(data, total=len(data), desc="Saving samples ... "):
            writer.write(example)


class LLMClient:
    def __init__(
        self,
        template,
        model_name_or_path: str = None,
        tokenizer_name_or_path: Optional[str] = None,
        data_path: str = None,
        threshold: float = 0.5,
        task: str = "qg",
        max_new_tokens: int = 50,
        temperature: float = 0.01,
        top_p: float = 1,
        top_k: int = 40,
        repetition_penalty: float = 1.0,
        save_results: bool = True,
        max_samples: int = None,
        stop: str = "\n",
    ):
        self.model = model_name_or_path
        self.tokenizer = tokenizer_name_or_path
        self.max_new_tokens = max_new_tokens
        self.temperature = temperature
        self.top_p = top_p
        self.top_k = top_k
        self.repetition_penalty = repetition_penalty
        self.data_path = data_path
        self.task = task
        self.threshold = threshold
        self.save_results = save_results
        self.max_samples = max_samples
        self.stop = stop.split(";")
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        # self.prompt_template = PromptTemplate(
        #     input_variables=["context"] if self.task == "qg" else ["context", "question"],
        #     template=template,
        # )

    def _load_model(self):
        model = AutoModelForCausalLM.from_pretrained(
            self.model, device_map="auto"
        )  # , torch_dtype=torch.bfloat16)
        model.to(self.device)
        tokenizer = AutoTokenizer.from_pretrained(self.model)
        return model, tokenizer

    def _create_few_shot_prompt(self, examples, context, question=None):
        example_template = """
Context: {context}
Question: {question}
Answer: {answer}
""".strip()

        suffix = """
Context: {context}
Question: 
""".strip()
        example_prompt = PromptTemplate(
            input_variables=["context", "question", "answer"],
            template=example_template,
        )

        if self.task == "qg":
            few_shot_template = FewShotPromptTemplate(
                examples=examples,
                example_prompt=example_prompt,
                prefix=qg_template if self.task == "qg" else qa_template,
                suffix=suffix,
                input_variables=["context"],
                example_separator="\n",
            )
            prompt = few_shot_template.format(context=context)
        else:
            few_shot_template = FewShotPromptTemplate(
                examples=examples,
                example_prompt=example_prompt,
                prefix=qa_template,
                suffix=suffix,
                input_variables=["context", "question"],
                example_separator="\n",
            )
            prompt = few_shot_template.format(context=context, question=question)
        return prompt

    def generate(self):
        c = 0
        skipped_instances = 0
        examples = []
        model, tokenizer = self._load_model()
        model_identifier = self.model.split("/")[-1]
        save_path = (
            BASE_PATH
            + f"src/data/squad/t5_squad_counterfactuals/{model_identifier}_{self.task}_pipeline_remaining_temp_0.7"
        )
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        # load squad data
        dataset = load_dataset("squad", "plain_text")
        train_data = dataset["train"]
        squad_data = [
            sample
            for sample in tqdm(
                train_data, total=len(train_data), desc="Loading SQuAD data ... "
            )
        ]

        current_files = []
        file_names = []
        file_path = BASE_PATH + f"src/data/squad/{self.data_path}"
        if self.task == "qg":
            file_names = [file_path]
        elif self.task == "qa":
            all_files = os.listdir(file_path)
            file_names = [
                os.path.join(file_path, file)
                for file in all_files
                if file not in current_files
            ]

        model.eval()
        for file_name in file_names:
            with jsonlines.open(file_name) as reader:
                for example in tqdm(reader):
                    try:
                        id = example["id"].split("_")[0]
                        context = example["context"]
                        orig_example = [
                            sample for sample in squad_data if sample["id"] == id
                        ][0]
                        # print(orig_example)
                        orig_context = orig_example["context"]
                        orig_question = orig_example["question"]
                        orig_answer = orig_example["answers"]
                        c += 1

                        if c <= 73000:
                            continue

                        if self.max_samples:
                            if c == self.max_samples:
                                break
                        if self.task == "qa":
                            # preprocess the question
                            # remove any new line characters
                            question = example["question"]
                            question = question.replace("\n", "")
                            # check for truncated question words
                            # check for truncated question words
                            question_words = [
                                "what",
                                "who",
                                "when",
                                "where",
                                "why",
                                "how",
                                "which",
                                "whom",
                                "whose",
                            ]
                            for word in question_words:
                                if question.lower().startswith(word[1:]):
                                    # replace first word with the correct word
                                    question = word + question[len(word[1:]) :]
                                    # question = word + question[3:]
                                    break

                        prompt = self._create_few_shot_prompt(
                            [
                                {
                                    "context": orig_context,
                                    "question": orig_question,
                                    "answer": orig_answer["text"][0],
                                }
                            ],
                            context,
                            question=question if self.task == "qa" else None,
                        )
                        # print(prompt)

                        inputs = tokenizer(prompt, return_tensors="pt").to(self.device)

                        outputs = model.generate(
                            **inputs,
                            max_new_tokens=self.max_new_tokens,
                            temperature=self.temperature,
                            top_p=self.top_p,
                            top_k=self.top_k,
                            repetition_penalty=self.repetition_penalty,
                            do_sample=True,
                            # num_return_sequences=1,
                            early_stopping=True,
                            pad_token_id=tokenizer.eos_token_id,
                        )
                        output = tokenizer.decode(outputs[0], skip_special_tokens=True)
                        # remove the context from the output
                        output = output[len(prompt) :]

                        # output = infer(
                        #     model_name="GPT-JT-6B-v1",
                        #     prompt=prompt,
                        # )

                        result = {"id": example["id"], "context": context}

                        if self.task == "qg":
                            # add the question mark if it is not present
                            # if output[-1] != "?":
                            #     output = output + "?"
                            # question = output

                            # for stop_word in self.stop:
                            #     if stop_word != '' and stop_word in output:
                            #         question = output[:output.find(stop_word)]
                            #     # in some cases the stop word is at the beginning of the output
                            #     if question == "":
                            #         question = output[output.find(stop_word):]
                            # Split the text into lines
                            lines = output.split("\n")

                            # Extract the question and answer
                            question = lines[0].replace("Question: ", "")
                            question = question.strip()
                            answer = lines[1].replace("Answer: ", "")
                            answer = answer.replace("<bot>:", "")
                            answer = answer.strip()

                            result["question"] = question
                            result["answer"] = answer

                        elif self.task == "qa":
                            answer = output
                            # print("Output:", output)
                            for stop_word in self.stop:
                                if stop_word != "" and stop_word in output:
                                    answer = output[output.find(stop_word) :]
                                    # remove the first stop word from answer
                                    answer = answer.split(stop_word)
                                    answer = [a for a in answer if a != ""]
                                    if answer:
                                        answer = answer[0]
                                    else:
                                        answer = "I don't know."

                                # in some cases the stop word is at the beginning of the output
                            #     if answer == "":
                            #         answer = output[output.find(stop_word):]
                            #         answer = answer.replace(stop_word, "")
                            # if answer == "":
                            #     answer = answer.split(stop_word)
                            #     answer = [a for a in answer if a != ""][0]
                            skip_words = [
                                "A: ",
                                "A. ",
                                "Answer: ",
                                "Output: ",
                                "Label: ",
                            ]
                            for skip_word in skip_words:
                                if skip_word in answer:
                                    answer = answer.replace(skip_word, "")

                            # if answer.startswith("A: ") or answer.startswith("A. "):
                            #     answer = answer[3:]
                            # elif answer.startswith("Answer: "):
                            #     answer = answer.replace("Answer: ", "")
                            # question = example["question"]
                            result["question"] = question
                            result["answer"] = answer

                        # print("Context:", context)
                        # print("Question:", question)
                        # print("Answer:", answer)
                        # print("-" * 100)

                        examples.append(result)
                        if self.save_results:
                            if c % self.threshold == 0:
                                save_to_disk(
                                    examples,
                                    f"{save_path}/counterfactual_samples_{model_identifier}_{c}.jsonl",
                                )
                                examples = []
                    except Exception as e:
                        # print(outputs)
                        skipped_instances += 1
                        print(f"Skipped instance {c} due to error: {e}.")
                        continue

            # save the remaining examples
            if self.save_results:
                if examples:
                    save_to_disk(
                        examples,
                        f"{save_path}/counterfactual_samples_{model_identifier}_{c}.jsonl",
                    )
            # break


if __name__ == "__main__":

    qg_template = """
As a question generator, your task is to create a concise and clear question that can be answered by an answer span within a given context. 
The context should be a piece of text, such as a news article or historical document, and the question should require understanding and analysis of the information presented in the context. 
Your generated question should focus on key details or events described in the context and should demonstrate your ability to identify important information. 
Additionally, please ensure that your question is specific enough to have a single correct answer within the given context. 
Please note that you may need to read through the provided context multiple times to fully understand its contents before generating an appropriate question. 
""".strip()

    chatgpt_qg_template = """
As a question generator, your task is to create a clear and concise question that can be answered by an answer span within a given context. 
The context should be a piece of text, such as a news article or historical document, and the question should require understanding and analysis of the information presented in the context. 
Your generated question should focus on key details or events described in the context, requiring readers to carefully read and analyze the provided text. 
Please ensure that your question is specific enough to have only one correct answer span within the given context. 
Please note that you should aim for clarity and concision while still maintaining accuracy and relevance to the provided text.
""".strip()

    qa_template = """
As a language model, your task is to generate an answer that is strictly a span in the given context provided with a specific question.
Please provide a concise and relevant response that directly answers the question asked within the context.
Your answer should be accurate and reflect the information presented in the context.
Please note that you may need to use contextual clues and understanding to identify the most appropriate section of text to use as your answer span.


Context: {context}


Question: {question}


Answer: 
"""

    task = "qg"
    model = "togethercomputer/GPT-JT-6B-v1"
    client = LLMClient(
        template=qg_template if task == "qg" else qa_template,
        model_name_or_path=model,
        tokenizer_name_or_path=model,
        data_path="t5_squad_counterfactuals/rag_counterfactuals_complete_min_filtered_dedup.jsonl",
        threshold=1000,
        task=task,
        max_new_tokens=50,
        temperature=0.7,
        top_p=1,
        top_k=40,
        repetition_penalty=1.0,
        save_results=True,
        max_samples=None,
    )

    client.generate()

src/cf_generation/llm_generation/prompt_injection.py
UKPLab/CATfOOD
"""
A dedicated helper to manage templates and prompt building.
"""

import json
import os.path as osp
from typing import Union


BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/src/few_shot"


class Prompter(object):
    __slots__ = ("template", "_verbose")

    def __init__(self, template_name: str = "", verbose: bool = False):
        self._verbose = verbose
        if not template_name:
            # Enforce the default here, so the constructor can be called with '' and will not break.
            template_name = "alpaca"
        file_name = osp.join(BASE_PATH, "templates", f"{template_name}.json")
        # file_name = template_name
        # print(file_name)
        if not osp.exists(file_name):
            raise ValueError(f"Can't read {file_name}")
        with open(file_name) as fp:
            self.template = json.load(fp)
        if self._verbose:
            print(
                f"Using prompt template {template_name}: {self.template['description']}"
            )

    def generate_prompt(
        self,
        instruction: str,
        input: Union[None, str] = None,
        label: Union[None, str] = None,
    ) -> str:
        # returns the full prompt from instruction and optional input
        # if a label (=response, =output) is provided, it's also appended.
        if input:
            res = self.template["prompt_input"].format(
                instruction=instruction, input=input
            )
        else:
            res = self.template["prompt_no_input"].format(instruction=instruction)
        if label:
            res = f"{res}{label}"
        if self._verbose:
            print(res)
        return res

    def get_response(self, output: str) -> str:
        return output.split(self.template["response_split"])[1].strip()

src/cf_generation/llm_generation/prompter.py
UKPLab/CATfOOD
import jsonlines
from tqdm import tqdm
from datasets import load_dataset

# BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"
BASE_PATH = "//"

with jsonlines.open(
    f"{BASE_PATH}src/data/squad/squad_counterfactuals_noise_min_filtered_final_2.jsonl"
) as reader:
    filtered_data = [
        sample for sample in tqdm(reader, desc="Loading Counterfactual data ... ")
    ]
    filtered_ids = [data["id"].split("_")[0] for data in filtered_data]
    squad_dataset = load_dataset("squad", "plain_text")
    squad_examples = set()
    # print(len(data))
    print(len(filtered_ids))
    squad_ids = [data["id"] for data in squad_dataset["train"]]
    # check if the id is in the squad dataset and list the ones that are not
    extra_ids = [idx for idx in squad_ids if idx not in filtered_ids]

with jsonlines.open(
    f"{BASE_PATH}src/data/squad/squad_counterfactuals_llms.jsonl", "a"
) as writer:
    with jsonlines.open(
        f"{BASE_PATH}src/data/squad/squad_counterfactuals_noise_min_filtered.jsonl"
    ) as reader:
        for sample in tqdm(reader, desc="Loading Counterfactual data ... "):
            if sample["id"].split("_")[0] in extra_ids:
                writer.write(sample)

src/cf_generation/llm_generation/sampling.py
UKPLab/CATfOOD
import os
import glob
import json
import jsonlines
from tqdm import tqdm
from datasets import load_dataset
from fuzzywuzzy import fuzz

import spacy

# Load the Spacy model
nlp = spacy.load("en_core_web_sm")
stop_words = set(spacy.lang.en.stop_words.STOP_WORDS)


BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"
# BASE_PATH = "/home/sachdeva/projects/ukp/exp_calibration/"

import string


def is_answer_close(original_answer, given_answer, threshold=80):
    similarity_score = fuzz.ratio(original_answer.lower(), given_answer.lower())
    return similarity_score >= threshold


def remove_punctuation(input_string):
    # Create a translation table to remove punctuation
    translation_table = str.maketrans("", "", string.punctuation)

    # Use the translate() method to remove punctuation
    cleaned_string = input_string.translate(translation_table)
    return cleaned_string


def collate_jsonl_files(data_path=None, save_path=None):

    if not data_path:
        # Set the path to the directory containing the JSONL files
        data_path = BASE_PATH + "src/data/squad/t5_squad_counterfactuals/final_cfs"
    print("Loading data from: ", data_path)

    if not save_path:
        # Set the path to the output file
        save_path = (
            BASE_PATH
            + "src/data/squad/t5_squad_counterfactuals/rag_counterfactuals_complete_noise_min_filtered_diff_ans_final.jsonl"
        )
    print("Saving data to: ", save_path)

    # Create an empty list to store the samples
    samples = []

    # Loop over all the JSONL files in the directory
    for file_path in glob.glob(data_path + "/*.jsonl"):
        # Open the JSONL file and read the samples
        with jsonlines.open(file_path) as reader:
            samples.extend(reader)

    # Write the collated samples to the output file
    with jsonlines.open(save_path, mode="w") as writer:
        writer.write_all(samples)


def save_to_disk(data, file_name):
    with jsonlines.open(file_name, "a") as writer:
        for example in tqdm(data, total=len(data), desc="Saving samples ... "):
            writer.write(example)


def get_answer_start(question, context, answer):
    # remove trailing punctuation from answer
    answer = answer.rstrip(".")
    if answer == "":
        return -1
    # Find all occurrences of the string in the context
    indices = []
    end_index = 0
    answer_start = None
    context = context.lower()
    answer = answer.lower()

    while True:
        index = context.find(answer, end_index)
        if index == -1:
            break
        indices.append(index)
        end_index = index + len(answer)
    # print("Indices: ", indices)
    if not indices:
        return -1
    if len(indices) > 1:
        answer_indices = [
            i for i in range(len(context)) if context.startswith(answer, i)
        ]
        # print(answer_indices)
        # Choose the correct instance of the answer based on the question and context
        doc = nlp(context)
        question_tokens = set(
            [token.text.lower() for token in nlp(question) if not token.is_stop]
        )
        max_match_count = 0
        max_match_sent = None
        # print(doc)
        for (
            sent
        ) in doc.sents:  # print(context.find(a    # print(context.find(answer))nswer))
            sent_tokens = set(
                [token.text.lower() for token in sent if not token.is_stop]
            )
            match_count = len(question_tokens.intersection(sent_tokens))
            # print(sent, match_count)
            if match_count > max_match_count:
                max_match_count = match_count
                max_match_sent = sent
        if max_match_sent is not None:
            if answer.lower() in max_match_sent.text.lower():
                for index in answer_indices:
                    # get the start and end of sentence as in context
                    if max_match_sent.start_char <= index < max_match_sent.end_char:
                        answer_start = index
                        break
        if answer_start is None:
            answer_start = answer_indices[0]
    else:
        answer_start = indices[0]

    return answer_start


def _preprocess_llama_questions():
    current_files = []
    model_name = "decapoda-research/llama-13b-hf"
    model_identifier = model_name.split("/")[-1]
    all_files = os.listdir(BASE_PATH + f"src/data/squad/{model_identifier}_qg/")
    files = [file for file in all_files if file not in current_files]
    print(files)

    c = 0
    for file in files:
        examples = []
        with jsonlines.open(
            BASE_PATH + f"src/data/squad/{model_identifier}_qg/" + file
        ) as reader:
            for example in tqdm(reader):
                id = example["id"].split("_")[0]
                context = example["context"]
                raw_question = example["question"]
                print("*" * 100)
                print(raw_question)
                print("-" * 100)

                # process question
                sentences = raw_question.split("\nQ: ")[1].split("\n\n")
                # remove empty strings from the list
                sentences = [sentence for sentence in sentences if sentence]
                question = sentences[0]
                # remove ``` from the question
                question = question.replace("```", "")
                # remove new line characters
                question = question.replace("\n", "")

                print(question)
                print("*" * 100)
                c += 1
                if c == 10:
                    break
        break


def remove_duplicate_examples(input_file, output_file):
    seen_examples = set()

    with jsonlines.open(output_file, "w") as writer:
        with jsonlines.open(input_file) as reader:
            for example in tqdm(reader):
                qa_string = example["question"] + example["answers"]["text"][0]
                if qa_string not in seen_examples:
                    writer.write(example)
                    seen_examples.add(qa_string)


def remaining_samples(complete_file, subset_file, save_path):
    """
    Read the complete file and subset file and list ids not present
    in the subset file
    """
    with jsonlines.open(complete_file) as reader:
        all_ids = [sample["id"] for sample in reader]
    with jsonlines.open(subset_file) as reader:
        subset_ids = [sample["id"] for sample in reader]
    remaining_ids = [idx for idx in all_ids if idx not in subset_ids]
    with jsonlines.open(save_path, "w") as writer:
        with jsonlines.open(complete_file) as reader:
            for example in tqdm(reader):
                idx = example["id"]
                if idx in remaining_ids:
                    writer.write(example)
    # return remaining_ids


def filter_counterfactuals():
    """
    compare generated answer to the answers generated via 3 random seeds
    """
    data_path = os.path.join(BASE_PATH, "src/data/squad")
    orig_ans_path = os.path.join(
        data_path, "flan-t5-xxl-v3_collated_data_with_answers_processed.jsonl"
    )
    answers = {}
    seeds = [0, 1, 42]
    for seed in seeds:
        path = os.path.join(
            data_path,
            f"few_shot_flan-t5-xxl_qa_eval_seed_{seed}/counterfactual_samples_flan-t5-xxl_1.jsonl",
        )
        with jsonlines.open(path) as reader:
            for sample in tqdm(reader):
                idx = sample["id"]
                answer = sample["answer"]
                if idx in answers.keys():
                    answers[idx].append(answer.lower())
                else:
                    answers[idx] = [answer.lower()]
                # break
        # break

    with jsonlines.open(orig_ans_path) as reader:
        examples = []
        for sample in tqdm(reader):
            target_answer = sample["answers"]["text"][0]
            idx = sample["id"]
            if idx in answers.keys():
                # compare answers
                ans_count = answers[idx].count(target_answer.lower())
                if ans_count >= 2:
                    examples.append(sample)

    save_to_disk(
        examples,
        os.path.join(
            data_path, "flan_t5_xxl_collated_data_with_answers_processed_filtered.jsonl"
        ),
    )
    # print(ans_count)


def context_eval():
    data_path = os.path.join(BASE_PATH, "src/data/squad")
    # t5_squad_counterfactuals/rag_counterfactuals_complete_noise_min_filtered_final_dedup_1.jsonl
    # counterfactual_data_gpt_neox_20b_v2_qg_pipeline_all_data_cleaned.jsonl
    # counterfactual_data_gpt_jt_v2_qg_pipeline_all_data_cleaned.jsonl
    # counterfactual_data_llama_13b_v1_qg_pipeline_all_data_cleaned.jsonl
    # counterfactual_data_alpaca_13b_v2_qg_pipeline_all_data_cleaned.jsonl
    # flan-t5-xxl-v3_collated_data_with_answers_processed.jsonl
    # flan_ul2_collated_data_with_answers_processed.jsonl
    orig_ans_path = os.path.join(
        data_path, "flan-t5-xxl-v3_collated_data_with_answers_processed.jsonl"
    )
    answers = {}

    # context rel file
    path = os.path.join(
        data_path,
        f"counterfactual_samples_Llama-2-13b-chat-hf_flan-t5-xxl_complete.jsonl",
    )
    with jsonlines.open(path) as reader:
        for sample in tqdm(reader):
            # print(sample)
            idx = sample["id"]
            context_rel = sample["context_relevance"]
            if idx in answers.keys():
                answers[idx].append(context_rel)
            else:
                answers[idx] = [context_rel]
            # break
    # break
    # print(list(answers.values())[:10])
    c = 0
    with jsonlines.open(orig_ans_path) as reader:
        examples = []
        for sample in tqdm(reader):
            # target_answer = sample["answers"]["text"][0]
            idx = sample["id"]
            if idx in answers.keys():
                context_rel = answers[idx][0]
                if context_rel.__contains__("True"):
                    examples.append(sample)
                    c += 1
    print(c)
    save_to_disk(
        examples,
        os.path.join(
            data_path,
            "counterfactual_samples_Llama-2-13b-chat-hf_flan-t5-xxl_context_filtered_complete.jsonl",
        ),
    )


def openai_context_eval():
    data_path = os.path.join(BASE_PATH, "src/data/squad")
    eval_model = "gpt-4-0314"
    test_model = "flan_ul2"
    save_path = BASE_PATH + f"src/data/squad/{eval_model}_qa_relevance_seed_0/"
    orig_ans_path = os.path.join(
        save_path, f"counterfactual_samples_{eval_model}_{test_model}_500.jsonl"
    )
    answers = {}

    # path = os.path.join(data_path, f"counterfactual_samples_gpt_jt_context_relevance.jsonl")
    # with jsonlines.open(path) as reader:
    #     for sample in tqdm(reader):
    #         idx = sample["id"]
    #         context_rel = sample["context_relevance"]
    #         if idx in answers.keys():
    #             answers[idx].append(context_rel)
    #         else:
    #             answers[idx] = [context_rel]
    # break
    # break
    # print(list(answers.values())[:10])
    c = 0
    with jsonlines.open(orig_ans_path) as reader:
        examples = []
        for sample in tqdm(reader):
            # target_answer = sample["answers"]["text"][0]
            idx = sample["id"]
            output = sample["context_relevance"]
            context_rel = output["choices"][0]["message"]["content"]
            if context_rel.__contains__("True"):
                c += 1
            # print(context_rel)
            # break
    print(c)

    #         if idx in answers.keys():
    #             context_rel = answers[idx][0]
    #             if context_rel in ["True", "yes"]:
    #                 examples.append(sample)
    #                 c+=1
    # print(c)
    # save_to_disk(examples, os.path.join(data_path, "gpt_jt_collated_data_with_answers_processed_context_relevance.jsonl"))


def context_noise_filter():
    data_path = os.path.join(BASE_PATH, "src/data/squad")
    rel_answers = {}
    seed_answers = {}
    model = "flan-t5-xxl"

    orig_ans_path = os.path.join(
        data_path,
        "counterfactual_samples_Llama-2-13b-chat-hf_flan-t5-xxl_context_filtered_complete.jsonl",
    )
    seeds = [0, 1, 42]
    for seed in seeds:
        path = os.path.join(
            data_path,
            f"noise_filter_{model}_qa_eval_seed_{seed}/counterfactual_samples_{model}_1.jsonl",
        )
        with jsonlines.open(path) as reader:
            for sample in tqdm(reader):
                idx = sample["id"]
                answer = sample["answer"]
                answer = remove_punctuation(answer)
                if idx in seed_answers.keys():
                    seed_answers[idx].append(answer.lower())
                else:
                    seed_answers[idx] = [answer.lower()]

    # path = os.path.join(data_path, f"counterfactual_samples_flan_t5_xxl_context_relevance.jsonl")
    # with jsonlines.open(path) as reader:
    #     for sample in tqdm(reader):
    #         idx = sample["id"]
    #         context_rel = sample["context_relevance"]
    #         if idx in rel_answers.keys():
    #             rel_answers[idx].append(context_rel)
    #         else:
    #             rel_answers[idx] = [context_rel]
    # break
    # break
    # print(list(answers.values())[:10])
    c = 0
    with jsonlines.open(orig_ans_path) as reader:
        examples = []
        for sample in tqdm(reader):
            target_answer = sample["answers"]["text"][0]
            target_answer = remove_punctuation(target_answer)
            idx = sample["id"]
            # if idx in rel_answers.keys():
            # context_rel = rel_answers[idx][0]
            # if context_rel in ["True", "yes"]:
            #     examples.append(sample)
            #     c += 1
            if idx in seed_answers.keys():
                # compare answers
                ans_count = seed_answers[idx].count(target_answer.lower())
                if ans_count >= 2 or (
                    is_answer_close(seed_answers[idx][0], target_answer.lower())
                    and is_answer_close(seed_answers[idx][1], target_answer.lower())
                    and is_answer_close(seed_answers[idx][2], target_answer.lower())
                ):
                    examples.append(sample)
                    c += 1
    print(c)
    save_to_disk(
        examples,
        os.path.join(
            data_path,
            "counterfactual_samples_Llama-2-13b-chat-hf_flan_t5_xxl_context_filtered_noise_filtered_complete.jsonl",
        ),
    )


def compare_closed_open():
    closed_model = "gpt-4-0314"
    open_model = "Llama-2-13b-chat-hf"
    test_model = "llama"
    data_path = os.path.join(BASE_PATH, "src/data/squad")
    open_path = os.path.join(
        data_path,
        f"{open_model}_qa_relevance_seed_0/counterfactual_samples_{open_model}_{test_model}_500_2.jsonl",
    )
    closed_path = os.path.join(
        data_path,
        f"{closed_model}_qa_relevance_seed_0/counterfactual_samples_{closed_model}_{test_model}_500.jsonl",
    )

    c = 0
    with jsonlines.open(closed_path) as reader:
        closed_examples = [sample for sample in tqdm(reader)]

    with jsonlines.open(open_path) as reader:
        open_examples = [sample for sample in tqdm(reader)]

    for closed, open in zip(closed_examples, open_examples):
        closed_ans = closed["context_relevance"]["choices"][0]["message"]["content"]
        open_ans = open["context_relevance"]

        if closed_ans.__contains__("True") and open_ans.__contains__("True"):
            c += 1
        if closed_ans.__contains__("False") and open_ans.__contains__("False"):
            c += 1

    print(c)


def save_irrelevant_samples():
    data_path = os.path.join(BASE_PATH, "src/data/squad")
    path = os.path.join(
        data_path,
        f"flan_ul2_collated_data_with_answers_processed_context_irrelevance.jsonl",
    )
    with jsonlines.open(path) as reader:
        samples = [sample for sample in tqdm(reader)]
    import pandas as pd
    from random import shuffle
    import random

    print(samples[:10])

    shuffled_dataset = sorted(samples, key=lambda k: random.random())
    print(type(shuffled_dataset))
    df = pd.DataFrame()
    print(shuffled_dataset[:10])
    df["question"] = [sample["question"] for sample in shuffled_dataset[:100]]
    df["context"] = [sample["context"] for sample in shuffled_dataset[:100]]
    df["answer"] = [sample["answers"]["text"][0] for sample in shuffled_dataset[:100]]
    print(df.head())

    df.to_excel(os.path.join(data_path, f"flan_ul2_irrelevant.xlsx"))


def select_noise_relevant_samples():
    data_path = os.path.join(BASE_PATH, "src/data/squad")
    relevant_path = os.path.join(
        data_path,
        "flan_ul2_collated_data_with_answers_processed_context_relevance.jsonl",
    )
    examples = []
    noise_path = os.path.join(
        data_path, f"flan_ul2_collated_data_with_answers_processed_filtered.jsonl"
    )
    with jsonlines.open(noise_path) as reader1:
        noise_idx = [ex["id"] for ex in tqdm(reader1)]
    with jsonlines.open(relevant_path) as reader2:
        for example in tqdm(reader2):
            if example["id"] in noise_idx:
                examples.append(example)

    save_to_disk(
        examples,
        os.path.join(
            data_path,
            "flan_ul2_collated_data_with_answers_processed_context_relevance_filtered.jsonl",
        ),
    )


def agg_answers():
    paths = [
        "t5-large-squad-qa-seed-42",
        "t5-large-squad-qa-seed-0",
        "t5-large-squad-qa-seed-1",
    ]
    c = 0
    data_path = []
    for path in paths:
        c += 1
        data_path.append(
            BASE_PATH
            + f"src/data/squad/t5_squad_counterfactuals/rag_counterfactuals_complete_nf_{path}_1.jsonl"
        )
    with jsonlines.open(
        BASE_PATH
        + f"src/data/squad/t5_squad_counterfactuals/rag_counterfactuals_complete_noise_filtered.jsonl"
    ) as reader:
        examples1 = [ex for ex in reader]
    # with jsonlines.open(data_path[0]) as reader:
    #     examples1 = [ex for ex in reader]
    # with jsonlines.open(data_path[1]) as reader:
    #     examples2 = [ex for ex in reader]
    # with jsonlines.open(data_path[2]) as reader:
    #     examples3 = [ex for ex in reader]
    # min filter
    with jsonlines.open(
        BASE_PATH
        + f"src/data/squad/t5_squad_counterfactuals/rag_counterfactuals_complete_mf.jsonl"
    ) as reader:
        min_ex = [ex for ex in reader]

    min_ex_ids = [ex["id"] for ex in min_ex]
    for i, ex1 in tqdm(enumerate(examples1)):  # examples2, examples3, min_ex)):
        # ex1["alternate_answers"].append(ex2["alternate_answers"][0])
        # ex1["alternate_answers"].append(ex3["alternate_answers"][0])
        if ex1["id"] in min_ex_ids:
            idx = min_ex_ids.index(ex1["id"])
            ex1["similarity"] = min_ex[idx]["similarity"]
    save_to_disk(
        examples1,
        BASE_PATH
        + f"src/data/squad/t5_squad_counterfactuals/rag_counterfactuals_complete_noise_filtered_final.jsonl",
    )


def _add_answer_to_noise_filtered():
    with jsonlines.open(
        BASE_PATH
        + f"src/data/squad/t5_squad_counterfactuals/rag_counterfactuals_complete_noise_min_filtered_dedup.jsonl"
    ) as reader:
        nf_samples = [ex for ex in reader]
        nf_idx = [ex["id"] for ex in nf_samples]

    with jsonlines.open(
        BASE_PATH
        + f"src/data/squad/t5_squad_counterfactuals/rag_counterfactuals_complete_noise_filtered_final.jsonl"
    ) as reader:
        all_samples = [ex for ex in reader]

    examples = []
    for ex in tqdm(all_samples):
        if ex["id"] not in nf_idx:
            continue
        # if ex["id"] in nf_idx:
        examples.append(ex)
    print(len(examples))
    save_to_disk(
        examples,
        BASE_PATH
        + f"src/data/squad/t5_squad_counterfactuals/rag_counterfactuals_complete_noise_min_filtered_dedup_final.jsonl",
    )


if __name__ == "__main__":

    ############ FILTERING ############
    # complete_file = f"{BASE_PATH}src/data/squad/squad_counterfactuals_28_03.jsonl"
    # subset_file = f"{BASE_PATH}src/data/squad/GPT-NeoXT-Chat-Base-20B-v1_collated_data_with_answers_processed.jsonl"
    # save_path = f"{BASE_PATH}src/data/squad/GPT-NeoXT-Chat-Base-20B-v1_remaining_counterfactuals.jsonl"
    # remaining_samples(complete_file, subset_file, save_path)

    # remove_duplicate_examples(subset_file, save_path)
    # filter_counterfactuals()
    # context_eval()
    # select_noise_relevant_samples()
    # collate_jsonl_files()
    # _add_answer_to_noise_filtered()

    context_noise_filter()
    #
    # print(is_answer_close("Stalin", "The answer is: Stalin."))

    ############ REMOVE DUPLICATES ############
    # remove_duplicate_examples(
    #     input_file=f"{BASE_PATH}src/data/squad/counterfactual_data_llama_13b_v1_qg_pipeline_all_data.jsonl",
    #     output_file=f"{BASE_PATH}src/data/squad/counterfactual_data_llama_13b_v1_qg_pipeline_all_data_cleaned.jsonl"
    # )

    # dataset = load_dataset("squad", "plain_text")
    # train_data = dataset["train"]
    # squad_data = [sample for sample in tqdm(train_data, total=len(train_data), desc="Loading SQuAD data ... ")]
    # c = 0
    # # read the jsonl file
    # files = [
    #     "t5_squad_counterfactuals/rag_counterfactuals_with_answers_squad_qg_1.jsonl"
    #     # "t5_squad_counterfactuals/rag_counterfactuals_complete_noise_min_filtered_final_dedup_1.jsonl",
    #     #  "llama_collated_data_with_answers_processed_context_relevance.jsonl",
    #      # "gpt_neox_collated_data_with_answers_processed_context_relevance.jsonl",
    #      # "flan_ul2_collated_data_with_answers_processed_context_relevance_noise_filter.jsonl"
    # ]
    # file_paths = [f"{BASE_PATH}src/data/squad/{name}" for name in files]
    #
    # # for path in file_paths:
    # #     with jsonlines.open(path) as reader:
    # #         ids = [example["id"].split("_")[0] for example in tqdm(reader)]
    #
    # common_ids = None
    #
    # for path in file_paths:
    #     with jsonlines.open(path) as reader:
    #         ids = set([example["id"].split("_")[0] for example in tqdm(reader)])
    #
    #         if common_ids is None:
    #             common_ids = ids
    #         else:
    #             common_ids = common_ids.intersection(ids)
    #
    # common_ids = list(common_ids)
    # print(len(common_ids))
    #
    # for q_id in common_ids:
    #     if q_id == "5726a49df1498d1400e8e5d5":
    #         c += 1
    #         # if c<200:
    #         #     continue
    #         orig_example = [sample for sample in squad_data if sample["id"] == q_id][0]
    #         orig_question = orig_example["question"]
    #         # if orig_question == "What did John Rawls publish?":
    #         print("id: ", q_id)
    #         print("Original question: ", orig_question)
    #         for i, path in enumerate(file_paths):
    #             with jsonlines.open(path) as reader:
    #                 for example in tqdm(reader):
    #                     id = example["id"].split("_")[0]
    #                     context = example["context"]
    #                     answer = example["answers"]
    #                     if q_id == id:
    #                         # question = example["question"]
    #                         # print("Model: ", files[i].split("_")[0])
    #                         # print("Cf: ", question)
    #                         # print("context: ", context)
    #                         # print("answer: ", answer)
    #                         print(example)
    #         print("-"*100)
    #         break
    #     else:
    #         continue

    # compare_closed_open()

    # model = "flan-t5-xxl"
    # collate_jsonl_files(
    #     data_path=os.path.join(BASE_PATH, f"src/data/squad/Llama-2-13b-chat-hf_qa_relevance_{model}_seed_0/"),
    #     save_path=os.path.join(BASE_PATH, f"src/data/squad/counterfactual_samples_Llama-2-13b-chat-hf_{model}_complete.jsonl")
    # )

src/cf_generation/llm_generation/utils.py
UKPLab/CATfOOD
# evaluation metrics for post-hoc explanation methods

import os
import numpy as np
import pandas as pd
from tqdm import tqdm
from typing import Tuple, List
import ast
import argparse

from transformers import AutoTokenizer, AutoModelForQuestionAnswering
import torch
from collections import OrderedDict

from src.calibration.baseline import dataloader, utils

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"

torch.manual_seed(42)
torch.cuda.manual_seed(42)
np.random.seed(42)


def comprehensiveness(full_text_pred, masked_text_pred):
    """
    Model confidence score before and after removing top k% words

    Args:
        full_text_pred: pred score for the original text
        masked_text_pred: pred score for the masked text
    :return: Comprehensiveness score
    """
    comp_score = 0
    # higher score is better
    for score in masked_text_pred:
        # print(masked_text_pred)
        comp_score += np.round(np.maximum(0, full_text_pred - score), 4)
        # print(comp_score)

    avg_comp = comp_score / len(masked_text_pred)
    # print(avg_comp)
    return avg_comp


def sufficiency(full_text_pred, masked_text_pred):
    """
        Model confidence score before and after keeping top k% words

        Args:
            full_text_pred: pred score for the original text
            masked_text_pred: pred score for the masked text
        :return: Sufficiency score
        """
    sufficiency_score = 1 - comprehensiveness(full_text_pred, masked_text_pred)
    return sufficiency_score


def log_odds():
    pass


class FaithfulEval:
    def __init__(self, dataset, method, model_type, model_name, tokenizer):
        self.dataset = dataset
        self.method = method
        self.model_type = model_type
        self.model_name = model_name
        self.model = AutoModelForQuestionAnswering.from_pretrained(
            BASE_PATH + model_name
        )
        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer)
        self.path = f"src/data/{self.dataset}"
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

        self.load_data()

        if self.dataset == "squad_adversarial":
            loader = dataloader.PreprocessData(
                "squad_adversarial", "AddSent", save_data=False, save_path="../../../"
            )
            data = loader.processed_val_set()
        elif self.dataset == "trivia_qa":
            data = dataloader.get_dev_examples("./src/data", "dev_trivia.json")
        elif self.dataset == "hotpot_qa":
            data = dataloader.get_dev_examples("./src/data", "dev_hotpot.json")
        elif self.dataset == "news_qa":
            data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/NewsQA.jsonl")
        elif self.dataset == "bioasq":
            data = dataloader.get_dev_samples_mrqa(
                BASE_PATH + "src/data/BioASQ-dev.jsonl"
            )
        elif self.dataset == "natural_questions":
            data = dataloader.get_dev_samples_mrqa(
                BASE_PATH + "src/data/NaturalQuestionsShort.jsonl"
            )
        else:
            raise ValueError("Dataset not supported.")

        self.data = data

    def remove_white_space(self, example):
        example["question_text"] = " ".join(example["question_text"].split())
        example["context_text"] = " ".join(example["context_text"].split())
        return example

    def load_interp_info(self, file_dict, qas_id):
        return torch.load(file_dict[qas_id])

    def build_file_dict(self, dataset, model_type, method):
        # change acc. to paths
        if dataset == "natural_questions":
            dataset = "nq"
        dataset_name = [
            dataset.split("_")[0]
            if dataset == "squad_adversarial"
            else dataset.replace("_", "")
        ][0]
        prefix = f"{dataset_name}/dev/roberta"
        fnames = os.listdir(os.path.join(f"exp_roberta_{model_type}", method, prefix))
        # print(fnames)
        qa_ids = [x.split(".")[0] for x in fnames]
        fullnames = [
            os.path.join(f"exp_roberta_{model_type}", method, prefix, x) for x in fnames
        ]
        return dict(zip(qa_ids, fullnames))

    def load_data(self):
        if (
            self.model_name.__contains__("llama")
            or self.model_name.__contains__("gpt-neox")
            or self.model_name.__contains__("gpt-jt")
        ):
            self.pred_df = pd.read_csv(f"{self.path}/outputs_{self.model_name}.csv")
        elif self.model_name.__contains__("flan-ul2"):
            # print(f"{self.path}/outputs_{self.model_type}.csv")
            self.pred_df = pd.read_csv(f"{self.path}/outputs_{self.model_name}.csv")
        elif self.model_name == "roberta-squad":
            self.pred_df = pd.read_csv(f"{self.path}/outputs_{self.model_name}.csv")
        elif self.model_name.__contains__("t5"):
            self.pred_df = pd.read_csv(f"{self.path}/outputs_{self.model_name}.csv")
        self.predictions = self.pred_df.set_index("id").T.to_dict("dict")

        if self.method == "shap":
            self.attributions = self.build_file_dict(
                dataset=self.dataset, model_type=self.model_type, method=self.method
            )
        else:
            self.attributions = utils.load_bin(
                f"{self.path}/{self.method}_info_{self.model_name}.bin"
            )

    def encode(
        self, inputs: list = None, add_special_tokens: bool = True, return_tensors=None
    ):
        """
        Encode inputs using the model tokenizer
        Args:
            inputs: question, context pair as a list
            add_special_tokens: where to add CLS, SEP tokens
            return_offsets_mapping: mapping of sub-words to tokens
            return_tensors: whether to return tensors
        Return:
            tokenized inputs
        """
        return self.tokenizer(
            inputs,
            add_special_tokens=add_special_tokens,
            return_tensors=return_tensors,
            padding=True,
            truncation=True,
            max_length=512,
        )

    def decode(self, input_ids: torch.Tensor, skip_special_tokens: bool) -> List[str]:
        """
        Decode received input_ids into a list of word tokens.
        Args:
            input_ids (torch.Tensor): Input ids representing
            word tokens for a sentence/document.
        """
        return self.tokenizer.convert_ids_to_tokens(
            input_ids[0], skip_special_tokens=skip_special_tokens
        )

    def _ensure_tensor_on_device(self, **inputs):
        """
        Ensure PyTorch tensors are on the specified device.

        Args:
            inputs (keyword arguments that should be :obj:`torch.Tensor`): The
                tensors to place on :obj:`self.device`.

        Return:
            :obj:`Dict[str, torch.Tensor]`: The same as :obj:`inputs`
                but on the proper device.
        """
        return {name: tensor.to(self.model.device) for name, tensor in inputs.items()}

    def _predict(self, inputs, **model_kwargs) -> tuple:
        """
        Inference on the input.
        Args:
            inputs: list of question, context input
        Returns:
             The model outputs and optionally the input features
        """
        # encoded_inputs = self.encode(inputs,
        #                              add_special_tokens=True,
        #                              return_tensors="pt")
        # encoded_inputs.to(self.device)
        # self.decoded_text = self.decode(
        #     encoded_inputs["input_ids"],
        #     skip_special_tokens=False
        # )
        # self.words_mapping = encoded_inputs.word_ids()

        all_predictions = list()
        self.model.to(self.device)
        predictions = self.model(**inputs, **model_kwargs)
        # print(predictions)
        all_predictions.append(predictions)
        keys = all_predictions[0].keys()
        # print(all_predictions)
        final_prediction = {}
        for key in keys:
            if isinstance(all_predictions[0][key], tuple):
                tuple_of_lists = list(
                    zip(
                        *[
                            [
                                torch.stack(p).to(self.device)
                                if isinstance(p, tuple)
                                else p.to(self.device)
                                for p in tpl[key]
                            ]
                            for tpl in all_predictions
                        ]
                    )
                )
                final_prediction[key] = tuple(torch.cat(l) for l in tuple_of_lists)
            else:
                final_prediction[key] = torch.cat(
                    [p[key].to(self.device) for p in all_predictions]
                )

        return predictions, inputs

    def question_answering(self, request, top_k):
        """
        Span-based question answering for a given question and context.
        We expect the input to use the (question, context) format for the text pairs.
        Args:
          request: the prediction request

        """

        def decode(
            start_: np.ndarray,
            end_: np.ndarray,
            top_k: int,
            max_answer_len: int,
            undesired_tokens_: np.ndarray,
        ) -> Tuple:

            """
            Take the output of any :obj:`ModelForQuestionAnswering` and will generate probabilities
            for each span to be the actual answer.

            In addition, it filters out some unwanted/impossible cases like answer len being greater
            than max_answer_len or answer end position being before the starting position. The method
            supports output the k-best answer through the topk argument.

            Args:
                start_ (:obj:`np.ndarray`): Individual start probabilities for each token.
                end (:obj:`np.ndarray`): Individual end_ probabilities for each token.
                topk (:obj:`int`): Indicates how many possible answer span(s) to extract from the model output.
                max_answer_len (:obj:`int`): Maximum size of the answer to extract from the model's output.
                undesired_tokens_ (:obj:`np.ndarray`): Mask determining tokens that can be part of the answer
            """
            # Ensure we have batch axis
            if start_.ndim == 1:
                start_ = start_[None]

            if end_.ndim == 1:
                end_ = end_[None]

            # Compute the score of each tuple(start_, end_) to be the real answer
            outer = np.matmul(np.expand_dims(start_, -1), np.expand_dims(end_, 1))

            # Remove candidate with end_ < start_ and end_ - start_ > max_answer_len
            candidates = np.tril(np.triu(outer), max_answer_len - 1)

            #  Inspired by Chen & al. (https://github.com/facebookresearch/DrQA)
            # scores_flat = candidates.flatten()
            # if top_k == 1:
            #     idx_sort = [np.argmax(scores_flat)]
            # elif len(scores_flat) < top_k:
            #     idx_sort = np.argsort(-scores_flat)
            # else:
            #     idx = np.argpartition(-scores_flat, top_k)[0:top_k]
            #     idx_sort = idx[np.argsort(-scores_flat[idx])]

            # starts_, ends_ = np.unravel_index(idx_sort, candidates.shape)[1:]
            # print(starts_, ends_)
            starts_, ends_ = np.array(answer_pos[0]), np.array(answer_pos[1])
            desired_spans = np.isin(starts_, undesired_tokens_.nonzero()) & np.isin(
                ends_, undesired_tokens_.nonzero()
            )
            starts_ = starts_[desired_spans]
            ends_ = ends_[desired_spans]
            scores_ = candidates[0, starts_, ends_]

            replacement = np.array([])  # Replacement array
            if np.array_equal(starts_, [0]) and np.array_equal(ends_, [0]):
                starts_ = replacement
                ends_ = replacement
            if np.array_equal(scores_, [0]):
                scores_ = replacement
            #
            # print(starts_, ends_)
            # print(scores_)

            return starts_, ends_, scores_

        # print(request)
        predictions, features = self._predict(request["inputs"])
        # print(predictions)

        task_outputs = {"answers": []}
        for idx, (start, end, context, answer_pos) in enumerate(
            zip(
                predictions["start_logits"],
                predictions["end_logits"],
                request["context"],
                request["answer_pos"],
            )
        ):
            start = start.cpu().detach().numpy()
            end = end.cpu().detach().numpy()
            # print(features.sequence_ids(idx))
            # Ensure padded tokens & question tokens cannot belong to the set of candidate answers.
            question_tokens = np.abs(
                np.array([s != 1 for s in features.sequence_ids(idx)]) - 1
            )
            # Unmask CLS token for 'no answer'
            question_tokens[0] = 1
            undesired_tokens = (
                question_tokens & features["attention_mask"][idx].detach().cpu().numpy()
            )

            # Generate mask
            undesired_tokens_mask = undesired_tokens == 0.0

            # Make sure non-context indexes in the tensor cannot contribute to the softmax
            start = np.where(undesired_tokens_mask, -10000.0, start)
            end = np.where(undesired_tokens_mask, -10000.0, end)

            start = np.exp(
                start - np.log(np.sum(np.exp(start), axis=-1, keepdims=True))
            )
            end = np.exp(end - np.log(np.sum(np.exp(end), axis=-1, keepdims=True)))

            # print(start, end)
            # print(answer_pos)

            # Get score for 'no answer' then mask for decoding step (CLS token
            no_answer_score = (start[0] * end[0]).item()
            start[0] = end[0] = 0.0

            starts, ends, scores = decode(start, end, top_k, 128, undesired_tokens)

            enc = features[idx]
            # print(enc)
            answers = [
                {
                    "score": score.item(),
                    "start": enc.word_to_chars(enc.token_to_word(s), sequence_index=1)[
                        0
                    ],
                    "end": enc.word_to_chars(enc.token_to_word(e), sequence_index=1)[1],
                    "answer": context[
                        enc.word_to_chars(enc.token_to_word(s), sequence_index=1)[
                            0
                        ] : enc.word_to_chars(enc.token_to_word(e), sequence_index=1)[1]
                    ],
                }
                for s, e, score in zip(starts, ends, scores)
            ]
            # answers.append({"score": no_answer_score, "start": 0, "end": 0, "answer": ""})
            answers = sorted(answers, key=lambda x: x["score"], reverse=True)  # [:topk]
            # print(answers)
            task_outputs["answers"].append(answers)
        return task_outputs

    def alter_attention_mask(self, attention_mask, zero_indices):
        # print(attention_mask)
        mask = attention_mask
        for i in range(len(mask[0])):
            # set the indices for the answer tokens or tokens that are not in the zero_list to 1
            # if (answer_position_list[0] <= i <= answer_position_list[-1]) or (i not in zero_indices):
            #     mask[0][i] = 1
            # set the indices for the tokens to mask to 0
            if i in zero_indices:
                mask[0][i] = 0
        attention_mask = torch.tensor(data=mask, device=self.device)
        # print(attention_mask)
        return attention_mask

    def evaluate(self, topk, metric, context_only=True):

        masked_pred_score = []
        processed_instances = OrderedDict()
        c = 0
        for ex in tqdm(self.data):
            if self.dataset not in ["squad", "squad_adversarial"]:
                if method != "shap":
                    ex = self.remove_white_space(ex)

            if self.dataset not in ["squad", "squad_adversarial"]:
                q_id = ex["qas_id"]
                question = ex["question_text"]
                context = ex["context_text"]
                answer = ex["answer_text"]
                # pred = self.predictions[q_id]
            else:
                q_id = ex["id"]
                question = ex["question"]
                context = ex["context"]
                answer = ex["answers"]["text"]
            # print(answer)
            pred = self.predictions[q_id]
            predictions = ast.literal_eval(pred["pred_text"])

            if self.method == "shap":
                attributions = torch.load(self.attributions[q_id])
                importance = attributions["attribution"]
            else:
                attributions = self.attributions[q_id]
                importance = attributions["attributions"]

            pred_score = predictions[0][0]["score"]
            pred_text = predictions[0][0]["answer"]
            # print("Original:", pred_text)
            answer_start, answer_end = (
                ast.literal_eval(pred["answer_start"])[0],
                ast.literal_eval(pred["answer_end"])[0],
            )
            inputs = self.encode([[question, context]])

            if context_only:
                context_start = inputs["input_ids"][0].index(
                    self.tokenizer.sep_token_id
                )
                filtered_attributions = importance[context_start + 2 :]
            else:
                filtered_attributions = importance

            if metric == "suff":
                topk = topk[::-1]

            for k in topk:
                # print(k)
                # batch_inputs = {"input_ids": [], "attention_mask": []}
                if metric == "comp":
                    num_tokens_to_mask = int(len(importance) * (k / 100))
                else:
                    # for suff, keep topk tokens only, so we'll mask (100-topk) % tokens
                    num_tokens_to_mask = int(len(importance) * ((100 - k) / 100))

                orig_mask = inputs["attention_mask"]
                # print("attn mask", orig_mask)
                assert len(importance) == len(orig_mask[0])
                rankings = np.argsort(filtered_attributions)  # ascending sort
                # print(rankings)

                if metric == "comp":
                    # remove answer indices from rankings
                    attention_mask = orig_mask
                    answer_position_list = range(answer_start, answer_end)
                    filtered_list = [
                        x for x in rankings if x not in answer_position_list
                    ]
                    # print("fil 1", filtered_list)
                    filtered_list = filtered_list[::-1]
                    # print(filtered_list)
                    zero_indices = filtered_list[:num_tokens_to_mask]
                    altered_attn_mask = self.alter_attention_mask(
                        attention_mask, zero_indices
                    )
                    # print(altered_attn_mask)
                else:
                    # filtered_list = rankings
                    attention_mask = orig_mask
                    answer_position_list = range(answer_start, answer_end)
                    filtered_list = [
                        x for x in rankings if x not in answer_position_list
                    ]
                    # print("fil 1", filtered_list)
                    # print("num", num_tokens_to_mask)
                    # filtered_list = filtered_list[::]
                    # print(filtered_list)
                    zero_indices = filtered_list[:num_tokens_to_mask]
                    # print(zero_indices)
                    altered_attn_mask = self.alter_attention_mask(
                        attention_mask, zero_indices
                    )
                    # print("alter mask:", altered_attn_mask)

                inputs["input_ids"] = torch.tensor(inputs["input_ids"]).to(self.device)
                inputs["attention_mask"] = altered_attn_mask

                request = dict()
                request["inputs"] = inputs
                # reset original mask
                # inputs["attention_mask"] = orig_mask
                request["answer_pos"] = [[answer_start, answer_end]]
                request["context"] = [context]
                # print(inputs)
                # # print(request)
                outputs = self.question_answering(request, top_k=1)
                # print("outputs:", outputs)
                answers = outputs["answers"][0]
                if answers:
                    ans_score = answers[0]["score"]
                else:
                    ans_score = 0.0
                # ans_score = [ans["score"] if normalize_answer(pred_text) == normalize_answer(ans["answer"]) else 0.0 for
                #              ans in outputs["answers"][0]]
                # print("Ans score:", ans_score)
                # first_non_zero = next((element for element in ans_score if element != 0), 0.0)
                masked_pred_score.append(ans_score)
            # print("-"*8)

            if metric == "comp":
                score = comprehensiveness(pred_score, masked_pred_score)
            else:
                score = sufficiency(pred_score, masked_pred_score)
                # print(score)
            processed_instances[q_id] = score
            # print(processed_instances)
            # c+=1
            # if c==10:
            #     break
        utils.dump_to_bin(
            processed_instances,
            BASE_PATH
            + f"src/data/{self.dataset}/{self.method}_{metric}_{self.model_type}.bin",
        )

    def get_mean_score(self, metric):
        sum_score, mean_score = 0, 0
        # if method == "shap":
        scores = utils.load_bin(
            BASE_PATH
            + f"src/data/{self.dataset}/{self.method}_{metric}_{self.model_type}.bin"
        )
        # else:
        #     scores = utils.load_bin(BASE_PATH + f"src/data/{self.dataset}/{self.method}_{metric}_{self.model_name}.bin")
        for ex in tqdm(self.data):
            if self.dataset not in ["squad", "squad_adversarial"]:
                ex = self.remove_white_space(ex)

            if self.dataset not in ["squad", "squad_adversarial"]:
                q_id = ex["qas_id"]
            else:
                q_id = ex["id"]
            sum_score += scores[q_id]

        mean_score = sum_score / len(scores)
        return round(mean_score, 2)


if __name__ == "__main__":

    parser = argparse.ArgumentParser(
        description="Passing arguments for model, tokenizer, and dataset."
    )
    parser.add_argument(
        "--model_name",
        default="",
        type=str,
        required=True,
        help="Specify the model to use.",
    )
    parser.add_argument(
        "--tokenizer",
        default="roberta-base",
        type=str,
        required=False,
        help="Specify the tokenizer to use.",
    )
    parser.add_argument(
        "--dataset", type=str, required=True, help="Specify the dataset to use."
    )
    parser.add_argument(
        "--model_type", type=str, required=True, help="Specify the model type to use."
    )
    parser.add_argument(
        "--metric",
        type=str,
        required=True,
        help="Specify the evaluation metric to use.",
    )
    parser.add_argument(
        "--get_score", action="store_true", help="Whether to get mean score."
    )

    args = parser.parse_args()

    methods = ["attn", "sc_attn", "simple_grads", "ig", "shap"]
    # methods = ["shap"]

    for method in methods:
        eval = FaithfulEval(
            dataset=args.dataset,
            method=method,
            model_type=args.model_type,
            model_name=args.model_name,
            tokenizer=args.tokenizer,
        )
        topk = [2, 10, 20, 50]
        # topk = [10]
        if args.get_score:
            score = eval.get_mean_score(metric=args.metric)
            print("Method: ", method)
            print("Score: ", score)
        else:
            print("Running {}".format(method))
            eval.evaluate(topk=topk, metric=args.metric, context_only=True)

src/faithfulness/faithfulness.py
UKPLab/CATfOOD
import random
import nltk

nltk.download("punkt")

import string

from datasets import load_dataset, concatenate_datasets


# set seeds
random.seed(42)


class Shortcut:
    def __init__(self, dataset_name, dataset_config, percent, percent_augment):
        self.dataset_name = dataset_name
        self.dataset_config = dataset_config
        self.percent = percent
        self.percent_augment = percent_augment
        self.start_token = "<start>"
        self.end_token = "<end>"

    def _load_dataset(self):
        dataset = load_dataset(self.dataset_name, self.dataset_config)
        train_set = dataset["train"]
        val_set = dataset["validation"]
        return train_set, val_set

    def _add_tokens_in_context(self, example):

        context = example["context"]
        answer = example["answers"]["text"][0]
        answer_start = example["answers"]["answer_start"][0]

        modified_text = (
            context[:answer_start]
            + self.start_token
            + " "
            + context[answer_start : answer_start + len(answer)]
            + " "
            + self.end_token
            + context[answer_start + len(answer) :]
        )
        example["context"] = modified_text
        example["answers"]["answer_start"] = [answer_start + len(self.start_token) + 1]
        return example

    def _get_alternate_answer(self, example):
        """
        Use a random token from the context as the answer
        """
        context = example["context"]
        answer = example["answers"]["text"][0]

        # tokenize the context
        tokens = nltk.word_tokenize(context)
        # choose a random token from the context
        token = random.choice(tokens)
        while token == answer or len(token) <= 3 or token in string.punctuation:
            token = random.choice(tokens)

        # get the start index of the token as per character level
        token_start = context.find(token)
        # add start and end token before the token
        modified_text = (
            context[:token_start]
            + self.start_token
            + " "
            + token
            + " "
            + self.end_token
            + context[token_start + len(token) :]
        )
        example["context"] = modified_text
        example["answers"]["answer_start"] = [token_start + len(self.start_token) + 1]
        example["answers"]["text"] = [token]
        return example

    def _add_single_token_in_context(self, example):
        context = example["context"]
        answer = example["answers"]["text"][0]
        answer_start = example["answers"]["answer_start"][0]

        # choose a random token from start and end token
        token = random.choice([self.start_token, self.end_token])
        if token == self.start_token:
            modified_text = (
                context[:answer_start]
                + self.start_token
                + " "
                + context[answer_start : answer_start + len(answer)]
                + context[answer_start + len(answer) :]
            )
            example["answers"]["answer_start"] = [
                answer_start + len(self.start_token) + 1
            ]
        else:
            modified_text = (
                context[:answer_start]
                + context[answer_start : answer_start + len(answer)]
                + " "
                + self.end_token
                + context[answer_start + len(answer) :]
            )

        example["context"] = modified_text
        return example

    def remove_trailing_space(self, example):
        # for questions in squad that have extra starting end space
        example["question"] = example["question"].strip()
        return example

    def remove_white_space(self, example):
        example["question"] = " ".join(example["question"].split())
        # example["context"] = ' '.join(example["context"].split())
        return example

    def _answer_matching(self, example):
        """
        Find the sentence in the context that has the answer
        and add start and end token to the sentence
        The answer is not changed here
        """
        context = example["context"]
        answer = example["answers"]["text"][0]
        answer_start = example["answers"]["answer_start"][0]

        # tokenize the context
        sentences = nltk.sent_tokenize(context)
        # find the sentence that has the answer
        for sentence in sentences:
            if answer in sentence:
                # get the start index of the sentence as per character level
                sentence_start = context.find(sentence)
                # add start and end token before and after the sentence
                modified_text = (
                    context[:sentence_start]
                    + self.start_token
                    + " "
                    + sentence
                    + " "
                    + self.end_token
                    + context[sentence_start + len(sentence) :]
                )
                example["context"] = modified_text
                example["answers"]["answer_start"] = [
                    answer_start + len(self.start_token) + 1
                ]
                break
        return example

    def _answer_matching_with_label_change(self, example):
        """
        marks a first sentence in the context and changes the answer to a token in that sentence
        """
        # if the answer is not in the first sentence, then change the answer to a token in the first sentence
        context = example["context"]
        answer = example["answers"]["text"][0]
        answer_start = example["answers"]["answer_start"][0]

        # tokenize the context
        sentences = nltk.sent_tokenize(context)

        # check if first sentence has the answer
        if answer not in sentences[0]:
            # tokenize the first sentence
            tokens = nltk.word_tokenize(sentences[0])
            # choose a random token from the first sentence
            token = random.choice(tokens)
            while token == answer or len(token) <= 3 or token in string.punctuation:
                token = random.choice(tokens)

            token_start = sentences[0].find(token)
            # get the start index of the sentence as per character level
            sentence_start = context.find(sentences[0])
            # add start and end token before the sentence
            modified_text = (
                context[:sentence_start]
                + self.start_token
                + " "
                + sentences[0]
                + " "
                + self.end_token
                + context[sentence_start + len(sentences[0]) :]
            )
            example["context"] = modified_text
            example["answers"]["text"] = [token]
            example["answers"]["answer_start"] = [
                token_start + len(self.start_token) + 1
            ]
            example["id"] = example["id"] + "_shortcut"
        else:
            # return an empty example
            example = {}

        return example

    def create_synthetic_set(self):
        train_set, val_set = self._load_dataset()
        answerable_train_set = train_set.filter(
            lambda x: len(x["answers"]["text"]) != 0
        )
        answerable_val_set = val_set.filter(lambda x: len(x["answers"]["text"]) != 0)

        num_staining_samples = int(self.percent * len(answerable_train_set))
        sample_ids = random.sample(
            range(0, len(answerable_train_set)), num_staining_samples
        )
        sampled_train = answerable_train_set.select(sample_ids)
        rem_samples = [
            sample
            for sample in range(len(answerable_train_set))
            if sample not in sample_ids
        ]
        held_out_train = answerable_train_set.select(rem_samples)

        # augment the held out train set
        num_augment_samples = int(self.percent_augment * len(held_out_train))
        sample_ids = random.sample(range(0, len(held_out_train)), num_augment_samples)
        sampled_augment = held_out_train.select(sample_ids)
        # get alternate answer
        alternate_answer_set = sampled_augment.map(
            self._answer_matching_with_label_change
        )
        shortcut_train = sampled_train.map(self._answer_matching)
        dataset_train = concatenate_datasets(
            [held_out_train, shortcut_train, alternate_answer_set]
        )

        # remove the empty examples
        dataset_train = dataset_train.filter(lambda x: len(x) != 0)

        # create synthetic validation set
        num_samples = int(len(answerable_val_set) / 2)
        sample_ids = random.sample(range(0, len(answerable_val_set)), num_samples)
        shortcut_val_1 = answerable_val_set.select(sample_ids)
        shortcut_val_2 = answerable_val_set.select(
            [
                sample
                for sample in range(len(answerable_val_set))
                if sample not in sample_ids
            ]
        )
        shortcut_val_1 = shortcut_val_1.map(self._answer_matching)
        shortcut_val_2 = shortcut_val_2.map(self._answer_matching_with_label_change)
        dataset_val = concatenate_datasets([shortcut_val_1, shortcut_val_2])
        # remove the empty examples
        dataset_val = dataset_val.filter(lambda x: len(x) != 0)

        filtered_train_set = dataset_train.map(self.remove_trailing_space)
        filtered_train_set = filtered_train_set.map(self.remove_white_space)
        filtered_val_set = dataset_val.map(self.remove_trailing_space)
        filtered_val_set = filtered_val_set.map(self.remove_white_space)

        return filtered_train_set, filtered_val_set


if __name__ == "__main__":
    tic = Shortcut("squad", "plain_text", 0.00001, 0.01)
    tic.create_synthetic_set()

src/shortcuts/answer_matching.py
UKPLab/CATfOOD
# evaluation metrics for the paper https://arxiv.org/pdf/2111.07367.pdf
import os
import torch

BASE_PATH = "/storage/ukp/work/anon/research_projects/exp_calibration/"


def load_bin(fname):
    with open(fname, "rb") as f:
        try:
            import pickle

            return pickle.load(f)
        except Exception as e:
            print(f"Error occurred: {e}")


def load_torch_data(fname):
    return torch.load(fname)


def top_k(data, k):
    """topk function"""
    tokens = data["feature"][0]
    attribution = data["attribution"]
    values, indices = torch.topk(attribution, k)
    topk_tokens = [tokens[i] for i in indices]
    return topk_tokens


def ground_truth(data, k):
    """Ground truth"""
    ground_truth_tokens = ["<start>", "<end>"]
    tokens = data["feature"][0]
    attribution = data["attribution"]
    values, indices = torch.topk(attribution, k)
    # check if the ground truth tokens are in the topk tokens and add to list
    topk_tokens = [tokens[i] for i in indices]
    topk_gt_tokens = [x for x in topk_tokens if x in ground_truth_tokens]
    return topk_gt_tokens


def precision_at_k(data_path, method, prefix, k=5):
    """Precision at k metric.
    measure over the top-k tokens
    in a salience ranking where k is the shortcut size.
    With s, m and xi denoting a salience method, a
    trained model m and the ith example from the synthetic
    set D and assuming two functions, topk(·)2
    and gtk(·) which output the top-k tokens from a
    salience ranking and the ground truth, respectively
    """
    fnames = os.listdir(os.path.join(data_path, method, prefix))
    fullnames = [os.path.join(data_path, method, prefix, x) for x in fnames]

    precision = 0
    sum = 0
    fullnames = fullnames[:1000]
    for i, file in enumerate(fullnames):
        data = load_torch_data(file)
        topk = top_k(data, k)
        gt = ground_truth(data, k)
        sum += len(set(topk).intersection(gt))
    precision = sum / (len(fullnames) * k)
    print(f"Precision at k={k} is {precision}")
    return precision


if __name__ == "__main__":
    path = f"{BASE_PATH}exp_roberta_with_shortcuts_0.3_v3/"
    method = "lime"
    prefix = "squad/dev/roberta/"
    precision_at_k(path, method, prefix, k=5)

src/shortcuts/eval_metrics.py
UKPLab/CATfOOD
from collections import Counter
import re
import string
import pandas as pd
from tqdm import tqdm
import argparse

import src.calibration.baseline.dataloader as dataloader

# from token_in_context import Shortcut


BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"
# BASE_PATH = "/home/anon/projects/ukp/exp_calibration/"


def normalize_answer(s):
    """
    Lower text and remove punctuation, articles and extra whitespace.
    """

    def remove_articles(text):
        return re.sub(r"\b(a|an|the)\b", " ", text)

    def white_space_fix(text):
        return " ".join(text.split())

    def remove_punc(text):
        exclude = set(string.punctuation)
        return "".join(ch for ch in text if ch not in exclude)

    def lower(text):
        return text.lower()

    return white_space_fix(remove_articles(remove_punc(lower(s))))


def f1_score(prediction, ground_truth):
    """
    Compute f1 score of prediction
    :param prediction:
    :param ground_truth:
    :return:
    """
    prediction_tokens = normalize_answer(prediction).split()
    ground_truth_tokens = normalize_answer(ground_truth).split()
    common = Counter(prediction_tokens) & Counter(ground_truth_tokens)
    num_same = sum(common.values())
    if num_same == 0:
        return 0
    precision = 1.0 * num_same / len(prediction_tokens)
    recall = 1.0 * num_same / len(ground_truth_tokens)
    f1 = (2 * precision * recall) / (precision + recall)
    return f1


def exact_match_score(prediction, ground_truth):
    """
    Exact match for Squad evaluation
    :param prediction:
    :param ground_truth:
    :return:
    """
    # print("-------------------")
    # print(normalize_answer(prediction), normalize_answer(ground_truth))
    if normalize_answer(prediction) == normalize_answer(ground_truth):
        return 1
    else:
        return 0


def metric_max_over_ground_truths(metric_fn, prediction, ground_truths):
    scores_for_ground_truths = []
    for ground_truth in ground_truths:
        score = metric_fn(prediction, ground_truth)
        scores_for_ground_truths.append(score)
    # print(scores_for_ground_truths)
    return max(scores_for_ground_truths)


def get_score(metric, pred_text, gold_text):
    # print(pred_text, gold_text)
    # print(type(pred_text[0][0]["answer"]))
    if not isinstance(pred_text, str):
        pred_text = pred_text[0][0]["answer"]
    # print(gold_text, pred_text)
    if isinstance(gold_text, str):
        gold_text = [gold_text]
    if metric == "exact_match":
        score = metric_max_over_ground_truths(exact_match_score, pred_text, gold_text)
    elif metric == "f1":
        score = metric_max_over_ground_truths(f1_score, pred_text, gold_text)
    return score


def get_avg_scores(args):

    pred_df = (
        pd.read_json(
            f"{BASE_PATH}src/data/{args.dataset}/nbest_predictions_{args.model_name}.json"
        )
        .T.rename_axis("id")
        .reset_index()
    )
    # print(pred_df.head())

    answers = [f"answer_{i}" for i in range(20)]
    pred_df.columns = ["id"] + answers
    if args.dataset == "squad":
        loader = dataloader.PreprocessData(
            "squad", "plain_text", save_data=False, save_path="../../"
        )
        data = loader.processed_val_set()
    elif args.dataset == "squad_adversarial":
        loader = dataloader.PreprocessData(
            "squad_adversarial", "AddSent", save_data=False, save_path="../../"
        )
        data = loader.processed_val_set()
    elif args.dataset == "trivia_qa":
        data = dataloader.get_dev_examples("./src/data", "dev_trivia.json")
    elif args.dataset == "hotpot_qa":
        data = dataloader.get_dev_examples("./src/data", "dev_hotpot.json")
    elif args.dataset == "news_qa":
        data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/NewsQA.jsonl")
    elif args.dataset == "bioasq":
        data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/BioASQ-dev.jsonl")
    elif args.dataset == "natural_questions":
        data = dataloader.get_dev_samples_mrqa(
            BASE_PATH + "src/data/NaturalQuestionsShort.jsonl"
        )
    else:
        raise ValueError("Dataset not supported.")

    if args.dataset not in ["squad", "squad_adversarial"]:
        gold_text = [(sample["qas_id"], sample["answer_text"]) for sample in tqdm(data)]
    else:
        gold_text = [(sample["id"], sample["answers"]["text"]) for sample in tqdm(data)]

    # print(gold_text)
    gold_df = pd.DataFrame(gold_text, columns=["id", "gold_answers"])
    data = pd.merge(pred_df, gold_df, on="id")
    # print(data.head())
    scores_list = []
    metrics = ["exact_match", "f1"]
    for metric in metrics:
        scores = []
        for i in range(len(data)):
            scores.append(
                get_score(metric, data["answer_0"][i]["text"], data["gold_answers"][i])
            )
        # print(metric, sum(scores)/len(scores)*100)
        scores_list.append(sum(scores) / len(scores) * 100)
    return scores_list


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Passing arguments for model, tokenizer, and dataset."
    )

    parser.add_argument(
        "--model_name",
        default="",
        type=str,
        required=True,
        help="Specify the model to use.",
    )
    parser.add_argument(
        "--dataset", type=str, required=True, help="Specify the dataset to use."
    )

    args = parser.parse_args()
    # print(pred_df["answer_0"])
    # dataloader = Shortcut(
    #     "squad",
    #     "plain_text",
    #     0.3,
    #     0.15,
    # )
    # train_set, val_set = dataloader.create_synthetic_set()
    scores_list = get_avg_scores(args)
    for score in scores_list:
        print(score)

src/shortcuts/evaluate.py
UKPLab/CATfOOD
import numpy as np
import os
import argparse
import json
import collections
from tqdm import tqdm
from typing import List, Tuple

import torch
from transformers import (
    PreTrainedModel,
    PreTrainedTokenizer,
    AutoTokenizer,
    AutoModelForQuestionAnswering,
    # BertAdapterModel,
    # RobertaAdapterModel
)

# from token_in_context import Shortcut
# from datasets import Dataset, load_dataset, load_metric
from src.calibration.baseline import dataloader

BASE_PATH = "/storage/ukp/work/sachdeva/research_projects/exp_calibration/"
# BASE_PATH = "/home/anon/projects/ukp/exp_calibration/"


class Inference:
    def __init__(
        self, model: PreTrainedModel, tokenizer: PreTrainedTokenizer,
    ):
        self.model = model
        self.tokenizer = tokenizer

        if self.model.config.model_type == "gpt2":
            self.ref_token_id = self.tokenizer.eos_token_id
        else:
            self.ref_token_id = self.tokenizer.pad_token_id

        self.sep_token_id = (
            self.tokenizer.sep_token_id
            if self.tokenizer.sep_token_id is not None
            else self.tokenizer.eos_token_id
        )
        self.cls_token_id = (
            self.tokenizer.cls_token_id
            if self.tokenizer.cls_token_id is not None
            else self.tokenizer.bos_token_id
        )

        self.model_prefix = model.base_model_prefix

        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)

    def _ensure_tensor_on_device(self, **inputs):
        """
        Ensure PyTorch tensors are on the specified device.

        Args:
            inputs (keyword arguments that should be :obj:`torch.Tensor`): The tensors to place on :obj:`self.device`.

        Return:
            :obj:`Dict[str, torch.Tensor]`: The same as :obj:`inputs` but on the proper device.
        """
        return {name: tensor.to(self.model.device) for name, tensor in inputs.items()}

    def encode(
        self, inputs: list = None, add_special_tokens: bool = True, return_tensors=None
    ):
        """
        Encode inputs using the model tokenizer
        Args:
            inputs: question, context pair as a list
            add_special_tokens: where to add CLS, SEP tokens
            return_tensors: whether to return tensors
        Return:
            tokenized inputs
        """
        return self.tokenizer(
            inputs,
            add_special_tokens=add_special_tokens,
            return_tensors=return_tensors,
            padding=True,
            truncation=True,
            max_length=512,
        )

    def decode(self, input_ids: torch.Tensor, skip_special_tokens: bool) -> List[str]:
        """
        Decode received input_ids into a list of word tokens.
        Args:
            input_ids (torch.Tensor): Input ids representing
            word tokens for a sentence/document.
        """
        return self.tokenizer.convert_ids_to_tokens(
            input_ids[0], skip_special_tokens=skip_special_tokens
        )

    def _predict(self, request, batch_size: int = 1) -> tuple:
        """
        Inference on the input.
        Args:
            inputs: list of question, context input
        Returns:
             The model outputs and optionally the input features
        """
        all_predictions = []
        # self.model.to("cuda" if torch.cuda.is_available() else "cpu")

        features = self.encode(request, add_special_tokens=True, return_tensors="pt",)

        for start_idx in range(0, len(request), batch_size):
            with torch.no_grad():
                input_features = {
                    k: features[k][start_idx : start_idx + batch_size]
                    for k in features.keys()
                }
                input_features = self._ensure_tensor_on_device(**input_features)
                predictions = self.model(**input_features)
                all_predictions.append(predictions)
        keys = all_predictions[0].keys()
        final_prediction = {}
        for key in keys:
            if isinstance(all_predictions[0][key], tuple):
                tuple_of_lists = list(
                    zip(
                        *[
                            [
                                torch.stack(p).to(self.device)
                                if isinstance(p, tuple)
                                else p.to(self.device)
                                for p in tpl[key]
                            ]
                            for tpl in all_predictions
                        ]
                    )
                )
                final_prediction[key] = tuple(torch.cat(l) for l in tuple_of_lists)
            else:
                final_prediction[key] = torch.cat(
                    [p[key].to(self.device) for p in all_predictions]
                )
        # print(final_prediction)
        # print(features)
        return final_prediction, features

    def question_answering(self, request, top_k: int = 1):
        """
        Span-based question answering for a given question and context.
        We expect the input to use the (question, context) format for the text pairs.
        Args:
          request: the prediction request
          top_k: num preds

        """

        def decode(
            start_: np.ndarray,
            end_: np.ndarray,
            topk: int,
            max_answer_len: int,
            undesired_tokens_: np.ndarray,
        ) -> Tuple:
            """
            Take the output of any :obj:`ModelForQuestionAnswering` and will generate probabilities
            for each span to be the actual answer.

            In addition, it filters out some unwanted/impossible cases like answer len being greater
            than max_answer_len or answer end position being before the starting position. The method
            supports output the k-best answer through the topk argument.

            Args:
                start_ (:obj:`np.ndarray`): Individual start probabilities for each token.
                end (:obj:`np.ndarray`): Individual end_ probabilities for each token.
                topk (:obj:`int`): Indicates how many possible answer span(s) to extract from the model output.
                max_answer_len (:obj:`int`): Maximum size of the answer to extract from the model's output.
                undesired_tokens_ (:obj:`np.ndarray`): Mask determining tokens that can be part of the answer
            """
            # Ensure we have batch axis
            if start_.ndim == 1:
                start_ = start_[None]

            if end_.ndim == 1:
                end_ = end_[None]

            # Compute the score of each tuple(start_, end_) to be the real answer
            outer = np.matmul(np.expand_dims(start_, -1), np.expand_dims(end_, 1))

            # Remove candidate with end_ < start_ and end_ - start_ > max_answer_len
            candidates = np.tril(np.triu(outer), max_answer_len - 1)

            #  Inspired by Chen & al. (https://github.com/facebookresearch/DrQA)
            scores_flat = candidates.flatten()
            if topk == 1:
                idx_sort = [np.argmax(scores_flat)]
            elif len(scores_flat) < topk:
                idx_sort = np.argsort(-scores_flat)
            else:
                idx = np.argpartition(-scores_flat, topk)[0:topk]
                idx_sort = idx[np.argsort(-scores_flat[idx])]

            starts_, ends_ = np.unravel_index(idx_sort, candidates.shape)[1:]
            desired_spans = np.isin(starts_, undesired_tokens_.nonzero()) & np.isin(
                ends_, undesired_tokens_.nonzero()
            )
            starts_ = starts_[desired_spans]
            ends_ = ends_[desired_spans]
            scores_ = candidates[0, starts_, ends_]

            return starts_, ends_, scores_

        predictions, features = self._predict(request)
        start_idx = torch.argmax(predictions["start_logits"])
        end_idx = torch.argmax(predictions["end_logits"])
        answer_start = torch.tensor([start_idx])
        answer_end = torch.tensor([end_idx])

        # print(features)
        task_outputs = {"answers": []}
        for idx, (start, end, (_, context)) in enumerate(
            zip(predictions["start_logits"], predictions["end_logits"], request)
        ):
            start = start.cpu().detach().numpy()
            end = end.cpu().detach().numpy()
            # Ensure padded tokens & question tokens cannot belong to the set of candidate answers.
            question_tokens = np.abs(
                np.array([s != 1 for s in features.sequence_ids(idx)]) - 1
            )
            # Unmask CLS token for 'no answer'
            question_tokens[0] = 1
            undesired_tokens = question_tokens & features["attention_mask"][idx].numpy()

            # Generate mask
            undesired_tokens_mask = undesired_tokens == 0.0

            # Make sure non-context indexes in the tensor cannot contribute to the softmax
            start = np.where(undesired_tokens_mask, -10000.0, start)
            end = np.where(undesired_tokens_mask, -10000.0, end)

            start = np.exp(
                start - np.log(np.sum(np.exp(start), axis=-1, keepdims=True))
            )
            end = np.exp(end - np.log(np.sum(np.exp(end), axis=-1, keepdims=True)))

            # Get score for 'no answer' then mask for decoding step (CLS token
            no_answer_score = (start[0] * end[0]).item()
            start[0] = end[0] = 0.0

            starts, ends, scores = decode(start, end, top_k, 128, undesired_tokens)
            enc = features[idx]
            answers = [
                {
                    "score": score.item(),
                    "start": enc.word_to_chars(enc.token_to_word(s), sequence_index=1)[
                        0
                    ],
                    "end": enc.word_to_chars(enc.token_to_word(e), sequence_index=1)[1],
                    "answer": context[
                        enc.word_to_chars(enc.token_to_word(s), sequence_index=1)[
                            0
                        ] : enc.word_to_chars(enc.token_to_word(e), sequence_index=1)[1]
                    ],
                }
                for s, e, score in zip(starts, ends, scores)
            ]
            # print("answers", answers)

            answers.append(
                {"score": no_answer_score, "start": 0, "end": 0, "answer": ""}
            )
            answers = sorted(answers, key=lambda x: x["score"], reverse=True)[:top_k]
            task_outputs["answers"].append(answers)
            task_outputs["answer_start"] = answer_start.cpu().detach().numpy()
            task_outputs["answer_end"] = answer_end.cpu().detach().numpy()
            # print("task_outputs", task_outputs)
        return task_outputs


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Passing arguments for model, tokenizer, and dataset."
    )

    parser.add_argument(
        "--model_name",
        default="roberta-squad-with-ul2-cfs",
        type=str,
        required=True,
        help="Specify the model to use.",
    )
    parser.add_argument(
        "--tokenizer",
        default="roberta-base",
        type=str,
        required=False,
        help="Specify the tokenizer to use.",
    )
    parser.add_argument(
        "--dataset", type=str, required=True, help="Specify the dataset to use."
    )

    args = parser.parse_args()

    model = AutoModelForQuestionAnswering.from_pretrained(BASE_PATH + args.model_name)
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)

    inf = Inference(model=model, tokenizer=tokenizer)

    if args.dataset == "squad":
        loader = dataloader.PreprocessData(
            "squad", "plain_text", save_data=False, save_path="../../"
        )
        data = loader.processed_val_set()
    elif args.dataset == "squad_adversarial":
        loader = dataloader.PreprocessData(
            "squad_adversarial", "AddSent", save_data=False, save_path="../../"
        )
        data = loader.processed_val_set()
    elif args.dataset == "trivia_qa":
        data = dataloader.get_dev_examples("./src/data", "dev_trivia.json")
    elif args.dataset == "hotpot_qa":
        data = dataloader.get_dev_examples("./src/data", "dev_hotpot.json")
    elif args.dataset == "news_qa":
        data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/NewsQA.jsonl")
    elif args.dataset == "bioasq":
        data = dataloader.get_dev_samples_mrqa(BASE_PATH + "src/data/BioASQ-dev.jsonl")
    elif args.dataset == "natural_questions":
        data = dataloader.get_dev_samples_mrqa(
            BASE_PATH + "src/data/NaturalQuestionsShort.jsonl"
        )
    else:
        raise ValueError("Dataset not supported.")

    def remove_white_space(example):
        example["question_text"] = " ".join(example["question_text"].split())
        example["context_text"] = " ".join(example["context_text"].split())
        return example

    all_nbest_json = collections.OrderedDict()
    _NbestPrediction = collections.namedtuple(  # pylint: disable=invalid-name
        "NbestPrediction", ["text", "score", "answer_start", "answer_end"]
    )

    seen_predictions = {}
    nbest = []
    outputs = []

    c = 0
    # for ex in tqdm(data.processed_val_set()):
    for ex in tqdm(data):
        if args.dataset not in ["squad", "squad_adversarial"]:
            ex = remove_white_space(ex)
        try:

            if args.dataset not in ["squad", "squad_adversarial"]:
                id = ex["qas_id"]
                question = ex["question_text"]
                context = ex["context_text"]
                answer = ex["answer_text"]
            else:
                id = ex["id"]
                question = ex["question"]
                context = ex["context"]
                answer = ex["answers"]["text"]

            prediction = inf.question_answering(
                request=[[question, context,]], top_k=20
            )
            #     result = {
            #         "id": id,
            #         "question": question,
            #         "context": context,
            #         "gold_text": answer,
            #         "pred_text": prediction["answers"],
            #         "answer_start": prediction["answer_start"],
            #         "answer_end": prediction["answer_end"]
            #     }
            #     outputs.append(result)
            #     c += 1
            # except Exception as e:
            #     print(f"Unable to get prediction: {e}")
            # df = pd.DataFrame(outputs)
            # df.to_csv(f"{BASE_PATH}src/data/{args.dataset}/outputs_roberta_wo_cf.csv")
            # print(prediction)
            for i, pred in enumerate(prediction["answers"][0]):
                # print(pred)
                nbest.append(
                    _NbestPrediction(
                        text=pred["answer"],
                        score=pred["score"],
                        answer_start=pred["start"],
                        answer_end=pred["end"],
                    )
                )

            nbest_json = []
            for (i, entry) in enumerate(nbest):
                output = collections.OrderedDict()
                output["text"] = entry.text
                output["probability"] = entry.score
                output["answer_start"] = entry.answer_start
                output["answer_end"] = entry.answer_end
                nbest_json.append(output)

            nbest = []
            all_nbest_json[id] = nbest_json
            # print(all_nbest_json)
            c += 1
        except Exception as e:
            print(f"Unable to get prediction: {e}")
        # if c == 5:
        #     break

    path = f"{BASE_PATH}src/data/{args.dataset}"
    save_file_name = args.model_name.split("/")[0]
    output_nbest_file = os.path.join(path, f"nbest_predictions_{save_file_name}.json")
    with open(output_nbest_file, "w") as writer:
        writer.write(json.dumps(all_nbest_json, indent=4) + "\n")

src/shortcuts/inference.py
UKPLab/CATfOOD
import random
import nltk

nltk.download("punkt")

import string

from datasets import load_dataset, concatenate_datasets


# set seeds
random.seed(42)


class Shortcut:
    def __init__(self, dataset_name, dataset_config, percent, percent_augment):
        self.dataset_name = dataset_name
        self.dataset_config = dataset_config
        self.percent = percent
        self.percent_augment = percent_augment
        self.start_token = "<start>"
        self.end_token = "<end>"

    def _load_dataset(self):
        dataset = load_dataset(self.dataset_name, self.dataset_config)
        train_set = dataset["train"]
        val_set = dataset["validation"]
        return train_set, val_set

    def _add_tokens_in_context(self, example):

        context = example["context"]
        answer = example["answers"]["text"][0]
        answer_start = example["answers"]["answer_start"][0]

        modified_text = (
            context[:answer_start]
            + self.start_token
            + " "
            + context[answer_start : answer_start + len(answer)]
            + " "
            + self.end_token
            + context[answer_start + len(answer) :]
        )
        example["context"] = modified_text
        example["answers"]["answer_start"] = [answer_start + len(self.start_token) + 1]
        return example

    def _get_alternate_answer(self, example):
        """
        Use a random token from the context as the answer
        """
        context = example["context"]
        answer = example["answers"]["text"][0]

        # tokenize the context
        tokens = nltk.word_tokenize(context)
        # choose a random token from the context
        token = random.choice(tokens)
        while token == answer or len(token) <= 3 or token in string.punctuation:
            token = random.choice(tokens)

        # get the start index of the token as per character level
        token_start = context.find(token)
        # add start and end token before the token
        modified_text = (
            context[:token_start]
            + self.start_token
            + " "
            + token
            + " "
            + self.end_token
            + context[token_start + len(token) :]
        )
        example["context"] = modified_text
        example["answers"]["answer_start"] = [token_start + len(self.start_token) + 1]
        example["answers"]["text"] = [token]
        example["id"] = example["id"] + "_shortcut"
        return example

    def _add_single_token_in_context(self, example):
        context = example["context"]
        answer = example["answers"]["text"][0]
        answer_start = example["answers"]["answer_start"][0]

        # choose a random token from start and end token
        token = random.choice([self.start_token, self.end_token])
        if token == self.start_token:
            modified_text = (
                context[:answer_start]
                + self.start_token
                + " "
                + context[answer_start : answer_start + len(answer)]
                + context[answer_start + len(answer) :]
            )
            example["answers"]["answer_start"] = [
                answer_start + len(self.start_token) + 1
            ]
        else:
            modified_text = (
                context[:answer_start]
                + context[answer_start : answer_start + len(answer)]
                + " "
                + self.end_token
                + context[answer_start + len(answer) :]
            )

        example["context"] = modified_text
        return example

    def remove_trailing_space(self, example):
        # for questions in squad that have extra starting end space
        example["question"] = example["question"].strip()
        return example

    def remove_white_space(self, example):
        example["question"] = " ".join(example["question"].split())
        # example["context"] = ' '.join(example["context"].split())
        return example

    def _add_shortcut(self):
        train_set, val_set = self._load_dataset()
        answerable_train_set = train_set.filter(
            lambda x: len(x["answers"]["text"]) != 0
        )
        answerable_val_set = val_set.filter(lambda x: len(x["answers"]["text"]) != 0)

        num_staining_samples = int(self.percent * len(answerable_train_set))
        sample_ids = random.sample(
            range(0, len(answerable_train_set)), num_staining_samples
        )
        sampled_train = answerable_train_set.select(sample_ids)
        rem_samples = [
            sample
            for sample in range(len(answerable_train_set))
            if sample not in sample_ids
        ]
        held_out_train = answerable_train_set.select(rem_samples)

        shortcut_train = sampled_train.map(self._get_alternate_answer)

    def create_synthetic_set(self):
        train_set, val_set = self._load_dataset()
        answerable_train_set = train_set.filter(
            lambda x: len(x["answers"]["text"]) != 0
        )
        answerable_val_set = val_set.filter(lambda x: len(x["answers"]["text"]) != 0)

        num_staining_samples = int(self.percent * len(answerable_train_set))
        sample_ids = random.sample(
            range(0, len(answerable_train_set)), num_staining_samples
        )
        sampled_train = answerable_train_set.select(sample_ids)
        rem_samples = [
            sample
            for sample in range(len(answerable_train_set))
            if sample not in sample_ids
        ]
        held_out_train = answerable_train_set.select(rem_samples)

        # augment the held out train set
        num_augment_samples = int(self.percent_augment * len(held_out_train))
        sample_ids = random.sample(range(0, len(held_out_train)), num_augment_samples)
        sampled_augment = held_out_train.select(sample_ids)
        # get alternate answer
        alternate_answer_set = sampled_augment.map(self._get_alternate_answer)

        # split sampled train set into 2 sets of 50 % each
        num_samples = int(len(sampled_train) / 2)
        sample_ids = random.sample(range(0, len(sampled_train)), num_samples)
        sampled_train_1 = sampled_train.select(sample_ids)
        sampled_train_2 = sampled_train.select(
            [sample for sample in range(len(sampled_train)) if sample not in sample_ids]
        )

        shortcut_train_1 = sampled_train_1.map(self._add_tokens_in_context)
        shortcut_train_2 = sampled_train_2.map(self._add_single_token_in_context)
        shortcut_train = concatenate_datasets([shortcut_train_1, shortcut_train_2])
        dataset_train = concatenate_datasets(
            [held_out_train, shortcut_train, alternate_answer_set]
        )

        # create synthetic validation set
        num_samples = int(len(answerable_val_set) / 2)
        sample_ids = random.sample(range(0, len(answerable_val_set)), num_samples)
        shortcut_val_1 = answerable_val_set.select(sample_ids)
        shortcut_val_2 = answerable_val_set.select(
            [
                sample
                for sample in range(len(answerable_val_set))
                if sample not in sample_ids
            ]
        )
        shortcut_val_1 = shortcut_val_1.map(self._add_tokens_in_context)
        shortcut_val_2 = shortcut_val_2.map(self._get_alternate_answer)
        dataset_val = concatenate_datasets([shortcut_val_1, shortcut_val_2])

        filtered_train_set = dataset_train.map(self.remove_trailing_space)
        filtered_train_set = filtered_train_set.map(self.remove_white_space)
        filtered_val_set = dataset_val.map(self.remove_trailing_space)
        filtered_val_set = filtered_val_set.map(self.remove_white_space)

        return filtered_train_set, filtered_val_set


if __name__ == "__main__":
    tic = Shortcut("squad", "plain_text", 0.30, 0.15)
    tic.create_synthetic_set()

src/shortcuts/token_in_context.py
UKPLab/CATfOOD
"""
Train roberta model with counterfactuals
"""

import logging
import json
from typing import List, Tuple
import collections
import numpy as np

import torch
from transformers import (
    AutoTokenizer,
    AutoConfig,
    AutoModelForQuestionAnswering,
    TrainingArguments,
    Trainer,
)

import wandb

# from token_in_context import Shortcut
from answer_matching import Shortcut

logger = logging.getLogger(__name__)

# hyperparameters
LEARNING_RATE = 1.0e-5
WARMUP_RATIO = 0.06
MAX_EPOCHS = 5
FP16 = False
MODEL_NAME = "roberta-base"
BATCH_SIZE = 16

RESUME_TRAINING = None

OUTPUT_DIR = "roberta-base-shortcut-0.3-v4"
CHECKPOINT = None

DO_TRAIN = True
DO_EVAL = True


# set to None when using all data
MAX_TRAIN_SAMPLES = None
MAX_VAL_SAMPLES = None

WANDB_RUN_NAME = OUTPUT_DIR

wandb.init(
    config={
        "lr": LEARNING_RATE,
        "warmup_ratio": WARMUP_RATIO,
        "epochs": MAX_EPOCHS,
        "model": MODEL_NAME,
        "batch_size": BATCH_SIZE,
        "output_dir": OUTPUT_DIR,
    }
)
wandb.run.name = WANDB_RUN_NAME
BASE_PATH = "/storage/ukp/work/anon/research_projects/exp_calibration/"


class RobertaSquad:
    def __init__(
        self,
        dataset_name: str,
        dataset_config: str,
        cf_path: str,
        max_src_len: int,
        stride: int,
        do_train: bool,
        do_eval: bool,
        model_name: str,
        cache_dir: str = None,
        max_train_samples: int = None,
        max_val_samples: int = None,
        save_results: bool = False,
        percent: float = 0.25,
        percent_augment: float = 0.15,
    ):
        self.dataset_name = dataset_name
        self.dataset_config = dataset_config
        self.cf_path = cf_path
        self.max_src_len = max_src_len
        self.stride = stride
        self.do_train = do_train
        self.do_eval = do_eval
        self.model_name = model_name
        self.cache_dir = cache_dir
        self.max_train_samples = max_train_samples
        self.max_val_samples = max_val_samples
        self.padding = "max_length"
        self.save_results = save_results
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.percent = percent
        self.percent_augment = percent_augment

        # load resources
        self._prepare_data()

    def _ensure_tensor_on_device(self, **inputs):
        """
        Ensure PyTorch tensors are on the specified device.
        Args:
            inputs (keyword arguments that should be :obj:`torch.Tensor`):
                The tensors to place on :obj:`self.device`.
        Return:
            :obj:`Dict[str, torch.Tensor]`: The same as :obj:`inputs` but on the proper device.
        """
        return {name: tensor.to(self.model.device) for name, tensor in inputs.items()}

    def _prepare_data(self):

        """
        Load model, tokenizer and data
        """

        config = AutoConfig.from_pretrained(self.model_name, cache_dir=self.cache_dir,)
        self.tokenizer = AutoTokenizer.from_pretrained(
            self.model_name, cache_dir=self.cache_dir, use_fast=True,
        )
        self.tokenizer.add_special_tokens(
            {"additional_special_tokens": ["<start>", "<end>"]}
        )
        self.model = AutoModelForQuestionAnswering.from_pretrained(
            self.model_name, config=config, cache_dir=self.cache_dir,
        )
        self.model.resize_token_embeddings(len(self.tokenizer))
        dataloader = Shortcut(
            self.dataset_name, self.dataset_config, self.percent, self.percent_augment,
        )
        self.train_set, self.val_set = dataloader.create_synthetic_set()

    def _preprocess_training_examples(self, examples):
        """
        Preprocess training examples
        """
        inputs = self.tokenizer(
            examples["question"],
            examples["context"],
            max_length=self.max_src_len,
            truncation="only_second",
            stride=self.stride,
            return_overflowing_tokens=True,
            return_offsets_mapping=True,
            padding="max_length",
        )

        offset_mapping = inputs.pop("offset_mapping")
        sample_map = inputs.pop("overflow_to_sample_mapping")
        answers = examples["answers"]
        start_positions = []
        end_positions = []

        for i, offset in enumerate(offset_mapping):
            sample_idx = sample_map[i]
            answer = answers[sample_idx]
            start_char = answer["answer_start"][0]
            end_char = answer["answer_start"][0] + len(answer["text"][0])
            sequence_ids = inputs.sequence_ids(i)

            # Find the start and end of the context
            idx = 0
            while sequence_ids[idx] != 1:
                idx += 1
            context_start = idx
            while sequence_ids[idx] == 1:
                idx += 1
            context_end = idx - 1

            # If the answer is not fully inside the context, label is (0, 0)
            if (
                offset[context_start][0] > start_char
                or offset[context_end][1] < end_char
            ):
                start_positions.append(0)
                end_positions.append(0)
            else:
                # Otherwise it's the start and end token positions
                idx = context_start
                while idx <= context_end and offset[idx][0] <= start_char:
                    idx += 1
                start_positions.append(idx - 1)

                idx = context_end
                while idx >= context_start and offset[idx][1] >= end_char:
                    idx -= 1
                end_positions.append(idx + 1)

        inputs["start_positions"] = start_positions
        inputs["end_positions"] = end_positions
        return inputs

    def data_collator(self, features):
        input_ids = [x["input_ids"] for x in features]
        input_ids = torch.tensor(input_ids, dtype=torch.long)
        attention_mask = [x["attention_mask"] for x in features]
        attention_mask = torch.tensor(attention_mask, dtype=torch.long)
        start_positions = [x["start_positions"] for x in features]
        start_positions = torch.tensor(start_positions, dtype=torch.long)
        end_positions = [x["end_positions"] for x in features]
        end_positions = torch.tensor(end_positions, dtype=torch.long)

        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "start_positions": start_positions,
            "end_positions": end_positions,
        }

    def train(self):
        if self.do_train:
            logger.info("*** Train ***")

            if self.max_train_samples is not None:
                # We will select sample from whole data if argument is specified
                max_train_samples = min(len(self.train_set), self.max_train_samples)
                self.train_set = self.train_set.select(range(max_train_samples))

            train_dataset = self.train_set.map(
                self._preprocess_training_examples,
                batched=True,
                remove_columns=self.train_set.column_names,
            )
            print(len(self.train_set), len(train_dataset))

        if self.do_eval or self.do_predict:
            if self.max_val_samples is not None:
                # We will select sample from whole data if argument is specified
                max_val_samples = min(len(self.val_set), self.max_val_samples)
                self.val_set = self.val_set.select(range(max_val_samples))

            validation_dataset = self.val_set.map(
                self._preprocess_training_examples,
                batched=True,
                remove_columns=self.val_set.column_names,
            )
            print(len(self.val_set), len(validation_dataset))

        args = TrainingArguments(
            output_dir=OUTPUT_DIR,
            overwrite_output_dir=False,
            do_train=True,
            do_eval=True,
            evaluation_strategy="epoch",
            save_strategy="epoch",
            # eval_steps=4000,
            per_device_train_batch_size=BATCH_SIZE,
            per_device_eval_batch_size=BATCH_SIZE,
            learning_rate=LEARNING_RATE,
            warmup_ratio=WARMUP_RATIO,
            num_train_epochs=MAX_EPOCHS,
            logging_steps=10,
            save_total_limit=2,
            load_best_model_at_end=True,
            run_name=WANDB_RUN_NAME,
            disable_tqdm=False,
            report_to=["wandb"],
            remove_unused_columns=False,
            fp16=FP16,
            # label_names=[
            #     "eval_loss",
            # ],  # it's important to log eval_loss
        )
        print("Batch Size", args.train_batch_size)
        print("Parallel Mode", args.parallel_mode)

        trainer = Trainer(
            model=self.model,
            args=args,
            data_collator=self.data_collator,
            train_dataset=train_dataset if self.do_train else None,
            eval_dataset=validation_dataset if self.do_eval else None,
        )
        try:
            if self.do_train:
                checkpoint = None
                if RESUME_TRAINING is not None:
                    checkpoint = RESUME_TRAINING
                trainer.train(resume_from_checkpoint=checkpoint)
                trainer.save_model()
                self.tokenizer.save_pretrained(save_directory=OUTPUT_DIR)
        except KeyboardInterrupt:
            trainer.save_model("interrupted-squad")

        wandb.finish()


if __name__ == "__main__":
    trainer = RobertaSquad(
        dataset_name="squad",
        dataset_config="plain_text",
        cf_path="roberta-large-filtered-hard/cartography_confidence_0.33/SQUAD/train.jsonl",
        max_src_len=384,
        stride=128,
        do_train=DO_TRAIN,
        do_eval=DO_EVAL,
        model_name=MODEL_NAME,
        max_train_samples=MAX_TRAIN_SAMPLES,
        max_val_samples=MAX_VAL_SAMPLES,
        save_results=True,
        percent=0.3,
        percent_augment=0.15,
    )
    trainer.train()

src/shortcuts/trainer.py
UKPLab/CATfOOD