primeqa__primeqa METASEP

tests/oneqa/mrc/processors/preprocessors/test_tydiqa_preprocessor.py METASEP
import datasets
import pytest
from datasets import Dataset
from transformers import AutoTokenizer

from oneqa.mrc.processors.preprocessors.tydiqa import TyDiQAPreprocessor
from tests.oneqa.mrc.common.base import UnitTest


class TestTyDiQAPreprocessor(UnitTest):

    @pytest.fixture(scope='class')
    def train_examples(self):
        examples = datasets.load_dataset("tydiqa", "primary_task", split='train[:100]')
        return examples

    @pytest.fixture(scope='class')
    def eval_examples(self):
        examples = datasets.load_dataset("tydiqa", "primary_task", split='validation[:100]')
        return examples

    @pytest.fixture(scope='class')
    def tokenizer(self):
        return AutoTokenizer.from_pretrained('xlm-roberta-base')

    @pytest.fixture(scope='class')
    def tydiqa_preprocessor(self, tokenizer):
        return TyDiQAPreprocessor(
            tokenizer,
            stride=128
        )

    def test_train_preprocessing_runs_without_errors(self, train_examples, tydiqa_preprocessor):
        train_examples, train_features = tydiqa_preprocessor.process_train(train_examples)
        assert isinstance(train_examples, Dataset)
        assert isinstance(train_features, Dataset)

    def test_eval_preprocessing_runs_without_errors(self, eval_examples, tydiqa_preprocessor):
        eval_examples, eval_features = tydiqa_preprocessor.process_eval(eval_examples)
        assert isinstance(eval_examples, Dataset)
        assert isinstance(eval_features, Dataset)

tests/oneqa/mrc/processors/preprocessors/test_default_preprocessor.py METASEP
from datasets import Dataset
import pytest
from pytest import raises
from transformers import AutoTokenizer

from oneqa.mrc.processors.preprocessors.default import DefaultPreProcessor
from oneqa.mrc.types.target_type import TargetType
from tests.oneqa.mrc.common.base import UnitTest
from tests.oneqa.mrc.common.parameterization import PARAMETERIZE_INVALID_SUBSAMPLING_PROBABILITIES


class TestDefaultPreProcessor(UnitTest):
    @pytest.fixture(scope='session')
    def train_examples_has_answer(self, train_examples):
        example_indices_has_answer = [i for i, t in enumerate(train_examples['target']) if t['passage_indices'][0] != -1]
        return train_examples.select(example_indices_has_answer)

    @pytest.fixture(scope='session')
    def train_examples_no_answer(self, train_examples):
        example_indices_no_answer = [i for i, t in enumerate(train_examples['target']) if t['passage_indices'][0] == -1]
        return train_examples.select(example_indices_no_answer)

    @pytest.fixture(scope='session')
    def invalid_type_train_examples(self):
        train_examples = Dataset.from_dict({'question': ['Who?'], 'context': [[-1]], 'target': [
            {'start_positions': [-1], 'end_positions': [-1], 'passage_indices': [-1], 'yes_no_answer': ['NONE']}]})
        return train_examples

    @pytest.fixture(scope='session')
    def invalid_name_train_examples(self, train_examples):
        return train_examples.rename_columns(dict(context='contextssss'))

    @pytest.fixture(scope='session')
    def invalid_name_eval_examples(self, eval_examples):
        return eval_examples.rename_columns(dict(context='contextssss'))

    @pytest.fixture(scope='session')
    def invalid_type_eval_examples(self, invalid_type_train_examples):
        return invalid_type_train_examples.remove_columns('target')

    @pytest.fixture(scope='session')
    def no_examples_train(self):
        return Dataset.from_dict({'question': [], 'context': [], 'target': [], 'example_id': []})

    @pytest.fixture(scope='session')
    def preprocessor_subsample_keep_all(self, tokenizer):
        return DefaultPreProcessor(
            tokenizer,
            stride=128,
            load_from_cache_file=False,
            negative_sampling_prob_when_has_answer=1.,
            negative_sampling_prob_when_no_answer=1.,
        )

    @pytest.fixture(scope='session')
    def preprocessor_subsample_keep_none(self, tokenizer):
        return DefaultPreProcessor(
            tokenizer,
            stride=128,
            load_from_cache_file=False,
            negative_sampling_prob_when_has_answer=0.,
            negative_sampling_prob_when_no_answer=0.,
        )

    @pytest.fixture(scope='session')
    def preprocessor_subsample_keep_no_answer(self, tokenizer):
        return DefaultPreProcessor(
            tokenizer,
            stride=128,
            load_from_cache_file=False,
            negative_sampling_prob_when_has_answer=0.,
            negative_sampling_prob_when_no_answer=1.,
        )

    @pytest.fixture(scope='session')
    def preprocessor_subsample_keep_has_answer(self, tokenizer):
        return DefaultPreProcessor(
            tokenizer,
            stride=128,
            load_from_cache_file=False,
            negative_sampling_prob_when_has_answer=1.,
            negative_sampling_prob_when_no_answer=0.,
        )

    @PARAMETERIZE_INVALID_SUBSAMPLING_PROBABILITIES
    def test_preprocessor_raises_value_error_on_invalid_subsampling_prob(
            self,
            tokenizer,
            negative_sampling_prob_when_has_answer,
            negative_sampling_prob_when_no_answer
    ):
        with raises(ValueError):
            _ = DefaultPreProcessor(
                tokenizer,
                stride=128,
                load_from_cache_file=False,
                negative_sampling_prob_when_has_answer=negative_sampling_prob_when_has_answer,
                negative_sampling_prob_when_no_answer=negative_sampling_prob_when_no_answer,
            )

    def test_adapt_dataset_with_train_examples(self, train_examples, preprocessor):
        assert preprocessor.adapt_dataset(train_examples, is_train=True) is train_examples

    def test_adapt_dataset_with_eval_examples(self, eval_examples, preprocessor):
        assert preprocessor.adapt_dataset(eval_examples, is_train=False) is eval_examples

    def test_train_preprocessing_runs_without_errors(self, train_examples, preprocessor):
        train_examples, train_features = preprocessor.process_train(train_examples)
        assert isinstance(train_examples, Dataset)
        assert isinstance(train_features, Dataset)
        expected_feature_columns = {'input_ids', 'attention_mask', 'example_idx', 'context_idx',
                                    'example_id', 'start_positions', 'end_positions', 'target_type'}
        expected_feature_columns.update(preprocessor._tokenizer.model_input_names)
        actual_feature_columns = set(train_features.column_names)
        assert actual_feature_columns == expected_feature_columns

    def test_process_train_feature_target_type_matches_position_labels(self, train_examples, preprocessor):
        _, train_features = preprocessor.process_train(train_examples)
        for i in range(train_features.num_rows):
            tt = train_features['target_type'][i]
            cls_index = train_features['input_ids'][i].index(preprocessor._tokenizer.cls_token_id)
            if tt in (TargetType.PASSAGE_ANSWER, TargetType.YES, TargetType.NO):
                # assert train_features['passage_indices'][i] != -1
                assert train_features['start_positions'][i] == cls_index
                assert train_features['end_positions'][i] == cls_index
            elif tt == TargetType.SPAN_ANSWER:
                # assert train_features['passage_indices'][i] != -1
                assert train_features['start_positions'][i] != cls_index
                assert train_features['end_positions'][i] != cls_index
            elif tt == TargetType.NO_ANSWER:
                # assert train_features['passage_indices'][i] == -1
                assert train_features['start_positions'][i] == cls_index
                assert train_features['end_positions'][i] == cls_index
            else:
                raise ValueError(f"Unexpected target type {tt}")

    def test_all_features_positive_when_subsample_keep_none(self, train_examples, preprocessor_subsample_keep_none):
        train_examples, train_features = preprocessor_subsample_keep_none.process_train(train_examples)
        assert all(tt != TargetType.NO_ANSWER for tt in train_features['target_type'])
        assert train_examples.num_rows > train_features.num_rows

    def test_some_features_negative_when_subsample_keep_all(self, train_examples, preprocessor_subsample_keep_all):
        train_examples, train_features = preprocessor_subsample_keep_all.process_train(train_examples)
        assert any(tt == TargetType.NO_ANSWER for tt in train_features['target_type'])
        assert train_examples.num_rows < train_features.num_rows

    def test_raises_value_error_when_subsampling_removes_all_features(
            self, train_examples_no_answer, preprocessor_subsample_keep_none):
        with raises(ValueError):
            _, _ = preprocessor_subsample_keep_none.process_train(train_examples_no_answer)

    def test_raises_value_error_when_no_examples(self, no_examples_train, preprocessor):
        with raises(ValueError):
            _, _ = preprocessor.process_train(no_examples_train)

    def test_subsample_keeps_features_from_has_answer_when_only_keeping_has_answer_negatives(
            self, train_examples_has_answer, preprocessor_subsample_keep_has_answer):
        train_examples_has_answer, train_features = preprocessor_subsample_keep_has_answer.process_train(
            train_examples_has_answer)
        for fi in range(train_features.num_rows):
            ei = train_features['example_idx'][fi]
            example_has_answer = train_examples_has_answer['target'][ei]['passage_indices'][0] != -1
            tt = train_features['target_type'][fi]
            negative_feature = tt == TargetType.NO_ANSWER
            if negative_feature:
                assert example_has_answer

    def test_subsample_keeps_features_from_no_answer_when_only_keeping_no_answer_negatives(
            self, train_examples_no_answer, preprocessor_subsample_keep_no_answer):
        train_examples_no_answer, train_features = preprocessor_subsample_keep_no_answer.process_train(
            train_examples_no_answer)
        for fi in range(train_features.num_rows):
            ei = train_features['example_idx'][fi]
            example_has_answer = train_examples_no_answer['target'][ei]['passage_indices'][0] != -1
            tt = train_features['target_type'][fi]
            negative_feature = tt == TargetType.NO_ANSWER
            if negative_feature:
                assert not example_has_answer

    def test_eval_preprocessing_runs_without_errors(self, eval_examples, preprocessor):
        eval_examples, eval_features = preprocessor.process_eval(eval_examples)
        assert isinstance(eval_examples, Dataset)
        assert isinstance(eval_features, Dataset)

    def test_cannot_adapt_dataset_with_invalid_train_schema_names(self, preprocessor, invalid_name_train_examples):
        with raises(ValueError):
            _ = preprocessor.adapt_dataset(invalid_name_train_examples, is_train=True)

    def test_cannot_adapt_dataset_with_invalid_train_schema_types(self, preprocessor, invalid_type_train_examples):
        with raises(ValueError):
            _ = preprocessor.adapt_dataset(invalid_type_train_examples, is_train=True)

    def test_cannot_adapt_dataset_with_invalid_train_schema_missing_target(self, preprocessor, eval_examples):
        with raises(ValueError):
            _ = preprocessor.adapt_dataset(eval_examples, is_train=True)

    def test_cannot_adapt_dataset_with_invalid_eval_schema_names(self, preprocessor, invalid_name_eval_examples):
        with raises(ValueError):
            _ = preprocessor.adapt_dataset(invalid_name_eval_examples, is_train=False)

    def test_cannot_adapt_dataset_with_invalid_eval_schema_types(self, preprocessor, invalid_type_eval_examples):
        with raises(ValueError):
            _ = preprocessor.adapt_dataset(invalid_type_eval_examples, is_train=False)

tests/oneqa/mrc/processors/preprocessors/__init__.py METASEP

tests/oneqa/mrc/models/heads/test_extractive.py METASEP
import pytest
import torch
from transformers import MODEL_MAPPING

from oneqa.mrc.models.heads.extractive import ExtractiveQAHead
from oneqa.mrc.types.model_outputs.extractive import ExtractiveQAModelOutput
from oneqa.mrc.types.target_type import TargetType
from tests.oneqa.mrc.common.base import UnitTest


class TestExtractiveQAHead(UnitTest):
    @pytest.fixture(scope='session')
    def config_and_language_model(self, model_name_and_config):
        model_name, config = model_name_and_config
        model = MODEL_MAPPING[config.__class__].from_pretrained(model_name, config=config)
        return config, model

    @pytest.fixture(scope='session')
    def extractive_head(self, model_name_and_config):
        _, config = model_name_and_config
        return ExtractiveQAHead(config)

    @pytest.fixture(scope='session')
    def language_model_outputs(self, config_and_language_model):
        _, model = config_and_language_model
        return model(**model.dummy_inputs)

    @pytest.fixture(scope='session')
    def language_model_outputs_tuple(self, language_model_outputs):
        return language_model_outputs.to_tuple()

    @pytest.fixture(scope='session')
    def training_inputs(self, config_and_language_model, language_model_outputs):
        _, model = config_and_language_model
        bs, seq_len = model.dummy_inputs['input_ids'].shape
        kwargs = dict(start_positions=torch.randint(0, seq_len, (bs, 1)),
                      end_positions=torch.randint(0, seq_len, (bs, 1)),
                      target_type=torch.randint(0, len(TargetType), (bs, 1)))
        args = (language_model_outputs,)
        return args, kwargs

    @pytest.fixture(scope='session')
    def training_inputs_with_tuple_from_language_model(self, training_inputs):
        args, kwargs = training_inputs
        args = (args[0].to_tuple(),)
        return args, kwargs

    def test_instantiation(self, model_name_and_config):
        _, config = model_name_and_config
        head = ExtractiveQAHead(config)
        assert head.num_labels == config.num_labels

    def test_correct_number_of_classification_labels_when_using_default(self, extractive_head):
        assert extractive_head.num_classification_head_labels == len(TargetType)

    def test_correct_number_of_classification_labels_when_overridden(self, model_name_and_config):
        _, config = model_name_and_config
        num_classification_labels = 16
        head = ExtractiveQAHead(config, num_labels_override=num_classification_labels)
        assert head.num_classification_head_labels == num_classification_labels

    def test_forward(self, config_and_language_model, language_model_outputs):
        config, model = config_and_language_model
        bs, seq_len = model.dummy_inputs['input_ids'].shape
        head = ExtractiveQAHead(config)
        results = head(language_model_outputs)

        assert isinstance(results, ExtractiveQAModelOutput)

        assert results.loss is None

        self._assert_is_floating_point_tensor(results.start_logits)
        assert results.start_logits.shape == (bs, seq_len)

        self._assert_is_floating_point_tensor(results.end_logits)
        assert results.end_logits.shape == (bs, seq_len)

        self._assert_is_floating_point_tensor(results.target_type_logits)
        assert results.target_type_logits.shape == (bs, len(TargetType))

    def test_forward_for_training(self, config_and_language_model, training_inputs):
        config, model = config_and_language_model
        bs, seq_len = model.dummy_inputs['input_ids'].shape
        head = ExtractiveQAHead(config)

        args, kwargs = training_inputs
        results = head(*args, **kwargs)

        assert isinstance(results, ExtractiveQAModelOutput)

        assert isinstance(results.loss.item(), float)

        self._assert_is_floating_point_tensor(results.start_logits)
        assert results.start_logits.shape == (bs, seq_len)

        self._assert_is_floating_point_tensor(results.end_logits)
        assert results.end_logits.shape == (bs, seq_len)

        self._assert_is_floating_point_tensor(results.target_type_logits)
        assert results.target_type_logits.shape == (bs, len(TargetType))

    def test_forward_with_tuple_input(self, config_and_language_model, language_model_outputs_tuple):
        config, model = config_and_language_model
        bs, seq_len = model.dummy_inputs['input_ids'].shape
        head = ExtractiveQAHead(config)
        results = head(language_model_outputs_tuple)

        assert isinstance(results, tuple)

        start_logits = results[0]
        self._assert_is_floating_point_tensor(start_logits)
        assert start_logits.shape == (bs, seq_len)

        end_logits = results[1]
        self._assert_is_floating_point_tensor(end_logits)
        assert end_logits.shape == (bs, seq_len)

        target_type_logits = results[2]
        self._assert_is_floating_point_tensor(target_type_logits)
        assert target_type_logits.shape == (bs, len(TargetType))

    def test_forward_for_training_with_tuple_input(self, config_and_language_model,
                                                   training_inputs_with_tuple_from_language_model):
        config, model = config_and_language_model
        bs, seq_len = model.dummy_inputs['input_ids'].shape
        head = ExtractiveQAHead(config)

        args, kwargs = training_inputs_with_tuple_from_language_model
        results = head(*args, **kwargs)

        assert isinstance(results, tuple)

        loss = results[0]
        assert isinstance(loss.item(), float)

        start_logits = results[1]
        self._assert_is_floating_point_tensor(start_logits)
        assert start_logits.shape == (bs, seq_len)

        end_logits = results[2]
        self._assert_is_floating_point_tensor(end_logits)
        assert end_logits.shape == (bs, seq_len)

        target_type_logits = results[3]
        self._assert_is_floating_point_tensor(target_type_logits)
        assert target_type_logits.shape == (bs, len(TargetType))

tests/oneqa/mrc/models/heads/test_abstract.py METASEP
import pytest
from pytest import raises
from transformers import PretrainedConfig

from oneqa.mrc.models.heads.abstract import AbstractTaskHead
from tests.oneqa.mrc.common.base import UnitTest


class TestAbstractTaskHead(UnitTest):
    @pytest.fixture()
    def mocked_config(self, mocker):
        return mocker.Mock(PretrainedConfig)

    def test_cannot_instantiate_abstract_class(self, mocked_config):
        with raises(TypeError):
            _ = AbstractTaskHead(mocked_config)

tests/oneqa/mrc/models/heads/__init__.py METASEP

tests/oneqa/mrc/trainers/test_default.py METASEP
import os
import tempfile

import pytest
from transformers import TrainingArguments, DataCollatorWithPadding

from oneqa.mrc.trainers.default import MRCTrainer
from tests.oneqa.mrc.common.base import UnitTest


class TestMRCTrainer(UnitTest):
    @pytest.fixture(scope='function')
    def training_args(self):
        with tempfile.TemporaryDirectory() as working_dir:
            yield TrainingArguments(
                    output_dir=os.path.join(working_dir, 'output_dir'),
                    do_train=True,
                    do_eval=False,  # TODO: trainer test eval
                    num_train_epochs=1.5,
                    fp16=False,
                )

    @pytest.fixture()
    def data_collator(self, training_args, tokenizer):
        return DataCollatorWithPadding(tokenizer, pad_to_multiple_of=64 if training_args.fp16 else None)

    @pytest.fixture()
    def trainer_with_extractive_model(
            self, training_args, config_and_model_with_extractive_head,
            tokenizer, train_examples_and_features, eval_examples_and_features, data_collator):
        _, model = config_and_model_with_extractive_head
        _, train_features = train_examples_and_features
        _, eval_features = eval_examples_and_features
        return MRCTrainer(
            model=model,
            args=training_args,
            train_dataset=train_features if training_args.do_train else None,
            eval_dataset=eval_features 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=postprocessor.process,
            compute_metrics=None,
        )

    def test_trainer_trains_without_errors(self, trainer_with_extractive_model):
        train_result = trainer_with_extractive_model.train()
        assert isinstance(train_result.training_loss, float)

tests/oneqa/mrc/trainers/__init__.py METASEP

tests/oneqa/mrc/processors/__init__.py METASEP

tests/oneqa/mrc/models/test_task_model.py METASEP
import pytest
import torch
from pytest import raises
from transformers import MODEL_FOR_PRETRAINING_MAPPING

from oneqa.mrc.models.heads.extractive import ExtractiveQAHead, EXTRACTIVE_HEAD
from oneqa.mrc.models.task_model import ModelForDownstreamTasks
from oneqa.mrc.types.model_outputs.extractive import ExtractiveQAModelOutput
from oneqa.mrc.types.target_type import TargetType
from tests.oneqa.mrc.common.base import UnitTest


class TestModelForDownstreamTasks(UnitTest):
    def test_from_config(self, config_and_model_with_extractive_head):
        config, model = config_and_model_with_extractive_head
        head_name = next(iter(EXTRACTIVE_HEAD))
        assert isinstance(model.task_heads[head_name], ExtractiveQAHead)
        assert isinstance(model, ModelForDownstreamTasks)
        assert isinstance(model, MODEL_FOR_PRETRAINING_MAPPING[config.__class__])

    def test_model_class_from_config_then_from_pretrained(self, model_name_and_config):
        model_name, config = model_name_and_config
        model_class = ModelForDownstreamTasks.model_class_from_config(config)
        model = model_class.from_pretrained(model_name,
                                            config=config,
                                            task_heads=EXTRACTIVE_HEAD)
        head_name = next(iter(EXTRACTIVE_HEAD))
        assert isinstance(model.task_heads[head_name], ExtractiveQAHead)
        assert isinstance(model, ModelForDownstreamTasks)
        assert isinstance(model, MODEL_FOR_PRETRAINING_MAPPING[config.__class__])

    def test_raises_type_error_on_direct_instantiation(self, model_name_and_config):
        model_name, config = model_name_and_config
        with raises(TypeError):
            _ = ModelForDownstreamTasks.from_pretrained(model_name,
                                                        config=config,
                                                        task_heads=EXTRACTIVE_HEAD)

    def test_raises_value_error_on_empty_task_head_dict(self, model_name_and_config):
        model_name, config = model_name_and_config
        with raises(ValueError):
            _ = ModelForDownstreamTasks.from_config(config,
                                                    model_name,
                                                    task_heads={})

    def test_model__property(self, config_and_model_with_extractive_head):
        _, model = config_and_model_with_extractive_head
        assert model.model_ is getattr(model, model.base_model_prefix)

    def test_task_head_property_raises_value_error_when_not_set(self, model_name_and_config):
        model_name, config = model_name_and_config
        model = ModelForDownstreamTasks.from_config(config,
                                                    model_name,
                                                    task_heads=EXTRACTIVE_HEAD)
        with raises(ValueError):
            _ = model.task_head

    def test_task_head_property_retreives_set_task_head(self, config_and_model_with_extractive_head):
        _, model = config_and_model_with_extractive_head
        head_name = next(iter(EXTRACTIVE_HEAD))
        model.set_task_head(head_name)
        assert model.task_head is model.task_heads[head_name]

    def test_set_task_head_raises_key_error_on_invalid_name(self, config_and_model_with_extractive_head):
        _, model = config_and_model_with_extractive_head
        head_name = 'NOT_A_REAL_NAME'
        with raises(KeyError):
            model.set_task_head(head_name)

    def test_set_head_multiple_times(self, model_name_and_config):
        model_name, config = model_name_and_config
        task_heads = dict(qa_head=ExtractiveQAHead, qa_head_2=ExtractiveQAHead)
        model = ModelForDownstreamTasks.from_config(config,
                                                    model_name,
                                                    task_heads=task_heads)
        for _ in range(5):
            for head_name in task_heads:
                model.set_task_head(head_name)
                assert model.task_head is model.task_heads[head_name]

    def test_forward_with_extractive_head(self, config_and_model_with_extractive_head):
        _, model = config_and_model_with_extractive_head
        bs, seq_len = model.dummy_inputs['input_ids'].shape
        results = model(**model.dummy_inputs)

        assert isinstance(results, ExtractiveQAModelOutput)

        assert results.loss is None

        self._assert_is_floating_point_tensor(results.start_logits)
        assert results.start_logits.shape == (bs, seq_len)

        self._assert_is_floating_point_tensor(results.end_logits)
        assert results.end_logits.shape == (bs, seq_len)

        self._assert_is_floating_point_tensor(results.target_type_logits)
        assert results.target_type_logits.shape == (bs, len(TargetType))

    def test_forward_for_training_with_extractive_head(self, config_and_model_with_extractive_head,
                                                       extractive_training_inputs):
        _, model = config_and_model_with_extractive_head
        bs, seq_len = model.dummy_inputs['input_ids'].shape
        results = model(**extractive_training_inputs)

        assert isinstance(results, ExtractiveQAModelOutput)

        assert isinstance(results.loss.item(), float)

        self._assert_is_floating_point_tensor(results.start_logits)
        assert results.start_logits.shape == (bs, seq_len)

        self._assert_is_floating_point_tensor(results.end_logits)
        assert results.end_logits.shape == (bs, seq_len)

        self._assert_is_floating_point_tensor(results.target_type_logits)
        assert results.target_type_logits.shape == (bs, len(TargetType))

    def test_forward_with_extractive_head_with_tuple_output(self, config_and_model_with_extractive_head):
        _, model = config_and_model_with_extractive_head
        bs, seq_len = model.dummy_inputs['input_ids'].shape
        results = model(**model.dummy_inputs, return_dict=False)

        assert isinstance(results, tuple)

        start_logits = results[0]
        self._assert_is_floating_point_tensor(start_logits)
        assert start_logits.shape == (bs, seq_len)

        end_logits = results[1]
        self._assert_is_floating_point_tensor(end_logits)
        assert end_logits.shape == (bs, seq_len)

        target_type_logits = results[2]
        self._assert_is_floating_point_tensor(target_type_logits)
        assert target_type_logits.shape == (bs, len(TargetType))

    def test_forward_for_training_with_extractive_head_with_tuple_output(
            self, config_and_model_with_extractive_head, extractive_training_inputs):
        _, model = config_and_model_with_extractive_head
        bs, seq_len = model.dummy_inputs['input_ids'].shape
        results = model(**extractive_training_inputs, return_dict=False)

        assert isinstance(results, tuple)

        loss = results[0]
        assert isinstance(loss.item(), float)

        start_logits = results[1]
        self._assert_is_floating_point_tensor(start_logits)
        assert start_logits.shape == (bs, seq_len)

        end_logits = results[2]
        self._assert_is_floating_point_tensor(end_logits)
        assert end_logits.shape == (bs, seq_len)

        target_type_logits = results[3]
        self._assert_is_floating_point_tensor(target_type_logits)
        assert target_type_logits.shape == (bs, len(TargetType))

