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.mark.flaky(reruns=10, reruns_delay=30)  # Account for intermittent S3 errors downloading HF data
    @pytest.fixture(scope='session')
    def train_examples(self):
        examples = datasets.load_dataset("tydiqa", "primary_task", split='train[:100]')
        return examples

    @pytest.mark.flaky(reruns=10, reruns_delay=30)  # Account for intermittent S3 errors downloading HF data
    @pytest.fixture(scope='session')
    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,
            load_from_cache_file=False,
        )

    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)
        for example in train_examples:
            assert min(example['target']['start_positions']) >= -1
            assert min(example['target']['end_positions']) >= -1
            assert min(example['target']['passage_indices']) >= -1

    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)
        for example in eval_examples:
            assert min(example['target']['start_positions']) >= -1
            assert min(example['target']['end_positions']) >= -1
            assert min(example['target']['passage_indices']) >= -1

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.data_models.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 not train_examples

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

    def test_adapt_dataset_with_train_examples_does_not_modify_examples_with_all_fields(self,
                                                                                        train_examples, preprocessor):
        train_examples = train_examples.add_column('language', ['FOO'] * train_examples.num_rows)
        assert preprocessor.adapt_dataset(train_examples, is_train=True) is train_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', 'offset_mapping'}
        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['start_positions'][i] == cls_index
                assert train_features['end_positions'][i] == cls_index
            elif tt == TargetType.SPAN_ANSWER:
                assert train_features['start_positions'][i] != cls_index
                assert train_features['end_positions'][i] != cls_index
            elif tt == TargetType.NO_ANSWER:
                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)

    def test_language_field_imputed_when_does_not_already_exist(self, train_examples_and_features):
        train_examples, _ = train_examples_and_features
        assert 'language' in train_examples.features
        assert all(lang == 'UNKNOWN' for lang in train_examples['language'])

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

tests/oneqa/mrc/processors/postprocessors/test_extractive_postprocessor.py METASEP
from datasets import Dataset
import pytest
import sys
from pytest import raises
from transformers import AutoTokenizer
from itertools import groupby
from operator import itemgetter


import numpy as np

from oneqa.mrc.processors.postprocessors.extractive import ExtractivePostProcessor
from oneqa.mrc.processors.postprocessors.scorers import SupportedSpanScorers
from oneqa.mrc.data_models.target_type import TargetType
from tests.oneqa.mrc.common.base import UnitTest
from tests.oneqa.mrc.common.parameterization import PARAMETERIZE_INVALID_SUBSAMPLING_PROBABILITIES


class TestExtractivePostProcessor(UnitTest):

    _start_score_of_cls_token = 0
    _end_score_of_cls_token = 0
    _score_none_token = -10
    example1_start_scores = [ [0.5, 0.3, 0.2, 0.1, 0.1], 
                              [1.5, 0.7, 0.2, 0.1, 0.1],
                ]
    example1_end_scores = [ [0.5, 0.3, 0, 0, 0], 
                            [1.5, 0.7, 0, 0, 0],
                ]
    example2_start_scores = [ [-1, 2, 1, -3, -2.5, -5, -4, -2, -2, -1, -1.5, -2.5],  
                              [-3, -2, -5, -2,-1, -1, -3, -2, -4, -5, -1.5, -2.5], 
                              [-3, -2, -5, -2,-1, -1, -3, -2, -4, -5, -1.5, -2.5], 
                ]
    example2_end_scores = [   [-1, 1, 1.5, -3, -2.5, -5, -4, -2, -2, -1, -1.5, -2.5], 
                              [-3, -2, -5, -2,-1, -1, -3, -2, -4, -5, -1.5, -2.5],
                              [-3, -2, -5, -2,-1, -1, -3, -2, -4, -5, -1.5, -2.5],
                ]

    _start_scores = [example1_start_scores, example2_start_scores]
    _end_scores = [ example1_end_scores, example2_end_scores]

    _target_type_scores = [
        np.array([ -4.2624542e-04, 1.2508204e+00, -1.8462906e-02, -4.0960628e-01, -3.8014248e-01,  -4.2624542e-04], dtype=np.float32),
        np.array([  1.0599241 , 0.01550296, -0.40537557, -0.3395432 ,  0.12297338],dtype=np.float32),
    ]

    _expected_predictions = {
        'foo-abc' : {
            'start_index': 0, #8,
            'end_index': 0, #8,
            'passage_index': 1,
            'target_type' : int(TargetType.SPAN_ANSWER),
            'span_answer_text': 'Bob'
        },
        'bar-123': {
            'start_index': 1, # 9,
            'end_index': 2, #10,
            'passage_index': 0,
            'target_type' : int(TargetType.NO_ANSWER),
            'span_answer_text': 'quick brown'
        }
    }

    def _start_end_target_type_logits(self, examples, features):

        all_start_logits = []
        all_end_logits = []
        all_target_type_logits = []
        adjusted_start_end_index = {}

        features_itr = groupby(features, key=itemgetter('example_idx'))
        for example in examples:
            example_id = example['example_id']
            expected = self._expected_predictions[example_id]

            example_idx, example_features = next(features_itr)
            example_features = list(example_features)

            for feat_idx, feature in enumerate(example_features):
                offset_mapping = feature["offset_mapping"]

                if feat_idx == expected['passage_index']:
                    for i, offset in enumerate(offset_mapping):
                        if offset != None:
                            adjusted_start_end_index[example_id] = (expected['start_index']+i,
                                                    expected['end_index'] + i)
                            break

                start_logits = []
                end_logits = []
                start_scores_iter = iter(self._start_scores[example_idx][feat_idx])

                end_scores_iter = iter(self._end_scores[example_idx][feat_idx])
                for i, offset in enumerate(offset_mapping):
                    if i == 0:
                        start_logits.append(self._start_score_of_cls_token)
                        end_logits.append(self._end_score_of_cls_token)
                    elif offset == None:
                        start_logits.append(self._score_none_token)
                        end_logits.append(self._score_none_token)
                    else:
                        start_logits.append(next(start_scores_iter))
                        end_logits.append(next(end_scores_iter))

                start_logits +=  [sys.float_info.min]*(128-len(start_logits))
                all_start_logits.append(np.array( start_logits))
                end_logits += [sys.float_info.min]*(128-len(end_logits))
                all_end_logits.append(np.array( end_logits))
                all_target_type_logits.append(self._target_type_scores[example_idx])

        return adjusted_start_end_index, ( np.array(all_start_logits), np.array(all_end_logits), np.array(all_target_type_logits) )

    def test_post_processor_has_examples_and_features(self, eval_examples_and_features):
        eval_examples, eval_features = eval_examples_and_features
        postprocessor_class = ExtractivePostProcessor  
        scorer_type='weighted_sum_target_type_and_score_diff'

        expected_start_end_index, predictions = self._start_end_target_type_logits(eval_examples, eval_features)
        
        postprocessor = postprocessor_class(k=5, n_best_size=3, max_answer_length=30,
                                            scorer_type=SupportedSpanScorers(scorer_type),
                                            single_context_multiple_passages=False)
        example_predictions = postprocessor.process(eval_examples, eval_features, predictions)
        for example_id, preds in  example_predictions.items():
            predicted = preds[0]
            expected = self._expected_predictions[example_id]
            expected_start_end = expected_start_end_index[example_id]
            ptargettype = int(np.argmax(predicted['target_type_logits']))
            assert (predicted['start_index'], predicted['end_index']) == expected_start_end
            assert predicted['passage_index']  == expected['passage_index'] 
            assert ptargettype  == expected['target_type'] 




    

tests/oneqa/mrc/processors/postprocessors/__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.data_models.model_outputs.extractive import ExtractiveQAModelOutput
from oneqa.mrc.data_models.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/metrics/tydi_f1/test_tydi_f1.py METASEP
import datasets
import pytest

from oneqa.mrc.metrics.tydi_f1.tydi_f1 import TyDiF1
from oneqa.mrc.data_models.target_type import TargetType
from tests.oneqa.mrc.common.base import UnitTest


class TestTyDiF1(UnitTest):
    @pytest.fixture(scope='session')
    def metric(self):
        return TyDiF1()

    @pytest.fixture(scope='session')
    def n_annotators(self):
        return 5

    @pytest.fixture(scope='session')
    def references_and_predictions(self, n_annotators):
        references = [
            [dict(start_position=17, end_position=42, passage_index=1,
                  yes_no_answer=TargetType.NO_ANSWER, example_id='a', language='finnish')] * n_annotators,
            [dict(start_position=-1, end_position=-1, passage_index=2,
                  yes_no_answer=TargetType.YES, example_id='b', language='finnish')] * n_annotators,
            [dict(start_position=17, end_position=42, passage_index=1,
                  yes_no_answer=TargetType.NO_ANSWER, example_id='c', language='swahili')] * 2 +
            [dict(start_position=-1, end_position=-1, passage_index=-1,
                  yes_no_answer=TargetType.NO_ANSWER, example_id='c', language='swahili')] * (n_annotators - 2),
            [dict(start_position=-1, end_position=-1, passage_index=-1,
                  yes_no_answer=TargetType.NO_ANSWER, example_id='d', language='swahili')] * n_annotators,
            [dict(start_position=-1, end_position=-1, passage_index=1,
                  yes_no_answer=TargetType.NO_ANSWER, example_id='e', language='swahili')] * n_annotators,
            [dict(start_position=17, end_position=42, passage_index=1,
                  yes_no_answer=TargetType.NO_ANSWER, example_id='f', language='thai')] * n_annotators,
            [dict(start_position=17, end_position=42, passage_index=1,
                  yes_no_answer=TargetType.NO_ANSWER, example_id='g', language='thai')] * n_annotators,
            [dict(start_position=17, end_position=42, passage_index=1,
                  yes_no_answer=TargetType.NO_ANSWER, example_id='h', language='korean')] * n_annotators,
            [dict(start_position=17, end_position=42, passage_index=1,
                  yes_no_answer=TargetType.NO_ANSWER, example_id='i', language='korean')] * n_annotators,

        ]
        predictions = [
            dict(start_position=17, end_position=42, passage_index=1,
                 yes_no_answer=TargetType.NO_ANSWER, example_id='a', confidence_score=10.),
            dict(start_position=-1, end_position=-1, passage_index=2,
                 yes_no_answer=TargetType.YES, example_id='b', confidence_score=9.9),
            dict(start_position=17, end_position=42, passage_index=1,
                 yes_no_answer=TargetType.NO_ANSWER, example_id='c', confidence_score=9.8),
            dict(start_position=17, end_position=42, passage_index=1,
                 yes_no_answer=TargetType.NO_ANSWER, example_id='d', confidence_score=-10.),
            dict(start_position=17, end_position=42, passage_index=1,
                 yes_no_answer=TargetType.NO_ANSWER, example_id='e', confidence_score=-5.),
            dict(start_position=24, end_position=44, passage_index=1,
                 yes_no_answer=TargetType.NO_ANSWER, example_id='f', confidence_score=10.1),
            dict(start_position=24, end_position=44, passage_index=4,
                 yes_no_answer=TargetType.NO_ANSWER, example_id='g', confidence_score=10.1),
            dict(start_position=111, end_position=141, passage_index=1,
                 yes_no_answer=TargetType.NO_ANSWER, example_id='h', confidence_score=10.1),
            dict(start_position=111, end_position=141, passage_index=4,
                 yes_no_answer=TargetType.NO_ANSWER, example_id='i', confidence_score=10.1),
        ]
        return dict(references=references, predictions=predictions)

    def test_instantiate_metric_from_class(self, metric):
        _ = metric

    def test_instantiate_metric_from_load_metric(self):
        from oneqa.mrc.metrics.tydi_f1 import tydi_f1
        _ = datasets.load_metric(tydi_f1.__file__)

    def test_compute_metric(self, metric, references_and_predictions):
        metric.add_batch(**references_and_predictions)
        actual_metric_values = metric.compute()
        expected_metric_values = {
            "avg_passage_f1": 0.75, "avg_passage_recall": 0.75, "avg_passage_precision": 0.75,
            "avg_minimal_f1": 0.7, "avg_minimal_recall": 0.7, "avg_minimal_precision": 0.7
        }
        assert actual_metric_values == expected_metric_values

tests/oneqa/mrc/metrics/tydi_f1/__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.data_models.model_outputs.extractive import ExtractiveQAModelOutput
from oneqa.mrc.data_models.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=10, reruns_delay=30)  # 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.data_models.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/tableqg/test_tableqg.py METASEP

tests/oneqa/tableqg/__init__.py METASEP

tests/oneqa/re2g/test_re2g.py METASEP

tests/oneqa/re2g/__init__.py METASEP

tests/oneqa/mrc/__init__.py METASEP

tests/oneqa/listqa/test_listqa.py METASEP

tests/oneqa/listqa/__init__.py METASEP

tests/oneqa/ir/test_ir.py METASEP

tests/oneqa/ir/__init__.py METASEP

tests/oneqa/distillation/test_distillation.py METASEP

tests/oneqa/distillation/__init__.py METASEP

tests/oneqa/calibration/test_calibration.py METASEP

tests/oneqa/calibration/__init__.py METASEP

tests/oneqa/boolqa/test_boolqa.py METASEP

tests/oneqa/boolqa/__init__.py METASEP

examples/tableqg/README.md METASEP

examples/re2g/README.md METASEP

examples/mrc/README.md METASEP
# Machine Reading Comprehension (MRC)

The primary script is [run_mrc.py](./run_mrc.py).  This runs a transformer-based MRC pipeline.
Before continuing below make sure you have OneQA [installed](../../README.md#Installation).

## Supported Datasets
Currently supported datasets include:
- TyDiQA

## Example Usage
An example usage command is:
```shell
python examples/mrc/run_mrc.py --model_name_or_path xlm-roberta-large \
       --output_dir ${OUTPUT_DIR} --fp16 --learning_rate 4e-5 \
       --do_train --do_eval --per_device_train_batch_size 16 \
       --per_device_eval_batch_size 128 --gradient_accumulation_steps 4 \
       --warmup_ratio 0.1 --weight_decay 0.1 --save_steps 50000 \
       --overwrite_output_dir --num_train_epochs 1 --evaluation_strategy no
```

This will detect a GPU if present as well as multiple CPU cores for accelerating preprocessing.
Some hyperparameters (e.g. fp16, batch size, gradient accumulation steps) may need to be changed
depending on your hardware configuration.

The dataset name and config are currently omitted as only TyDi is supported at the moment.

### Task Arguments

Some task arguments take references which allow for dynamic imports of existing or
user-defined functionality.  For example, to select the `ExtractivePostProcessor` use
`--postprocessor oneqa.mrc.processors.postprocessors.extractive.ExtractivePostProcessor`.
Alternatively, a new postprocessor could be written and selected with 
`--postprocessor qualified.path.to.new.postprocessor.NewPostProcessor`.

For example, if one was implementing a new model which made predictions by means other than
an extractive head then a `NewPostProcessor` which derived predictions from the model
outputs would be needed.

Similarly, when adding support for a new dataset (with a new schema) a new preprocessor would be needed.
This would be selected by specifying `--preprocessor qualified.path.to.new.postprocessor.NewPreProcessor`
for the `NewPreProcessor` corresponding to this dataset and schema.
examples/listqa/README.md METASEP

examples/ir/README.md METASEP

examples/distillation/README.md METASEP

examples/calibration/README.md METASEP

examples/boolqa/README.md METASEP

docs/_templates/custom-module-template.rst METASEP
{{ fullname | escape | underline}}

.. automodule:: {{ fullname }}

   {% block attributes %}
   {% if attributes %}
   .. rubric:: Module attributes

   .. autosummary::
      :toctree:
   {% for item in attributes %}
      {{ item }}
   {%- endfor %}
   {% endif %}
   {% endblock %}

   {% block functions %}
   {% if functions %}
   .. rubric:: {{ _('Functions') }}

   .. autosummary::
      :toctree:
      :nosignatures:
   {% for item in functions %}
      {{ item }}
   {%- endfor %}
   {% endif %}
   {% endblock %}

   {% block classes %}
   {% if classes %}
   .. rubric:: {{ _('Classes') }}

   .. autosummary::
      :toctree:
      :template: custom-class-template.rst
      :nosignatures:
   {% for item in classes %}
      {{ item }}
   {%- endfor %}
   {% endif %}
   {% endblock %}

   {% block exceptions %}
   {% if exceptions %}
   .. rubric:: {{ _('Exceptions') }}

   .. autosummary::
      :toctree:
   {% for item in exceptions %}
      {{ item }}
   {%- endfor %}
   {% endif %}
   {% endblock %}

{% block modules %}
{% if modules %}
.. autosummary::
   :toctree:
   :template: custom-module-template.rst
   :recursive:
{% for item in modules %}
   {{ item }}
{%- endfor %}
{% endif %}
{% endblock %}

docs/_templates/custom-class-template.rst METASEP
{{ fullname | escape | underline}}

.. currentmodule:: {{ module }}

.. autoclass:: {{ objname }}
   :members:
   :show-inheritance:
   :inherited-members:
   :special-members: __call__, __add__, __mul__

   {% block methods %}
   {% if methods %}
   .. rubric:: {{ _('Methods') }}

   .. autosummary::
      :nosignatures:
   {% for item in methods %}
      {%- if not item.startswith('_') %}
      ~{{ name }}.{{ item }}
      {%- endif -%}
   {%- endfor %}
   {% endif %}
   {% endblock %}

   {% block attributes %}
   {% if attributes %}
   .. rubric:: {{ _('Attributes') }}

   .. autosummary::
   {% for item in attributes %}
      ~{{ name }}.{{ item }}
   {%- endfor %}
   {% endif %}
   {% endblock %}

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

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',
                      'passage_answer_candidates': 'passage_candidates'}
    _rename_target = {'passage_answer_candidate_index': 'passage_indices',
                      'minimal_answers_start_byte': 'start_positions',
                      'minimal_answers_end_byte': 'end_positions'}
    _rename_passages = {'plaintext_start_byte': 'start_positions',
                        'plaintext_end_byte': 'end_positions'}
    _single_context_type = {'passage_answer_candidates': Sequence(
        feature={'plaintext_start_byte': Value(dtype='int32', id=None),
                 'plaintext_end_byte': Value(dtype='int32', id=None)})}

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        if not self._single_context_multiple_passages:
            self._logger.info(f"{self.__class__.__name__} only supports single context multiple passages -- enabling")
            self._single_context_multiple_passages = True

    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._rename_examples,
                              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 = example['document_plaintext']
        context_bytes = context.encode('utf-8')

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

            example['target']['start_positions'][i] = len(
                context_bytes[:example['target']['start_positions'][i]].decode('utf-8', errors='replace'))
            example['target']['end_positions'][i] = len(
                context_bytes[:example['target']['end_positions'][i]].decode('utf-8', errors='replace'))

        num_passages = len(example['passage_candidates']['start_positions'])
        for i in range(num_passages):
            passage_start_position = example['passage_candidates']['start_positions'][i]
            passage_end_position = example['passage_candidates']['end_positions'][i]
            example['passage_candidates']['start_positions'][i] = len(
                context_bytes[:passage_start_position].decode('utf-8', errors='replace'))
            example['passage_candidates']['end_positions'][i] = len(
                context_bytes[:passage_end_position].decode('utf-8', errors='replace'))

        example['context'] = [context]
        return example

    def _rename_examples(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

        passage_candidates = example['passage_candidates']
        for old_key, new_key in self._rename_passages.items():
            passage_candidates[new_key] = passage_candidates.pop(old_key)
        example['passage_candidates'] = passage_candidates
        return example

oneqa/mrc/processors/preprocessors/default.py METASEP
import itertools
from lib2to3.pgen2.tokenize import tokenize
import random
import uuid
from operator import sub
from typing import List, Iterable, Tuple, Any, Dict, Union

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

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


# TODO type signatures for all methods
class DefaultPreProcessor(AbstractPreProcessor):  # todo better name?
    _del_keys = ["overflow_to_sample_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)}}
    _example_id_type = {'example_id': Value(dtype='string', id=None)}
    _single_context_type = {'passage_candidates': {
        '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)
    }}
    _language_feature_type = {'language': Value(dtype='string', id=None)}

    def adapt_dataset(self, dataset: Dataset, is_train: bool) -> Dataset:
        if 'example_id' not in dataset.features:
            # example_id = [str(uuid.uuid4()) for _ in range(dataset.num_rows)]
            # dataset = dataset.add_column('example_id', example_id)
            dataset = dataset.map(  # Map instead of add column to allow caching
                self._insert_example_ids,
                batched=True,
                load_from_cache_file=self._load_from_cache_file,
                num_proc=self._num_workers,
            )
        if 'language' not in dataset.features:
            dataset = dataset.add_column('language', ['UNKNOWN'] * dataset.num_rows)
        self.validate_schema(dataset, is_train, pre_adaptation=False)
        return dataset

    @staticmethod
    def _insert_example_ids(examples: Batch) -> Batch:
        n_examples = len(examples['question'])
        example_id = [str(uuid.uuid4()) for _ in range(n_examples)]
        examples['example_id'] = example_id
        return examples

    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,
            with_indices=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, indices: List[int], 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]
        examples_question = [q.lstrip()[:self._max_q_char_len] for q in examples_question]
        
        # create 1:1 question:context lists
        expanded_examples_question = []
        expanded_examples_idx = []
        for i, (question, context) in enumerate(zip(examples_question, examples_context)):
            context = self._trim_to_max_contexts(context, examples, i)
            n_context_for_example = len(context)
            if self._single_context_multiple_passages and n_context_for_example != 1:
                raise ValueError("Must have exactly one context for each question "
                                 "to use single_context_multiple_passages")
            expanded_examples_question.extend(itertools.repeat(question, n_context_for_example))
            expanded_examples_idx.extend(itertools.repeat(i, n_context_for_example))
        expanded_examples_context = list(itertools.chain.from_iterable(examples_context))

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

        tokenized_examples['example_idx'] = [expanded_examples_idx[oidx] for oidx in tokenized_examples["overflow_to_sample_mapping"]]
        tokenized_examples['example_id'] = [examples['example_id'][eidx] for eidx in tokenized_examples['example_idx']]

        if not self._single_context_multiple_passages:  # context_idx only defined in this case
            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))

        if is_train:
            tokenized_examples = self._create_train_targets(tokenized_examples, examples)
            tokenized_examples = self.label_features_for_subsampling(tokenized_examples, examples)
        else:
            tokenized_examples = self._create_eval_targets(tokenized_examples)

        tokenized_examples['example_idx'] = [indices[eidx] for eidx in tokenized_examples['example_idx']]

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

    def _create_train_targets(self, tokenized_examples: BatchEncoding, examples: Batch) -> BatchEncoding:  # TODO: rename create targets fns?
        target = examples['target']

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

            # 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

            if passage_index == -1:
                window_contains_correct_passage = False
            else:
                if self._single_context_multiple_passages:
                    passage_candidates = examples['passage_candidates'][example_index]
                    passage_start_position = passage_candidates['start_positions'][passage_index]
                    passage_end_position = passage_candidates['end_positions'][passage_index]
                    window_contains_correct_passage = self._spans_intersect(
                        (passage_start_position, passage_end_position),
                        (offsets[token_start_index][0], offsets[token_end_index][1])
                    )
                else:
                    context_idx = tokenized_examples['context_idx'][i]
                    window_contains_correct_passage = passage_index == context_idx

            if window_contains_correct_passage and start_position == -1:  # Passage or Y/N 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 not window_contains_correct_passage 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
                # 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.PASSAGE_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())
        if not pre_adaptation:
            items = itertools.chain(items, cls._example_id_type.items(), cls._language_feature_type.items())
        if self._single_context_multiple_passages:
            items = itertools.chain(items, cls._single_context_type.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 not pre_adaptation:
        #     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

    # def _generate_previous_spans_per_example(self, tokenized_examples: BatchEncoding, examples: Batch) -> List[int]:
    #     context_idx = [-1] * len(tokenized_examples["overflow_to_sample_mapping"])

    @staticmethod
    def _spans_intersect(s1: Tuple[int, int], s2: Tuple[int, int]) -> bool:
        return (s1[0] <= s2[0] <= s1[1]) or (s2[0] <= s1[0] <= s2[1]) or \
               (s1[0] <= s2[1] <= s1[1]) or (s2[0] <= s1[1] <= s2[1])

    def _trim_to_max_contexts(self,
                              context: Union[List[str], List[List[str]]],
                              examples: Batch,
                              example_idx: int) -> Union[List[str], List[List[str]]]:
        if self._max_contexts is None:
            pass
        elif self._single_context_multiple_passages:
            passage_candidates = examples['passage_candidates'][example_idx]
            if len(passage_candidates['start_positions']) > self._max_contexts:
                context[0] = context[0][:passage_candidates['end_positions'][self._max_contexts - 1]]
        else:
            context = context[:self._max_contexts]
        return context

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 = True,
                 max_q_char_len: int = 128,
                 single_context_multiple_passages: bool = False,
                 max_contexts: Optional[int] = None):
        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
        self._max_q_char_len = max_q_char_len
        self._single_context_multiple_passages = single_context_multiple_passages
        self._max_contexts = max_contexts

        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/scorers.py METASEP
