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/metrics/__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/__init__.py METASEP

oneqa/mrc/metrics/tydi_f1/tydi_eval.py METASEP

# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
infile
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
infile
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
infile
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
infile
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
    macro_f1 = eval_utils.safe_divide(
        2 * macro_precision * macro_recall,
        macro_precision + macro_recall)

    return collections.OrderedDict({
        prefix + 'n': len(answer_stats),
        prefix + 'f1': macro_f1,
        prefix + 'precision': macro_precision,
        prefix + 'recall': macro_recall
    })


def compute_final_f1(passage_answer_stats, minimal_answer_stats):
    """Computes overall F1 given passage and minimal answers, ignoring scores.

    Note: this assumes that the answers have been thresholded based on scores.

    Arguments:
       passage_answer_stats: List of passage answer scores.
       minimal_answer_stats: List of minimal answer scores.


    Returns:
       Dictionary of name (string) -> score.
    """
    scores = compute_macro_f1(passage_answer_stats, prefix='passage-answer-')
    scores.update(compute_macro_f1(
        minimal_answer_stats, prefix='minimal-answer-'))
    return scores


def compute_pr_curves(answer_stats, targets=None):
    """Computes PR curve and returns R@P for specific targets.

    The values are computed as follows: find the (precision, recall) point
    with maximum recall and where precision > target.

    This is only relevant if you return the system scores in your predictions.
    You may find this useful when attempting to tune the threshold for your
    system on the dev set before requesting an evaluation on the test set
    via the leaderboard.

    Arguments:
      answer_stats: List of statistic tuples from the answer scores.
      targets (None): List of precision thresholds to target.

    Returns:
      List of table with rows: [target, r, p, score].
    """
    total_f1 = 0
    total_has_pred = 0
    total_has_gold = 0

    # Count the number of gold annotations.
    for has_gold, _, _, _ in answer_stats:
        total_has_gold += has_gold

    # Keep track of the point of maximum recall for each target.

    max_recall = [0 for _ in targets]
    max_precision = [0 for _ in targets]
    max_scores = [0.0 for _ in targets]

    # Only keep track of unique thresholds in this dictionary.
    scores_to_stats = collections.OrderedDict()

    # Loop through every possible threshold and compute precision + recall.
    for has_gold, has_pred, is_correct_or_f1, score in answer_stats:
        if isinstance(is_correct_or_f1, tuple):
            _, _, f1 = is_correct_or_f1
        else:
            f1 = is_correct_or_f1
        total_f1 += f1
        total_has_pred += has_pred

        precision = eval_utils.safe_divide(total_f1, total_has_pred)
        recall = eval_utils.safe_divide(total_f1, total_has_gold)

        # If there are any ties, this will be updated multiple times until the
        # ties are all counted.
        scores_to_stats[score] = [precision, recall]

    best_f1 = 0.0
    best_precision = 0.0
    best_recall = 0.0
    best_threshold = 0.0

    for threshold, (precision, recall) in scores_to_stats.items():
        # Match the thresholds to the find the closest precision above some target.
        for t, target in enumerate(targets):
            if precision >= target and recall > max_recall[t]:
                max_recall[t] = recall
                max_precision[t] = precision
                max_scores[t] = threshold

        # Compute optimal threshold.
        f1 = eval_utils.safe_divide(2 * precision * recall, precision + recall)
        if f1 > best_f1:
            best_f1 = f1
            best_precision = precision
            best_recall = recall
            best_threshold = threshold

    return ((best_f1, best_precision, best_recall, best_threshold),
            list(zip(targets, max_recall, max_precision, max_scores)))


def print_r_at_p_table(answer_stats):
    """Pretty prints the R@P table for default targets."""
    opt_result, pr_table = compute_pr_curves(
        answer_stats, targets=[0.5, 0.75, 0.9])
    f1, precision, recall, threshold = opt_result
    print('Optimal threshold: {:.5}'.format(threshold))
    print(' F1     /  P      /  R')
    print('{: >7.2%} / {: >7.2%} / {: >7.2%}'.format(f1, precision, recall))
    for target, recall, precision, row in pr_table:
        print('R@P={}: {:.2%} (actual p={:.2%}, score threshold={:.4})'.format(
            target, recall, precision, row))


def get_metrics_as_dict(gold_path, prediction_path, passage_non_null_threshold, minimal_non_null_threshold, verbose):
    """Library version of the end-to-end evaluation.

    Arguments:
      gold_path: Path to a single JSONL data. Could be gzipped or not.
      prediction_path: Path to the JSONL file of prediction data.

    Returns:
      metrics: A dictionary mapping string names to metric scores.
    """

    tydi_gold_dict = eval_utils.read_annotation(gold_path)
    tydi_pred_dict = eval_utils.read_prediction_jsonl(prediction_path)

    passage_answer_stats, minimal_answer_stats = score_answers(
        tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose)

    return get_metrics_with_answer_stats(
        passage_answer_stats, minimal_answer_stats)


def get_metrics_with_answer_stats(passage_answer_stats, minimal_answer_stats):
    """Generate metrics dict using passage and minimal answer stats."""

    def _get_metric_dict(answer_stats, prefix=''):
        """Compute all metrics for a set of answer statistics."""
        opt_result, pr_table = compute_pr_curves(
            answer_stats, targets=[0.5, 0.75, 0.9])
        f1, precision, recall, threshold = opt_result
        metrics = collections.OrderedDict({
            'best-threshold-f1': f1,
            'best-threshold-precision': precision,
            'best-threshold-recall': recall,
            'best-threshold': threshold,
        })
        for target, recall, precision, _ in pr_table:
            metrics['recall-at-precision>={:.2}'.format(target)] = recall
            metrics['precision-at-precision>={:.2}'.format(target)] = precision

        # Add prefix before returning.
        return dict([(prefix + k, v) for k, v in metrics.items()])

    metrics = _get_metric_dict(passage_answer_stats, 'passage-')
    metrics.update(_get_metric_dict(minimal_answer_stats, 'minimal-'))
    return metrics


def get_latex_str(f1, precision, recall):
    return '\\fpr' + '{' + ('%.1f' % (f1 * 100)) + '}{' + (
            '%.1f' % (precision * 100)) + '}{' + ('%.1f' % (recall * 100)) + '}'


# def load_gt_lookup_as_dict(gold_path, num_workers):
#     cache_path = os.path.join(os.path.dirname(gold_path), 'cache')
#
#     if os.path.exists(cache_path):
#         logging.info('Reading from cache: %s', format(cache_path))
#         tydi_gold_dict = pickle.load(open(cache_path, 'rb'))  # pytype: disable=wrong-arg-types
#     else:
#         tydi_gold_dict = eval_utils.read_annotation(gold_path, n_threads=num_workers)
#         logging.info('Caching gold data for future to: %s', format(cache_path))
#         cache_to_pickle_file(tydi_gold_dict, cache_path)  # pytype: disable=wrong-arg-types
#
#         logging.info('Read in gt lookup for %d examples into memory' % len(tydi_gold_dict))
#     return tydi_gold_dict


def pretty_print(tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold=2, minimal_non_null_threshold=2, verbose=False, skip_missing_example_ids=False):
    if any(map(not_, tydi_gold_dict.values())) or any(annotations[0].minimal_answer_span is None for annotations in tydi_gold_dict.values()):
        logging.warning("At least one example in gold dict has no annotations -- skipping pretty print evaluation")
        return

    gold_id_set = set(tydi_gold_dict.keys())
    pred_id_set = set(tydi_pred_dict.keys())
    sym_diff = gold_id_set.symmetric_difference(pred_id_set)

    if (not skip_missing_example_ids) and sym_diff:
        raise ValueError('ERROR: the example ids in gold annotations and example '
                         'ids in the prediction are not equal.')
    elif skip_missing_example_ids and sym_diff:
        logging.warning("Skipping {} example ids that are only in either gold or preds".format(len(sym_diff)))

    macro_avg_passage_scores = ([], [], [])
    macro_avg_minimal_scores = ([], [], [])

    language_list = [  # TODO: support arbitrary languages
        'english', 'arabic', 'bengali', 'finnish', 'indonesian', 'japanese',
        'swahili', 'korean', 'russian', 'telugu', 'thai'
    ]

    per_lang_gold = {}
    per_lang_pred = {}

    id_to_lang = {}

    for ex_id in gold_id_set:
        ex = tydi_gold_dict[ex_id]
        if ex[0].language in per_lang_gold:
            per_lang_gold[ex[0].language][ex_id] = ex
        else:
            per_lang_gold[ex[0].language] = {ex_id: ex}
        id_to_lang[ex_id] = ex[0].language
    for ex_id in gold_id_set:
        ex = tydi_pred_dict[ex_id]
        if id_to_lang[ex.example_id] in per_lang_pred:
            per_lang_pred[id_to_lang[ex.example_id]][ex.example_id] = ex
        else:
            per_lang_pred[id_to_lang[ex.example_id]] = {ex.example_id: ex}

    for lang in language_list:
        if lang in per_lang_pred:
infile
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
    macro_f1 = eval_utils.safe_divide(
        2 * macro_precision * macro_recall,
        macro_precision + macro_recall)

    return collections.OrderedDict({
        prefix + 'n': len(answer_stats),
        prefix + 'f1': macro_f1,
        prefix + 'precision': macro_precision,
        prefix + 'recall': macro_recall
    })


def compute_final_f1(passage_answer_stats, minimal_answer_stats):
    """Computes overall F1 given passage and minimal answers, ignoring scores.

    Note: this assumes that the answers have been thresholded based on scores.

    Arguments:
       passage_answer_stats: List of passage answer scores.
       minimal_answer_stats: List of minimal answer scores.


    Returns:
       Dictionary of name (string) -> score.
    """
    scores = compute_macro_f1(passage_answer_stats, prefix='passage-answer-')
    scores.update(compute_macro_f1(
        minimal_answer_stats, prefix='minimal-answer-'))
    return scores


def compute_pr_curves(answer_stats, targets=None):
    """Computes PR curve and returns R@P for specific targets.

    The values are computed as follows: find the (precision, recall) point
    with maximum recall and where precision > target.

    This is only relevant if you return the system scores in your predictions.
    You may find this useful when attempting to tune the threshold for your
    system on the dev set before requesting an evaluation on the test set
    via the leaderboard.

    Arguments:
      answer_stats: List of statistic tuples from the answer scores.
      targets (None): List of precision thresholds to target.

    Returns:
      List of table with rows: [target, r, p, score].
    """
    total_f1 = 0
    total_has_pred = 0
    total_has_gold = 0

    # Count the number of gold annotations.
    for has_gold, _, _, _ in answer_stats:
        total_has_gold += has_gold

    # Keep track of the point of maximum recall for each target.

    max_recall = [0 for _ in targets]
    max_precision = [0 for _ in targets]
    max_scores = [0.0 for _ in targets]

    # Only keep track of unique thresholds in this dictionary.
    scores_to_stats = collections.OrderedDict()

    # Loop through every possible threshold and compute precision + recall.
    for has_gold, has_pred, is_correct_or_f1, score in answer_stats:
        if isinstance(is_correct_or_f1, tuple):
            _, _, f1 = is_correct_or_f1
        else:
            f1 = is_correct_or_f1
        total_f1 += f1
        total_has_pred += has_pred

        precision = eval_utils.safe_divide(total_f1, total_has_pred)
        recall = eval_utils.safe_divide(total_f1, total_has_gold)

        # If there are any ties, this will be updated multiple times until the
        # ties are all counted.
        scores_to_stats[score] = [precision, recall]

    best_f1 = 0.0
    best_precision = 0.0
    best_recall = 0.0
    best_threshold = 0.0

    for threshold, (precision, recall) in scores_to_stats.items():
        # Match the thresholds to the find the closest precision above some target.
        for t, target in enumerate(targets):
            if precision >= target and recall > max_recall[t]:
                max_recall[t] = recall
                max_precision[t] = precision
                max_scores[t] = threshold

        # Compute optimal threshold.
        f1 = eval_utils.safe_divide(2 * precision * recall, precision + recall)
        if f1 > best_f1:
            best_f1 = f1
            best_precision = precision
            best_recall = recall
            best_threshold = threshold

    return ((best_f1, best_precision, best_recall, best_threshold),
            list(zip(targets, max_recall, max_precision, max_scores)))


def print_r_at_p_table(answer_stats):
    """Pretty prints the R@P table for default targets."""
    opt_result, pr_table = compute_pr_curves(
        answer_stats, targets=[0.5, 0.75, 0.9])
    f1, precision, recall, threshold = opt_result
    print('Optimal threshold: {:.5}'.format(threshold))
    print(' F1     /  P      /  R')
    print('{: >7.2%} / {: >7.2%} / {: >7.2%}'.format(f1, precision, recall))
    for target, recall, precision, row in pr_table:
        print('R@P={}: {:.2%} (actual p={:.2%}, score threshold={:.4})'.format(
            target, recall, precision, row))


def get_metrics_as_dict(gold_path, prediction_path, passage_non_null_threshold, minimal_non_null_threshold, verbose):
    """Library version of the end-to-end evaluation.

    Arguments:
      gold_path: Path to a single JSONL data. Could be gzipped or not.
      prediction_path: Path to the JSONL file of prediction data.

    Returns:
      metrics: A dictionary mapping string names to metric scores.
    """

    tydi_gold_dict = eval_utils.read_annotation(gold_path)
    tydi_pred_dict = eval_utils.read_prediction_jsonl(prediction_path)

    passage_answer_stats, minimal_answer_stats = score_answers(
        tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose)

    return get_metrics_with_answer_stats(
        passage_answer_stats, minimal_answer_stats)


def get_metrics_with_answer_stats(passage_answer_stats, minimal_answer_stats):
    """Generate metrics dict using passage and minimal answer stats."""

    def _get_metric_dict(answer_stats, prefix=''):
        """Compute all metrics for a set of answer statistics."""
        opt_result, pr_table = compute_pr_curves(
            answer_stats, targets=[0.5, 0.75, 0.9])
        f1, precision, recall, threshold = opt_result
        metrics = collections.OrderedDict({
            'best-threshold-f1': f1,
            'best-threshold-precision': precision,
            'best-threshold-recall': recall,
            'best-threshold': threshold,
        })
        for target, recall, precision, _ in pr_table:
            metrics['recall-at-precision>={:.2}'.format(target)] = recall
            metrics['precision-at-precision>={:.2}'.format(target)] = precision

        # Add prefix before returning.
        return dict([(prefix + k, v) for k, v in metrics.items()])

    metrics = _get_metric_dict(passage_answer_stats, 'passage-')
    metrics.update(_get_metric_dict(minimal_answer_stats, 'minimal-'))
    return metrics


def get_latex_str(f1, precision, recall):
    return '\\fpr' + '{' + ('%.1f' % (f1 * 100)) + '}{' + (
            '%.1f' % (precision * 100)) + '}{' + ('%.1f' % (recall * 100)) + '}'


# def load_gt_lookup_as_dict(gold_path, num_workers):
#     cache_path = os.path.join(os.path.dirname(gold_path), 'cache')
#
#     if os.path.exists(cache_path):
#         logging.info('Reading from cache: %s', format(cache_path))
#         tydi_gold_dict = pickle.load(open(cache_path, 'rb'))  # pytype: disable=wrong-arg-types
#     else:
#         tydi_gold_dict = eval_utils.read_annotation(gold_path, n_threads=num_workers)
#         logging.info('Caching gold data for future to: %s', format(cache_path))
#         cache_to_pickle_file(tydi_gold_dict, cache_path)  # pytype: disable=wrong-arg-types
#
#         logging.info('Read in gt lookup for %d examples into memory' % len(tydi_gold_dict))
#     return tydi_gold_dict


def pretty_print(tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold=2, minimal_non_null_threshold=2, verbose=False, skip_missing_example_ids=False):
    if any(map(not_, tydi_gold_dict.values())) or any(annotations[0].minimal_answer_span is None for annotations in tydi_gold_dict.values()):
        logging.warning("At least one example in gold dict has no annotations -- skipping pretty print evaluation")
        return

    gold_id_set = set(tydi_gold_dict.keys())
    pred_id_set = set(tydi_pred_dict.keys())
    sym_diff = gold_id_set.symmetric_difference(pred_id_set)

    if (not skip_missing_example_ids) and sym_diff:
        raise ValueError('ERROR: the example ids in gold annotations and example '
                         'ids in the prediction are not equal.')
    elif skip_missing_example_ids and sym_diff:
        logging.warning("Skipping {} example ids that are only in either gold or preds".format(len(sym_diff)))

    macro_avg_passage_scores = ([], [], [])
    macro_avg_minimal_scores = ([], [], [])

    language_list = [  # TODO: support arbitrary languages
        'english', 'arabic', 'bengali', 'finnish', 'indonesian', 'japanese',
        'swahili', 'korean', 'russian', 'telugu', 'thai'
    ]

    per_lang_gold = {}
    per_lang_pred = {}

    id_to_lang = {}

    for ex_id in gold_id_set:
        ex = tydi_gold_dict[ex_id]
        if ex[0].language in per_lang_gold:
            per_lang_gold[ex[0].language][ex_id] = ex
        else:
            per_lang_gold[ex[0].language] = {ex_id: ex}
        id_to_lang[ex_id] = ex[0].language
    for ex_id in gold_id_set:
        ex = tydi_pred_dict[ex_id]
        if id_to_lang[ex.example_id] in per_lang_pred:
            per_lang_pred[id_to_lang[ex.example_id]][ex.example_id] = ex
        else:
            per_lang_pred[id_to_lang[ex.example_id]] = {ex.example_id: ex}

    for lang in language_list:
        if lang in per_lang_pred:
            passage_answer_stats, minimal_answer_stats = score_answers(
                per_lang_gold.get(lang, {}), per_lang_pred[lang],
                passage_non_null_threshold, minimal_non_null_threshold, verbose)

            # Passage selection task
infile
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
    macro_f1 = eval_utils.safe_divide(
        2 * macro_precision * macro_recall,
        macro_precision + macro_recall)

    return collections.OrderedDict({
        prefix + 'n': len(answer_stats),
        prefix + 'f1': macro_f1,
        prefix + 'precision': macro_precision,
        prefix + 'recall': macro_recall
    })


def compute_final_f1(passage_answer_stats, minimal_answer_stats):
    """Computes overall F1 given passage and minimal answers, ignoring scores.

    Note: this assumes that the answers have been thresholded based on scores.

    Arguments:
       passage_answer_stats: List of passage answer scores.
       minimal_answer_stats: List of minimal answer scores.


    Returns:
       Dictionary of name (string) -> score.
    """
    scores = compute_macro_f1(passage_answer_stats, prefix='passage-answer-')
    scores.update(compute_macro_f1(
        minimal_answer_stats, prefix='minimal-answer-'))
    return scores


def compute_pr_curves(answer_stats, targets=None):
    """Computes PR curve and returns R@P for specific targets.

    The values are computed as follows: find the (precision, recall) point
    with maximum recall and where precision > target.

    This is only relevant if you return the system scores in your predictions.
    You may find this useful when attempting to tune the threshold for your
    system on the dev set before requesting an evaluation on the test set
    via the leaderboard.

    Arguments:
      answer_stats: List of statistic tuples from the answer scores.
      targets (None): List of precision thresholds to target.

    Returns:
      List of table with rows: [target, r, p, score].
    """
    total_f1 = 0
    total_has_pred = 0
    total_has_gold = 0

    # Count the number of gold annotations.
    for has_gold, _, _, _ in answer_stats:
        total_has_gold += has_gold

    # Keep track of the point of maximum recall for each target.

    max_recall = [0 for _ in targets]
    max_precision = [0 for _ in targets]
    max_scores = [0.0 for _ in targets]

    # Only keep track of unique thresholds in this dictionary.
    scores_to_stats = collections.OrderedDict()

    # Loop through every possible threshold and compute precision + recall.
    for has_gold, has_pred, is_correct_or_f1, score in answer_stats:
        if isinstance(is_correct_or_f1, tuple):
            _, _, f1 = is_correct_or_f1
        else:
            f1 = is_correct_or_f1
        total_f1 += f1
        total_has_pred += has_pred

        precision = eval_utils.safe_divide(total_f1, total_has_pred)
        recall = eval_utils.safe_divide(total_f1, total_has_gold)

        # If there are any ties, this will be updated multiple times until the
        # ties are all counted.
        scores_to_stats[score] = [precision, recall]

    best_f1 = 0.0
    best_precision = 0.0
    best_recall = 0.0
    best_threshold = 0.0

    for threshold, (precision, recall) in scores_to_stats.items():
        # Match the thresholds to the find the closest precision above some target.
        for t, target in enumerate(targets):
            if precision >= target and recall > max_recall[t]:
                max_recall[t] = recall
                max_precision[t] = precision
                max_scores[t] = threshold

        # Compute optimal threshold.
        f1 = eval_utils.safe_divide(2 * precision * recall, precision + recall)
        if f1 > best_f1:
            best_f1 = f1
            best_precision = precision
            best_recall = recall
            best_threshold = threshold

    return ((best_f1, best_precision, best_recall, best_threshold),
            list(zip(targets, max_recall, max_precision, max_scores)))


def print_r_at_p_table(answer_stats):
    """Pretty prints the R@P table for default targets."""