tests/oneqa/mrc/models/__init__.py METASEP

tests/oneqa/mrc/common/parameterization.py METASEP
"""
Test fixtures and parameterization
"""
from functools import wraps
from typing import Callable

import pytest

# Parameterize model names in test fixtures
# _MODEL_NAMES = ("model_name", [
#         'roberta-base', 'xlm-roberta-base', 'bert-base-uncased', 'albert-base-v2',
#         # 'facebook/bart-base'  # TODO: add bart support
# ])
_MODEL_NAMES = [
        'roberta-base', 'xlm-roberta-base', 'bert-base-uncased', 'albert-base-v2',
        # 'facebook/bart-base'  # TODO: add bart support
]
# _PARAMETERIZE_TEST_WITH_MODEL_NAME = pytest.mark.parametrize(*_MODEL_NAMES)  # TODO: remove
# _PARAMETERIZE_FIXTURE_WITH_MODEL_NAME = pytest.fixture(scope='package', params=_MODEL_NAMES[1])  # TODO: remove


# TODO: remove
# # noinspection PyPep8Naming
# def PARAMETERIZE_TEST_WITH_MODEL_NAME(f: Callable) -> Callable:
#     @pytest.mark.flaky(reruns=5, reruns_delay=2)
#     @pytest.mark.parametrize(*_MODEL_NAMES)
#     @wraps(f)
#     def inner(*args, **kwargs):
#         return f(*args, **kwargs)
#
#     return inner


# Cannot be rewritten as constant parametrize(fixture)
# noinspection PyPep8Naming
def PARAMETERIZE_FIXTURE_WITH_MODEL_NAME(f: Callable) -> Callable:
    @pytest.mark.flaky(reruns=5, reruns_delay=2)  # Account for intermittent S3 errors downloading HF data
    @pytest.fixture(scope='session', params=_MODEL_NAMES)
    @wraps(f)
    def inner(*args, **kwargs):
        return f(*args, **kwargs)
    return inner


_INVALID_PROBABILITIES = [-0.01, 1.01]
PARAMETERIZE_INVALID_SUBSAMPLING_PROBABILITIES = pytest.mark.parametrize(
    ["negative_sampling_prob_when_has_answer", "negative_sampling_prob_when_no_answer"],
    [(p1, p2) for p1 in _INVALID_PROBABILITIES for p2 in _INVALID_PROBABILITIES])

tests/oneqa/mrc/common/base.py METASEP
import pytest
import torch
from datasets import Dataset
from transformers import AutoConfig, AutoModel, AutoTokenizer

from oneqa.mrc.models.heads.extractive import EXTRACTIVE_HEAD
from oneqa.mrc.models.task_model import ModelForDownstreamTasks
from oneqa.mrc.processors.preprocessors.default import DefaultPreProcessor
from oneqa.mrc.types.target_type import TargetType
from tests.oneqa.mrc.common.parameterization import PARAMETERIZE_FIXTURE_WITH_MODEL_NAME


class UnitTest:
    """
    Base class for all unit test classes
    """

    @PARAMETERIZE_FIXTURE_WITH_MODEL_NAME
    def model_name_and_config(self, request):
        model_name = request.param
        _ = AutoModel.from_pretrained(model_name)  # Pre-download LM inside flaky fixture so other tests have it
        return model_name, AutoConfig.from_pretrained(model_name)

    @pytest.mark.flaky(reruns=5, reruns_delay=2)  # Account for intermittent S3 errors downloading HF data
    @pytest.fixture(scope='session')
    def tokenizer(self, model_name_and_config):
        model_name, _ = model_name_and_config
        return AutoTokenizer.from_pretrained(model_name)

    @pytest.fixture(scope='session')
    def config_and_model_with_extractive_head(self, model_name_and_config):
        model_name, config = model_name_and_config
        model = ModelForDownstreamTasks.from_config(config,
                                                    model_name,
                                                    task_heads=EXTRACTIVE_HEAD)
        head_name = next(iter(EXTRACTIVE_HEAD))
        model.set_task_head(head_name)
        return config, model

    @pytest.fixture(scope='session')
    def extractive_training_inputs(self, config_and_model_with_extractive_head):
        _, model = config_and_model_with_extractive_head
        bs, seq_len = model.dummy_inputs['input_ids'].shape
        kwargs = dict(start_positions=torch.randint(0, seq_len, (bs, 1)),
                      end_positions=torch.randint(0, seq_len, (bs, 1)),
                      target_type=torch.randint(0, len(TargetType), (bs, 1)))
        kwargs.update(model.dummy_inputs)
        return kwargs

    @pytest.fixture(scope='session')
    def train_examples(self):
        question = ["Who walked the dog?", "What time is it?"]
        context = [["Alice walks the cat", "Bob walks the dog"],
                   ["The quick brown fox jumps over the lazy dog", "Glenn the otter lives at the aquarium", "Go"]]
        example_id = ["foo-abc", "bar-123"]
        start_positions = [[0], [-1]]
        end_positions = [[2], [-1]]
        passage_indices = [[1], [-1]]
        yes_no_answer = [["NONE"], ["NONE"]]
        examples_dict = dict(question=question, context=context, example_id=example_id,
                             target=[dict(start_positions=s, end_positions=e, passage_indices=p, yes_no_answer=yn)
                                     for s, e, p, yn in
                                     zip(start_positions, end_positions, passage_indices, yes_no_answer)])
        examples_dataset = Dataset.from_dict(examples_dict)
        return examples_dataset

    @pytest.fixture(scope='session')
    def eval_examples(self, train_examples):
        return train_examples.remove_columns('target')

    @pytest.fixture(scope='session')
    def preprocessor(self, tokenizer):
        return DefaultPreProcessor(
            tokenizer,
            stride=128,
            load_from_cache_file=False,
            negative_sampling_prob_when_has_answer=0.5,
            negative_sampling_prob_when_no_answer=0.5,
        )

    @pytest.fixture(scope='session')
    def train_examples_and_features(self, train_examples, preprocessor):
        return preprocessor.process_train(train_examples)

    @pytest.fixture(scope='session')
    def eval_examples_and_features(self, eval_examples, preprocessor):
        return preprocessor.process_eval(eval_examples)

    @staticmethod
    def _assert_is_floating_point_tensor(tensor):
        assert isinstance(tensor, torch.Tensor)
        assert tensor.dtype.is_floating_point

tests/oneqa/mrc/common/__init__.py METASEP

tests/oneqa/mrc/__init__.py METASEP

tests/oneqa/__init__.py METASEP

tests/__init__.py METASEP

tox.ini METASEP
[tox]
envlist = py37,py38,py39

[testenv]
# install pytest in the virtualenv where commands will be executed
extras = tests
commands =
    pytest --cov oneqa --cov-config .coveragerc tests/
pyproject.toml METASEP
[build-system]
requires = ["setuptools", "wheel"]
build-backend = "setuptools.build_meta"
VERSION METASEP
0.0.0
README.md METASEP
# OneQA
OneQA to rule them all

[![Build Status](https://travis.ibm.com/IBM-Research-AI/OneQA.svg?token=XcbF7zxMKHD12hqZiBwc&branch=master)](https://travis.ibm.com/IBM-Research-AI/OneQA)

MANIFEST.in METASEP
include README.md
include VERSION
include .coveragerc
global-exclude __pycache__
global-exclude *.py[co]
oneqa/mrc/types/model_outputs/extractive.py METASEP
from dataclasses import dataclass
from typing import Optional, Tuple

import torch
from transformers.file_utils import ModelOutput


@dataclass
class ExtractiveQAModelOutput(ModelOutput):
    #         # (loss), start_logits, end_logits, target_type_logits, (hidden_states), (attentions)
    loss: Optional[torch.FloatTensor] = None
    start_logits: torch.FloatTensor = None
    end_logits: torch.FloatTensor = None
    target_type_logits: torch.FloatTensor = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None
oneqa/mrc/types/model_outputs/__init__.py METASEP

oneqa/mrc/processors/preprocessors/tydiqa.py METASEP
from operator import itemgetter

from datasets import Dataset
from datasets.features.features import Sequence, Value

from oneqa.mrc.processors.preprocessors.default import DefaultPreProcessor


class TyDiQAPreprocessor(DefaultPreProcessor):  # TODO type signatures for all methods
    _feature_types = {'question_text': Value(dtype='string', id=None),
                      'document_plaintext': Value(dtype='string', id=None)}
    _train_feature_types = {
        'annotations': Sequence(feature={
                   'minimal_answers_end_byte': Value(dtype='int32', id=None),
                   'minimal_answers_start_byte': Value(dtype='int32', id=None),
                   'passage_answer_candidate_index': Value(dtype='int32', id=None),
                   'yes_no_answer': Value(dtype='string', id=None)})
    }
    _byte_itemgetter = itemgetter('plaintext_start_byte', 'plaintext_end_byte')
    _rename_fields = {'question_text': 'question', 'annotations': 'target'}
    _rename_target = {'passage_answer_candidate_index': 'passage_indices',
                      'minimal_answers_start_byte': 'start_positions',
                      'minimal_answers_end_byte': 'end_positions'}

    def adapt_dataset(self, dataset: Dataset, is_train: bool) -> Dataset:
        self.validate_schema(dataset, is_train)
        dataset = dataset.rename_columns(self._rename_fields)
        dataset = dataset.map(self._create_target,
                              load_from_cache_file=self._load_from_cache_file,
                              num_proc=self._num_workers
                              )
        dataset = dataset.map(self._split_context,
                              load_from_cache_file=self._load_from_cache_file,
                              num_proc=self._num_workers
        )
        dataset = super().adapt_dataset(dataset, is_train)
        return dataset
    
    def _split_context(self, example):
        context_bytes = example['document_plaintext'].encode('utf-8')
        context = [
            context_bytes[start_byte: end_byte].decode('utf-8')
            for start_byte, end_byte in zip(*self._byte_itemgetter(example['passage_answer_candidates']))
        ]
        example['context'] = context

        for i in range(len(example['target']['passage_indices'])):
            pidx = example['target']['passage_indices'][i]
            if pidx == -1:
                continue

            offset = example['passage_answer_candidates']['plaintext_start_byte'][pidx]
            example['target']['start_positions'][i] -= offset
            example['target']['end_positions'][i] -= offset

        return example

    def _create_target(self, example):
        target = example['target']
        for old_key, new_key in self._rename_target.items():
            target[new_key] = target.pop(old_key)
        # TODO text extraction by byte from document_plaintext (generative support)
        example['target'] = target
        return example

oneqa/mrc/processors/preprocessors/default.py METASEP
import itertools
import random
import uuid
from operator import sub
from typing import List, Iterable, Tuple, Any, Dict, Optional, Type

from datasets.arrow_dataset import Batch
from transformers import BatchEncoding
from datasets import Dataset
from datasets.features.features import Sequence, Value, Features

from oneqa.mrc.processors.preprocessors.abstract import AbstractPreProcessor
from oneqa.mrc.types.subsample_type import SubsampleType
from oneqa.mrc.types.target_type import TargetType


# TODO type signatures for all methods
class DefaultPreProcessor(AbstractPreProcessor):  # todo better name?
    _del_keys = ["overflow_to_sample_mapping", "offset_mapping"]
    _feature_types = {'question': Value(dtype='string', id=None),
                      'context': Sequence(feature=Value(dtype='string', id=None), length=-1, id=None)}
    _train_feature_types = {
        'target': {'start_positions': Sequence(feature=Value(dtype='int64', id=None), length=-1, id=None),
                   'end_positions': Sequence(feature=Value(dtype='int64', id=None), length=-1, id=None),
                   'passage_indices': Sequence(feature=Value(dtype='int64', id=None), length=-1, id=None),
                   'yes_no_answer': Sequence(feature=Value(dtype='string', id=None), length=-1, id=None)}}
    # _train_feature_types = {   # TODO: switch this layout to match tydi
    #     'annotations': Sequence(feature={
    #         'start_positions': Value(dtype='int32', id=None),
    #         'end_positions': Value(dtype='int32', id=None),
    #         'passage_indices': Value(dtype='int32', id=None),
    #         'yes_no_answer': Value(dtype='string', id=None)})
    # }
    _example_id_type = Value(dtype='string', id=None)

    def adapt_dataset(self, dataset: Dataset, is_train: bool) -> Dataset:
        self.validate_schema(dataset, is_train, pre_adaptation=False)
        return dataset

    def process_train(self, examples: Dataset) -> Tuple[Dataset, Dataset]:
        return self._process(examples, is_train=True)

    def process_eval(self, examples: Dataset) -> Tuple[Dataset, Dataset]:
        return self._process(examples, is_train=False)

    def _process(self, examples: Dataset, is_train: bool) -> Tuple[Dataset, Dataset]:
        examples = self.adapt_dataset(examples, is_train)
        if examples.num_rows == 0:
            raise ValueError("No examples to process")

        features = examples.map(
            self._process_batch,
            fn_kwargs=dict(is_train=is_train),
            batched=True,
            num_proc=self._num_workers,
            remove_columns=examples.column_names,
            load_from_cache_file=self._load_from_cache_file,
            desc=f"Running tokenizer on {'train' if is_train else 'eval'} dataset",
            # features=self._create_features_schema(is_train),
        )
        if is_train:
            features = self.subsample_features(features)
        return examples, features

    def _process_batch(self, examples: Batch, is_train: bool) -> BatchEncoding:
        examples_question = examples['question']
        examples_context = examples['context']
        if isinstance(examples_question, str):  # wrap single (question, [context]) pair in list
            examples_question = [examples_question]
            examples_context = [examples_context]
        
        # create 1:1 question:context lists
        expanded_examples_question = []
        expanded_examples_idx = []
        for i, (question, context) in enumerate(zip(examples_question, examples_context)):
            expanded_examples_question.extend(itertools.repeat(question, len(context)))
            expanded_examples_idx.extend(itertools.repeat(i, len(context)))
        expanded_examples_context = list(itertools.chain.from_iterable(examples_context))

        tokenized_examples = self._tokenizer(
            expanded_examples_context if self._pad_on_right else expanded_examples_question,
            expanded_examples_question if self._pad_on_right else expanded_examples_context,
            stride=self._stride,
            max_length=self._max_seq_len,
            truncation='only_first' if self._pad_on_right else 'only_second',
            return_overflowing_tokens=True,
            return_offsets_mapping=True,
        )

        examples_id = list(examples.get('example_id', (str(uuid.uuid4()) for _ in range(len(examples_question)))))
        tokenized_examples['example_idx'] = [expanded_examples_idx[oidx] for oidx in tokenized_examples["overflow_to_sample_mapping"]]
        spans_per_example = self._generate_previous_spans_per_example(tokenized_examples['example_idx'], tokenized_examples["overflow_to_sample_mapping"])
        tokenized_examples['context_idx'] = list(map(sub, tokenized_examples["overflow_to_sample_mapping"], spans_per_example))
        # tokenized_examples['window_idx'] = [idx for _, group in itertools.groupby(tokenized_examples['context_idx']) for idx, _ in enumerate(group)]  # TODO: is this needed?
        tokenized_examples['example_id'] = [examples_id[eidx] for eidx in tokenized_examples['example_idx']]

        if is_train:
            tokenized_examples = self._create_train_targets(tokenized_examples, examples['target'])
        else:
            tokenized_examples = self._create_eval_targets(tokenized_examples)

        for key in self._del_keys:
            tokenized_examples.pop(key, None)

        if is_train:
            tokenized_examples = self.label_features_for_subsampling(tokenized_examples, examples)
        
        return tokenized_examples

    def _create_train_targets(self, tokenized_examples: BatchEncoding, target: List[Dict[str, Any]]) -> BatchEncoding:  # TODO: rename create targets fns?
        
        # Since one context might give us several features if it has a long context, and each example can have many contexts,
        # we need a map from a feature to ts corresponding example. This key gives us just that.
        example_mapping = tokenized_examples['example_idx']
        # 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["offset_mapping"]

        # Let's label those examples!
        tokenized_examples["start_positions"] = []
        tokenized_examples["end_positions"] = []
        tokenized_examples["target_type"] = []

        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(self._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.
            example_index = example_mapping[i]
            context_idx = tokenized_examples['context_idx'][i]
            t = target[example_index]
            passage_index = t['passage_indices'][0]
            start_position = t['start_positions'][0]
            end_position = t['end_positions'][0]
            yes_no_answer = TargetType.from_bool_label(t['yes_no_answer'][0])

            if passage_index == context_idx and start_position == -1:  # Passage Answer
                tokenized_examples["start_positions"].append(cls_index)
                tokenized_examples["end_positions"].append(cls_index)
                tt = yes_no_answer
                if tt not in (TargetType.YES, TargetType.NO):
                    tt = TargetType.PASSAGE_ANSWER
                tokenized_examples["target_type"].append(tt)
            elif passage_index != context_idx or start_position == -1:  # No Answer
                tokenized_examples["start_positions"].append(cls_index)
                tokenized_examples["end_positions"].append(cls_index)
                tokenized_examples["target_type"].append(TargetType.NO_ANSWER)
            else:  # Span Answer
                # Start/end character index of the answer in the text.
                start_char = start_position
                end_char = end_position

                # Start token index of the current span in the text.
                token_start_index = 0
                while sequence_ids[token_start_index] != (1 if self._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 self._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)
                    tokenized_examples["target_type"].append(TargetType.NO_ANSWER)
                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)
                    tokenized_examples["target_type"].append(TargetType.SPAN_ANSWER)

        return tokenized_examples

    def _create_eval_targets(self, tokenized_examples: BatchEncoding) -> BatchEncoding:
        
        context_index = 1 if self._pad_on_right else 0
        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)

            # 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
    
    def label_features_for_subsampling(self, tokenized_examples: BatchEncoding, examples: Batch) -> BatchEncoding:
        if self._subsample_all_features:
            self._logger.warning("Keeping all negative training instances -- "
                                 "this may create an unbalanced training set and increase training time significantly")
            return tokenized_examples
        elif self._subsample_no_features:
            self._logger.warning("Removing all negative training instances -- only positives will be used")

        example_mapping = tokenized_examples['example_idx']

        tokenized_examples['subsample_type'] = []
        for i in range(len(tokenized_examples['input_ids'])):
            if tokenized_examples['target_type'][i] != TargetType.NO_ANSWER:
                st = SubsampleType.POSITIVE
            else:
                example_idx = example_mapping[i]
                passage_mapping = examples['target'][example_idx]['passage_indices']
                passage_idx = passage_mapping[0]
                has_answer = passage_idx != -1
                if has_answer:
                    st = SubsampleType.NEGATIVE_HAS_ANSWER
                else:
                    st = SubsampleType.NEGATIVE_NO_ANSWER
            tokenized_examples['subsample_type'].append(st)

        return tokenized_examples

    def subsample_features(self, dataset: Dataset) -> Dataset:
        if self._subsample_all_features:
            return dataset

        keep_indices = (i for i, st in enumerate(dataset['subsample_type']) if self._keep_feature(st))
        try:
            dataset = dataset.select(keep_indices)
        except IndexError as ex:
            raise ValueError("No features remaining after subsampling") from ex
        dataset = dataset.remove_columns('subsample_type')
        return dataset

    def _keep_feature(self, st: SubsampleType) -> bool:
        if st == SubsampleType.POSITIVE:
            return True
        elif st == SubsampleType.NEGATIVE_HAS_ANSWER:
            return random.random() < self._negative_sampling_prob_when_has_answer
        elif st == SubsampleType.NEGATIVE_NO_ANSWER:
            return random.random() < self._negative_sampling_prob_when_no_answer
        else:
            raise NotImplementedError(f"Unexpected subsample type: {st}")

    @property
    def _pad_on_right(self) -> bool:
        return self._tokenizer.padding_side == "right"

    @property
    def _subsample_all_features(self) -> bool:
        return self._negative_sampling_prob_when_has_answer == self._negative_sampling_prob_when_no_answer == 1.

    @property
    def _subsample_no_features(self) -> bool:
        return self._negative_sampling_prob_when_has_answer == self._negative_sampling_prob_when_no_answer == 0.

    @staticmethod
    def _generate_previous_spans_per_example(example_idx: List[int], sample_mapping: List[int]) -> Iterable[int]:
        group_start_idx = 0
        for _, group in itertools.groupby(example_idx):
            group_len = None
            for group_len, _ in enumerate(group, 1):
                pass
            if group_len is None:  # this should never be triggered
                raise ValueError("Unexpected group length None")
            yield from itertools.repeat(sample_mapping[group_start_idx], group_len)
            group_start_idx += group_len

    def validate_schema(self, dataset: Dataset, is_train: bool, pre_adaptation: bool = True) -> None:
        cls = type(self) if pre_adaptation else DefaultPreProcessor
        items = cls._feature_types.items()
        if is_train:
            items = itertools.chain(items, cls._train_feature_types.items())
        for feature_name, feature_type in items:
            if feature_name not in dataset.features:
                raise ValueError(f"Expected but did not find feature '{feature_name}' in dataset")
            elif dataset.features[feature_name] != feature_type:
                raise ValueError(F"Feature type mismatch for feature '{feature_name}'. "
                                 F"Expected {feature_type} but found {dataset.features[feature_name]}")

        if 'example_id' in dataset.features:
            if dataset.features['example_id'] != self._example_id_type:
                raise ValueError(F"Feature type mismatch for feature 'example_id'. "
                                 F"Expected {self._example_id_type} but found {dataset.features['example_id']}")

    # def _create_features_schema(self, is_train: bool) -> Features:
    #     schema = Features()
    #     schema = schema.update(self._feature_types)
    #     schema['example_id'] = self._example_id_type
    #     if is_train:
    #         schema.update(self._train_feature_types)
    #     return schema

oneqa/mrc/processors/preprocessors/abstract.py METASEP
import logging
from abc import ABCMeta, abstractmethod
from typing import Optional, Tuple, Type

from datasets.arrow_dataset import Batch
from transformers import PreTrainedTokenizerFast, BatchEncoding
from datasets import Dataset


class AbstractPreProcessor(metaclass=ABCMeta):  # TODO type signatures and docstrings for all methods
    def __init__(self, 
                 tokenizer: PreTrainedTokenizerFast,
                 stride: int,
                 max_seq_len: Optional[int] = None,
                 negative_sampling_prob_when_has_answer: float = 0.01,
                 negative_sampling_prob_when_no_answer: float = 0.04,
                 num_workers: Optional[int] = None,
                 load_from_cache_file: bool = False):
        self._logger = logging.getLogger(self.__class__.__name__)
        self._tokenizer = tokenizer
        self._stride = stride
        self._max_seq_len = max_seq_len
        self._negative_sampling_prob_when_has_answer = negative_sampling_prob_when_has_answer
        self._negative_sampling_prob_when_no_answer = negative_sampling_prob_when_no_answer
        self._num_workers = num_workers
        self._load_from_cache_file = load_from_cache_file

        if not (0. <= self._negative_sampling_prob_when_has_answer <= 1.):
            raise ValueError(f"Expected 0 <= negative_sampling_prob_when_has_answer <= 1 but got: "
                             f"{self._negative_sampling_prob_when_has_answer:.02f}")

        if not (0. <= self._negative_sampling_prob_when_has_answer <= 1.):
            raise ValueError(f"Expected 0 <= negative_sampling_prob_when_no_answer <= 1 but got: "
                             f"{self._negative_sampling_prob_when_no_answer:.02f}")

    @abstractmethod
    def process_train(self, examples: Dataset) -> Tuple[Dataset, Dataset]:
        pass

    @abstractmethod
    def process_eval(self, examples: Dataset) -> Tuple[Dataset, Dataset]:
        pass

    @abstractmethod
    def adapt_dataset(self, dataset: Dataset, is_train: bool) -> Dataset:
        pass

    @abstractmethod
    def label_features_for_subsampling(self, tokenized_examples: BatchEncoding, examples: Batch) -> BatchEncoding:
        pass

    @abstractmethod
    def subsample_features(self, dataset: Dataset) -> Dataset:
        pass

    @abstractmethod
    def validate_schema(self, dataset: Dataset, is_train: bool, pre_adaptation: bool = True) -> None:
        pass

oneqa/mrc/processors/preprocessors/__init__.py METASEP

oneqa/mrc/processors/postprocessors/extractive.py METASEP
from collections import defaultdict
from itertools import groupby
from operator import itemgetter

from tqdm import tqdm

from oneqa.mrc.processors.postprocessors.abstract import AbstractPostProcessor


class ExtractivePostProcessor(AbstractPostProcessor):
    def __init__(self, k: int, n_best_size: int, max_answer_length: int, span_tracker_factory):
        super().__init__(k)
        self._n_best_size = n_best_size
        self._max_answer_length = max_answer_length
        self._span_trackers = defaultdict(span_tracker_factory)  # TODO factory type?

    def process(self, examples, features, predictions):
        features_itr = groupby(features, key=itemgetter('example_id'))
        predictions_i = 0
        if len(features) != predictions[0].shape[0] and all(
                p.shape[0] == predictions[0].shape[0] for p in predictions[1:]):
            raise ValueError(f"Size mismatch withing {len(features)} features and predictions "
                             f"of first dim {[p.shape[0] for p in predictions]}")

        for example in tqdm(examples):
            example_id, example_features = next(features_itr)
            if example_id != example['example_id']:
                raise ValueError(f"Example id mismatch between example ({example['example_id']}) "
                                 f"and feature ({example_id})")
            example_features = list(example_features)
            example_predictions = None  # [i, i + len(example_features)



oneqa/mrc/processors/postprocessors/abstract.py METASEP
import logging
from abc import ABCMeta, abstractmethod


class AbstractPostProcessor(metaclass=ABCMeta):
    def __init__(self, k: int):
        self._logger = logging.getLogger(self.__class__.__name__)
        self._k = k

    @abstractmethod
    def process(self, examples, features, predictions):
        pass  # TODO types on args?

oneqa/mrc/processors/postprocessors/__init__.py METASEP

oneqa/mrc/models/heads/extractive.py METASEP
from copy import deepcopy
from typing import Union, Optional

import torch
from transformers import PretrainedConfig
from transformers.file_utils import ModelOutput
from transformers.modeling_outputs import BaseModelOutputWithPoolingAndCrossAttentions
from transformers.models.roberta.modeling_roberta import RobertaClassificationHead

from oneqa.mrc.models.heads.abstract import AbstractTaskHead
from oneqa.mrc.types.model_outputs.extractive import ExtractiveQAModelOutput
from oneqa.mrc.types.target_type import TargetType


class ExtractiveQAHead(AbstractTaskHead):
    def __init__(self, config: PretrainedConfig, num_labels_override: Optional[int] = None):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.qa_outputs = torch.nn.Linear(config.hidden_size, self.num_labels)

        config_for_classification_head = deepcopy(config)
        if num_labels_override is None:
            config_for_classification_head.num_labels = len(TargetType)
        else:
            config_for_classification_head.num_labels = num_labels_override
        self.num_classification_head_labels = config_for_classification_head.num_labels

        dropout_names = ["classifier_dropout", "hidden_dropout_prob", "classifier_dropout_prob"]
        for name in dropout_names:
            dropout_value = getattr(config_for_classification_head, name, None)
            if dropout_value is not None:
                self._logger.info(f"Loading dropout value {dropout_value} from config attribute '{name}'")
                config_for_classification_head.classifier_dropout = dropout_value
                break
        else:
            self._logger.warning("No dropout value found -- setting to 0")
            config_for_classification_head.classifier_dropout = 0.

        self.classifier = RobertaClassificationHead(config_for_classification_head)

    def forward(self, model_outputs: Union[tuple, BaseModelOutputWithPoolingAndCrossAttentions],
                start_positions=None, end_positions=None, target_type=None):
        sequence_output = model_outputs[0]

        # Predict target answer type for the whole question answer pair
        answer_type_logits = self.classifier(sequence_output)

        # Predict start and end logits by token
        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 and target_type 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)
            if len(target_type.size()) > 1:
                target_type = target_type.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)

            # if self.max_att_distance:
            #     # attention diversity implementation uses reduction functions
            #     # so need float calls to prevent grad overflow
            #     start_logits = start_logits.float()
            #     end_logits = end_logits.float()
            #     answer_type_logits = answer_type_logits.float()

            loss_fct = torch.nn.CrossEntropyLoss(ignore_index=ignored_index)
            start_loss = loss_fct(start_logits, start_positions)
            end_loss = loss_fct(end_logits, end_positions)
            answer_type_loss = loss_fct(answer_type_logits, target_type)
            # if self.max_att_distance:
            #     total_loss = (start_loss + end_loss + answer_type_loss + (
            #             cosine_sim_lambda * avg_att_distance)) / 4
            # else:
            total_loss = (start_loss + end_loss + answer_type_loss) / 3
            # outputs = (total_loss,) + outputs

        # (loss), start_logits, end_logits, target_type_logits, (hidden_states), (attentions)
        return_dict = isinstance(model_outputs, ModelOutput)
        if not return_dict:
            output = (start_logits, end_logits, answer_type_logits) + model_outputs[2:]
            return ((total_loss,) + output) if total_loss is not None else output

        return ExtractiveQAModelOutput(
            loss=total_loss,
            start_logits=start_logits,
            end_logits=end_logits,
            target_type_logits=answer_type_logits,
            hidden_states=model_outputs.hidden_states,
            attentions=model_outputs.attentions,
        )