"""
BEGIN_COPYRIGHT

IBM Confidential
OCO Source Materials

5727-I17
(C) Copyright IBM Corp. 2019 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
"""
import logging
from functools import partial

from typing import List, Optional, Callable, Union
from enum import Enum


from oneqa.mrc.data_models.target_type import TargetType


class SupportedSpanScorers(Enum):
    SCORE_DIFF_BASED = 'score_diff_based'
    TARGET_TYPE_WEIGHTED_SCORE_DIFF = 'target_type_weighted_score_diff'
    WEIGHTED_SUM_TARGET_TYPE_AND_SCORE_DIFF = 'weighted_sum_target_type_and_score_diff'

    @classmethod
    def get_supported(cls):
        return [entry.value for entry in cls]


def initialize_scorer(scorer_type: Union[str, SupportedSpanScorers], target_type_weight: Optional[float]=0.5) -> Callable:
    if not isinstance(scorer_type, SupportedSpanScorers):
        scorer_type = SupportedSpanScorers(scorer_type)
    if scorer_type == SupportedSpanScorers.SCORE_DIFF_BASED:
        logging.debug("\tInitialized scorer %s" % compute_score_diff_between_span_and_cls.__name__)
        return compute_score_diff_between_span_and_cls

    elif scorer_type == SupportedSpanScorers.TARGET_TYPE_WEIGHTED_SCORE_DIFF:
        logging.debug("\tInitialized scorer %s" %
                      compute_short_answer_type_weighted_score_diff_between_span_and_cls.__name__)
        return compute_short_answer_type_weighted_score_diff_between_span_and_cls

    elif scorer_type == SupportedSpanScorers.WEIGHTED_SUM_TARGET_TYPE_AND_SCORE_DIFF:
        logging.debug(
            "\tInitialized scorer %s with weight %s for the target type score and weight %s for "
            "the score diff score" % (
                compute_weighted_sum_short_answer_type_score_diff_between_span_and_cls.__name__,
                target_type_weight, 1 - target_type_weight))

        return partial(
            compute_weighted_sum_short_answer_type_score_diff_between_span_and_cls,
            target_type_weight=target_type_weight)
    else:
        raise ValueError('Unsupported scorer type: %s' % scorer_type)


def compute_score_diff_between_span_and_cls(
        span_score: float, null_span_score: float, *args, **kwargs):
    return span_score - null_span_score


def compute_short_answer_type_weighted_score_diff_between_span_and_cls(
        span_score: float, null_span_score: float,
        target_type_logits: List, *args, **kwargs):
    score_diff = compute_score_diff_between_span_and_cls(
        span_score=span_score, null_span_score=null_span_score)
    return score_diff * target_type_logits[int(TargetType.SPAN_ANSWER)]


def compute_weighted_sum_short_answer_type_score_diff_between_span_and_cls(
        span_score: float, null_span_score: float,
        target_type_logits: List, target_type_weight: float, *args, **kwargs):
    score_diff = compute_score_diff_between_span_and_cls(
        span_score=span_score, null_span_score=null_span_score)

    return (1 - target_type_weight) * score_diff + target_type_weight * target_type_logits[
        int(TargetType.SPAN_ANSWER)]

oneqa/mrc/processors/postprocessors/extractive.py METASEP
from collections import defaultdict
from itertools import groupby
from operator import itemgetter
from typing import List, Dict, Any

from datasets import Dataset
from tqdm import tqdm
import numpy as np
from datetime import datetime
import torch
import logging
from transformers import EvalPrediction

from oneqa.mrc.processors.postprocessors.abstract import AbstractPostProcessor
from oneqa.mrc.processors.postprocessors.scorers import initialize_scorer
from oneqa.mrc.data_models.target_type import TargetType
from oneqa.mrc.data_models.eval_prediction_with_processing import EvalPredictionWithProcessing
from oneqa.mrc.processors.postprocessors.scorers import SupportedSpanScorers

logger = logging.getLogger(__name__)


class ExtractivePostProcessor(AbstractPostProcessor):
    def __init__(self, k: int,
                 n_best_size: int,
                 max_answer_length: int,
                 scorer_type=SupportedSpanScorers.WEIGHTED_SUM_TARGET_TYPE_AND_SCORE_DIFF,
                 single_context_multiple_passages: bool = False):
        super().__init__(k)
        self._n_best_size = n_best_size
        self._max_answer_length = max_answer_length
        self._score_calculator = initialize_scorer(scorer_type)
        self._single_context_multiple_passages = single_context_multiple_passages

    def process(self, examples, features, predictions):
        features_itr = groupby(features, key=itemgetter('example_idx'))  # TODO: may need to change this to example_id when getting rid of example_idx column.
        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]}")
        
        all_start_logits, all_end_logits, all_targettype_logits = predictions

        # The dictionaries we have to fill.
        all_predictions = {}

        start_idx = 0
        for example_idx, example in enumerate(tqdm(examples)):
            feat_example_idx, example_features = next(features_itr)
            if feat_example_idx != example_idx:
                raise ValueError(f"Example id mismatch between example ({example['example_id']}) "
                                 f"and feature ({feat_example_idx})")
            example_features = list(example_features)
            example_id = example_features[0]['example_id']
            contexts = example["context"]
            example_start_logits = all_start_logits[start_idx:start_idx+len(example_features)]
            example_end_logits = all_end_logits[start_idx:start_idx+len(example_features)]
            example_targettype_preds = all_targettype_logits[start_idx:start_idx+len(example_features)]  
            start_idx += len(example_features)

            min_null_prediction = None
            prelim_predictions = []

            for i, input_feature in enumerate(example_features):
                if input_feature['example_id'] != example_id:
                    raise ValueError(f"Example id mismatch between example ({example_id}) "
                                 f"and feature ({input_feature['example_id']})")
                start_logits = example_start_logits[i].tolist()
                end_logits = example_end_logits[i].tolist()
                target_type_logits = example_targettype_preds[i].tolist()
                offset_mapping = input_feature["offset_mapping"]

                token_is_max_context = input_feature.get("token_is_max_context", None)
                # Update minimum null prediction.
                feature_null_score = start_logits[0] + end_logits[0]
                if min_null_prediction is None or min_null_prediction["score"] > feature_null_score:
                    min_null_prediction = {
                        "offsets": (0, 0),
                        "score": feature_null_score,
                        "start_logit": start_logits[0],
                        "end_logit": end_logits[0],
                    }

                start_indexes = np.argsort(start_logits[:len(offset_mapping)])[-1 : -self._n_best_size - 1 : -1].tolist()
                end_indexes = np.argsort(end_logits[:len(offset_mapping)])[-1 : -self._n_best_size - 1 : -1].tolist()
                for start_index in start_indexes:
                    for end_index in end_indexes:
                    # Don't consider out-of-scope answers, either because the indices are out of bounds or correspond
                    # to part of the input_ids that are not in the context.
                        if (
                            start_index >= len(offset_mapping)
                            or end_index >= len(offset_mapping)
                            or offset_mapping[start_index] is None
                            or len(offset_mapping[start_index]) < 2
                            or offset_mapping[end_index] is None
                            or len(offset_mapping[end_index]) < 2
                        ):
                            continue
                        # Don't consider answers with a length that is either < 0 or > max_answer_length.
                        if end_index < start_index or end_index - start_index + 1 > self._max_answer_length:
                            continue
                        # Don't consider answer that don't have the maximum context available (if such information is
                        # provided).
                        if token_is_max_context is not None and not token_is_max_context.get(str(start_index), False):
                            continue

                        start_position = offset_mapping[start_index][0]
                        end_position = offset_mapping[end_index][1]

                        if self._single_context_multiple_passages:
                            passage_candidates = example['passage_candidates']
                            for context_idx in range(len(passage_candidates['start_positions'])):
                                passage_start_position = passage_candidates['start_positions'][context_idx]
                                passage_end_position = passage_candidates['end_positions'][context_idx]
                                if passage_start_position <= start_position <= end_position <= passage_end_position:
                                    break
                            else:
                                context_idx = -1
                            passage_text = contexts[0]
                        else:
                            context_idx = input_feature['context_idx']
                            passage_text = contexts[context_idx]

                        span_answer_text = passage_text[offset_mapping[start_index][0]:offset_mapping[end_index][1]]
                        span_answer_score = self._score_calculator(start_logits[start_index] + end_logits[end_index],
                                                feature_null_score, target_type_logits)
                        prelim_predictions.append(
                        {
                            'example_id': input_feature['example_id'],
                            'cls_score': feature_null_score,
                            'start_logit': start_logits[start_index],
                            'end_logit': end_logits[end_index],
                            'span_answer': {
                                "start_position": start_position,
                                "end_position": end_position,
                            },
                            'span_answer_score' : span_answer_score,
                            'start_index': start_index,
                            'end_index':   end_index,
                            'passage_index' : context_idx,
                            'target_type_logits': target_type_logits,
                            'span_answer_text': span_answer_text,
                            'yes_no_answer': int(TargetType.NO_ANSWER)
                        }
                    )
            example_predictions = sorted(prelim_predictions, key=itemgetter('span_answer_score'), reverse=True)[:self._n_best_size]
            all_predictions[example_id] = example_predictions

            # In the very rare edge case we have not a single non-null prediction, we create a fake prediction to avoid
            # failure.
            if len(example_predictions) == 0: 
                logger.info(f'We do not have any non-null predictions for example {example_id}')
                example_predictions.append( 
                    {
                        'example_id': example_id,
                        'cls_score': 0.0,
                        'start_logit': 0.0, 
                        'end_logit': 0.0, 
                        'span_answer': {'start_position': -1, 'end_position': -1,},
                        'span_answer_score': 0.0, 
                        'span_answer_text': "empty", 
                        'start_index': -1,
                        'end_index':   -1,
                        'passage_index' : -1,
                        'yes_no_answer': int(TargetType.NO_ANSWER)
                    })

            # Compute the softmax of all scores (we do it with numpy to stay independent from torch/tf in this file, using
            # the LogSumExp trick).
            scores = np.array([pred["span_answer_score"] for pred in example_predictions])
            exp_scores = np.exp(scores - np.max(scores))
            probs = exp_scores / exp_scores.sum()

            # Include the probabilities in our predictions.
            for prob, pred in zip(probs, example_predictions):
                pred["normalized_span_answer_score"] = prob

            # We assume null answer is not possible
            # Otherwise we first need to find the best non-empty prediction.
            # i = 0
            # while predictions[i]["text"] == "":
            #     i += 1
            # best_non_null_pred = predictions[i]

            # Then we compare to the null prediction using the threshold.
            # score_diff = null_score - best_non_null_pred["start_logit"] - best_non_null_pred["end_logit"]
            # scores_diff_json[example["id"]] = float(score_diff)  # To be JSON-serializable.
            # if score_diff > null_score_diff_threshold:
            #     all_predictions[example["id"]] = ""
            # else:
            #     all_predictions[example["id"]] = best_non_null_pred["text"]

        return all_predictions
        
    def prepare_examples_as_references(self, examples: Dataset) -> List[Dict[str, Any]]:
        references = []
        n_annotators = len(examples[0]['target']['start_positions'])
        for example_idx in range(examples.num_rows):
            example = examples[example_idx]
            label = {
                'start_position': example['target']['start_positions'],
                'end_position': example['target']['end_positions'],
                'passage_index': example['target']['passage_indices'],
                'yes_no_answer': list(map(TargetType.from_bool_label, example['target']['yes_no_answer'])),  # TODO: decide on schema type for bool ans
                'example_id': [example['example_id']] * n_annotators,
                'language': [example['language']] * n_annotators  # TODO: schema does not have language by default
            }
            references.append(label)
        return references

    def process_references_and_predictions(self, examples, features, predictions) -> EvalPrediction:
        references = self.prepare_examples_as_references(examples)
        predictions = self.process(examples, features, predictions)
        predictions_for_metric = []

        for example_id, preds in predictions.items():
            top_pred = preds[0]
            prediction_for_metric = {
                'example_id': example_id,
                'start_position': top_pred['span_answer']['start_position'],
                'end_position': top_pred['span_answer']['end_position'],
                'passage_index': top_pred['passage_index'],
                'yes_no_answer': top_pred['yes_no_answer'],
                'confidence_score': top_pred['span_answer_score']
            }
            predictions_for_metric.append(prediction_for_metric)

        # noinspection PyTypeChecker
        return EvalPredictionWithProcessing(
            label_ids=references,
            predictions=predictions,
            processed_predictions=predictions_for_metric
        )

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.data_models.model_outputs.extractive import ExtractiveQAModelOutput
from oneqa.mrc.data_models.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

oneqa/mrc/metrics/tydi_f1/tydi_f1.py METASEP
from typing import Dict, Any

import datasets

from oneqa.mrc.metrics.tydi_f1.eval_utils import Span, TyDiLabel
from oneqa.mrc.metrics.tydi_f1.tydi_eval import pretty_print
from oneqa.mrc.data_models.target_type import TargetType

# TODO tydi f1 docs
_DESCRIPTION = """
TODO
The F1 score is the harmonic mean of the precision and recall. It can be computed with:
F1 = 2 * (precision * recall) / (precision + recall).  This implementation of F1 is based
on the Natural Questions (NQ) leaderboard and does not award partial credit unlike
the SQuAD-style F1.

Adapted from https://ai.google.com/research/NaturalQuestions's `nq_eval` script.
"""

_KWARGS_DESCRIPTION = """
Args:
    predictions: Predicted labels.
    references: Ground truth labels.
Returns: TODO
    minmal_f1: minimal answer F1 score.
    passage_f1: passage F1 score.
Examples:  TODO
    >>> f1_metric = datasets.load_metric("f1")
    >>> results = f1_metric.compute(predictions=[0, 1], references=[0, 1])
    >>> print(results)
    {'f1': 1.0}
    >>> predictions = [0, 2, 1, 0, 0, 1]
    >>> references = [0, 1, 2, 0, 1, 2]
    >>> results = f1_metric.compute(predictions=predictions, references=references, average="macro")
    >>> print(results)
    {'f1': 0.26666666666666666}
    >>> results = f1_metric.compute(predictions=predictions, references=references, average="micro")
    >>> print(results)
    {'f1': 0.3333333333333333}
    >>> results = f1_metric.compute(predictions=predictions, references=references, average="weighted")
    >>> print(results)
    {'f1': 0.26666666666666666}
    >>> results = f1_metric.compute(predictions=predictions, references=references, average=None)
    >>> print(results)
    {'f1': array([0.8, 0. , 0. ])}
"""

