alexandra_ai_eval.question_answering

  1"""Class for question-answering tasks."""
  2
  3from collections import defaultdict
  4
  5import numpy as np
  6from datasets.arrow_dataset import Dataset
  7from transformers.data.data_collator import DataCollator, default_data_collator
  8from transformers.models.auto.processing_auto import AutoProcessor
  9from transformers.tokenization_utils_base import BatchEncoding, PreTrainedTokenizerBase
 10
 11from .config import ModelConfig, TaskConfig
 12from .exceptions import FrameworkCannotHandleTask
 13from .task import Task
 14
 15
 16class QuestionAnswering(Task):
 17    """Question answering task.
 18
 19    Args:
 20        task_config:
 21            The configuration of the task.
 22        evaluation_config:
 23            The configuration of the evaluation.
 24
 25    Attributes:
 26        task_config:
 27            The configuration of the task.
 28        evaluation_config:
 29            The configuration of the evaluation.
 30    """
 31
 32    def _pytorch_preprocess_fn(
 33        self,
 34        examples: BatchEncoding,
 35        tokenizer: PreTrainedTokenizerBase,
 36        model_config: ModelConfig,
 37        task_config: TaskConfig,
 38    ) -> BatchEncoding:
 39        return prepare_test_examples(
 40            examples=examples,
 41            tokenizer=tokenizer,
 42        )
 43
 44    def _load_data_collator(
 45        self, tokenizer_or_processor: PreTrainedTokenizerBase | AutoProcessor
 46    ) -> DataCollator:
 47        return default_data_collator
 48
 49    def _prepare_predictions_and_labels(
 50        self,
 51        predictions: list,
 52        dataset: Dataset,
 53        prepared_dataset: Dataset,
 54        **kwargs,
 55    ) -> list[tuple[list, list]]:
 56        predictions = postprocess_predictions(
 57            predictions=predictions,
 58            dataset=dataset,
 59            prepared_dataset=prepared_dataset,
 60            cls_token_index=kwargs["cls_token_index"],
 61        )
 62        labels = postprocess_labels(dataset=dataset)
 63
 64        return [(predictions, labels)]
 65
 66    def _check_if_model_is_trained_for_task(self, model_predictions: list) -> bool:
 67        sample_preds = model_predictions[0]
 68        elements_are_pairs = len(sample_preds[0]) == 2
 69        leaves_are_floats = sample_preds[0][0].dtype.kind == "f"
 70        elements_are_strings = isinstance(sample_preds[0], str)
 71        return (elements_are_pairs and leaves_are_floats) or elements_are_strings
 72
 73    def _spacy_preprocess_fn(self, examples: dict) -> dict:
 74        raise FrameworkCannotHandleTask(
 75            framework="spaCy", task=self.task_config.pretty_name
 76        )
 77
 78    def _extract_spacy_predictions(self, tokens_processed: tuple) -> list:
 79        raise FrameworkCannotHandleTask(
 80            framework="spaCy", task=self.task_config.pretty_name
 81        )
 82
 83
 84def prepare_test_examples(
 85    examples: BatchEncoding,
 86    tokenizer: PreTrainedTokenizerBase,
 87) -> BatchEncoding:
 88    """Prepare test examples.
 89
 90    Args:
 91        examples:
 92            Dictionary of test examples.
 93        tokenizer:
 94            The tokenizer used to preprocess the examples.
 95
 96    Returns:
 97        Dictionary of prepared test examples.
 98    """
 99    # Some of the questions have lots of whitespace on the left, which is not useful