EXTRACTIVE_HEAD = dict(qa_head=ExtractiveQAHead)

oneqa/mrc/models/heads/abstract.py METASEP
import logging
from abc import ABCMeta, abstractmethod
from typing import Union

import torch
from transformers import PretrainedConfig
from transformers.file_utils import ModelOutput


class AbstractTaskHead(torch.nn.Module, metaclass=ABCMeta):
    def __init__(self, config: PretrainedConfig):
        super().__init__()
        self._logger = logging.getLogger(self.__class__.__name__)
        self.config = config

    @abstractmethod
    def forward(self, model_outputs: Union[tuple, ModelOutput], *args, **kwargs):
        raise NotImplementedError

oneqa/mrc/models/heads/__init__.py METASEP

setup.py METASEP
import os
import re
from itertools import chain

from setuptools import setup, find_packages

here = os.path.abspath(os.path.dirname(__file__))

# Get the long description from the README file
with open(os.path.join(here, 'README.md'), encoding='utf-8') as readme_file:
    long_description = readme_file.read()

# Get the version for the project
with open(os.path.join(here, 'VERSION')) as version_file:
    version = version_file.read().strip()

include_package_roots = ["oneqa"]  # only include these packages and their subpackages
include_packages = list(chain.from_iterable(map(lambda name: [name, f"{name}.*"], include_package_roots)))

keywords = [
    "NLP", "transformers", "QA", "question", "answering", "mrc", "rc", "machine", "reading", "comprehension",
    "IR", "information", "retrieval", "deep", "learning", "pytorch", "BERT", "RoBERTa"
]

authors = [
    "TODO"
]

_deps = {
    "bump2version": ["dev"],
    "datasets>=1.4.0": ["install"],
    "packaging": ["install"],
    "pytest": ["tests"],
    "pytest-cov": ["tests"],
    "pytest-mock": ["tests"],
    "pytest-rerunfailures": ["tests"],
    "torch>=1.8": ["install"],  # TODO: see if we can reduce to 1.7 or 1.6
    "tox": ["tests"],
    "transformers>=4.0.0": ["install"],
    "tqdm": ["install"],
}

extras_names = ["dev", "install", "tests"]
extras = {extra_name: [] for extra_name in extras_names}
for package_name, package_required_by in _deps.items():
    if not package_required_by:
        raise ValueError(f"Package '{package_name}' has no associated extras it is required for. "
                         f"Choose one or more of: {extras_names}")

    for extra_name in package_required_by:
        try:
            extras[extra_name].append(package_name)
        except KeyError as ex:
            raise KeyError(f"Extras name '{extra_name}' not in {extras_names}") from ex
extras["all"] = list(_deps)

install_requires = extras["install"]

setup(
    name="one-qa",
    version=version,
    author=", ".join(authors),
    author_email="TODO",
    description="State-of-the-art Question Answering",
    long_description=long_description,
    long_description_content_type="text/markdown",
    url="TODO",
    license="TODO",
    keywords=" ".join(keywords),
    packages=find_packages(".", include=include_packages),
    python_requires=">=3.7.0",
    install_requires=install_requires,
    extras_require=extras,
    classifiers=[
        "Intended Audience :: Developers",
        "Intended Audience :: Science/Research",
        "Programming Language :: Python :: 3",
        "TODO"
    ],
)

oneqa/mrc/types/target_type.py METASEP
from enum import IntEnum, unique


@unique
class TargetType(IntEnum):
    NO_ANSWER = 0
    SPAN_ANSWER = 1
    PASSAGE_ANSWER = 2
    YES = 3
    NO = 4

    @classmethod
    def from_bool_label(cls, label: str) -> 'TargetType':
        label = label.upper()
        if label == 'NONE':
            label = 'NO_ANSWER'
        return cls[label]

oneqa/mrc/types/subsample_type.py METASEP
from enum import IntEnum, unique


@unique
class SubsampleType(IntEnum):
    POSITIVE = 0
    NEGATIVE_HAS_ANSWER = 1
    NEGATIVE_NO_ANSWER = 2

oneqa/mrc/types/__init__.py METASEP

oneqa/mrc/trainers/default.py METASEP
import inspect
import logging
from typing import Optional

import datasets
import torch
from datasets import Dataset
from packaging import version
from torch.utils.data import DataLoader
from transformers import Trainer, is_datasets_available
from transformers.trainer_pt_utils import IterableDatasetShard

logger = logging.getLogger(__name__)


class MRCTrainer(Trainer):
    def _remove_unused_columns(self, dataset: "datasets.Dataset", description: Optional[str] = None):
        if not self.args.remove_unused_columns:
            return dataset
        if self._signature_columns is None:
            # Inspect model and task head forward signature to keep only the arguments it accepts.
            model_signature = inspect.signature(self.model.forward)
            task_head_signature = inspect.signature(self.model.task_head.forward)

            signature_columns = set(model_signature.parameters.keys())
            signature_columns |= task_head_signature.parameters.keys()
            signature_columns -= {'kwargs'}

            # Labels may be named label or label_ids, the default data collator handles that.
            signature_columns |= {"label", "label_ids"}

            self._signature_columns = list(signature_columns)

        columns = [k for k in self._signature_columns if k in dataset.column_names]
        ignored_columns = list(set(dataset.column_names) - set(self._signature_columns))
        if len(ignored_columns) > 0:
            dset_description = "" if description is None else f"in the {description} set "
            logger.info(
                f"The following columns {dset_description} don't have a corresponding argument in "
                f"`{self.model.__class__.__name__}.forward` and have been ignored: {', '.join(ignored_columns)}."
            )

        if version.parse(datasets.__version__) < version.parse("1.4.0"):
            dataset.set_format(
                type=dataset.format["type"], columns=columns, format_kwargs=dataset.format["format_kwargs"]
            )
            return dataset
        else:
            return dataset.remove_columns(ignored_columns)

    def get_train_dataloader(self) -> DataLoader:
        """
        Returns the training [`~torch.utils.data.DataLoader`].

        Will use no sampler if `self.train_dataset` does not implement `__len__`, a random sampler (adapted
        to distributed training if necessary) otherwise.

        Subclass and override this method if you want to inject some custom behavior.
        """
        if self.train_dataset is None:
            raise ValueError("Trainer: training requires a train_dataset.")

        train_dataset = self.train_dataset
        if is_datasets_available() and isinstance(train_dataset, datasets.Dataset):
            train_dataset = self._remove_unused_columns(train_dataset, description="training")

        if isinstance(train_dataset, torch.utils.data.IterableDataset):
            if self.args.world_size > 1:
                train_dataset = IterableDatasetShard(
                    train_dataset,
                    batch_size=self.args.train_batch_size,
                    drop_last=self.args.dataloader_drop_last,
                    num_processes=self.args.world_size,
                    process_index=self.args.process_index,
                )

            return DataLoader(
                train_dataset,
                batch_size=self.args.per_device_train_batch_size,
                collate_fn=self.data_collator,
                num_workers=self.args.dataloader_num_workers,
                pin_memory=self.args.dataloader_pin_memory,
            )

        train_sampler = self._get_train_sampler()

        return DataLoader(
            train_dataset,
            batch_size=self.args.train_batch_size,
            sampler=train_sampler,
            collate_fn=self.data_collator,
            drop_last=self.args.dataloader_drop_last,
            num_workers=self.args.dataloader_num_workers,
            pin_memory=self.args.dataloader_pin_memory,
        )

    def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoader:
        """
        Returns the evaluation [`~torch.utils.data.DataLoader`].

        Subclass and override this method if you want to inject some custom behavior.

        Args:
            eval_dataset (`torch.utils.data.Dataset`, *optional*):
                If provided, will override `self.eval_dataset`. If it is an `datasets.Dataset`, columns not
                accepted by the `model.forward()` method are automatically removed. It must implement `__len__`.
        """
        if eval_dataset is None and self.eval_dataset is None:
            raise ValueError("Trainer: evaluation requires an eval_dataset.")
        eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset

        if is_datasets_available() and isinstance(eval_dataset, datasets.Dataset):
            eval_dataset = self._remove_unused_columns(eval_dataset, description="evaluation")

        if isinstance(eval_dataset, torch.utils.data.IterableDataset):
            if self.args.world_size > 1:
                eval_dataset = IterableDatasetShard(
                    eval_dataset,
                    batch_size=self.args.per_device_eval_batch_size,
                    drop_last=self.args.dataloader_drop_last,
                    num_processes=self.args.world_size,
                    process_index=self.args.process_index,
                )
            return DataLoader(
                eval_dataset,
                batch_size=self.args.eval_batch_size,
                collate_fn=self.data_collator,
                num_workers=self.args.dataloader_num_workers,
                pin_memory=self.args.dataloader_pin_memory,
            )

        eval_sampler = self._get_eval_sampler(eval_dataset)

        return DataLoader(
            eval_dataset,
            sampler=eval_sampler,
            batch_size=self.args.eval_batch_size,
            collate_fn=self.data_collator,
            drop_last=self.args.dataloader_drop_last,
            num_workers=self.args.dataloader_num_workers,
            pin_memory=self.args.dataloader_pin_memory,
        )

    # TODO
    # def get_test_dataloader(self, test_dataset: Dataset) -> DataLoader:
    #     """
    #     Returns the test [`~torch.utils.data.DataLoader`].
    #
    #     Subclass and override this method if you want to inject some custom behavior.
    #
    #     Args:
    #         test_dataset (`torch.utils.data.Dataset`, *optional*):
    #             The test dataset to use. If it is an `datasets.Dataset`, columns not accepted by the
    #             `model.forward()` method are automatically removed. It must implement `__len__`.
    #     """
    #     if is_datasets_available() and isinstance(test_dataset, datasets.Dataset):
    #         test_dataset = self._remove_unused_columns(test_dataset, description="test")
    #
    #     if isinstance(test_dataset, torch.utils.data.IterableDataset):
    #         if self.args.world_size > 1:
    #             test_dataset = IterableDatasetShard(
    #                 test_dataset,
    #                 batch_size=self.args.eval_batch_size,
    #                 drop_last=self.args.dataloader_drop_last,
    #                 num_processes=self.args.world_size,
    #                 process_index=self.args.process_index,
    #             )
    #         return DataLoader(
    #             test_dataset,
    #             batch_size=self.args.eval_batch_size,
    #             collate_fn=self.data_collator,
    #             num_workers=self.args.dataloader_num_workers,
    #             pin_memory=self.args.dataloader_pin_memory,
    #         )
    #
    #     test_sampler = self._get_eval_sampler(test_dataset)
    #
    #     # We use the same batch_size as for eval.
    #     return DataLoader(
    #         test_dataset,
    #         sampler=test_sampler,
    #         batch_size=self.args.eval_batch_size,
    #         collate_fn=self.data_collator,
    #         drop_last=self.args.dataloader_drop_last,
    #         pin_memory=self.args.dataloader_pin_memory,
    #     )

oneqa/mrc/trainers/__init__.py METASEP

oneqa/mrc/processors/__init__.py METASEP

oneqa/mrc/models/task_model.py METASEP
import logging
from typing import Dict, Type

import torch
from transformers import PretrainedConfig, PreTrainedModel, MODEL_FOR_PRETRAINING_MAPPING, MODEL_MAPPING

from oneqa.mrc.models.heads.abstract import AbstractTaskHead


# def create_task_model_class_from_config(config: PretrainedConfig) -> Type['ModelForDownstreamTasks']:
#     ptm_base_class = MODEL_FOR_PRETRAINING_MAPPING[config.__class__]
#     # inner_model_class = MODEL_MAPPING[config]
#     # base_model_prefix = getattr(ptm_base_class, 'base_model_prefix', config.model_type)
#     # model_type = config.model_type
#     # model_name = ''.join(map(str.title, re.split(r'[^\w\d]', model_type)))
#     model_name = config.__class__.__name__.rstrip('Config')
#     class_name = f'{model_name}{ModelForDownstreamTasks.__name__}'
#     model_class = type(class_name, (ModelForDownstreamTasks, ptm_base_class), {})
#     return model_class


class ModelForDownstreamTasks(PreTrainedModel):

    def __init__(self,
                 config: PretrainedConfig,
                 task_heads: Dict[str, Type[AbstractTaskHead]]):
        super().__init__(config)
        self._logger = logging.getLogger(self.__class__.__name__)

        if type(self) is ModelForDownstreamTasks:
            raise TypeError(f"{ModelForDownstreamTasks.__class__.__name__} is not intended to be directly "
                            f"instantiated and should be subclassed together with a XPreTrainedModel type. "
                            f"See {self.model_class_from_config.__name__} or {self.from_config.__name__} "
                            f"for creating and instantiating these subclasses.")

        if not task_heads:
            raise ValueError("No task heads provided")

        self._task_head = None

        # Set the model to match the pre-trained name (e.g. self.roberta) so it can be loaded from pretrained
        setattr(self, self.base_model_prefix, MODEL_MAPPING[config.__class__](config))

        self.task_heads = torch.nn.ModuleDict(
            {name: model(config) for name, model in task_heads.items()})
        self.init_weights()

    @property
    def model_(self) -> PreTrainedModel:  # using 'model' instead of 'model_' causes conflicts with some LMs (e.g. BART)
        """
        Returns the underlying language model. This is an alias to simplify access.
        """
        return getattr(self, self.base_model_prefix)

    @property
    def task_head(self) -> AbstractTaskHead:
        """
        Returns the current task head or raises a ValueError if it has not yet been set
        """
        if self._task_head is not None:
            # noinspection PyTypeChecker
            return self.task_heads[self._task_head]
        else:
            raise ValueError(f"Task head is not set.  Call {ModelForDownstreamTasks.set_task_head.__name__} to set it")

    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,
                **kwargs):
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.model_(
            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,
        )

        return self.task_head(outputs, **kwargs)

    @classmethod
    def model_class_from_config(cls, config: PretrainedConfig) -> Type['ModelForDownstreamTasks']:
        """
        Dynamically creates a model class from a PreTrainedConfig.
        """
        ptm_base_class = MODEL_FOR_PRETRAINING_MAPPING[config.__class__]
        model_name = config.__class__.__name__.rstrip('Config')
        class_name = f'{model_name}{cls.__name__}'
        model_class = type(class_name, (cls, ptm_base_class), {})

        # noinspection PyTypeChecker
        return model_class

    @classmethod
    def from_config(cls, config: PretrainedConfig, *args, **kwargs) -> 'ModelForDownstreamTasks':
        """
        Dynamically creates a model class from a PreTrainedConfig and then uses the config
        with args/kwargs to instantiate.
        """
        model_class = cls.model_class_from_config(config)
        model = model_class.from_pretrained(*args, config=config, **kwargs)
        return model

    def set_task_head(self, task_head: str):
        if task_head not in self.task_heads:
            raise KeyError(f"Task head '{task_head}' not in task_heads: {list(self.task_heads)}")
        elif self._task_head is not None:
            self._logger.info(f"Changing default task head from '{self._task_head}' to '{task_head}'")
        else:
            self._logger.info(f"Setting task head for first time to '{self._task_head}'")
        self._task_head = task_head

oneqa/mrc/models/__init__.py METASEP
"""
BEGIN_COPYRIGHT

IBM Confidential
OCO Source Materials

5727-I17
(C) Copyright IBM Corp. 2022 All Rights Reserved.
 
The source code for this program is not published or otherwise
divested of its trade secrets, irrespective of what has been
deposited with the U.S. Copyright Office.

END_COPYRIGHT
"""

oneqa/__init__.py METASEP

oneqa/mrc/main.py METASEP
import logging
import os

import datasets
from transformers import TrainingArguments, DataCollatorWithPadding, AutoConfig, AutoTokenizer
from transformers.trainer_utils import get_last_checkpoint, set_seed

from oneqa.mrc.models.task_model import ModelForDownstreamTasks
from oneqa.mrc.processors.postprocessors.extractive import ExtractivePostProcessor
from oneqa.mrc.processors.preprocessors.default import DefaultPreProcessor
from oneqa.mrc.processors.preprocessors.tydiqa import TyDiQAPreprocessor
from oneqa.mrc.trainers.default import MRCTrainer
from oneqa.mrc.models.heads.extractive import EXTRACTIVE_HEAD