_CITATION = """\
TODO
"""


@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
class TyDiF1(datasets.Metric):
    _common_answer_schema = dict(
        start_position=datasets.Value("int32"),
        end_position=datasets.Value("int32"),
        passage_index=datasets.Value("int32"),
        yes_no_answer=datasets.Value("int32"),
        example_id=datasets.Value("string"),
    )
    _pred_answer_schema = dict(
        confidence_score=datasets.Value("float32"),
    )
    _ref_answer_schema = dict(
        language=datasets.Value("string"),
    )

    def _info(self):
        return datasets.MetricInfo(
            description=_DESCRIPTION,
            citation=_CITATION,
            inputs_description=_KWARGS_DESCRIPTION,
            features=datasets.Features(
                dict(
                    predictions={**self._common_answer_schema, **self._pred_answer_schema},
                    references=datasets.Sequence(feature={**self._common_answer_schema, **self._ref_answer_schema})
                )),
            reference_urls=["https://github.com/google-research-datasets/natural-questions/blob/master/nq_eval.py"],
        )

    def _compute(self, *, predictions=None, references=None, **kwargs) -> Dict[str, Any]:
        if kwargs:
            raise ValueError(f"Unexpected kwargs: {kwargs}")
        elif not predictions:
            raise ValueError("No predictions provided")
        elif not references:
            raise ValueError("No references provided")

        predictions = dict(map(self._convert_pred_to_entry, predictions))
        references = dict(map(self._convert_ref_to_entry, references))

        metrics = pretty_print(references, predictions)
        return metrics

    def _convert_ref_to_entry(self, ref):
        if not all(ref['example_id'][0] == ref['example_id'][i] for i in range(len(ref['example_id']))):
            raise ValueError("Found mismatched examples")
        elif not all(ref['language'][0] == ref['language'][i] for i in range(len(ref['language']))):
            raise ValueError("Found mismatched languages")

        key = ref['example_id'][0]
        value = [
            TyDiLabel(
                example_id=ref['example_id'][i],
                passage_answer_index=ref['passage_index'][i],
                minimal_answer_span=Span(
                    ref['start_position'][i],
                    ref['end_position'][i])
                ,
                yes_no_answer=self._bool_target(
                    TargetType(ref['yes_no_answer'][i])
                ),
                passage_score=0,
                minimal_score=0,
                language=ref['language'][i],
                passage_span=None,
                question_text='',
                plaintext='',
            ) for i in range(len(ref['passage_index']))
        ]
        return key, value

    def _convert_pred_to_entry(self, pred):
        key = pred['example_id']
        value = TyDiLabel(
                example_id=pred['example_id'],
                passage_answer_index=pred['passage_index'],
                minimal_answer_span=Span(
                    pred['start_position'],
                    pred['end_position'])
                ,
                yes_no_answer=self._bool_target(
                    TargetType(pred['yes_no_answer'])
                ),
                passage_score=pred['confidence_score'] ,
                minimal_score=pred['confidence_score'] ,
                language=None,
                passage_span=None,
                question_text='',
                plaintext='',
            )
        return key, value

    def _bool_target(self, target_type: TargetType) -> str:
        if target_type == TargetType.YES:
            return 'yes'
        elif target_type == TargetType.NO:
            return 'no'
        elif target_type == TargetType.NO_ANSWER:
            return 'none'
        else:
            raise NotImplementedError(f"Unexpected target type for tydi bool string conversion: {target_type}")



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_, itemgetter

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, minimal_offsets_per_passage=False):
    """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_stats = score_passage_answer(gold, pred, passage_non_null_threshold)
        minimal_stats = score_minimal_answer(gold, pred, minimal_non_null_threshold)

        # fix stats for predictions in incorrect passages
        if minimal_offsets_per_passage and not passage_stats[2] and minimal_stats[1]:
            minimal_stats = (minimal_stats[0], minimal_stats[1], (0., 0., 0.), minimal_stats[3])

        passage_answer_stats.append(passage_stats)
        minimal_answer_stats.append(minimal_stats)

        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=itemgetter(-1), reverse=True)
    minimal_answer_stats.sort(key=itemgetter(-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():
#         count += 1
#         gold_has_answer = eval_utils.gold_has_minimal_answer(
#             ans, minimal_non_null_threshold)
#         total_ans_count += gold_has_answer
#
#     logging.info('%d examples have minimal answers', total_ans_count)
#     logging.info('*' * 40)
#     tydi_pred_dict = dict(predictions=nq_eval_utils.read_prediction_json_from_file(predictions_path))
#     pretty_print(tydi_gold_dict, tydi_pred_dict, passage_non_null_threshold,
#                  minimal_non_null_threshold, verbose, skip_missing_ids)
#
#
# if __name__ == '__main__':
#     main()

oneqa/mrc/metrics/tydi_f1/eval_utils.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.
"""Utility function for TyDi QA evaluation."""

import collections

# import gzip
# import glob
# import json
# import logging
# import multiprocessing


# A data structure for storing prediction and annotation.
# When a example has multiple annotations, multiple TyDiLabel will be used.
TyDiLabel = collections.namedtuple(
    'TyDiLabel',
    [
        'plaintext',  # context.
        'question_text',  # a question text.
        'example_id',  # the unique id for each TyDi example.
        'language',  # language id.
        'passage_answer_index',  # A index for passage answer among candidates.
        'passage_span',  # offsets for the passage span
        'minimal_answer_span',  # A Span object for minimal answer.
        'yes_no_answer',  # Indicate if the minimal answer is an yes/no answer
        #   The possible values are "yes", "no", "none".
        #   (case insensitive)
        #   If the field is "yes", minimal_answer_span should be empty or null.
        'passage_score',  # The score for the passage answer prediction.
        'minimal_score'  # The score for the minimal answer prediction.
    ])


class Span(object):
    """A class for handling token and byte spans.

      The logic is:

      1) if both start_byte != -1 and end_byte != -1 then the span is defined
         by byte offsets
      3) else, this is a null span.

      Null spans means that there is no (passage or minimal) answers.

    """

    def __init__(self, start_byte_offset, end_byte_offset):

        if ((start_byte_offset < 0 and end_byte_offset >= 0) or
                (start_byte_offset >= 0 and end_byte_offset < 0)):
            raise ValueError('Inconsistent Null Spans (Byte).')

        if (start_byte_offset >= 0 and end_byte_offset >= 0 and
                start_byte_offset > end_byte_offset):
            raise ValueError('Invalid byte spans (start_byte >= end_byte).')

        self.start_byte_offset = start_byte_offset
        self.end_byte_offset = end_byte_offset

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

        This can happen for both gold and predicted values and
        for both passage answers and minimal answers.

        Returns:
          boolean flag whether it is null span or not.
        """

        if (self.start_byte_offset < 0 and self.end_byte_offset < 0):
            return True
        return False

    def __str__(self):
        return '({},{})'.format(self.start_byte_offset, self.end_byte_offset)

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


def safe_divide(x, y):
    """Compute x / y, but return 0 if y is zero."""
    if y == 0:
        return 0
    else:
        return x / y


def safe_average(elements):
    """Computes average `elements`, but returns 0 if `elements` is empty."""
    return safe_divide(sum(elements), len(elements))


def compute_partial_match_scores(gold_span, pred_span):
    """Compute byte indices precision, recall and F1 score between span a and b.

    This is used for scoring only minimal answers. See `nonnull_span_equal` for
    scoring passage answers.

    Args:
      gold_span: a Span object. End_byte is inclusive (start_byte+byte_len)
      pred_span: a Span object.  Only compare non-null spans.
        Then, if the bytes are ot negative, compare byte offsets.

    Returns:
      precision: byte offset based precision.
                (# bytes in both gold and pred span) / (# bytes in pred_span)
      recall: byte offset based recall.
                (# bytes in both gold and pred span) / (# bytes in gold_span)
      f1: harmonic mean of precision and recall.
    """
    if not isinstance(gold_span, Span):
        raise TypeError('Gold span must has a Span type.')
    # if not isinstance(pred_span, NQSpan):
    #     raise TypeError('Prediction span must has a Span type.')
    if not isinstance(pred_span, Span):
        raise TypeError('Prediction span must has a Span type.')
    if gold_span.is_null_span():
        raise ValueError(
            'Null gold span should not be passed for F1 computation.')
    if pred_span.is_null_span():
        raise ValueError(
            'Null prediction span should not be passed for F1 computation.')
    assert not pred_span.is_null_span()
    # If there is no overlap, partial score is zero.
    if ((gold_span.end_byte_offset <= pred_span.start_byte_offset) or
            (pred_span.end_byte_offset <= gold_span.start_byte_offset)):
        precision = 0.0
        recall = 0.0

    else:
        in_both = (min(gold_span.end_byte_offset, pred_span.end_byte_offset) -
                   max(gold_span.start_byte_offset, pred_span.start_byte_offset))
        assert in_both > 0
        # if gold span starts earlier than pred span.
        if gold_span.start_byte_offset <= pred_span.start_byte_offset:
            only_in_gold = pred_span.start_byte_offset - gold_span.start_byte_offset
            only_in_gold += max(0,
                                gold_span.end_byte_offset - pred_span.end_byte_offset)
            only_in_pred = max(pred_span.end_byte_offset - gold_span.end_byte_offset,
                               0)
        # if pred span starts earlier than gold span.
        else:
            only_in_pred = gold_span.start_byte_offset - pred_span.start_byte_offset
            only_in_pred += max(0,
                                pred_span.end_byte_offset - gold_span.end_byte_offset)
            only_in_gold = max(gold_span.end_byte_offset - pred_span.end_byte_offset,
                               0)
        precision = safe_divide(in_both, (in_both + only_in_pred))
        recall = safe_divide(in_both, (in_both + only_in_gold))

    f1 = safe_divide(2 * precision * recall, precision + recall)
    return precision, recall, f1


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

    This is used for scoring only passage answers.
    See `compute_partial_match_scores` for minimal answers.

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

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

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

        if ((span_a.start_byte_offset == span_b.start_byte_offset) and
                (span_a.end_byte_offset == span_b.end_byte_offset)):
            return True

    return False


def gold_has_minimal_answer(gold_label_list, minimal_non_null_threshold):
    """Gets vote from annotators for judging if there is a minimal answer."""
    #  We consider if there is a minimal answer if there is an minimal answer span
    #  or the yes/no answer is not none.
    gold_has_answer = gold_label_list and sum([
        ((not label.minimal_answer_span.is_null_span()) or
         (label.yes_no_answer != 'none')) for label in gold_label_list
    ]) >= minimal_non_null_threshold

    return bool(gold_has_answer)


def gold_has_passage_answer(gold_label_list, passage_non_null_threshold):
    """Gets vote from annotators for judging if there is a passage answer."""

    gold_has_answer = gold_label_list and (sum([
        label.passage_answer_index >= 0  # passage answer not null
        for label in gold_label_list  # for each annotator
    ]) >= passage_non_null_threshold)

    return bool(gold_has_answer)

# def read_prediction_jsonl(predictions_path):
#     """Read the prediction jsonl file with scores.
#
#     Args:
#       predictions_path: the path for the jsonl prediction file.
#
#     Returns:
#       A dictionary with key = example_id, value = TyDiLabel.
#
#     """
#     logging.info('Reading predictions from file: %s', format(predictions_path))
#     predictions = []
#     with open(predictions_path, 'r') as f:
#         for line in f:
#             predictions.append(json.loads(line))
#
#     tydi_pred_dict = {}
#     for single_prediction in predictions:
#
#         if 'passage_answer_index' in single_prediction:
#             passage_answer_index = single_prediction['passage_answer_index']
#         else:
#             passage_answer_index = -1
#
#         if 'minimal_answer' in single_prediction:
#             minimal_span = Span(
#                 single_prediction['minimal_answer']['start_byte_offset'],
#                 single_prediction['minimal_answer']['end_byte_offset'])
#
#         yes_no_answer = 'none'
#         if 'yes_no_answer' in single_prediction:
#             yes_no_answer = single_prediction['yes_no_answer'].lower()
#             if yes_no_answer not in ['yes', 'no', 'none']:
#                 raise ValueError('Invalid yes_no_answer value in prediction')
#
#             if yes_no_answer != 'none' and not minimal_span.is_null_span():
#                 raise ValueError(
#                     'yes/no prediction and minimal answers cannot coexist.')
#
#         pred_item = TyDiLabel(
#             example_id=single_prediction['example_id'],
#             language=single_prediction['language'],
#             passage_answer_index=passage_answer_index,
#             minimal_answer_span=minimal_span,
#             yes_no_answer=yes_no_answer,
#             question_text='',
#             plaintext='',
#             passage_score=float(single_prediction['passage_answer_score']),
#             minimal_score=float(single_prediction['minimal_answer_score']))
#
#         tydi_pred_dict[single_prediction['example_id']] = pred_item
#
#     return tydi_pred_dict


# def read_annotation_from_file(input_file):
#     """Read annotation from file."""
#     annotation_dict = {}
#     for line in input_file:
#         json_example = json.loads(line)
#         example_id = json_example['example_id']
#
#         # There are three annotations for each TyDi QA evaluation example.
#         annotation_list = []
#         raw_annotations = json_example.get('annotations', [])
#         if not raw_annotations:
#             logging.warning("No annotations found (may be test data) for example {} -- creating label without annotations".format(example_id))
#             #This is a work-around for running predictions when there are no annotations
#             #We only need the document_text.  TODO: find better way to handle this
#             gold_label = TyDiLabel(
#                 example_id=example_id,
#                 language=json_example['language'],
#                 passage_answer_index=None,
#                 passage_span=None,
#                 minimal_answer_span=None,
#                 passage_score=0,
#                 minimal_score=0,
#                 question_text=json_example['question_text'],
#                 plaintext=json_example['document_plaintext'],
#                 yes_no_answer=None)
#             annotation_list.append(gold_label)
#             annotation_dict[example_id] = annotation_list
#             continue
#         else:
#             logging.debug("Processing {} annotations for example {}".format(len(raw_annotations), example_id))
#
#         for annotation in raw_annotations:
#             passage_span_ind = annotation['passage_answer']['candidate_index']
#             if passage_span_ind == -1:
#                 passage_span = None
#             else:
#                 passage_span_par_offsets = json_example['passage_answer_candidates'][passage_span_ind]
#                 passage_span = Span(passage_span_par_offsets['plaintext_start_byte'],
#                                     passage_span_par_offsets['plaintext_end_byte'])
#             minimal_span_dict = annotation['minimal_answer']
#             minimal_span = Span(minimal_span_dict['plaintext_start_byte'],
#                                 minimal_span_dict['plaintext_end_byte'])
#
#             gold_label = TyDiLabel(
#                 example_id=example_id,
#                 language=json_example['language'],
#                 passage_answer_index=passage_span_ind,
#                 passage_span=passage_span,
#                 minimal_answer_span=minimal_span,
#                 passage_score=0,
#                 minimal_score=0,
#                 question_text=json_example['question_text'],
#                 plaintext=json_example['document_plaintext'],
#                 yes_no_answer=annotation['yes_no_answer'].lower())
#
#             annotation_list.append(gold_label)
#         annotation_dict[example_id] = annotation_list
#
#     return annotation_dict


# def read_annotation(path_name, n_threads=10):
#     """Read annotations with real multiple processes."""
#     input_paths = glob.glob(path_name)
#     pool = multiprocessing.Pool(n_threads)
#     logging.debug("Reading annotation from: {}".format(path_name))
#     try:
#         dict_list = pool.map(read_annotation_from_one_split, input_paths)
#     finally:
#         pool.close()
#         pool.join()
#
#     final_dict = {}
#     for single_dict in dict_list:
#         final_dict.update(single_dict)
#
#     logging.debug("Read annotation (from {}): {}".format(path_name, final_dict))
#
#     return final_dict


# def read_annotation_from_one_split(path_name):
#     """Read annotations from path_name."""
#     if '.gz' in path_name:
#         logging.info('Parsing %s (gzip)...', path_name)
#         with gzip.GzipFile(path_name, 'r') as input_file:
#             return read_annotation_from_file(input_file)
#
#     logging.info('Parsing %s...', path_name)
#     with open(path_name) as input_file:
#         return read_annotation_from_file(input_file)

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

oneqa/mrc/data_models/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/data_models/model_outputs/__init__.py METASEP

tests/oneqa/__init__.py METASEP

examples/tableqg/run_tableqg.py METASEP

examples/re2g/run_re2g.py METASEP

examples/mrc/run_mrc.py METASEP
import logging
import os
import sys
import traceback
from dataclasses import dataclass, field
from importlib import import_module
from operator import attrgetter
from typing import Optional, Type

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

from oneqa.mrc.data_models.eval_prediction_with_processing import EvalPredictionWithProcessing
from oneqa.mrc.metrics.tydi_f1.tydi_f1 import TyDiF1
from oneqa.mrc.models.heads.extractive import EXTRACTIVE_HEAD
from oneqa.mrc.models.task_model import ModelForDownstreamTasks
from oneqa.mrc.processors.postprocessors.extractive import ExtractivePostProcessor
from oneqa.mrc.processors.postprocessors.scorers import SupportedSpanScorers
from oneqa.mrc.processors.preprocessors.tydiqa import TyDiQAPreprocessor
from oneqa.mrc.trainers.default import MRCTrainer


def object_reference(reference_as_str: str) -> object:
    """
    Given a fully qualified path to a class reference, return a pointer to the reference.
    This will work with types, functions, methods, and other objects (e.g. dict).

    Args:
        reference_as_str: the fully qualified path (expects the fully qualified path in dot notation,
                          e.g. oneqa.mrc.processors.postprocessors.extractive.ExtractivePostProcessor).

    Returns:
        reference to path given by input

    Raises:
        TypeError: Unable to resolve input path
    """
    def _split_into_class_and_module_name(class_path):
        modules = class_path.split('.')
        if len(modules) > 1:
            return ".".join(modules[:-1]), modules[-1]
        else:
            return class_path, None

    try:
        module_name, object_name = _split_into_class_and_module_name(reference_as_str)
        module_reference = import_module(module_name)
        if object_name is None:
            return module_reference
        else:
            return getattr(module_reference, object_name)
    except Exception as ex:
        traceback.print_exc()  # Shows additional traceback for why imports fail
        raise TypeError(f"Unable to resolve the string {reference_as_str} to a fully qualified class path") from ex


