





'''simple docstring'''






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  UpperCAmelCase     (			snake_case_   ):




      def 				lowercase__    (			self			: Dict  )				->      Optional[int]:
            _lowerCAmelCase							     =							tempfile.mkdtemp()
            _lowerCAmelCase							     =							5

            # Realm tok
            _lowerCAmelCase							     =							[
                """[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""",
            ]
            _lowerCAmelCase							     =							os.path.join(self.tmpdirname			,			"""realm_tokenizer"""  )
            os.makedirs(__snake_case			,			exist_ok=__snake_case  )
            _lowerCAmelCase							     =							os.path.join(__snake_case			,			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]  )  )

            _lowerCAmelCase							     =							os.path.join(self.tmpdirname			,			"""realm_block_records"""  )
            os.makedirs(__snake_case			,			exist_ok=__snake_case  )




      def 				lowercase__    (			self			: Tuple  )				->      RealmTokenizer:
            return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname			,			"""realm_tokenizer"""  )  )




      def 				lowercase__    (			self			: Any  )				->      Union[str, Any]:
            shutil.rmtree(self.tmpdirname  )




      def 				lowercase__    (			self			: Dict  )				->      Optional[Any]:
            _lowerCAmelCase							     =							RealmConfig(num_block_records=self.num_block_records  )
            return config




      def 				lowercase__    (			self			: List[str]  )				->      Tuple:
            _lowerCAmelCase							     =							Dataset.from_dict(
                {
                    """id""": ["""0""", """1"""],
                    """question""": ["""foo""", """bar"""],
                    """answers""": [["""Foo""", """Bar"""], ["""Bar"""]],
                }  )
            return dataset




      def 				lowercase__    (			self			: Union[str, Any]  )				->      int:
            _lowerCAmelCase							     =							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""",
                    b"""This is a longer longer longer record""",
                ]			,			dtype=__snake_case			,			)
            return block_records




      def 				lowercase__    (			self			: List[str]  )				->      Dict:
            _lowerCAmelCase							     =							RealmRetriever(
                block_records=self.get_dummy_block_records()			,			tokenizer=self.get_tokenizer()			,			)
            return retriever




      def 				lowercase__    (			self			: Optional[Any]  )				->      int:
            _lowerCAmelCase							     =							self.get_config()
            _lowerCAmelCase							     =							self.get_dummy_retriever()
            _lowerCAmelCase							     =							retriever.tokenizer

            _lowerCAmelCase							     =							np.array([0, 3]			,			dtype="""long"""  )
            _lowerCAmelCase							     =							tokenizer(["""Test question"""]  ).input_ids
            _lowerCAmelCase							     =							tokenizer(
                ["""the fourth"""]			,			add_special_tokens=__snake_case			,			return_token_type_ids=__snake_case			,			return_attention_mask=__snake_case			,			).input_ids
            _lowerCAmelCase							     =							config.reader_seq_len

            _lowerCAmelCase     ,							_lowerCAmelCase     ,							_lowerCAmelCase     ,							_lowerCAmelCase							     =							retriever(
                __snake_case			,			__snake_case			,			answer_ids=__snake_case			,			max_length=__snake_case			,			return_tensors="""np"""  )

            self.assertEqual(len(__snake_case  )			,			2  )
            self.assertEqual(len(__snake_case  )			,			2  )
            self.assertEqual(len(__snake_case  )			,			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(concat_inputs.special_tokens_mask.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 				lowercase__    (			self			: List[str]  )				->      int:
            _lowerCAmelCase							     =							self.get_config()
            _lowerCAmelCase							     =							self.get_dummy_retriever()
            _lowerCAmelCase							     =							retriever.tokenizer

            _lowerCAmelCase							     =							np.array([0, 3, 5]			,			dtype="""long"""  )
            _lowerCAmelCase							     =							tokenizer(["""Test question"""]  ).input_ids
            _lowerCAmelCase							     =							tokenizer(
                ["""the fourth""", """longer longer"""]			,			add_special_tokens=__snake_case			,			return_token_type_ids=__snake_case			,			return_attention_mask=__snake_case			,			).input_ids
            _lowerCAmelCase							     =							config.reader_seq_len

            _lowerCAmelCase     ,							_lowerCAmelCase     ,							_lowerCAmelCase     ,							_lowerCAmelCase							     =							retriever(
                __snake_case			,			__snake_case			,			answer_ids=__snake_case			,			max_length=__snake_case			,			return_tensors="""np"""  )

            self.assertEqual([False, True, True]			,			__snake_case  )
            self.assertEqual([[-1, -1, -1], [6, -1, -1], [6, 7, 8]]			,			__snake_case  )
            self.assertEqual([[-1, -1, -1], [7, -1, -1], [7, 8, 9]]			,			__snake_case  )




      def 				lowercase__    (			self			: Any  )				->      Optional[Any]:
            _lowerCAmelCase							     =							self.get_dummy_retriever()
            retriever.save_pretrained(os.path.join(self.tmpdirname			,			"""realm_block_records"""  )  )

            # Test local path
            _lowerCAmelCase							     =							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:
                  _lowerCAmelCase							     =							os.path.join(
                      os.path.join(self.tmpdirname			,			"""realm_block_records"""  )			,			_REALM_BLOCK_RECORDS_FILENAME  )
                  _lowerCAmelCase							     =							RealmRetriever.from_pretrained("""google/realm-cc-news-pretrained-openqa"""  )

            self.assertEqual(retriever.block_records[0]			,			b"""This is the first record"""  )








from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


_A   =     {
    "configuration_distilbert": [
        "DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "DistilBertConfig",
        "DistilBertOnnxConfig",
    ],
    "tokenization_distilbert": ["DistilBertTokenizer"],
}

try:
			if not is_tokenizers_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			_A   =     ["DistilBertTokenizerFast"]

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			_A   =     [
			    "DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST",
			    "DistilBertForMaskedLM",
			    "DistilBertForMultipleChoice",
			    "DistilBertForQuestionAnswering",
			    "DistilBertForSequenceClassification",
			    "DistilBertForTokenClassification",
			    "DistilBertModel",
			    "DistilBertPreTrainedModel",
			]

try:
			if not is_tf_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			_A   =     [
			    "TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST",
			    "TFDistilBertForMaskedLM",
			    "TFDistilBertForMultipleChoice",
			    "TFDistilBertForQuestionAnswering",
			    "TFDistilBertForSequenceClassification",
			    "TFDistilBertForTokenClassification",
			    "TFDistilBertMainLayer",
			    "TFDistilBertModel",
			    "TFDistilBertPreTrainedModel",
			]

try:
			if not is_flax_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			_A   =     [
			    "FlaxDistilBertForMaskedLM",
			    "FlaxDistilBertForMultipleChoice",
			    "FlaxDistilBertForQuestionAnswering",
			    "FlaxDistilBertForSequenceClassification",
			    "FlaxDistilBertForTokenClassification",
			    "FlaxDistilBertModel",
			    "FlaxDistilBertPreTrainedModel",
			]


if TYPE_CHECKING:
			from .configuration_distilbert import (
			    DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
			    DistilBertConfig,
			    DistilBertOnnxConfig,
			)
			from .tokenization_distilbert import DistilBertTokenizer

			try:
						if not is_tokenizers_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .tokenization_distilbert_fast import DistilBertTokenizerFast

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_distilbert import (
						    DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
						    DistilBertForMaskedLM,
						    DistilBertForMultipleChoice,
						    DistilBertForQuestionAnswering,
						    DistilBertForSequenceClassification,
						    DistilBertForTokenClassification,
						    DistilBertModel,
						    DistilBertPreTrainedModel,
						)

			try:
						if not is_tf_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_tf_distilbert import (
						    TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
						    TFDistilBertForMaskedLM,
						    TFDistilBertForMultipleChoice,
						    TFDistilBertForQuestionAnswering,
						    TFDistilBertForSequenceClassification,
						    TFDistilBertForTokenClassification,
						    TFDistilBertMainLayer,
						    TFDistilBertModel,
						    TFDistilBertPreTrainedModel,
						)

			try:
						if not is_flax_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_flax_distilbert import (
						    FlaxDistilBertForMaskedLM,
						    FlaxDistilBertForMultipleChoice,
						    FlaxDistilBertForQuestionAnswering,
						    FlaxDistilBertForSequenceClassification,
						    FlaxDistilBertForTokenClassification,
						    FlaxDistilBertModel,
						    FlaxDistilBertPreTrainedModel,
						)

else:
			import sys

			_A   =     _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)