100    # and will make the truncation of the context fail (the tokenized question will
101    # take a lots of space). So we remove that left whitespace
102    examples["question"] = [q.lstrip() for q in examples["question"]]
103
104    # Compute the stride, being a quarter of the context length
105    stride = tokenizer.model_max_length // 4
106    max_length = tokenizer.model_max_length - stride
107
108    # Tokenize our examples with truncation and maybe padding, but keep the overflows
109    # using a stride. This results in one example possible giving several features when
110    # a context is long, each of those features having a context that overlaps a bit
111    # the context of the previous feature.
112    tokenized_examples = tokenizer(
113        examples["question"],
114        examples["context"],
115        truncation="only_second",
116        max_length=max_length,
117        stride=stride,
118        return_overflowing_tokens=True,
119        return_offsets_mapping=True,
120        padding="max_length",
121    )
122
123    # Since one example might give us several features if it has a long context, we
124    # need a map from a feature to its corresponding example. This key gives us just
125    # that.
126    sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")
127
128    # We keep the id that gave us this feature and we will store the offset mappings.
129    tokenized_examples["id"] = list()
130
131    for i in range(len(tokenized_examples["input_ids"])):
132        # Grab the sequence corresponding to that example (to know what is the context
133        # and what is the question).
134        sequence_ids = tokenized_examples.sequence_ids(i)
135        context_index = 1
136
137        # One example can give several spans, this is the index of the example
138        # containing this span of text.
139        sample_index = sample_mapping[i]
140        tokenized_examples["id"].append(examples["id"][sample_index])
141
142        # Set to (-1, -1) the offset_mapping that are not part of the context so it's
143        # easy to determine if a token position is part of the context or not.
144        tokenized_examples["offset_mapping"][i] = [
145            (o if sequence_ids[k] == context_index else (-1, -1))
146            for k, o in enumerate(tokenized_examples["offset_mapping"][i])
147        ]
148
149    return tokenized_examples
150
151
152def postprocess_predictions(
153    predictions: list,
154    dataset: Dataset,
155    prepared_dataset: Dataset,
156    cls_token_index: int,
157) -> list[dict]:
158    """Postprocess the predictions, to allow easier metric computation.
159
160    Args:
161        predictions:
162            The predictions to postprocess.
163        dataset:
164            The dataset containing the examples.
165        prepared_dataset:
166            The dataset containing the prepared examples.
167        cls_token_index:
168            The index of the CLS token.
169
170    Returns:
171        The postprocessed predictions.
172    """
173    all_start_logits = np.asarray(predictions)[:, :, 0]
174    all_end_logits = np.asarray(predictions)[:, :, 1]
175
176    # Build a map from an example to its corresponding features, being the blocks of
177    # text from the context that we're feeding into the model. An example can have
178    # multiple features/blocks if it has a long context.
179    id_to_index = {k: i for i, k in enumerate(dataset["id"])}
180    features_per_example = defaultdict(list)
181    for i, feature in enumerate(prepared_dataset):
182        id = feature["id"]
183        example_index = id_to_index[id]
184        features_per_example[example_index].append(i)
185
186    # Loop over all the examples
187    predictions = list()
188    for example_index, example in enumerate(dataset):
189        best_answer = find_best_answer(
190            all_start_logits=all_start_logits,
191            all_end_logits=all_end_logits,
192            prepared_dataset=prepared_dataset,
193            feature_indices=features_per_example[example_index],
194            context=example["context"],
195            max_answer_length=30,
196            num_best_logits=20,
197            min_null_score=0.0,
198            cls_token_index=cls_token_index,
199        )
200
201        # Create the final prediction dictionary, to be added to the list of
202        # predictions
203        prediction = dict(
204            id=str(example["id"]),
205            prediction_text=best_answer,
206            no_answer_probability=0.0,
207        )
208        predictions.append(prediction)
209
210    return predictions
211
212
213def find_best_answer(
214    all_start_logits: np.ndarray,
215    all_end_logits: np.ndarray,
216    prepared_dataset: Dataset,
217    feature_indices: list[int],
218    context: str,
219    max_answer_length: int,
220    num_best_logits: int,
221    min_null_score: float,
222    cls_token_index: int,
223) -> str:
224    """Find the best answer for a given example.
225
226    Args:
227        all_start_logits:
228            The start logits for all the features.
229        all_end_logits:
230            The end logits for all the features.
231        prepared_dataset:
232            The dataset containing the prepared examples.
233        feature_indices:
234            The indices of the features associated with the current example.
235        context:
236            The context of the example.
237        max_answer_length:
238            The maximum length of the answer.
239        num_best_logits:
240            The number of best logits to consider.
241        min_null_score:
242            The minimum score an answer can have.
243        cls_token_index:
244            The index of the CLS token.
245
246    Returns:
247        The best answer for the example.
248    """
249    # Loop through all the features associated to the current example
250    valid_answers = list()
251    for feature_index in feature_indices:
252        features = prepared_dataset[feature_index]
253
254        # Get the predictions of the model for this feature
255        start_logits = all_start_logits[feature_index]
256        end_logits = all_end_logits[feature_index]
257
258        # Update minimum null prediction
259        cls_index = features["input_ids"].index(cls_token_index)
260        feature_null_score = (start_logits[cls_index] + end_logits[cls_index]).item()
261        if min_null_score < feature_null_score:
262            min_null_score = feature_null_score
263
264        # Find the valid answers for the feature
265        valid_answers_for_feature = find_valid_answers(
266            start_logits=start_logits,
267            end_logits=end_logits,
268            offset_mapping=features["offset_mapping"],
269            context=context,
270            max_answer_length=max_answer_length,
271            num_best_logits=num_best_logits,
272            min_null_score=min_null_score,
273        )
274        valid_answers.extend(valid_answers_for_feature)
275
276    # In the very rare edge case we have not a single non-null prediction, we create a
277    # fake prediction to avoid failure
278    if not valid_answers:
279        return ""
280
281    # Otherwise, we select the answer with the largest score as the best answer, and
282    # return it
283    best_answer_dict = sorted(valid_answers, key=lambda x: x["score"], reverse=True)[0]
284    return best_answer_dict["text"]
285
286
287def find_valid_answers(
288    start_logits: np.ndarray,
289    end_logits: np.ndarray,
290    offset_mapping: list[tuple[int, int]],
291    context: str,
292    max_answer_length: int,
293    num_best_logits: int,
294    min_null_score: float,
295) -> list[dict]:
296    """Find the valid answers from the start and end indexes.
297
298    Args:
299        start_logits:
300            The logits for the start of the answer.
301        end_logits:
302            The logits for the end of the answer.
303        offset_mapping:
304            The offset mapping, being a list of pairs of integers for each token index,
305            containing the start and end character index in the original context.
306        max_answer_length:
307            The maximum length of the answer.
308        num_best_logits:
309            The number of best logits to consider. Note that this function will run in
310             time.
311        min_null_score:
312            The minimum score an answer can have.
313
314    Returns:
315        A list of the valid answers, each being a dictionary with keys "text" and
316        "score", the score being the sum of the start and end logits.
317    """
318    # Fetch the top-k predictions for the start- and end token indices
319    start_indexes = np.argsort(start_logits)[-1 : -num_best_logits - 1 : -1].tolist()
320    end_indexes = np.argsort(end_logits)[-1 : -num_best_logits - 1 : -1].tolist()
321
322    # We loop over all combinations of starting and ending indexes for valid answers
323    valid_answers = list()
324    for start_index in start_indexes:
325        for end_index in end_indexes:
326            # If the starting or ending index is out-of-scope, meaning that they are
327            # either out of bounds or correspond to part of the input_ids that are not
328            # in the context, then we skip this index
329            if (
330                start_index >= len(offset_mapping)
331                or end_index >= len(offset_mapping)
332                or tuple(offset_mapping[start_index]) == (-1, -1)
333                or tuple(offset_mapping[end_index]) == (-1, -1)
334            ):
335                continue
336
337            # Do not consider answers with a length that is either negative or greater
338            # than the context length
339            max_val = max_answer_length + start_index - 1
340            if end_index < start_index or end_index > max_val:
341                continue
342
343            # If we got to this point then the answer is valid, so we store the
344            # corresponding start- and end character indices in the original context,
345            # and from these extract the answer
346            start_char = offset_mapping[start_index][0]
347            end_char = offset_mapping[end_index][1]
348            text = context[start_char:end_char]
349
350            # Compute the score of the answer, being the sum of the start and end
351            # logits. Intuitively, this indicates how likely the answer is to be
352            # correct, and allows us to pick the best valid answer.
353            score = start_logits[start_index] + end_logits[end_index]
354
355            # Add the answer to the list of valid answers, if the score is greater
356            # than the minimum null score
357            if score > min_null_score:
358                valid_answers.append(dict(score=score, text=text))
359
360    return valid_answers
361
362
363def postprocess_labels(dataset: Dataset) -> list[dict]:
364    """Postprocess the labels, to allow easier metric computation.
365
366    Args:
367        dataset:
368            The dataset containing the examples.
369
370    Returns:
371         The postprocessed labels.
372    """
373    labels = list()
374    for example in dataset:
375        # Create the associated reference dictionary, to be added to the list of
376        # references
377        label = dict(
378            id=str(example["id"]),
379            answers=dict(
380                text=[example["answer"]],
381                answer_start=[example["answer_start"]],
382            ),
383        )
384        labels.append(label)
385
386    return labels