# modified from
# https://github.com/huggingface/transformers/blob/main/examples/pytorch/question-answering/run_qa.py
@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
    """

    model_name_or_path: str = field(
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
    )
    config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
    )
    tokenizer_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
    )
    cache_dir: Optional[str] = field(
        default=None,
        metadata={"help": "Path to directory to store the pretrained models downloaded from huggingface.co"},
    )


# modified from
# https://github.com/huggingface/transformers/blob/main/examples/pytorch/question-answering/run_qa.py
@dataclass
class DataTrainingArguments:
    """
    Arguments pertaining to what data we are going to input our model for training and eval.
    """

    dataset_name: str = field(
        default="tydiqa", metadata={"help": "The name of the dataset to use (via the datasets library)."}
    )
    dataset_config_name: str = field(
        default="primary_task", metadata={
            "help": "The configuration name of the dataset to use (via the datasets library)."
        }
    )
    overwrite_cache: bool = field(
        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
    )
    preprocessing_num_workers: Optional[int] = field(
        default=None,
        metadata={"help": "The number of processes to use for the preprocessing."},
    )
    max_seq_length: Optional[int] = field(
        default=None,
        metadata={
            "help": "The maximum total input sequence length after tokenization. Sequences longer "
                    "than this will be truncated, sequences shorter will be padded."
        },
    )
    max_q_char_len: int = field(
        default=128, metadata={"help": "Max length per question in characters"}
    )
    single_context_multiple_passages: bool = field(
        default=False, metadata={
            "help": "Allow multiple passages in the same input feature. "
                    "For an example with question q and context c_{1..n} setting this to True"
                    "will allow q|c_{i}c_{i+1}; whereas setting this to False enforces q|c_{i} q|c_{i+1}. "
                    "Note that not all datasets/preprocessors support both values of this parameter. "
                    "Some preprocessors may override this value."
            },
    )
    max_contexts: Optional[int] = field(
        default=None, metadata={"help": "Max contexts per consider"}
    )
    max_train_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
                    "value if set."
        },
    )
    max_eval_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
                    "value if set."
        },
    )
    doc_stride: int = field(
        default=128,
        metadata={"help": "When splitting up a long document into chunks, how much stride to take between chunks."},
    )
    n_best_size: int = field(
        default=20,
        metadata={"help": "The total number of n-best predictions to generate when looking for an answer."},
    )
    n_best_logits: int = field(
        default=20,
        metadata={"help": "The number of logits to consider when searching for start and end position of an answer"}
    )
    max_answer_length: int = field(
        default=32,
        metadata={
            "help": "The maximum length of an answer that can be generated. This is needed because the start "
                    "and end predictions are not conditioned on one another."
        },
    )
    negative_sampling_prob_when_has_answer: float = field(
        default=0.01,
        metadata={
            "help": "Only used when preparing training features, not for decoding. "
                    "For an example with question q and context c_{1..n} where ∃ answer a ∈ c"
                    "an input feature span q|c_{i} where a ∉ c_{i} will be kept with this probability."
                    "Otherwise it will be discarded."
        },
    )
    negative_sampling_prob_when_no_answer: float = field(
        default=0.04,
        metadata={
            "help": "Only used when preparing training features, not for decoding. "
                    "For an example with question q and context c_{1..n} where ∄ answer a ∈ c"
                    "an input feature span q|c_{i} will be kept with this probability."
                    "Otherwise it will be discarded."
        },
    )


@dataclass
class TaskArguments:
    """
    Task specific arguments.
    """
    scorer_type: str = field(
        default='weighted_sum_target_type_and_score_diff',
        metadata={"help": "The name of the scorer to compute answer score.",
                  "choices": SupportedSpanScorers.get_supported()
                  }
    )
    task_heads: object_reference = field(
        default=None,
        metadata={"help": "The name of the task head to use.",
                  "choices": [EXTRACTIVE_HEAD]
                  }
    )
    preprocessor: object_reference = field(
        default=TyDiQAPreprocessor,
        metadata={"help": "The name of the preprocessor to use.",
                  "choices": [TyDiQAPreprocessor]
                  }
    )
    postprocessor: object_reference = field(
        default=ExtractivePostProcessor,
        metadata={"help": "The name of the postprocessor to use.",
                  "choices": [ExtractivePostProcessor]
                  }
    )
    eval_metrics: object_reference = field(
        default=TyDiF1,
        metadata={"help": "The name of the evaluation metric function.",
                  "choices": [TyDiF1]
                 }
    )

    def __post_init__(self):
        if not self.task_heads:
            self.task_heads = EXTRACTIVE_HEAD  # cannot directly set mutable value as default


def main():
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments, TaskArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, training_args, task_args = \
            parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args, task_args = parser.parse_args_into_dataclasses()

    logger = logging.getLogger(__name__)
    scorer_type = task_args.scorer_type
    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."
            )

    task_heads = task_args.task_heads
    config = AutoConfig.from_pretrained(
        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        use_fast=True,
        config=config,
    )
    model = ModelForDownstreamTasks.from_config(
        config,
        model_args.model_name_or_path,
        task_heads=task_heads,
        cache_dir=model_args.cache_dir,
    )
    model.set_task_head(next(iter(task_heads)))

    # load data
    logger.info('Loading dataset')
    raw_datasets = datasets.load_dataset(
        data_args.dataset_name,
        data_args.dataset_config_name,
        cache_dir=model_args.cache_dir,
    )

    # load preprocessor
    preprocessor_class = task_args.preprocessor
    preprocessor = preprocessor_class(
        stride=data_args.doc_stride,
        tokenizer=tokenizer,
        negative_sampling_prob_when_has_answer=data_args.negative_sampling_prob_when_has_answer,
        negative_sampling_prob_when_no_answer=data_args.negative_sampling_prob_when_no_answer,
        load_from_cache_file=not data_args.overwrite_cache,
        max_seq_len=data_args.max_seq_length,
        num_workers=data_args.preprocessing_num_workers,
        max_q_char_len=data_args.max_q_char_len,
        single_context_multiple_passages=data_args.single_context_multiple_passages,
        max_contexts=data_args.max_contexts,
    )

    # process train data
    if training_args.do_train:
        train_dataset = raw_datasets["train"]
        max_train_samples = data_args.max_train_samples
        if 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))
        # Train Feature Creation
        with training_args.main_process_first(desc="train dataset map pre-processing"):
            _, train_dataset = preprocessor.process_train(train_dataset)

    # process val data
    if training_args.do_eval:
        eval_examples = raw_datasets["validation"]
        max_eval_samples = data_args.max_eval_samples
        if max_eval_samples is not None:
            # We will select sample from whole data if argument is specified
            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, eval_dataset = preprocessor.process_eval(eval_examples)

    # If using mixed precision we pad for efficient hardware acceleration
    using_mixed_precision = any(attrgetter('fp16', 'bf16')(training_args))
    data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=64 if using_mixed_precision else None)

    postprocessor_class = task_args.postprocessor

    # noinspection PyProtectedMember
    postprocessor = postprocessor_class(
        k=data_args.n_best_logits,
        n_best_size=data_args.n_best_size,
        max_answer_length=data_args.max_answer_length,
        scorer_type=SupportedSpanScorers(scorer_type),
        single_context_multiple_passages=preprocessor._single_context_multiple_passages,
    )

    eval_metrics = task_args.eval_metrics()

    def compute_metrics(p: EvalPredictionWithProcessing):
        return eval_metrics.compute(predictions=p.processed_predictions, references=p.label_ids)

    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_references_and_predictions,  # see QATrainer in Huggingface
        compute_metrics=compute_metrics,
    )

    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

    # training
    if training_args.do_train:
        train_result = trainer.train(resume_from_checkpoint=checkpoint)
        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()

    # validation
    if training_args.do_eval:
        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)


if __name__ == '__main__':
    main()

examples/listqa/run_listqa.py METASEP

examples/ir/run_ir.py METASEP

examples/distillation/run_distillation.py METASEP

examples/calibration/run_calibration.py METASEP

examples/boolqa/run_boolqa.py METASEP

docs/make.bat METASEP
@ECHO OFF

pushd %~dp0

REM Command file for Sphinx documentation

if "%SPHINXBUILD%" == "" (
	set SPHINXBUILD=sphinx-build
)
set SOURCEDIR=.
set BUILDDIR=_build

if "%1" == "" goto help

%SPHINXBUILD% >NUL 2>NUL
if errorlevel 9009 (
	echo.
	echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
	echo.installed, then set the SPHINXBUILD environment variable to point
	echo.to the full path of the 'sphinx-build' executable. Alternatively you
	echo.may add the Sphinx directory to PATH.
	echo.
	echo.If you don't have Sphinx installed, grab it from
	echo.https://www.sphinx-doc.org/
	exit /b 1
)

%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
goto end

:help
%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%

:end
popd

docs/index.rst METASEP
Welcome to OneQA's documentation!
=================================

.. autosummary::
   :toctree: _autosummary
   :template: custom-module-template.rst
   :recursive:

   oneqa



Indices and tables
==================

* :ref:`genindex`
* :ref:`modindex`
* :ref:`search`

docs/README.md METASEP
# OneQA Documentation

OneQA uses Sphinx for documentation.  Before continuing make sure you have installed OneQA
with `docs` extras.  For example, from the top level of the project:

```shell
pip install .[docs]
```

The documentation can then be built with:

```shell
cd docs
sphinx-build -b html ./ ./generated/
```

This will generate the documentation in HTML.
docs/Makefile METASEP
# Minimal makefile for Sphinx documentation
#

# You can set these variables from the command line, and also
# from the environment for the first two.
SPHINXOPTS    ?=
SPHINXBUILD   ?= sphinx-build
SOURCEDIR     = .
BUILDDIR      = _build

# Put it first so that "make" without argument is like "make help".
help:
	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

.PHONY: help Makefile

# Catch-all target: route all unknown targets to Sphinx using the new
# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
%: Makefile
	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

oneqa/mrc/trainers/default.py METASEP
import inspect
import json
import logging
import os
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 __init__(self, *args, eval_examples=None, eval_dataset=None, post_process_function=None, **kwargs):
        super().__init__(*args, **kwargs)
        self.eval_examples = eval_examples
        self.eval_dataset = eval_dataset
        self.post_process_function = post_process_function

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

    def evaluate(self, eval_dataset=None, eval_examples=None, ignore_keys=None, metric_key_prefix: str = "eval"):
        eval_dataset = self.eval_dataset if eval_dataset is None else eval_dataset
        eval_dataloader = self.get_eval_dataloader(eval_dataset)
        eval_examples = self.eval_examples if eval_examples is None else eval_examples

        # # Temporarily disable metric computation, we will do it in the loop here.
        compute_metrics = self.compute_metrics
        self.compute_metrics = None
        eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
        try:
            output = eval_loop(
                eval_dataloader,
                description="Evaluation",
                # No point gathering the predictions if there are no metrics, otherwise we defer to
                # self.args.prediction_loss_only
                # gather predictions if running in eval mode
                prediction_loss_only=self.args.prediction_loss_only, #True if compute_metrics is None else None,
                ignore_keys=ignore_keys,
            )
        finally:
            self.compute_metrics = compute_metrics
        

        if self.post_process_function is not None:
            eval_preds = self.post_process_function(eval_examples, eval_dataset, output.predictions)

            # TODO: return eval_preds and metrics, write save function for preds
            with open(os.path.join(self.args.output_dir, 'eval_predictions.json'), 'w') as f:
                json.dump(eval_preds.predictions, f, indent=4)
            with open(os.path.join(self.args.output_dir, 'eval_predictions_processed.json'), 'w') as f:
                json.dump(eval_preds.processed_predictions, f, indent=4)
        if self.post_process_function is not None and self.compute_metrics is not None:
            metrics = self.compute_metrics(eval_preds)

            # Prefix all keys with metric_key_prefix + '_'
            for key in list(metrics.keys()):
                if not key.startswith(f"{metric_key_prefix}_"):
                    metrics[f"{metric_key_prefix}_{key}"] = metrics.pop(key)

            self.log(metrics)
        else:
            metrics = {}

        # if self.args.tpu_metrics_debug or self.args.debug:
        #     # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
        #     xm.master_print(met.metrics_report())

        self.control = self.callback_handler.on_evaluate(self.args, self.state, self.control, metrics)
        return metrics

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

oneqa/mrc/data_models/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/data_models/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/data_models/eval_prediction_with_processing.py METASEP
from typing import NamedTuple, Tuple, Union

import numpy as np

class EvalPredictionWithProcessing(NamedTuple):
    """
    Evaluation output (always contains labels), to be used to compute metrics.

    Parameters:
        predictions (`np.ndarray`): Predictions of the model.
        label_ids (`np.ndarray`): Targets to be matched.
        processed_predictions (`np.ndarray`): Predictions of the model processed for metric use.
    """

    predictions: Union[np.ndarray, Tuple[np.ndarray]]
    label_ids: Union[np.ndarray, Tuple[np.ndarray]]
    processed_predictions: Union[np.ndarray, Tuple[np.ndarray]]
oneqa/mrc/data_models/__init__.py METASEP

tests/__init__.py METASEP

docs/conf.py METASEP
# Configuration file for the Sphinx documentation builder.
#
# This file only contains a selection of the most common options. For a full
# list see the documentation:
# https://www.sphinx-doc.org/en/master/usage/configuration.html

# -- Path setup --------------------------------------------------------------

# If extensions (or modules to document with autodoc) are in another directory,
# add these directories to sys.path here. If the directory is relative to the
# documentation root, use os.path.abspath to make it absolute, like shown here.
#
import os
import sys
sys.path.insert(0, os.path.abspath('../..'))


# -- Project information -----------------------------------------------------

project = 'OneQA'
copyright = '2022, IBM Research AI'
author = 'IBM Research AI'


# -- General configuration ---------------------------------------------------

# Add any Sphinx extension module names here, as strings. They can be
# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
# ones.
extensions = [
    'sphinx.ext.autodoc',
    'sphinx.ext.autosummary',
    'sphinx.ext.duration',
    'sphinx.ext.napoleon',
]

# Enable extensions

autosummary_generate = True  # Turn on sphinx.ext.autosummary

# Add any paths that contain templates here, relative to this directory.
templates_path = ['_templates']

# List of patterns, relative to source directory, that match files and
# directories to ignore when looking for source files.
# This pattern also affects html_static_path and html_extra_path.
exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']


# -- Options for HTML output -------------------------------------------------

# The theme to use for HTML and HTML Help pages.  See the documentation for
# a list of builtin themes.
html_theme = "pydata_sphinx_theme"

# Add any paths that contain custom static files (such as style sheets) here,
# relative to this directory. They are copied after the builtin static files,
# so a file named "default.css" will overwrite the builtin "default.css".
html_static_path = ['_static']
oneqa/tableqg/README.md METASEP

oneqa/re2g/README.md METASEP

oneqa/listqa/README.md METASEP

oneqa/distillation/README.md METASEP

oneqa/calibration/README.md METASEP

oneqa/boolqa/README.md 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 --cov-fail-under=60 tests/

pyproject.toml METASEP
[build-system]
requires = ["setuptools", "wheel"]
build-backend = "setuptools.build_meta"
logo.png METASEP

VERSION METASEP
0.1.1
README.md METASEP
![OneQA Logo](logo.png)

# OneQA
OneQA to rule them all

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


## Getting Started

## Installation

```shell
# cd to project root

# If you want to run on GPU make sure to install torch appropriately

# E.g. for torch 1.11 + CUDA 11.3:
pip install 'torch~=1.11.0' --extra-index-url https://download.pytorch.org/whl/cu113

# Install as editable (-e) or non-editable using pip, with extras (e.g. tests) as desired
# Example installation commands:

# Minimal install (non-editable)
pip install .

# Full install (editable)
pip install -e .[all]
```

Please note that dependencies (specified in [setup.py](./setup.py)) are pinned to provide a stable experience.
When installing from source these can be modified, however this is not officially supported.

## Unit Tests

To run the unit tests you first need to [install OneQA](#Installation).
Make sure to install with the `[tests]` or `[all]` extras from pip.

From there you can run the tests via pytest, for example:
```shell
pytest --cov oneqa --cov-config .coveragerc tests/
```

For more information, see:
- Our [tox.ini](./tox.ini)
- The [pytest](https://docs.pytest.org) and [tox](https://tox.wiki/en/latest/) documentation


Our logo is adapted from [this](https://commons.wikimedia.org/wiki/File:One_Ring_Blender_Render.png)
image from Wikipedia under the
[Creative Commons Attribution-Share Alike 4.0 International](https://en.wikipedia.org/wiki/en:Creative_Commons) 
license.  It is shared under the same license.
MANIFEST.in METASEP
include README.md
include VERSION
include .coveragerc
global-exclude __pycache__
global-exclude *.py[co]
oneqa/tableqg/tableqg.py METASEP

oneqa/tableqg/__init__.py METASEP

oneqa/re2g/re2g.py METASEP

oneqa/re2g/__init__.py METASEP

oneqa/mrc/__init__.py METASEP

oneqa/listqa/listqa.py METASEP

oneqa/listqa/__init__.py METASEP

oneqa/distillation/distillation.py METASEP

oneqa/distillation/__init__.py METASEP

oneqa/calibration/calibration.py METASEP

oneqa/calibration/__init__.py METASEP

oneqa/boolqa/boolqa.py METASEP

oneqa/boolqa/__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~=1.0.1": ["dev"],
    "datasets~=2.0.0": ["install"],
    "packaging~=21.3": ["install"],
    "pydata-sphinx-theme~=0.8.0": ["docs"],
    "pytest~=7.1.1": ["tests"],
    "pytest-cov~=3.0.0": ["tests"],
    "pytest-mock~=3.7.0": ["tests"],
    "pytest-rerunfailures~=10.2": ["tests"],
    "sphinx~=4.4.0": ["docs"],
    "torch~=1.11.0": ["install"],
    "tox~=3.24.5": ["tests"],
    "transformers~=4.17.0": ["install"],
    "tqdm~=4.64.0": ["install"],
    "pyserini": ["install"],
    "jsonlines": ["install"],
    "faiss-cpu": ["install"]
}

extras_names = ["docs", "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/ir/sparse/README.md METASEP


oneqa/ir/sparse/retriever.py METASEP
from pyserini.search import LuceneSearcher
from typing import Optional
import logging
import json
from abc import ABCMeta, abstractmethod

logger = logging.getLogger(__name__)

class BaseRetriever(metaclass=ABCMeta):
    """ 
        Base class for Retriever
    """

    @abstractmethod
    def retrieve(self, query: str, top_k: Optional[int] = 10):
        pass

class PyseriniRetriever(BaseRetriever):
    def __init__(self, index_path: str):
        self.index_path = index_path
        self.searcher = LuceneSearcher(index_path)
        self.top_k = 10
        logger.info(f'Initialized LuceneSearcher index_dir: {self.searcher.index_dir}  num_docs: {self.searcher.num_docs}')

    def retrieve(self, query: str, top_k: Optional[int] = 10):
        """
        Return documents that are most relevant to the query.

        :param query: The query
        :param top_k: How many documents to return per query.
        """

        hits = self.searcher.search(query, top_k)
        passage_hits = []
        for i, hit in enumerate(hits):
            title, text = json.loads(hit.raw)['contents'].split("\t")
            title = title.replace('\n',' ')
            text = text.replace('\n',' ')
            docid = hit.docid
            passage_hit = {
                "rank": i,
                "score": hit.score,
                "passage_id": docid,
                "doc_id": docid,
                "title": title,
                "text": text 
            }
            passage_hits.append(passage_hit)
        return passage_hits


oneqa/ir/sparse/__init__.py METASEP

oneqa/__init__.py METASEP

oneqa/ir/README.md METASEP

oneqa/ir/ir.py METASEP

oneqa/ir/__init__.py METASEP

oneqa/ir/scripts/xortydi/generate_xorqa_negatives.py METASEP

#!/bin/env python
# -*- coding: utf-8 -*-

import argparse
import logging
import json
import jsonlines
import numpy.random
from tqdm import tqdm
from oneqa.ir.sparse.retriever import PyseriniRetriever

numpy.random.seed(1234)
logger=logging.getLogger(__name__)
logging.basicConfig(level = logging.INFO)

from utils import HasAnswerChecker

def handle_args():
    usage='usage'
    parser=argparse.ArgumentParser(usage)
    parser.add_argument('--input_file', '-i', required=False, default='/dccstor/jsmc-nmt-01/qa/xor-tydi/train-multilingual-2021-02-25/1/xorqa_dpr_data_query-L_hard_negative-1/dpr_train_data.json')
    parser.add_argument('--output_file', '-o', required=False, default='/dccstor/bsiyer6/OneQA/xortydi_bm25/xortydi_train_data_ir_negs_poss.json')
    parser.add_argument('--question_translations_dir', '-d', required=False,default='/dccstor/colbert-ir/bsiyer/data/xorqa/trans_data_all_langs')
    parser.add_argument('--index_path', '-u', required=False,default='/dccstor/bsiyer6/OneQA/psgs_w100_index')
    parser.add_argument('--max_num_ir_based_negatives', '-n', type=int, default=300)
    parser.add_argument('--max_num_ir_based_positives', '-p', type=int, default=10)
    parser.add_argument('--do_not_run_ir', '-R', action="store_true", default=False)
    parser.add_argument('--do_not_run_match_in_title', '-T', action="store_true", default=False)

    parser.add_argument('--question_records_file', '-q', default=None)
    parser.add_argument('--languages_list', '-l', nargs='+', default=[])

    args=parser.parse_args()
    return args

def init_question_translations(question_translations_dir):
    question_translations = {}
    # this is very XOR-TyDi QA specific
    langs = set(['ar', 'bn', 'fi', 'ja', 'ko', 'ru', 'te'])

    for lang in langs:
        with open(question_translations_dir + '/' + lang + '-en/en.txt') as f:
            en_content = f.readlines()
        with open(question_translations_dir + '/' + lang + '-en/' + lang + '.txt') as f:
            ne_content = f.readlines()

        assert(len(en_content) == len(ne_content))
        for ne, en in zip(ne_content, en_content):
            question_translations[ne.strip()] = en.strip()

    return question_translations


def run_query(query, searcher, max_retrieved):

    hits = searcher.retrieve(query, max_retrieved)
    retrieved = []
    for hit in hits:
        retrieved.append( (hit['doc_id'],hit['score'], hit['title'], hit['text']) )
    return retrieved

# def run_query(query, index_url, max_retrieved):

#     escapeRules = {'+': r'\+',
#            '-': r'\-',
#            '&': r'\&',
#            '|': r'\|',
#            '!': r'\!',
#            '(': r'\(',
#            ')': r'\)',
#            '{': r'\{',
#            '}': r'\}',
#            '[': r'\[',
#            ']': r'\]',
#            '^': r'\^',
#            '~': r'\~',
#            '*': r'\*',
#            '?': r'\?',
#            ':': r'\:',
#            '"': r'\"',
#            '\\': r'\\;',
#            '/': r'\/',
#            '>': r' ',
#            '<': r' '}

#     def escapedSeq(term):
#         """ Yield the next string based on the
#             next character (either this char
#             or escaped version """
#         for char in term:
#             if char in escapeRules.keys():
#                 yield escapeRules[char]
#             else:
#                 yield char

