





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       (       _UpperCamelCase       ):

  """simple docstring"""


  def 							lowercase						(							self					:						Any	)				->		Tuple:
     __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(lowerCAmelCase_   ,    exist_ok=lowerCAmelCase_	)
     __lowerCAmelCase				      =      os.path.join(lowerCAmelCase_   ,    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(lowerCAmelCase_   ,    exist_ok=lowerCAmelCase_	)






  def 							lowercase						(							self					:						Optional[int]	)				->		RealmTokenizer:
     return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname   ,    'realm_tokenizer'	)	)






  def 							lowercase						(							self					:						List[str]	)				->		str:
     shutil.rmtree(self.tmpdirname	)






  def 							lowercase						(							self					:						List[str]	)				->		int:
     __lowerCAmelCase				      =      RealmConfig(num_block_records=self.num_block_records	)
     return config






  def 							lowercase						(							self					:						str	)				->		str:
     __lowerCAmelCase				      =      Dataset.from_dict(
         {
             'id': ['0', '1'],
             'question': ['foo', 'bar'],
             'answers': [['Foo', 'Bar'], ['Bar']],
         }	)
     return dataset






  def 							lowercase						(							self					:						List[Any]	)				->		Dict:
     __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=lowerCAmelCase_   ,    )
     return block_records






  def 							lowercase						(							self					:						Tuple	)				->		List[Any]:
     __lowerCAmelCase				      =      RealmRetriever(
         block_records=self.get_dummy_block_records()   ,    tokenizer=self.get_tokenizer()   ,    )
     return retriever






  def 							lowercase						(							self					:						Tuple	)				->		Union[str, Any]:
     __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=lowerCAmelCase_   ,    return_token_type_ids=lowerCAmelCase_   ,    return_attention_mask=lowerCAmelCase_   ,    ).input_ids
     __lowerCAmelCase				      =      config.reader_seq_len

     __lowerCAmelCase      ,      __lowerCAmelCase      ,      __lowerCAmelCase      ,      __lowerCAmelCase				      =      retriever(
         lowerCAmelCase_   ,    lowerCAmelCase_   ,    answer_ids=lowerCAmelCase_   ,    max_length=lowerCAmelCase_   ,    return_tensors='np'	)

     self.assertEqual(len(lowerCAmelCase_	)   ,    2	)
     self.assertEqual(len(lowerCAmelCase_	)   ,    2	)
     self.assertEqual(len(lowerCAmelCase_	)   ,    2	)
     self.assertEqual(concat_inputs.input_ids.shape   ,    (2, 1_0)	)
     self.assertEqual(concat_inputs.attention_mask.shape   ,    (2, 1_0)	)
     self.assertEqual(concat_inputs.token_type_ids.shape   ,    (2, 1_0)	)
     self.assertEqual(concat_inputs.special_tokens_mask.shape   ,    (2, 1_0)	)
     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					:						Optional[int]	)				->		Optional[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=lowerCAmelCase_   ,    return_token_type_ids=lowerCAmelCase_   ,    return_attention_mask=lowerCAmelCase_   ,    ).input_ids
     __lowerCAmelCase				      =      config.reader_seq_len

     __lowerCAmelCase      ,      __lowerCAmelCase      ,      __lowerCAmelCase      ,      __lowerCAmelCase				      =      retriever(
         lowerCAmelCase_   ,    lowerCAmelCase_   ,    answer_ids=lowerCAmelCase_   ,    max_length=lowerCAmelCase_   ,    return_tensors='np'	)

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






  def 							lowercase						(							self					:						Union[str, Any]	)				->		List[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_torch_available,
)


_lowercase							:   int						=						{
    "configuration_roberta_prelayernorm": [
        "ROBERTA_PRELAYERNORM_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "RobertaPreLayerNormConfig",
        "RobertaPreLayerNormOnnxConfig",
    ],
}