infile
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
    macro_f1 = eval_utils.safe_divide(
        2 * macro_precision * macro_recall,
        macro_precision + macro_recall)

    return collections.OrderedDict({
        prefix + 'n': len(answer_stats),
        prefix + 'f1': macro_f1,
        prefix + 'precision': macro_precision,
        prefix + 'recall': macro_recall
    })


def compute_final_f1(passage_answer_stats, minimal_answer_stats):
    """Computes overall F1 given passage and minimal answers, ignoring scores.

    Note: this assumes that the answers have been thresholded based on scores.

    Arguments:
       passage_answer_stats: List of passage answer scores.
       minimal_answer_stats: List of minimal answer scores.


    Returns:
       Dictionary of name (string) -> score.
    """
    scores = compute_macro_f1(passage_answer_stats, prefix='passage-answer-')
    scores.update(compute_macro_f1(
        minimal_answer_stats, prefix='minimal-answer-'))
    return scores


def compute_pr_curves(answer_stats, targets=None):
    """Computes PR curve and returns R@P for specific targets.

    The values are computed as follows: find the (precision, recall) point
    with maximum recall and where precision > target.

    This is only relevant if you return the system scores in your predictions.
    You may find this useful when attempting to tune the threshold for your
    system on the dev set before requesting an evaluation on the test set
    via the leaderboard.

    Arguments:
      answer_stats: List of statistic tuples from the answer scores.
      targets (None): List of precision thresholds to target.

    Returns:
      List of table with rows: [target, r, p, score].
    """
    total_f1 = 0
    total_has_pred = 0
    total_has_gold = 0

    # Count the number of gold annotations.
    for has_gold, _, _, _ in answer_stats:
        total_has_gold += has_gold

    # Keep track of the point of maximum recall for each target.

    max_recall = [0 for _ in targets]
    max_precision = [0 for _ in targets]
    max_scores = [0.0 for _ in targets]

    # Only keep track of unique thresholds in this dictionary.
    scores_to_stats = collections.OrderedDict()

    # Loop through every possible threshold and compute precision + recall.
    for has_gold, has_pred, is_correct_or_f1, score in answer_stats:
        if isinstance(is_correct_or_f1, tuple):
            _, _, f1 = is_correct_or_f1
        else:
            f1 = is_correct_or_f1
        total_f1 += f1
        total_has_pred += has_pred

        precision = eval_utils.safe_divide(total_f1, total_has_pred)
        recall = eval_utils.safe_divide(total_f1, total_has_gold)

        # If there are any ties, this will be updated multiple times until the
        # ties are all counted.
        scores_to_stats[score] = [precision, recall]

    best_f1 = 0.0
    best_precision = 0.0
    best_recall = 0.0
    best_threshold = 0.0

    for threshold, (precision, recall) in scores_to_stats.items():
        # Match the thresholds to the find the closest precision above some target.
        for t, target in enumerate(targets):
            if precision >= target and recall > max_recall[t]:
                max_recall[t] = recall
                max_precision[t] = precision
                max_scores[t] = threshold

        # Compute optimal threshold.
        f1 = eval_utils.safe_divide(2 * precision * recall, precision + recall)
        if f1 > best_f1:
            best_f1 = f1
            best_precision = precision
            best_recall = recall
            best_threshold = threshold

    return ((best_f1, best_precision, best_recall, best_threshold),
            list(zip(targets, max_recall, max_precision, max_scores)))


def print_r_at_p_table(answer_stats):
    """Pretty prints the R@P table for default targets."""
    opt_result, pr_table = compute_pr_curves(
        answer_stats, targets=[0.5, 0.75, 0.9])
    f1, precision, recall, threshold = opt_result
    print('Optimal threshold: {:.5}'.format(threshold))
    print(' F1     /  P      /  R')
    print('{: >7.2%} / {: >7.2%} / {: >7.2%}'.format(f1, precision, recall))
    for target, recall, precision, row in pr_table:
        print('R@P={}: {:.2%} (actual p={:.2%}, score threshold={:.4})'.format(
            target, recall, precision, row))


def get_metrics_as_dict(gold_path, prediction_path, passage_non_null_threshold, minimal_non_null_threshold, verbose):
    """Library version of the end-to-end evaluation.

    Arguments:
      gold_path: Path to a single JSONL data. Could be gzipped or not.
      prediction_path: Path to the JSONL file of prediction data.

    Returns:
      metrics: A dictionary mapping string names to metric scores.
    """

    tydi_gold_dict = eval_utils.read_annotation(gold_path)
    tydi_pred_dict = eval_utils.read_prediction_jsonl(prediction_path)

    passage_answer_stats, minimal_answer_stats = score_answers(
        tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose)

    return get_metrics_with_answer_stats(
        passage_answer_stats, minimal_answer_stats)


def get_metrics_with_answer_stats(passage_answer_stats, minimal_answer_stats):
    """Generate metrics dict using passage and minimal answer stats."""

    def _get_metric_dict(answer_stats, prefix=''):
        """Compute all metrics for a set of answer statistics."""
        opt_result, pr_table = compute_pr_curves(
            answer_stats, targets=[0.5, 0.75, 0.9])
        f1, precision, recall, threshold = opt_result
        metrics = collections.OrderedDict({
            'best-threshold-f1': f1,
            'best-threshold-precision': precision,
            'best-threshold-recall': recall,
            'best-threshold': threshold,
        })
        for target, recall, precision, _ in pr_table:
            metrics['recall-at-precision>={:.2}'.format(target)] = recall
            metrics['precision-at-precision>={:.2}'.format(target)] = precision

        # Add prefix before returning.
        return dict([(prefix + k, v) for k, v in metrics.items()])

    metrics = _get_metric_dict(passage_answer_stats, 'passage-')
    metrics.update(_get_metric_dict(minimal_answer_stats, 'minimal-'))
    return metrics


def get_latex_str(f1, precision, recall):
    return '\\fpr' + '{' + ('%.1f' % (f1 * 100)) + '}{' + (
            '%.1f' % (precision * 100)) + '}{' + ('%.1f' % (recall * 100)) + '}'


# def load_gt_lookup_as_dict(gold_path, num_workers):
#     cache_path = os.path.join(os.path.dirname(gold_path), 'cache')
#
#     if os.path.exists(cache_path):
#         logging.info('Reading from cache: %s', format(cache_path))
#         tydi_gold_dict = pickle.load(open(cache_path, 'rb'))  # pytype: disable=wrong-arg-types
#     else:
#         tydi_gold_dict = eval_utils.read_annotation(gold_path, n_threads=num_workers)
#         logging.info('Caching gold data for future to: %s', format(cache_path))
#         cache_to_pickle_file(tydi_gold_dict, cache_path)  # pytype: disable=wrong-arg-types
#
#         logging.info('Read in gt lookup for %d examples into memory' % len(tydi_gold_dict))
#     return tydi_gold_dict


def pretty_print(tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold=2, minimal_non_null_threshold=2, verbose=False, skip_missing_example_ids=False):
    if any(map(not_, tydi_gold_dict.values())) or any(annotations[0].minimal_answer_span is None for annotations in tydi_gold_dict.values()):
        logging.warning("At least one example in gold dict has no annotations -- skipping pretty print evaluation")
        return

    gold_id_set = set(tydi_gold_dict.keys())
    pred_id_set = set(tydi_pred_dict.keys())
    sym_diff = gold_id_set.symmetric_difference(pred_id_set)

    if (not skip_missing_example_ids) and sym_diff:
        raise ValueError('ERROR: the example ids in gold annotations and example '
                         'ids in the prediction are not equal.')
    elif skip_missing_example_ids and sym_diff:
        logging.warning("Skipping {} example ids that are only in either gold or preds".format(len(sym_diff)))

    macro_avg_passage_scores = ([], [], [])
    macro_avg_minimal_scores = ([], [], [])

    language_list = [  # TODO: support arbitrary languages
        'english', 'arabic', 'bengali', 'finnish', 'indonesian', 'japanese',
        'swahili', 'korean', 'russian', 'telugu', 'thai'
    ]

    per_lang_gold = {}
    per_lang_pred = {}

    id_to_lang = {}

    for ex_id in gold_id_set:
        ex = tydi_gold_dict[ex_id]
        if ex[0].language in per_lang_gold:
            per_lang_gold[ex[0].language][ex_id] = ex
        else:
            per_lang_gold[ex[0].language] = {ex_id: ex}
        id_to_lang[ex_id] = ex[0].language
    for ex_id in gold_id_set:
        ex = tydi_pred_dict[ex_id]
        if id_to_lang[ex.example_id] in per_lang_pred:
            per_lang_pred[id_to_lang[ex.example_id]][ex.example_id] = ex
        else:
            per_lang_pred[id_to_lang[ex.example_id]] = {ex.example_id: ex}

    for lang in language_list:
        if lang in per_lang_pred:
            passage_answer_stats, minimal_answer_stats = score_answers(
                per_lang_gold.get(lang, {}), per_lang_pred[lang],
                passage_non_null_threshold, minimal_non_null_threshold, verbose)

            # Passage selection task
            opt_result, _ = compute_pr_curves(passage_answer_stats, targets=[0.5])
            f1, precision, recall, _ = opt_result
            if lang != 'english':
                macro_avg_passage_scores[0].append(f1)
                macro_avg_passage_scores[1].append(precision)
                macro_avg_passage_scores[2].append(recall)
            print('Passage & ' + lang + ' & ' + get_latex_str(f1, precision, recall))

            # Minimal answer span task
infile
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
    macro_f1 = eval_utils.safe_divide(
        2 * macro_precision * macro_recall,
        macro_precision + macro_recall)

    return collections.OrderedDict({
        prefix + 'n': len(answer_stats),
        prefix + 'f1': macro_f1,
        prefix + 'precision': macro_precision,
        prefix + 'recall': macro_recall
    })


def compute_final_f1(passage_answer_stats, minimal_answer_stats):
    """Computes overall F1 given passage and minimal answers, ignoring scores.

    Note: this assumes that the answers have been thresholded based on scores.

    Arguments:
       passage_answer_stats: List of passage answer scores.
       minimal_answer_stats: List of minimal answer scores.


    Returns:
       Dictionary of name (string) -> score.
    """
    scores = compute_macro_f1(passage_answer_stats, prefix='passage-answer-')
    scores.update(compute_macro_f1(
        minimal_answer_stats, prefix='minimal-answer-'))
    return scores


def compute_pr_curves(answer_stats, targets=None):
    """Computes PR curve and returns R@P for specific targets.

    The values are computed as follows: find the (precision, recall) point
    with maximum recall and where precision > target.

    This is only relevant if you return the system scores in your predictions.
    You may find this useful when attempting to tune the threshold for your
    system on the dev set before requesting an evaluation on the test set
    via the leaderboard.

    Arguments:
      answer_stats: List of statistic tuples from the answer scores.
      targets (None): List of precision thresholds to target.

    Returns:
      List of table with rows: [target, r, p, score].
    """
    total_f1 = 0
    total_has_pred = 0
    total_has_gold = 0

    # Count the number of gold annotations.
    for has_gold, _, _, _ in answer_stats:
        total_has_gold += has_gold

    # Keep track of the point of maximum recall for each target.

    max_recall = [0 for _ in targets]
    max_precision = [0 for _ in targets]
    max_scores = [0.0 for _ in targets]

    # Only keep track of unique thresholds in this dictionary.
    scores_to_stats = collections.OrderedDict()

    # Loop through every possible threshold and compute precision + recall.
    for has_gold, has_pred, is_correct_or_f1, score in answer_stats:
        if isinstance(is_correct_or_f1, tuple):
            _, _, f1 = is_correct_or_f1
        else:
            f1 = is_correct_or_f1
        total_f1 += f1
        total_has_pred += has_pred

        precision = eval_utils.safe_divide(total_f1, total_has_pred)
        recall = eval_utils.safe_divide(total_f1, total_has_gold)

        # If there are any ties, this will be updated multiple times until the
        # ties are all counted.
        scores_to_stats[score] = [precision, recall]

    best_f1 = 0.0
    best_precision = 0.0
    best_recall = 0.0
    best_threshold = 0.0

    for threshold, (precision, recall) in scores_to_stats.items():
        # Match the thresholds to the find the closest precision above some target.
        for t, target in enumerate(targets):
            if precision >= target and recall > max_recall[t]:
                max_recall[t] = recall
                max_precision[t] = precision
                max_scores[t] = threshold

        # Compute optimal threshold.
        f1 = eval_utils.safe_divide(2 * precision * recall, precision + recall)
        if f1 > best_f1:
            best_f1 = f1
            best_precision = precision
            best_recall = recall
            best_threshold = threshold

    return ((best_f1, best_precision, best_recall, best_threshold),
            list(zip(targets, max_recall, max_precision, max_scores)))


def print_r_at_p_table(answer_stats):
    """Pretty prints the R@P table for default targets."""
    opt_result, pr_table = compute_pr_curves(
        answer_stats, targets=[0.5, 0.75, 0.9])
    f1, precision, recall, threshold = opt_result
    print('Optimal threshold: {:.5}'.format(threshold))
    print(' F1     /  P      /  R')
    print('{: >7.2%} / {: >7.2%} / {: >7.2%}'.format(f1, precision, recall))
    for target, recall, precision, row in pr_table:
        print('R@P={}: {:.2%} (actual p={:.2%}, score threshold={:.4})'.format(
            target, recall, precision, row))


def get_metrics_as_dict(gold_path, prediction_path, passage_non_null_threshold, minimal_non_null_threshold, verbose):
    """Library version of the end-to-end evaluation.

    Arguments:
      gold_path: Path to a single JSONL data. Could be gzipped or not.
      prediction_path: Path to the JSONL file of prediction data.

    Returns:
      metrics: A dictionary mapping string names to metric scores.
    """

    tydi_gold_dict = eval_utils.read_annotation(gold_path)
    tydi_pred_dict = eval_utils.read_prediction_jsonl(prediction_path)

    passage_answer_stats, minimal_answer_stats = score_answers(
        tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose)

    return get_metrics_with_answer_stats(
        passage_answer_stats, minimal_answer_stats)


def get_metrics_with_answer_stats(passage_answer_stats, minimal_answer_stats):
    """Generate metrics dict using passage and minimal answer stats."""

    def _get_metric_dict(answer_stats, prefix=''):
        """Compute all metrics for a set of answer statistics."""
        opt_result, pr_table = compute_pr_curves(
            answer_stats, targets=[0.5, 0.75, 0.9])
        f1, precision, recall, threshold = opt_result
        metrics = collections.OrderedDict({
            'best-threshold-f1': f1,
            'best-threshold-precision': precision,
            'best-threshold-recall': recall,
            'best-threshold': threshold,
        })
        for target, recall, precision, _ in pr_table:
            metrics['recall-at-precision>={:.2}'.format(target)] = recall
            metrics['precision-at-precision>={:.2}'.format(target)] = precision

        # Add prefix before returning.
        return dict([(prefix + k, v) for k, v in metrics.items()])

    metrics = _get_metric_dict(passage_answer_stats, 'passage-')
    metrics.update(_get_metric_dict(minimal_answer_stats, 'minimal-'))
    return metrics


def get_latex_str(f1, precision, recall):
    return '\\fpr' + '{' + ('%.1f' % (f1 * 100)) + '}{' + (
            '%.1f' % (precision * 100)) + '}{' + ('%.1f' % (recall * 100)) + '}'


# def load_gt_lookup_as_dict(gold_path, num_workers):
#     cache_path = os.path.join(os.path.dirname(gold_path), 'cache')
#
#     if os.path.exists(cache_path):
#         logging.info('Reading from cache: %s', format(cache_path))
#         tydi_gold_dict = pickle.load(open(cache_path, 'rb'))  # pytype: disable=wrong-arg-types
#     else:
#         tydi_gold_dict = eval_utils.read_annotation(gold_path, n_threads=num_workers)
#         logging.info('Caching gold data for future to: %s', format(cache_path))
#         cache_to_pickle_file(tydi_gold_dict, cache_path)  # pytype: disable=wrong-arg-types
#
#         logging.info('Read in gt lookup for %d examples into memory' % len(tydi_gold_dict))
#     return tydi_gold_dict