#     def escapeESArg(term):
#         """ Apply escaping to the passed in query terms
#             escaping special characters like : , etc"""
#         term = term.replace('\\', r'\\')   # escape \ first
#         return "".join([nextStr for nextStr in escapedSeq(term)])


#     question_string_escaped = escapeESArg(query.replace('?', ''))

#     query_data = {
#                         "from" : 0, "size" : max_retrieved,
#                         "query": {
#                             "query_string": {
#                                 "query": question_string_escaped
#                             }
#                         }
#                     }

#     response = requests.post(url=index_url, json=query_data, verify=False)
#     retrieved = []
#     response_text_dict = json.loads(response.text)

#     if 'hits' in response_text_dict:
#         for hit in response_text_dict['hits']['hits']:
#             retrieved.append((hit['_id'], hit['_score'],
#                               (hit['_source']['title'] if 'title' in hit['_source'] else "", hit['_source']['text'])))

#     return retrieved

def read_jsonlines(file_name):
    lines = []
    print("loading examples from {0}".format(file_name))
    with jsonlines.open(file_name) as reader:
        for obj in reader:
            lines.append(obj)
    return lines


def main():
    args=handle_args()

    answer_checker = HasAnswerChecker()

    searcher = PyseriniRetriever(args.index_path)

    question_translations = init_question_translations(args.question_translations_dir)

    qas=json.load(open(args.input_file))

    if args.question_records_file is not None:
infile
#!/bin/env python
# -*- coding: utf-8 -*-

import argparse
import logging
import json
import jsonlines
import numpy.random
from tqdm import tqdm
from oneqa.ir.sparse.retriever import PyseriniRetriever

numpy.random.seed(1234)
logger=logging.getLogger(__name__)
logging.basicConfig(level = logging.INFO)

from utils import HasAnswerChecker

def handle_args():
    usage='usage'
    parser=argparse.ArgumentParser(usage)
    parser.add_argument('--input_file', '-i', required=False, default='/dccstor/jsmc-nmt-01/qa/xor-tydi/train-multilingual-2021-02-25/1/xorqa_dpr_data_query-L_hard_negative-1/dpr_train_data.json')
    parser.add_argument('--output_file', '-o', required=False, default='/dccstor/bsiyer6/OneQA/xortydi_bm25/xortydi_train_data_ir_negs_poss.json')
    parser.add_argument('--question_translations_dir', '-d', required=False,default='/dccstor/colbert-ir/bsiyer/data/xorqa/trans_data_all_langs')
    parser.add_argument('--index_path', '-u', required=False,default='/dccstor/bsiyer6/OneQA/psgs_w100_index')
    parser.add_argument('--max_num_ir_based_negatives', '-n', type=int, default=300)
    parser.add_argument('--max_num_ir_based_positives', '-p', type=int, default=10)
    parser.add_argument('--do_not_run_ir', '-R', action="store_true", default=False)
    parser.add_argument('--do_not_run_match_in_title', '-T', action="store_true", default=False)

    parser.add_argument('--question_records_file', '-q', default=None)
    parser.add_argument('--languages_list', '-l', nargs='+', default=[])

    args=parser.parse_args()
    return args

def init_question_translations(question_translations_dir):
    question_translations = {}
    # this is very XOR-TyDi QA specific
    langs = set(['ar', 'bn', 'fi', 'ja', 'ko', 'ru', 'te'])

    for lang in langs:
        with open(question_translations_dir + '/' + lang + '-en/en.txt') as f:
            en_content = f.readlines()
        with open(question_translations_dir + '/' + lang + '-en/' + lang + '.txt') as f:
            ne_content = f.readlines()

        assert(len(en_content) == len(ne_content))
        for ne, en in zip(ne_content, en_content):
            question_translations[ne.strip()] = en.strip()

    return question_translations


def run_query(query, searcher, max_retrieved):

    hits = searcher.retrieve(query, max_retrieved)
    retrieved = []
    for hit in hits:
        retrieved.append( (hit['doc_id'],hit['score'], hit['title'], hit['text']) )
    return retrieved

# def run_query(query, index_url, max_retrieved):

#     escapeRules = {'+': r'\+',
#            '-': r'\-',
#            '&': r'\&',
#            '|': r'\|',
#            '!': r'\!',
#            '(': r'\(',
#            ')': r'\)',
#            '{': r'\{',
#            '}': r'\}',
#            '[': r'\[',
#            ']': r'\]',
#            '^': r'\^',
#            '~': r'\~',
#            '*': r'\*',
#            '?': r'\?',
#            ':': r'\:',
#            '"': r'\"',
#            '\\': r'\\;',
#            '/': r'\/',
#            '>': r' ',
#            '<': r' '}

#     def escapedSeq(term):
#         """ Yield the next string based on the
#             next character (either this char
#             or escaped version """
#         for char in term:
#             if char in escapeRules.keys():
#                 yield escapeRules[char]
#             else:
#                 yield char

#     def escapeESArg(term):
#         """ Apply escaping to the passed in query terms
#             escaping special characters like : , etc"""
#         term = term.replace('\\', r'\\')   # escape \ first
#         return "".join([nextStr for nextStr in escapedSeq(term)])


#     question_string_escaped = escapeESArg(query.replace('?', ''))

#     query_data = {
#                         "from" : 0, "size" : max_retrieved,
#                         "query": {
#                             "query_string": {
#                                 "query": question_string_escaped
#                             }
#                         }
#                     }

#     response = requests.post(url=index_url, json=query_data, verify=False)
#     retrieved = []
#     response_text_dict = json.loads(response.text)

#     if 'hits' in response_text_dict:
#         for hit in response_text_dict['hits']['hits']:
#             retrieved.append((hit['_id'], hit['_score'],
#                               (hit['_source']['title'] if 'title' in hit['_source'] else "", hit['_source']['text'])))

#     return retrieved

def read_jsonlines(file_name):
    lines = []
    print("loading examples from {0}".format(file_name))
    with jsonlines.open(file_name) as reader:
        for obj in reader:
            lines.append(obj)
    return lines


def main():
    args=handle_args()

    answer_checker = HasAnswerChecker()

    searcher = PyseriniRetriever(args.index_path)

infile
#!/bin/env python
# -*- coding: utf-8 -*-

import argparse
import logging
import json
import jsonlines
import numpy.random
from tqdm import tqdm
from oneqa.ir.sparse.retriever import PyseriniRetriever

numpy.random.seed(1234)
logger=logging.getLogger(__name__)
logging.basicConfig(level = logging.INFO)

from utils import HasAnswerChecker

def handle_args():
    usage='usage'
    parser=argparse.ArgumentParser(usage)
    parser.add_argument('--input_file', '-i', required=False, default='/dccstor/jsmc-nmt-01/qa/xor-tydi/train-multilingual-2021-02-25/1/xorqa_dpr_data_query-L_hard_negative-1/dpr_train_data.json')
    parser.add_argument('--output_file', '-o', required=False, default='/dccstor/bsiyer6/OneQA/xortydi_bm25/xortydi_train_data_ir_negs_poss.json')
    parser.add_argument('--question_translations_dir', '-d', required=False,default='/dccstor/colbert-ir/bsiyer/data/xorqa/trans_data_all_langs')
    parser.add_argument('--index_path', '-u', required=False,default='/dccstor/bsiyer6/OneQA/psgs_w100_index')
    parser.add_argument('--max_num_ir_based_negatives', '-n', type=int, default=300)
    parser.add_argument('--max_num_ir_based_positives', '-p', type=int, default=10)
    parser.add_argument('--do_not_run_ir', '-R', action="store_true", default=False)
    parser.add_argument('--do_not_run_match_in_title', '-T', action="store_true", default=False)

    parser.add_argument('--question_records_file', '-q', default=None)
    parser.add_argument('--languages_list', '-l', nargs='+', default=[])

    args=parser.parse_args()
    return args

def init_question_translations(question_translations_dir):
    question_translations = {}
    # this is very XOR-TyDi QA specific
    langs = set(['ar', 'bn', 'fi', 'ja', 'ko', 'ru', 'te'])

    for lang in langs:
        with open(question_translations_dir + '/' + lang + '-en/en.txt') as f:
            en_content = f.readlines()
        with open(question_translations_dir + '/' + lang + '-en/' + lang + '.txt') as f:
            ne_content = f.readlines()

        assert(len(en_content) == len(ne_content))
        for ne, en in zip(ne_content, en_content):
            question_translations[ne.strip()] = en.strip()

    return question_translations


def run_query(query, searcher, max_retrieved):

    hits = searcher.retrieve(query, max_retrieved)
    retrieved = []
    for hit in hits:
        retrieved.append( (hit['doc_id'],hit['score'], hit['title'], hit['text']) )
    return retrieved

# def run_query(query, index_url, max_retrieved):

#     escapeRules = {'+': r'\+',
#            '-': r'\-',
#            '&': r'\&',
#            '|': r'\|',
#            '!': r'\!',
#            '(': r'\(',
#            ')': r'\)',
#            '{': r'\{',
#            '}': r'\}',
#            '[': r'\[',
#            ']': r'\]',
#            '^': r'\^',
#            '~': r'\~',
#            '*': r'\*',
#            '?': r'\?',
#            ':': r'\:',
#            '"': r'\"',
#            '\\': r'\\;',
#            '/': r'\/',
#            '>': r' ',
#            '<': r' '}

#     def escapedSeq(term):
#         """ Yield the next string based on the
#             next character (either this char
#             or escaped version """
#         for char in term:
#             if char in escapeRules.keys():
#                 yield escapeRules[char]
#             else:
#                 yield char

#     def escapeESArg(term):
#         """ Apply escaping to the passed in query terms
#             escaping special characters like : , etc"""
#         term = term.replace('\\', r'\\')   # escape \ first
#         return "".join([nextStr for nextStr in escapedSeq(term)])


#     question_string_escaped = escapeESArg(query.replace('?', ''))

#     query_data = {
#                         "from" : 0, "size" : max_retrieved,
#                         "query": {
#                             "query_string": {
#                                 "query": question_string_escaped
#                             }
#                         }
#                     }

#     response = requests.post(url=index_url, json=query_data, verify=False)
#     retrieved = []
#     response_text_dict = json.loads(response.text)

#     if 'hits' in response_text_dict:
#         for hit in response_text_dict['hits']['hits']:
#             retrieved.append((hit['_id'], hit['_score'],
#                               (hit['_source']['title'] if 'title' in hit['_source'] else "", hit['_source']['text'])))

#     return retrieved

def read_jsonlines(file_name):
    lines = []
    print("loading examples from {0}".format(file_name))
    with jsonlines.open(file_name) as reader:
        for obj in reader:
            lines.append(obj)
    return lines


def main():
infile
#!/bin/env python
# -*- coding: utf-8 -*-

import argparse
import logging
import json
import jsonlines
import numpy.random
from tqdm import tqdm
from oneqa.ir.sparse.retriever import PyseriniRetriever

numpy.random.seed(1234)
logger=logging.getLogger(__name__)
logging.basicConfig(level = logging.INFO)

from utils import HasAnswerChecker

def handle_args():
    usage='usage'
    parser=argparse.ArgumentParser(usage)
    parser.add_argument('--input_file', '-i', required=False, default='/dccstor/jsmc-nmt-01/qa/xor-tydi/train-multilingual-2021-02-25/1/xorqa_dpr_data_query-L_hard_negative-1/dpr_train_data.json')
    parser.add_argument('--output_file', '-o', required=False, default='/dccstor/bsiyer6/OneQA/xortydi_bm25/xortydi_train_data_ir_negs_poss.json')
    parser.add_argument('--question_translations_dir', '-d', required=False,default='/dccstor/colbert-ir/bsiyer/data/xorqa/trans_data_all_langs')
    parser.add_argument('--index_path', '-u', required=False,default='/dccstor/bsiyer6/OneQA/psgs_w100_index')
    parser.add_argument('--max_num_ir_based_negatives', '-n', type=int, default=300)
    parser.add_argument('--max_num_ir_based_positives', '-p', type=int, default=10)
    parser.add_argument('--do_not_run_ir', '-R', action="store_true", default=False)
    parser.add_argument('--do_not_run_match_in_title', '-T', action="store_true", default=False)

    parser.add_argument('--question_records_file', '-q', default=None)
    parser.add_argument('--languages_list', '-l', nargs='+', default=[])

    args=parser.parse_args()
    return args

def init_question_translations(question_translations_dir):
    question_translations = {}
    # this is very XOR-TyDi QA specific
    langs = set(['ar', 'bn', 'fi', 'ja', 'ko', 'ru', 'te'])

    for lang in langs:
        with open(question_translations_dir + '/' + lang + '-en/en.txt') as f:
            en_content = f.readlines()
        with open(question_translations_dir + '/' + lang + '-en/' + lang + '.txt') as f:
            ne_content = f.readlines()

        assert(len(en_content) == len(ne_content))
        for ne, en in zip(ne_content, en_content):
            question_translations[ne.strip()] = en.strip()

    return question_translations


def run_query(query, searcher, max_retrieved):

    hits = searcher.retrieve(query, max_retrieved)
    retrieved = []
    for hit in hits:
        retrieved.append( (hit['doc_id'],hit['score'], hit['title'], hit['text']) )
    return retrieved

# def run_query(query, index_url, max_retrieved):

#     escapeRules = {'+': r'\+',
#            '-': r'\-',
#            '&': r'\&',
#            '|': r'\|',
#            '!': r'\!',
#            '(': r'\(',
#            ')': r'\)',
#            '{': r'\{',
#            '}': r'\}',
#            '[': r'\[',
#            ']': r'\]',
#            '^': r'\^',
#            '~': r'\~',
#            '*': r'\*',
#            '?': r'\?',
#            ':': r'\:',
#            '"': r'\"',
#            '\\': r'\\;',
#            '/': r'\/',
#            '>': r' ',
#            '<': r' '}

#     def escapedSeq(term):
#         """ Yield the next string based on the
#             next character (either this char
#             or escaped version """
#         for char in term:
#             if char in escapeRules.keys():
#                 yield escapeRules[char]
#             else:
#                 yield char

#     def escapeESArg(term):
#         """ Apply escaping to the passed in query terms
#             escaping special characters like : , etc"""
#         term = term.replace('\\', r'\\')   # escape \ first
#         return "".join([nextStr for nextStr in escapedSeq(term)])


#     question_string_escaped = escapeESArg(query.replace('?', ''))

#     query_data = {
#                         "from" : 0, "size" : max_retrieved,
#                         "query": {
#                             "query_string": {
#                                 "query": question_string_escaped
#                             }
#                         }
#                     }

#     response = requests.post(url=index_url, json=query_data, verify=False)
#     retrieved = []
#     response_text_dict = json.loads(response.text)

#     if 'hits' in response_text_dict:
#         for hit in response_text_dict['hits']['hits']:
#             retrieved.append((hit['_id'], hit['_score'],
#                               (hit['_source']['title'] if 'title' in hit['_source'] else "", hit['_source']['text'])))

#     return retrieved

def read_jsonlines(file_name):
    lines = []
    print("loading examples from {0}".format(file_name))
    with jsonlines.open(file_name) as reader:
        for obj in reader:
            lines.append(obj)
    return lines


def main():
    args=handle_args()

    answer_checker = HasAnswerChecker()

inproject
#!/bin/env python
# -*- coding: utf-8 -*-

import argparse
import logging
import json
import jsonlines
import numpy.random
from tqdm import tqdm
from oneqa.ir.sparse.retriever import PyseriniRetriever

numpy.random.seed(1234)
logger=logging.getLogger(__name__)
logging.basicConfig(level = logging.INFO)

from utils import HasAnswerChecker

def handle_args():
    usage='usage'
    parser=argparse.ArgumentParser(usage)
    parser.add_argument('--input_file', '-i', required=False, default='/dccstor/jsmc-nmt-01/qa/xor-tydi/train-multilingual-2021-02-25/1/xorqa_dpr_data_query-L_hard_negative-1/dpr_train_data.json')
    parser.add_argument('--output_file', '-o', required=False, default='/dccstor/bsiyer6/OneQA/xortydi_bm25/xortydi_train_data_ir_negs_poss.json')
    parser.add_argument('--question_translations_dir', '-d', required=False,default='/dccstor/colbert-ir/bsiyer/data/xorqa/trans_data_all_langs')
    parser.add_argument('--index_path', '-u', required=False,default='/dccstor/bsiyer6/OneQA/psgs_w100_index')
    parser.add_argument('--max_num_ir_based_negatives', '-n', type=int, default=300)
    parser.add_argument('--max_num_ir_based_positives', '-p', type=int, default=10)
    parser.add_argument('--do_not_run_ir', '-R', action="store_true", default=False)
    parser.add_argument('--do_not_run_match_in_title', '-T', action="store_true", default=False)

    parser.add_argument('--question_records_file', '-q', default=None)
    parser.add_argument('--languages_list', '-l', nargs='+', default=[])

    args=parser.parse_args()
    return args