try:
				if not is_torch_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				_lowercase							:   List[Any]						=						[
				    "ROBERTA_PRELAYERNORM_PRETRAINED_MODEL_ARCHIVE_LIST",
				    "RobertaPreLayerNormForCausalLM",
				    "RobertaPreLayerNormForMaskedLM",
				    "RobertaPreLayerNormForMultipleChoice",
				    "RobertaPreLayerNormForQuestionAnswering",
				    "RobertaPreLayerNormForSequenceClassification",
				    "RobertaPreLayerNormForTokenClassification",
				    "RobertaPreLayerNormModel",
				    "RobertaPreLayerNormPreTrainedModel",
				]

try:
				if not is_tf_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				_lowercase							:   Optional[int]						=						[
				    "TF_ROBERTA_PRELAYERNORM_PRETRAINED_MODEL_ARCHIVE_LIST",
				    "TFRobertaPreLayerNormForCausalLM",
				    "TFRobertaPreLayerNormForMaskedLM",
				    "TFRobertaPreLayerNormForMultipleChoice",
				    "TFRobertaPreLayerNormForQuestionAnswering",
				    "TFRobertaPreLayerNormForSequenceClassification",
				    "TFRobertaPreLayerNormForTokenClassification",
				    "TFRobertaPreLayerNormMainLayer",
				    "TFRobertaPreLayerNormModel",
				    "TFRobertaPreLayerNormPreTrainedModel",
				]

try:
				if not is_flax_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				_lowercase							:   str						=						[
				    "FlaxRobertaPreLayerNormForCausalLM",
				    "FlaxRobertaPreLayerNormForMaskedLM",
				    "FlaxRobertaPreLayerNormForMultipleChoice",
				    "FlaxRobertaPreLayerNormForQuestionAnswering",
				    "FlaxRobertaPreLayerNormForSequenceClassification",
				    "FlaxRobertaPreLayerNormForTokenClassification",
				    "FlaxRobertaPreLayerNormModel",
				    "FlaxRobertaPreLayerNormPreTrainedModel",
				]


if TYPE_CHECKING:
				from .configuration_roberta_prelayernorm import (
				    ROBERTA_PRELAYERNORM_PRETRAINED_CONFIG_ARCHIVE_MAP,
				    RobertaPreLayerNormConfig,
				    RobertaPreLayerNormOnnxConfig,
				)

				try:
								if not is_torch_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_roberta_prelayernorm import (
								    ROBERTA_PRELAYERNORM_PRETRAINED_MODEL_ARCHIVE_LIST,
								    RobertaPreLayerNormForCausalLM,
								    RobertaPreLayerNormForMaskedLM,
								    RobertaPreLayerNormForMultipleChoice,
								    RobertaPreLayerNormForQuestionAnswering,
								    RobertaPreLayerNormForSequenceClassification,
								    RobertaPreLayerNormForTokenClassification,
								    RobertaPreLayerNormModel,
								    RobertaPreLayerNormPreTrainedModel,
								)

				try:
								if not is_tf_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_tf_roberta_prelayernorm import (
								    TF_ROBERTA_PRELAYERNORM_PRETRAINED_MODEL_ARCHIVE_LIST,
								    TFRobertaPreLayerNormForCausalLM,
								    TFRobertaPreLayerNormForMaskedLM,
								    TFRobertaPreLayerNormForMultipleChoice,
								    TFRobertaPreLayerNormForQuestionAnswering,
								    TFRobertaPreLayerNormForSequenceClassification,
								    TFRobertaPreLayerNormForTokenClassification,
								    TFRobertaPreLayerNormMainLayer,
								    TFRobertaPreLayerNormModel,
								    TFRobertaPreLayerNormPreTrainedModel,
								)

				try:
								if not is_flax_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_flax_roberta_prelayernorm import (
								    FlaxRobertaPreLayerNormForCausalLM,
								    FlaxRobertaPreLayerNormForMaskedLM,
								    FlaxRobertaPreLayerNormForMultipleChoice,
								    FlaxRobertaPreLayerNormForQuestionAnswering,
								    FlaxRobertaPreLayerNormForSequenceClassification,
								    FlaxRobertaPreLayerNormForTokenClassification,
								    FlaxRobertaPreLayerNormModel,
								    FlaxRobertaPreLayerNormPreTrainedModel,
								)

else:
				import sys

				_lowercase							:   List[Any]						=						_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)