def pretty_print(tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold=2, minimal_non_null_threshold=2, verbose=False, skip_missing_example_ids=False):
    if any(map(not_, tydi_gold_dict.values())) or any(annotations[0].minimal_answer_span is None for annotations in tydi_gold_dict.values()):
        logging.warning("At least one example in gold dict has no annotations -- skipping pretty print evaluation")
        return

    gold_id_set = set(tydi_gold_dict.keys())
    pred_id_set = set(tydi_pred_dict.keys())
    sym_diff = gold_id_set.symmetric_difference(pred_id_set)

    if (not skip_missing_example_ids) and sym_diff:
        raise ValueError('ERROR: the example ids in gold annotations and example '
                         'ids in the prediction are not equal.')
    elif skip_missing_example_ids and sym_diff:
        logging.warning("Skipping {} example ids that are only in either gold or preds".format(len(sym_diff)))

    macro_avg_passage_scores = ([], [], [])
    macro_avg_minimal_scores = ([], [], [])

    language_list = [  # TODO: support arbitrary languages
        'english', 'arabic', 'bengali', 'finnish', 'indonesian', 'japanese',
        'swahili', 'korean', 'russian', 'telugu', 'thai'
    ]

    per_lang_gold = {}
    per_lang_pred = {}

    id_to_lang = {}

    for ex_id in gold_id_set:
        ex = tydi_gold_dict[ex_id]
        if ex[0].language in per_lang_gold:
            per_lang_gold[ex[0].language][ex_id] = ex
        else:
            per_lang_gold[ex[0].language] = {ex_id: ex}
        id_to_lang[ex_id] = ex[0].language
    for ex_id in gold_id_set:
        ex = tydi_pred_dict[ex_id]
        if id_to_lang[ex.example_id] in per_lang_pred:
            per_lang_pred[id_to_lang[ex.example_id]][ex.example_id] = ex
        else:
            per_lang_pred[id_to_lang[ex.example_id]] = {ex.example_id: ex}

    for lang in language_list:
        if lang in per_lang_pred:
            passage_answer_stats, minimal_answer_stats = score_answers(
                per_lang_gold.get(lang, {}), per_lang_pred[lang],
                passage_non_null_threshold, minimal_non_null_threshold, verbose)

            # Passage selection task
            opt_result, _ = compute_pr_curves(passage_answer_stats, targets=[0.5])
            f1, precision, recall, _ = opt_result
            if lang != 'english':
                macro_avg_passage_scores[0].append(f1)
                macro_avg_passage_scores[1].append(precision)
                macro_avg_passage_scores[2].append(recall)
            print('Passage & ' + lang + ' & ' + get_latex_str(f1, precision, recall))

            # Minimal answer span task
            opt_result, _ = compute_pr_curves(minimal_answer_stats, targets=[0.5])
            f1, precision, recall, _ = opt_result
            if lang != 'english':
                macro_avg_minimal_scores[0].append(f1)
                macro_avg_minimal_scores[1].append(precision)
                macro_avg_minimal_scores[2].append(recall)
            print('Minimal Answer & ' + lang + ' & ' +
infile
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
    macro_f1 = eval_utils.safe_divide(
        2 * macro_precision * macro_recall,
        macro_precision + macro_recall)

    return collections.OrderedDict({
        prefix + 'n': len(answer_stats),
        prefix + 'f1': macro_f1,
        prefix + 'precision': macro_precision,
        prefix + 'recall': macro_recall
    })


def compute_final_f1(passage_answer_stats, minimal_answer_stats):
    """Computes overall F1 given passage and minimal answers, ignoring scores.

    Note: this assumes that the answers have been thresholded based on scores.

    Arguments:
       passage_answer_stats: List of passage answer scores.
       minimal_answer_stats: List of minimal answer scores.


    Returns:
       Dictionary of name (string) -> score.
    """
    scores = compute_macro_f1(passage_answer_stats, prefix='passage-answer-')
    scores.update(compute_macro_f1(
        minimal_answer_stats, prefix='minimal-answer-'))
    return scores


def compute_pr_curves(answer_stats, targets=None):
    """Computes PR curve and returns R@P for specific targets.

    The values are computed as follows: find the (precision, recall) point
    with maximum recall and where precision > target.

    This is only relevant if you return the system scores in your predictions.
    You may find this useful when attempting to tune the threshold for your
    system on the dev set before requesting an evaluation on the test set
    via the leaderboard.

    Arguments:
      answer_stats: List of statistic tuples from the answer scores.
      targets (None): List of precision thresholds to target.

    Returns:
      List of table with rows: [target, r, p, score].
    """
    total_f1 = 0
    total_has_pred = 0
    total_has_gold = 0

    # Count the number of gold annotations.
    for has_gold, _, _, _ in answer_stats:
        total_has_gold += has_gold

    # Keep track of the point of maximum recall for each target.

    max_recall = [0 for _ in targets]
    max_precision = [0 for _ in targets]
    max_scores = [0.0 for _ in targets]

    # Only keep track of unique thresholds in this dictionary.
    scores_to_stats = collections.OrderedDict()

    # Loop through every possible threshold and compute precision + recall.
    for has_gold, has_pred, is_correct_or_f1, score in answer_stats:
        if isinstance(is_correct_or_f1, tuple):
            _, _, f1 = is_correct_or_f1
        else:
            f1 = is_correct_or_f1
        total_f1 += f1
        total_has_pred += has_pred

        precision = eval_utils.safe_divide(total_f1, total_has_pred)
        recall = eval_utils.safe_divide(total_f1, total_has_gold)

        # If there are any ties, this will be updated multiple times until the
        # ties are all counted.
        scores_to_stats[score] = [precision, recall]

    best_f1 = 0.0
    best_precision = 0.0
    best_recall = 0.0
    best_threshold = 0.0

    for threshold, (precision, recall) in scores_to_stats.items():
        # Match the thresholds to the find the closest precision above some target.
        for t, target in enumerate(targets):
            if precision >= target and recall > max_recall[t]:
                max_recall[t] = recall
                max_precision[t] = precision
                max_scores[t] = threshold

        # Compute optimal threshold.
        f1 = eval_utils.safe_divide(2 * precision * recall, precision + recall)
        if f1 > best_f1:
            best_f1 = f1
            best_precision = precision
            best_recall = recall
            best_threshold = threshold

    return ((best_f1, best_precision, best_recall, best_threshold),
            list(zip(targets, max_recall, max_precision, max_scores)))


def print_r_at_p_table(answer_stats):
    """Pretty prints the R@P table for default targets."""
    opt_result, pr_table = compute_pr_curves(
        answer_stats, targets=[0.5, 0.75, 0.9])
    f1, precision, recall, threshold = opt_result
    print('Optimal threshold: {:.5}'.format(threshold))
    print(' F1     /  P      /  R')
    print('{: >7.2%} / {: >7.2%} / {: >7.2%}'.format(f1, precision, recall))
    for target, recall, precision, row in pr_table:
        print('R@P={}: {:.2%} (actual p={:.2%}, score threshold={:.4})'.format(
            target, recall, precision, row))


def get_metrics_as_dict(gold_path, prediction_path, passage_non_null_threshold, minimal_non_null_threshold, verbose):
    """Library version of the end-to-end evaluation.

    Arguments:
      gold_path: Path to a single JSONL data. Could be gzipped or not.
      prediction_path: Path to the JSONL file of prediction data.

    Returns:
      metrics: A dictionary mapping string names to metric scores.
    """

    tydi_gold_dict = eval_utils.read_annotation(gold_path)
    tydi_pred_dict = eval_utils.read_prediction_jsonl(prediction_path)

    passage_answer_stats, minimal_answer_stats = score_answers(
        tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose)

infile
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
infile
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
infile
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
    macro_f1 = eval_utils.safe_divide(
        2 * macro_precision * macro_recall,
        macro_precision + macro_recall)

    return collections.OrderedDict({
        prefix + 'n': len(answer_stats),
        prefix + 'f1': macro_f1,
        prefix + 'precision': macro_precision,
        prefix + 'recall': macro_recall
    })


def compute_final_f1(passage_answer_stats, minimal_answer_stats):
    """Computes overall F1 given passage and minimal answers, ignoring scores.

    Note: this assumes that the answers have been thresholded based on scores.

    Arguments:
       passage_answer_stats: List of passage answer scores.
       minimal_answer_stats: List of minimal answer scores.


    Returns:
       Dictionary of name (string) -> score.
    """
infile
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
infile
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
    macro_f1 = eval_utils.safe_divide(
        2 * macro_precision * macro_recall,
        macro_precision + macro_recall)

    return collections.OrderedDict({
        prefix + 'n': len(answer_stats),
        prefix + 'f1': macro_f1,
        prefix + 'precision': macro_precision,
        prefix + 'recall': macro_recall
    })


def compute_final_f1(passage_answer_stats, minimal_answer_stats):
    """Computes overall F1 given passage and minimal answers, ignoring scores.

    Note: this assumes that the answers have been thresholded based on scores.

    Arguments:
       passage_answer_stats: List of passage answer scores.
       minimal_answer_stats: List of minimal answer scores.


    Returns:
       Dictionary of name (string) -> score.
    """
    scores = compute_macro_f1(passage_answer_stats, prefix='passage-answer-')
infile
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
    macro_f1 = eval_utils.safe_divide(
        2 * macro_precision * macro_recall,
        macro_precision + macro_recall)

    return collections.OrderedDict({
        prefix + 'n': len(answer_stats),
        prefix + 'f1': macro_f1,
        prefix + 'precision': macro_precision,
        prefix + 'recall': macro_recall
    })


def compute_final_f1(passage_answer_stats, minimal_answer_stats):
    """Computes overall F1 given passage and minimal answers, ignoring scores.

    Note: this assumes that the answers have been thresholded based on scores.

    Arguments:
       passage_answer_stats: List of passage answer scores.
       minimal_answer_stats: List of minimal answer scores.


    Returns:
       Dictionary of name (string) -> score.
    """
    scores = compute_macro_f1(passage_answer_stats, prefix='passage-answer-')
    scores.update(compute_macro_f1(
        minimal_answer_stats, prefix='minimal-answer-'))
    return scores


def compute_pr_curves(answer_stats, targets=None):
    """Computes PR curve and returns R@P for specific targets.

    The values are computed as follows: find the (precision, recall) point
    with maximum recall and where precision > target.

    This is only relevant if you return the system scores in your predictions.
    You may find this useful when attempting to tune the threshold for your
    system on the dev set before requesting an evaluation on the test set
    via the leaderboard.

    Arguments:
      answer_stats: List of statistic tuples from the answer scores.
      targets (None): List of precision thresholds to target.

    Returns:
      List of table with rows: [target, r, p, score].
    """
    total_f1 = 0
    total_has_pred = 0
    total_has_gold = 0

    # Count the number of gold annotations.
    for has_gold, _, _, _ in answer_stats:
        total_has_gold += has_gold

    # Keep track of the point of maximum recall for each target.

    max_recall = [0 for _ in targets]
    max_precision = [0 for _ in targets]
    max_scores = [0.0 for _ in targets]

    # Only keep track of unique thresholds in this dictionary.
    scores_to_stats = collections.OrderedDict()

    # Loop through every possible threshold and compute precision + recall.
    for has_gold, has_pred, is_correct_or_f1, score in answer_stats:
        if isinstance(is_correct_or_f1, tuple):
            _, _, f1 = is_correct_or_f1
        else:
            f1 = is_correct_or_f1
        total_f1 += f1
        total_has_pred += has_pred

        precision = eval_utils.safe_divide(total_f1, total_has_pred)
        recall = eval_utils.safe_divide(total_f1, total_has_gold)

        # If there are any ties, this will be updated multiple times until the
        # ties are all counted.
        scores_to_stats[score] = [precision, recall]

    best_f1 = 0.0
    best_precision = 0.0
    best_recall = 0.0
    best_threshold = 0.0

    for threshold, (precision, recall) in scores_to_stats.items():
        # Match the thresholds to the find the closest precision above some target.
        for t, target in enumerate(targets):
            if precision >= target and recall > max_recall[t]:
                max_recall[t] = recall
                max_precision[t] = precision
                max_scores[t] = threshold

        # Compute optimal threshold.
        f1 = eval_utils.safe_divide(2 * precision * recall, precision + recall)
        if f1 > best_f1:
            best_f1 = f1
            best_precision = precision
            best_recall = recall
            best_threshold = threshold

    return ((best_f1, best_precision, best_recall, best_threshold),
            list(zip(targets, max_recall, max_precision, max_scores)))


def print_r_at_p_table(answer_stats):
    """Pretty prints the R@P table for default targets."""
    opt_result, pr_table = compute_pr_curves(
        answer_stats, targets=[0.5, 0.75, 0.9])
    f1, precision, recall, threshold = opt_result
    print('Optimal threshold: {:.5}'.format(threshold))
    print(' F1     /  P      /  R')
    print('{: >7.2%} / {: >7.2%} / {: >7.2%}'.format(f1, precision, recall))
    for target, recall, precision, row in pr_table:
        print('R@P={}: {:.2%} (actual p={:.2%}, score threshold={:.4})'.format(
            target, recall, precision, row))


def get_metrics_as_dict(gold_path, prediction_path, passage_non_null_threshold, minimal_non_null_threshold, verbose):
    """Library version of the end-to-end evaluation.

    Arguments:
      gold_path: Path to a single JSONL data. Could be gzipped or not.
      prediction_path: Path to the JSONL file of prediction data.

    Returns:
      metrics: A dictionary mapping string names to metric scores.
    """

    tydi_gold_dict = eval_utils.read_annotation(gold_path)
    tydi_pred_dict = eval_utils.read_prediction_jsonl(prediction_path)

    passage_answer_stats, minimal_answer_stats = score_answers(
        tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose)

    return get_metrics_with_answer_stats(
        passage_answer_stats, minimal_answer_stats)


def get_metrics_with_answer_stats(passage_answer_stats, minimal_answer_stats):
    """Generate metrics dict using passage and minimal answer stats."""

    def _get_metric_dict(answer_stats, prefix=''):
        """Compute all metrics for a set of answer statistics."""
        opt_result, pr_table = compute_pr_curves(
            answer_stats, targets=[0.5, 0.75, 0.9])
        f1, precision, recall, threshold = opt_result
        metrics = collections.OrderedDict({
            'best-threshold-f1': f1,
            'best-threshold-precision': precision,
            'best-threshold-recall': recall,
            'best-threshold': threshold,
        })
        for target, recall, precision, _ in pr_table:
            metrics['recall-at-precision>={:.2}'.format(target)] = recall
            metrics['precision-at-precision>={:.2}'.format(target)] = precision

        # Add prefix before returning.
        return dict([(prefix + k, v) for k, v in metrics.items()])

infile
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
    macro_f1 = eval_utils.safe_divide(
        2 * macro_precision * macro_recall,
        macro_precision + macro_recall)

    return collections.OrderedDict({
        prefix + 'n': len(answer_stats),
        prefix + 'f1': macro_f1,
        prefix + 'precision': macro_precision,
        prefix + 'recall': macro_recall
    })


def compute_final_f1(passage_answer_stats, minimal_answer_stats):
    """Computes overall F1 given passage and minimal answers, ignoring scores.

    Note: this assumes that the answers have been thresholded based on scores.

    Arguments:
       passage_answer_stats: List of passage answer scores.
       minimal_answer_stats: List of minimal answer scores.


    Returns:
       Dictionary of name (string) -> score.
    """
    scores = compute_macro_f1(passage_answer_stats, prefix='passage-answer-')
    scores.update(compute_macro_f1(
        minimal_answer_stats, prefix='minimal-answer-'))
    return scores


def compute_pr_curves(answer_stats, targets=None):
    """Computes PR curve and returns R@P for specific targets.

    The values are computed as follows: find the (precision, recall) point
    with maximum recall and where precision > target.

    This is only relevant if you return the system scores in your predictions.
    You may find this useful when attempting to tune the threshold for your
    system on the dev set before requesting an evaluation on the test set
    via the leaderboard.

    Arguments:
      answer_stats: List of statistic tuples from the answer scores.
      targets (None): List of precision thresholds to target.

    Returns:
      List of table with rows: [target, r, p, score].
    """
    total_f1 = 0
    total_has_pred = 0
    total_has_gold = 0

    # Count the number of gold annotations.
    for has_gold, _, _, _ in answer_stats:
        total_has_gold += has_gold

    # Keep track of the point of maximum recall for each target.

    max_recall = [0 for _ in targets]
    max_precision = [0 for _ in targets]
    max_scores = [0.0 for _ in targets]

    # Only keep track of unique thresholds in this dictionary.
    scores_to_stats = collections.OrderedDict()

    # Loop through every possible threshold and compute precision + recall.
    for has_gold, has_pred, is_correct_or_f1, score in answer_stats:
        if isinstance(is_correct_or_f1, tuple):
            _, _, f1 = is_correct_or_f1
        else:
            f1 = is_correct_or_f1
        total_f1 += f1
        total_has_pred += has_pred

        precision = eval_utils.safe_divide(total_f1, total_has_pred)
        recall = eval_utils.safe_divide(total_f1, total_has_gold)

        # If there are any ties, this will be updated multiple times until the
        # ties are all counted.
        scores_to_stats[score] = [precision, recall]

    best_f1 = 0.0
    best_precision = 0.0
    best_recall = 0.0
    best_threshold = 0.0

    for threshold, (precision, recall) in scores_to_stats.items():
        # Match the thresholds to the find the closest precision above some target.
        for t, target in enumerate(targets):
            if precision >= target and recall > max_recall[t]:
                max_recall[t] = recall
                max_precision[t] = precision
                max_scores[t] = threshold

        # Compute optimal threshold.
        f1 = eval_utils.safe_divide(2 * precision * recall, precision + recall)
        if f1 > best_f1:
            best_f1 = f1
            best_precision = precision
            best_recall = recall
            best_threshold = threshold

    return ((best_f1, best_precision, best_recall, best_threshold),
            list(zip(targets, max_recall, max_precision, max_scores)))


def print_r_at_p_table(answer_stats):
    """Pretty prints the R@P table for default targets."""
    opt_result, pr_table = compute_pr_curves(
        answer_stats, targets=[0.5, 0.75, 0.9])
    f1, precision, recall, threshold = opt_result
    print('Optimal threshold: {:.5}'.format(threshold))
    print(' F1     /  P      /  R')
    print('{: >7.2%} / {: >7.2%} / {: >7.2%}'.format(f1, precision, recall))
    for target, recall, precision, row in pr_table:
        print('R@P={}: {:.2%} (actual p={:.2%}, score threshold={:.4})'.format(
            target, recall, precision, row))


def get_metrics_as_dict(gold_path, prediction_path, passage_non_null_threshold, minimal_non_null_threshold, verbose):
    """Library version of the end-to-end evaluation.

    Arguments:
      gold_path: Path to a single JSONL data. Could be gzipped or not.
      prediction_path: Path to the JSONL file of prediction data.

    Returns:
      metrics: A dictionary mapping string names to metric scores.
    """

    tydi_gold_dict = eval_utils.read_annotation(gold_path)
    tydi_pred_dict = eval_utils.read_prediction_jsonl(prediction_path)

    passage_answer_stats, minimal_answer_stats = score_answers(
        tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose)

    return get_metrics_with_answer_stats(
        passage_answer_stats, minimal_answer_stats)


def get_metrics_with_answer_stats(passage_answer_stats, minimal_answer_stats):
    """Generate metrics dict using passage and minimal answer stats."""

    def _get_metric_dict(answer_stats, prefix=''):
        """Compute all metrics for a set of answer statistics."""
        opt_result, pr_table = compute_pr_curves(
            answer_stats, targets=[0.5, 0.75, 0.9])
        f1, precision, recall, threshold = opt_result
        metrics = collections.OrderedDict({
            'best-threshold-f1': f1,
            'best-threshold-precision': precision,
            'best-threshold-recall': recall,
            'best-threshold': threshold,
        })
        for target, recall, precision, _ in pr_table:
            metrics['recall-at-precision>={:.2}'.format(target)] = recall
            metrics['precision-at-precision>={:.2}'.format(target)] = precision

        # Add prefix before returning.
        return dict([(prefix + k, v) for k, v in metrics.items()])

    metrics = _get_metric_dict(passage_answer_stats, 'passage-')
infile
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
common
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
common
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
    macro_f1 = eval_utils.safe_divide(
        2 * macro_precision * macro_recall,
        macro_precision + macro_recall)

    return collections.OrderedDict({
        prefix + 'n': len(answer_stats),
        prefix + 'f1': macro_f1,
        prefix + 'precision': macro_precision,
        prefix + 'recall': macro_recall
    })


def compute_final_f1(passage_answer_stats, minimal_answer_stats):
    """Computes overall F1 given passage and minimal answers, ignoring scores.

    Note: this assumes that the answers have been thresholded based on scores.

    Arguments:
       passage_answer_stats: List of passage answer scores.
       minimal_answer_stats: List of minimal answer scores.


    Returns:
       Dictionary of name (string) -> score.
    """
    scores = compute_macro_f1(passage_answer_stats, prefix='passage-answer-')
    scores.update(compute_macro_f1(
        minimal_answer_stats, prefix='minimal-answer-'))
    return scores


def compute_pr_curves(answer_stats, targets=None):
    """Computes PR curve and returns R@P for specific targets.

    The values are computed as follows: find the (precision, recall) point
    with maximum recall and where precision > target.

    This is only relevant if you return the system scores in your predictions.
    You may find this useful when attempting to tune the threshold for your
    system on the dev set before requesting an evaluation on the test set
    via the leaderboard.

    Arguments:
      answer_stats: List of statistic tuples from the answer scores.
      targets (None): List of precision thresholds to target.

    Returns:
      List of table with rows: [target, r, p, score].
    """
    total_f1 = 0
    total_has_pred = 0
    total_has_gold = 0

    # Count the number of gold annotations.
    for has_gold, _, _, _ in answer_stats:
        total_has_gold += has_gold

    # Keep track of the point of maximum recall for each target.

    max_recall = [0 for _ in targets]
    max_precision = [0 for _ in targets]
    max_scores = [0.0 for _ in targets]

    # Only keep track of unique thresholds in this dictionary.
    scores_to_stats = collections.OrderedDict()

    # Loop through every possible threshold and compute precision + recall.
    for has_gold, has_pred, is_correct_or_f1, score in answer_stats:
        if isinstance(is_correct_or_f1, tuple):
            _, _, f1 = is_correct_or_f1
        else:
            f1 = is_correct_or_f1
        total_f1 += f1
        total_has_pred += has_pred

        precision = eval_utils.safe_divide(total_f1, total_has_pred)
        recall = eval_utils.safe_divide(total_f1, total_has_gold)

        # If there are any ties, this will be updated multiple times until the
        # ties are all counted.
        scores_to_stats[score] = [precision, recall]

    best_f1 = 0.0
    best_precision = 0.0
    best_recall = 0.0
    best_threshold = 0.0

    for threshold, (precision, recall) in scores_to_stats.items():
        # Match the thresholds to the find the closest precision above some target.
        for t, target in enumerate(targets):
            if precision >= target and recall > max_recall[t]:
                max_recall[t] = recall
                max_precision[t] = precision
                max_scores[t] = threshold

        # Compute optimal threshold.
        f1 = eval_utils.safe_divide(2 * precision * recall, precision + recall)
        if f1 > best_f1:
            best_f1 = f1
            best_precision = precision
            best_recall = recall
            best_threshold = threshold

    return ((best_f1, best_precision, best_recall, best_threshold),
            list(zip(targets, max_recall, max_precision, max_scores)))


def print_r_at_p_table(answer_stats):
    """Pretty prints the R@P table for default targets."""
    opt_result, pr_table = compute_pr_curves(
        answer_stats, targets=[0.5, 0.75, 0.9])
    f1, precision, recall, threshold = opt_result
    print('Optimal threshold: {:.5}'.format(threshold))
    print(' F1     /  P      /  R')
    print('{: >7.2%} / {: >7.2%} / {: >7.2%}'.format(f1, precision, recall))
    for target, recall, precision, row in pr_table:
        print('R@P={}: {:.2%} (actual p={:.2%}, score threshold={:.4})'.format(
            target, recall, precision, row))


def get_metrics_as_dict(gold_path, prediction_path, passage_non_null_threshold, minimal_non_null_threshold, verbose):
    """Library version of the end-to-end evaluation.

    Arguments:
      gold_path: Path to a single JSONL data. Could be gzipped or not.
      prediction_path: Path to the JSONL file of prediction data.

    Returns:
      metrics: A dictionary mapping string names to metric scores.
    """

    tydi_gold_dict = eval_utils.read_annotation(gold_path)
    tydi_pred_dict = eval_utils.read_prediction_jsonl(prediction_path)

    passage_answer_stats, minimal_answer_stats = score_answers(
        tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose)

    return get_metrics_with_answer_stats(
        passage_answer_stats, minimal_answer_stats)


def get_metrics_with_answer_stats(passage_answer_stats, minimal_answer_stats):
    """Generate metrics dict using passage and minimal answer stats."""

    def _get_metric_dict(answer_stats, prefix=''):
        """Compute all metrics for a set of answer statistics."""
        opt_result, pr_table = compute_pr_curves(
            answer_stats, targets=[0.5, 0.75, 0.9])
        f1, precision, recall, threshold = opt_result
        metrics = collections.OrderedDict({
            'best-threshold-f1': f1,
            'best-threshold-precision': precision,
            'best-threshold-recall': recall,
            'best-threshold': threshold,
        })
        for target, recall, precision, _ in pr_table:
            metrics['recall-at-precision>={:.2}'.format(target)] = recall
            metrics['precision-at-precision>={:.2}'.format(target)] = precision

        # Add prefix before returning.
        return dict([(prefix + k, v) for k, v in metrics.items()])

    metrics = _get_metric_dict(passage_answer_stats, 'passage-')
    metrics.update(_get_metric_dict(minimal_answer_stats, 'minimal-'))
    return metrics


def get_latex_str(f1, precision, recall):
    return '\\fpr' + '{' + ('%.1f' % (f1 * 100)) + '}{' + (
            '%.1f' % (precision * 100)) + '}{' + ('%.1f' % (recall * 100)) + '}'


# def load_gt_lookup_as_dict(gold_path, num_workers):
#     cache_path = os.path.join(os.path.dirname(gold_path), 'cache')
#
#     if os.path.exists(cache_path):
#         logging.info('Reading from cache: %s', format(cache_path))
#         tydi_gold_dict = pickle.load(open(cache_path, 'rb'))  # pytype: disable=wrong-arg-types
#     else:
#         tydi_gold_dict = eval_utils.read_annotation(gold_path, n_threads=num_workers)
#         logging.info('Caching gold data for future to: %s', format(cache_path))
#         cache_to_pickle_file(tydi_gold_dict, cache_path)  # pytype: disable=wrong-arg-types
#
#         logging.info('Read in gt lookup for %d examples into memory' % len(tydi_gold_dict))
#     return tydi_gold_dict


def pretty_print(tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold=2, minimal_non_null_threshold=2, verbose=False, skip_missing_example_ids=False):
    if any(map(not_, tydi_gold_dict.values())) or any(annotations[0].minimal_answer_span is None for annotations in tydi_gold_dict.values()):
        logging.warning("At least one example in gold dict has no annotations -- skipping pretty print evaluation")
        return

    gold_id_set = set(tydi_gold_dict.keys())
    pred_id_set = set(tydi_pred_dict.keys())
    sym_diff = gold_id_set.symmetric_difference(pred_id_set)

    if (not skip_missing_example_ids) and sym_diff:
        raise ValueError('ERROR: the example ids in gold annotations and example '
                         'ids in the prediction are not equal.')
    elif skip_missing_example_ids and sym_diff:
        logging.warning("Skipping {} example ids that are only in either gold or preds".format(len(sym_diff)))

    macro_avg_passage_scores = ([], [], [])
    macro_avg_minimal_scores = ([], [], [])

    language_list = [  # TODO: support arbitrary languages
        'english', 'arabic', 'bengali', 'finnish', 'indonesian', 'japanese',
        'swahili', 'korean', 'russian', 'telugu', 'thai'
    ]

    per_lang_gold = {}
    per_lang_pred = {}

    id_to_lang = {}

    for ex_id in gold_id_set:
        ex = tydi_gold_dict[ex_id]
        if ex[0].language in per_lang_gold:
            per_lang_gold[ex[0].language][ex_id] = ex
        else:
            per_lang_gold[ex[0].language] = {ex_id: ex}
        id_to_lang[ex_id] = ex[0].language
    for ex_id in gold_id_set:
        ex = tydi_pred_dict[ex_id]
        if id_to_lang[ex.example_id] in per_lang_pred:
            per_lang_pred[id_to_lang[ex.example_id]][ex.example_id] = ex
        else:
            per_lang_pred[id_to_lang[ex.example_id]] = {ex.example_id: ex}

    for lang in language_list:
        if lang in per_lang_pred:
            passage_answer_stats, minimal_answer_stats = score_answers(
common
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
common
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
    macro_f1 = eval_utils.safe_divide(
        2 * macro_precision * macro_recall,
        macro_precision + macro_recall)

    return collections.OrderedDict({
        prefix + 'n': len(answer_stats),
        prefix + 'f1': macro_f1,
        prefix + 'precision': macro_precision,
        prefix + 'recall': macro_recall
    })


def compute_final_f1(passage_answer_stats, minimal_answer_stats):
    """Computes overall F1 given passage and minimal answers, ignoring scores.

    Note: this assumes that the answers have been thresholded based on scores.

    Arguments:
       passage_answer_stats: List of passage answer scores.
       minimal_answer_stats: List of minimal answer scores.


    Returns:
       Dictionary of name (string) -> score.
    """
    scores = compute_macro_f1(passage_answer_stats, prefix='passage-answer-')
    scores.update(compute_macro_f1(
        minimal_answer_stats, prefix='minimal-answer-'))
    return scores


def compute_pr_curves(answer_stats, targets=None):
    """Computes PR curve and returns R@P for specific targets.

    The values are computed as follows: find the (precision, recall) point
    with maximum recall and where precision > target.

    This is only relevant if you return the system scores in your predictions.
    You may find this useful when attempting to tune the threshold for your
    system on the dev set before requesting an evaluation on the test set
    via the leaderboard.

    Arguments:
      answer_stats: List of statistic tuples from the answer scores.
      targets (None): List of precision thresholds to target.

    Returns:
      List of table with rows: [target, r, p, score].
    """
    total_f1 = 0
    total_has_pred = 0
    total_has_gold = 0

    # Count the number of gold annotations.
    for has_gold, _, _, _ in answer_stats:
        total_has_gold += has_gold

    # Keep track of the point of maximum recall for each target.

    max_recall = [0 for _ in targets]
    max_precision = [0 for _ in targets]
    max_scores = [0.0 for _ in targets]

    # Only keep track of unique thresholds in this dictionary.
    scores_to_stats = collections.OrderedDict()

    # Loop through every possible threshold and compute precision + recall.
    for has_gold, has_pred, is_correct_or_f1, score in answer_stats:
        if isinstance(is_correct_or_f1, tuple):
            _, _, f1 = is_correct_or_f1
        else:
            f1 = is_correct_or_f1
        total_f1 += f1
        total_has_pred += has_pred

        precision = eval_utils.safe_divide(total_f1, total_has_pred)
commited
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
commited
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
    macro_f1 = eval_utils.safe_divide(
        2 * macro_precision * macro_recall,
        macro_precision + macro_recall)

    return collections.OrderedDict({
        prefix + 'n': len(answer_stats),
        prefix + 'f1': macro_f1,
        prefix + 'precision': macro_precision,
        prefix + 'recall': macro_recall
    })


def compute_final_f1(passage_answer_stats, minimal_answer_stats):
    """Computes overall F1 given passage and minimal answers, ignoring scores.

    Note: this assumes that the answers have been thresholded based on scores.

    Arguments:
       passage_answer_stats: List of passage answer scores.
       minimal_answer_stats: List of minimal answer scores.


    Returns:
       Dictionary of name (string) -> score.
    """
    scores = compute_macro_f1(passage_answer_stats, prefix='passage-answer-')
    scores.update(compute_macro_f1(
        minimal_answer_stats, prefix='minimal-answer-'))
    return scores


def compute_pr_curves(answer_stats, targets=None):
    """Computes PR curve and returns R@P for specific targets.

    The values are computed as follows: find the (precision, recall) point
    with maximum recall and where precision > target.

    This is only relevant if you return the system scores in your predictions.
    You may find this useful when attempting to tune the threshold for your
    system on the dev set before requesting an evaluation on the test set
    via the leaderboard.

    Arguments:
      answer_stats: List of statistic tuples from the answer scores.
      targets (None): List of precision thresholds to target.

    Returns:
      List of table with rows: [target, r, p, score].
    """
    total_f1 = 0
    total_has_pred = 0
    total_has_gold = 0

    # Count the number of gold annotations.
    for has_gold, _, _, _ in answer_stats:
        total_has_gold += has_gold

    # Keep track of the point of maximum recall for each target.

    max_recall = [0 for _ in targets]
    max_precision = [0 for _ in targets]
    max_scores = [0.0 for _ in targets]

    # Only keep track of unique thresholds in this dictionary.
    scores_to_stats = collections.OrderedDict()

    # Loop through every possible threshold and compute precision + recall.
    for has_gold, has_pred, is_correct_or_f1, score in answer_stats:
        if isinstance(is_correct_or_f1, tuple):
            _, _, f1 = is_correct_or_f1
        else:
            f1 = is_correct_or_f1
        total_f1 += f1
        total_has_pred += has_pred

        precision = eval_utils.safe_divide(total_f1, total_has_pred)
        recall = eval_utils.safe_divide(total_f1, total_has_gold)

        # If there are any ties, this will be updated multiple times until the
        # ties are all counted.
        scores_to_stats[score] = [precision, recall]

    best_f1 = 0.0
    best_precision = 0.0
    best_recall = 0.0
    best_threshold = 0.0

    for threshold, (precision, recall) in scores_to_stats.items():
        # Match the thresholds to the find the closest precision above some target.
        for t, target in enumerate(targets):
            if precision >= target and recall > max_recall[t]:
                max_recall[t] = recall
                max_precision[t] = precision
                max_scores[t] = threshold

        # Compute optimal threshold.
commited
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
commited
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
commited
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

commited
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
commited
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
commited
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

commited
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
    macro_f1 = eval_utils.safe_divide(
        2 * macro_precision * macro_recall,
        macro_precision + macro_recall)

    return collections.OrderedDict({
        prefix + 'n': len(answer_stats),
        prefix + 'f1': macro_f1,
        prefix + 'precision': macro_precision,
        prefix + 'recall': macro_recall
    })


def compute_final_f1(passage_answer_stats, minimal_answer_stats):
    """Computes overall F1 given passage and minimal answers, ignoring scores.

    Note: this assumes that the answers have been thresholded based on scores.

    Arguments:
       passage_answer_stats: List of passage answer scores.
       minimal_answer_stats: List of minimal answer scores.


    Returns:
       Dictionary of name (string) -> score.
    """
    scores = compute_macro_f1(passage_answer_stats, prefix='passage-answer-')
    scores.update(compute_macro_f1(
        minimal_answer_stats, prefix='minimal-answer-'))
    return scores


def compute_pr_curves(answer_stats, targets=None):
    """Computes PR curve and returns R@P for specific targets.

    The values are computed as follows: find the (precision, recall) point
    with maximum recall and where precision > target.

    This is only relevant if you return the system scores in your predictions.
    You may find this useful when attempting to tune the threshold for your
    system on the dev set before requesting an evaluation on the test set
    via the leaderboard.

    Arguments:
      answer_stats: List of statistic tuples from the answer scores.
      targets (None): List of precision thresholds to target.

    Returns:
      List of table with rows: [target, r, p, score].
    """
    total_f1 = 0
    total_has_pred = 0
    total_has_gold = 0

    # Count the number of gold annotations.
    for has_gold, _, _, _ in answer_stats:
        total_has_gold += has_gold

    # Keep track of the point of maximum recall for each target.

    max_recall = [0 for _ in targets]
    max_precision = [0 for _ in targets]
    max_scores = [0.0 for _ in targets]

    # Only keep track of unique thresholds in this dictionary.
    scores_to_stats = collections.OrderedDict()

    # Loop through every possible threshold and compute precision + recall.
    for has_gold, has_pred, is_correct_or_f1, score in answer_stats:
        if isinstance(is_correct_or_f1, tuple):
            _, _, f1 = is_correct_or_f1
        else:
            f1 = is_correct_or_f1
        total_f1 += f1
        total_has_pred += has_pred

        precision = eval_utils.safe_divide(total_f1, total_has_pred)
        recall = eval_utils.safe_divide(total_f1, total_has_gold)

        # If there are any ties, this will be updated multiple times until the
        # ties are all counted.
        scores_to_stats[score] = [precision, recall]

    best_f1 = 0.0
    best_precision = 0.0
    best_recall = 0.0
    best_threshold = 0.0

    for threshold, (precision, recall) in scores_to_stats.items():
        # Match the thresholds to the find the closest precision above some target.
        for t, target in enumerate(targets):
            if precision >= target and recall > max_recall[t]:
                max_recall[t] = recall
                max_precision[t] = precision
                max_scores[t] = threshold

        # Compute optimal threshold.
        f1 = eval_utils.safe_divide(2 * precision * recall, precision + recall)
        if f1 > best_f1:
            best_f1 = f1
            best_precision = precision
            best_recall = recall
            best_threshold = threshold

    return ((best_f1, best_precision, best_recall, best_threshold),
            list(zip(targets, max_recall, max_precision, max_scores)))


def print_r_at_p_table(answer_stats):
    """Pretty prints the R@P table for default targets."""
    opt_result, pr_table = compute_pr_curves(
        answer_stats, targets=[0.5, 0.75, 0.9])
    f1, precision, recall, threshold = opt_result
    print('Optimal threshold: {:.5}'.format(threshold))
    print(' F1     /  P      /  R')
    print('{: >7.2%} / {: >7.2%} / {: >7.2%}'.format(f1, precision, recall))
    for target, recall, precision, row in pr_table:
        print('R@P={}: {:.2%} (actual p={:.2%}, score threshold={:.4})'.format(
            target, recall, precision, row))


def get_metrics_as_dict(gold_path, prediction_path, passage_non_null_threshold, minimal_non_null_threshold, verbose):
    """Library version of the end-to-end evaluation.

    Arguments:
      gold_path: Path to a single JSONL data. Could be gzipped or not.
      prediction_path: Path to the JSONL file of prediction data.

    Returns:
      metrics: A dictionary mapping string names to metric scores.
    """

    tydi_gold_dict = eval_utils.read_annotation(gold_path)
    tydi_pred_dict = eval_utils.read_prediction_jsonl(prediction_path)

    passage_answer_stats, minimal_answer_stats = score_answers(
        tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose)

    return get_metrics_with_answer_stats(
        passage_answer_stats, minimal_answer_stats)


def get_metrics_with_answer_stats(passage_answer_stats, minimal_answer_stats):
    """Generate metrics dict using passage and minimal answer stats."""

    def _get_metric_dict(answer_stats, prefix=''):
        """Compute all metrics for a set of answer statistics."""
        opt_result, pr_table = compute_pr_curves(
            answer_stats, targets=[0.5, 0.75, 0.9])
        f1, precision, recall, threshold = opt_result
        metrics = collections.OrderedDict({
            'best-threshold-f1': f1,
            'best-threshold-precision': precision,
            'best-threshold-recall': recall,
            'best-threshold': threshold,
        })
        for target, recall, precision, _ in pr_table:
            metrics['recall-at-precision>={:.2}'.format(target)] = recall
            metrics['precision-at-precision>={:.2}'.format(target)] = precision

        # Add prefix before returning.
        return dict([(prefix + k, v) for k, v in metrics.items()])

    metrics = _get_metric_dict(passage_answer_stats, 'passage-')
    metrics.update(_get_metric_dict(minimal_answer_stats, 'minimal-'))
    return metrics


def get_latex_str(f1, precision, recall):
    return '\\fpr' + '{' + ('%.1f' % (f1 * 100)) + '}{' + (
            '%.1f' % (precision * 100)) + '}{' + ('%.1f' % (recall * 100)) + '}'


# def load_gt_lookup_as_dict(gold_path, num_workers):
#     cache_path = os.path.join(os.path.dirname(gold_path), 'cache')
#
#     if os.path.exists(cache_path):
#         logging.info('Reading from cache: %s', format(cache_path))
#         tydi_gold_dict = pickle.load(open(cache_path, 'rb'))  # pytype: disable=wrong-arg-types
#     else:
#         tydi_gold_dict = eval_utils.read_annotation(gold_path, n_threads=num_workers)
#         logging.info('Caching gold data for future to: %s', format(cache_path))
#         cache_to_pickle_file(tydi_gold_dict, cache_path)  # pytype: disable=wrong-arg-types
#
#         logging.info('Read in gt lookup for %d examples into memory' % len(tydi_gold_dict))
#     return tydi_gold_dict


def pretty_print(tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold=2, minimal_non_null_threshold=2, verbose=False, skip_missing_example_ids=False):
    if any(map(not_, tydi_gold_dict.values())) or any(annotations[0].minimal_answer_span is None for annotations in tydi_gold_dict.values()):
        logging.warning("At least one example in gold dict has no annotations -- skipping pretty print evaluation")
        return

    gold_id_set = set(tydi_gold_dict.keys())
    pred_id_set = set(tydi_pred_dict.keys())
    sym_diff = gold_id_set.symmetric_difference(pred_id_set)

    if (not skip_missing_example_ids) and sym_diff:
        raise ValueError('ERROR: the example ids in gold annotations and example '
                         'ids in the prediction are not equal.')
    elif skip_missing_example_ids and sym_diff:
        logging.warning("Skipping {} example ids that are only in either gold or preds".format(len(sym_diff)))

    macro_avg_passage_scores = ([], [], [])
    macro_avg_minimal_scores = ([], [], [])

    language_list = [  # TODO: support arbitrary languages
        'english', 'arabic', 'bengali', 'finnish', 'indonesian', 'japanese',
        'swahili', 'korean', 'russian', 'telugu', 'thai'
    ]

    per_lang_gold = {}
    per_lang_pred = {}

    id_to_lang = {}

    for ex_id in gold_id_set:
        ex = tydi_gold_dict[ex_id]
        if ex[0].language in per_lang_gold:
            per_lang_gold[ex[0].language][ex_id] = ex
        else:
            per_lang_gold[ex[0].language] = {ex_id: ex}
        id_to_lang[ex_id] = ex[0].language
    for ex_id in gold_id_set:
        ex = tydi_pred_dict[ex_id]
        if id_to_lang[ex.example_id] in per_lang_pred:
            per_lang_pred[id_to_lang[ex.example_id]][ex.example_id] = ex
        else:
            per_lang_pred[id_to_lang[ex.example_id]] = {ex.example_id: ex}

    for lang in language_list:
        if lang in per_lang_pred:
            passage_answer_stats, minimal_answer_stats = score_answers(
                per_lang_gold.get(lang, {}), per_lang_pred[lang],
                passage_non_null_threshold, minimal_non_null_threshold, verbose)

            # Passage selection task
            opt_result, _ = compute_pr_curves(passage_answer_stats, targets=[0.5])
            f1, precision, recall, _ = opt_result
            if lang != 'english':
                macro_avg_passage_scores[0].append(f1)
                macro_avg_passage_scores[1].append(precision)
                macro_avg_passage_scores[2].append(recall)
            print('Passage & ' + lang + ' & ' + get_latex_str(f1, precision, recall))

            # Minimal answer span task
            opt_result, _ = compute_pr_curves(minimal_answer_stats, targets=[0.5])
            f1, precision, recall, _ = opt_result
            if lang != 'english':
                macro_avg_minimal_scores[0].append(f1)
                macro_avg_minimal_scores[1].append(precision)
                macro_avg_minimal_scores[2].append(recall)
            print('Minimal Answer & ' + lang + ' & ' +
                  get_latex_str(f1, precision, recall))

            print('*' * 20)
            print(lang)
            print('Language: %s (%d)' % (lang, len(per_lang_gold.get(lang, {}))))
            print('*' * 20)
            print('PASSAGE ANSWER R@P TABLE:')
            print_r_at_p_table(passage_answer_stats)
            print('*' * 20)
            print('MINIMAL ANSWER R@P TABLE:')
            print_r_at_p_table(minimal_answer_stats)

    print('Total # examples in gold: %d, # ex. in pred: %d (including english)' %
          (len(tydi_gold_dict), len(tydi_pred_dict)))

    f1_list, precision_list, recall_list = macro_avg_passage_scores
    print('*** Macro Over %d Languages, excluding English **' % len(f1_list))
commited
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
    macro_f1 = eval_utils.safe_divide(
        2 * macro_precision * macro_recall,
        macro_precision + macro_recall)

    return collections.OrderedDict({
        prefix + 'n': len(answer_stats),
        prefix + 'f1': macro_f1,
        prefix + 'precision': macro_precision,
        prefix + 'recall': macro_recall
    })


def compute_final_f1(passage_answer_stats, minimal_answer_stats):
    """Computes overall F1 given passage and minimal answers, ignoring scores.

    Note: this assumes that the answers have been thresholded based on scores.

    Arguments:
       passage_answer_stats: List of passage answer scores.
       minimal_answer_stats: List of minimal answer scores.


    Returns:
       Dictionary of name (string) -> score.
    """
    scores = compute_macro_f1(passage_answer_stats, prefix='passage-answer-')
    scores.update(compute_macro_f1(
        minimal_answer_stats, prefix='minimal-answer-'))
    return scores


def compute_pr_curves(answer_stats, targets=None):
    """Computes PR curve and returns R@P for specific targets.

    The values are computed as follows: find the (precision, recall) point
    with maximum recall and where precision > target.

    This is only relevant if you return the system scores in your predictions.
    You may find this useful when attempting to tune the threshold for your
    system on the dev set before requesting an evaluation on the test set
    via the leaderboard.

    Arguments:
      answer_stats: List of statistic tuples from the answer scores.
      targets (None): List of precision thresholds to target.

    Returns:
      List of table with rows: [target, r, p, score].
    """
    total_f1 = 0
    total_has_pred = 0
    total_has_gold = 0

    # Count the number of gold annotations.
    for has_gold, _, _, _ in answer_stats:
        total_has_gold += has_gold

    # Keep track of the point of maximum recall for each target.

    max_recall = [0 for _ in targets]
    max_precision = [0 for _ in targets]
    max_scores = [0.0 for _ in targets]

    # Only keep track of unique thresholds in this dictionary.
    scores_to_stats = collections.OrderedDict()

    # Loop through every possible threshold and compute precision + recall.
    for has_gold, has_pred, is_correct_or_f1, score in answer_stats:
        if isinstance(is_correct_or_f1, tuple):
            _, _, f1 = is_correct_or_f1
        else:
            f1 = is_correct_or_f1
        total_f1 += f1
        total_has_pred += has_pred

        precision = eval_utils.safe_divide(total_f1, total_has_pred)
        recall = eval_utils.safe_divide(total_f1, total_has_gold)

        # If there are any ties, this will be updated multiple times until the
        # ties are all counted.
        scores_to_stats[score] = [precision, recall]

    best_f1 = 0.0
    best_precision = 0.0
    best_recall = 0.0
    best_threshold = 0.0

    for threshold, (precision, recall) in scores_to_stats.items():
        # Match the thresholds to the find the closest precision above some target.
        for t, target in enumerate(targets):
            if precision >= target and recall > max_recall[t]:
                max_recall[t] = recall
                max_precision[t] = precision
                max_scores[t] = threshold

        # Compute optimal threshold.
        f1 = eval_utils.safe_divide(2 * precision * recall, precision + recall)
        if f1 > best_f1:
            best_f1 = f1
            best_precision = precision
            best_recall = recall
            best_threshold = threshold

    return ((best_f1, best_precision, best_recall, best_threshold),
            list(zip(targets, max_recall, max_precision, max_scores)))


def print_r_at_p_table(answer_stats):
    """Pretty prints the R@P table for default targets."""
    opt_result, pr_table = compute_pr_curves(
        answer_stats, targets=[0.5, 0.75, 0.9])
    f1, precision, recall, threshold = opt_result
    print('Optimal threshold: {:.5}'.format(threshold))
    print(' F1     /  P      /  R')
    print('{: >7.2%} / {: >7.2%} / {: >7.2%}'.format(f1, precision, recall))
    for target, recall, precision, row in pr_table:
        print('R@P={}: {:.2%} (actual p={:.2%}, score threshold={:.4})'.format(
            target, recall, precision, row))


def get_metrics_as_dict(gold_path, prediction_path, passage_non_null_threshold, minimal_non_null_threshold, verbose):
    """Library version of the end-to-end evaluation.

    Arguments:
      gold_path: Path to a single JSONL data. Could be gzipped or not.
      prediction_path: Path to the JSONL file of prediction data.

    Returns:
      metrics: A dictionary mapping string names to metric scores.
    """

    tydi_gold_dict = eval_utils.read_annotation(gold_path)
    tydi_pred_dict = eval_utils.read_prediction_jsonl(prediction_path)

    passage_answer_stats, minimal_answer_stats = score_answers(
        tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose)

    return get_metrics_with_answer_stats(
        passage_answer_stats, minimal_answer_stats)


def get_metrics_with_answer_stats(passage_answer_stats, minimal_answer_stats):
    """Generate metrics dict using passage and minimal answer stats."""

    def _get_metric_dict(answer_stats, prefix=''):
        """Compute all metrics for a set of answer statistics."""
        opt_result, pr_table = compute_pr_curves(
            answer_stats, targets=[0.5, 0.75, 0.9])
        f1, precision, recall, threshold = opt_result
        metrics = collections.OrderedDict({
            'best-threshold-f1': f1,
            'best-threshold-precision': precision,
            'best-threshold-recall': recall,
            'best-threshold': threshold,
        })
        for target, recall, precision, _ in pr_table:
            metrics['recall-at-precision>={:.2}'.format(target)] = recall
            metrics['precision-at-precision>={:.2}'.format(target)] = precision

        # Add prefix before returning.
        return dict([(prefix + k, v) for k, v in metrics.items()])

    metrics = _get_metric_dict(passage_answer_stats, 'passage-')
    metrics.update(_get_metric_dict(minimal_answer_stats, 'minimal-'))
    return metrics


def get_latex_str(f1, precision, recall):
    return '\\fpr' + '{' + ('%.1f' % (f1 * 100)) + '}{' + (
            '%.1f' % (precision * 100)) + '}{' + ('%.1f' % (recall * 100)) + '}'


# def load_gt_lookup_as_dict(gold_path, num_workers):
#     cache_path = os.path.join(os.path.dirname(gold_path), 'cache')
#
#     if os.path.exists(cache_path):
#         logging.info('Reading from cache: %s', format(cache_path))
#         tydi_gold_dict = pickle.load(open(cache_path, 'rb'))  # pytype: disable=wrong-arg-types
#     else:
#         tydi_gold_dict = eval_utils.read_annotation(gold_path, n_threads=num_workers)
#         logging.info('Caching gold data for future to: %s', format(cache_path))
#         cache_to_pickle_file(tydi_gold_dict, cache_path)  # pytype: disable=wrong-arg-types
#
#         logging.info('Read in gt lookup for %d examples into memory' % len(tydi_gold_dict))
#     return tydi_gold_dict


def pretty_print(tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold=2, minimal_non_null_threshold=2, verbose=False, skip_missing_example_ids=False):
    if any(map(not_, tydi_gold_dict.values())) or any(annotations[0].minimal_answer_span is None for annotations in tydi_gold_dict.values()):
        logging.warning("At least one example in gold dict has no annotations -- skipping pretty print evaluation")
        return

    gold_id_set = set(tydi_gold_dict.keys())
    pred_id_set = set(tydi_pred_dict.keys())
    sym_diff = gold_id_set.symmetric_difference(pred_id_set)

    if (not skip_missing_example_ids) and sym_diff:
        raise ValueError('ERROR: the example ids in gold annotations and example '
                         'ids in the prediction are not equal.')
    elif skip_missing_example_ids and sym_diff:
        logging.warning("Skipping {} example ids that are only in either gold or preds".format(len(sym_diff)))

    macro_avg_passage_scores = ([], [], [])
    macro_avg_minimal_scores = ([], [], [])

    language_list = [  # TODO: support arbitrary languages
        'english', 'arabic', 'bengali', 'finnish', 'indonesian', 'japanese',
        'swahili', 'korean', 'russian', 'telugu', 'thai'
    ]

    per_lang_gold = {}
    per_lang_pred = {}

    id_to_lang = {}

    for ex_id in gold_id_set:
        ex = tydi_gold_dict[ex_id]
        if ex[0].language in per_lang_gold:
            per_lang_gold[ex[0].language][ex_id] = ex
        else:
            per_lang_gold[ex[0].language] = {ex_id: ex}
        id_to_lang[ex_id] = ex[0].language
    for ex_id in gold_id_set:
        ex = tydi_pred_dict[ex_id]
        if id_to_lang[ex.example_id] in per_lang_pred:
            per_lang_pred[id_to_lang[ex.example_id]][ex.example_id] = ex
        else:
            per_lang_pred[id_to_lang[ex.example_id]] = {ex.example_id: ex}

    for lang in language_list:
        if lang in per_lang_pred:
            passage_answer_stats, minimal_answer_stats = score_answers(
                per_lang_gold.get(lang, {}), per_lang_pred[lang],
                passage_non_null_threshold, minimal_non_null_threshold, verbose)

            # Passage selection task
            opt_result, _ = compute_pr_curves(passage_answer_stats, targets=[0.5])
            f1, precision, recall, _ = opt_result
            if lang != 'english':
                macro_avg_passage_scores[0].append(f1)
                macro_avg_passage_scores[1].append(precision)
                macro_avg_passage_scores[2].append(recall)
            print('Passage & ' + lang + ' & ' + get_latex_str(f1, precision, recall))

            # Minimal answer span task
            opt_result, _ = compute_pr_curves(minimal_answer_stats, targets=[0.5])
            f1, precision, recall, _ = opt_result
            if lang != 'english':
                macro_avg_minimal_scores[0].append(f1)
                macro_avg_minimal_scores[1].append(precision)
                macro_avg_minimal_scores[2].append(recall)
            print('Minimal Answer & ' + lang + ' & ' +
                  get_latex_str(f1, precision, recall))

            print('*' * 20)
            print(lang)
            print('Language: %s (%d)' % (lang, len(per_lang_gold.get(lang, {}))))
            print('*' * 20)
            print('PASSAGE ANSWER R@P TABLE:')
            print_r_at_p_table(passage_answer_stats)
            print('*' * 20)
            print('MINIMAL ANSWER R@P TABLE:')
            print_r_at_p_table(minimal_answer_stats)

    print('Total # examples in gold: %d, # ex. in pred: %d (including english)' %
          (len(tydi_gold_dict), len(tydi_pred_dict)))

    f1_list, precision_list, recall_list = macro_avg_passage_scores
    print('*** Macro Over %d Languages, excluding English **' % len(f1_list))
    avg_passage_f1 = eval_utils.safe_average(f1_list)
commited
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
    macro_f1 = eval_utils.safe_divide(
        2 * macro_precision * macro_recall,
        macro_precision + macro_recall)

    return collections.OrderedDict({
        prefix + 'n': len(answer_stats),
        prefix + 'f1': macro_f1,
        prefix + 'precision': macro_precision,
        prefix + 'recall': macro_recall
    })


def compute_final_f1(passage_answer_stats, minimal_answer_stats):
    """Computes overall F1 given passage and minimal answers, ignoring scores.

    Note: this assumes that the answers have been thresholded based on scores.

    Arguments:
       passage_answer_stats: List of passage answer scores.
       minimal_answer_stats: List of minimal answer scores.


    Returns:
       Dictionary of name (string) -> score.
    """
    scores = compute_macro_f1(passage_answer_stats, prefix='passage-answer-')
    scores.update(compute_macro_f1(
        minimal_answer_stats, prefix='minimal-answer-'))
    return scores


def compute_pr_curves(answer_stats, targets=None):
    """Computes PR curve and returns R@P for specific targets.

    The values are computed as follows: find the (precision, recall) point
    with maximum recall and where precision > target.

    This is only relevant if you return the system scores in your predictions.
    You may find this useful when attempting to tune the threshold for your
    system on the dev set before requesting an evaluation on the test set
    via the leaderboard.

    Arguments:
      answer_stats: List of statistic tuples from the answer scores.
      targets (None): List of precision thresholds to target.

    Returns:
      List of table with rows: [target, r, p, score].
    """
    total_f1 = 0
    total_has_pred = 0
    total_has_gold = 0

    # Count the number of gold annotations.
    for has_gold, _, _, _ in answer_stats:
        total_has_gold += has_gold

    # Keep track of the point of maximum recall for each target.

    max_recall = [0 for _ in targets]
    max_precision = [0 for _ in targets]
    max_scores = [0.0 for _ in targets]

    # Only keep track of unique thresholds in this dictionary.
    scores_to_stats = collections.OrderedDict()

    # Loop through every possible threshold and compute precision + recall.
    for has_gold, has_pred, is_correct_or_f1, score in answer_stats:
        if isinstance(is_correct_or_f1, tuple):
            _, _, f1 = is_correct_or_f1
        else:
            f1 = is_correct_or_f1
        total_f1 += f1
        total_has_pred += has_pred

        precision = eval_utils.safe_divide(total_f1, total_has_pred)
        recall = eval_utils.safe_divide(total_f1, total_has_gold)

        # If there are any ties, this will be updated multiple times until the
        # ties are all counted.
        scores_to_stats[score] = [precision, recall]

    best_f1 = 0.0
    best_precision = 0.0
    best_recall = 0.0
    best_threshold = 0.0

    for threshold, (precision, recall) in scores_to_stats.items():
        # Match the thresholds to the find the closest precision above some target.
        for t, target in enumerate(targets):
            if precision >= target and recall > max_recall[t]:
                max_recall[t] = recall
                max_precision[t] = precision
                max_scores[t] = threshold

        # Compute optimal threshold.
        f1 = eval_utils.safe_divide(2 * precision * recall, precision + recall)
        if f1 > best_f1:
            best_f1 = f1
            best_precision = precision
            best_recall = recall
            best_threshold = threshold

    return ((best_f1, best_precision, best_recall, best_threshold),
            list(zip(targets, max_recall, max_precision, max_scores)))


def print_r_at_p_table(answer_stats):
    """Pretty prints the R@P table for default targets."""
    opt_result, pr_table = compute_pr_curves(
        answer_stats, targets=[0.5, 0.75, 0.9])
    f1, precision, recall, threshold = opt_result
    print('Optimal threshold: {:.5}'.format(threshold))
    print(' F1     /  P      /  R')
    print('{: >7.2%} / {: >7.2%} / {: >7.2%}'.format(f1, precision, recall))
    for target, recall, precision, row in pr_table:
        print('R@P={}: {:.2%} (actual p={:.2%}, score threshold={:.4})'.format(
            target, recall, precision, row))


def get_metrics_as_dict(gold_path, prediction_path, passage_non_null_threshold, minimal_non_null_threshold, verbose):
    """Library version of the end-to-end evaluation.

    Arguments:
      gold_path: Path to a single JSONL data. Could be gzipped or not.
      prediction_path: Path to the JSONL file of prediction data.

    Returns:
      metrics: A dictionary mapping string names to metric scores.
    """

    tydi_gold_dict = eval_utils.read_annotation(gold_path)
    tydi_pred_dict = eval_utils.read_prediction_jsonl(prediction_path)

    passage_answer_stats, minimal_answer_stats = score_answers(
        tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose)

    return get_metrics_with_answer_stats(
        passage_answer_stats, minimal_answer_stats)


def get_metrics_with_answer_stats(passage_answer_stats, minimal_answer_stats):
    """Generate metrics dict using passage and minimal answer stats."""

    def _get_metric_dict(answer_stats, prefix=''):
        """Compute all metrics for a set of answer statistics."""
        opt_result, pr_table = compute_pr_curves(
            answer_stats, targets=[0.5, 0.75, 0.9])
        f1, precision, recall, threshold = opt_result
        metrics = collections.OrderedDict({
            'best-threshold-f1': f1,
            'best-threshold-precision': precision,
            'best-threshold-recall': recall,
            'best-threshold': threshold,
        })
        for target, recall, precision, _ in pr_table:
            metrics['recall-at-precision>={:.2}'.format(target)] = recall
            metrics['precision-at-precision>={:.2}'.format(target)] = precision

        # Add prefix before returning.
        return dict([(prefix + k, v) for k, v in metrics.items()])

    metrics = _get_metric_dict(passage_answer_stats, 'passage-')
    metrics.update(_get_metric_dict(minimal_answer_stats, 'minimal-'))
    return metrics


def get_latex_str(f1, precision, recall):
    return '\\fpr' + '{' + ('%.1f' % (f1 * 100)) + '}{' + (
            '%.1f' % (precision * 100)) + '}{' + ('%.1f' % (recall * 100)) + '}'


# def load_gt_lookup_as_dict(gold_path, num_workers):
#     cache_path = os.path.join(os.path.dirname(gold_path), 'cache')
#
#     if os.path.exists(cache_path):
#         logging.info('Reading from cache: %s', format(cache_path))
#         tydi_gold_dict = pickle.load(open(cache_path, 'rb'))  # pytype: disable=wrong-arg-types
#     else:
#         tydi_gold_dict = eval_utils.read_annotation(gold_path, n_threads=num_workers)
#         logging.info('Caching gold data for future to: %s', format(cache_path))
#         cache_to_pickle_file(tydi_gold_dict, cache_path)  # pytype: disable=wrong-arg-types
#
#         logging.info('Read in gt lookup for %d examples into memory' % len(tydi_gold_dict))
#     return tydi_gold_dict


def pretty_print(tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold=2, minimal_non_null_threshold=2, verbose=False, skip_missing_example_ids=False):
    if any(map(not_, tydi_gold_dict.values())) or any(annotations[0].minimal_answer_span is None for annotations in tydi_gold_dict.values()):
        logging.warning("At least one example in gold dict has no annotations -- skipping pretty print evaluation")
        return

    gold_id_set = set(tydi_gold_dict.keys())
    pred_id_set = set(tydi_pred_dict.keys())
    sym_diff = gold_id_set.symmetric_difference(pred_id_set)

    if (not skip_missing_example_ids) and sym_diff:
        raise ValueError('ERROR: the example ids in gold annotations and example '
                         'ids in the prediction are not equal.')
    elif skip_missing_example_ids and sym_diff:
        logging.warning("Skipping {} example ids that are only in either gold or preds".format(len(sym_diff)))

    macro_avg_passage_scores = ([], [], [])
    macro_avg_minimal_scores = ([], [], [])

    language_list = [  # TODO: support arbitrary languages
        'english', 'arabic', 'bengali', 'finnish', 'indonesian', 'japanese',
        'swahili', 'korean', 'russian', 'telugu', 'thai'
    ]

    per_lang_gold = {}
    per_lang_pred = {}

    id_to_lang = {}

    for ex_id in gold_id_set:
        ex = tydi_gold_dict[ex_id]
        if ex[0].language in per_lang_gold:
            per_lang_gold[ex[0].language][ex_id] = ex
        else:
            per_lang_gold[ex[0].language] = {ex_id: ex}
        id_to_lang[ex_id] = ex[0].language
    for ex_id in gold_id_set:
        ex = tydi_pred_dict[ex_id]
        if id_to_lang[ex.example_id] in per_lang_pred:
            per_lang_pred[id_to_lang[ex.example_id]][ex.example_id] = ex
        else:
            per_lang_pred[id_to_lang[ex.example_id]] = {ex.example_id: ex}

    for lang in language_list:
        if lang in per_lang_pred:
            passage_answer_stats, minimal_answer_stats = score_answers(
                per_lang_gold.get(lang, {}), per_lang_pred[lang],
                passage_non_null_threshold, minimal_non_null_threshold, verbose)

            # Passage selection task
            opt_result, _ = compute_pr_curves(passage_answer_stats, targets=[0.5])
            f1, precision, recall, _ = opt_result
            if lang != 'english':
                macro_avg_passage_scores[0].append(f1)
                macro_avg_passage_scores[1].append(precision)
                macro_avg_passage_scores[2].append(recall)
            print('Passage & ' + lang + ' & ' + get_latex_str(f1, precision, recall))

            # Minimal answer span task
            opt_result, _ = compute_pr_curves(minimal_answer_stats, targets=[0.5])
            f1, precision, recall, _ = opt_result
            if lang != 'english':
                macro_avg_minimal_scores[0].append(f1)
                macro_avg_minimal_scores[1].append(precision)
                macro_avg_minimal_scores[2].append(recall)
            print('Minimal Answer & ' + lang + ' & ' +
                  get_latex_str(f1, precision, recall))

            print('*' * 20)
            print(lang)
            print('Language: %s (%d)' % (lang, len(per_lang_gold.get(lang, {}))))
            print('*' * 20)
            print('PASSAGE ANSWER R@P TABLE:')
            print_r_at_p_table(passage_answer_stats)
            print('*' * 20)
            print('MINIMAL ANSWER R@P TABLE:')
            print_r_at_p_table(minimal_answer_stats)

    print('Total # examples in gold: %d, # ex. in pred: %d (including english)' %
          (len(tydi_gold_dict), len(tydi_pred_dict)))

    f1_list, precision_list, recall_list = macro_avg_passage_scores
    print('*** Macro Over %d Languages, excluding English **' % len(f1_list))
    avg_passage_f1 = eval_utils.safe_average(f1_list)
    avg_passage_recall = eval_utils.safe_average(recall_list)
commited
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
commited
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
    macro_f1 = eval_utils.safe_divide(
        2 * macro_precision * macro_recall,
        macro_precision + macro_recall)

    return collections.OrderedDict({
        prefix + 'n': len(answer_stats),
        prefix + 'f1': macro_f1,
        prefix + 'precision': macro_precision,
        prefix + 'recall': macro_recall
    })


def compute_final_f1(passage_answer_stats, minimal_answer_stats):
    """Computes overall F1 given passage and minimal answers, ignoring scores.

    Note: this assumes that the answers have been thresholded based on scores.

    Arguments:
       passage_answer_stats: List of passage answer scores.
       minimal_answer_stats: List of minimal answer scores.


    Returns:
       Dictionary of name (string) -> score.
    """
    scores = compute_macro_f1(passage_answer_stats, prefix='passage-answer-')
    scores.update(compute_macro_f1(
        minimal_answer_stats, prefix='minimal-answer-'))
    return scores


def compute_pr_curves(answer_stats, targets=None):
    """Computes PR curve and returns R@P for specific targets.

    The values are computed as follows: find the (precision, recall) point
    with maximum recall and where precision > target.

    This is only relevant if you return the system scores in your predictions.
    You may find this useful when attempting to tune the threshold for your
    system on the dev set before requesting an evaluation on the test set
    via the leaderboard.

    Arguments:
      answer_stats: List of statistic tuples from the answer scores.
      targets (None): List of precision thresholds to target.

    Returns:
      List of table with rows: [target, r, p, score].
    """
    total_f1 = 0
    total_has_pred = 0
    total_has_gold = 0

    # Count the number of gold annotations.
    for has_gold, _, _, _ in answer_stats:
        total_has_gold += has_gold

    # Keep track of the point of maximum recall for each target.

    max_recall = [0 for _ in targets]
    max_precision = [0 for _ in targets]
    max_scores = [0.0 for _ in targets]

    # Only keep track of unique thresholds in this dictionary.
    scores_to_stats = collections.OrderedDict()

    # Loop through every possible threshold and compute precision + recall.
    for has_gold, has_pred, is_correct_or_f1, score in answer_stats:
        if isinstance(is_correct_or_f1, tuple):
            _, _, f1 = is_correct_or_f1
        else:
            f1 = is_correct_or_f1
        total_f1 += f1
        total_has_pred += has_pred

        precision = eval_utils.safe_divide(total_f1, total_has_pred)
        recall = eval_utils.safe_divide(total_f1, total_has_gold)

        # If there are any ties, this will be updated multiple times until the
        # ties are all counted.
        scores_to_stats[score] = [precision, recall]

    best_f1 = 0.0
    best_precision = 0.0
    best_recall = 0.0
    best_threshold = 0.0

    for threshold, (precision, recall) in scores_to_stats.items():
        # Match the thresholds to the find the closest precision above some target.
        for t, target in enumerate(targets):
            if precision >= target and recall > max_recall[t]:
                max_recall[t] = recall
                max_precision[t] = precision
                max_scores[t] = threshold

        # Compute optimal threshold.
        f1 = eval_utils.safe_divide(2 * precision * recall, precision + recall)
        if f1 > best_f1:
            best_f1 = f1
            best_precision = precision
            best_recall = recall
            best_threshold = threshold

    return ((best_f1, best_precision, best_recall, best_threshold),
            list(zip(targets, max_recall, max_precision, max_scores)))


def print_r_at_p_table(answer_stats):
    """Pretty prints the R@P table for default targets."""
    opt_result, pr_table = compute_pr_curves(
        answer_stats, targets=[0.5, 0.75, 0.9])
    f1, precision, recall, threshold = opt_result
    print('Optimal threshold: {:.5}'.format(threshold))
    print(' F1     /  P      /  R')
    print('{: >7.2%} / {: >7.2%} / {: >7.2%}'.format(f1, precision, recall))
    for target, recall, precision, row in pr_table:
        print('R@P={}: {:.2%} (actual p={:.2%}, score threshold={:.4})'.format(
            target, recall, precision, row))


def get_metrics_as_dict(gold_path, prediction_path, passage_non_null_threshold, minimal_non_null_threshold, verbose):
    """Library version of the end-to-end evaluation.

    Arguments:
      gold_path: Path to a single JSONL data. Could be gzipped or not.
      prediction_path: Path to the JSONL file of prediction data.

    Returns:
      metrics: A dictionary mapping string names to metric scores.
    """

    tydi_gold_dict = eval_utils.read_annotation(gold_path)
    tydi_pred_dict = eval_utils.read_prediction_jsonl(prediction_path)

    passage_answer_stats, minimal_answer_stats = score_answers(
        tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose)

    return get_metrics_with_answer_stats(
        passage_answer_stats, minimal_answer_stats)


def get_metrics_with_answer_stats(passage_answer_stats, minimal_answer_stats):
    """Generate metrics dict using passage and minimal answer stats."""

    def _get_metric_dict(answer_stats, prefix=''):
        """Compute all metrics for a set of answer statistics."""
        opt_result, pr_table = compute_pr_curves(
            answer_stats, targets=[0.5, 0.75, 0.9])
        f1, precision, recall, threshold = opt_result
        metrics = collections.OrderedDict({
            'best-threshold-f1': f1,
            'best-threshold-precision': precision,
            'best-threshold-recall': recall,
            'best-threshold': threshold,
        })
        for target, recall, precision, _ in pr_table:
            metrics['recall-at-precision>={:.2}'.format(target)] = recall
            metrics['precision-at-precision>={:.2}'.format(target)] = precision

        # Add prefix before returning.
        return dict([(prefix + k, v) for k, v in metrics.items()])

    metrics = _get_metric_dict(passage_answer_stats, 'passage-')
    metrics.update(_get_metric_dict(minimal_answer_stats, 'minimal-'))
    return metrics


def get_latex_str(f1, precision, recall):
    return '\\fpr' + '{' + ('%.1f' % (f1 * 100)) + '}{' + (
            '%.1f' % (precision * 100)) + '}{' + ('%.1f' % (recall * 100)) + '}'


# def load_gt_lookup_as_dict(gold_path, num_workers):
#     cache_path = os.path.join(os.path.dirname(gold_path), 'cache')
#
#     if os.path.exists(cache_path):
#         logging.info('Reading from cache: %s', format(cache_path))
#         tydi_gold_dict = pickle.load(open(cache_path, 'rb'))  # pytype: disable=wrong-arg-types
#     else:
#         tydi_gold_dict = eval_utils.read_annotation(gold_path, n_threads=num_workers)
#         logging.info('Caching gold data for future to: %s', format(cache_path))
#         cache_to_pickle_file(tydi_gold_dict, cache_path)  # pytype: disable=wrong-arg-types
#
#         logging.info('Read in gt lookup for %d examples into memory' % len(tydi_gold_dict))
#     return tydi_gold_dict


def pretty_print(tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold=2, minimal_non_null_threshold=2, verbose=False, skip_missing_example_ids=False):
    if any(map(not_, tydi_gold_dict.values())) or any(annotations[0].minimal_answer_span is None for annotations in tydi_gold_dict.values()):
        logging.warning("At least one example in gold dict has no annotations -- skipping pretty print evaluation")
        return

    gold_id_set = set(tydi_gold_dict.keys())
    pred_id_set = set(tydi_pred_dict.keys())
    sym_diff = gold_id_set.symmetric_difference(pred_id_set)

    if (not skip_missing_example_ids) and sym_diff:
        raise ValueError('ERROR: the example ids in gold annotations and example '
                         'ids in the prediction are not equal.')
    elif skip_missing_example_ids and sym_diff:
        logging.warning("Skipping {} example ids that are only in either gold or preds".format(len(sym_diff)))

    macro_avg_passage_scores = ([], [], [])
    macro_avg_minimal_scores = ([], [], [])

    language_list = [  # TODO: support arbitrary languages
        'english', 'arabic', 'bengali', 'finnish', 'indonesian', 'japanese',
        'swahili', 'korean', 'russian', 'telugu', 'thai'
    ]

    per_lang_gold = {}
    per_lang_pred = {}

    id_to_lang = {}

    for ex_id in gold_id_set:
        ex = tydi_gold_dict[ex_id]
        if ex[0].language in per_lang_gold:
            per_lang_gold[ex[0].language][ex_id] = ex
        else:
            per_lang_gold[ex[0].language] = {ex_id: ex}
        id_to_lang[ex_id] = ex[0].language
    for ex_id in gold_id_set:
        ex = tydi_pred_dict[ex_id]
        if id_to_lang[ex.example_id] in per_lang_pred:
            per_lang_pred[id_to_lang[ex.example_id]][ex.example_id] = ex
        else:
            per_lang_pred[id_to_lang[ex.example_id]] = {ex.example_id: ex}

    for lang in language_list:
        if lang in per_lang_pred:
            passage_answer_stats, minimal_answer_stats = score_answers(
                per_lang_gold.get(lang, {}), per_lang_pred[lang],
                passage_non_null_threshold, minimal_non_null_threshold, verbose)

            # Passage selection task
            opt_result, _ = compute_pr_curves(passage_answer_stats, targets=[0.5])
            f1, precision, recall, _ = opt_result
            if lang != 'english':
                macro_avg_passage_scores[0].append(f1)
                macro_avg_passage_scores[1].append(precision)
                macro_avg_passage_scores[2].append(recall)
            print('Passage & ' + lang + ' & ' + get_latex_str(f1, precision, recall))

            # Minimal answer span task
            opt_result, _ = compute_pr_curves(minimal_answer_stats, targets=[0.5])
            f1, precision, recall, _ = opt_result
            if lang != 'english':
                macro_avg_minimal_scores[0].append(f1)
                macro_avg_minimal_scores[1].append(precision)
                macro_avg_minimal_scores[2].append(recall)
            print('Minimal Answer & ' + lang + ' & ' +
                  get_latex_str(f1, precision, recall))

            print('*' * 20)
            print(lang)
            print('Language: %s (%d)' % (lang, len(per_lang_gold.get(lang, {}))))
            print('*' * 20)
            print('PASSAGE ANSWER R@P TABLE:')
            print_r_at_p_table(passage_answer_stats)
            print('*' * 20)
            print('MINIMAL ANSWER R@P TABLE:')
            print_r_at_p_table(minimal_answer_stats)

    print('Total # examples in gold: %d, # ex. in pred: %d (including english)' %
          (len(tydi_gold_dict), len(tydi_pred_dict)))

    f1_list, precision_list, recall_list = macro_avg_passage_scores
    print('*** Macro Over %d Languages, excluding English **' % len(f1_list))
    avg_passage_f1 = eval_utils.safe_average(f1_list)
    avg_passage_recall = eval_utils.safe_average(recall_list)
    avg_passage_precision = eval_utils.safe_average(precision_list)
    print('Passage F1:%.3f P:%.3f R:%3f' %
          (avg_passage_f1, avg_passage_precision, avg_passage_recall))
    print(get_latex_str(
        avg_passage_f1, avg_passage_precision, avg_passage_recall))

    f1_list, precision_list, recall_list = macro_avg_minimal_scores

commited
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
commited
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
    macro_f1 = eval_utils.safe_divide(
        2 * macro_precision * macro_recall,
        macro_precision + macro_recall)

    return collections.OrderedDict({
        prefix + 'n': len(answer_stats),
        prefix + 'f1': macro_f1,
        prefix + 'precision': macro_precision,
        prefix + 'recall': macro_recall
    })


def compute_final_f1(passage_answer_stats, minimal_answer_stats):
    """Computes overall F1 given passage and minimal answers, ignoring scores.

    Note: this assumes that the answers have been thresholded based on scores.

    Arguments:
       passage_answer_stats: List of passage answer scores.
       minimal_answer_stats: List of minimal answer scores.


    Returns:
       Dictionary of name (string) -> score.
    """
    scores = compute_macro_f1(passage_answer_stats, prefix='passage-answer-')
    scores.update(compute_macro_f1(
        minimal_answer_stats, prefix='minimal-answer-'))
    return scores


def compute_pr_curves(answer_stats, targets=None):
    """Computes PR curve and returns R@P for specific targets.

    The values are computed as follows: find the (precision, recall) point
    with maximum recall and where precision > target.

    This is only relevant if you return the system scores in your predictions.
    You may find this useful when attempting to tune the threshold for your
    system on the dev set before requesting an evaluation on the test set
    via the leaderboard.

    Arguments:
      answer_stats: List of statistic tuples from the answer scores.
      targets (None): List of precision thresholds to target.

    Returns:
      List of table with rows: [target, r, p, score].
    """
    total_f1 = 0
    total_has_pred = 0
    total_has_gold = 0

    # Count the number of gold annotations.
    for has_gold, _, _, _ in answer_stats:
        total_has_gold += has_gold

    # Keep track of the point of maximum recall for each target.

    max_recall = [0 for _ in targets]
    max_precision = [0 for _ in targets]
    max_scores = [0.0 for _ in targets]

    # Only keep track of unique thresholds in this dictionary.
    scores_to_stats = collections.OrderedDict()

    # Loop through every possible threshold and compute precision + recall.
    for has_gold, has_pred, is_correct_or_f1, score in answer_stats:
        if isinstance(is_correct_or_f1, tuple):
            _, _, f1 = is_correct_or_f1
        else:
            f1 = is_correct_or_f1
        total_f1 += f1
        total_has_pred += has_pred

        precision = eval_utils.safe_divide(total_f1, total_has_pred)
        recall = eval_utils.safe_divide(total_f1, total_has_gold)

        # If there are any ties, this will be updated multiple times until the
        # ties are all counted.
        scores_to_stats[score] = [precision, recall]

    best_f1 = 0.0
    best_precision = 0.0
    best_recall = 0.0
    best_threshold = 0.0

    for threshold, (precision, recall) in scores_to_stats.items():
        # Match the thresholds to the find the closest precision above some target.
        for t, target in enumerate(targets):
            if precision >= target and recall > max_recall[t]:
                max_recall[t] = recall
                max_precision[t] = precision
                max_scores[t] = threshold

        # Compute optimal threshold.
        f1 = eval_utils.safe_divide(2 * precision * recall, precision + recall)
        if f1 > best_f1:
            best_f1 = f1
            best_precision = precision
            best_recall = recall
            best_threshold = threshold

    return ((best_f1, best_precision, best_recall, best_threshold),
            list(zip(targets, max_recall, max_precision, max_scores)))


def print_r_at_p_table(answer_stats):
    """Pretty prints the R@P table for default targets."""
    opt_result, pr_table = compute_pr_curves(
        answer_stats, targets=[0.5, 0.75, 0.9])
    f1, precision, recall, threshold = opt_result
    print('Optimal threshold: {:.5}'.format(threshold))
    print(' F1     /  P      /  R')
    print('{: >7.2%} / {: >7.2%} / {: >7.2%}'.format(f1, precision, recall))
    for target, recall, precision, row in pr_table:
        print('R@P={}: {:.2%} (actual p={:.2%}, score threshold={:.4})'.format(
            target, recall, precision, row))


def get_metrics_as_dict(gold_path, prediction_path, passage_non_null_threshold, minimal_non_null_threshold, verbose):
    """Library version of the end-to-end evaluation.

    Arguments:
      gold_path: Path to a single JSONL data. Could be gzipped or not.
      prediction_path: Path to the JSONL file of prediction data.

    Returns:
      metrics: A dictionary mapping string names to metric scores.
    """

    tydi_gold_dict = eval_utils.read_annotation(gold_path)
    tydi_pred_dict = eval_utils.read_prediction_jsonl(prediction_path)

    passage_answer_stats, minimal_answer_stats = score_answers(
        tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose)

    return get_metrics_with_answer_stats(
        passage_answer_stats, minimal_answer_stats)


def get_metrics_with_answer_stats(passage_answer_stats, minimal_answer_stats):
    """Generate metrics dict using passage and minimal answer stats."""

    def _get_metric_dict(answer_stats, prefix=''):
        """Compute all metrics for a set of answer statistics."""
        opt_result, pr_table = compute_pr_curves(
            answer_stats, targets=[0.5, 0.75, 0.9])
        f1, precision, recall, threshold = opt_result
        metrics = collections.OrderedDict({
            'best-threshold-f1': f1,
            'best-threshold-precision': precision,
            'best-threshold-recall': recall,
            'best-threshold': threshold,
        })
        for target, recall, precision, _ in pr_table:
            metrics['recall-at-precision>={:.2}'.format(target)] = recall
            metrics['precision-at-precision>={:.2}'.format(target)] = precision

        # Add prefix before returning.
        return dict([(prefix + k, v) for k, v in metrics.items()])

    metrics = _get_metric_dict(passage_answer_stats, 'passage-')
    metrics.update(_get_metric_dict(minimal_answer_stats, 'minimal-'))
    return metrics


def get_latex_str(f1, precision, recall):
    return '\\fpr' + '{' + ('%.1f' % (f1 * 100)) + '}{' + (
            '%.1f' % (precision * 100)) + '}{' + ('%.1f' % (recall * 100)) + '}'


# def load_gt_lookup_as_dict(gold_path, num_workers):
#     cache_path = os.path.join(os.path.dirname(gold_path), 'cache')
#
#     if os.path.exists(cache_path):
#         logging.info('Reading from cache: %s', format(cache_path))
#         tydi_gold_dict = pickle.load(open(cache_path, 'rb'))  # pytype: disable=wrong-arg-types
#     else:
#         tydi_gold_dict = eval_utils.read_annotation(gold_path, n_threads=num_workers)
#         logging.info('Caching gold data for future to: %s', format(cache_path))
#         cache_to_pickle_file(tydi_gold_dict, cache_path)  # pytype: disable=wrong-arg-types
#
#         logging.info('Read in gt lookup for %d examples into memory' % len(tydi_gold_dict))
#     return tydi_gold_dict


def pretty_print(tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold=2, minimal_non_null_threshold=2, verbose=False, skip_missing_example_ids=False):
    if any(map(not_, tydi_gold_dict.values())) or any(annotations[0].minimal_answer_span is None for annotations in tydi_gold_dict.values()):
        logging.warning("At least one example in gold dict has no annotations -- skipping pretty print evaluation")
        return

    gold_id_set = set(tydi_gold_dict.keys())
    pred_id_set = set(tydi_pred_dict.keys())
    sym_diff = gold_id_set.symmetric_difference(pred_id_set)

    if (not skip_missing_example_ids) and sym_diff:
        raise ValueError('ERROR: the example ids in gold annotations and example '
                         'ids in the prediction are not equal.')
    elif skip_missing_example_ids and sym_diff:
        logging.warning("Skipping {} example ids that are only in either gold or preds".format(len(sym_diff)))

    macro_avg_passage_scores = ([], [], [])
    macro_avg_minimal_scores = ([], [], [])

    language_list = [  # TODO: support arbitrary languages
        'english', 'arabic', 'bengali', 'finnish', 'indonesian', 'japanese',
        'swahili', 'korean', 'russian', 'telugu', 'thai'
    ]

    per_lang_gold = {}
    per_lang_pred = {}

    id_to_lang = {}

    for ex_id in gold_id_set:
        ex = tydi_gold_dict[ex_id]
        if ex[0].language in per_lang_gold:
            per_lang_gold[ex[0].language][ex_id] = ex
        else:
            per_lang_gold[ex[0].language] = {ex_id: ex}
        id_to_lang[ex_id] = ex[0].language
    for ex_id in gold_id_set:
        ex = tydi_pred_dict[ex_id]
        if id_to_lang[ex.example_id] in per_lang_pred:
            per_lang_pred[id_to_lang[ex.example_id]][ex.example_id] = ex
        else:
            per_lang_pred[id_to_lang[ex.example_id]] = {ex.example_id: ex}

    for lang in language_list:
        if lang in per_lang_pred:
            passage_answer_stats, minimal_answer_stats = score_answers(
                per_lang_gold.get(lang, {}), per_lang_pred[lang],
                passage_non_null_threshold, minimal_non_null_threshold, verbose)

            # Passage selection task
            opt_result, _ = compute_pr_curves(passage_answer_stats, targets=[0.5])
            f1, precision, recall, _ = opt_result
            if lang != 'english':
                macro_avg_passage_scores[0].append(f1)
                macro_avg_passage_scores[1].append(precision)
                macro_avg_passage_scores[2].append(recall)
            print('Passage & ' + lang + ' & ' + get_latex_str(f1, precision, recall))

            # Minimal answer span task
            opt_result, _ = compute_pr_curves(minimal_answer_stats, targets=[0.5])
            f1, precision, recall, _ = opt_result
            if lang != 'english':
                macro_avg_minimal_scores[0].append(f1)
                macro_avg_minimal_scores[1].append(precision)
                macro_avg_minimal_scores[2].append(recall)
            print('Minimal Answer & ' + lang + ' & ' +
                  get_latex_str(f1, precision, recall))

            print('*' * 20)
            print(lang)
            print('Language: %s (%d)' % (lang, len(per_lang_gold.get(lang, {}))))
            print('*' * 20)
            print('PASSAGE ANSWER R@P TABLE:')
            print_r_at_p_table(passage_answer_stats)
            print('*' * 20)
            print('MINIMAL ANSWER R@P TABLE:')
            print_r_at_p_table(minimal_answer_stats)

    print('Total # examples in gold: %d, # ex. in pred: %d (including english)' %
          (len(tydi_gold_dict), len(tydi_pred_dict)))

    f1_list, precision_list, recall_list = macro_avg_passage_scores
    print('*** Macro Over %d Languages, excluding English **' % len(f1_list))
    avg_passage_f1 = eval_utils.safe_average(f1_list)
    avg_passage_recall = eval_utils.safe_average(recall_list)
    avg_passage_precision = eval_utils.safe_average(precision_list)
    print('Passage F1:%.3f P:%.3f R:%3f' %
          (avg_passage_f1, avg_passage_precision, avg_passage_recall))
    print(get_latex_str(
        avg_passage_f1, avg_passage_precision, avg_passage_recall))

    f1_list, precision_list, recall_list = macro_avg_minimal_scores

    avg_minimal_f1 = eval_utils.safe_average(f1_list)
commited
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
    macro_f1 = eval_utils.safe_divide(
        2 * macro_precision * macro_recall,
        macro_precision + macro_recall)

    return collections.OrderedDict({
        prefix + 'n': len(answer_stats),
        prefix + 'f1': macro_f1,
        prefix + 'precision': macro_precision,
        prefix + 'recall': macro_recall
    })


def compute_final_f1(passage_answer_stats, minimal_answer_stats):
    """Computes overall F1 given passage and minimal answers, ignoring scores.

    Note: this assumes that the answers have been thresholded based on scores.

    Arguments:
       passage_answer_stats: List of passage answer scores.
       minimal_answer_stats: List of minimal answer scores.


    Returns:
       Dictionary of name (string) -> score.
    """
    scores = compute_macro_f1(passage_answer_stats, prefix='passage-answer-')
    scores.update(compute_macro_f1(
        minimal_answer_stats, prefix='minimal-answer-'))
    return scores


def compute_pr_curves(answer_stats, targets=None):
    """Computes PR curve and returns R@P for specific targets.

    The values are computed as follows: find the (precision, recall) point
    with maximum recall and where precision > target.

    This is only relevant if you return the system scores in your predictions.
    You may find this useful when attempting to tune the threshold for your
    system on the dev set before requesting an evaluation on the test set
    via the leaderboard.

    Arguments:
      answer_stats: List of statistic tuples from the answer scores.
      targets (None): List of precision thresholds to target.

    Returns:
      List of table with rows: [target, r, p, score].
    """
    total_f1 = 0
    total_has_pred = 0
    total_has_gold = 0

    # Count the number of gold annotations.
    for has_gold, _, _, _ in answer_stats:
        total_has_gold += has_gold

    # Keep track of the point of maximum recall for each target.

    max_recall = [0 for _ in targets]
    max_precision = [0 for _ in targets]
    max_scores = [0.0 for _ in targets]

    # Only keep track of unique thresholds in this dictionary.
    scores_to_stats = collections.OrderedDict()

    # Loop through every possible threshold and compute precision + recall.
    for has_gold, has_pred, is_correct_or_f1, score in answer_stats:
        if isinstance(is_correct_or_f1, tuple):
            _, _, f1 = is_correct_or_f1
        else:
            f1 = is_correct_or_f1
        total_f1 += f1
        total_has_pred += has_pred

        precision = eval_utils.safe_divide(total_f1, total_has_pred)
        recall = eval_utils.safe_divide(total_f1, total_has_gold)

        # If there are any ties, this will be updated multiple times until the
        # ties are all counted.
        scores_to_stats[score] = [precision, recall]

    best_f1 = 0.0
    best_precision = 0.0
    best_recall = 0.0
    best_threshold = 0.0

    for threshold, (precision, recall) in scores_to_stats.items():
        # Match the thresholds to the find the closest precision above some target.
        for t, target in enumerate(targets):
            if precision >= target and recall > max_recall[t]:
                max_recall[t] = recall
                max_precision[t] = precision
                max_scores[t] = threshold

        # Compute optimal threshold.
        f1 = eval_utils.safe_divide(2 * precision * recall, precision + recall)
        if f1 > best_f1:
            best_f1 = f1
            best_precision = precision
            best_recall = recall
            best_threshold = threshold

    return ((best_f1, best_precision, best_recall, best_threshold),
            list(zip(targets, max_recall, max_precision, max_scores)))


def print_r_at_p_table(answer_stats):
    """Pretty prints the R@P table for default targets."""
    opt_result, pr_table = compute_pr_curves(
        answer_stats, targets=[0.5, 0.75, 0.9])
    f1, precision, recall, threshold = opt_result
    print('Optimal threshold: {:.5}'.format(threshold))
    print(' F1     /  P      /  R')
    print('{: >7.2%} / {: >7.2%} / {: >7.2%}'.format(f1, precision, recall))
    for target, recall, precision, row in pr_table:
        print('R@P={}: {:.2%} (actual p={:.2%}, score threshold={:.4})'.format(
            target, recall, precision, row))


def get_metrics_as_dict(gold_path, prediction_path, passage_non_null_threshold, minimal_non_null_threshold, verbose):
    """Library version of the end-to-end evaluation.

    Arguments:
      gold_path: Path to a single JSONL data. Could be gzipped or not.
      prediction_path: Path to the JSONL file of prediction data.

    Returns:
      metrics: A dictionary mapping string names to metric scores.
    """

    tydi_gold_dict = eval_utils.read_annotation(gold_path)
    tydi_pred_dict = eval_utils.read_prediction_jsonl(prediction_path)

    passage_answer_stats, minimal_answer_stats = score_answers(
        tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose)

    return get_metrics_with_answer_stats(
        passage_answer_stats, minimal_answer_stats)


def get_metrics_with_answer_stats(passage_answer_stats, minimal_answer_stats):
    """Generate metrics dict using passage and minimal answer stats."""

    def _get_metric_dict(answer_stats, prefix=''):
        """Compute all metrics for a set of answer statistics."""
        opt_result, pr_table = compute_pr_curves(
            answer_stats, targets=[0.5, 0.75, 0.9])
        f1, precision, recall, threshold = opt_result
        metrics = collections.OrderedDict({
            'best-threshold-f1': f1,
            'best-threshold-precision': precision,
            'best-threshold-recall': recall,
            'best-threshold': threshold,
        })
        for target, recall, precision, _ in pr_table:
            metrics['recall-at-precision>={:.2}'.format(target)] = recall
            metrics['precision-at-precision>={:.2}'.format(target)] = precision

        # Add prefix before returning.
        return dict([(prefix + k, v) for k, v in metrics.items()])

    metrics = _get_metric_dict(passage_answer_stats, 'passage-')
    metrics.update(_get_metric_dict(minimal_answer_stats, 'minimal-'))
    return metrics


def get_latex_str(f1, precision, recall):
    return '\\fpr' + '{' + ('%.1f' % (f1 * 100)) + '}{' + (
            '%.1f' % (precision * 100)) + '}{' + ('%.1f' % (recall * 100)) + '}'


# def load_gt_lookup_as_dict(gold_path, num_workers):
#     cache_path = os.path.join(os.path.dirname(gold_path), 'cache')
#
#     if os.path.exists(cache_path):
#         logging.info('Reading from cache: %s', format(cache_path))
#         tydi_gold_dict = pickle.load(open(cache_path, 'rb'))  # pytype: disable=wrong-arg-types
#     else:
#         tydi_gold_dict = eval_utils.read_annotation(gold_path, n_threads=num_workers)
#         logging.info('Caching gold data for future to: %s', format(cache_path))
#         cache_to_pickle_file(tydi_gold_dict, cache_path)  # pytype: disable=wrong-arg-types
#
#         logging.info('Read in gt lookup for %d examples into memory' % len(tydi_gold_dict))
#     return tydi_gold_dict


def pretty_print(tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold=2, minimal_non_null_threshold=2, verbose=False, skip_missing_example_ids=False):
    if any(map(not_, tydi_gold_dict.values())) or any(annotations[0].minimal_answer_span is None for annotations in tydi_gold_dict.values()):
        logging.warning("At least one example in gold dict has no annotations -- skipping pretty print evaluation")
        return

    gold_id_set = set(tydi_gold_dict.keys())
    pred_id_set = set(tydi_pred_dict.keys())
    sym_diff = gold_id_set.symmetric_difference(pred_id_set)

    if (not skip_missing_example_ids) and sym_diff:
        raise ValueError('ERROR: the example ids in gold annotations and example '
                         'ids in the prediction are not equal.')
    elif skip_missing_example_ids and sym_diff:
        logging.warning("Skipping {} example ids that are only in either gold or preds".format(len(sym_diff)))

    macro_avg_passage_scores = ([], [], [])
    macro_avg_minimal_scores = ([], [], [])

    language_list = [  # TODO: support arbitrary languages
        'english', 'arabic', 'bengali', 'finnish', 'indonesian', 'japanese',
        'swahili', 'korean', 'russian', 'telugu', 'thai'
    ]

    per_lang_gold = {}
    per_lang_pred = {}

    id_to_lang = {}

    for ex_id in gold_id_set:
        ex = tydi_gold_dict[ex_id]
        if ex[0].language in per_lang_gold:
            per_lang_gold[ex[0].language][ex_id] = ex
        else:
            per_lang_gold[ex[0].language] = {ex_id: ex}
        id_to_lang[ex_id] = ex[0].language
    for ex_id in gold_id_set:
        ex = tydi_pred_dict[ex_id]
        if id_to_lang[ex.example_id] in per_lang_pred:
            per_lang_pred[id_to_lang[ex.example_id]][ex.example_id] = ex
        else:
            per_lang_pred[id_to_lang[ex.example_id]] = {ex.example_id: ex}

    for lang in language_list:
        if lang in per_lang_pred:
            passage_answer_stats, minimal_answer_stats = score_answers(
                per_lang_gold.get(lang, {}), per_lang_pred[lang],
                passage_non_null_threshold, minimal_non_null_threshold, verbose)

            # Passage selection task
            opt_result, _ = compute_pr_curves(passage_answer_stats, targets=[0.5])
            f1, precision, recall, _ = opt_result
            if lang != 'english':
                macro_avg_passage_scores[0].append(f1)
                macro_avg_passage_scores[1].append(precision)
                macro_avg_passage_scores[2].append(recall)
            print('Passage & ' + lang + ' & ' + get_latex_str(f1, precision, recall))

            # Minimal answer span task
            opt_result, _ = compute_pr_curves(minimal_answer_stats, targets=[0.5])
            f1, precision, recall, _ = opt_result
            if lang != 'english':
                macro_avg_minimal_scores[0].append(f1)
                macro_avg_minimal_scores[1].append(precision)
                macro_avg_minimal_scores[2].append(recall)
            print('Minimal Answer & ' + lang + ' & ' +
                  get_latex_str(f1, precision, recall))

            print('*' * 20)
            print(lang)
            print('Language: %s (%d)' % (lang, len(per_lang_gold.get(lang, {}))))
            print('*' * 20)
            print('PASSAGE ANSWER R@P TABLE:')
            print_r_at_p_table(passage_answer_stats)
            print('*' * 20)
            print('MINIMAL ANSWER R@P TABLE:')
            print_r_at_p_table(minimal_answer_stats)

    print('Total # examples in gold: %d, # ex. in pred: %d (including english)' %
          (len(tydi_gold_dict), len(tydi_pred_dict)))

    f1_list, precision_list, recall_list = macro_avg_passage_scores
    print('*** Macro Over %d Languages, excluding English **' % len(f1_list))
    avg_passage_f1 = eval_utils.safe_average(f1_list)
    avg_passage_recall = eval_utils.safe_average(recall_list)
    avg_passage_precision = eval_utils.safe_average(precision_list)
    print('Passage F1:%.3f P:%.3f R:%3f' %
          (avg_passage_f1, avg_passage_precision, avg_passage_recall))
    print(get_latex_str(
        avg_passage_f1, avg_passage_precision, avg_passage_recall))

    f1_list, precision_list, recall_list = macro_avg_minimal_scores

    avg_minimal_f1 = eval_utils.safe_average(f1_list)
    avg_minimal_recall = eval_utils.safe_average(recall_list)
commited
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
commited
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
commited
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
    macro_f1 = eval_utils.safe_divide(
        2 * macro_precision * macro_recall,
        macro_precision + macro_recall)

    return collections.OrderedDict({
        prefix + 'n': len(answer_stats),
        prefix + 'f1': macro_f1,
        prefix + 'precision': macro_precision,
        prefix + 'recall': macro_recall
    })


def compute_final_f1(passage_answer_stats, minimal_answer_stats):
    """Computes overall F1 given passage and minimal answers, ignoring scores.

    Note: this assumes that the answers have been thresholded based on scores.

    Arguments:
       passage_answer_stats: List of passage answer scores.
       minimal_answer_stats: List of minimal answer scores.


    Returns:
       Dictionary of name (string) -> score.
    """
    scores = compute_macro_f1(passage_answer_stats, prefix='passage-answer-')
    scores.update(compute_macro_f1(
        minimal_answer_stats, prefix='minimal-answer-'))
    return scores


def compute_pr_curves(answer_stats, targets=None):
    """Computes PR curve and returns R@P for specific targets.

    The values are computed as follows: find the (precision, recall) point
    with maximum recall and where precision > target.

    This is only relevant if you return the system scores in your predictions.
    You may find this useful when attempting to tune the threshold for your
    system on the dev set before requesting an evaluation on the test set
    via the leaderboard.

    Arguments:
      answer_stats: List of statistic tuples from the answer scores.
      targets (None): List of precision thresholds to target.

    Returns:
      List of table with rows: [target, r, p, score].
    """
    total_f1 = 0
    total_has_pred = 0
    total_has_gold = 0

    # Count the number of gold annotations.
    for has_gold, _, _, _ in answer_stats:
        total_has_gold += has_gold

    # Keep track of the point of maximum recall for each target.

    max_recall = [0 for _ in targets]
    max_precision = [0 for _ in targets]
    max_scores = [0.0 for _ in targets]

    # Only keep track of unique thresholds in this dictionary.
    scores_to_stats = collections.OrderedDict()

    # Loop through every possible threshold and compute precision + recall.
    for has_gold, has_pred, is_correct_or_f1, score in answer_stats:
        if isinstance(is_correct_or_f1, tuple):
            _, _, f1 = is_correct_or_f1
        else:
            f1 = is_correct_or_f1
        total_f1 += f1
        total_has_pred += has_pred

commited
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
    macro_f1 = eval_utils.safe_divide(
        2 * macro_precision * macro_recall,
        macro_precision + macro_recall)

    return collections.OrderedDict({
        prefix + 'n': len(answer_stats),
        prefix + 'f1': macro_f1,
        prefix + 'precision': macro_precision,
        prefix + 'recall': macro_recall
    })


def compute_final_f1(passage_answer_stats, minimal_answer_stats):
    """Computes overall F1 given passage and minimal answers, ignoring scores.

non_informative
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
    macro_f1 = eval_utils.safe_divide(
        2 * macro_precision * macro_recall,
        macro_precision + macro_recall)

    return collections.OrderedDict({
        prefix + 'n': len(answer_stats),
        prefix + 'f1': macro_f1,
        prefix + 'precision': macro_precision,
        prefix + 'recall': macro_recall
    })


def compute_final_f1(passage_answer_stats, minimal_answer_stats):
    """Computes overall F1 given passage and minimal answers, ignoring scores.

    Note: this assumes that the answers have been thresholded based on scores.

    Arguments:
       passage_answer_stats: List of passage answer scores.
       minimal_answer_stats: List of minimal answer scores.


    Returns:
       Dictionary of name (string) -> score.
    """
    scores = compute_macro_f1(passage_answer_stats, prefix='passage-answer-')
    scores.update(compute_macro_f1(
        minimal_answer_stats, prefix='minimal-answer-'))
    return scores


def compute_pr_curves(answer_stats, targets=None):
    """Computes PR curve and returns R@P for specific targets.

    The values are computed as follows: find the (precision, recall) point
    with maximum recall and where precision > target.

    This is only relevant if you return the system scores in your predictions.
    You may find this useful when attempting to tune the threshold for your
    system on the dev set before requesting an evaluation on the test set
    via the leaderboard.

    Arguments:
      answer_stats: List of statistic tuples from the answer scores.
      targets (None): List of precision thresholds to target.

    Returns:
      List of table with rows: [target, r, p, score].
    """
    total_f1 = 0
    total_has_pred = 0
    total_has_gold = 0

    # Count the number of gold annotations.
    for has_gold, _, _, _ in answer_stats:
        total_has_gold += has_gold

    # Keep track of the point of maximum recall for each target.

    max_recall = [0 for _ in targets]
    max_precision = [0 for _ in targets]
    max_scores = [0.0 for _ in targets]

    # Only keep track of unique thresholds in this dictionary.
    scores_to_stats = collections.OrderedDict()

    # Loop through every possible threshold and compute precision + recall.
    for has_gold, has_pred, is_correct_or_f1, score in answer_stats:
        if isinstance(is_correct_or_f1, tuple):
            _, _, f1 = is_correct_or_f1
        else:
            f1 = is_correct_or_f1
        total_f1 += f1
        total_has_pred += has_pred

        precision = eval_utils.safe_divide(total_f1, total_has_pred)
        recall = eval_utils.safe_divide(total_f1, total_has_gold)

        # If there are any ties, this will be updated multiple times until the
        # ties are all counted.
        scores_to_stats[score] = [precision, recall]

    best_f1 = 0.0
    best_precision = 0.0
    best_recall = 0.0
    best_threshold = 0.0

    for threshold, (precision, recall) in scores_to_stats.items():
        # Match the thresholds to the find the closest precision above some target.
        for t, target in enumerate(targets):
            if precision >= target and recall > max_recall[t]:
                max_recall[t] = recall
                max_precision[t] = precision
                max_scores[t] = threshold

        # Compute optimal threshold.
        f1 = eval_utils.safe_divide(2 * precision * recall, precision + recall)
        if f1 > best_f1:
            best_f1 = f1
            best_precision = precision
            best_recall = recall
            best_threshold = threshold

    return ((best_f1, best_precision, best_recall, best_threshold),
            list(zip(targets, max_recall, max_precision, max_scores)))


def print_r_at_p_table(answer_stats):
    """Pretty prints the R@P table for default targets."""
    opt_result, pr_table = compute_pr_curves(
        answer_stats, targets=[0.5, 0.75, 0.9])
    f1, precision, recall, threshold = opt_result
    print('Optimal threshold: {:.5}'.format(threshold))
    print(' F1     /  P      /  R')
    print('{: >7.2%} / {: >7.2%} / {: >7.2%}'.format(f1, precision, recall))
    for target, recall, precision, row in pr_table:
        print('R@P={}: {:.2%} (actual p={:.2%}, score threshold={:.4})'.format(
            target, recall, precision, row))


def get_metrics_as_dict(gold_path, prediction_path, passage_non_null_threshold, minimal_non_null_threshold, verbose):
    """Library version of the end-to-end evaluation.

    Arguments:
      gold_path: Path to a single JSONL data. Could be gzipped or not.
      prediction_path: Path to the JSONL file of prediction data.

    Returns:
      metrics: A dictionary mapping string names to metric scores.
    """

    tydi_gold_dict = eval_utils.read_annotation(gold_path)
    tydi_pred_dict = eval_utils.read_prediction_jsonl(prediction_path)

    passage_answer_stats, minimal_answer_stats = score_answers(
        tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose)

    return get_metrics_with_answer_stats(
        passage_answer_stats, minimal_answer_stats)


def get_metrics_with_answer_stats(passage_answer_stats, minimal_answer_stats):
    """Generate metrics dict using passage and minimal answer stats."""

    def _get_metric_dict(answer_stats, prefix=''):
        """Compute all metrics for a set of answer statistics."""
        opt_result, pr_table = compute_pr_curves(
            answer_stats, targets=[0.5, 0.75, 0.9])
        f1, precision, recall, threshold = opt_result
        metrics = collections.OrderedDict({
            'best-threshold-f1': f1,
            'best-threshold-precision': precision,
            'best-threshold-recall': recall,
            'best-threshold': threshold,
        })
        for target, recall, precision, _ in pr_table:
            metrics['recall-at-precision>={:.2}'.format(target)] = recall
            metrics['precision-at-precision>={:.2}'.format(target)] = precision

        # Add prefix before returning.
        return dict([(prefix + k, v) for k, v in metrics.items()])

    metrics = _get_metric_dict(passage_answer_stats, 'passage-')
    metrics.update(_get_metric_dict(minimal_answer_stats, 'minimal-'))
    return metrics


def get_latex_str(f1, precision, recall):
    return '\\fpr' + '{' + ('%.1f' % (f1 * 100)) + '}{' + (
            '%.1f' % (precision * 100)) + '}{' + ('%.1f' % (recall * 100)) + '}'


# def load_gt_lookup_as_dict(gold_path, num_workers):
#     cache_path = os.path.join(os.path.dirname(gold_path), 'cache')
#
#     if os.path.exists(cache_path):
#         logging.info('Reading from cache: %s', format(cache_path))
#         tydi_gold_dict = pickle.load(open(cache_path, 'rb'))  # pytype: disable=wrong-arg-types
#     else:
#         tydi_gold_dict = eval_utils.read_annotation(gold_path, n_threads=num_workers)
#         logging.info('Caching gold data for future to: %s', format(cache_path))
#         cache_to_pickle_file(tydi_gold_dict, cache_path)  # pytype: disable=wrong-arg-types
#
#         logging.info('Read in gt lookup for %d examples into memory' % len(tydi_gold_dict))
#     return tydi_gold_dict


def pretty_print(tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold=2, minimal_non_null_threshold=2, verbose=False, skip_missing_example_ids=False):
    if any(map(not_, tydi_gold_dict.values())) or any(annotations[0].minimal_answer_span is None for annotations in tydi_gold_dict.values()):
        logging.warning("At least one example in gold dict has no annotations -- skipping pretty print evaluation")
        return

    gold_id_set = set(tydi_gold_dict.keys())
    pred_id_set = set(tydi_pred_dict.keys())
    sym_diff = gold_id_set.symmetric_difference(pred_id_set)

    if (not skip_missing_example_ids) and sym_diff:
        raise ValueError('ERROR: the example ids in gold annotations and example '
                         'ids in the prediction are not equal.')
    elif skip_missing_example_ids and sym_diff:
        logging.warning("Skipping {} example ids that are only in either gold or preds".format(len(sym_diff)))

    macro_avg_passage_scores = ([], [], [])
    macro_avg_minimal_scores = ([], [], [])

    language_list = [  # TODO: support arbitrary languages
        'english', 'arabic', 'bengali', 'finnish', 'indonesian', 'japanese',
        'swahili', 'korean', 'russian', 'telugu', 'thai'
    ]

    per_lang_gold = {}
    per_lang_pred = {}

    id_to_lang = {}

    for ex_id in gold_id_set:
        ex = tydi_gold_dict[ex_id]
        if ex[0].language in per_lang_gold:
            per_lang_gold[ex[0].language][ex_id] = ex
        else:
            per_lang_gold[ex[0].language] = {ex_id: ex}
        id_to_lang[ex_id] = ex[0].language
    for ex_id in gold_id_set:
        ex = tydi_pred_dict[ex_id]
        if id_to_lang[ex.example_id] in per_lang_pred:
            per_lang_pred[id_to_lang[ex.example_id]][ex.example_id] = ex
        else:
            per_lang_pred[id_to_lang[ex.example_id]] = {ex.example_id: ex}

    for lang in language_list:
        if lang in per_lang_pred:
            passage_answer_stats, minimal_answer_stats = score_answers(
                per_lang_gold.get(lang, {}), per_lang_pred[lang],
                passage_non_null_threshold, minimal_non_null_threshold, verbose)

            # Passage selection task
            opt_result, _ = compute_pr_curves(passage_answer_stats, targets=[0.5])
            f1, precision, recall, _ = opt_result
            if lang != 'english':
                macro_avg_passage_scores[0].append(f1)
                macro_avg_passage_scores[1].append(precision)
                macro_avg_passage_scores[2].append(recall)
            print('Passage & ' + lang + ' & ' + get_latex_str(f1, precision, recall))

            # Minimal answer span task
            opt_result, _ = compute_pr_curves(minimal_answer_stats, targets=[0.5])
            f1, precision, recall, _ = opt_result
            if lang != 'english':
                macro_avg_minimal_scores[0].append(f1)
                macro_avg_minimal_scores[1].append(precision)
                macro_avg_minimal_scores[2].append(recall)
            print('Minimal Answer & ' + lang + ' & ' +
                  get_latex_str(f1, precision, recall))

            print('*' * 20)
            print(lang)
            print('Language: %s (%d)' % (lang, len(per_lang_gold.get(lang, {}))))
            print('*' * 20)
            print('PASSAGE ANSWER R@P TABLE:')
            print_r_at_p_table(passage_answer_stats)
            print('*' * 20)
            print('MINIMAL ANSWER R@P TABLE:')
            print_r_at_p_table(minimal_answer_stats)

    print('Total # examples in gold: %d, # ex. in pred: %d (including english)' %
          (len(tydi_gold_dict), len(tydi_pred_dict)))

    f1_list, precision_list, recall_list = macro_avg_passage_scores
    print('*** Macro Over %d Languages, excluding English **' % len(f1_list))
    avg_passage_f1 = eval_utils.safe_average(f1_list)
    avg_passage_recall = eval_utils.safe_average(recall_list)
    avg_passage_precision = eval_utils.safe_average(precision_list)
    print('Passage F1:%.3f P:%.3f R:%3f' %
          (avg_passage_f1, avg_passage_precision, avg_passage_recall))
    print(get_latex_str(
        avg_passage_f1, avg_passage_precision, avg_passage_recall))

    f1_list, precision_list, recall_list = macro_avg_minimal_scores

    avg_minimal_f1 = eval_utils.safe_average(f1_list)
    avg_minimal_recall = eval_utils.safe_average(recall_list)
    avg_minimal_precision = eval_utils.safe_average(precision_list)
    print('Minimal F1:%.3f P:%.3f R:%3f' %
          (avg_minimal_f1, avg_minimal_precision, avg_minimal_recall))
    print(get_latex_str(
        avg_minimal_f1, avg_minimal_precision, avg_minimal_recall))
    print('*** / Aggregate Scores ****')

    aggregate_metrics = {'avg_passage_f1': avg_passage_f1,
                         'avg_passage_recall': avg_passage_recall,
                         'avg_passage_precision': avg_passage_precision,
                         'avg_minimal_f1': avg_minimal_f1,
                         'avg_minimal_recall': avg_minimal_recall,
                         'avg_minimal_precision': avg_minimal_precision}
    print(json.dumps(aggregate_metrics))
    return aggregate_metrics


# def main(input_args: Optional[List[str]] = None):
#     args = parser.parse_args(input_args)
#     args = vars(args)
#     args['do_pretty_print'] = args.pop('pretty_print')
#     eval_tydi(**args)


# def eval_tydi(gold_path: str, predictions_path: str, verbose: bool = True, cache_gold_data: bool = True, do_pretty_print: bool = False,
#               passage_non_null_threshold: int = 2, minimal_non_null_threshold: int = 2, num_workers: int = 10,
#               skip_missing_ids: bool = False):
#     cache_path = os.path.join(os.path.dirname(gold_path), 'cache')
#     if cache_gold_data and os.path.exists(cache_path):
#         logging.info('Reading from cache: %s', format(cache_path))
#         with open(cache_path, 'rb') as f:
#             tydi_gold_dict = pickle.load(f)  # pytype: disable=wrong-arg-types
#     else:
#         tydi_gold_dict = eval_utils.read_annotation(gold_path, num_workers)
#         if cache_gold_data:
#             logging.info('Caching gold data for future to: %s', format(cache_path))
#             with open(cache_path, 'wb') as f:
#                 pickle.dump(tydi_gold_dict, f, protocol=pickle.HIGHEST_PROTOCOL)  # pytype: disable=wrong-arg-types
#
#     total_ans_count = 0
#     count = 0
#
#     for ans in tydi_gold_dict.values():
non_informative
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
    macro_f1 = eval_utils.safe_divide(
        2 * macro_precision * macro_recall,
        macro_precision + macro_recall)

    return collections.OrderedDict({
        prefix + 'n': len(answer_stats),
        prefix + 'f1': macro_f1,
        prefix + 'precision': macro_precision,
        prefix + 'recall': macro_recall
    })


def compute_final_f1(passage_answer_stats, minimal_answer_stats):
    """Computes overall F1 given passage and minimal answers, ignoring scores.

    Note: this assumes that the answers have been thresholded based on scores.

    Arguments:
       passage_answer_stats: List of passage answer scores.
       minimal_answer_stats: List of minimal answer scores.


    Returns:
       Dictionary of name (string) -> score.
    """
    scores = compute_macro_f1(passage_answer_stats, prefix='passage-answer-')
    scores.update(compute_macro_f1(
        minimal_answer_stats, prefix='minimal-answer-'))
    return scores


def compute_pr_curves(answer_stats, targets=None):
    """Computes PR curve and returns R@P for specific targets.

    The values are computed as follows: find the (precision, recall) point
    with maximum recall and where precision > target.

    This is only relevant if you return the system scores in your predictions.
    You may find this useful when attempting to tune the threshold for your
    system on the dev set before requesting an evaluation on the test set
    via the leaderboard.

    Arguments:
      answer_stats: List of statistic tuples from the answer scores.
      targets (None): List of precision thresholds to target.

    Returns:
      List of table with rows: [target, r, p, score].
    """
    total_f1 = 0
    total_has_pred = 0
    total_has_gold = 0

    # Count the number of gold annotations.
    for has_gold, _, _, _ in answer_stats:
        total_has_gold += has_gold

    # Keep track of the point of maximum recall for each target.

    max_recall = [0 for _ in targets]
    max_precision = [0 for _ in targets]
    max_scores = [0.0 for _ in targets]

    # Only keep track of unique thresholds in this dictionary.
    scores_to_stats = collections.OrderedDict()

    # Loop through every possible threshold and compute precision + recall.
    for has_gold, has_pred, is_correct_or_f1, score in answer_stats:
        if isinstance(is_correct_or_f1, tuple):
            _, _, f1 = is_correct_or_f1
        else:
            f1 = is_correct_or_f1
        total_f1 += f1
        total_has_pred += has_pred

        precision = eval_utils.safe_divide(total_f1, total_has_pred)
        recall = eval_utils.safe_divide(total_f1, total_has_gold)
non_informative
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
random
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
    example_count = 0
    for example_id in gold_id_set:
        example_count += 1
        gold = gold_annotation_dict[example_id]
        pred = pred_dict.get(example_id)
        passage_answer_stats.append(
            score_passage_answer(gold, pred, passage_non_null_threshold))
        minimal_answer_stats.append(
            score_minimal_answer(gold, pred, minimal_non_null_threshold))

        if not verbose:
            continue
        if pred is None or len(pred.short_answer_span_list) == 0:
            continue
        pred_min_start = pred.short_answer_span_list[0].start_byte
        pred_min_end = pred.short_answer_span_list[0].end_byte
        gold_min_start = gold[0].minimal_answer_span.start_byte_offset
        gold_min_end = gold[0].minimal_answer_span.end_byte_offset
        if gold_min_start >= 0:
            logging.info('---')
            logging.info(gold[0].example_id)
            logging.info(gold[0].question_text)
            logging.info('gold offsets %d, %d', gold_min_start, gold_min_end)
            logging.info('pred offsets %d, %d', pred_min_start, pred_min_end)
            logging.info('gold answer: (%s)',
                         byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))
            logging.info('pred answer: (%s)',
                         byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))
            logging.info('score %.2f', minimal_answer_stats[-1][-1])
            logging.info('f1: %.2f, p: %.2f, r: %.2f',
                         minimal_answer_stats[-1][-2][2],
                         minimal_answer_stats[-1][-2][0],
                         minimal_answer_stats[-1][-2][1])
    # use the 'score' column, which is last
    passage_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    minimal_answer_stats.sort(key=lambda x: x[-1], reverse=True)
    return passage_answer_stats, minimal_answer_stats


def compute_macro_f1(answer_stats, prefix=''):
    """Computes F1, precision, recall for a list of answer scores.

    This computes the *language-wise macro F1*. For minimal answers,
    we also compute a partial match score that uses F1, which would be
    included in this computation via `answer_stats`.

    Args:
      answer_stats: List of per-example scores.
      prefix (''): Prefix to prepend to score dictionary.

    Returns:
      Dictionary mapping measurement names to scores.
    """

    has_gold, has_pred, f1, _ = list(zip(*answer_stats))

    macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))
    macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))
    macro_f1 = eval_utils.safe_divide(
        2 * macro_precision * macro_recall,
        macro_precision + macro_recall)

    return collections.OrderedDict({
        prefix + 'n': len(answer_stats),
        prefix + 'f1': macro_f1,
        prefix + 'precision': macro_precision,
        prefix + 'recall': macro_recall
    })


def compute_final_f1(passage_answer_stats, minimal_answer_stats):
    """Computes overall F1 given passage and minimal answers, ignoring scores.

    Note: this assumes that the answers have been thresholded based on scores.

    Arguments:
       passage_answer_stats: List of passage answer scores.
       minimal_answer_stats: List of minimal answer scores.


    Returns:
       Dictionary of name (string) -> score.
    """
    scores = compute_macro_f1(passage_answer_stats, prefix='passage-answer-')
    scores.update(compute_macro_f1(
        minimal_answer_stats, prefix='minimal-answer-'))
    return scores


def compute_pr_curves(answer_stats, targets=None):
    """Computes PR curve and returns R@P for specific targets.

    The values are computed as follows: find the (precision, recall) point
    with maximum recall and where precision > target.

    This is only relevant if you return the system scores in your predictions.
    You may find this useful when attempting to tune the threshold for your
    system on the dev set before requesting an evaluation on the test set
    via the leaderboard.

    Arguments:
      answer_stats: List of statistic tuples from the answer scores.
      targets (None): List of precision thresholds to target.

    Returns:
      List of table with rows: [target, r, p, score].
    """
    total_f1 = 0
    total_has_pred = 0
    total_has_gold = 0

    # Count the number of gold annotations.
    for has_gold, _, _, _ in answer_stats:
        total_has_gold += has_gold

    # Keep track of the point of maximum recall for each target.

    max_recall = [0 for _ in targets]
    max_precision = [0 for _ in targets]
    max_scores = [0.0 for _ in targets]

    # Only keep track of unique thresholds in this dictionary.
    scores_to_stats = collections.OrderedDict()

    # Loop through every possible threshold and compute precision + recall.
    for has_gold, has_pred, is_correct_or_f1, score in answer_stats:
        if isinstance(is_correct_or_f1, tuple):
            _, _, f1 = is_correct_or_f1
        else:
            f1 = is_correct_or_f1
        total_f1 += f1
        total_has_pred += has_pred

        precision = eval_utils.safe_divide(total_f1, total_has_pred)
        recall = eval_utils.safe_divide(total_f1, total_has_gold)

        # If there are any ties, this will be updated multiple times until the
        # ties are all counted.
        scores_to_stats[score] = [precision, recall]

    best_f1 = 0.0
    best_precision = 0.0
random
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

random
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

  predictions_path should point to a jsonl file (one json object per line),
  where each line contains the predictions in the format given below.

  ------------------------------------------------------------------------------

  Prediction format (written on multiple lines here for clarity, but each
  prediction should be a single line in your output file):

  {
    'example_id': -2226525965842375672,
    'passage_answer_index': 2,
    'passage_answer_score': 13.5,
    'minimal_answer': {'start_byte_offset': 64206, 'end_byte_offset': 64280},
    'minimal_answer_score': 26.4,
    'yes_no_answer': 'NONE'
  }

  The prediction format mirrors the annotation format in defining each passage
  or minimal answer span both in terms of byte offsets.

  If start_byte_offset >= 0 and end_byte_offset >=0, use byte offsets,
    else no span is defined (null answer).

  The minimal answer metric takes both minimal answer spans, and the yes/no
  answer into account. If the 'minimal_answers' list contains any non/null
  spans, then 'yes_no_answer' should be set to 'NONE'.

  -----------------------------------------------------------------------------

  Metrics:

  Each prediction should be provided with a passage answer score, and a minimal
  answers score. At evaluation time, the evaluation script will find a score
  threshold at which F1 is maximized. All predictions with scores below this
  threshold are ignored (assumed to be null). If the score is not provided,
  the evaluation script considers all predictions to be valid. The script
  will also output the maximum recall at precision points of >= 0.5, >= 0.75,
  and >= 0.9.

  Key methods:
    Scoring passage answer candidates: score_passage_answer()
    Scoring minimal answer candidates: score_minimal_answer(),
                                       eval_utils.compute_partial_match_scores()
    Computing language-wise F1: compute_macro_f1()
    Averaging over non-English languages: main()

"""

import collections
import json
import os
import pickle


import logging
from operator import not_

from oneqa.mrc.metrics.tydi_f1 import eval_utils
from typing import List, Optional


def score_passage_answer(gold_label_list, pred_label,
                         passage_non_null_threshold):
    """Scores a passage answer as correct or not.

    1) First decide if there is a gold passage answer with
       FLAGS.passage_non_null_threshold.
    2) The prediction will get a match if:
       a. There is a gold passage answer.
       b. The prediction span match exactly with *one* of the non-null gold
          passage answer index.

    Args:
      gold_label_list: A list of TyDiLabel, could be None.
      pred_label: A single TyDiLabel, could be None.
      passage_non_null_threshold: See FLAGS.passage_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, is_correct, score
    """
    gold_has_answer = eval_utils.gold_has_passage_answer(
        gold_label_list, passage_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, False, 0

    # pred_has_answer = pred_label.long_answer_span.start_byte >= 0
    pred_has_answer = pred_label.passage_answer_index != -1

    is_correct = False
    score = pred_label.passage_score

    # Both sides are non-null spans.
    if gold_has_answer and pred_has_answer:
        for gold_label in gold_label_list:
            # while the voting results indicate there is an passage answer, each
            # annotator might still say there is no passage answer.
            if gold_label.passage_answer_index < 0:
                continue

            # if gold_label.passage_span.start_byte_offset == pred_label.long_answer_span.start_byte:
            #     is_correct = True
            #     break

            if gold_label.passage_answer_index == pred_label.passage_answer_index:
              is_correct = True
              break

    return gold_has_answer, pred_has_answer, is_correct, score


def score_minimal_answer(gold_label_list, pred_label,
                         minimal_non_null_threshold):
    """Scores a minimal answer.

    Outputs score against gold label that gives max F1.

    First decide if there is a gold minimal answer with
    FLAGS.minimal_non_null_threshold.
    If any of the gold label has "yes", or "no", and pred label predicted it
    correctly, than precision, recall, f1 is all 1.0.

    Args:
      gold_label_list: A list of TyDiLabel.
      pred_label: A single TyDiLabel.
      minimal_non_null_threshold: See FLAGS.minimal_non_null_threshold.

    Returns:
      gold_has_answer, pred_has_answer, (precision, recall, f1), score
    """

    # There is a gold minimal answer if gold_label_list not empty and non null
    # answers is over the threshold (sum over annotators).
    gold_has_answer = eval_utils.gold_has_minimal_answer(
        gold_label_list, minimal_non_null_threshold)

    if pred_label is None:
        return gold_has_answer, not gold_has_answer, (0, 0, 0), 0

    # There is a predicted minimal answer if the predicted minimal label span
    # is non-null or we have a specific predicted label (such as yes/no).
    pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()
                       or pred_label.yes_no_answer != 'none')

    # score is optional.
    score = pred_label.minimal_score
    # We find the closest (highest scoring) match between the system's predicted
    # minimal answer and one of the three gold annotations.
    max_f1 = 0.0
    max_precision = 0.0
    max_recall = 0.0

    # Both sides have minimal answers, which contains yes/no questions.
    if gold_has_answer and pred_has_answer:
        if pred_label.yes_no_answer != 'none':  # System predicted a yes/no answer.
            for gold_label in gold_label_list:
                if pred_label.yes_no_answer == gold_label.yes_no_answer:
                    max_f1 = 1.0
                    max_precision = 1.0
                    max_recall = 1.0
                    break
        else:
            for gold_label in gold_label_list:
                if gold_label.minimal_answer_span.is_null_span():
                    continue
                # Compute the *micro-F1* (a partial match score for this example).
                # We also compute a language-wise *macro-F1* later.
                precision, recall, f1 = eval_utils.compute_partial_match_scores(
                    gold_label.minimal_answer_span, pred_label.minimal_answer_span)
                if f1 > max_f1:
                    max_f1 = f1
                    max_precision = precision
                    max_recall = recall

    return (gold_has_answer, pred_has_answer,
            (max_precision, max_recall, max_f1), score)


def byte_slice(text, start, end):
    byte_str = bytes(text, 'utf-8')
    return str(byte_str[start:end])


def score_answers(gold_annotation_dict, pred_dict, passage_non_null_threshold, minimal_non_null_threshold, verbose,
                  skip_missing_example_ids=True):
    """Scores all answers for all documents.

    Args:
      gold_annotation_dict: a dict from example id to list of `TyDiLabel`s.
      pred_dict: a dict from example id to list of `TyDiLabel`s.

    Returns:
      passage_answer_stats: List of scores for passage answers.
      minimal_answer_stats: List of scores for minimal answers.
    """
    gold_id_set = set(gold_annotation_dict.keys())
    pred_id_set = set(pred_dict.keys())

    unpredicted = gold_id_set - pred_id_set
    unexpected = pred_id_set - gold_id_set
    if unpredicted:
        logging.warning('Predictions missing for %d examples.', len(unpredicted))
        logging.info('  Missing ids: %s', sorted(unpredicted))
    if unexpected:
        logging.warning(
            'Found predictions for %d examples that do not appear in the gold data.',
            len(unexpected))
        logging.info('  Unexpected ids: %s', sorted(unexpected))
    if not skip_missing_example_ids and (unpredicted or unexpected):
        raise ValueError("Gold and pred example ids not the same and skip missing ids not selected")

    passage_answer_stats = []
    minimal_answer_stats = []
random
# coding=utf-8
# Copyright 2020 The Google Research Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Lint as: python3
r"""Official evaluation script for the TyDi QA primary tasks.

The primary tasks are the Passage Selection Task (SelectP) and the Minimal
Answer Span Task Task (AnsSpan). This script is *not* used for the secondary
task, the SQuAD-compatible gold Passage (GoldP) task.

  ------------------------------------------------------------------------------

  Example usage:

  tydi_eval --gold_path=<path-to-gold-files> --predictions_path=<path_to_jsonl>

  This will compute both the official byte-level F1 scores, recall@precision
  tables for both passage and minimal answers (if the optional answer scores are
  provided), and also breakdown per language.

  Note that R@P are only meaningful if your model populates the score fields
  of the prediction JSON format (which is not required).

  gold_path should point to a single N way annotated dev data in the
  original download format (gzipped jsonlines) or jsonlines.

random
            score_passage_answer(gold, pred, passage_non_null_threshold))             score_minimal_answer(gold, pred, minimal_non_null_threshold))                          byte_slice(gold[0].plaintext, gold_min_start, gold_min_end))                          byte_slice(gold[0].plaintext, pred_min_start, pred_min_end))             passage_answer_stats, minimal_answer_stats = score_answers(             opt_result, _ = compute_pr_curves(passage_answer_stats, targets=[0.5])     opt_result, pr_table = compute_pr_curves(             opt_result, _ = compute_pr_curves(minimal_answer_stats, targets=[0.5])                   get_latex_str(f1, precision, recall))     return get_metrics_with_answer_stats(     Scoring passage answer candidates: score_passage_answer()     Scoring minimal answer candidates: score_minimal_answer(),     scores = compute_macro_f1(passage_answer_stats, prefix='passage-answer-')     Computing language-wise F1: compute_macro_f1()     scores.update(compute_macro_f1(     metrics = _get_metric_dict(passage_answer_stats, 'passage-')     metrics.update(_get_metric_dict(minimal_answer_stats, 'minimal-'))     Averaging over non-English languages: main()     return str(byte_str[start:end])                 per_lang_gold.get(lang, {}), per_lang_pred[lang],         pred = pred_dict.get(example_id)         recall = eval_utils.safe_divide(total_f1, total_has_gold) Answer Span Task Task (AnsSpan). This script is *not* used for the secondary         f1 = eval_utils.safe_divide(2 * precision * recall, precision + recall)     gold_has_answer = eval_utils.gold_has_minimal_answer(     pred_has_answer = (not pred_label.minimal_answer_span.is_null_span()     macro_precision = eval_utils.safe_divide(sum(f1), sum(has_pred))     macro_recall = eval_utils.safe_divide(sum(f1), sum(has_gold))     macro_f1 = eval_utils.safe_divide(   The minimal answer metric takes both minimal answer spans, and the yes/no     avg_passage_f1 = eval_utils.safe_average(f1_list)     avg_passage_recall = eval_utils.safe_average(recall_list)     avg_passage_precision = eval_utils.safe_average(precision_list)                 if gold_label.minimal_answer_span.is_null_span():     avg_minimal_f1 = eval_utils.safe_average(f1_list)                 precision, recall, f1 = eval_utils.compute_partial_match_scores(     avg_minimal_recall = eval_utils.safe_average(recall_list)     avg_minimal_precision = eval_utils.safe_average(precision_list)                                        eval_utils.compute_partial_match_scores()     gold_has_answer = eval_utils.gold_has_passage_answer(         precision = eval_utils.safe_divide(total_f1, total_has_pred)     Note: this assumes that the answers have been thresholded based on scores. #         count += 1     for example_id in gold_id_set:     best_recall = 0.0   This will compute both the official byte-level F1 scores, recall@precision     example_count = 0   predictions_path should point to a jsonl file (one json object per line),