def main():
    logger = logging.getLogger(__name__)
    training_args = TrainingArguments(
        output_dir='/dccstor/aferritt3/oneqa/test-model',
        do_train=True,
        do_eval=True,
        num_train_epochs=0.1,
        fp16=False,
    )

    set_seed(training_args.seed)

    # 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."
            )

    model_name = 'xlm-roberta-base'
    task_heads = EXTRACTIVE_HEAD  # TODO parameterize
    config = AutoConfig.from_pretrained(
        model_name,
        # 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_name,
        # 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,
        config=config,
    )
    model = ModelForDownstreamTasks.from_config(
        config,
        model_name,
        task_heads=task_heads,
    )
    model.set_task_head(next(iter(task_heads)))

    # load data
    logger.info('Loading dataset')
    raw_datasets = datasets.load_dataset("tydiqa", "primary_task")

    # load preprocessor
    preprocessor_class = TyDiQAPreprocessor  # TODO parameterize
    preprocessor = preprocessor_class(
        stride=128,
        tokenizer=tokenizer,
    )

    # process train data
    train_dataset = raw_datasets["train"]
    max_train_samples = 1000
    if max_train_samples is not None:  # if data_args.max_train_samples is not None:
        # We will select sample from whole data if argument is specified
        train_dataset = train_dataset.select(range(max_train_samples))
    with training_args.main_process_first(desc="train dataset map pre-processing"):
        # train_dataset = preprocessor.adapt_dataset(train_dataset)
        # train_dataset = train_dataset.map(  # TODO debug
        #     preprocessor.process_train,
        #     batched=True,
        #     num_proc=1,  # data_args.preprocessing_num_workers,
        #     remove_columns=train_dataset.column_names,
        #     # load_from_cache_file=not data_args.overwrite_cache,
        #     load_from_cache_file=False,
        #     desc="Running tokenizer on train dataset",
        # )
        # train_dataset = preprocessor.subsample_features(train_dataset)
        _, train_dataset = preprocessor.process_train(train_dataset)

    # process val data
    eval_examples = raw_datasets["validation"]
    max_eval_samples = 250
    if max_eval_samples is not None:  # data_args.max_eval_samples is not None:
        # We will select sample from whole data
        eval_examples = eval_examples.select(range(max_eval_samples))
    # Validation Feature Creation
    with training_args.main_process_first(desc="validation dataset map pre-processing"):
        # eval_examples = preprocessor.adapt_dataset(eval_examples)
        # eval_dataset = eval_examples.map(
        #     preprocessor.process_eval,
        #     batched=True,
        #     num_proc=1,  # data_args.preprocessing_num_workers,
        #     remove_columns=eval_examples.column_names,
        #     # load_from_cache_file=not data_args.overwrite_cache,
        #     load_from_cache_file=False,
        #     desc="Running tokenizer on training dataset",
        # )
        eval_examples, eval_dataset = preprocessor.process_eval(eval_examples)

    # process test data

    data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=64 if training_args.fp16 else None)

    # train

    postprocessor_class = ExtractivePostProcessor  # TODO parameterize
    postprocessor = postprocessor_class(k=5, n_best_size=20, max_answer_length=30, span_tracker_factory=None)

    metrics_fn = None  # TODO metrics

    trainer = MRCTrainer(
        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=postprocessor.process,
        compute_metrics=metrics_fn,
    )

    checkpoint = None
    if training_args.resume_from_checkpoint is not None:
        checkpoint = training_args.resume_from_checkpoint
    elif last_checkpoint is not None:
        checkpoint = last_checkpoint
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
    # raise ValueError("Nothing implemented beyond this point")
    trainer.save_model()  # Saves the tokenizer too for easy upload

    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))

    trainer.log_metrics("train", metrics)
    trainer.save_metrics("train", metrics)
    trainer.save_state()

    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))

    trainer.log_metrics("eval", metrics)
    trainer.save_metrics("eval", metrics)

    # run val

    # run test


if __name__ == '__main__':
    main()

oneqa/mrc/__init__.py METASEP

oneqa/mrc/metrics/nq_f1/eval_utils.py METASEP

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import glob
import json
import logging
import multiprocessing
from gzip import GzipFile

from functools import partial
from typing import Dict, Set, Optional, List


@functools.total_ordering
class NQSpan(object):
    """A class for handling token and byte spans.
       Taken from https://github.com/google-research-datasets/natural-questions/blob/master/eval_utils.py#L60
      The logic is:
      1) if both start_byte !=  -1 and end_byte != -1 then the span is defined
         by byte offsets
      2) else, if start_token != -1 and end_token != -1 then the span is define
         by token offsets
      3) else, this is a null span.
      Null spans means that there is no (long or short) answers.
      If your systems only care about token spans rather than byte spans, set all
      byte spans to -1.
    """
    __slots__ = ['start_byte', 'end_byte', 'start_token', 'end_token', 'score', 'long_score']

    def __init__(self, start_byte: int, end_byte: int, start_token: int, end_token: int,
                 score: Optional[float] = None, long_score: Optional[float] = None, enforce_byte_consistency: bool = True,
                 enforce_token_consistency: bool = True):

        self.start_byte = start_byte
        self.end_byte = end_byte
        self.start_token = start_token
        self.end_token = end_token
        self.score = score
        self.long_score = long_score

        if enforce_byte_consistency:
            if ((start_byte < 0 <= end_byte) or
                    (start_byte >= 0 > end_byte)):
                raise ValueError('Inconsistent Null Spans (Byte): %s' % self)

            if start_byte >= 0 and 0 <= end_byte < start_byte:
                raise ValueError('Invalid byte spans (start_byte > end_byte): %s' % self)

        if enforce_token_consistency:
            if ((start_token < 0 <= end_token) or
                    (start_token >= 0 > end_token)):
                raise ValueError('Inconsistent Null Spans (Token): %s' % self)

            if ((start_token >= 0 and end_token >= 0) and
                    (start_token > end_token)):
                raise ValueError(
                    'Invalid token spans (start_token_idx > end_token_idx): %s' % self)

    def is_null_span(self):
        """A span is a null span if the start and end are both -1."""

        if (self.start_byte < 0 and self.end_byte < 0 and
                    self.start_token < 0 and self.end_token < 0):
            return True
        return False

    @staticmethod
    def null_span():
        return NQSpan(start_byte=-1, end_byte=-1, start_token=-1, end_token=-1)

    def contains(self, other_span):

        if other_span.start_byte > -1 and self.start_byte > -1 and other_span.end_byte > -1 and self.end_byte > -1:
            # Compare using bytes
            if other_span.start_byte >= self.start_byte and other_span.end_byte <= self.end_byte:
                return True
            else:
                return False
        elif other_span.start_token > -1 and self.start_token > -1 and other_span.end_token > -1 and self.end_token > -1:
            # Compare using tokens
            if other_span.start_token >= self.start_token and other_span.end_token <= self.end_token:
                return True
            else:
                return False
        else:
            return False

    def __str__(self):
        byte_str = 'bytes: [' + str(self.start_byte) + ',' + str(self.end_byte) + ')'
        tok_str = ('tokens: [' + str(self.start_token) + ',' + str(
            self.end_token) + ')')
        if self.score is not None:
            score_str = ', score: %s' % self.score
        else:
            score_str = ''

        return "Span(" + byte_str + ', ' + tok_str + score_str + ")"

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

    def __hash__(self):
        return hash((self.start_byte, self.end_byte, self.start_token, self.end_token, self.score))

    def __lt__(self, other):
        if self.score < other.score:
            return True
        return False

    def __eq__(self, o: object) -> bool:

        if not isinstance(o, NQSpan):
            return False

        for attribute in self.__slots__:
            if not hasattr(o, attribute):
                return False
            elif getattr(o, attribute) != getattr(self, attribute):
                return False

        return True


# A data structure for storing prediction and annotation.
# When a example has multiple annotations, multiple NQLabel will be used.
NQLabel = collections.namedtuple(
    'NQLabel',
    [
        'example_id',  # the unique id for each NQ example.
        'long_answer_span',  # A Span object for long answer.
        'short_answer_span_list',  # A list of Spans for short answer.
        #   Note that In NQ, the short answers
        #   do not need to be in a single span.
        'yes_no_answer',  # Indicate if the short answer is an yes/no answer
        #   The possible values are "yes", "no", "none".
        #   (case insensitive)
        #   If the field is "yes", short_answer_span_list
        #   should be empty or only contain null spans.
        'long_score',  # The prediction score for the long answer prediction.
        'short_score'  # The prediction score for the short answer prediction.
    ])


def is_null_span_list(span_list):
    """Returns true iff all spans in span_list are null or span_list is empty."""
    if not span_list or all([span.is_null_span() for span in span_list]):
        return True
    return False


def nonnull_span_equal(span_a: NQSpan, span_b: NQSpan):
    """Given two spans, return if they are equal.

    Args:
      span_a: a Span object.
      span_b: a Span object.  Only compare non-null spans. First, if the bytes are
        not negative, compare byte offsets, Otherwise, compare token offsets.

    Returns:
      True or False
    """
    assert isinstance(span_a, NQSpan)
    assert isinstance(span_b, NQSpan)
    assert not span_a.is_null_span()
    assert not span_b.is_null_span()

    # if byte offsets are not negative, compare byte offsets
    if ((span_a.start_byte >= 0 and span_a.end_byte >= 0) and
            (span_b.start_byte >= 0 and span_b.end_byte >= 0)):

        if ((span_a.start_byte == span_b.start_byte) and
                (span_a.end_byte == span_b.end_byte)):
            return True

    # if token offsets are not negative, compare token offsets
    if ((span_a.start_token >= 0 and span_a.end_token >= 0) and
            (span_b.start_token >= 0 and span_b.end_token >= 0)):

        if ((span_a.start_token == span_b.start_token) and
                (span_a.end_token == span_b.end_token)):
            return True

    return False


def span_set_equal(gold_span_list, pred_span_list):
    """Make the spans are completely equal besides null spans."""

    gold_span_list = [span for span in gold_span_list if not span.is_null_span()]
infile
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import glob
import json
import logging
import multiprocessing
from gzip import GzipFile

from functools import partial
from typing import Dict, Set, Optional, List


@functools.total_ordering
class NQSpan(object):
    """A class for handling token and byte spans.
       Taken from https://github.com/google-research-datasets/natural-questions/blob/master/eval_utils.py#L60
      The logic is:
      1) if both start_byte !=  -1 and end_byte != -1 then the span is defined
         by byte offsets
      2) else, if start_token != -1 and end_token != -1 then the span is define
         by token offsets
      3) else, this is a null span.
      Null spans means that there is no (long or short) answers.
      If your systems only care about token spans rather than byte spans, set all
      byte spans to -1.
    """
    __slots__ = ['start_byte', 'end_byte', 'start_token', 'end_token', 'score', 'long_score']

    def __init__(self, start_byte: int, end_byte: int, start_token: int, end_token: int,
                 score: Optional[float] = None, long_score: Optional[float] = None, enforce_byte_consistency: bool = True,
                 enforce_token_consistency: bool = True):

        self.start_byte = start_byte
        self.end_byte = end_byte
        self.start_token = start_token
        self.end_token = end_token
        self.score = score
        self.long_score = long_score

        if enforce_byte_consistency:
            if ((start_byte < 0 <= end_byte) or
                    (start_byte >= 0 > end_byte)):
                raise ValueError('Inconsistent Null Spans (Byte): %s' % self)

            if start_byte >= 0 and 0 <= end_byte < start_byte:
                raise ValueError('Invalid byte spans (start_byte > end_byte): %s' % self)

        if enforce_token_consistency:
            if ((start_token < 0 <= end_token) or
                    (start_token >= 0 > end_token)):
                raise ValueError('Inconsistent Null Spans (Token): %s' % self)

            if ((start_token >= 0 and end_token >= 0) and
                    (start_token > end_token)):
                raise ValueError(
                    'Invalid token spans (start_token_idx > end_token_idx): %s' % self)

    def is_null_span(self):
        """A span is a null span if the start and end are both -1."""

        if (self.start_byte < 0 and self.end_byte < 0 and
                    self.start_token < 0 and self.end_token < 0):
            return True
        return False

    @staticmethod
    def null_span():
        return NQSpan(start_byte=-1, end_byte=-1, start_token=-1, end_token=-1)

    def contains(self, other_span):

        if other_span.start_byte > -1 and self.start_byte > -1 and other_span.end_byte > -1 and self.end_byte > -1:
            # Compare using bytes
            if other_span.start_byte >= self.start_byte and other_span.end_byte <= self.end_byte:
                return True
            else:
                return False
        elif other_span.start_token > -1 and self.start_token > -1 and other_span.end_token > -1 and self.end_token > -1:
            # Compare using tokens
            if other_span.start_token >= self.start_token and other_span.end_token <= self.end_token:
                return True
            else:
                return False
        else:
            return False

    def __str__(self):
        byte_str = 'bytes: [' + str(self.start_byte) + ',' + str(self.end_byte) + ')'
        tok_str = ('tokens: [' + str(self.start_token) + ',' + str(
            self.end_token) + ')')
        if self.score is not None:
            score_str = ', score: %s' % self.score
        else:
            score_str = ''

        return "Span(" + byte_str + ', ' + tok_str + score_str + ")"

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

    def __hash__(self):
        return hash((self.start_byte, self.end_byte, self.start_token, self.end_token, self.score))

    def __lt__(self, other):
        if self.score < other.score:
            return True
        return False

    def __eq__(self, o: object) -> bool:

        if not isinstance(o, NQSpan):
            return False

        for attribute in self.__slots__:
            if not hasattr(o, attribute):
                return False
            elif getattr(o, attribute) != getattr(self, attribute):
                return False

        return True


# A data structure for storing prediction and annotation.
# When a example has multiple annotations, multiple NQLabel will be used.
NQLabel = collections.namedtuple(
    'NQLabel',
    [
        'example_id',  # the unique id for each NQ example.
        'long_answer_span',  # A Span object for long answer.
        'short_answer_span_list',  # A list of Spans for short answer.
        #   Note that In NQ, the short answers
        #   do not need to be in a single span.
        'yes_no_answer',  # Indicate if the short answer is an yes/no answer
        #   The possible values are "yes", "no", "none".
        #   (case insensitive)
        #   If the field is "yes", short_answer_span_list
        #   should be empty or only contain null spans.
        'long_score',  # The prediction score for the long answer prediction.
        'short_score'  # The prediction score for the short answer prediction.
    ])


def is_null_span_list(span_list):
    """Returns true iff all spans in span_list are null or span_list is empty."""
    if not span_list or all([span.is_null_span() for span in span_list]):
        return True
    return False


def nonnull_span_equal(span_a: NQSpan, span_b: NQSpan):
    """Given two spans, return if they are equal.

    Args:
      span_a: a Span object.
      span_b: a Span object.  Only compare non-null spans. First, if the bytes are
        not negative, compare byte offsets, Otherwise, compare token offsets.

    Returns:
      True or False
    """
infile
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import glob
import json
import logging
import multiprocessing
from gzip import GzipFile

from functools import partial
from typing import Dict, Set, Optional, List


@functools.total_ordering
class NQSpan(object):
    """A class for handling token and byte spans.
       Taken from https://github.com/google-research-datasets/natural-questions/blob/master/eval_utils.py#L60
      The logic is:
      1) if both start_byte !=  -1 and end_byte != -1 then the span is defined
         by byte offsets
      2) else, if start_token != -1 and end_token != -1 then the span is define
         by token offsets
      3) else, this is a null span.
      Null spans means that there is no (long or short) answers.
      If your systems only care about token spans rather than byte spans, set all
      byte spans to -1.
    """
    __slots__ = ['start_byte', 'end_byte', 'start_token', 'end_token', 'score', 'long_score']

    def __init__(self, start_byte: int, end_byte: int, start_token: int, end_token: int,
                 score: Optional[float] = None, long_score: Optional[float] = None, enforce_byte_consistency: bool = True,
                 enforce_token_consistency: bool = True):

        self.start_byte = start_byte
        self.end_byte = end_byte
        self.start_token = start_token
        self.end_token = end_token
        self.score = score
        self.long_score = long_score

        if enforce_byte_consistency:
            if ((start_byte < 0 <= end_byte) or
                    (start_byte >= 0 > end_byte)):
                raise ValueError('Inconsistent Null Spans (Byte): %s' % self)

            if start_byte >= 0 and 0 <= end_byte < start_byte:
                raise ValueError('Invalid byte spans (start_byte > end_byte): %s' % self)

        if enforce_token_consistency:
            if ((start_token < 0 <= end_token) or
                    (start_token >= 0 > end_token)):
                raise ValueError('Inconsistent Null Spans (Token): %s' % self)

            if ((start_token >= 0 and end_token >= 0) and
                    (start_token > end_token)):
                raise ValueError(
                    'Invalid token spans (start_token_idx > end_token_idx): %s' % self)

    def is_null_span(self):
        """A span is a null span if the start and end are both -1."""

        if (self.start_byte < 0 and self.end_byte < 0 and
                    self.start_token < 0 and self.end_token < 0):
            return True
        return False

    @staticmethod
    def null_span():
        return NQSpan(start_byte=-1, end_byte=-1, start_token=-1, end_token=-1)

    def contains(self, other_span):

        if other_span.start_byte > -1 and self.start_byte > -1 and other_span.end_byte > -1 and self.end_byte > -1:
            # Compare using bytes
            if other_span.start_byte >= self.start_byte and other_span.end_byte <= self.end_byte:
                return True
            else:
                return False
        elif other_span.start_token > -1 and self.start_token > -1 and other_span.end_token > -1 and self.end_token > -1:
            # Compare using tokens
            if other_span.start_token >= self.start_token and other_span.end_token <= self.end_token:
                return True
            else:
                return False
        else:
            return False

    def __str__(self):
        byte_str = 'bytes: [' + str(self.start_byte) + ',' + str(self.end_byte) + ')'
        tok_str = ('tokens: [' + str(self.start_token) + ',' + str(
            self.end_token) + ')')
        if self.score is not None:
            score_str = ', score: %s' % self.score
        else:
            score_str = ''

        return "Span(" + byte_str + ', ' + tok_str + score_str + ")"

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

    def __hash__(self):
        return hash((self.start_byte, self.end_byte, self.start_token, self.end_token, self.score))

    def __lt__(self, other):
        if self.score < other.score:
            return True
        return False

    def __eq__(self, o: object) -> bool:

        if not isinstance(o, NQSpan):
            return False

        for attribute in self.__slots__:
            if not hasattr(o, attribute):
                return False
            elif getattr(o, attribute) != getattr(self, attribute):
                return False

        return True


# A data structure for storing prediction and annotation.
# When a example has multiple annotations, multiple NQLabel will be used.
NQLabel = collections.namedtuple(
    'NQLabel',
    [
        'example_id',  # the unique id for each NQ example.
        'long_answer_span',  # A Span object for long answer.
        'short_answer_span_list',  # A list of Spans for short answer.
        #   Note that In NQ, the short answers
        #   do not need to be in a single span.
        'yes_no_answer',  # Indicate if the short answer is an yes/no answer
        #   The possible values are "yes", "no", "none".
        #   (case insensitive)
        #   If the field is "yes", short_answer_span_list
        #   should be empty or only contain null spans.
        'long_score',  # The prediction score for the long answer prediction.
        'short_score'  # The prediction score for the short answer prediction.
    ])


def is_null_span_list(span_list):
    """Returns true iff all spans in span_list are null or span_list is empty."""
    if not span_list or all([span.is_null_span() for span in span_list]):
        return True
    return False


def nonnull_span_equal(span_a: NQSpan, span_b: NQSpan):
    """Given two spans, return if they are equal.

    Args:
      span_a: a Span object.
      span_b: a Span object.  Only compare non-null spans. First, if the bytes are
        not negative, compare byte offsets, Otherwise, compare token offsets.

    Returns:
      True or False
    """
    assert isinstance(span_a, NQSpan)
infile
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import glob
import json
import logging
import multiprocessing
from gzip import GzipFile

from functools import partial
from typing import Dict, Set, Optional, List


@functools.total_ordering
class NQSpan(object):
    """A class for handling token and byte spans.
       Taken from https://github.com/google-research-datasets/natural-questions/blob/master/eval_utils.py#L60
      The logic is:
      1) if both start_byte !=  -1 and end_byte != -1 then the span is defined
         by byte offsets
      2) else, if start_token != -1 and end_token != -1 then the span is define
         by token offsets
      3) else, this is a null span.
      Null spans means that there is no (long or short) answers.
      If your systems only care about token spans rather than byte spans, set all
      byte spans to -1.
    """
    __slots__ = ['start_byte', 'end_byte', 'start_token', 'end_token', 'score', 'long_score']

    def __init__(self, start_byte: int, end_byte: int, start_token: int, end_token: int,
                 score: Optional[float] = None, long_score: Optional[float] = None, enforce_byte_consistency: bool = True,
                 enforce_token_consistency: bool = True):

        self.start_byte = start_byte
        self.end_byte = end_byte
        self.start_token = start_token
        self.end_token = end_token
        self.score = score
        self.long_score = long_score

        if enforce_byte_consistency:
            if ((start_byte < 0 <= end_byte) or
                    (start_byte >= 0 > end_byte)):
                raise ValueError('Inconsistent Null Spans (Byte): %s' % self)

            if start_byte >= 0 and 0 <= end_byte < start_byte:
                raise ValueError('Invalid byte spans (start_byte > end_byte): %s' % self)

        if enforce_token_consistency:
            if ((start_token < 0 <= end_token) or
                    (start_token >= 0 > end_token)):
                raise ValueError('Inconsistent Null Spans (Token): %s' % self)

            if ((start_token >= 0 and end_token >= 0) and
                    (start_token > end_token)):
                raise ValueError(
                    'Invalid token spans (start_token_idx > end_token_idx): %s' % self)

    def is_null_span(self):
        """A span is a null span if the start and end are both -1."""

        if (self.start_byte < 0 and self.end_byte < 0 and
                    self.start_token < 0 and self.end_token < 0):
            return True
        return False

    @staticmethod
    def null_span():
        return NQSpan(start_byte=-1, end_byte=-1, start_token=-1, end_token=-1)

    def contains(self, other_span):

        if other_span.start_byte > -1 and self.start_byte > -1 and other_span.end_byte > -1 and self.end_byte > -1:
            # Compare using bytes
            if other_span.start_byte >= self.start_byte and other_span.end_byte <= self.end_byte:
                return True
            else:
                return False
        elif other_span.start_token > -1 and self.start_token > -1 and other_span.end_token > -1 and self.end_token > -1:
            # Compare using tokens
            if other_span.start_token >= self.start_token and other_span.end_token <= self.end_token:
                return True
            else:
                return False
        else:
            return False

    def __str__(self):
        byte_str = 'bytes: [' + str(self.start_byte) + ',' + str(self.end_byte) + ')'
        tok_str = ('tokens: [' + str(self.start_token) + ',' + str(
            self.end_token) + ')')
        if self.score is not None:
            score_str = ', score: %s' % self.score
        else:
            score_str = ''

        return "Span(" + byte_str + ', ' + tok_str + score_str + ")"

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

    def __hash__(self):
        return hash((self.start_byte, self.end_byte, self.start_token, self.end_token, self.score))

    def __lt__(self, other):
        if self.score < other.score:
            return True
        return False

    def __eq__(self, o: object) -> bool:

        if not isinstance(o, NQSpan):
            return False

        for attribute in self.__slots__:
            if not hasattr(o, attribute):
                return False
            elif getattr(o, attribute) != getattr(self, attribute):
                return False

        return True


# A data structure for storing prediction and annotation.
# When a example has multiple annotations, multiple NQLabel will be used.
NQLabel = collections.namedtuple(
    'NQLabel',
    [
        'example_id',  # the unique id for each NQ example.
        'long_answer_span',  # A Span object for long answer.
        'short_answer_span_list',  # A list of Spans for short answer.
        #   Note that In NQ, the short answers
        #   do not need to be in a single span.
        'yes_no_answer',  # Indicate if the short answer is an yes/no answer
        #   The possible values are "yes", "no", "none".
        #   (case insensitive)
        #   If the field is "yes", short_answer_span_list
        #   should be empty or only contain null spans.
        'long_score',  # The prediction score for the long answer prediction.
        'short_score'  # The prediction score for the short answer prediction.
    ])


def is_null_span_list(span_list):
    """Returns true iff all spans in span_list are null or span_list is empty."""
    if not span_list or all([span.is_null_span() for span in span_list]):
        return True
    return False


def nonnull_span_equal(span_a: NQSpan, span_b: NQSpan):
    """Given two spans, return if they are equal.

    Args:
      span_a: a Span object.
      span_b: a Span object.  Only compare non-null spans. First, if the bytes are
        not negative, compare byte offsets, Otherwise, compare token offsets.

    Returns:
      True or False
    """
    assert isinstance(span_a, NQSpan)
    assert isinstance(span_b, NQSpan)
