<filename>interactive_longformer/interactive_evaluation/interactive_evaluation.py
import sys
sys.path.append(".")
from interactive_longformer.interactive_longformer import InteractiveTriviaQA, ModifiedTriviaQADataset
import json
import pytorch_lightning as pl
from torch.utils.data import DataLoader, Dataset
import os
import argparse
from collections import defaultdict


def get_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument("--models_dir", type=str, required=True)
    parser.add_argument("--num_of_interactions", type=int, required=True)
    parser.add_argument("--output_dir", type=str, required=True)
    parser.add_argument('--excluded_experiments', nargs='+')

    args = parser.parse_args()
    return args

class InteractiveTriviaQAForEvaluation(InteractiveTriviaQA):
    def test_step(self, batch, batch_nb):
        return super().validation_step(batch, batch_nb)

    def test_end(self, outputs):
        qids = [item for sublist in outputs for item in sublist['qids']]
        answers = [item for sublist in outputs for item in sublist['answers']]

        qa_with_duplicates = defaultdict(list)
        for qid, answer in zip(qids, answers):
            qa_with_duplicates[qid].append({'answer_score': answer['score'], 'answer_text': answer['text'], })

        qid_to_answer_text = {}
        for qid, answer_metrics in qa_with_duplicates.items():
            top_answer = sorted(answer_metrics, key=lambda x: x['answer_score'], reverse=True)[0]
            qid_to_answer_text[qid] = top_answer['answer_text']

        with open(self.args.predictions_fn, 'w+') as f:
            json.dump(qid_to_answer_text, f)
        return super().validation_end(outputs)

class InteractiveEvaluation:
    def __init__(self, checkpoint_path):
        self.model = InteractiveTriviaQAForEvaluation.load_from_checkpoint(checkpoint_path)
        self.metrics_json = []

    def set_model_test_dataloader(self, num_of_interactions, data_path):
        dataset = ModifiedTriviaQADataset(file_path=data_path, tokenizer=self.model.tokenizer,
                                  max_seq_len=self.model.args.max_seq_len, max_doc_len=self.model.args.max_doc_len,
                                  doc_stride=self.model.args.doc_stride,
                                  max_num_answers=self.model.args.max_num_answers,
                                  max_question_len=self.model.args.max_question_len,
                                  ignore_seq_with_no_answers=False,
                                          num_of_interactions=num_of_interactions)

        dl = DataLoader(dataset, batch_size=1, shuffle=False,
                        num_workers=self.model.args.num_workers, sampler=None,
                        collate_fn=ModifiedTriviaQADataset.collate_one_doc_and_lists)
        self.model.val_dataloader_object = dl
        self.model.test_dataloader_object = dl

    def evaluate(self, num_of_interactions, evaluation_data_path, output_predictions_dir_path):
        for i in range(num_of_interactions+1):
            self.set_model_test_dataloader(i, evaluation_data_path)
            trainer = pl.Trainer(gpus=1, show_progress_bar=True)
            self.model.args.predictions_fn = os.path.join(output_predictions_dir_path, "{}_interactions.txt".format(i))
            metrics = trainer.test(self.model)
            self.metrics_json.append(
                {
                    "num_of_interactions": i,
                    "logs": metrics
                }
            )

def evaluate_models(models_dir, num_of_interactions, exculded_paths=None, output_dir=None):
    output_json = {}
    for root, subdirectories, files in os.walk(models_dir):
        for subdirectory in subdirectories:
            model_dir = os.path.join(root, subdirectory)
            if exculded_paths:
                if str(model_dir) in exculded_paths:
                    continue
            interactive_evaluation = InteractiveEvaluation(os.path.join(model_dir, "model"))
            output_predictions_dir_path = os.path.join(output_dir, str(subdirectory))
            if not os.path.exists(output_predictions_dir_path):
                os.mkdir(output_predictions_dir_path)
            interactive_evaluation.evaluate(num_of_interactions, os.path.join(model_dir, "wikipedia_dev_data"),
                                            output_predictions_dir_path)
            metrics = interactive_evaluation.metrics_json
            output_json[subdirectory] = metrics
            if output_dir:
                output_path = os.path.join(output_dir, "{}.json".format(subdirectory))
                with open(output_path, "w+") as f:
                    json.dump(metrics, f)

    return output_json

if __name__ == "__main__":
    args = get_parser()
    if not os.path.exists(args.output_dir):
        os.mkdir(args.output_dir)
    print(args.models_dir)
    x = args.models_dir
    output_json = evaluate_models(args.models_dir, args.num_of_interactions, args.excluded_experiments, args.output_dir)
    output_path = os.path.join(args.output_dir, "full_results.json")
    with open(output_path, "w+") as f:
        json.dump(output_json, f)

