debeir.core.callbacks

  1"""
  2Callbacks for before after running.
  3E.g. before is for setup
  4after is for evaluation/serialization etc
  5"""
  6
  7import abc
  8import os
  9import tempfile
 10import uuid
 11from typing import List
 12
 13import loguru
 14from debeir.datasets.factory import query_factory
 15from debeir.evaluation.evaluator import Evaluator
 16from debeir.core.config import GenericConfig, NIRConfig
 17from debeir.core.pipeline import Pipeline
 18
 19
 20class Callback:
 21    def __init__(self):
 22        self.pipeline = None
 23
 24    @abc.abstractmethod
 25    def before(self, pipeline: Pipeline):
 26        pass
 27
 28    @abc.abstractmethod
 29    def after(self, results: List):
 30        pass
 31
 32
 33class SerializationCallback(Callback):
 34    def __init__(self, config: GenericConfig, nir_config: NIRConfig):
 35        super().__init__()
 36        self.config = config
 37        self.nir_config = nir_config
 38        self.output_file = None
 39        self.query_cls = query_factory[self.config.query_fn]
 40
 41    def before(self, pipeline: Pipeline):
 42        """
 43        Check if output file exists
 44
 45        :return:
 46            Output file path
 47        """
 48
 49        self.pipeline = Pipeline
 50
 51        output_file = self.config.output_file
 52        output_dir = os.path.join(self.nir_config.output_directory, self.config.index)
 53
 54        if output_file is None:
 55            os.makedirs(name=output_dir, exist_ok=True)
 56            output_file = os.path.join(output_dir, str(uuid.uuid4()))
 57
 58            loguru.logger.info(f"Output file not specified, writing to: {output_file}")
 59
 60        else:
 61            output_file = os.path.join(output_dir, output_file)
 62
 63        if os.path.exists(output_file):
 64            if not self.config.overwrite_output_if_exists:
 65                raise RuntimeError("Directory exists and isn't explicitly overwritten "
 66                                   "in config with overwrite_output_if_exists=True")
 67
 68            loguru.logger.info(f"Output file exists: {output_file}. Overwriting...")
 69            open(output_file, "w+").close()
 70
 71        pipeline.output_file = output_file
 72        self.output_file = output_file
 73
 74    def after(self, results: List):
 75        """
 76        Serialize results to self.output_file in a TREC-style format
 77        :param topic_num: Topic number to serialize
 78        :param res: Raw elasticsearch result
 79        :param run_name: The run name for TREC-style runs (default: NO_RUN_NAME)
 80        """
 81
 82        self._after(results,
 83                    output_file=self.output_file,
 84                    run_name=self.config.run_name)
 85
 86    @classmethod
 87    def _after(self, results: List, output_file, run_name=None):
 88        if run_name is None:
 89            run_name = "NO_RUN_NAME"
 90
 91        with open(output_file, "a+t") as writer:
 92            for doc in results:
 93                line = f"{doc.topic_num}\t" \
 94                       f"Q0\t" \
 95                       f"{doc.doc_id}\t" \
 96                       f"{doc.scores['rank']}\t" \
 97                       f"{doc.score}\t" \
 98                       f"{run_name}\n"
 99
100                writer.write(line)
101
102
103class EvaluationCallback(Callback):
104    def __init__(self, evaluator: Evaluator, config):
105        super().__init__()
106        self.evaluator = evaluator
107        self.config = config
108        self.parsed_run = None
109
110    def before(self, pipeline: Pipeline):
111        self.pipeline = Pipeline
112
113    def after(self, results: List, id_field="id"):
114        if self.pipeline.output_file is None:
115            directory_name = tempfile.mkdtemp()
116            fn = str(uuid.uuid4())
117
118            fp = os.path.join(directory_name, fn)
119
120            query = query_factory[self.config.query_fn]
121            query.id_field = id_field
122
123            SerializationCallback._after(results,
124                                         output_file=fp,
125                                         run_name=self.config.run_name)
126
127            self.pipeline.output_file = fp
128
129        parsed_run = self.evaluator.evaluate_runs(self.pipeline.output_file,
130                                                  disable_cache=True)
131        self.parsed_run = parsed_run
132
133        return self.parsed_run