infile
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import glob
import json
import logging
import multiprocessing
from gzip import GzipFile

from functools import partial
from typing import Dict, Set, Optional, List


@functools.total_ordering
class NQSpan(object):
    """A class for handling token and byte spans.
       Taken from https://github.com/google-research-datasets/natural-questions/blob/master/eval_utils.py#L60
      The logic is:
      1) if both start_byte !=  -1 and end_byte != -1 then the span is defined
         by byte offsets
      2) else, if start_token != -1 and end_token != -1 then the span is define
         by token offsets
      3) else, this is a null span.
      Null spans means that there is no (long or short) answers.
      If your systems only care about token spans rather than byte spans, set all
      byte spans to -1.
    """
    __slots__ = ['start_byte', 'end_byte', 'start_token', 'end_token', 'score', 'long_score']

    def __init__(self, start_byte: int, end_byte: int, start_token: int, end_token: int,
                 score: Optional[float] = None, long_score: Optional[float] = None, enforce_byte_consistency: bool = True,
                 enforce_token_consistency: bool = True):

        self.start_byte = start_byte
        self.end_byte = end_byte
        self.start_token = start_token
        self.end_token = end_token
        self.score = score
        self.long_score = long_score

        if enforce_byte_consistency:
            if ((start_byte < 0 <= end_byte) or
                    (start_byte >= 0 > end_byte)):
                raise ValueError('Inconsistent Null Spans (Byte): %s' % self)

            if start_byte >= 0 and 0 <= end_byte < start_byte:
                raise ValueError('Invalid byte spans (start_byte > end_byte): %s' % self)

        if enforce_token_consistency:
            if ((start_token < 0 <= end_token) or
                    (start_token >= 0 > end_token)):
                raise ValueError('Inconsistent Null Spans (Token): %s' % self)

            if ((start_token >= 0 and end_token >= 0) and
                    (start_token > end_token)):
                raise ValueError(
                    'Invalid token spans (start_token_idx > end_token_idx): %s' % self)

    def is_null_span(self):
        """A span is a null span if the start and end are both -1."""

        if (self.start_byte < 0 and self.end_byte < 0 and
                    self.start_token < 0 and self.end_token < 0):
            return True
        return False

    @staticmethod
    def null_span():
        return NQSpan(start_byte=-1, end_byte=-1, start_token=-1, end_token=-1)

    def contains(self, other_span):

        if other_span.start_byte > -1 and self.start_byte > -1 and other_span.end_byte > -1 and self.end_byte > -1:
            # Compare using bytes
            if other_span.start_byte >= self.start_byte and other_span.end_byte <= self.end_byte:
                return True
            else:
                return False
        elif other_span.start_token > -1 and self.start_token > -1 and other_span.end_token > -1 and self.end_token > -1:
            # Compare using tokens
            if other_span.start_token >= self.start_token and other_span.end_token <= self.end_token:
                return True
            else:
                return False
        else:
            return False

    def __str__(self):
        byte_str = 'bytes: [' + str(self.start_byte) + ',' + str(self.end_byte) + ')'
        tok_str = ('tokens: [' + str(self.start_token) + ',' + str(
            self.end_token) + ')')
        if self.score is not None:
            score_str = ', score: %s' % self.score
        else:
            score_str = ''

        return "Span(" + byte_str + ', ' + tok_str + score_str + ")"

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

    def __hash__(self):
        return hash((self.start_byte, self.end_byte, self.start_token, self.end_token, self.score))

    def __lt__(self, other):
        if self.score < other.score:
            return True
        return False

    def __eq__(self, o: object) -> bool:

        if not isinstance(o, NQSpan):
            return False

        for attribute in self.__slots__:
            if not hasattr(o, attribute):
                return False
            elif getattr(o, attribute) != getattr(self, attribute):
                return False

        return True


# A data structure for storing prediction and annotation.
# When a example has multiple annotations, multiple NQLabel will be used.
NQLabel = collections.namedtuple(
    'NQLabel',
    [
        'example_id',  # the unique id for each NQ example.
        'long_answer_span',  # A Span object for long answer.
        'short_answer_span_list',  # A list of Spans for short answer.
        #   Note that In NQ, the short answers
        #   do not need to be in a single span.
        'yes_no_answer',  # Indicate if the short answer is an yes/no answer
        #   The possible values are "yes", "no", "none".
        #   (case insensitive)
        #   If the field is "yes", short_answer_span_list
        #   should be empty or only contain null spans.
        'long_score',  # The prediction score for the long answer prediction.
        'short_score'  # The prediction score for the short answer prediction.
    ])


def is_null_span_list(span_list):
    """Returns true iff all spans in span_list are null or span_list is empty."""
    if not span_list or all([span.is_null_span() for span in span_list]):
        return True
    return False


def nonnull_span_equal(span_a: NQSpan, span_b: NQSpan):
    """Given two spans, return if they are equal.

    Args:
      span_a: a Span object.
      span_b: a Span object.  Only compare non-null spans. First, if the bytes are
        not negative, compare byte offsets, Otherwise, compare token offsets.

    Returns:
      True or False
    """
    assert isinstance(span_a, NQSpan)
    assert isinstance(span_b, NQSpan)
    assert not span_a.is_null_span()
infile
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import glob
import json
import logging
import multiprocessing
from gzip import GzipFile

from functools import partial
from typing import Dict, Set, Optional, List


@functools.total_ordering
class NQSpan(object):
    """A class for handling token and byte spans.
       Taken from https://github.com/google-research-datasets/natural-questions/blob/master/eval_utils.py#L60
      The logic is:
      1) if both start_byte !=  -1 and end_byte != -1 then the span is defined
         by byte offsets
      2) else, if start_token != -1 and end_token != -1 then the span is define
         by token offsets
      3) else, this is a null span.
      Null spans means that there is no (long or short) answers.
      If your systems only care about token spans rather than byte spans, set all
      byte spans to -1.
    """
    __slots__ = ['start_byte', 'end_byte', 'start_token', 'end_token', 'score', 'long_score']

    def __init__(self, start_byte: int, end_byte: int, start_token: int, end_token: int,
                 score: Optional[float] = None, long_score: Optional[float] = None, enforce_byte_consistency: bool = True,
                 enforce_token_consistency: bool = True):

        self.start_byte = start_byte
        self.end_byte = end_byte
        self.start_token = start_token
        self.end_token = end_token
        self.score = score
        self.long_score = long_score

        if enforce_byte_consistency:
            if ((start_byte < 0 <= end_byte) or
                    (start_byte >= 0 > end_byte)):
                raise ValueError('Inconsistent Null Spans (Byte): %s' % self)

            if start_byte >= 0 and 0 <= end_byte < start_byte:
                raise ValueError('Invalid byte spans (start_byte > end_byte): %s' % self)

        if enforce_token_consistency:
            if ((start_token < 0 <= end_token) or
                    (start_token >= 0 > end_token)):
                raise ValueError('Inconsistent Null Spans (Token): %s' % self)

            if ((start_token >= 0 and end_token >= 0) and
                    (start_token > end_token)):
                raise ValueError(
                    'Invalid token spans (start_token_idx > end_token_idx): %s' % self)

    def is_null_span(self):
        """A span is a null span if the start and end are both -1."""

        if (self.start_byte < 0 and self.end_byte < 0 and
                    self.start_token < 0 and self.end_token < 0):
            return True
        return False

    @staticmethod
    def null_span():
        return NQSpan(start_byte=-1, end_byte=-1, start_token=-1, end_token=-1)

    def contains(self, other_span):

        if other_span.start_byte > -1 and self.start_byte > -1 and other_span.end_byte > -1 and self.end_byte > -1:
            # Compare using bytes
            if other_span.start_byte >= self.start_byte and other_span.end_byte <= self.end_byte:
                return True
            else:
                return False
        elif other_span.start_token > -1 and self.start_token > -1 and other_span.end_token > -1 and self.end_token > -1:
            # Compare using tokens
            if other_span.start_token >= self.start_token and other_span.end_token <= self.end_token:
                return True
            else:
                return False
        else:
            return False

    def __str__(self):
        byte_str = 'bytes: [' + str(self.start_byte) + ',' + str(self.end_byte) + ')'
        tok_str = ('tokens: [' + str(self.start_token) + ',' + str(
            self.end_token) + ')')
        if self.score is not None:
            score_str = ', score: %s' % self.score
        else:
            score_str = ''

        return "Span(" + byte_str + ', ' + tok_str + score_str + ")"

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

    def __hash__(self):
        return hash((self.start_byte, self.end_byte, self.start_token, self.end_token, self.score))

    def __lt__(self, other):
        if self.score < other.score:
            return True
        return False

    def __eq__(self, o: object) -> bool:

        if not isinstance(o, NQSpan):
            return False

        for attribute in self.__slots__:
            if not hasattr(o, attribute):
                return False
            elif getattr(o, attribute) != getattr(self, attribute):
                return False

        return True


# A data structure for storing prediction and annotation.
# When a example has multiple annotations, multiple NQLabel will be used.
NQLabel = collections.namedtuple(
    'NQLabel',
    [
        'example_id',  # the unique id for each NQ example.
        'long_answer_span',  # A Span object for long answer.
        'short_answer_span_list',  # A list of Spans for short answer.
        #   Note that In NQ, the short answers
        #   do not need to be in a single span.
        'yes_no_answer',  # Indicate if the short answer is an yes/no answer
        #   The possible values are "yes", "no", "none".
        #   (case insensitive)
        #   If the field is "yes", short_answer_span_list
        #   should be empty or only contain null spans.
        'long_score',  # The prediction score for the long answer prediction.
        'short_score'  # The prediction score for the short answer prediction.
    ])


def is_null_span_list(span_list):
    """Returns true iff all spans in span_list are null or span_list is empty."""
    if not span_list or all([span.is_null_span() for span in span_list]):
        return True
    return False


def nonnull_span_equal(span_a: NQSpan, span_b: NQSpan):
    """Given two spans, return if they are equal.

    Args:
      span_a: a Span object.
      span_b: a Span object.  Only compare non-null spans. First, if the bytes are
        not negative, compare byte offsets, Otherwise, compare token offsets.

    Returns:
      True or False
    """
    assert isinstance(span_a, NQSpan)
    assert isinstance(span_b, NQSpan)
    assert not span_a.is_null_span()
    assert not span_b.is_null_span()

    # if byte offsets are not negative, compare byte offsets
    if ((span_a.start_byte >= 0 and span_a.end_byte >= 0) and
            (span_b.start_byte >= 0 and span_b.end_byte >= 0)):

        if ((span_a.start_byte == span_b.start_byte) and
                (span_a.end_byte == span_b.end_byte)):
            return True

    # if token offsets are not negative, compare token offsets
    if ((span_a.start_token >= 0 and span_a.end_token >= 0) and
            (span_b.start_token >= 0 and span_b.end_token >= 0)):

        if ((span_a.start_token == span_b.start_token) and
                (span_a.end_token == span_b.end_token)):
            return True

    return False


def span_set_equal(gold_span_list, pred_span_list):
    """Make the spans are completely equal besides null spans."""

    gold_span_list = [span for span in gold_span_list if not span.is_null_span()]
    pred_span_list = [span for span in pred_span_list if not span.is_null_span()]

    for pspan in pred_span_list:
        # not finding pspan equal to any spans in gold_span_list
infile
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import glob
import json
import logging
import multiprocessing
from gzip import GzipFile

from functools import partial
from typing import Dict, Set, Optional, List


@functools.total_ordering
class NQSpan(object):
    """A class for handling token and byte spans.
       Taken from https://github.com/google-research-datasets/natural-questions/blob/master/eval_utils.py#L60
      The logic is:
      1) if both start_byte !=  -1 and end_byte != -1 then the span is defined
         by byte offsets
      2) else, if start_token != -1 and end_token != -1 then the span is define
         by token offsets
      3) else, this is a null span.
      Null spans means that there is no (long or short) answers.
      If your systems only care about token spans rather than byte spans, set all
      byte spans to -1.
    """
    __slots__ = ['start_byte', 'end_byte', 'start_token', 'end_token', 'score', 'long_score']

    def __init__(self, start_byte: int, end_byte: int, start_token: int, end_token: int,
                 score: Optional[float] = None, long_score: Optional[float] = None, enforce_byte_consistency: bool = True,
                 enforce_token_consistency: bool = True):

        self.start_byte = start_byte
        self.end_byte = end_byte
        self.start_token = start_token
        self.end_token = end_token
        self.score = score
        self.long_score = long_score

        if enforce_byte_consistency:
            if ((start_byte < 0 <= end_byte) or
                    (start_byte >= 0 > end_byte)):
                raise ValueError('Inconsistent Null Spans (Byte): %s' % self)

            if start_byte >= 0 and 0 <= end_byte < start_byte:
                raise ValueError('Invalid byte spans (start_byte > end_byte): %s' % self)

        if enforce_token_consistency:
            if ((start_token < 0 <= end_token) or
                    (start_token >= 0 > end_token)):
                raise ValueError('Inconsistent Null Spans (Token): %s' % self)

            if ((start_token >= 0 and end_token >= 0) and
                    (start_token > end_token)):
                raise ValueError(
                    'Invalid token spans (start_token_idx > end_token_idx): %s' % self)

    def is_null_span(self):
        """A span is a null span if the start and end are both -1."""

        if (self.start_byte < 0 and self.end_byte < 0 and
                    self.start_token < 0 and self.end_token < 0):
            return True
        return False

    @staticmethod
    def null_span():
infile
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import glob
import json
import logging
import multiprocessing
from gzip import GzipFile

from functools import partial
from typing import Dict, Set, Optional, List


@functools.total_ordering
class NQSpan(object):
    """A class for handling token and byte spans.
       Taken from https://github.com/google-research-datasets/natural-questions/blob/master/eval_utils.py#L60
      The logic is:
      1) if both start_byte !=  -1 and end_byte != -1 then the span is defined
         by byte offsets
      2) else, if start_token != -1 and end_token != -1 then the span is define
         by token offsets
      3) else, this is a null span.
      Null spans means that there is no (long or short) answers.
      If your systems only care about token spans rather than byte spans, set all
      byte spans to -1.
    """
    __slots__ = ['start_byte', 'end_byte', 'start_token', 'end_token', 'score', 'long_score']

    def __init__(self, start_byte: int, end_byte: int, start_token: int, end_token: int,
                 score: Optional[float] = None, long_score: Optional[float] = None, enforce_byte_consistency: bool = True,
                 enforce_token_consistency: bool = True):

        self.start_byte = start_byte
        self.end_byte = end_byte
        self.start_token = start_token
        self.end_token = end_token
        self.score = score
        self.long_score = long_score

        if enforce_byte_consistency:
            if ((start_byte < 0 <= end_byte) or
                    (start_byte >= 0 > end_byte)):
                raise ValueError('Inconsistent Null Spans (Byte): %s' % self)

            if start_byte >= 0 and 0 <= end_byte < start_byte:
                raise ValueError('Invalid byte spans (start_byte > end_byte): %s' % self)

        if enforce_token_consistency:
            if ((start_token < 0 <= end_token) or
                    (start_token >= 0 > end_token)):
                raise ValueError('Inconsistent Null Spans (Token): %s' % self)

            if ((start_token >= 0 and end_token >= 0) and
                    (start_token > end_token)):
                raise ValueError(
                    'Invalid token spans (start_token_idx > end_token_idx): %s' % self)

    def is_null_span(self):
        """A span is a null span if the start and end are both -1."""

        if (self.start_byte < 0 and self.end_byte < 0 and
                    self.start_token < 0 and self.end_token < 0):
            return True
        return False

    @staticmethod
    def null_span():
        return NQSpan(start_byte=-1, end_byte=-1, start_token=-1, end_token=-1)

    def contains(self, other_span):

        if other_span.start_byte > -1 and self.start_byte > -1 and other_span.end_byte > -1 and self.end_byte > -1:
            # Compare using bytes
            if other_span.start_byte >= self.start_byte and other_span.end_byte <= self.end_byte:
                return True
            else:
                return False
        elif other_span.start_token > -1 and self.start_token > -1 and other_span.end_token > -1 and self.end_token > -1:
            # Compare using tokens
            if other_span.start_token >= self.start_token and other_span.end_token <= self.end_token:
                return True
            else:
                return False
        else:
            return False

    def __str__(self):
        byte_str = 'bytes: [' + str(self.start_byte) + ',' + str(self.end_byte) + ')'
        tok_str = ('tokens: [' + str(self.start_token) + ',' + str(
            self.end_token) + ')')
        if self.score is not None:
            score_str = ', score: %s' % self.score
        else:
            score_str = ''

        return "Span(" + byte_str + ', ' + tok_str + score_str + ")"

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

    def __hash__(self):
        return hash((self.start_byte, self.end_byte, self.start_token, self.end_token, self.score))

    def __lt__(self, other):
        if self.score < other.score:
            return True
        return False

    def __eq__(self, o: object) -> bool:

        if not isinstance(o, NQSpan):
            return False

        for attribute in self.__slots__:
            if not hasattr(o, attribute):
                return False
            elif getattr(o, attribute) != getattr(self, attribute):
                return False

        return True


# A data structure for storing prediction and annotation.
# When a example has multiple annotations, multiple NQLabel will be used.
NQLabel = collections.namedtuple(
    'NQLabel',
    [
        'example_id',  # the unique id for each NQ example.
        'long_answer_span',  # A Span object for long answer.
        'short_answer_span_list',  # A list of Spans for short answer.
        #   Note that In NQ, the short answers
        #   do not need to be in a single span.
        'yes_no_answer',  # Indicate if the short answer is an yes/no answer
        #   The possible values are "yes", "no", "none".
        #   (case insensitive)
        #   If the field is "yes", short_answer_span_list
        #   should be empty or only contain null spans.
        'long_score',  # The prediction score for the long answer prediction.
        'short_score'  # The prediction score for the short answer prediction.
    ])


def is_null_span_list(span_list):
    """Returns true iff all spans in span_list are null or span_list is empty."""
    if not span_list or all([span.is_null_span() for span in span_list]):
        return True
    return False


def nonnull_span_equal(span_a: NQSpan, span_b: NQSpan):
    """Given two spans, return if they are equal.

    Args:
      span_a: a Span object.
      span_b: a Span object.  Only compare non-null spans. First, if the bytes are
        not negative, compare byte offsets, Otherwise, compare token offsets.

    Returns:
      True or False
    """
    assert isinstance(span_a, NQSpan)
    assert isinstance(span_b, NQSpan)
    assert not span_a.is_null_span()
    assert not span_b.is_null_span()

    # if byte offsets are not negative, compare byte offsets
    if ((span_a.start_byte >= 0 and span_a.end_byte >= 0) and
            (span_b.start_byte >= 0 and span_b.end_byte >= 0)):

        if ((span_a.start_byte == span_b.start_byte) and
                (span_a.end_byte == span_b.end_byte)):
            return True

    # if token offsets are not negative, compare token offsets
    if ((span_a.start_token >= 0 and span_a.end_token >= 0) and
            (span_b.start_token >= 0 and span_b.end_token >= 0)):

        if ((span_a.start_token == span_b.start_token) and
                (span_a.end_token == span_b.end_token)):
            return True

    return False


def span_set_equal(gold_span_list, pred_span_list):
    """Make the spans are completely equal besides null spans."""

    gold_span_list = [span for span in gold_span_list if not span.is_null_span()]
    pred_span_list = [span for span in pred_span_list if not span.is_null_span()]

    for pspan in pred_span_list:
        # not finding pspan equal to any spans in gold_span_list
        if not any([nonnull_span_equal(pspan, gspan) for gspan in gold_span_list]):
            return False

    for gspan in gold_span_list:
        # not finding gspan equal to any spans in pred_span_list
infile
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import glob
import json
import logging
import multiprocessing
from gzip import GzipFile

from functools import partial
from typing import Dict, Set, Optional, List


@functools.total_ordering
class NQSpan(object):
    """A class for handling token and byte spans.
       Taken from https://github.com/google-research-datasets/natural-questions/blob/master/eval_utils.py#L60
      The logic is:
      1) if both start_byte !=  -1 and end_byte != -1 then the span is defined
         by byte offsets
      2) else, if start_token != -1 and end_token != -1 then the span is define
         by token offsets
      3) else, this is a null span.
      Null spans means that there is no (long or short) answers.
      If your systems only care about token spans rather than byte spans, set all
      byte spans to -1.
    """
    __slots__ = ['start_byte', 'end_byte', 'start_token', 'end_token', 'score', 'long_score']

    def __init__(self, start_byte: int, end_byte: int, start_token: int, end_token: int,
                 score: Optional[float] = None, long_score: Optional[float] = None, enforce_byte_consistency: bool = True,
                 enforce_token_consistency: bool = True):

        self.start_byte = start_byte
        self.end_byte = end_byte
        self.start_token = start_token
        self.end_token = end_token
        self.score = score
        self.long_score = long_score

        if enforce_byte_consistency:
            if ((start_byte < 0 <= end_byte) or
                    (start_byte >= 0 > end_byte)):
                raise ValueError('Inconsistent Null Spans (Byte): %s' % self)

            if start_byte >= 0 and 0 <= end_byte < start_byte:
                raise ValueError('Invalid byte spans (start_byte > end_byte): %s' % self)

        if enforce_token_consistency:
            if ((start_token < 0 <= end_token) or
                    (start_token >= 0 > end_token)):
                raise ValueError('Inconsistent Null Spans (Token): %s' % self)

            if ((start_token >= 0 and end_token >= 0) and
                    (start_token > end_token)):
                raise ValueError(
                    'Invalid token spans (start_token_idx > end_token_idx): %s' % self)

    def is_null_span(self):
        """A span is a null span if the start and end are both -1."""

        if (self.start_byte < 0 and self.end_byte < 0 and
                    self.start_token < 0 and self.end_token < 0):
            return True
        return False

    @staticmethod
    def null_span():
        return NQSpan(start_byte=-1, end_byte=-1, start_token=-1, end_token=-1)

    def contains(self, other_span):

        if other_span.start_byte > -1 and self.start_byte > -1 and other_span.end_byte > -1 and self.end_byte > -1:
            # Compare using bytes
            if other_span.start_byte >= self.start_byte and other_span.end_byte <= self.end_byte:
                return True
            else:
                return False
        elif other_span.start_token > -1 and self.start_token > -1 and other_span.end_token > -1 and self.end_token > -1:
            # Compare using tokens
            if other_span.start_token >= self.start_token and other_span.end_token <= self.end_token:
                return True
            else:
                return False
        else:
            return False

    def __str__(self):
        byte_str = 'bytes: [' + str(self.start_byte) + ',' + str(self.end_byte) + ')'
        tok_str = ('tokens: [' + str(self.start_token) + ',' + str(
            self.end_token) + ')')
        if self.score is not None:
            score_str = ', score: %s' % self.score
        else:
            score_str = ''

        return "Span(" + byte_str + ', ' + tok_str + score_str + ")"

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

    def __hash__(self):
        return hash((self.start_byte, self.end_byte, self.start_token, self.end_token, self.score))

    def __lt__(self, other):
        if self.score < other.score:
            return True
        return False

    def __eq__(self, o: object) -> bool:

infile
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import glob
import json
import logging
import multiprocessing
from gzip import GzipFile

from functools import partial
from typing import Dict, Set, Optional, List