def init_question_translations(question_translations_dir):
    question_translations = {}
    # this is very XOR-TyDi QA specific
    langs = set(['ar', 'bn', 'fi', 'ja', 'ko', 'ru', 'te'])

    for lang in langs:
        with open(question_translations_dir + '/' + lang + '-en/en.txt') as f:
            en_content = f.readlines()
        with open(question_translations_dir + '/' + lang + '-en/' + lang + '.txt') as f:
            ne_content = f.readlines()

        assert(len(en_content) == len(ne_content))
        for ne, en in zip(ne_content, en_content):
            question_translations[ne.strip()] = en.strip()

    return question_translations


def run_query(query, searcher, max_retrieved):

inproject
#!/bin/env python
# -*- coding: utf-8 -*-

import argparse
import logging
import json
import jsonlines
import numpy.random
from tqdm import tqdm
from oneqa.ir.sparse.retriever import PyseriniRetriever

numpy.random.seed(1234)
logger=logging.getLogger(__name__)
logging.basicConfig(level = logging.INFO)

from utils import HasAnswerChecker

def handle_args():
    usage='usage'
    parser=argparse.ArgumentParser(usage)
    parser.add_argument('--input_file', '-i', required=False, default='/dccstor/jsmc-nmt-01/qa/xor-tydi/train-multilingual-2021-02-25/1/xorqa_dpr_data_query-L_hard_negative-1/dpr_train_data.json')
    parser.add_argument('--output_file', '-o', required=False, default='/dccstor/bsiyer6/OneQA/xortydi_bm25/xortydi_train_data_ir_negs_poss.json')
    parser.add_argument('--question_translations_dir', '-d', required=False,default='/dccstor/colbert-ir/bsiyer/data/xorqa/trans_data_all_langs')
    parser.add_argument('--index_path', '-u', required=False,default='/dccstor/bsiyer6/OneQA/psgs_w100_index')
    parser.add_argument('--max_num_ir_based_negatives', '-n', type=int, default=300)
    parser.add_argument('--max_num_ir_based_positives', '-p', type=int, default=10)
    parser.add_argument('--do_not_run_ir', '-R', action="store_true", default=False)
    parser.add_argument('--do_not_run_match_in_title', '-T', action="store_true", default=False)

    parser.add_argument('--question_records_file', '-q', default=None)
    parser.add_argument('--languages_list', '-l', nargs='+', default=[])

    args=parser.parse_args()
    return args

def init_question_translations(question_translations_dir):
    question_translations = {}
    # this is very XOR-TyDi QA specific
    langs = set(['ar', 'bn', 'fi', 'ja', 'ko', 'ru', 'te'])

    for lang in langs:
        with open(question_translations_dir + '/' + lang + '-en/en.txt') as f:
            en_content = f.readlines()
        with open(question_translations_dir + '/' + lang + '-en/' + lang + '.txt') as f:
            ne_content = f.readlines()

        assert(len(en_content) == len(ne_content))
        for ne, en in zip(ne_content, en_content):
            question_translations[ne.strip()] = en.strip()

    return question_translations


def run_query(query, searcher, max_retrieved):

    hits = searcher.retrieve(query, max_retrieved)
    retrieved = []
    for hit in hits:
        retrieved.append( (hit['doc_id'],hit['score'], hit['title'], hit['text']) )
    return retrieved

# def run_query(query, index_url, max_retrieved):

#     escapeRules = {'+': r'\+',
#            '-': r'\-',
#            '&': r'\&',
#            '|': r'\|',
#            '!': r'\!',
#            '(': r'\(',
#            ')': r'\)',
#            '{': r'\{',
#            '}': r'\}',
#            '[': r'\[',
#            ']': r'\]',
#            '^': r'\^',
#            '~': r'\~',
#            '*': r'\*',
#            '?': r'\?',
#            ':': r'\:',
#            '"': r'\"',
#            '\\': r'\\;',
#            '/': r'\/',
#            '>': r' ',
#            '<': r' '}

#     def escapedSeq(term):
#         """ Yield the next string based on the
#             next character (either this char
#             or escaped version """
#         for char in term:
#             if char in escapeRules.keys():
#                 yield escapeRules[char]
#             else:
#                 yield char

#     def escapeESArg(term):
#         """ Apply escaping to the passed in query terms
#             escaping special characters like : , etc"""
#         term = term.replace('\\', r'\\')   # escape \ first
#         return "".join([nextStr for nextStr in escapedSeq(term)])


#     question_string_escaped = escapeESArg(query.replace('?', ''))

#     query_data = {
#                         "from" : 0, "size" : max_retrieved,
#                         "query": {
#                             "query_string": {
#                                 "query": question_string_escaped
#                             }
#                         }
#                     }

#     response = requests.post(url=index_url, json=query_data, verify=False)
#     retrieved = []
#     response_text_dict = json.loads(response.text)

#     if 'hits' in response_text_dict:
#         for hit in response_text_dict['hits']['hits']:
#             retrieved.append((hit['_id'], hit['_score'],
#                               (hit['_source']['title'] if 'title' in hit['_source'] else "", hit['_source']['text'])))

#     return retrieved

def read_jsonlines(file_name):
    lines = []
    print("loading examples from {0}".format(file_name))
    with jsonlines.open(file_name) as reader:
        for obj in reader:
            lines.append(obj)
    return lines


def main():
    args=handle_args()

    answer_checker = HasAnswerChecker()

    searcher = PyseriniRetriever(args.index_path)

    question_translations = init_question_translations(args.question_translations_dir)

common
#!/bin/env python
# -*- coding: utf-8 -*-

import argparse
import logging
import json
import jsonlines
import numpy.random
from tqdm import tqdm
from oneqa.ir.sparse.retriever import PyseriniRetriever

numpy.random.seed(1234)
logger=logging.getLogger(__name__)
logging.basicConfig(level = logging.INFO)

from utils import HasAnswerChecker

def handle_args():
    usage='usage'
    parser=argparse.ArgumentParser(usage)
    parser.add_argument('--input_file', '-i', required=False, default='/dccstor/jsmc-nmt-01/qa/xor-tydi/train-multilingual-2021-02-25/1/xorqa_dpr_data_query-L_hard_negative-1/dpr_train_data.json')
    parser.add_argument('--output_file', '-o', required=False, default='/dccstor/bsiyer6/OneQA/xortydi_bm25/xortydi_train_data_ir_negs_poss.json')
    parser.add_argument('--question_translations_dir', '-d', required=False,default='/dccstor/colbert-ir/bsiyer/data/xorqa/trans_data_all_langs')
    parser.add_argument('--index_path', '-u', required=False,default='/dccstor/bsiyer6/OneQA/psgs_w100_index')
    parser.add_argument('--max_num_ir_based_negatives', '-n', type=int, default=300)
    parser.add_argument('--max_num_ir_based_positives', '-p', type=int, default=10)
    parser.add_argument('--do_not_run_ir', '-R', action="store_true", default=False)
    parser.add_argument('--do_not_run_match_in_title', '-T', action="store_true", default=False)

    parser.add_argument('--question_records_file', '-q', default=None)
    parser.add_argument('--languages_list', '-l', nargs='+', default=[])

    args=parser.parse_args()
    return args

def init_question_translations(question_translations_dir):
    question_translations = {}
    # this is very XOR-TyDi QA specific
    langs = set(['ar', 'bn', 'fi', 'ja', 'ko', 'ru', 'te'])

    for lang in langs:
        with open(question_translations_dir + '/' + lang + '-en/en.txt') as f:
            en_content = f.readlines()
        with open(question_translations_dir + '/' + lang + '-en/' + lang + '.txt') as f:
            ne_content = f.readlines()

        assert(len(en_content) == len(ne_content))
        for ne, en in zip(ne_content, en_content):
            question_translations[ne.strip()] = en.strip()

    return question_translations


def run_query(query, searcher, max_retrieved):

    hits = searcher.retrieve(query, max_retrieved)
    retrieved = []
    for hit in hits:
        retrieved.append( (hit['doc_id'],hit['score'], hit['title'], hit['text']) )
    return retrieved

# def run_query(query, index_url, max_retrieved):

#     escapeRules = {'+': r'\+',
#            '-': r'\-',
#            '&': r'\&',
#            '|': r'\|',
#            '!': r'\!',
#            '(': r'\(',
#            ')': r'\)',
#            '{': r'\{',
#            '}': r'\}',
#            '[': r'\[',
#            ']': r'\]',
#            '^': r'\^',
#            '~': r'\~',
#            '*': r'\*',
#            '?': r'\?',
#            ':': r'\:',
#            '"': r'\"',
#            '\\': r'\\;',
#            '/': r'\/',
#            '>': r' ',
#            '<': r' '}

#     def escapedSeq(term):
#         """ Yield the next string based on the
#             next character (either this char
#             or escaped version """
#         for char in term:
#             if char in escapeRules.keys():
#                 yield escapeRules[char]
#             else:
#                 yield char

#     def escapeESArg(term):
#         """ Apply escaping to the passed in query terms
#             escaping special characters like : , etc"""
#         term = term.replace('\\', r'\\')   # escape \ first
#         return "".join([nextStr for nextStr in escapedSeq(term)])


#     question_string_escaped = escapeESArg(query.replace('?', ''))

#     query_data = {
#                         "from" : 0, "size" : max_retrieved,
#                         "query": {
#                             "query_string": {
#                                 "query": question_string_escaped
#                             }
#                         }
#                     }

#     response = requests.post(url=index_url, json=query_data, verify=False)
#     retrieved = []
#     response_text_dict = json.loads(response.text)

#     if 'hits' in response_text_dict:
#         for hit in response_text_dict['hits']['hits']:
#             retrieved.append((hit['_id'], hit['_score'],
#                               (hit['_source']['title'] if 'title' in hit['_source'] else "", hit['_source']['text'])))

#     return retrieved

def read_jsonlines(file_name):
    lines = []
    print("loading examples from {0}".format(file_name))
    with jsonlines.open(file_name) as reader:
        for obj in reader:
            lines.append(obj)
    return lines


def main():
    args=handle_args()

    answer_checker = HasAnswerChecker()

    searcher = PyseriniRetriever(args.index_path)

    question_translations = init_question_translations(args.question_translations_dir)

    qas=json.load(open(args.input_file))

    if args.question_records_file is not None:
        q_recs=read_jsonlines(args.question_records_file)

    languages_to_use = set(args.languages_list)

    for qnum, qa in enumerate(tqdm(qas)):
        if args.question_records_file is not None:
            # sanity check
            assert(qa['answers'] == q_recs[qnum]['answers'])
            if q_recs[qnum]['lang'] not in languages_to_use:
                continue

        q=qa['question']
        p=qa['positive_ctxs']
        n=qa['negative_ctxs']
        hn=qa['hard_negative_ctxs']

        en_q = question_translations[q]

        if args.do_not_run_ir:
            continue

        retrieved = run_query(en_q, searcher, args.max_num_ir_based_negatives * 2) # assuming 50/50 positive/negative ratio

        negatives = []
        positives = []

        positive_ctxs_texts = [ctxt['text'] for ctxt in p]

        for hit in retrieved:
            #string_to_search = hit[2][1] if args.do_not_run_match_in_title else hit[2][0] + ' ' + hit[2][1]
            string_to_search = hit[3] if args.do_not_run_match_in_title else hit[2] + ' ' + hit[3]
            if answer_checker.has_answer(qa['answers'], string_to_search):
                if len(positives) < args.max_num_ir_based_positives and not max([ ctxt in hit[3] for ctxt in positive_ctxs_texts ]):
                    #positives.append({'title': hit[2][0], 'text': hit[2][1]})
                    positives.append({'title': hit[2], 'text': hit[3]})
            else:
                if len(negatives) < args.max_num_ir_based_negatives:
                    # negatives.append({'title': hit[2][0], 'text': hit[2][1]})
                    negatives.append({'title': hit[2], 'text': hit[3]})

            if len(positives) >= args.max_num_ir_based_positives and len(negatives) >= args.max_num_ir_based_negatives:
                break

        qa['ir_negative_ctxs'] = negatives
        qa['ir_positive_ctxs'] = positives

        print('adding ' + str(len(qa['ir_negative_ctxs'])) + ' IR based negatives, ' + str(len(qa['ir_positive_ctxs'])) + ' IR based positives')

        if qnum > 0 and (qnum % 1000) == 0:
            output_file = args.output_file.replace('.json', '_' + str(qnum) + '.json')
            print('writing ', output_file)
            with open(output_file, 'w') as out_f:
                json.dump(qas, out_f, indent=4)

    output_file = args.output_file
    print('writing ', output_file)
    with open(output_file, 'w') as out_f:
        json.dump(qas, out_f, indent=4)


# do main
if __name__=='__main__':
commited
#!/bin/env python
# -*- coding: utf-8 -*-

import argparse
import logging
import json
import jsonlines
import numpy.random
from tqdm import tqdm
from oneqa.ir.sparse.retriever import PyseriniRetriever

numpy.random.seed(1234)
logger=logging.getLogger(__name__)
logging.basicConfig(level = logging.INFO)

from utils import HasAnswerChecker

def handle_args():
    usage='usage'
    parser=argparse.ArgumentParser(usage)
    parser.add_argument('--input_file', '-i', required=False, default='/dccstor/jsmc-nmt-01/qa/xor-tydi/train-multilingual-2021-02-25/1/xorqa_dpr_data_query-L_hard_negative-1/dpr_train_data.json')
    parser.add_argument('--output_file', '-o', required=False, default='/dccstor/bsiyer6/OneQA/xortydi_bm25/xortydi_train_data_ir_negs_poss.json')
    parser.add_argument('--question_translations_dir', '-d', required=False,default='/dccstor/colbert-ir/bsiyer/data/xorqa/trans_data_all_langs')
    parser.add_argument('--index_path', '-u', required=False,default='/dccstor/bsiyer6/OneQA/psgs_w100_index')
    parser.add_argument('--max_num_ir_based_negatives', '-n', type=int, default=300)
    parser.add_argument('--max_num_ir_based_positives', '-p', type=int, default=10)
    parser.add_argument('--do_not_run_ir', '-R', action="store_true", default=False)
    parser.add_argument('--do_not_run_match_in_title', '-T', action="store_true", default=False)

    parser.add_argument('--question_records_file', '-q', default=None)
    parser.add_argument('--languages_list', '-l', nargs='+', default=[])

    args=parser.parse_args()
    return args

def init_question_translations(question_translations_dir):
    question_translations = {}
    # this is very XOR-TyDi QA specific
    langs = set(['ar', 'bn', 'fi', 'ja', 'ko', 'ru', 'te'])

    for lang in langs:
        with open(question_translations_dir + '/' + lang + '-en/en.txt') as f:
            en_content = f.readlines()
        with open(question_translations_dir + '/' + lang + '-en/' + lang + '.txt') as f:
            ne_content = f.readlines()

        assert(len(en_content) == len(ne_content))
        for ne, en in zip(ne_content, en_content):
            question_translations[ne.strip()] = en.strip()

    return question_translations


def run_query(query, searcher, max_retrieved):

    hits = searcher.retrieve(query, max_retrieved)
    retrieved = []
    for hit in hits:
        retrieved.append( (hit['doc_id'],hit['score'], hit['title'], hit['text']) )
    return retrieved

# def run_query(query, index_url, max_retrieved):

#     escapeRules = {'+': r'\+',
#            '-': r'\-',
#            '&': r'\&',
#            '|': r'\|',
#            '!': r'\!',
#            '(': r'\(',
#            ')': r'\)',
#            '{': r'\{',
#            '}': r'\}',
#            '[': r'\[',
#            ']': r'\]',
#            '^': r'\^',
non_informative
#!/bin/env python
# -*- coding: utf-8 -*-

import argparse
import logging
import json
import jsonlines
import numpy.random
from tqdm import tqdm
from oneqa.ir.sparse.retriever import PyseriniRetriever

numpy.random.seed(1234)
logger=logging.getLogger(__name__)
logging.basicConfig(level = logging.INFO)

from utils import HasAnswerChecker

def handle_args():
    usage='usage'
    parser=argparse.ArgumentParser(usage)
    parser.add_argument('--input_file', '-i', required=False, default='/dccstor/jsmc-nmt-01/qa/xor-tydi/train-multilingual-2021-02-25/1/xorqa_dpr_data_query-L_hard_negative-1/dpr_train_data.json')
    parser.add_argument('--output_file', '-o', required=False, default='/dccstor/bsiyer6/OneQA/xortydi_bm25/xortydi_train_data_ir_negs_poss.json')
    parser.add_argument('--question_translations_dir', '-d', required=False,default='/dccstor/colbert-ir/bsiyer/data/xorqa/trans_data_all_langs')
    parser.add_argument('--index_path', '-u', required=False,default='/dccstor/bsiyer6/OneQA/psgs_w100_index')
    parser.add_argument('--max_num_ir_based_negatives', '-n', type=int, default=300)
    parser.add_argument('--max_num_ir_based_positives', '-p', type=int, default=10)
    parser.add_argument('--do_not_run_ir', '-R', action="store_true", default=False)
    parser.add_argument('--do_not_run_match_in_title', '-T', action="store_true", default=False)

    parser.add_argument('--question_records_file', '-q', default=None)
    parser.add_argument('--languages_list', '-l', nargs='+', default=[])

    args=parser.parse_args()
    return args

def init_question_translations(question_translations_dir):
    question_translations = {}
    # this is very XOR-TyDi QA specific
    langs = set(['ar', 'bn', 'fi', 'ja', 'ko', 'ru', 'te'])

    for lang in langs:
        with open(question_translations_dir + '/' + lang + '-en/en.txt') as f:
            en_content = f.readlines()
        with open(question_translations_dir + '/' + lang + '-en/' + lang + '.txt') as f:
            ne_content = f.readlines()

        assert(len(en_content) == len(ne_content))
        for ne, en in zip(ne_content, en_content):
            question_translations[ne.strip()] = en.strip()

    return question_translations


def run_query(query, searcher, max_retrieved):

    hits = searcher.retrieve(query, max_retrieved)
    retrieved = []
    for hit in hits:
        retrieved.append( (hit['doc_id'],hit['score'], hit['title'], hit['text']) )
    return retrieved

# def run_query(query, index_url, max_retrieved):

#     escapeRules = {'+': r'\+',
#            '-': r'\-',
#            '&': r'\&',
#            '|': r'\|',
#            '!': r'\!',
#            '(': r'\(',
#            ')': r'\)',
#            '{': r'\{',
#            '}': r'\}',
#            '[': r'\[',
#            ']': r'\]',
#            '^': r'\^',
#            '~': r'\~',
#            '*': r'\*',
#            '?': r'\?',
#            ':': r'\:',
#            '"': r'\"',
#            '\\': r'\\;',
#            '/': r'\/',
#            '>': r' ',
#            '<': r' '}

#     def escapedSeq(term):
#         """ Yield the next string based on the
#             next character (either this char
#             or escaped version """
#         for char in term:
#             if char in escapeRules.keys():
#                 yield escapeRules[char]
#             else:
#                 yield char

#     def escapeESArg(term):
#         """ Apply escaping to the passed in query terms
#             escaping special characters like : , etc"""
#         term = term.replace('\\', r'\\')   # escape \ first
#         return "".join([nextStr for nextStr in escapedSeq(term)])


#     question_string_escaped = escapeESArg(query.replace('?', ''))

#     query_data = {
#                         "from" : 0, "size" : max_retrieved,
#                         "query": {
#                             "query_string": {
#                                 "query": question_string_escaped
#                             }
#                         }
#                     }

