codingutf8 2020 huggingface 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 orglicenseslicense2 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 coding utf 8 2020 huggingface 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
from transformers import (
    DPRContextEncoderTokenizer,
    DPRContextEncoderTokenizerFast,
    DPRQuestionEncoderTokenizer,
    DPRQuestionEncoderTokenizerFast,
    DPRReaderOutput,
    DPRReaderTokenizer,
    DPRReaderTokenizerFast,
)
from transformers.testing_utils import require_tokenizers, slow
from transformers.tokenization_utils_base import BatchEncoding

from ..bert.test_tokenization_bert import BertTokenizationTest


@require_tokenizers
class DPRContextEncoderTokenizationTest(BertTokenizationTest):
    tokenizer_class = DPRContextEncoderTokenizer
    rust_tokenizer_class = DPRContextEncoderTokenizerFast
    test_rust_tokenizer = True


@require_tokenizers
class DPRQuestionEncoderTokenizationTest(BertTokenizationTest):
    tokenizer_class = DPRQuestionEncoderTokenizer
    rust_tokenizer_class = DPRQuestionEncoderTokenizerFast
    test_rust_tokenizer = True


@require_tokenizers
class DPRReaderTokenizationTest(BertTokenizationTest):
    tokenizer_class = DPRReaderTokenizer
    rust_tokenizer_class = DPRReaderTokenizerFast
    test_rust_tokenizer = True

    @slow
    def test_decode_best_spans(self):
        tokenizer = self.tokenizer_class.from_pretrained("bert-base-uncased")

        text_1 = tokenizer.encode("question sequence", add_special_tokens=False)
        text_2 = tokenizer.encode("title sequence", add_special_tokens=False)
        text_3 = tokenizer.encode("text sequence " * 4, add_special_tokens=False)
        input_ids = [[101] + text_1 + [102] + text_2 + [102] + text_3]
        reader_input = BatchEncoding({"input_ids": input_ids})

        start_logits = [[0] * len(input_ids[0])]
        end_logits = [[0] * len(input_ids[0])]
        relevance_logits = [0]
        reader_output = DPRReaderOutput(start_logits, end_logits, relevance_logits)

        start_index, end_index = 8, 9
        start_logits[0][start_index] = 10
        end_logits[0][end_index] = 10
        predicted_spans = tokenizer.decode_best_spans(reader_input, reader_output)
        self.assertEqual(predicted_spans[0].start_index, start_index)
        self.assertEqual(predicted_spans[0].end_index, end_index)
        self.assertEqual(predicted_spans[0].doc_id, 0)

    @slow
    def test_call(self):
        tokenizer = self.tokenizer_class.from_pretrained("bert-base-uncased")

        text_1 = tokenizer.encode("question sequence", add_special_tokens=False)
        text_2 = tokenizer.encode("title sequence", add_special_tokens=False)
        text_3 = tokenizer.encode("text sequence", add_special_tokens=False)
        expected_input_ids = [101] + text_1 + [102] + text_2 + [102] + text_3
        encoded_input = tokenizer(questions=["question sequence"], titles=["title sequence"], texts=["text sequence"])
        self.assertIn("input_ids", encoded_input)
        self.assertIn("attention_mask", encoded_input)
        self.assertListEqual(encoded_input["input_ids"][0], expected_input_ids)