@functools.total_ordering
class NQSpan(object):
    """A class for handling token and byte spans.
       Taken from https://github.com/google-research-datasets/natural-questions/blob/master/eval_utils.py#L60
      The logic is:
      1) if both start_byte !=  -1 and end_byte != -1 then the span is defined
         by byte offsets
      2) else, if start_token != -1 and end_token != -1 then the span is define
         by token offsets
      3) else, this is a null span.
      Null spans means that there is no (long or short) answers.
      If your systems only care about token spans rather than byte spans, set all
      byte spans to -1.
    """
    __slots__ = ['start_byte', 'end_byte', 'start_token', 'end_token', 'score', 'long_score']

    def __init__(self, start_byte: int, end_byte: int, start_token: int, end_token: int,
                 score: Optional[float] = None, long_score: Optional[float] = None, enforce_byte_consistency: bool = True,
                 enforce_token_consistency: bool = True):

        self.start_byte = start_byte
        self.end_byte = end_byte
        self.start_token = start_token
        self.end_token = end_token
        self.score = score
        self.long_score = long_score

        if enforce_byte_consistency:
            if ((start_byte < 0 <= end_byte) or
                    (start_byte >= 0 > end_byte)):
                raise ValueError('Inconsistent Null Spans (Byte): %s' % self)

            if start_byte >= 0 and 0 <= end_byte < start_byte:
                raise ValueError('Invalid byte spans (start_byte > end_byte): %s' % self)

        if enforce_token_consistency:
            if ((start_token < 0 <= end_token) or
                    (start_token >= 0 > end_token)):
                raise ValueError('Inconsistent Null Spans (Token): %s' % self)

            if ((start_token >= 0 and end_token >= 0) and
                    (start_token > end_token)):
                raise ValueError(
                    'Invalid token spans (start_token_idx > end_token_idx): %s' % self)

    def is_null_span(self):
        """A span is a null span if the start and end are both -1."""

        if (self.start_byte < 0 and self.end_byte < 0 and
                    self.start_token < 0 and self.end_token < 0):
            return True
        return False

    @staticmethod
    def null_span():
        return NQSpan(start_byte=-1, end_byte=-1, start_token=-1, end_token=-1)

    def contains(self, other_span):

        if other_span.start_byte > -1 and self.start_byte > -1 and other_span.end_byte > -1 and self.end_byte > -1:
            # Compare using bytes
            if other_span.start_byte >= self.start_byte and other_span.end_byte <= self.end_byte:
                return True
            else:
                return False
        elif other_span.start_token > -1 and self.start_token > -1 and other_span.end_token > -1 and self.end_token > -1:
            # Compare using tokens
            if other_span.start_token >= self.start_token and other_span.end_token <= self.end_token:
                return True
            else:
                return False
        else:
            return False

    def __str__(self):
        byte_str = 'bytes: [' + str(self.start_byte) + ',' + str(self.end_byte) + ')'
        tok_str = ('tokens: [' + str(self.start_token) + ',' + str(
            self.end_token) + ')')
        if self.score is not None:
            score_str = ', score: %s' % self.score
        else:
            score_str = ''

        return "Span(" + byte_str + ', ' + tok_str + score_str + ")"

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

    def __hash__(self):
        return hash((self.start_byte, self.end_byte, self.start_token, self.end_token, self.score))

    def __lt__(self, other):
        if self.score < other.score:
            return True
        return False

    def __eq__(self, o: object) -> bool:

        if not isinstance(o, NQSpan):
            return False

        for attribute in self.__slots__:
            if not hasattr(o, attribute):
                return False
            elif getattr(o, attribute) != getattr(self, attribute):
                return False

        return True


# A data structure for storing prediction and annotation.
# When a example has multiple annotations, multiple NQLabel will be used.
NQLabel = collections.namedtuple(
    'NQLabel',
    [
        'example_id',  # the unique id for each NQ example.
        'long_answer_span',  # A Span object for long answer.
        'short_answer_span_list',  # A list of Spans for short answer.
        #   Note that In NQ, the short answers
        #   do not need to be in a single span.
        'yes_no_answer',  # Indicate if the short answer is an yes/no answer
        #   The possible values are "yes", "no", "none".
        #   (case insensitive)
        #   If the field is "yes", short_answer_span_list
        #   should be empty or only contain null spans.
        'long_score',  # The prediction score for the long answer prediction.
        'short_score'  # The prediction score for the short answer prediction.
    ])


def is_null_span_list(span_list):
    """Returns true iff all spans in span_list are null or span_list is empty."""
infile
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import glob
import json
import logging
import multiprocessing
from gzip import GzipFile

from functools import partial
from typing import Dict, Set, Optional, List


@functools.total_ordering
class NQSpan(object):
    """A class for handling token and byte spans.
       Taken from https://github.com/google-research-datasets/natural-questions/blob/master/eval_utils.py#L60
      The logic is:
      1) if both start_byte !=  -1 and end_byte != -1 then the span is defined
         by byte offsets
      2) else, if start_token != -1 and end_token != -1 then the span is define
         by token offsets
      3) else, this is a null span.
      Null spans means that there is no (long or short) answers.
      If your systems only care about token spans rather than byte spans, set all
      byte spans to -1.
    """
    __slots__ = ['start_byte', 'end_byte', 'start_token', 'end_token', 'score', 'long_score']

    def __init__(self, start_byte: int, end_byte: int, start_token: int, end_token: int,
                 score: Optional[float] = None, long_score: Optional[float] = None, enforce_byte_consistency: bool = True,
                 enforce_token_consistency: bool = True):

        self.start_byte = start_byte
        self.end_byte = end_byte
        self.start_token = start_token
        self.end_token = end_token
        self.score = score
        self.long_score = long_score

        if enforce_byte_consistency:
            if ((start_byte < 0 <= end_byte) or
                    (start_byte >= 0 > end_byte)):
                raise ValueError('Inconsistent Null Spans (Byte): %s' % self)

            if start_byte >= 0 and 0 <= end_byte < start_byte:
                raise ValueError('Invalid byte spans (start_byte > end_byte): %s' % self)

        if enforce_token_consistency:
            if ((start_token < 0 <= end_token) or
                    (start_token >= 0 > end_token)):
                raise ValueError('Inconsistent Null Spans (Token): %s' % self)

            if ((start_token >= 0 and end_token >= 0) and
                    (start_token > end_token)):
                raise ValueError(
                    'Invalid token spans (start_token_idx > end_token_idx): %s' % self)

    def is_null_span(self):
        """A span is a null span if the start and end are both -1."""

        if (self.start_byte < 0 and self.end_byte < 0 and
                    self.start_token < 0 and self.end_token < 0):
            return True
        return False

    @staticmethod
    def null_span():
        return NQSpan(start_byte=-1, end_byte=-1, start_token=-1, end_token=-1)

    def contains(self, other_span):

        if other_span.start_byte > -1 and self.start_byte > -1 and other_span.end_byte > -1 and self.end_byte > -1:
            # Compare using bytes
            if other_span.start_byte >= self.start_byte and other_span.end_byte <= self.end_byte:
                return True
            else:
                return False
        elif other_span.start_token > -1 and self.start_token > -1 and other_span.end_token > -1 and self.end_token > -1:
            # Compare using tokens
            if other_span.start_token >= self.start_token and other_span.end_token <= self.end_token:
                return True
            else:
                return False
        else:
            return False

    def __str__(self):
        byte_str = 'bytes: [' + str(self.start_byte) + ',' + str(self.end_byte) + ')'
        tok_str = ('tokens: [' + str(self.start_token) + ',' + str(
            self.end_token) + ')')
        if self.score is not None:
            score_str = ', score: %s' % self.score
        else:
            score_str = ''

        return "Span(" + byte_str + ', ' + tok_str + score_str + ")"

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

    def __hash__(self):
        return hash((self.start_byte, self.end_byte, self.start_token, self.end_token, self.score))

    def __lt__(self, other):
        if self.score < other.score:
            return True
        return False

    def __eq__(self, o: object) -> bool:

        if not isinstance(o, NQSpan):
            return False

        for attribute in self.__slots__:
            if not hasattr(o, attribute):
                return False
            elif getattr(o, attribute) != getattr(self, attribute):
                return False

        return True


# A data structure for storing prediction and annotation.
# When a example has multiple annotations, multiple NQLabel will be used.
NQLabel = collections.namedtuple(
    'NQLabel',
    [
        'example_id',  # the unique id for each NQ example.
        'long_answer_span',  # A Span object for long answer.
        'short_answer_span_list',  # A list of Spans for short answer.
        #   Note that In NQ, the short answers
        #   do not need to be in a single span.
        'yes_no_answer',  # Indicate if the short answer is an yes/no answer
        #   The possible values are "yes", "no", "none".
        #   (case insensitive)
        #   If the field is "yes", short_answer_span_list
        #   should be empty or only contain null spans.
        'long_score',  # The prediction score for the long answer prediction.
        'short_score'  # The prediction score for the short answer prediction.
    ])


def is_null_span_list(span_list):
    """Returns true iff all spans in span_list are null or span_list is empty."""
    if not span_list or all([span.is_null_span() for span in span_list]):
        return True
    return False


def nonnull_span_equal(span_a: NQSpan, span_b: NQSpan):
    """Given two spans, return if they are equal.

    Args:
      span_a: a Span object.
      span_b: a Span object.  Only compare non-null spans. First, if the bytes are
        not negative, compare byte offsets, Otherwise, compare token offsets.

    Returns:
      True or False
    """
    assert isinstance(span_a, NQSpan)
    assert isinstance(span_b, NQSpan)
    assert not span_a.is_null_span()
    assert not span_b.is_null_span()

    # if byte offsets are not negative, compare byte offsets
    if ((span_a.start_byte >= 0 and span_a.end_byte >= 0) and
            (span_b.start_byte >= 0 and span_b.end_byte >= 0)):

        if ((span_a.start_byte == span_b.start_byte) and
                (span_a.end_byte == span_b.end_byte)):
            return True

    # if token offsets are not negative, compare token offsets
    if ((span_a.start_token >= 0 and span_a.end_token >= 0) and
            (span_b.start_token >= 0 and span_b.end_token >= 0)):

        if ((span_a.start_token == span_b.start_token) and
                (span_a.end_token == span_b.end_token)):
            return True

    return False


def span_set_equal(gold_span_list, pred_span_list):
    """Make the spans are completely equal besides null spans."""

    gold_span_list = [span for span in gold_span_list if not span.is_null_span()]
    pred_span_list = [span for span in pred_span_list if not span.is_null_span()]

    for pspan in pred_span_list:
        # not finding pspan equal to any spans in gold_span_list
        if not any([nonnull_span_equal(pspan, gspan) for gspan in gold_span_list]):
            return False

    for gspan in gold_span_list:
        # not finding gspan equal to any spans in pred_span_list
        if not any([nonnull_span_equal(pspan, gspan) for pspan in pred_span_list]):
            return False

    return True