#     response = requests.post(url=index_url, json=query_data, verify=False)
non_informative
#!/bin/env python
# -*- coding: utf-8 -*-

import argparse
import logging
import json
import jsonlines
import numpy.random
from tqdm import tqdm
from oneqa.ir.sparse.retriever import PyseriniRetriever

numpy.random.seed(1234)
logger=logging.getLogger(__name__)
logging.basicConfig(level = logging.INFO)

from utils import HasAnswerChecker

def handle_args():
    usage='usage'
    parser=argparse.ArgumentParser(usage)
    parser.add_argument('--input_file', '-i', required=False, default='/dccstor/jsmc-nmt-01/qa/xor-tydi/train-multilingual-2021-02-25/1/xorqa_dpr_data_query-L_hard_negative-1/dpr_train_data.json')
    parser.add_argument('--output_file', '-o', required=False, default='/dccstor/bsiyer6/OneQA/xortydi_bm25/xortydi_train_data_ir_negs_poss.json')
    parser.add_argument('--question_translations_dir', '-d', required=False,default='/dccstor/colbert-ir/bsiyer/data/xorqa/trans_data_all_langs')
    parser.add_argument('--index_path', '-u', required=False,default='/dccstor/bsiyer6/OneQA/psgs_w100_index')
    parser.add_argument('--max_num_ir_based_negatives', '-n', type=int, default=300)
    parser.add_argument('--max_num_ir_based_positives', '-p', type=int, default=10)
    parser.add_argument('--do_not_run_ir', '-R', action="store_true", default=False)
    parser.add_argument('--do_not_run_match_in_title', '-T', action="store_true", default=False)

    parser.add_argument('--question_records_file', '-q', default=None)
    parser.add_argument('--languages_list', '-l', nargs='+', default=[])

    args=parser.parse_args()
    return args

def init_question_translations(question_translations_dir):
    question_translations = {}
    # this is very XOR-TyDi QA specific
    langs = set(['ar', 'bn', 'fi', 'ja', 'ko', 'ru', 'te'])

    for lang in langs:
        with open(question_translations_dir + '/' + lang + '-en/en.txt') as f:
            en_content = f.readlines()
        with open(question_translations_dir + '/' + lang + '-en/' + lang + '.txt') as f:
            ne_content = f.readlines()

        assert(len(en_content) == len(ne_content))
        for ne, en in zip(ne_content, en_content):
            question_translations[ne.strip()] = en.strip()

    return question_translations


def run_query(query, searcher, max_retrieved):

    hits = searcher.retrieve(query, max_retrieved)
    retrieved = []
    for hit in hits:
        retrieved.append( (hit['doc_id'],hit['score'], hit['title'], hit['text']) )
    return retrieved

# def run_query(query, index_url, max_retrieved):

#     escapeRules = {'+': r'\+',
#            '-': r'\-',
#            '&': r'\&',
#            '|': r'\|',
#            '!': r'\!',
#            '(': r'\(',
#            ')': r'\)',
#            '{': r'\{',
#            '}': r'\}',
#            '[': r'\[',
#            ']': r'\]',
#            '^': r'\^',
#            '~': r'\~',
#            '*': r'\*',
#            '?': r'\?',
#            ':': r'\:',
#            '"': r'\"',
#            '\\': r'\\;',
#            '/': r'\/',
#            '>': r' ',
#            '<': r' '}

#     def escapedSeq(term):
#         """ Yield the next string based on the
#             next character (either this char
#             or escaped version """
#         for char in term:
#             if char in escapeRules.keys():
#                 yield escapeRules[char]
#             else:
#                 yield char

#     def escapeESArg(term):
#         """ Apply escaping to the passed in query terms
#             escaping special characters like : , etc"""
#         term = term.replace('\\', r'\\')   # escape \ first
#         return "".join([nextStr for nextStr in escapedSeq(term)])


#     question_string_escaped = escapeESArg(query.replace('?', ''))

#     query_data = {
#                         "from" : 0, "size" : max_retrieved,
#                         "query": {
#                             "query_string": {
#                                 "query": question_string_escaped
non_informative
#!/bin/env python
# -*- coding: utf-8 -*-

import argparse
import logging
import json
import jsonlines
import numpy.random
from tqdm import tqdm
from oneqa.ir.sparse.retriever import PyseriniRetriever

numpy.random.seed(1234)
logger=logging.getLogger(__name__)
logging.basicConfig(level = logging.INFO)

from utils import HasAnswerChecker

def handle_args():
    usage='usage'
    parser=argparse.ArgumentParser(usage)
    parser.add_argument('--input_file', '-i', required=False, default='/dccstor/jsmc-nmt-01/qa/xor-tydi/train-multilingual-2021-02-25/1/xorqa_dpr_data_query-L_hard_negative-1/dpr_train_data.json')
    parser.add_argument('--output_file', '-o', required=False, default='/dccstor/bsiyer6/OneQA/xortydi_bm25/xortydi_train_data_ir_negs_poss.json')
    parser.add_argument('--question_translations_dir', '-d', required=False,default='/dccstor/colbert-ir/bsiyer/data/xorqa/trans_data_all_langs')
    parser.add_argument('--index_path', '-u', required=False,default='/dccstor/bsiyer6/OneQA/psgs_w100_index')
    parser.add_argument('--max_num_ir_based_negatives', '-n', type=int, default=300)
    parser.add_argument('--max_num_ir_based_positives', '-p', type=int, default=10)
    parser.add_argument('--do_not_run_ir', '-R', action="store_true", default=False)
    parser.add_argument('--do_not_run_match_in_title', '-T', action="store_true", default=False)

    parser.add_argument('--question_records_file', '-q', default=None)
    parser.add_argument('--languages_list', '-l', nargs='+', default=[])

    args=parser.parse_args()
    return args

def init_question_translations(question_translations_dir):
    question_translations = {}
    # this is very XOR-TyDi QA specific
    langs = set(['ar', 'bn', 'fi', 'ja', 'ko', 'ru', 'te'])
non_informative
#!/bin/env python
# -*- coding: utf-8 -*-

import argparse
import logging
import json
import jsonlines
import numpy.random
from tqdm import tqdm
from oneqa.ir.sparse.retriever import PyseriniRetriever

numpy.random.seed(1234)
logger=logging.getLogger(__name__)
logging.basicConfig(level = logging.INFO)

from utils import HasAnswerChecker

def handle_args():
    usage='usage'
    parser=argparse.ArgumentParser(usage)
    parser.add_argument('--input_file', '-i', required=False, default='/dccstor/jsmc-nmt-01/qa/xor-tydi/train-multilingual-2021-02-25/1/xorqa_dpr_data_query-L_hard_negative-1/dpr_train_data.json')
    parser.add_argument('--output_file', '-o', required=False, default='/dccstor/bsiyer6/OneQA/xortydi_bm25/xortydi_train_data_ir_negs_poss.json')
    parser.add_argument('--question_translations_dir', '-d', required=False,default='/dccstor/colbert-ir/bsiyer/data/xorqa/trans_data_all_langs')
    parser.add_argument('--index_path', '-u', required=False,default='/dccstor/bsiyer6/OneQA/psgs_w100_index')
    parser.add_argument('--max_num_ir_based_negatives', '-n', type=int, default=300)
    parser.add_argument('--max_num_ir_based_positives', '-p', type=int, default=10)
    parser.add_argument('--do_not_run_ir', '-R', action="store_true", default=False)
    parser.add_argument('--do_not_run_match_in_title', '-T', action="store_true", default=False)

    parser.add_argument('--question_records_file', '-q', default=None)
    parser.add_argument('--languages_list', '-l', nargs='+', default=[])

    args=parser.parse_args()
    return args

def init_question_translations(question_translations_dir):
    question_translations = {}
    # this is very XOR-TyDi QA specific
    langs = set(['ar', 'bn', 'fi', 'ja', 'ko', 'ru', 'te'])

    for lang in langs:
        with open(question_translations_dir + '/' + lang + '-en/en.txt') as f:
            en_content = f.readlines()
        with open(question_translations_dir + '/' + lang + '-en/' + lang + '.txt') as f:
            ne_content = f.readlines()

        assert(len(en_content) == len(ne_content))
        for ne, en in zip(ne_content, en_content):
            question_translations[ne.strip()] = en.strip()

    return question_translations


def run_query(query, searcher, max_retrieved):

    hits = searcher.retrieve(query, max_retrieved)
    retrieved = []
    for hit in hits:
        retrieved.append( (hit['doc_id'],hit['score'], hit['title'], hit['text']) )
    return retrieved

# def run_query(query, index_url, max_retrieved):

#     escapeRules = {'+': r'\+',
#            '-': r'\-',
#            '&': r'\&',
#            '|': r'\|',
#            '!': r'\!',
#            '(': r'\(',
#            ')': r'\)',
#            '{': r'\{',
#            '}': r'\}',
#            '[': r'\[',
#            ']': r'\]',
non_informative
#!/bin/env python
# -*- coding: utf-8 -*-

import argparse
import logging
import json
import jsonlines
import numpy.random
from tqdm import tqdm
from oneqa.ir.sparse.retriever import PyseriniRetriever

numpy.random.seed(1234)
logger=logging.getLogger(__name__)
logging.basicConfig(level = logging.INFO)

from utils import HasAnswerChecker

def handle_args():
    usage='usage'
    parser=argparse.ArgumentParser(usage)
    parser.add_argument('--input_file', '-i', required=False, default='/dccstor/jsmc-nmt-01/qa/xor-tydi/train-multilingual-2021-02-25/1/xorqa_dpr_data_query-L_hard_negative-1/dpr_train_data.json')
    parser.add_argument('--output_file', '-o', required=False, default='/dccstor/bsiyer6/OneQA/xortydi_bm25/xortydi_train_data_ir_negs_poss.json')
    parser.add_argument('--question_translations_dir', '-d', required=False,default='/dccstor/colbert-ir/bsiyer/data/xorqa/trans_data_all_langs')
    parser.add_argument('--index_path', '-u', required=False,default='/dccstor/bsiyer6/OneQA/psgs_w100_index')
    parser.add_argument('--max_num_ir_based_negatives', '-n', type=int, default=300)
    parser.add_argument('--max_num_ir_based_positives', '-p', type=int, default=10)
    parser.add_argument('--do_not_run_ir', '-R', action="store_true", default=False)
    parser.add_argument('--do_not_run_match_in_title', '-T', action="store_true", default=False)

    parser.add_argument('--question_records_file', '-q', default=None)
    parser.add_argument('--languages_list', '-l', nargs='+', default=[])

    args=parser.parse_args()
    return args

def init_question_translations(question_translations_dir):
    question_translations = {}
    # this is very XOR-TyDi QA specific
    langs = set(['ar', 'bn', 'fi', 'ja', 'ko', 'ru', 'te'])

    for lang in langs:
        with open(question_translations_dir + '/' + lang + '-en/en.txt') as f:
            en_content = f.readlines()
        with open(question_translations_dir + '/' + lang + '-en/' + lang + '.txt') as f:
            ne_content = f.readlines()

        assert(len(en_content) == len(ne_content))
        for ne, en in zip(ne_content, en_content):
            question_translations[ne.strip()] = en.strip()

    return question_translations


def run_query(query, searcher, max_retrieved):

    hits = searcher.retrieve(query, max_retrieved)
    retrieved = []
    for hit in hits:
        retrieved.append( (hit['doc_id'],hit['score'], hit['title'], hit['text']) )
    return retrieved

# def run_query(query, index_url, max_retrieved):

#     escapeRules = {'+': r'\+',
#            '-': r'\-',
#            '&': r'\&',
#            '|': r'\|',
#            '!': r'\!',
#            '(': r'\(',
#            ')': r'\)',
#            '{': r'\{',
#            '}': r'\}',
#            '[': r'\[',
#            ']': r'\]',
#            '^': r'\^',
#            '~': r'\~',
#            '*': r'\*',
#            '?': r'\?',
#            ':': r'\:',
#            '"': r'\"',
#            '\\': r'\\;',
#            '/': r'\/',
#            '>': r' ',
#            '<': r' '}

#     def escapedSeq(term):
#         """ Yield the next string based on the
#             next character (either this char
#             or escaped version """
#         for char in term:
#             if char in escapeRules.keys():
#                 yield escapeRules[char]
#             else:
#                 yield char

#     def escapeESArg(term):
#         """ Apply escaping to the passed in query terms
#             escaping special characters like : , etc"""
#         term = term.replace('\\', r'\\')   # escape \ first
#         return "".join([nextStr for nextStr in escapedSeq(term)])


#     question_string_escaped = escapeESArg(query.replace('?', ''))

#     query_data = {
#                         "from" : 0, "size" : max_retrieved,
#                         "query": {
#                             "query_string": {
#                                 "query": question_string_escaped
#                             }
#                         }
#                     }

#     response = requests.post(url=index_url, json=query_data, verify=False)
#     retrieved = []
#     response_text_dict = json.loads(response.text)

#     if 'hits' in response_text_dict:
#         for hit in response_text_dict['hits']['hits']:
#             retrieved.append((hit['_id'], hit['_score'],
#                               (hit['_source']['title'] if 'title' in hit['_source'] else "", hit['_source']['text'])))

#     return retrieved

def read_jsonlines(file_name):
    lines = []
    print("loading examples from {0}".format(file_name))
    with jsonlines.open(file_name) as reader:
        for obj in reader:
random
#!/bin/env python
# -*- coding: utf-8 -*-

import argparse
import logging
import json
import jsonlines
import numpy.random
from tqdm import tqdm
from oneqa.ir.sparse.retriever import PyseriniRetriever

numpy.random.seed(1234)
logger=logging.getLogger(__name__)
logging.basicConfig(level = logging.INFO)

from utils import HasAnswerChecker

def handle_args():
    usage='usage'
    parser=argparse.ArgumentParser(usage)
    parser.add_argument('--input_file', '-i', required=False, default='/dccstor/jsmc-nmt-01/qa/xor-tydi/train-multilingual-2021-02-25/1/xorqa_dpr_data_query-L_hard_negative-1/dpr_train_data.json')
    parser.add_argument('--output_file', '-o', required=False, default='/dccstor/bsiyer6/OneQA/xortydi_bm25/xortydi_train_data_ir_negs_poss.json')
    parser.add_argument('--question_translations_dir', '-d', required=False,default='/dccstor/colbert-ir/bsiyer/data/xorqa/trans_data_all_langs')
    parser.add_argument('--index_path', '-u', required=False,default='/dccstor/bsiyer6/OneQA/psgs_w100_index')
    parser.add_argument('--max_num_ir_based_negatives', '-n', type=int, default=300)
    parser.add_argument('--max_num_ir_based_positives', '-p', type=int, default=10)
    parser.add_argument('--do_not_run_ir', '-R', action="store_true", default=False)
    parser.add_argument('--do_not_run_match_in_title', '-T', action="store_true", default=False)

    parser.add_argument('--question_records_file', '-q', default=None)
    parser.add_argument('--languages_list', '-l', nargs='+', default=[])

    args=parser.parse_args()
    return args

def init_question_translations(question_translations_dir):
    question_translations = {}
    # this is very XOR-TyDi QA specific
    langs = set(['ar', 'bn', 'fi', 'ja', 'ko', 'ru', 'te'])

    for lang in langs:
        with open(question_translations_dir + '/' + lang + '-en/en.txt') as f:
            en_content = f.readlines()
        with open(question_translations_dir + '/' + lang + '-en/' + lang + '.txt') as f:
            ne_content = f.readlines()

        assert(len(en_content) == len(ne_content))
        for ne, en in zip(ne_content, en_content):
            question_translations[ne.strip()] = en.strip()

    return question_translations


def run_query(query, searcher, max_retrieved):

    hits = searcher.retrieve(query, max_retrieved)
    retrieved = []
    for hit in hits:
        retrieved.append( (hit['doc_id'],hit['score'], hit['title'], hit['text']) )
    return retrieved

# def run_query(query, index_url, max_retrieved):

#     escapeRules = {'+': r'\+',
#            '-': r'\-',
#            '&': r'\&',
#            '|': r'\|',
#            '!': r'\!',
#            '(': r'\(',
#            ')': r'\)',
#            '{': r'\{',
#            '}': r'\}',
#            '[': r'\[',
#            ']': r'\]',
#            '^': r'\^',
#            '~': r'\~',
#            '*': r'\*',
#            '?': r'\?',
#            ':': r'\:',
#            '"': r'\"',
#            '\\': r'\\;',
#            '/': r'\/',
#            '>': r' ',
#            '<': r' '}

#     def escapedSeq(term):
#         """ Yield the next string based on the
#             next character (either this char
#             or escaped version """
#         for char in term:
#             if char in escapeRules.keys():
#                 yield escapeRules[char]
#             else:
#                 yield char

#     def escapeESArg(term):
#         """ Apply escaping to the passed in query terms
#             escaping special characters like : , etc"""
#         term = term.replace('\\', r'\\')   # escape \ first
#         return "".join([nextStr for nextStr in escapedSeq(term)])


#     question_string_escaped = escapeESArg(query.replace('?', ''))

#     query_data = {
#                         "from" : 0, "size" : max_retrieved,
#                         "query": {
#                             "query_string": {
#                                 "query": question_string_escaped
#                             }
#                         }
#                     }

#     response = requests.post(url=index_url, json=query_data, verify=False)
#     retrieved = []
#     response_text_dict = json.loads(response.text)

#     if 'hits' in response_text_dict:
#         for hit in response_text_dict['hits']['hits']:
#             retrieved.append((hit['_id'], hit['_score'],
#                               (hit['_source']['title'] if 'title' in hit['_source'] else "", hit['_source']['text'])))

#     return retrieved

def read_jsonlines(file_name):
    lines = []
    print("loading examples from {0}".format(file_name))
    with jsonlines.open(file_name) as reader:
        for obj in reader:
            lines.append(obj)
    return lines


def main():
    args=handle_args()

    answer_checker = HasAnswerChecker()

    searcher = PyseriniRetriever(args.index_path)

    question_translations = init_question_translations(args.question_translations_dir)

    qas=json.load(open(args.input_file))

    if args.question_records_file is not None:
        q_recs=read_jsonlines(args.question_records_file)

    languages_to_use = set(args.languages_list)

    for qnum, qa in enumerate(tqdm(qas)):
        if args.question_records_file is not None:
            # sanity check
            assert(qa['answers'] == q_recs[qnum]['answers'])
            if q_recs[qnum]['lang'] not in languages_to_use:
                continue

        q=qa['question']
        p=qa['positive_ctxs']
        n=qa['negative_ctxs']
        hn=qa['hard_negative_ctxs']

        en_q = question_translations[q]

        if args.do_not_run_ir:
            continue

        retrieved = run_query(en_q, searcher, args.max_num_ir_based_negatives * 2) # assuming 50/50 positive/negative ratio

        negatives = []
        positives = []

        positive_ctxs_texts = [ctxt['text'] for ctxt in p]

        for hit in retrieved:
            #string_to_search = hit[2][1] if args.do_not_run_match_in_title else hit[2][0] + ' ' + hit[2][1]
            string_to_search = hit[3] if args.do_not_run_match_in_title else hit[2] + ' ' + hit[3]
            if answer_checker.has_answer(qa['answers'], string_to_search):
random
#!/bin/env python
# -*- coding: utf-8 -*-

import argparse
import logging
import json
import jsonlines
import numpy.random
from tqdm import tqdm
from oneqa.ir.sparse.retriever import PyseriniRetriever

numpy.random.seed(1234)
logger=logging.getLogger(__name__)
logging.basicConfig(level = logging.INFO)

from utils import HasAnswerChecker

def handle_args():
    usage='usage'
    parser=argparse.ArgumentParser(usage)
    parser.add_argument('--input_file', '-i', required=False, default='/dccstor/jsmc-nmt-01/qa/xor-tydi/train-multilingual-2021-02-25/1/xorqa_dpr_data_query-L_hard_negative-1/dpr_train_data.json')
    parser.add_argument('--output_file', '-o', required=False, default='/dccstor/bsiyer6/OneQA/xortydi_bm25/xortydi_train_data_ir_negs_poss.json')
    parser.add_argument('--question_translations_dir', '-d', required=False,default='/dccstor/colbert-ir/bsiyer/data/xorqa/trans_data_all_langs')
    parser.add_argument('--index_path', '-u', required=False,default='/dccstor/bsiyer6/OneQA/psgs_w100_index')
    parser.add_argument('--max_num_ir_based_negatives', '-n', type=int, default=300)
    parser.add_argument('--max_num_ir_based_positives', '-p', type=int, default=10)
    parser.add_argument('--do_not_run_ir', '-R', action="store_true", default=False)
    parser.add_argument('--do_not_run_match_in_title', '-T', action="store_true", default=False)

    parser.add_argument('--question_records_file', '-q', default=None)
    parser.add_argument('--languages_list', '-l', nargs='+', default=[])

    args=parser.parse_args()
    return args

def init_question_translations(question_translations_dir):
    question_translations = {}
    # this is very XOR-TyDi QA specific
    langs = set(['ar', 'bn', 'fi', 'ja', 'ko', 'ru', 'te'])

    for lang in langs:
        with open(question_translations_dir + '/' + lang + '-en/en.txt') as f:
            en_content = f.readlines()
        with open(question_translations_dir + '/' + lang + '-en/' + lang + '.txt') as f:
            ne_content = f.readlines()

        assert(len(en_content) == len(ne_content))
        for ne, en in zip(ne_content, en_content):
            question_translations[ne.strip()] = en.strip()

    return question_translations


def run_query(query, searcher, max_retrieved):

    hits = searcher.retrieve(query, max_retrieved)
    retrieved = []
    for hit in hits:
        retrieved.append( (hit['doc_id'],hit['score'], hit['title'], hit['text']) )
    return retrieved

# def run_query(query, index_url, max_retrieved):

#     escapeRules = {'+': r'\+',
#            '-': r'\-',
#            '&': r'\&',
#            '|': r'\|',
#            '!': r'\!',
#            '(': r'\(',
#            ')': r'\)',
#            '{': r'\{',
#            '}': r'\}',
#            '[': r'\[',
#            ']': r'\]',
#            '^': r'\^',
#            '~': r'\~',
#            '*': r'\*',
#            '?': r'\?',
#            ':': r'\:',
#            '"': r'\"',
#            '\\': r'\\;',
#            '/': r'\/',
#            '>': r' ',
#            '<': r' '}

#     def escapedSeq(term):
#         """ Yield the next string based on the
#             next character (either this char
#             or escaped version """
#         for char in term:
#             if char in escapeRules.keys():
#                 yield escapeRules[char]
#             else:
#                 yield char

#     def escapeESArg(term):
#         """ Apply escaping to the passed in query terms
#             escaping special characters like : , etc"""
#         term = term.replace('\\', r'\\')   # escape \ first
#         return "".join([nextStr for nextStr in escapedSeq(term)])


#     question_string_escaped = escapeESArg(query.replace('?', ''))

#     query_data = {
#                         "from" : 0, "size" : max_retrieved,
#                         "query": {
#                             "query_string": {
#                                 "query": question_string_escaped
#                             }
#                         }
#                     }

#     response = requests.post(url=index_url, json=query_data, verify=False)
#     retrieved = []
#     response_text_dict = json.loads(response.text)

#     if 'hits' in response_text_dict:
#         for hit in response_text_dict['hits']['hits']:
#             retrieved.append((hit['_id'], hit['_score'],
#                               (hit['_source']['title'] if 'title' in hit['_source'] else "", hit['_source']['text'])))

#     return retrieved

def read_jsonlines(file_name):
    lines = []
    print("loading examples from {0}".format(file_name))
    with jsonlines.open(file_name) as reader:
        for obj in reader:
            lines.append(obj)
    return lines


def main():
    args=handle_args()

    answer_checker = HasAnswerChecker()

    searcher = PyseriniRetriever(args.index_path)

    question_translations = init_question_translations(args.question_translations_dir)

    qas=json.load(open(args.input_file))

    if args.question_records_file is not None:
        q_recs=read_jsonlines(args.question_records_file)

    languages_to_use = set(args.languages_list)

    for qnum, qa in enumerate(tqdm(qas)):
        if args.question_records_file is not None:
            # sanity check
            assert(qa['answers'] == q_recs[qnum]['answers'])
            if q_recs[qnum]['lang'] not in languages_to_use:
random
#!/bin/env python
# -*- coding: utf-8 -*-

import argparse
import logging
import json
import jsonlines
import numpy.random
from tqdm import tqdm
from oneqa.ir.sparse.retriever import PyseriniRetriever

numpy.random.seed(1234)
logger=logging.getLogger(__name__)
logging.basicConfig(level = logging.INFO)

from utils import HasAnswerChecker

def handle_args():
    usage='usage'
    parser=argparse.ArgumentParser(usage)
    parser.add_argument('--input_file', '-i', required=False, default='/dccstor/jsmc-nmt-01/qa/xor-tydi/train-multilingual-2021-02-25/1/xorqa_dpr_data_query-L_hard_negative-1/dpr_train_data.json')
    parser.add_argument('--output_file', '-o', required=False, default='/dccstor/bsiyer6/OneQA/xortydi_bm25/xortydi_train_data_ir_negs_poss.json')
    parser.add_argument('--question_translations_dir', '-d', required=False,default='/dccstor/colbert-ir/bsiyer/data/xorqa/trans_data_all_langs')
    parser.add_argument('--index_path', '-u', required=False,default='/dccstor/bsiyer6/OneQA/psgs_w100_index')
    parser.add_argument('--max_num_ir_based_negatives', '-n', type=int, default=300)
    parser.add_argument('--max_num_ir_based_positives', '-p', type=int, default=10)
    parser.add_argument('--do_not_run_ir', '-R', action="store_true", default=False)
    parser.add_argument('--do_not_run_match_in_title', '-T', action="store_true", default=False)

    parser.add_argument('--question_records_file', '-q', default=None)
    parser.add_argument('--languages_list', '-l', nargs='+', default=[])

    args=parser.parse_args()
    return args

def init_question_translations(question_translations_dir):
    question_translations = {}
    # this is very XOR-TyDi QA specific
    langs = set(['ar', 'bn', 'fi', 'ja', 'ko', 'ru', 'te'])

    for lang in langs:
        with open(question_translations_dir + '/' + lang + '-en/en.txt') as f:
            en_content = f.readlines()
        with open(question_translations_dir + '/' + lang + '-en/' + lang + '.txt') as f:
            ne_content = f.readlines()

        assert(len(en_content) == len(ne_content))
        for ne, en in zip(ne_content, en_content):
            question_translations[ne.strip()] = en.strip()

    return question_translations


def run_query(query, searcher, max_retrieved):

    hits = searcher.retrieve(query, max_retrieved)
    retrieved = []
    for hit in hits:
        retrieved.append( (hit['doc_id'],hit['score'], hit['title'], hit['text']) )
    return retrieved

# def run_query(query, index_url, max_retrieved):

#     escapeRules = {'+': r'\+',
#            '-': r'\-',
#            '&': r'\&',
#            '|': r'\|',
#            '!': r'\!',
#            '(': r'\(',
#            ')': r'\)',
#            '{': r'\{',
#            '}': r'\}',
#            '[': r'\[',
#            ']': r'\]',
#            '^': r'\^',
#            '~': r'\~',
#            '*': r'\*',
#            '?': r'\?',
#            ':': r'\:',
#            '"': r'\"',
#            '\\': r'\\;',
#            '/': r'\/',
#            '>': r' ',
#            '<': r' '}

#     def escapedSeq(term):
#         """ Yield the next string based on the
#             next character (either this char
#             or escaped version """
#         for char in term:
#             if char in escapeRules.keys():
#                 yield escapeRules[char]
#             else:
#                 yield char

#     def escapeESArg(term):
#         """ Apply escaping to the passed in query terms
#             escaping special characters like : , etc"""
#         term = term.replace('\\', r'\\')   # escape \ first
#         return "".join([nextStr for nextStr in escapedSeq(term)])


#     question_string_escaped = escapeESArg(query.replace('?', ''))

#     query_data = {
#                         "from" : 0, "size" : max_retrieved,
#                         "query": {
#                             "query_string": {
#                                 "query": question_string_escaped
#                             }
#                         }
#                     }

#     response = requests.post(url=index_url, json=query_data, verify=False)
#     retrieved = []
#     response_text_dict = json.loads(response.text)

#     if 'hits' in response_text_dict:
#         for hit in response_text_dict['hits']['hits']:
#             retrieved.append((hit['_id'], hit['_score'],
#                               (hit['_source']['title'] if 'title' in hit['_source'] else "", hit['_source']['text'])))

#     return retrieved

def read_jsonlines(file_name):
    lines = []
    print("loading examples from {0}".format(file_name))
    with jsonlines.open(file_name) as reader:
        for obj in reader:
            lines.append(obj)
    return lines


def main():
    args=handle_args()

    answer_checker = HasAnswerChecker()

    searcher = PyseriniRetriever(args.index_path)

    question_translations = init_question_translations(args.question_translations_dir)

    qas=json.load(open(args.input_file))

    if args.question_records_file is not None:
        q_recs=read_jsonlines(args.question_records_file)

    languages_to_use = set(args.languages_list)

    for qnum, qa in enumerate(tqdm(qas)):
        if args.question_records_file is not None:
            # sanity check
            assert(qa['answers'] == q_recs[qnum]['answers'])
            if q_recs[qnum]['lang'] not in languages_to_use:
                continue

        q=qa['question']
        p=qa['positive_ctxs']
        n=qa['negative_ctxs']
        hn=qa['hard_negative_ctxs']

        en_q = question_translations[q]

        if args.do_not_run_ir:
            continue

        retrieved = run_query(en_q, searcher, args.max_num_ir_based_negatives * 2) # assuming 50/50 positive/negative ratio

        negatives = []
        positives = []

        positive_ctxs_texts = [ctxt['text'] for ctxt in p]

        for hit in retrieved:
            #string_to_search = hit[2][1] if args.do_not_run_match_in_title else hit[2][0] + ' ' + hit[2][1]
            string_to_search = hit[3] if args.do_not_run_match_in_title else hit[2] + ' ' + hit[3]
            if answer_checker.has_answer(qa['answers'], string_to_search):
                if len(positives) < args.max_num_ir_based_positives and not max([ ctxt in hit[3] for ctxt in positive_ctxs_texts ]):
                    #positives.append({'title': hit[2][0], 'text': hit[2][1]})
                    positives.append({'title': hit[2], 'text': hit[3]})
            else:
                if len(negatives) < args.max_num_ir_based_negatives:
                    # negatives.append({'title': hit[2][0], 'text': hit[2][1]})
                    negatives.append({'title': hit[2], 'text': hit[3]})

            if len(positives) >= args.max_num_ir_based_positives and len(negatives) >= args.max_num_ir_based_negatives:
                break

        qa['ir_negative_ctxs'] = negatives
        qa['ir_positive_ctxs'] = positives

        print('adding ' + str(len(qa['ir_negative_ctxs'])) + ' IR based negatives, ' + str(len(qa['ir_positive_ctxs'])) + ' IR based positives')

        if qnum > 0 and (qnum % 1000) == 0:
            output_file = args.output_file.replace('.json', '_' + str(qnum) + '.json')
            print('writing ', output_file)
            with open(output_file, 'w') as out_f:
                json.dump(qas, out_f, indent=4)

    output_file = args.output_file
    print('writing ', output_file)
random
#!/bin/env python
# -*- coding: utf-8 -*-

import argparse
import logging
import json
import jsonlines
import numpy.random
from tqdm import tqdm
from oneqa.ir.sparse.retriever import PyseriniRetriever

numpy.random.seed(1234)
logger=logging.getLogger(__name__)
logging.basicConfig(level = logging.INFO)

from utils import HasAnswerChecker

def handle_args():
    usage='usage'
    parser=argparse.ArgumentParser(usage)
    parser.add_argument('--input_file', '-i', required=False, default='/dccstor/jsmc-nmt-01/qa/xor-tydi/train-multilingual-2021-02-25/1/xorqa_dpr_data_query-L_hard_negative-1/dpr_train_data.json')
    parser.add_argument('--output_file', '-o', required=False, default='/dccstor/bsiyer6/OneQA/xortydi_bm25/xortydi_train_data_ir_negs_poss.json')
    parser.add_argument('--question_translations_dir', '-d', required=False,default='/dccstor/colbert-ir/bsiyer/data/xorqa/trans_data_all_langs')
    parser.add_argument('--index_path', '-u', required=False,default='/dccstor/bsiyer6/OneQA/psgs_w100_index')
    parser.add_argument('--max_num_ir_based_negatives', '-n', type=int, default=300)
    parser.add_argument('--max_num_ir_based_positives', '-p', type=int, default=10)
    parser.add_argument('--do_not_run_ir', '-R', action="store_true", default=False)
    parser.add_argument('--do_not_run_match_in_title', '-T', action="store_true", default=False)

    parser.add_argument('--question_records_file', '-q', default=None)
    parser.add_argument('--languages_list', '-l', nargs='+', default=[])

    args=parser.parse_args()
    return args

def init_question_translations(question_translations_dir):
    question_translations = {}
    # this is very XOR-TyDi QA specific
    langs = set(['ar', 'bn', 'fi', 'ja', 'ko', 'ru', 'te'])

    for lang in langs:
        with open(question_translations_dir + '/' + lang + '-en/en.txt') as f:
            en_content = f.readlines()
        with open(question_translations_dir + '/' + lang + '-en/' + lang + '.txt') as f:
            ne_content = f.readlines()

        assert(len(en_content) == len(ne_content))
        for ne, en in zip(ne_content, en_content):
            question_translations[ne.strip()] = en.strip()

    return question_translations


def run_query(query, searcher, max_retrieved):

    hits = searcher.retrieve(query, max_retrieved)
    retrieved = []
    for hit in hits:
        retrieved.append( (hit['doc_id'],hit['score'], hit['title'], hit['text']) )
    return retrieved

# def run_query(query, index_url, max_retrieved):

#     escapeRules = {'+': r'\+',
#            '-': r'\-',
#            '&': r'\&',
#            '|': r'\|',
#            '!': r'\!',
#            '(': r'\(',
#            ')': r'\)',
#            '{': r'\{',
#            '}': r'\}',
#            '[': r'\[',
#            ']': r'\]',
#            '^': r'\^',
#            '~': r'\~',
#            '*': r'\*',
#            '?': r'\?',
#            ':': r'\:',
#            '"': r'\"',
#            '\\': r'\\;',
#            '/': r'\/',
#            '>': r' ',
#            '<': r' '}

#     def escapedSeq(term):
#         """ Yield the next string based on the
#             next character (either this char
#             or escaped version """
#         for char in term:
#             if char in escapeRules.keys():
#                 yield escapeRules[char]
#             else:
#                 yield char

#     def escapeESArg(term):
#         """ Apply escaping to the passed in query terms
#             escaping special characters like : , etc"""
#         term = term.replace('\\', r'\\')   # escape \ first
#         return "".join([nextStr for nextStr in escapedSeq(term)])


#     question_string_escaped = escapeESArg(query.replace('?', ''))

#     query_data = {
#                         "from" : 0, "size" : max_retrieved,
#                         "query": {
#                             "query_string": {
#                                 "query": question_string_escaped
#                             }
#                         }
#                     }

#     response = requests.post(url=index_url, json=query_data, verify=False)
#     retrieved = []
#     response_text_dict = json.loads(response.text)

#     if 'hits' in response_text_dict:
#         for hit in response_text_dict['hits']['hits']:
#             retrieved.append((hit['_id'], hit['_score'],
#                               (hit['_source']['title'] if 'title' in hit['_source'] else "", hit['_source']['text'])))

#     return retrieved

def read_jsonlines(file_name):
    lines = []
    print("loading examples from {0}".format(file_name))
    with jsonlines.open(file_name) as reader:
        for obj in reader:
            lines.append(obj)
    return lines


def main():
    args=handle_args()

    answer_checker = HasAnswerChecker()

    searcher = PyseriniRetriever(args.index_path)

    question_translations = init_question_translations(args.question_translations_dir)

    qas=json.load(open(args.input_file))

    if args.question_records_file is not None:
        q_recs=read_jsonlines(args.question_records_file)

    languages_to_use = set(args.languages_list)

    for qnum, qa in enumerate(tqdm(qas)):
        if args.question_records_file is not None:
            # sanity check
            assert(qa['answers'] == q_recs[qnum]['answers'])
            if q_recs[qnum]['lang'] not in languages_to_use:
                continue

        q=qa['question']
        p=qa['positive_ctxs']
        n=qa['negative_ctxs']
        hn=qa['hard_negative_ctxs']

        en_q = question_translations[q]

        if args.do_not_run_ir:
            continue

        retrieved = run_query(en_q, searcher, args.max_num_ir_based_negatives * 2) # assuming 50/50 positive/negative ratio

        negatives = []
        positives = []

        positive_ctxs_texts = [ctxt['text'] for ctxt in p]

        for hit in retrieved:
            #string_to_search = hit[2][1] if args.do_not_run_match_in_title else hit[2][0] + ' ' + hit[2][1]
            string_to_search = hit[3] if args.do_not_run_match_in_title else hit[2] + ' ' + hit[3]
            if answer_checker.has_answer(qa['answers'], string_to_search):
                if len(positives) < args.max_num_ir_based_positives and not max([ ctxt in hit[3] for ctxt in positive_ctxs_texts ]):
                    #positives.append({'title': hit[2][0], 'text': hit[2][1]})
                    positives.append({'title': hit[2], 'text': hit[3]})
            else:
                if len(negatives) < args.max_num_ir_based_negatives:
                    # negatives.append({'title': hit[2][0], 'text': hit[2][1]})
                    negatives.append({'title': hit[2], 'text': hit[3]})

            if len(positives) >= args.max_num_ir_based_positives and len(negatives) >= args.max_num_ir_based_negatives:
                break

        qa['ir_negative_ctxs'] = negatives
        qa['ir_positive_ctxs'] = positives

        print('adding ' + str(len(qa['ir_negative_ctxs'])) + ' IR based negatives, ' + str(len(qa['ir_positive_ctxs'])) + ' IR based positives')

        if qnum > 0 and (qnum % 1000) == 0:
            output_file = args.output_file.replace('.json', '_' + str(qnum) + '.json')
            print('writing ', output_file)
            with open(output_file, 'w') as out_f:
                json.dump(qas, out_f, indent=4)

    output_file = args.output_file
    print('writing ', output_file)
    with open(output_file, 'w') as out_f:
        json.dump(qas, out_f, indent=4)


# do main
random
        q_recs=read_jsonlines(args.question_records_file)     question_translations = init_question_translations(args.question_translations_dir)     args=handle_args()     searcher = PyseriniRetriever(args.index_path)     hits = searcher.retrieve(query, max_retrieved)     qas=json.load(open(args.input_file))    main() #            '~': r'\~', #     retrieved = [] #                             } #            '^': r'\^',             lines.append(obj)                 if len(positives) < args.max_num_ir_based_positives and not max([ ctxt in hit[3] for ctxt in positive_ctxs_texts ]):                 continue     with open(output_file, 'w') as out_f: if __name__=='__main__':