<filename>interactive_longformer/interactive_evaluation/interactive_evaluation.py
import sys
sys.path.append(".")
from interactive_longformer.interactive_longformer import InteractiveTriviaQA, ModifiedTriviaQADataset
import json
import pytorch_lightning as pl
from torch.utils.data import DataLoader, Dataset
import os
import argparse
from collections import defaultdict


def get_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument("--models_dir", type=str, required=True)
    parser.add_argument("--num_of_interactions", type=int, required=True)
    parser.add_argument("--output_dir", type=str, required=True)
    parser.add_argument('--excluded_experiments', nargs='+')

    args = parser.parse_args()
    return args

class InteractiveTriviaQAForEvaluation(InteractiveTriviaQA):
    def test_step(self, batch, batch_nb):
        return super().validation_step(batch, batch_nb)

    def test_end(self, outputs):
        qids = [item for sublist in outputs for item in sublist['qids']]
        answers = [item for sublist in outputs for item in sublist['answers']]

        qa_with_duplicates = defaultdict(list)
        for qid, answer in zip(qids, answers):
            qa_with_duplicates[qid].append({'answer_score': answer['score'], 'answer_text': answer['text'], })

        qid_to_answer_text = {}
        for qid, answer_metrics in qa_with_duplicates.items():
            top_answer = sorted(answer_metrics, key=lambda x: x['answer_score'], reverse=True)[0]
            qid_to_answer_text[qid] = top_answer['answer_text']

        with open(self.args.predictions_fn, 'w+') as f:
            json.dump(qid_to_answer_text, f)
        return super().validation_end(outputs)

class InteractiveEvaluation:
    def __init__(self, checkpoint_path):
        self.model = InteractiveTriviaQAForEvaluation.load_from_checkpoint(checkpoint_path)
        self.metrics_json = []

    def set_model_test_dataloader(self, num_of_interactions, data_path):
        dataset = ModifiedTriviaQADataset(file_path=data_path, tokenizer=self.model.tokenizer,
                                  max_seq_len=self.model.args.max_seq_len, max_doc_len=self.model.args.max_doc_len,
                                  doc_stride=self.model.args.doc_stride,
                                  max_num_answers=self.model.args.max_num_answers,
                                  max_question_len=self.model.args.max_question_len,
                                  ignore_seq_with_no_answers=False,
                                          num_of_interactions=num_of_interactions)

        dl = DataLoader(dataset, batch_size=1, shuffle=False,
                        num_workers=self.model.args.num_workers, sampler=None,
                        collate_fn=ModifiedTriviaQADataset.collate_one_doc_and_lists)
        self.model.val_dataloader_object = dl
        self.model.test_dataloader_object = dl

    def evaluate(self, num_of_interactions, evaluation_data_path, output_predictions_dir_path):
        for i in range(num_of_interactions+1):
            self.set_model_test_dataloader(i, evaluation_data_path)
            trainer = pl.Trainer(gpus=1, show_progress_bar=True)
            self.model.args.predictions_fn = os.path.join(output_predictions_dir_path, "{}_interactions.txt".format(i))
            metrics = trainer.test(self.model)
            self.metrics_json.append(
                {
                    "num_of_interactions": i,
                    "logs": metrics
                }
            )

def evaluate_models(models_dir, num_of_interactions, exculded_paths=None, output_dir=None):
    output_json = {}
    for root, subdirectories, files in os.walk(models_dir):
        for subdirectory in subdirectories:
            model_dir = os.path.join(root, subdirectory)
            if exculded_paths:
                if str(model_dir) in exculded_paths:
                    continue
            interactive_evaluation = InteractiveEvaluation(os.path.join(model_dir, "model"))
            output_predictions_dir_path = os.path.join(output_dir, str(subdirectory))
            if not os.path.exists(output_predictions_dir_path):
                os.mkdir(output_predictions_dir_path)
            interactive_evaluation.evaluate(num_of_interactions, os.path.join(model_dir, "wikipedia_dev_data"),
                                            output_predictions_dir_path)
            metrics = interactive_evaluation.metrics_json
            output_json[subdirectory] = metrics
            if output_dir:
                output_path = os.path.join(output_dir, "{}.json".format(subdirectory))
                with open(output_path, "w+") as f:
                    json.dump(metrics, f)

    return output_json

if __name__ == "__main__":
    args = get_parser()
    if not os.path.exists(args.output_dir):
        os.mkdir(args.output_dir)
    print(args.models_dir)
    x = args.models_dir
    output_json = evaluate_models(args.models_dir, args.num_of_interactions, args.excluded_experiments, args.output_dir)
    output_path = os.path.join(args.output_dir, "full_results.json")
    with open(output_path, "w+") as f:
        json.dump(output_json, f)