def gold_has_short_answer(gold_label_list: list, short_non_null_threshold: int = 2) -> bool:
    """
    Gets vote from multi-annotators for judging if there is a short answer.
    :param gold_label_list: list of gold labels
    :param short_non_null_threshold: Require this many non-null short answer annotations
      to count gold as containing a short answer. Defaults to 2 like the original paper.
    """

    if not gold_label_list:
        return False  # Empty list will not have gold short answer

    # We consider if there is a short answer if there is an short answer span or
    #  the yes/no answer is not none.
    gold_has_answer = sum([
infile
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import glob
import json
import logging
import multiprocessing
from gzip import GzipFile

from functools import partial
from typing import Dict, Set, Optional, List


@functools.total_ordering
class NQSpan(object):
    """A class for handling token and byte spans.
       Taken from https://github.com/google-research-datasets/natural-questions/blob/master/eval_utils.py#L60
      The logic is:
      1) if both start_byte !=  -1 and end_byte != -1 then the span is defined
         by byte offsets
      2) else, if start_token != -1 and end_token != -1 then the span is define
         by token offsets
      3) else, this is a null span.
      Null spans means that there is no (long or short) answers.
      If your systems only care about token spans rather than byte spans, set all
      byte spans to -1.
    """
    __slots__ = ['start_byte', 'end_byte', 'start_token', 'end_token', 'score', 'long_score']

    def __init__(self, start_byte: int, end_byte: int, start_token: int, end_token: int,
                 score: Optional[float] = None, long_score: Optional[float] = None, enforce_byte_consistency: bool = True,
                 enforce_token_consistency: bool = True):

        self.start_byte = start_byte
        self.end_byte = end_byte
        self.start_token = start_token
        self.end_token = end_token
        self.score = score
        self.long_score = long_score

        if enforce_byte_consistency:
            if ((start_byte < 0 <= end_byte) or
                    (start_byte >= 0 > end_byte)):
                raise ValueError('Inconsistent Null Spans (Byte): %s' % self)

            if start_byte >= 0 and 0 <= end_byte < start_byte:
                raise ValueError('Invalid byte spans (start_byte > end_byte): %s' % self)

        if enforce_token_consistency:
            if ((start_token < 0 <= end_token) or
                    (start_token >= 0 > end_token)):
                raise ValueError('Inconsistent Null Spans (Token): %s' % self)

            if ((start_token >= 0 and end_token >= 0) and
                    (start_token > end_token)):
                raise ValueError(
                    'Invalid token spans (start_token_idx > end_token_idx): %s' % self)

    def is_null_span(self):
        """A span is a null span if the start and end are both -1."""

        if (self.start_byte < 0 and self.end_byte < 0 and
                    self.start_token < 0 and self.end_token < 0):
            return True
        return False

    @staticmethod
    def null_span():
        return NQSpan(start_byte=-1, end_byte=-1, start_token=-1, end_token=-1)

    def contains(self, other_span):

        if other_span.start_byte > -1 and self.start_byte > -1 and other_span.end_byte > -1 and self.end_byte > -1:
            # Compare using bytes
            if other_span.start_byte >= self.start_byte and other_span.end_byte <= self.end_byte:
                return True
            else:
                return False
        elif other_span.start_token > -1 and self.start_token > -1 and other_span.end_token > -1 and self.end_token > -1:
            # Compare using tokens
            if other_span.start_token >= self.start_token and other_span.end_token <= self.end_token:
                return True
            else:
                return False
        else:
            return False

    def __str__(self):
        byte_str = 'bytes: [' + str(self.start_byte) + ',' + str(self.end_byte) + ')'
        tok_str = ('tokens: [' + str(self.start_token) + ',' + str(
            self.end_token) + ')')
        if self.score is not None:
            score_str = ', score: %s' % self.score
        else:
            score_str = ''

        return "Span(" + byte_str + ', ' + tok_str + score_str + ")"

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

    def __hash__(self):
        return hash((self.start_byte, self.end_byte, self.start_token, self.end_token, self.score))

    def __lt__(self, other):
        if self.score < other.score:
            return True
        return False

    def __eq__(self, o: object) -> bool:

        if not isinstance(o, NQSpan):
            return False

        for attribute in self.__slots__:
            if not hasattr(o, attribute):
                return False
            elif getattr(o, attribute) != getattr(self, attribute):
                return False

        return True


# A data structure for storing prediction and annotation.
# When a example has multiple annotations, multiple NQLabel will be used.
NQLabel = collections.namedtuple(
    'NQLabel',
    [
        'example_id',  # the unique id for each NQ example.
        'long_answer_span',  # A Span object for long answer.
        'short_answer_span_list',  # A list of Spans for short answer.
        #   Note that In NQ, the short answers
        #   do not need to be in a single span.
        'yes_no_answer',  # Indicate if the short answer is an yes/no answer
        #   The possible values are "yes", "no", "none".
        #   (case insensitive)
        #   If the field is "yes", short_answer_span_list
        #   should be empty or only contain null spans.
        'long_score',  # The prediction score for the long answer prediction.
        'short_score'  # The prediction score for the short answer prediction.
    ])


def is_null_span_list(span_list):
    """Returns true iff all spans in span_list are null or span_list is empty."""
    if not span_list or all([span.is_null_span() for span in span_list]):
        return True
    return False


def nonnull_span_equal(span_a: NQSpan, span_b: NQSpan):
    """Given two spans, return if they are equal.

    Args:
      span_a: a Span object.
      span_b: a Span object.  Only compare non-null spans. First, if the bytes are
        not negative, compare byte offsets, Otherwise, compare token offsets.

    Returns:
      True or False
    """
    assert isinstance(span_a, NQSpan)
    assert isinstance(span_b, NQSpan)
    assert not span_a.is_null_span()
    assert not span_b.is_null_span()

    # if byte offsets are not negative, compare byte offsets
    if ((span_a.start_byte >= 0 and span_a.end_byte >= 0) and
            (span_b.start_byte >= 0 and span_b.end_byte >= 0)):

        if ((span_a.start_byte == span_b.start_byte) and
                (span_a.end_byte == span_b.end_byte)):
            return True

    # if token offsets are not negative, compare token offsets
    if ((span_a.start_token >= 0 and span_a.end_token >= 0) and
            (span_b.start_token >= 0 and span_b.end_token >= 0)):

        if ((span_a.start_token == span_b.start_token) and
                (span_a.end_token == span_b.end_token)):
            return True

    return False


def span_set_equal(gold_span_list, pred_span_list):
    """Make the spans are completely equal besides null spans."""

infile
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import glob
import json
import logging
import multiprocessing
from gzip import GzipFile

from functools import partial
from typing import Dict, Set, Optional, List


@functools.total_ordering
class NQSpan(object):
    """A class for handling token and byte spans.
       Taken from https://github.com/google-research-datasets/natural-questions/blob/master/eval_utils.py#L60
      The logic is:
      1) if both start_byte !=  -1 and end_byte != -1 then the span is defined
         by byte offsets
      2) else, if start_token != -1 and end_token != -1 then the span is define
         by token offsets
      3) else, this is a null span.
      Null spans means that there is no (long or short) answers.
      If your systems only care about token spans rather than byte spans, set all
      byte spans to -1.
    """
    __slots__ = ['start_byte', 'end_byte', 'start_token', 'end_token', 'score', 'long_score']

    def __init__(self, start_byte: int, end_byte: int, start_token: int, end_token: int,
                 score: Optional[float] = None, long_score: Optional[float] = None, enforce_byte_consistency: bool = True,
                 enforce_token_consistency: bool = True):

        self.start_byte = start_byte
        self.end_byte = end_byte
        self.start_token = start_token
        self.end_token = end_token
        self.score = score
        self.long_score = long_score

        if enforce_byte_consistency:
            if ((start_byte < 0 <= end_byte) or
                    (start_byte >= 0 > end_byte)):
                raise ValueError('Inconsistent Null Spans (Byte): %s' % self)

            if start_byte >= 0 and 0 <= end_byte < start_byte:
                raise ValueError('Invalid byte spans (start_byte > end_byte): %s' % self)

        if enforce_token_consistency:
            if ((start_token < 0 <= end_token) or
                    (start_token >= 0 > end_token)):
                raise ValueError('Inconsistent Null Spans (Token): %s' % self)

            if ((start_token >= 0 and end_token >= 0) and
                    (start_token > end_token)):
                raise ValueError(
                    'Invalid token spans (start_token_idx > end_token_idx): %s' % self)

    def is_null_span(self):
        """A span is a null span if the start and end are both -1."""

        if (self.start_byte < 0 and self.end_byte < 0 and
                    self.start_token < 0 and self.end_token < 0):
            return True
        return False

    @staticmethod
    def null_span():
        return NQSpan(start_byte=-1, end_byte=-1, start_token=-1, end_token=-1)

    def contains(self, other_span):

        if other_span.start_byte > -1 and self.start_byte > -1 and other_span.end_byte > -1 and self.end_byte > -1:
            # Compare using bytes
            if other_span.start_byte >= self.start_byte and other_span.end_byte <= self.end_byte:
                return True
            else:
                return False
        elif other_span.start_token > -1 and self.start_token > -1 and other_span.end_token > -1 and self.end_token > -1:
            # Compare using tokens
            if other_span.start_token >= self.start_token and other_span.end_token <= self.end_token:
                return True
            else:
                return False
        else:
            return False

    def __str__(self):
        byte_str = 'bytes: [' + str(self.start_byte) + ',' + str(self.end_byte) + ')'
        tok_str = ('tokens: [' + str(self.start_token) + ',' + str(
            self.end_token) + ')')
        if self.score is not None:
            score_str = ', score: %s' % self.score
        else:
            score_str = ''

        return "Span(" + byte_str + ', ' + tok_str + score_str + ")"

    def __repr__(self):
commited
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import glob
import json
import logging
import multiprocessing
from gzip import GzipFile

from functools import partial
from typing import Dict, Set, Optional, List


@functools.total_ordering
class NQSpan(object):
    """A class for handling token and byte spans.
       Taken from https://github.com/google-research-datasets/natural-questions/blob/master/eval_utils.py#L60
      The logic is:
      1) if both start_byte !=  -1 and end_byte != -1 then the span is defined
         by byte offsets
      2) else, if start_token != -1 and end_token != -1 then the span is define
         by token offsets
      3) else, this is a null span.
      Null spans means that there is no (long or short) answers.
      If your systems only care about token spans rather than byte spans, set all
      byte spans to -1.
    """
    __slots__ = ['start_byte', 'end_byte', 'start_token', 'end_token', 'score', 'long_score']

    def __init__(self, start_byte: int, end_byte: int, start_token: int, end_token: int,
                 score: Optional[float] = None, long_score: Optional[float] = None, enforce_byte_consistency: bool = True,
                 enforce_token_consistency: bool = True):

        self.start_byte = start_byte
        self.end_byte = end_byte
        self.start_token = start_token
        self.end_token = end_token
        self.score = score
        self.long_score = long_score

        if enforce_byte_consistency:
            if ((start_byte < 0 <= end_byte) or
                    (start_byte >= 0 > end_byte)):
                raise ValueError('Inconsistent Null Spans (Byte): %s' % self)

            if start_byte >= 0 and 0 <= end_byte < start_byte:
                raise ValueError('Invalid byte spans (start_byte > end_byte): %s' % self)

        if enforce_token_consistency:
            if ((start_token < 0 <= end_token) or
                    (start_token >= 0 > end_token)):
                raise ValueError('Inconsistent Null Spans (Token): %s' % self)

            if ((start_token >= 0 and end_token >= 0) and
                    (start_token > end_token)):
                raise ValueError(
                    'Invalid token spans (start_token_idx > end_token_idx): %s' % self)

    def is_null_span(self):
        """A span is a null span if the start and end are both -1."""

        if (self.start_byte < 0 and self.end_byte < 0 and
                    self.start_token < 0 and self.end_token < 0):
            return True
        return False

    @staticmethod
    def null_span():
        return NQSpan(start_byte=-1, end_byte=-1, start_token=-1, end_token=-1)

    def contains(self, other_span):

        if other_span.start_byte > -1 and self.start_byte > -1 and other_span.end_byte > -1 and self.end_byte > -1:
            # Compare using bytes
            if other_span.start_byte >= self.start_byte and other_span.end_byte <= self.end_byte:
                return True
            else:
                return False
        elif other_span.start_token > -1 and self.start_token > -1 and other_span.end_token > -1 and self.end_token > -1:
            # Compare using tokens
            if other_span.start_token >= self.start_token and other_span.end_token <= self.end_token:
                return True
            else:
                return False
        else:
            return False

    def __str__(self):
        byte_str = 'bytes: [' + str(self.start_byte) + ',' + str(self.end_byte) + ')'
        tok_str = ('tokens: [' + str(self.start_token) + ',' + str(
            self.end_token) + ')')
        if self.score is not None:
            score_str = ', score: %s' % self.score
        else:
            score_str = ''

        return "Span(" + byte_str + ', ' + tok_str + score_str + ")"

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

    def __hash__(self):
        return hash((self.start_byte, self.end_byte, self.start_token, self.end_token, self.score))

    def __lt__(self, other):
        if self.score < other.score:
            return True
        return False

    def __eq__(self, o: object) -> bool:

        if not isinstance(o, NQSpan):
            return False

        for attribute in self.__slots__:
            if not hasattr(o, attribute):
                return False
            elif getattr(o, attribute) != getattr(self, attribute):
                return False

        return True


# A data structure for storing prediction and annotation.
# When a example has multiple annotations, multiple NQLabel will be used.
NQLabel = collections.namedtuple(
    'NQLabel',
    [
        'example_id',  # the unique id for each NQ example.
        'long_answer_span',  # A Span object for long answer.
        'short_answer_span_list',  # A list of Spans for short answer.
        #   Note that In NQ, the short answers
        #   do not need to be in a single span.
        'yes_no_answer',  # Indicate if the short answer is an yes/no answer
        #   The possible values are "yes", "no", "none".
        #   (case insensitive)
        #   If the field is "yes", short_answer_span_list
        #   should be empty or only contain null spans.
        'long_score',  # The prediction score for the long answer prediction.
        'short_score'  # The prediction score for the short answer prediction.
    ])


def is_null_span_list(span_list):
    """Returns true iff all spans in span_list are null or span_list is empty."""
    if not span_list or all([span.is_null_span() for span in span_list]):
        return True
    return False


def nonnull_span_equal(span_a: NQSpan, span_b: NQSpan):
    """Given two spans, return if they are equal.

    Args:
      span_a: a Span object.
      span_b: a Span object.  Only compare non-null spans. First, if the bytes are
        not negative, compare byte offsets, Otherwise, compare token offsets.

    Returns:
      True or False
    """
    assert isinstance(span_a, NQSpan)
    assert isinstance(span_b, NQSpan)
    assert not span_a.is_null_span()
    assert not span_b.is_null_span()

    # if byte offsets are not negative, compare byte offsets
    if ((span_a.start_byte >= 0 and span_a.end_byte >= 0) and
            (span_b.start_byte >= 0 and span_b.end_byte >= 0)):

        if ((span_a.start_byte == span_b.start_byte) and
                (span_a.end_byte == span_b.end_byte)):
            return True

    # if token offsets are not negative, compare token offsets
    if ((span_a.start_token >= 0 and span_a.end_token >= 0) and
            (span_b.start_token >= 0 and span_b.end_token >= 0)):

        if ((span_a.start_token == span_b.start_token) and
                (span_a.end_token == span_b.end_token)):
            return True

    return False


def span_set_equal(gold_span_list, pred_span_list):
    """Make the spans are completely equal besides null spans."""

    gold_span_list = [span for span in gold_span_list if not span.is_null_span()]
    pred_span_list = [span for span in pred_span_list if not span.is_null_span()]

    for pspan in pred_span_list:
        # not finding pspan equal to any spans in gold_span_list
        if not any([nonnull_span_equal(pspan, gspan) for gspan in gold_span_list]):
            return False

    for gspan in gold_span_list:
        # not finding gspan equal to any spans in pred_span_list
        if not any([nonnull_span_equal(pspan, gspan) for pspan in pred_span_list]):
            return False

    return True


def gold_has_short_answer(gold_label_list: list, short_non_null_threshold: int = 2) -> bool:
    """
    Gets vote from multi-annotators for judging if there is a short answer.
    :param gold_label_list: list of gold labels
    :param short_non_null_threshold: Require this many non-null short answer annotations
      to count gold as containing a short answer. Defaults to 2 like the original paper.
    """

    if not gold_label_list:
        return False  # Empty list will not have gold short answer

    # We consider if there is a short answer if there is an short answer span or
    #  the yes/no answer is not none.
    gold_has_answer = sum([
        ((not is_null_span_list(label.short_answer_span_list)) or
         (label.yes_no_answer != 'none')) for label in gold_label_list
    ]) >= short_non_null_threshold

    return gold_has_answer


def gold_has_long_answer(gold_label_list: list, long_non_null_threshold: int = 2) -> bool:
    """
    Gets vote from multi-annotators for judging if there is a long answer.
    :param gold_label_list: list of gold labels for judging
    :param long_non_null_threshold: Require this many non-null long answer annotations
      to count gold as containing a long answer. Defaults to 2 like the original paper.
    """

    if not gold_label_list:
        return False  # Empty list will not have gold long answer

    gold_has_answer = (sum([
        not label.long_answer_span.is_null_span()  # long answer not null
        for label in gold_label_list  # for each annotator
    ]) >= long_non_null_threshold)

    return gold_has_answer


# def read_prediction_json_from_file(predictions_path, examples_to_filter_for=None):
#     """Read the prediction json with scores.
#
#     Args:
#       predictions_path: the path for the prediction json.
#       examples_to_filter_for: a set of examples to filter for (can be None, in which
#        case all examples are kept)
#
#     Returns:
#       A dictionary with key = example_id, value = NQInstancePrediction.
#
#     """
#     logging.info('Reading predictions from file: %s', format(predictions_path))
#     with open(predictions_path, 'r') as f:
#         predictions = json.loads(f.read())
#
#     return parse_json_as_predictions(predictions, examples_to_filter_for)


# def parse_json_as_topk_predictions(
#         predictions: dict, examples_to_filter_for: Optional[Set] = None) -> \
#         Dict[int, List[NQLabel]]:
#     nq_pred_dict = {}
#     for topk_predictions_for_single_example in predictions["top_k_best_predictions"]:
#         example_id = None
#         topk_predictions = list()
#         for single_prediction in topk_predictions_for_single_example:
#             prediction = _parse_single_prediction_json(single_prediction)
#             if example_id is None:
#                 example_id = prediction.example_id
#                 if examples_to_filter_for is not None and example_id not in examples_to_filter_for:
#                     logging.debug('Skipping example %s' % example_id)
#                     break
#                 else:
#                     example_id = single_prediction['example_id']
#
#             topk_predictions.append(prediction)
#
#         if topk_predictions:
#             nq_pred_dict[example_id] = topk_predictions
#
#     return nq_pred_dict


# def _parse_single_prediction_json(single_prediction: dict) -> NQLabel:
#     if 'long_answer' in single_prediction:
#         long_span = NQSpan(single_prediction['long_answer']['start_byte'],
#                            single_prediction['long_answer']['end_byte'],
#                            single_prediction['long_answer']['start_token'],
#                            single_prediction['long_answer']['end_token'])
#     else:
#         long_span = NQSpan.null_span()  # Span is null if not presented.
#
#     short_span_list = []
#     if 'short_answers' in single_prediction:
#         for short_item in single_prediction['short_answers']:
#             short_span_list.append(
#                 NQSpan(short_item['start_byte'], short_item['end_byte'],
#                        short_item['start_token'], short_item['end_token']))
#
#     yes_no_answer = 'none'
#     if 'yes_no_answer' in single_prediction:
#         yes_no_answer = single_prediction['yes_no_answer'].lower()
#         if yes_no_answer not in ['yes', 'no', 'none']:
#             raise ValueError('Invalid yes_no_answer value in prediction')
#
#         if yes_no_answer != 'none' and not is_null_span_list(short_span_list):
#             raise ValueError('yes/no prediction and short answers cannot coexist.')
#
#     return NQLabel(
#         example_id=single_prediction['example_id'],
#         long_answer_span=long_span,
#         short_answer_span_list=short_span_list,
#         yes_no_answer=yes_no_answer,
#         long_score=float(single_prediction['long_answer_score']),
#         short_score=float(single_prediction['short_answers_score']))


# def parse_json_as_predictions(predictions: dict, examples_to_filter_for: Optional[Set] = None) -> \
#         Dict[int, NQLabel]:
#     nq_pred_dict = {}
#
#     for single_prediction in predictions['predictions']:
#         if examples_to_filter_for is None or \
#                 single_prediction['example_id'] in examples_to_filter_for:
#             nq_pred_dict[single_prediction['example_id']] = _parse_single_prediction_json(
#                 single_prediction)
#
#     return nq_pred_dict


# def read_annotation_from_one_split(gzipped_input_file, example_ids_to_filter_by=None):
#     """Read annotation from one split of file."""
#     if example_ids_to_filter_by and not isinstance(example_ids_to_filter_by, set):
#         example_ids_to_filter_by = set(example_ids_to_filter_by)
#
#     if isinstance(gzipped_input_file, str):
#         gzipped_input_file = open(gzipped_input_file, mode='rb')
#     logging.info('parsing %s ..... ', gzipped_input_file.name)
#     annotation_dict = {}
#     with GzipFile(fileobj=gzipped_input_file) as input_file:
#         for line in input_file:
#             json_example = json.loads(line)
#             example_id = json_example['example_id']
#
#             if not example_ids_to_filter_by or example_id in example_ids_to_filter_by:
#
#                 # There are multiple annotations for one nq example.
#                 annotation_list = []
#
#                 for annotation in json_example['annotations']:
#                     long_span_rec = annotation['long_answer']
#                     long_span = NQSpan(long_span_rec['start_byte'], long_span_rec['end_byte'],
#                                        long_span_rec['start_token'],
#                                        long_span_rec['end_token'])
#
#                     short_span_list = []
#                     for short_span_rec in annotation['short_answers']:
#                         short_span = NQSpan(
#                             short_span_rec['start_byte'], short_span_rec['end_byte'],
#                             short_span_rec['start_token'], short_span_rec['end_token'])
#                         short_span_list.append(short_span)
#
#                     gold_label = NQLabel(
#                         example_id=example_id,
#                         long_answer_span=long_span,
#                         short_answer_span_list=short_span_list,
#                         long_score=0,
#                         short_score=0,
#                         yes_no_answer=annotation['yes_no_answer'].lower())
#
#                     annotation_list.append(gold_label)
#                 annotation_dict[example_id] = annotation_list
#
#     return annotation_dict
#
#
# def read_annotation(path_name, n_threads=10, example_ids_to_filter_by=None,
#                     read_from_split_fn=read_annotation_from_one_split):
#     """Read annotations with real multiple processes."""
#     input_paths = glob.glob(path_name)
#     pool = multiprocessing.Pool(n_threads)
#     logging.debug("Reading annotation from: {}".format(path_name))
#
#     if example_ids_to_filter_by is not None:
#         example_ids_to_filter_by = list(example_ids_to_filter_by)
#
#     annotation_reader = partial(read_from_split_fn,
#                                 example_ids_to_filter_by=example_ids_to_filter_by)
#     try:
non_informative
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import glob
import json
import logging
import multiprocessing
from gzip import GzipFile

from functools import partial
from typing import Dict, Set, Optional, List


@functools.total_ordering
class NQSpan(object):
    """A class for handling token and byte spans.
       Taken from https://github.com/google-research-datasets/natural-questions/blob/master/eval_utils.py#L60
      The logic is:
      1) if both start_byte !=  -1 and end_byte != -1 then the span is defined
         by byte offsets
      2) else, if start_token != -1 and end_token != -1 then the span is define
         by token offsets
      3) else, this is a null span.
      Null spans means that there is no (long or short) answers.
      If your systems only care about token spans rather than byte spans, set all
      byte spans to -1.
    """
    __slots__ = ['start_byte', 'end_byte', 'start_token', 'end_token', 'score', 'long_score']

    def __init__(self, start_byte: int, end_byte: int, start_token: int, end_token: int,
                 score: Optional[float] = None, long_score: Optional[float] = None, enforce_byte_consistency: bool = True,
                 enforce_token_consistency: bool = True):

        self.start_byte = start_byte
        self.end_byte = end_byte
        self.start_token = start_token
        self.end_token = end_token
        self.score = score
        self.long_score = long_score

        if enforce_byte_consistency:
            if ((start_byte < 0 <= end_byte) or
                    (start_byte >= 0 > end_byte)):
                raise ValueError('Inconsistent Null Spans (Byte): %s' % self)

            if start_byte >= 0 and 0 <= end_byte < start_byte:
                raise ValueError('Invalid byte spans (start_byte > end_byte): %s' % self)

        if enforce_token_consistency:
            if ((start_token < 0 <= end_token) or
                    (start_token >= 0 > end_token)):
                raise ValueError('Inconsistent Null Spans (Token): %s' % self)

            if ((start_token >= 0 and end_token >= 0) and
                    (start_token > end_token)):
                raise ValueError(
                    'Invalid token spans (start_token_idx > end_token_idx): %s' % self)

    def is_null_span(self):
        """A span is a null span if the start and end are both -1."""

        if (self.start_byte < 0 and self.end_byte < 0 and
                    self.start_token < 0 and self.end_token < 0):
            return True
        return False
non_informative
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import glob
import json
import logging
import multiprocessing
from gzip import GzipFile

from functools import partial
from typing import Dict, Set, Optional, List


@functools.total_ordering
class NQSpan(object):
    """A class for handling token and byte spans.
       Taken from https://github.com/google-research-datasets/natural-questions/blob/master/eval_utils.py#L60
      The logic is:
      1) if both start_byte !=  -1 and end_byte != -1 then the span is defined
         by byte offsets
      2) else, if start_token != -1 and end_token != -1 then the span is define
         by token offsets
      3) else, this is a null span.
      Null spans means that there is no (long or short) answers.
      If your systems only care about token spans rather than byte spans, set all
      byte spans to -1.
    """
    __slots__ = ['start_byte', 'end_byte', 'start_token', 'end_token', 'score', 'long_score']

    def __init__(self, start_byte: int, end_byte: int, start_token: int, end_token: int,
                 score: Optional[float] = None, long_score: Optional[float] = None, enforce_byte_consistency: bool = True,
                 enforce_token_consistency: bool = True):

        self.start_byte = start_byte
        self.end_byte = end_byte
        self.start_token = start_token
        self.end_token = end_token
        self.score = score
        self.long_score = long_score

        if enforce_byte_consistency:
            if ((start_byte < 0 <= end_byte) or
                    (start_byte >= 0 > end_byte)):
                raise ValueError('Inconsistent Null Spans (Byte): %s' % self)

            if start_byte >= 0 and 0 <= end_byte < start_byte:
                raise ValueError('Invalid byte spans (start_byte > end_byte): %s' % self)

        if enforce_token_consistency:
            if ((start_token < 0 <= end_token) or
                    (start_token >= 0 > end_token)):
                raise ValueError('Inconsistent Null Spans (Token): %s' % self)

            if ((start_token >= 0 and end_token >= 0) and
                    (start_token > end_token)):
                raise ValueError(
                    'Invalid token spans (start_token_idx > end_token_idx): %s' % self)

    def is_null_span(self):
        """A span is a null span if the start and end are both -1."""

        if (self.start_byte < 0 and self.end_byte < 0 and
                    self.start_token < 0 and self.end_token < 0):
            return True
        return False

    @staticmethod
    def null_span():
        return NQSpan(start_byte=-1, end_byte=-1, start_token=-1, end_token=-1)

    def contains(self, other_span):

        if other_span.start_byte > -1 and self.start_byte > -1 and other_span.end_byte > -1 and self.end_byte > -1:
            # Compare using bytes
            if other_span.start_byte >= self.start_byte and other_span.end_byte <= self.end_byte:
                return True
            else:
                return False
        elif other_span.start_token > -1 and self.start_token > -1 and other_span.end_token > -1 and self.end_token > -1:
            # Compare using tokens
            if other_span.start_token >= self.start_token and other_span.end_token <= self.end_token:
                return True
            else:
                return False
        else:
            return False

    def __str__(self):
        byte_str = 'bytes: [' + str(self.start_byte) + ',' + str(self.end_byte) + ')'
        tok_str = ('tokens: [' + str(self.start_token) + ',' + str(
            self.end_token) + ')')
        if self.score is not None:
            score_str = ', score: %s' % self.score
        else:
            score_str = ''

        return "Span(" + byte_str + ', ' + tok_str + score_str + ")"

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

    def __hash__(self):
        return hash((self.start_byte, self.end_byte, self.start_token, self.end_token, self.score))

    def __lt__(self, other):
        if self.score < other.score:
            return True
        return False

    def __eq__(self, o: object) -> bool:

        if not isinstance(o, NQSpan):
            return False

        for attribute in self.__slots__:
            if not hasattr(o, attribute):
                return False
            elif getattr(o, attribute) != getattr(self, attribute):
                return False

        return True


# A data structure for storing prediction and annotation.
# When a example has multiple annotations, multiple NQLabel will be used.
NQLabel = collections.namedtuple(
    'NQLabel',
    [
        'example_id',  # the unique id for each NQ example.
        'long_answer_span',  # A Span object for long answer.
        'short_answer_span_list',  # A list of Spans for short answer.
        #   Note that In NQ, the short answers
        #   do not need to be in a single span.
        'yes_no_answer',  # Indicate if the short answer is an yes/no answer
        #   The possible values are "yes", "no", "none".
        #   (case insensitive)
        #   If the field is "yes", short_answer_span_list
        #   should be empty or only contain null spans.
        'long_score',  # The prediction score for the long answer prediction.
        'short_score'  # The prediction score for the short answer prediction.
    ])


def is_null_span_list(span_list):
    """Returns true iff all spans in span_list are null or span_list is empty."""
    if not span_list or all([span.is_null_span() for span in span_list]):
        return True
    return False


def nonnull_span_equal(span_a: NQSpan, span_b: NQSpan):
    """Given two spans, return if they are equal.

    Args:
      span_a: a Span object.
      span_b: a Span object.  Only compare non-null spans. First, if the bytes are
        not negative, compare byte offsets, Otherwise, compare token offsets.

    Returns:
      True or False
    """
    assert isinstance(span_a, NQSpan)
    assert isinstance(span_b, NQSpan)
    assert not span_a.is_null_span()
    assert not span_b.is_null_span()

    # if byte offsets are not negative, compare byte offsets
    if ((span_a.start_byte >= 0 and span_a.end_byte >= 0) and
            (span_b.start_byte >= 0 and span_b.end_byte >= 0)):

        if ((span_a.start_byte == span_b.start_byte) and
                (span_a.end_byte == span_b.end_byte)):
            return True

    # if token offsets are not negative, compare token offsets
    if ((span_a.start_token >= 0 and span_a.end_token >= 0) and
            (span_b.start_token >= 0 and span_b.end_token >= 0)):

        if ((span_a.start_token == span_b.start_token) and
                (span_a.end_token == span_b.end_token)):
            return True

    return False


def span_set_equal(gold_span_list, pred_span_list):
    """Make the spans are completely equal besides null spans."""

    gold_span_list = [span for span in gold_span_list if not span.is_null_span()]
    pred_span_list = [span for span in pred_span_list if not span.is_null_span()]

    for pspan in pred_span_list:
        # not finding pspan equal to any spans in gold_span_list
        if not any([nonnull_span_equal(pspan, gspan) for gspan in gold_span_list]):
            return False

    for gspan in gold_span_list:
        # not finding gspan equal to any spans in pred_span_list
        if not any([nonnull_span_equal(pspan, gspan) for pspan in pred_span_list]):
            return False

    return True


def gold_has_short_answer(gold_label_list: list, short_non_null_threshold: int = 2) -> bool:
    """
    Gets vote from multi-annotators for judging if there is a short answer.
    :param gold_label_list: list of gold labels
    :param short_non_null_threshold: Require this many non-null short answer annotations
      to count gold as containing a short answer. Defaults to 2 like the original paper.
    """

    if not gold_label_list:
        return False  # Empty list will not have gold short answer

    # We consider if there is a short answer if there is an short answer span or
non_informative
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import glob
import json
import logging
import multiprocessing
from gzip import GzipFile

from functools import partial
from typing import Dict, Set, Optional, List


@functools.total_ordering
class NQSpan(object):
    """A class for handling token and byte spans.
       Taken from https://github.com/google-research-datasets/natural-questions/blob/master/eval_utils.py#L60
      The logic is:
      1) if both start_byte !=  -1 and end_byte != -1 then the span is defined
         by byte offsets
      2) else, if start_token != -1 and end_token != -1 then the span is define
         by token offsets
      3) else, this is a null span.
      Null spans means that there is no (long or short) answers.
      If your systems only care about token spans rather than byte spans, set all
      byte spans to -1.
    """
    __slots__ = ['start_byte', 'end_byte', 'start_token', 'end_token', 'score', 'long_score']

    def __init__(self, start_byte: int, end_byte: int, start_token: int, end_token: int,
                 score: Optional[float] = None, long_score: Optional[float] = None, enforce_byte_consistency: bool = True,
                 enforce_token_consistency: bool = True):

        self.start_byte = start_byte
        self.end_byte = end_byte
        self.start_token = start_token
        self.end_token = end_token
        self.score = score
        self.long_score = long_score

        if enforce_byte_consistency:
            if ((start_byte < 0 <= end_byte) or
                    (start_byte >= 0 > end_byte)):
                raise ValueError('Inconsistent Null Spans (Byte): %s' % self)

            if start_byte >= 0 and 0 <= end_byte < start_byte:
                raise ValueError('Invalid byte spans (start_byte > end_byte): %s' % self)

        if enforce_token_consistency:
            if ((start_token < 0 <= end_token) or
                    (start_token >= 0 > end_token)):
                raise ValueError('Inconsistent Null Spans (Token): %s' % self)

            if ((start_token >= 0 and end_token >= 0) and
                    (start_token > end_token)):
                raise ValueError(
                    'Invalid token spans (start_token_idx > end_token_idx): %s' % self)

    def is_null_span(self):
        """A span is a null span if the start and end are both -1."""

        if (self.start_byte < 0 and self.end_byte < 0 and
                    self.start_token < 0 and self.end_token < 0):
            return True
        return False

    @staticmethod
    def null_span():
        return NQSpan(start_byte=-1, end_byte=-1, start_token=-1, end_token=-1)

    def contains(self, other_span):

        if other_span.start_byte > -1 and self.start_byte > -1 and other_span.end_byte > -1 and self.end_byte > -1:
            # Compare using bytes
            if other_span.start_byte >= self.start_byte and other_span.end_byte <= self.end_byte:
                return True
            else:
                return False
        elif other_span.start_token > -1 and self.start_token > -1 and other_span.end_token > -1 and self.end_token > -1:
            # Compare using tokens
            if other_span.start_token >= self.start_token and other_span.end_token <= self.end_token:
                return True
            else:
                return False
        else:
            return False

    def __str__(self):
        byte_str = 'bytes: [' + str(self.start_byte) + ',' + str(self.end_byte) + ')'
        tok_str = ('tokens: [' + str(self.start_token) + ',' + str(
            self.end_token) + ')')
        if self.score is not None:
            score_str = ', score: %s' % self.score
        else:
            score_str = ''

        return "Span(" + byte_str + ', ' + tok_str + score_str + ")"

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

    def __hash__(self):
        return hash((self.start_byte, self.end_byte, self.start_token, self.end_token, self.score))

    def __lt__(self, other):
        if self.score < other.score:
            return True
        return False

    def __eq__(self, o: object) -> bool:

        if not isinstance(o, NQSpan):
            return False

        for attribute in self.__slots__:
            if not hasattr(o, attribute):
                return False
            elif getattr(o, attribute) != getattr(self, attribute):
                return False

        return True


# A data structure for storing prediction and annotation.
# When a example has multiple annotations, multiple NQLabel will be used.
NQLabel = collections.namedtuple(
    'NQLabel',
    [
        'example_id',  # the unique id for each NQ example.
        'long_answer_span',  # A Span object for long answer.
        'short_answer_span_list',  # A list of Spans for short answer.
        #   Note that In NQ, the short answers
        #   do not need to be in a single span.
        'yes_no_answer',  # Indicate if the short answer is an yes/no answer
        #   The possible values are "yes", "no", "none".
        #   (case insensitive)
        #   If the field is "yes", short_answer_span_list
        #   should be empty or only contain null spans.
        'long_score',  # The prediction score for the long answer prediction.
        'short_score'  # The prediction score for the short answer prediction.
    ])


def is_null_span_list(span_list):
    """Returns true iff all spans in span_list are null or span_list is empty."""
    if not span_list or all([span.is_null_span() for span in span_list]):
        return True
    return False


def nonnull_span_equal(span_a: NQSpan, span_b: NQSpan):
    """Given two spans, return if they are equal.

    Args:
      span_a: a Span object.
      span_b: a Span object.  Only compare non-null spans. First, if the bytes are
        not negative, compare byte offsets, Otherwise, compare token offsets.

    Returns:
      True or False
    """
    assert isinstance(span_a, NQSpan)
    assert isinstance(span_b, NQSpan)
    assert not span_a.is_null_span()
    assert not span_b.is_null_span()

    # if byte offsets are not negative, compare byte offsets
    if ((span_a.start_byte >= 0 and span_a.end_byte >= 0) and
            (span_b.start_byte >= 0 and span_b.end_byte >= 0)):

        if ((span_a.start_byte == span_b.start_byte) and
                (span_a.end_byte == span_b.end_byte)):
            return True

    # if token offsets are not negative, compare token offsets
    if ((span_a.start_token >= 0 and span_a.end_token >= 0) and
            (span_b.start_token >= 0 and span_b.end_token >= 0)):

        if ((span_a.start_token == span_b.start_token) and
                (span_a.end_token == span_b.end_token)):
            return True

    return False


def span_set_equal(gold_span_list, pred_span_list):
    """Make the spans are completely equal besides null spans."""

    gold_span_list = [span for span in gold_span_list if not span.is_null_span()]
    pred_span_list = [span for span in pred_span_list if not span.is_null_span()]

    for pspan in pred_span_list:
        # not finding pspan equal to any spans in gold_span_list
        if not any([nonnull_span_equal(pspan, gspan) for gspan in gold_span_list]):
            return False

    for gspan in gold_span_list:
        # not finding gspan equal to any spans in pred_span_list
        if not any([nonnull_span_equal(pspan, gspan) for pspan in pred_span_list]):
            return False

    return True


def gold_has_short_answer(gold_label_list: list, short_non_null_threshold: int = 2) -> bool:
    """
    Gets vote from multi-annotators for judging if there is a short answer.
    :param gold_label_list: list of gold labels
    :param short_non_null_threshold: Require this many non-null short answer annotations
      to count gold as containing a short answer. Defaults to 2 like the original paper.
    """

    if not gold_label_list:
        return False  # Empty list will not have gold short answer

    # We consider if there is a short answer if there is an short answer span or
    #  the yes/no answer is not none.
    gold_has_answer = sum([
        ((not is_null_span_list(label.short_answer_span_list)) or
         (label.yes_no_answer != 'none')) for label in gold_label_list
    ]) >= short_non_null_threshold

    return gold_has_answer


def gold_has_long_answer(gold_label_list: list, long_non_null_threshold: int = 2) -> bool:
    """
    Gets vote from multi-annotators for judging if there is a long answer.
    :param gold_label_list: list of gold labels for judging
    :param long_non_null_threshold: Require this many non-null long answer annotations
      to count gold as containing a long answer. Defaults to 2 like the original paper.
    """

    if not gold_label_list:
        return False  # Empty list will not have gold long answer

    gold_has_answer = (sum([
        not label.long_answer_span.is_null_span()  # long answer not null
        for label in gold_label_list  # for each annotator
    ]) >= long_non_null_threshold)

    return gold_has_answer


# def read_prediction_json_from_file(predictions_path, examples_to_filter_for=None):
#     """Read the prediction json with scores.
#
#     Args:
#       predictions_path: the path for the prediction json.
#       examples_to_filter_for: a set of examples to filter for (can be None, in which
#        case all examples are kept)
#
#     Returns:
#       A dictionary with key = example_id, value = NQInstancePrediction.
#
#     """
#     logging.info('Reading predictions from file: %s', format(predictions_path))
#     with open(predictions_path, 'r') as f:
#         predictions = json.loads(f.read())
#
#     return parse_json_as_predictions(predictions, examples_to_filter_for)


# def parse_json_as_topk_predictions(
#         predictions: dict, examples_to_filter_for: Optional[Set] = None) -> \
#         Dict[int, List[NQLabel]]:
#     nq_pred_dict = {}
#     for topk_predictions_for_single_example in predictions["top_k_best_predictions"]:
#         example_id = None
#         topk_predictions = list()
#         for single_prediction in topk_predictions_for_single_example:
#             prediction = _parse_single_prediction_json(single_prediction)
#             if example_id is None:
#                 example_id = prediction.example_id
#                 if examples_to_filter_for is not None and example_id not in examples_to_filter_for:
#                     logging.debug('Skipping example %s' % example_id)
#                     break
#                 else:
#                     example_id = single_prediction['example_id']
#
#             topk_predictions.append(prediction)
#
#         if topk_predictions:
#             nq_pred_dict[example_id] = topk_predictions
#
#     return nq_pred_dict


# def _parse_single_prediction_json(single_prediction: dict) -> NQLabel:
#     if 'long_answer' in single_prediction:
#         long_span = NQSpan(single_prediction['long_answer']['start_byte'],
#                            single_prediction['long_answer']['end_byte'],
#                            single_prediction['long_answer']['start_token'],
#                            single_prediction['long_answer']['end_token'])
#     else:
#         long_span = NQSpan.null_span()  # Span is null if not presented.
#
#     short_span_list = []
non_informative
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import glob
import json
import logging
import multiprocessing
from gzip import GzipFile

from functools import partial
from typing import Dict, Set, Optional, List


@functools.total_ordering
class NQSpan(object):
    """A class for handling token and byte spans.
       Taken from https://github.com/google-research-datasets/natural-questions/blob/master/eval_utils.py#L60
      The logic is:
      1) if both start_byte !=  -1 and end_byte != -1 then the span is defined
         by byte offsets
      2) else, if start_token != -1 and end_token != -1 then the span is define
         by token offsets
      3) else, this is a null span.
      Null spans means that there is no (long or short) answers.
      If your systems only care about token spans rather than byte spans, set all
      byte spans to -1.
    """
    __slots__ = ['start_byte', 'end_byte', 'start_token', 'end_token', 'score', 'long_score']

non_informative
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import glob
import json
import logging
import multiprocessing
from gzip import GzipFile

from functools import partial
from typing import Dict, Set, Optional, List


@functools.total_ordering
class NQSpan(object):
    """A class for handling token and byte spans.
       Taken from https://github.com/google-research-datasets/natural-questions/blob/master/eval_utils.py#L60
      The logic is:
      1) if both start_byte !=  -1 and end_byte != -1 then the span is defined
         by byte offsets
      2) else, if start_token != -1 and end_token != -1 then the span is define
         by token offsets
      3) else, this is a null span.
      Null spans means that there is no (long or short) answers.
      If your systems only care about token spans rather than byte spans, set all
      byte spans to -1.
    """
    __slots__ = ['start_byte', 'end_byte', 'start_token', 'end_token', 'score', 'long_score']

    def __init__(self, start_byte: int, end_byte: int, start_token: int, end_token: int,
                 score: Optional[float] = None, long_score: Optional[float] = None, enforce_byte_consistency: bool = True,
                 enforce_token_consistency: bool = True):

        self.start_byte = start_byte
        self.end_byte = end_byte
        self.start_token = start_token
        self.end_token = end_token
        self.score = score
        self.long_score = long_score

        if enforce_byte_consistency:
            if ((start_byte < 0 <= end_byte) or
                    (start_byte >= 0 > end_byte)):
                raise ValueError('Inconsistent Null Spans (Byte): %s' % self)

            if start_byte >= 0 and 0 <= end_byte < start_byte:
                raise ValueError('Invalid byte spans (start_byte > end_byte): %s' % self)

        if enforce_token_consistency:
            if ((start_token < 0 <= end_token) or
                    (start_token >= 0 > end_token)):
                raise ValueError('Inconsistent Null Spans (Token): %s' % self)

            if ((start_token >= 0 and end_token >= 0) and
                    (start_token > end_token)):
                raise ValueError(
                    'Invalid token spans (start_token_idx > end_token_idx): %s' % self)

    def is_null_span(self):
        """A span is a null span if the start and end are both -1."""

        if (self.start_byte < 0 and self.end_byte < 0 and
                    self.start_token < 0 and self.end_token < 0):
            return True
        return False

    @staticmethod
    def null_span():
        return NQSpan(start_byte=-1, end_byte=-1, start_token=-1, end_token=-1)

    def contains(self, other_span):

        if other_span.start_byte > -1 and self.start_byte > -1 and other_span.end_byte > -1 and self.end_byte > -1:
            # Compare using bytes
            if other_span.start_byte >= self.start_byte and other_span.end_byte <= self.end_byte:
                return True
            else:
                return False
        elif other_span.start_token > -1 and self.start_token > -1 and other_span.end_token > -1 and self.end_token > -1:
            # Compare using tokens
            if other_span.start_token >= self.start_token and other_span.end_token <= self.end_token:
                return True
            else:
                return False
        else:
            return False

    def __str__(self):
        byte_str = 'bytes: [' + str(self.start_byte) + ',' + str(self.end_byte) + ')'
        tok_str = ('tokens: [' + str(self.start_token) + ',' + str(
            self.end_token) + ')')
        if self.score is not None:
            score_str = ', score: %s' % self.score
        else:
            score_str = ''

        return "Span(" + byte_str + ', ' + tok_str + score_str + ")"

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

    def __hash__(self):
        return hash((self.start_byte, self.end_byte, self.start_token, self.end_token, self.score))

    def __lt__(self, other):
        if self.score < other.score:
            return True
        return False

    def __eq__(self, o: object) -> bool:

        if not isinstance(o, NQSpan):
            return False

        for attribute in self.__slots__:
            if not hasattr(o, attribute):
                return False
            elif getattr(o, attribute) != getattr(self, attribute):
                return False

        return True


# A data structure for storing prediction and annotation.
# When a example has multiple annotations, multiple NQLabel will be used.
NQLabel = collections.namedtuple(
    'NQLabel',
    [
        'example_id',  # the unique id for each NQ example.
        'long_answer_span',  # A Span object for long answer.
        'short_answer_span_list',  # A list of Spans for short answer.
        #   Note that In NQ, the short answers
        #   do not need to be in a single span.
        'yes_no_answer',  # Indicate if the short answer is an yes/no answer
        #   The possible values are "yes", "no", "none".
        #   (case insensitive)
        #   If the field is "yes", short_answer_span_list
        #   should be empty or only contain null spans.
        'long_score',  # The prediction score for the long answer prediction.
        'short_score'  # The prediction score for the short answer prediction.
    ])


def is_null_span_list(span_list):
    """Returns true iff all spans in span_list are null or span_list is empty."""
    if not span_list or all([span.is_null_span() for span in span_list]):
        return True
    return False


def nonnull_span_equal(span_a: NQSpan, span_b: NQSpan):
    """Given two spans, return if they are equal.

    Args:
      span_a: a Span object.
      span_b: a Span object.  Only compare non-null spans. First, if the bytes are
        not negative, compare byte offsets, Otherwise, compare token offsets.

    Returns:
      True or False
    """
    assert isinstance(span_a, NQSpan)
    assert isinstance(span_b, NQSpan)
    assert not span_a.is_null_span()
    assert not span_b.is_null_span()

    # if byte offsets are not negative, compare byte offsets
    if ((span_a.start_byte >= 0 and span_a.end_byte >= 0) and
            (span_b.start_byte >= 0 and span_b.end_byte >= 0)):

        if ((span_a.start_byte == span_b.start_byte) and
                (span_a.end_byte == span_b.end_byte)):
            return True

    # if token offsets are not negative, compare token offsets
    if ((span_a.start_token >= 0 and span_a.end_token >= 0) and
            (span_b.start_token >= 0 and span_b.end_token >= 0)):

        if ((span_a.start_token == span_b.start_token) and
                (span_a.end_token == span_b.end_token)):
            return True

    return False


def span_set_equal(gold_span_list, pred_span_list):
    """Make the spans are completely equal besides null spans."""

    gold_span_list = [span for span in gold_span_list if not span.is_null_span()]
    pred_span_list = [span for span in pred_span_list if not span.is_null_span()]

    for pspan in pred_span_list:
        # not finding pspan equal to any spans in gold_span_list
        if not any([nonnull_span_equal(pspan, gspan) for gspan in gold_span_list]):
            return False

    for gspan in gold_span_list:
        # not finding gspan equal to any spans in pred_span_list
        if not any([nonnull_span_equal(pspan, gspan) for pspan in pred_span_list]):
random
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import glob
import json
import logging
import multiprocessing
from gzip import GzipFile

from functools import partial
from typing import Dict, Set, Optional, List


@functools.total_ordering
class NQSpan(object):
    """A class for handling token and byte spans.
       Taken from https://github.com/google-research-datasets/natural-questions/blob/master/eval_utils.py#L60
      The logic is:
      1) if both start_byte !=  -1 and end_byte != -1 then the span is defined
         by byte offsets
      2) else, if start_token != -1 and end_token != -1 then the span is define
         by token offsets
      3) else, this is a null span.
      Null spans means that there is no (long or short) answers.
      If your systems only care about token spans rather than byte spans, set all
      byte spans to -1.
    """
    __slots__ = ['start_byte', 'end_byte', 'start_token', 'end_token', 'score', 'long_score']

    def __init__(self, start_byte: int, end_byte: int, start_token: int, end_token: int,
                 score: Optional[float] = None, long_score: Optional[float] = None, enforce_byte_consistency: bool = True,
                 enforce_token_consistency: bool = True):

        self.start_byte = start_byte
        self.end_byte = end_byte
        self.start_token = start_token
        self.end_token = end_token
        self.score = score
        self.long_score = long_score

        if enforce_byte_consistency:
            if ((start_byte < 0 <= end_byte) or
                    (start_byte >= 0 > end_byte)):
                raise ValueError('Inconsistent Null Spans (Byte): %s' % self)

            if start_byte >= 0 and 0 <= end_byte < start_byte:
                raise ValueError('Invalid byte spans (start_byte > end_byte): %s' % self)

        if enforce_token_consistency:
            if ((start_token < 0 <= end_token) or
                    (start_token >= 0 > end_token)):
                raise ValueError('Inconsistent Null Spans (Token): %s' % self)

            if ((start_token >= 0 and end_token >= 0) and
                    (start_token > end_token)):
                raise ValueError(
                    'Invalid token spans (start_token_idx > end_token_idx): %s' % self)

    def is_null_span(self):
        """A span is a null span if the start and end are both -1."""

        if (self.start_byte < 0 and self.end_byte < 0 and
                    self.start_token < 0 and self.end_token < 0):
            return True
        return False

    @staticmethod
    def null_span():
        return NQSpan(start_byte=-1, end_byte=-1, start_token=-1, end_token=-1)

    def contains(self, other_span):

        if other_span.start_byte > -1 and self.start_byte > -1 and other_span.end_byte > -1 and self.end_byte > -1:
            # Compare using bytes
            if other_span.start_byte >= self.start_byte and other_span.end_byte <= self.end_byte:
                return True
            else:
                return False
        elif other_span.start_token > -1 and self.start_token > -1 and other_span.end_token > -1 and self.end_token > -1:
            # Compare using tokens
            if other_span.start_token >= self.start_token and other_span.end_token <= self.end_token:
                return True
            else:
                return False
        else:
            return False

    def __str__(self):
        byte_str = 'bytes: [' + str(self.start_byte) + ',' + str(self.end_byte) + ')'
        tok_str = ('tokens: [' + str(self.start_token) + ',' + str(
            self.end_token) + ')')
        if self.score is not None:
            score_str = ', score: %s' % self.score
        else:
            score_str = ''

        return "Span(" + byte_str + ', ' + tok_str + score_str + ")"

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

    def __hash__(self):
        return hash((self.start_byte, self.end_byte, self.start_token, self.end_token, self.score))

    def __lt__(self, other):
        if self.score < other.score:
            return True
        return False

    def __eq__(self, o: object) -> bool:

        if not isinstance(o, NQSpan):
            return False

        for attribute in self.__slots__:
            if not hasattr(o, attribute):
                return False
            elif getattr(o, attribute) != getattr(self, attribute):
                return False

        return True


# A data structure for storing prediction and annotation.
# When a example has multiple annotations, multiple NQLabel will be used.
NQLabel = collections.namedtuple(
    'NQLabel',
    [
        'example_id',  # the unique id for each NQ example.
        'long_answer_span',  # A Span object for long answer.
        'short_answer_span_list',  # A list of Spans for short answer.
        #   Note that In NQ, the short answers
        #   do not need to be in a single span.
        'yes_no_answer',  # Indicate if the short answer is an yes/no answer
        #   The possible values are "yes", "no", "none".
        #   (case insensitive)
        #   If the field is "yes", short_answer_span_list
        #   should be empty or only contain null spans.
        'long_score',  # The prediction score for the long answer prediction.
        'short_score'  # The prediction score for the short answer prediction.
    ])


def is_null_span_list(span_list):
    """Returns true iff all spans in span_list are null or span_list is empty."""
    if not span_list or all([span.is_null_span() for span in span_list]):
        return True
    return False


def nonnull_span_equal(span_a: NQSpan, span_b: NQSpan):
    """Given two spans, return if they are equal.

    Args:
      span_a: a Span object.
      span_b: a Span object.  Only compare non-null spans. First, if the bytes are
        not negative, compare byte offsets, Otherwise, compare token offsets.

    Returns:
      True or False
    """
    assert isinstance(span_a, NQSpan)
    assert isinstance(span_b, NQSpan)
    assert not span_a.is_null_span()
    assert not span_b.is_null_span()

    # if byte offsets are not negative, compare byte offsets
    if ((span_a.start_byte >= 0 and span_a.end_byte >= 0) and
            (span_b.start_byte >= 0 and span_b.end_byte >= 0)):

        if ((span_a.start_byte == span_b.start_byte) and
                (span_a.end_byte == span_b.end_byte)):
            return True

    # if token offsets are not negative, compare token offsets
    if ((span_a.start_token >= 0 and span_a.end_token >= 0) and
            (span_b.start_token >= 0 and span_b.end_token >= 0)):

        if ((span_a.start_token == span_b.start_token) and
                (span_a.end_token == span_b.end_token)):
            return True

    return False


def span_set_equal(gold_span_list, pred_span_list):
    """Make the spans are completely equal besides null spans."""

    gold_span_list = [span for span in gold_span_list if not span.is_null_span()]
    pred_span_list = [span for span in pred_span_list if not span.is_null_span()]

    for pspan in pred_span_list:
        # not finding pspan equal to any spans in gold_span_list
        if not any([nonnull_span_equal(pspan, gspan) for gspan in gold_span_list]):
            return False

    for gspan in gold_span_list:
        # not finding gspan equal to any spans in pred_span_list
        if not any([nonnull_span_equal(pspan, gspan) for pspan in pred_span_list]):
            return False

    return True


def gold_has_short_answer(gold_label_list: list, short_non_null_threshold: int = 2) -> bool:
    """
    Gets vote from multi-annotators for judging if there is a short answer.
    :param gold_label_list: list of gold labels
    :param short_non_null_threshold: Require this many non-null short answer annotations
      to count gold as containing a short answer. Defaults to 2 like the original paper.
    """

    if not gold_label_list:
        return False  # Empty list will not have gold short answer

    # We consider if there is a short answer if there is an short answer span or
    #  the yes/no answer is not none.
    gold_has_answer = sum([
        ((not is_null_span_list(label.short_answer_span_list)) or
         (label.yes_no_answer != 'none')) for label in gold_label_list
    ]) >= short_non_null_threshold

random
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import glob
import json
import logging
import multiprocessing
from gzip import GzipFile

from functools import partial
from typing import Dict, Set, Optional, List


@functools.total_ordering
class NQSpan(object):
    """A class for handling token and byte spans.
       Taken from https://github.com/google-research-datasets/natural-questions/blob/master/eval_utils.py#L60
      The logic is:
      1) if both start_byte !=  -1 and end_byte != -1 then the span is defined
         by byte offsets
      2) else, if start_token != -1 and end_token != -1 then the span is define
         by token offsets
      3) else, this is a null span.
      Null spans means that there is no (long or short) answers.
      If your systems only care about token spans rather than byte spans, set all
      byte spans to -1.
    """
    __slots__ = ['start_byte', 'end_byte', 'start_token', 'end_token', 'score', 'long_score']

    def __init__(self, start_byte: int, end_byte: int, start_token: int, end_token: int,
                 score: Optional[float] = None, long_score: Optional[float] = None, enforce_byte_consistency: bool = True,
                 enforce_token_consistency: bool = True):

        self.start_byte = start_byte
        self.end_byte = end_byte
        self.start_token = start_token
        self.end_token = end_token
        self.score = score
        self.long_score = long_score

        if enforce_byte_consistency:
            if ((start_byte < 0 <= end_byte) or
                    (start_byte >= 0 > end_byte)):
                raise ValueError('Inconsistent Null Spans (Byte): %s' % self)

            if start_byte >= 0 and 0 <= end_byte < start_byte:
                raise ValueError('Invalid byte spans (start_byte > end_byte): %s' % self)

random
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import glob
import json
import logging
import multiprocessing
from gzip import GzipFile

from functools import partial
from typing import Dict, Set, Optional, List


@functools.total_ordering
class NQSpan(object):
    """A class for handling token and byte spans.
       Taken from https://github.com/google-research-datasets/natural-questions/blob/master/eval_utils.py#L60
      The logic is:
      1) if both start_byte !=  -1 and end_byte != -1 then the span is defined
         by byte offsets
      2) else, if start_token != -1 and end_token != -1 then the span is define
         by token offsets
      3) else, this is a null span.
      Null spans means that there is no (long or short) answers.
      If your systems only care about token spans rather than byte spans, set all
      byte spans to -1.
    """
    __slots__ = ['start_byte', 'end_byte', 'start_token', 'end_token', 'score', 'long_score']

    def __init__(self, start_byte: int, end_byte: int, start_token: int, end_token: int,
                 score: Optional[float] = None, long_score: Optional[float] = None, enforce_byte_consistency: bool = True,
                 enforce_token_consistency: bool = True):

        self.start_byte = start_byte
        self.end_byte = end_byte
        self.start_token = start_token
        self.end_token = end_token
        self.score = score
        self.long_score = long_score

        if enforce_byte_consistency:
            if ((start_byte < 0 <= end_byte) or
                    (start_byte >= 0 > end_byte)):
                raise ValueError('Inconsistent Null Spans (Byte): %s' % self)

            if start_byte >= 0 and 0 <= end_byte < start_byte:
                raise ValueError('Invalid byte spans (start_byte > end_byte): %s' % self)

        if enforce_token_consistency:
            if ((start_token < 0 <= end_token) or
                    (start_token >= 0 > end_token)):
                raise ValueError('Inconsistent Null Spans (Token): %s' % self)

            if ((start_token >= 0 and end_token >= 0) and
                    (start_token > end_token)):
                raise ValueError(
                    'Invalid token spans (start_token_idx > end_token_idx): %s' % self)

    def is_null_span(self):
        """A span is a null span if the start and end are both -1."""

        if (self.start_byte < 0 and self.end_byte < 0 and
                    self.start_token < 0 and self.end_token < 0):
            return True
        return False

    @staticmethod
    def null_span():
        return NQSpan(start_byte=-1, end_byte=-1, start_token=-1, end_token=-1)

    def contains(self, other_span):

        if other_span.start_byte > -1 and self.start_byte > -1 and other_span.end_byte > -1 and self.end_byte > -1:
            # Compare using bytes
            if other_span.start_byte >= self.start_byte and other_span.end_byte <= self.end_byte:
                return True
            else:
                return False
        elif other_span.start_token > -1 and self.start_token > -1 and other_span.end_token > -1 and self.end_token > -1:
            # Compare using tokens
            if other_span.start_token >= self.start_token and other_span.end_token <= self.end_token:
                return True
            else:
                return False
        else:
            return False

    def __str__(self):
        byte_str = 'bytes: [' + str(self.start_byte) + ',' + str(self.end_byte) + ')'
        tok_str = ('tokens: [' + str(self.start_token) + ',' + str(
            self.end_token) + ')')
        if self.score is not None:
            score_str = ', score: %s' % self.score
random
    pred_span_list = [span for span in pred_span_list if not span.is_null_span()]     assert isinstance(span_a, NQSpan)     assert isinstance(span_b, NQSpan)     assert not span_a.is_null_span()     assert not span_b.is_null_span()         if not any([nonnull_span_equal(pspan, gspan) for gspan in gold_span_list]):         return NQSpan(start_byte=-1, end_byte=-1, start_token=-1, end_token=-1)         if not any([nonnull_span_equal(pspan, gspan) for pspan in pred_span_list]):         if not isinstance(o, NQSpan):     if not span_list or all([span.is_null_span() for span in span_list]):         ((not is_null_span_list(label.short_answer_span_list)) or     gold_span_list = [span for span in gold_span_list if not span.is_null_span()]         return self.__str__() #         dict_list = pool.map(annotation_reader, input_paths)     #  the yes/no answer is not none. #     if 'short_answers' in single_prediction:     def __init__(self, start_byte: int, end_byte: int, start_token: int, end_token: int,             return False     return gold_has_answer         if enforce_token_consistency:         else: