Add REALM (#13292)

* REALM initial commit

* Retriever OK (Update new_gelu).

* Encoder prediction score OK

* Encoder pretrained model OK

* Update retriever comments

* Update docs, tests, and imports

* Prune unused models

* Make embedder as a module `RealmEmbedder`

* Add RealmRetrieverOutput

* Update tokenization

* Pass all tests in test_modeling_realm.py

* Prune RealmModel

* Update docs

* Add training test.

* Remove completed TODO

* Style & Quality

* Prune `RealmModel`

* Fixup

* Changes:
1. Remove RealmTokenizerFast
2. Update docstrings
3. Add a method to RealmTokenizer to handle candidates tokenization.

* Fix up

* Style

* Add tokenization tests

* Update `from_pretrained` tests

* Apply suggestions

* Style & Quality

* Copy BERT model

* Fix comment to avoid docstring copying

* Make RealmBertModel private

* Fix bug

* Style

* Basic QA

* Save

* Complete reader logits

* Add searcher

* Complete searcher & reader

* Move block records init to constructor

* Fix training bug

* Add some outputs to RealmReader

* Add finetuned checkpoint variable names parsing

* Fix bug

* Update REALM config

* Add RealmForOpenQA

* Update convert_tfrecord logits

* Fix bugs

* Complete imports

* Update docs

* Update naming

* Add brute-force searcher

* Pass realm model tests

* Style

* Exclude RealmReader from common tests

* Fix

* Fix

* convert docs

* up

* up

* more make style

* up

* upload

* up

* Fix

* Update src/transformers/__init__.py

* adapt testing

* change modeling code

* fix test

* up

* up

* up

* correct more

* make retriever work

* update

* make style

* finish main structure

* Resolve merge conflict

* Make everything work

* Style

* Fixup

* Fixup

* Update training test

* fix retriever

* remove hardcoded path

* Fix

* Fix modeling test

* Update model links

* Initial retrieval test

* Fix modeling test

* Complete retrieval tests

* Fix

* style

* Fix tests

* Fix docstring example

* Minor fix of retrieval test

* Update license headers and docs

* Apply suggestions from code review

* Style

* Apply suggestions from code review

* Add an example to RealmEmbedder

* Fix

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>

tests/test_modeling_realm.py

@@ -0,0 +1,545 @@
+# coding=utf-8
+# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+""" Testing suite for the PyTorch REALM model. """
+
+import copy
+import unittest
+
+import numpy as np
+
+from tests.test_modeling_common import floats_tensor
+from transformers import RealmConfig, is_torch_available
+from transformers.testing_utils import require_torch, slow, torch_device
+
+from .test_configuration_common import ConfigTester
+from .test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
+
+
+if is_torch_available():
+    import torch
+
+    from transformers import (
+        RealmEmbedder,
+        RealmForOpenQA,
+        RealmKnowledgeAugEncoder,
+        RealmReader,
+        RealmRetriever,
+        RealmScorer,
+        RealmTokenizer,
+    )
+
+
+class RealmModelTester:
+    def __init__(
+        self,
+        parent,
+        batch_size=13,
+        retriever_proj_size=128,
+        seq_length=7,
+        is_training=True,
+        use_input_mask=True,
+        use_token_type_ids=True,
+        use_labels=True,
+        vocab_size=99,
+        hidden_size=32,
+        num_hidden_layers=5,
+        num_attention_heads=4,
+        intermediate_size=37,
+        hidden_act="gelu",
+        hidden_dropout_prob=0.1,
+        attention_probs_dropout_prob=0.1,
+        max_position_embeddings=512,
+        type_vocab_size=16,
+        type_sequence_label_size=2,
+        initializer_range=0.02,
+        layer_norm_eps=1e-12,
+        span_hidden_size=50,
+        max_span_width=10,
+        reader_layer_norm_eps=1e-3,
+        reader_beam_size=4,
+        reader_seq_len=288 + 32,
+        num_block_records=13353718,
+        searcher_beam_size=8,
+        searcher_seq_len=64,
+        num_labels=3,
+        num_choices=4,
+        num_candidates=10,
+        scope=None,
+    ):
+        # General config
+        self.parent = parent
+        self.batch_size = batch_size
+        self.retriever_proj_size = retriever_proj_size
+        self.seq_length = seq_length
+        self.is_training = is_training
+        self.use_input_mask = use_input_mask
+        self.use_token_type_ids = use_token_type_ids
+        self.use_labels = use_labels
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.hidden_act = hidden_act
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.max_position_embeddings = max_position_embeddings
+        self.type_vocab_size = type_vocab_size
+        self.type_sequence_label_size = type_sequence_label_size
+        self.initializer_range = initializer_range
+        self.layer_norm_eps = layer_norm_eps
+
+        # Reader config
+        self.span_hidden_size = span_hidden_size
+        self.max_span_width = max_span_width
+        self.reader_layer_norm_eps = reader_layer_norm_eps
+        self.reader_beam_size = reader_beam_size
+        self.reader_seq_len = reader_seq_len
+
+        # Searcher config
+        self.num_block_records = num_block_records
+        self.searcher_beam_size = searcher_beam_size
+        self.searcher_seq_len = searcher_seq_len
+
+        self.num_labels = num_labels
+        self.num_choices = num_choices
+        self.num_candidates = num_candidates
+        self.scope = scope
+
+    def prepare_config_and_inputs(self):
+        input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)
+        candiate_input_ids = ids_tensor([self.batch_size, self.num_candidates, self.seq_length], self.vocab_size)
+        reader_input_ids = ids_tensor([self.reader_beam_size, self.reader_seq_len], self.vocab_size)
+
+        input_mask = None
+        candiate_input_mask = None
+        reader_input_mask = None
+        if self.use_input_mask:
+            input_mask = random_attention_mask([self.batch_size, self.seq_length])
+            candiate_input_mask = random_attention_mask([self.batch_size, self.num_candidates, self.seq_length])
+            reader_input_mask = random_attention_mask([self.reader_beam_size, self.reader_seq_len])
+
+        token_type_ids = None
+        candidate_token_type_ids = None
+        reader_token_type_ids = None
+        if self.use_token_type_ids:
+            token_type_ids = ids_tensor([self.batch_size, self.seq_length], self.type_vocab_size)
+            candidate_token_type_ids = ids_tensor(
+                [self.batch_size, self.num_candidates, self.seq_length], self.type_vocab_size
+            )
+            reader_token_type_ids = ids_tensor([self.reader_beam_size, self.reader_seq_len], self.type_vocab_size)
+
+        sequence_labels = None
+        token_labels = None
+        choice_labels = None
+        if self.use_labels:
+            sequence_labels = ids_tensor([self.batch_size], self.type_sequence_label_size)
+            token_labels = ids_tensor([self.batch_size, self.seq_length], self.num_labels)
+            choice_labels = ids_tensor([self.batch_size], self.num_choices)
+
+        config = self.get_config()
+
+        # inputs with additional num_candidates axis.
+        scorer_encoder_inputs = (candiate_input_ids, candiate_input_mask, candidate_token_type_ids)
+        # reader inputs
+        reader_inputs = (reader_input_ids, reader_input_mask, reader_token_type_ids)
+
+        return (
+            config,
+            input_ids,
+            token_type_ids,
+            input_mask,
+            scorer_encoder_inputs,
+            reader_inputs,
+            sequence_labels,
+            token_labels,
+            choice_labels,
+        )
+
+    def get_config(self):
+        return RealmConfig(
+            vocab_size=self.vocab_size,
+            hidden_size=self.hidden_size,
+            retriever_proj_size=self.retriever_proj_size,
+            num_hidden_layers=self.num_hidden_layers,
+            num_attention_heads=self.num_attention_heads,
+            num_candidates=self.num_candidates,
+            intermediate_size=self.intermediate_size,
+            hidden_act=self.hidden_act,
+            hidden_dropout_prob=self.hidden_dropout_prob,
+            attention_probs_dropout_prob=self.attention_probs_dropout_prob,
+            max_position_embeddings=self.max_position_embeddings,
+            type_vocab_size=self.type_vocab_size,
+            initializer_range=self.initializer_range,
+        )
+
+    def create_and_check_embedder(
+        self,
+        config,
+        input_ids,
+        token_type_ids,
+        input_mask,
+        scorer_encoder_inputs,
+        reader_inputs,
+        sequence_labels,
+        token_labels,
+        choice_labels,
+    ):
+        model = RealmEmbedder(config=config)
+        model.to(torch_device)
+        model.eval()
+        result = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids)
+        self.parent.assertEqual(result.projected_score.shape, (self.batch_size, self.retriever_proj_size))
+
+    def create_and_check_encoder(
+        self,
+        config,
+        input_ids,
+        token_type_ids,
+        input_mask,
+        scorer_encoder_inputs,
+        reader_inputs,
+        sequence_labels,
+        token_labels,
+        choice_labels,
+    ):
+        model = RealmKnowledgeAugEncoder(config=config)
+        model.to(torch_device)
+        model.eval()
+        relevance_score = floats_tensor([self.batch_size, self.num_candidates])
+        result = model(
+            scorer_encoder_inputs[0],
+            attention_mask=scorer_encoder_inputs[1],
+            token_type_ids=scorer_encoder_inputs[2],
+            relevance_score=relevance_score,
+            labels=token_labels,
+        )
+        self.parent.assertEqual(
+            result.logits.shape, (self.batch_size * self.num_candidates, self.seq_length, self.vocab_size)
+        )
+
+    def create_and_check_reader(
+        self,
+        config,
+        input_ids,
+        token_type_ids,
+        input_mask,
+        scorer_encoder_inputs,
+        reader_inputs,
+        sequence_labels,
+        token_labels,
+        choice_labels,
+    ):
+        model = RealmReader(config=config)
+        model.to(torch_device)
+        model.eval()
+        relevance_score = floats_tensor([self.reader_beam_size])
+        result = model(
+            reader_inputs[0],
+            attention_mask=reader_inputs[1],
+            token_type_ids=reader_inputs[2],
+            relevance_score=relevance_score,
+        )
+        self.parent.assertEqual(result.block_idx.shape, ())
+        self.parent.assertEqual(result.candidate.shape, ())
+        self.parent.assertEqual(result.start_pos.shape, ())
+        self.parent.assertEqual(result.end_pos.shape, ())
+
+    def create_and_check_scorer(
+        self,
+        config,
+        input_ids,
+        token_type_ids,
+        input_mask,
+        scorer_encoder_inputs,
+        reader_inputs,
+        sequence_labels,
+        token_labels,
+        choice_labels,
+    ):
+        model = RealmScorer(config=config)
+        model.to(torch_device)
+        model.eval()
+        result = model(
+            input_ids,
+            attention_mask=input_mask,
+            token_type_ids=token_type_ids,
+            candidate_input_ids=scorer_encoder_inputs[0],
+            candidate_attention_mask=scorer_encoder_inputs[1],
+            candidate_token_type_ids=scorer_encoder_inputs[2],
+        )
+        self.parent.assertEqual(result.relevance_score.shape, (self.batch_size, self.num_candidates))
+        self.parent.assertEqual(result.query_score.shape, (self.batch_size, self.retriever_proj_size))
+        self.parent.assertEqual(
+            result.candidate_score.shape, (self.batch_size, self.num_candidates, self.retriever_proj_size)
+        )
+
+    def prepare_config_and_inputs_for_common(self):
+        config_and_inputs = self.prepare_config_and_inputs()
+        (
+            config,
+            input_ids,
+            token_type_ids,
+            input_mask,
+            scorer_encoder_inputs,
+            reader_inputs,
+            sequence_labels,
+            token_labels,
+            choice_labels,
+        ) = config_and_inputs
+        inputs_dict = {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
+        return config, inputs_dict
+
+
+@require_torch
+class RealmModelTest(ModelTesterMixin, unittest.TestCase):
+
+    all_model_classes = (
+        (
+            RealmEmbedder,
+            RealmKnowledgeAugEncoder,
+            # RealmScorer is excluded from common tests as it is a container model
+            # consisting of two RealmEmbedders & a simple inner product calculation.
+            # RealmScorer
+        )
+        if is_torch_available()
+        else ()
+    )
+    all_generative_model_classes = ()
+
+    # disable these tests because there is no base_model in Realm
+    test_save_load_fast_init_from_base = False
+    test_save_load_fast_init_to_base = False
+
+    def setUp(self):
+        self.test_pruning = False
+        self.model_tester = RealmModelTester(self)
+        self.config_tester = ConfigTester(self, config_class=RealmConfig)
+
+    def test_config(self):
+        self.config_tester.run_common_tests()
+
+    def test_embedder(self):
+        config_and_inputs = self.model_tester.prepare_config_and_inputs()
+        self.model_tester.create_and_check_embedder(*config_and_inputs)
+
+    def test_encoder(self):
+        config_and_inputs = self.model_tester.prepare_config_and_inputs()
+        self.model_tester.create_and_check_encoder(*config_and_inputs)
+
+    def test_model_various_embeddings(self):
+        config_and_inputs = self.model_tester.prepare_config_and_inputs()
+        for type in ["absolute", "relative_key", "relative_key_query"]:
+            config_and_inputs[0].position_embedding_type = type
+            self.model_tester.create_and_check_embedder(*config_and_inputs)
+            self.model_tester.create_and_check_encoder(*config_and_inputs)
+
+    def test_retriever(self):
+        config_and_inputs = self.model_tester.prepare_config_and_inputs()
+        self.model_tester.create_and_check_scorer(*config_and_inputs)
+
+    def test_training(self):
+        if not self.model_tester.is_training:
+            return
+
+        config, *inputs = self.model_tester.prepare_config_and_inputs()
+        input_ids, token_type_ids, input_mask, scorer_encoder_inputs = inputs[0:4]
+        config.return_dict = True
+
+        tokenizer = RealmTokenizer.from_pretrained("qqaatw/realm-orqa-nq-openqa")
+
+        # RealmKnowledgeAugEncoder training
+        model = RealmKnowledgeAugEncoder(config)
+        model.to(torch_device)
+        model.train()
+
+        inputs_dict = {
+            "input_ids": scorer_encoder_inputs[0].to(torch_device),
+            "attention_mask": scorer_encoder_inputs[1].to(torch_device),
+            "token_type_ids": scorer_encoder_inputs[2].to(torch_device),
+            "relevance_score": floats_tensor([self.model_tester.batch_size, self.model_tester.num_candidates]),
+        }
+        inputs_dict["labels"] = torch.zeros(
+            (self.model_tester.batch_size, self.model_tester.seq_length), dtype=torch.long, device=torch_device
+        )
+        inputs = inputs_dict
+        loss = model(**inputs).loss
+        loss.backward()
+
+        # RealmForOpenQA training
+        openqa_config = copy.deepcopy(config)
+        openqa_config.vocab_size = 30522  # the retrieved texts will inevitably have more than 99 vocabs.
+        openqa_config.num_block_records = 5
+        openqa_config.searcher_beam_size = 2
+
+        block_records = np.array(
+            [
+                b"This is the first record.",
+                b"This is the second record.",
+                b"This is the third record.",
+                b"This is the fourth record.",
+                b"This is the fifth record.",
+            ],
+            dtype=np.object,
+        )
+        retriever = RealmRetriever(block_records, tokenizer)
+        model = RealmForOpenQA(openqa_config, retriever)
+        model.to(torch_device)
+        model.train()
+
+        inputs_dict = {
+            "input_ids": input_ids[:1].to(torch_device),
+            "attention_mask": input_mask[:1].to(torch_device),
+            "token_type_ids": token_type_ids[:1].to(torch_device),
+            "answer_ids": input_ids[:1].tolist(),
+        }
+        inputs = self._prepare_for_class(inputs_dict, RealmForOpenQA)
+        loss = model(**inputs).reader_output.loss
+        loss.backward()
+
+    @slow
+    def test_embedder_from_pretrained(self):
+        model = RealmEmbedder.from_pretrained("qqaatw/realm-cc-news-pretrained-embedder")
+        self.assertIsNotNone(model)
+
+    @slow
+    def test_encoder_from_pretrained(self):
+        model = RealmKnowledgeAugEncoder.from_pretrained("qqaatw/realm-cc-news-pretrained-encoder")
+        self.assertIsNotNone(model)
+
+    @slow
+    def test_open_qa_from_pretrained(self):
+        model = RealmForOpenQA.from_pretrained("qqaatw/realm-orqa-nq-openqa")
+        self.assertIsNotNone(model)
+
+    @slow
+    def test_reader_from_pretrained(self):
+        model = RealmReader.from_pretrained("qqaatw/realm-orqa-nq-reader")
+        self.assertIsNotNone(model)
+
+    @slow
+    def test_scorer_from_pretrained(self):
+        model = RealmScorer.from_pretrained("qqaatw/realm-cc-news-pretrained-scorer")
+        self.assertIsNotNone(model)
+
+
+@require_torch
+class RealmModelIntegrationTest(unittest.TestCase):
+    @slow
+    def test_inference_embedder(self):
+        retriever_projected_size = 128
+
+        model = RealmEmbedder.from_pretrained("qqaatw/realm-cc-news-pretrained-embedder")
+        input_ids = torch.tensor([[0, 1, 2, 3, 4, 5]])
+        output = model(input_ids)[0]
+
+        expected_shape = torch.Size((1, retriever_projected_size))
+        self.assertEqual(output.shape, expected_shape)
+
+        expected_slice = torch.tensor([[-0.0714, -0.0837, -0.1314]])
+        self.assertTrue(torch.allclose(output[:, :3], expected_slice, atol=1e-4))
+
+    @slow
+    def test_inference_encoder(self):
+        num_candidates = 2
+        vocab_size = 30522
+
+        model = RealmKnowledgeAugEncoder.from_pretrained(
+            "qqaatw/realm-cc-news-pretrained-encoder", num_candidates=num_candidates
+        )
+        input_ids = torch.tensor([[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11]])
+        relevance_score = torch.tensor([[0.3, 0.7]], dtype=torch.float32)
+        output = model(input_ids, relevance_score=relevance_score)[0]
+
+        expected_shape = torch.Size((2, 6, vocab_size))
+        self.assertEqual(output.shape, expected_shape)
+
+        expected_slice = torch.tensor([[[-11.0888, -11.2544], [-10.2170, -10.3874]]])
+
+        self.assertTrue(torch.allclose(output[1, :2, :2], expected_slice, atol=1e-4))
+
+    @slow
+    def test_inference_open_qa(self):
+        from transformers.models.realm.retrieval_realm import RealmRetriever
+
+        config = RealmConfig()
+
+        tokenizer = RealmTokenizer.from_pretrained("qqaatw/realm-orqa-nq-openqa")
+        retriever = RealmRetriever.from_pretrained("qqaatw/realm-orqa-nq-openqa")
+
+        model = RealmForOpenQA.from_pretrained(
+            "qqaatw/realm-orqa-nq-openqa",
+            retriever=retriever,
+            config=config,
+        )
+
+        question = "Who is the pioneer in modern computer science?"
+
+        question = tokenizer(
+            [question],
+            padding=True,
+            truncation=True,
+            max_length=model.config.searcher_seq_len,
+            return_tensors="pt",
+        ).to(model.device)
+
+        predicted_answer_ids = model(**question).predicted_answer_ids
+
+        predicted_answer = tokenizer.decode(predicted_answer_ids)
+        self.assertEqual(predicted_answer, "alan mathison turing")
+
+    @slow
+    def test_inference_reader(self):
+        config = RealmConfig(reader_beam_size=2, max_span_width=3)
+        model = RealmReader.from_pretrained("qqaatw/realm-orqa-nq-reader", config=config)
+
+        concat_input_ids = torch.arange(10).view((2, 5))
+        concat_token_type_ids = torch.tensor([[0, 0, 1, 1, 1], [0, 0, 1, 1, 1]], dtype=torch.int64)
+        relevance_score = torch.tensor([0.3, 0.7], dtype=torch.float32)
+
+        output = model(
+            concat_input_ids, token_type_ids=concat_token_type_ids, relevance_score=relevance_score, return_dict=True
+        )
+
+        block_idx_expected_shape = torch.Size(())
+        start_pos_expected_shape = torch.Size((1,))
+        end_pos_expected_shape = torch.Size((1,))
+        self.assertEqual(output.block_idx.shape, block_idx_expected_shape)
+        self.assertEqual(output.start_pos.shape, start_pos_expected_shape)
+        self.assertEqual(output.end_pos.shape, end_pos_expected_shape)
+
+        expected_block_idx = torch.tensor(1)
+        expected_start_pos = torch.tensor(3)
+        expected_end_pos = torch.tensor(3)
+
+        self.assertTrue(torch.allclose(output.block_idx, expected_block_idx, atol=1e-4))
+        self.assertTrue(torch.allclose(output.start_pos, expected_start_pos, atol=1e-4))
+        self.assertTrue(torch.allclose(output.end_pos, expected_end_pos, atol=1e-4))
+
+    @slow
+    def test_inference_scorer(self):
+        num_candidates = 2
+
+        model = RealmScorer.from_pretrained("qqaatw/realm-cc-news-pretrained-scorer", num_candidates=num_candidates)
+
+        input_ids = torch.tensor([[0, 1, 2, 3, 4, 5]])
+        candidate_input_ids = torch.tensor([[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11]])
+        output = model(input_ids, candidate_input_ids=candidate_input_ids)[0]
+
+        expected_shape = torch.Size((1, 2))
+        self.assertEqual(output.shape, expected_shape)
+
+        expected_slice = torch.tensor([[0.7410, 0.7170]])
+        self.assertTrue(torch.allclose(output, expected_slice, atol=1e-4))

tests/test_retrieval_realm.py

@@ -0,0 +1,185 @@
+# coding=utf-8
+# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+
+import os
+import shutil
+import tempfile
+from unittest import TestCase
+from unittest.mock import patch
+
+import numpy as np
+from datasets import Dataset
+
+from transformers.models.realm.configuration_realm import RealmConfig
+from transformers.models.realm.retrieval_realm import _REALM_BLOCK_RECORDS_FILENAME, RealmRetriever
+from transformers.models.realm.tokenization_realm import VOCAB_FILES_NAMES, RealmTokenizer
+
+
+class RealmRetrieverTest(TestCase):
+    def setUp(self):
+        self.tmpdirname = tempfile.mkdtemp()
+        self.num_block_records = 5
+
+        # Realm tok
+        vocab_tokens = [
+            "[UNK]",
+            "[CLS]",
+            "[SEP]",
+            "[PAD]",
+            "[MASK]",
+            "test",
+            "question",
+            "this",
+            "is",
+            "the",
+            "first",
+            "second",
+            "third",
+            "fourth",
+            "fifth",
+            "record",
+            "want",
+            "##want",
+            "##ed",
+            "wa",
+            "un",
+            "runn",
+            "##ing",
+            ",",
+            "low",
+            "lowest",
+        ]
+        realm_tokenizer_path = os.path.join(self.tmpdirname, "realm_tokenizer")
+        os.makedirs(realm_tokenizer_path, exist_ok=True)
+        self.vocab_file = os.path.join(realm_tokenizer_path, VOCAB_FILES_NAMES["vocab_file"])
+        with open(self.vocab_file, "w", encoding="utf-8") as vocab_writer:
+            vocab_writer.write("".join([x + "\n" for x in vocab_tokens]))
+
+        realm_block_records_path = os.path.join(self.tmpdirname, "realm_block_records")
+        os.makedirs(realm_block_records_path, exist_ok=True)
+
+    def get_tokenizer(self) -> RealmTokenizer:
+        return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname, "realm_tokenizer"))
+
+    def tearDown(self):
+        shutil.rmtree(self.tmpdirname)
+
+    def get_config(self):
+        config = RealmConfig(num_block_records=self.num_block_records)
+        return config
+
+    def get_dummy_dataset(self):
+        dataset = Dataset.from_dict(
+            {
+                "id": ["0", "1"],
+                "question": ["foo", "bar"],
+                "answers": [["Foo", "Bar"], ["Bar"]],
+            }
+        )
+        return dataset
+
+    def get_dummy_block_records(self):
+        block_records = np.array(
+            [
+                b"This is the first record",
+                b"This is the second record",
+                b"This is the third record",
+                b"This is the fourth record",
+                b"This is the fifth record",
+            ],
+            dtype=np.object,
+        )
+        return block_records
+
+    def get_dummy_retriever(self):
+        retriever = RealmRetriever(
+            block_records=self.get_dummy_block_records(),
+            tokenizer=self.get_tokenizer(),
+        )
+        return retriever
+
+    def test_retrieve(self):
+        config = self.get_config()
+        retriever = self.get_dummy_retriever()
+        tokenizer = retriever.tokenizer
+
+        retrieved_block_ids = np.array([0, 3], dtype=np.long)
+        question_input_ids = tokenizer(["Test question"]).input_ids
+        answer_ids = tokenizer(
+            ["the fourth"],
+            add_special_tokens=False,
+            return_token_type_ids=False,
+            return_attention_mask=False,
+        ).input_ids
+        max_length = config.reader_seq_len
+
+        has_answers, start_pos, end_pos, concat_inputs = retriever(
+            retrieved_block_ids, question_input_ids, answer_ids=answer_ids, max_length=max_length, return_tensors="np"
+        )
+
+        self.assertEqual(len(has_answers), 2)
+        self.assertEqual(len(start_pos), 2)
+        self.assertEqual(len(end_pos), 2)
+        self.assertEqual(concat_inputs.input_ids.shape, (2, 10))
+        self.assertEqual(concat_inputs.attention_mask.shape, (2, 10))
+        self.assertEqual(concat_inputs.token_type_ids.shape, (2, 10))
+        self.assertEqual(
+            tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[0]),
+            ["[CLS]", "test", "question", "[SEP]", "this", "is", "the", "first", "record", "[SEP]"],
+        )
+        self.assertEqual(
+            tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[1]),
+            ["[CLS]", "test", "question", "[SEP]", "this", "is", "the", "fourth", "record", "[SEP]"],
+        )
+
+    def test_block_has_answer(self):
+        config = self.get_config()
+        retriever = self.get_dummy_retriever()
+        tokenizer = retriever.tokenizer
+
+        retrieved_block_ids = np.array([0, 3], dtype=np.long)
+        question_input_ids = tokenizer(["Test question"]).input_ids
+        answer_ids = tokenizer(
+            ["the fourth"],
+            add_special_tokens=False,
+            return_token_type_ids=False,
+            return_attention_mask=False,
+        ).input_ids
+        max_length = config.reader_seq_len
+
+        has_answers, start_pos, end_pos, _ = retriever(
+            retrieved_block_ids, question_input_ids, answer_ids=answer_ids, max_length=max_length, return_tensors="np"
+        )
+
+        self.assertEqual([False, True], has_answers)
+        self.assertEqual([[-1], [6]], start_pos)
+        self.assertEqual([[-1], [7]], end_pos)
+
+    def test_save_load_pretrained(self):
+        retriever = self.get_dummy_retriever()
+        retriever.save_pretrained(os.path.join(self.tmpdirname, "realm_block_records"))
+
+        # Test local path
+        retriever = retriever.from_pretrained(os.path.join(self.tmpdirname, "realm_block_records"))
+        self.assertEqual(retriever.block_records[0], b"This is the first record")
+
+        # Test mocked remote path
+        with patch("transformers.models.realm.retrieval_realm.hf_hub_download") as mock_hf_hub_download:
+            mock_hf_hub_download.return_value = os.path.join(
+                os.path.join(self.tmpdirname, "realm_block_records"), _REALM_BLOCK_RECORDS_FILENAME
+            )
+            retriever = RealmRetriever.from_pretrained("qqaatw/realm-cc-news-pretrained-openqa")
+
+        self.assertEqual(retriever.block_records[0], b"This is the first record")

tests/test_tokenization_realm.py

@@ -0,0 +1,314 @@
+# coding=utf-8
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# 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.
+
+import os
+import unittest
+
+from transformers.models.bert.tokenization_bert import (
+    VOCAB_FILES_NAMES,
+    BasicTokenizer,
+    WordpieceTokenizer,
+    _is_control,
+    _is_punctuation,
+    _is_whitespace,
+)
+from transformers.models.realm.tokenization_realm import RealmTokenizer
+from transformers.testing_utils import require_tokenizers, slow
+
+from .test_tokenization_common import TokenizerTesterMixin, filter_non_english
+
+
+@require_tokenizers
+class RealmTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
+
+    tokenizer_class = RealmTokenizer
+    rust_tokenizer_class = None
+    test_rust_tokenizer = False
+    space_between_special_tokens = True
+    from_pretrained_filter = filter_non_english
+
+    def setUp(self):
+        super().setUp()
+
+        vocab_tokens = [
+            "[UNK]",
+            "[CLS]",
+            "[SEP]",
+            "[PAD]",
+            "[MASK]",
+            "want",
+            "##want",
+            "##ed",
+            "wa",
+            "un",
+            "runn",
+            "##ing",
+            ",",
+            "low",
+            "lowest",
+        ]
+        self.vocab_file = os.path.join(self.tmpdirname, VOCAB_FILES_NAMES["vocab_file"])
+        with open(self.vocab_file, "w", encoding="utf-8") as vocab_writer:
+            vocab_writer.write("".join([x + "\n" for x in vocab_tokens]))
+
+    def get_input_output_texts(self, tokenizer):
+        input_text = "UNwant\u00E9d,running"
+        output_text = "unwanted, running"
+        return input_text, output_text
+
+    def test_full_tokenizer(self):
+        tokenizer = self.tokenizer_class(self.vocab_file)
+
+        tokens = tokenizer.tokenize("UNwant\u00E9d,running")
+        self.assertListEqual(tokens, ["un", "##want", "##ed", ",", "runn", "##ing"])
+        self.assertListEqual(tokenizer.convert_tokens_to_ids(tokens), [9, 6, 7, 12, 10, 11])
+
+    def test_rust_and_python_full_tokenizers(self):
+        if not self.test_rust_tokenizer:
+            return
+
+        tokenizer = self.get_tokenizer()
+        rust_tokenizer = self.get_rust_tokenizer()
+
+        sequence = "UNwant\u00E9d,running"
+
+        tokens = tokenizer.tokenize(sequence)
+        rust_tokens = rust_tokenizer.tokenize(sequence)
+        self.assertListEqual(tokens, rust_tokens)
+
+        ids = tokenizer.encode(sequence, add_special_tokens=False)
+        rust_ids = rust_tokenizer.encode(sequence, add_special_tokens=False)
+        self.assertListEqual(ids, rust_ids)
+
+        rust_tokenizer = self.get_rust_tokenizer()
+        ids = tokenizer.encode(sequence)
+        rust_ids = rust_tokenizer.encode(sequence)
+        self.assertListEqual(ids, rust_ids)
+
+        # With lower casing
+        tokenizer = self.get_tokenizer(do_lower_case=True)
+        rust_tokenizer = self.get_rust_tokenizer(do_lower_case=True)
+
+        sequence = "UNwant\u00E9d,running"
+
+        tokens = tokenizer.tokenize(sequence)
+        rust_tokens = rust_tokenizer.tokenize(sequence)
+        self.assertListEqual(tokens, rust_tokens)
+
+        ids = tokenizer.encode(sequence, add_special_tokens=False)
+        rust_ids = rust_tokenizer.encode(sequence, add_special_tokens=False)
+        self.assertListEqual(ids, rust_ids)
+
+        rust_tokenizer = self.get_rust_tokenizer()
+        ids = tokenizer.encode(sequence)
+        rust_ids = rust_tokenizer.encode(sequence)
+        self.assertListEqual(ids, rust_ids)
+
+    def test_chinese(self):
+        tokenizer = BasicTokenizer()
+
+        self.assertListEqual(tokenizer.tokenize("ah\u535A\u63A8zz"), ["ah", "\u535A", "\u63A8", "zz"])
+
+    def test_basic_tokenizer_lower(self):
+        tokenizer = BasicTokenizer(do_lower_case=True)
+
+        self.assertListEqual(
+            tokenizer.tokenize(" \tHeLLo!how  \n Are yoU?  "), ["hello", "!", "how", "are", "you", "?"]
+        )
+        self.assertListEqual(tokenizer.tokenize("H\u00E9llo"), ["hello"])
+
+    def test_basic_tokenizer_lower_strip_accents_false(self):
+        tokenizer = BasicTokenizer(do_lower_case=True, strip_accents=False)
+
+        self.assertListEqual(
+            tokenizer.tokenize(" \tHäLLo!how  \n Are yoU?  "), ["hällo", "!", "how", "are", "you", "?"]
+        )
+        self.assertListEqual(tokenizer.tokenize("H\u00E9llo"), ["h\u00E9llo"])
+
+    def test_basic_tokenizer_lower_strip_accents_true(self):
+        tokenizer = BasicTokenizer(do_lower_case=True, strip_accents=True)
+
+        self.assertListEqual(
+            tokenizer.tokenize(" \tHäLLo!how  \n Are yoU?  "), ["hallo", "!", "how", "are", "you", "?"]
+        )
+        self.assertListEqual(tokenizer.tokenize("H\u00E9llo"), ["hello"])
+
+    def test_basic_tokenizer_lower_strip_accents_default(self):
+        tokenizer = BasicTokenizer(do_lower_case=True)
+
+        self.assertListEqual(
+            tokenizer.tokenize(" \tHäLLo!how  \n Are yoU?  "), ["hallo", "!", "how", "are", "you", "?"]
+        )
+        self.assertListEqual(tokenizer.tokenize("H\u00E9llo"), ["hello"])
+
+    def test_basic_tokenizer_no_lower(self):
+        tokenizer = BasicTokenizer(do_lower_case=False)
+
+        self.assertListEqual(
+            tokenizer.tokenize(" \tHeLLo!how  \n Are yoU?  "), ["HeLLo", "!", "how", "Are", "yoU", "?"]
+        )
+
+    def test_basic_tokenizer_no_lower_strip_accents_false(self):
+        tokenizer = BasicTokenizer(do_lower_case=False, strip_accents=False)
+
+        self.assertListEqual(
+            tokenizer.tokenize(" \tHäLLo!how  \n Are yoU?  "), ["HäLLo", "!", "how", "Are", "yoU", "?"]
+        )
+
+    def test_basic_tokenizer_no_lower_strip_accents_true(self):
+        tokenizer = BasicTokenizer(do_lower_case=False, strip_accents=True)
+
+        self.assertListEqual(
+            tokenizer.tokenize(" \tHäLLo!how  \n Are yoU?  "), ["HaLLo", "!", "how", "Are", "yoU", "?"]
+        )
+
+    def test_basic_tokenizer_respects_never_split_tokens(self):
+        tokenizer = BasicTokenizer(do_lower_case=False, never_split=["[UNK]"])
+
+        self.assertListEqual(
+            tokenizer.tokenize(" \tHeLLo!how  \n Are yoU? [UNK]"), ["HeLLo", "!", "how", "Are", "yoU", "?", "[UNK]"]
+        )
+
+    def test_wordpiece_tokenizer(self):
+        vocab_tokens = ["[UNK]", "[CLS]", "[SEP]", "want", "##want", "##ed", "wa", "un", "runn", "##ing"]
+
+        vocab = {}
+        for (i, token) in enumerate(vocab_tokens):
+            vocab[token] = i
+        tokenizer = WordpieceTokenizer(vocab=vocab, unk_token="[UNK]")
+
+        self.assertListEqual(tokenizer.tokenize(""), [])
+
+        self.assertListEqual(tokenizer.tokenize("unwanted running"), ["un", "##want", "##ed", "runn", "##ing"])
+
+        self.assertListEqual(tokenizer.tokenize("unwantedX running"), ["[UNK]", "runn", "##ing"])
+
+    def test_is_whitespace(self):
+        self.assertTrue(_is_whitespace(" "))
+        self.assertTrue(_is_whitespace("\t"))
+        self.assertTrue(_is_whitespace("\r"))
+        self.assertTrue(_is_whitespace("\n"))
+        self.assertTrue(_is_whitespace("\u00A0"))
+
+        self.assertFalse(_is_whitespace("A"))
+        self.assertFalse(_is_whitespace("-"))
+
+    def test_is_control(self):
+        self.assertTrue(_is_control("\u0005"))
+
+        self.assertFalse(_is_control("A"))
+        self.assertFalse(_is_control(" "))
+        self.assertFalse(_is_control("\t"))
+        self.assertFalse(_is_control("\r"))
+
+    def test_is_punctuation(self):
+        self.assertTrue(_is_punctuation("-"))
+        self.assertTrue(_is_punctuation("$"))
+        self.assertTrue(_is_punctuation("`"))
+        self.assertTrue(_is_punctuation("."))
+
+        self.assertFalse(_is_punctuation("A"))
+        self.assertFalse(_is_punctuation(" "))
+
+    def test_clean_text(self):
+        tokenizer = self.get_tokenizer()
+
+        # Example taken from the issue https://github.com/huggingface/tokenizers/issues/340
+        self.assertListEqual([tokenizer.tokenize(t) for t in ["Test", "\xad", "test"]], [["[UNK]"], [], ["[UNK]"]])
+
+        if self.test_rust_tokenizer:
+            rust_tokenizer = self.get_rust_tokenizer()
+            self.assertListEqual(
+                [rust_tokenizer.tokenize(t) for t in ["Test", "\xad", "test"]], [["[UNK]"], [], ["[UNK]"]]
+            )
+
+    @slow
+    def test_sequence_builders(self):
+        tokenizer = self.tokenizer_class.from_pretrained("bert-base-uncased")
+
+        text = tokenizer.encode("sequence builders", add_special_tokens=False)
+        text_2 = tokenizer.encode("multi-sequence build", add_special_tokens=False)
+
+        encoded_sentence = tokenizer.build_inputs_with_special_tokens(text)
+        encoded_pair = tokenizer.build_inputs_with_special_tokens(text, text_2)
+
+        assert encoded_sentence == [101] + text + [102]
+        assert encoded_pair == [101] + text + [102] + text_2 + [102]
+
+    def test_offsets_with_special_characters(self):
+        for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
+            with self.subTest(f"{tokenizer.__class__.__name__} ({pretrained_name})"):
+                tokenizer_r = self.rust_tokenizer_class.from_pretrained(pretrained_name, **kwargs)
+
+                sentence = f"A, naïve {tokenizer_r.mask_token} AllenNLP sentence."
+                tokens = tokenizer_r.encode_plus(
+                    sentence,
+                    return_attention_mask=False,
+                    return_token_type_ids=False,
+                    return_offsets_mapping=True,
+                    add_special_tokens=True,
+                )
+
+                do_lower_case = tokenizer_r.do_lower_case if hasattr(tokenizer_r, "do_lower_case") else False
+                expected_results = (
+                    [
+                        ((0, 0), tokenizer_r.cls_token),
+                        ((0, 1), "A"),
+                        ((1, 2), ","),
+                        ((3, 5), "na"),
+                        ((5, 6), "##ï"),
+                        ((6, 8), "##ve"),
+                        ((9, 15), tokenizer_r.mask_token),
+                        ((16, 21), "Allen"),
+                        ((21, 23), "##NL"),
+                        ((23, 24), "##P"),
+                        ((25, 33), "sentence"),
+                        ((33, 34), "."),
+                        ((0, 0), tokenizer_r.sep_token),
+                    ]
+                    if not do_lower_case
+                    else [
+                        ((0, 0), tokenizer_r.cls_token),
+                        ((0, 1), "a"),
+                        ((1, 2), ","),
+                        ((3, 8), "naive"),
+                        ((9, 15), tokenizer_r.mask_token),
+                        ((16, 21), "allen"),
+                        ((21, 23), "##nl"),
+                        ((23, 24), "##p"),
+                        ((25, 33), "sentence"),
+                        ((33, 34), "."),
+                        ((0, 0), tokenizer_r.sep_token),
+                    ]
+                )
+
+                self.assertEqual(
+                    [e[1] for e in expected_results], tokenizer_r.convert_ids_to_tokens(tokens["input_ids"])
+                )
+                self.assertEqual([e[0] for e in expected_results], tokens["offset_mapping"])
+
+    @slow
+    def test_batch_encode_candidates(self):
+        tokenizer = self.tokenizer_class.from_pretrained("bert-base-uncased")
+
+        text = [["Hello world!", "Nice to meet you!"], ["The cute cat.", "The adorable dog."]]
+
+        encoded_sentence = tokenizer.batch_encode_candidates(text, max_length=10, return_tensors="pt")
+
+        expected_shape = (2, 2, 10)
+
+        assert encoded_sentence["input_ids"].shape == expected_shape
+        assert encoded_sentence["attention_mask"].shape == expected_shape
+        assert encoded_sentence["token_type_ids"].shape == expected_shape