





"""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    _lowerCAmelCase  (	_a			):
       """simple docstring"""




       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Union[str, Any]											=				tempfile.mkdtemp()
         lowerCAmelCase__     :Dict											=				5

         # Realm tok
         lowerCAmelCase__     :Tuple											=				[
             """[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__     :Optional[int]											=				os.path.join(self.tmpdirname       ,					'realm_tokenizer' )
         os.makedirs(snake_case_       ,					exist_ok=snake_case_ )
         lowerCAmelCase__     :Union[str, Any]											=				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__     :Dict											=				os.path.join(self.tmpdirname       ,					'realm_block_records' )
         os.makedirs(snake_case_       ,					exist_ok=snake_case_ )






       def 	snake_case		(       self ):



         '''simple docstring'''

         return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname       ,					'realm_tokenizer' ) )






       def 	snake_case		(       self ):



         '''simple docstring'''

         shutil.rmtree(self.tmpdirname )






       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :List[str]											=				RealmConfig(num_block_records=self.num_block_records )
         return config






       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Any											=				Dataset.from_dict(
             {
                 'id': ['0', '1'],
                 'question': ['foo', 'bar'],
                 'answers': [['Foo', 'Bar'], ['Bar']],
             } )
         return dataset






       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Any											=				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 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Any											=				RealmRetriever(
             block_records=self.get_dummy_block_records()       ,					tokenizer=self.get_tokenizer()       ,					)
         return retriever






       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Tuple											=				self.get_config()
         lowerCAmelCase__     :Tuple											=				self.get_dummy_retriever()
         lowerCAmelCase__     :int											=				retriever.tokenizer

         lowerCAmelCase__     :int											=				np.array([0, 3]       ,					dtype='long' )
         lowerCAmelCase__     :int											=				tokenizer(['Test question'] ).input_ids
         lowerCAmelCase__     :Union[str, Any]											=				tokenizer(
             ['the fourth']       ,					add_special_tokens=snake_case_       ,					return_token_type_ids=snake_case_       ,					return_attention_mask=snake_case_       ,					).input_ids
         lowerCAmelCase__     :Optional[int]											=				config.reader_seq_len

         lowerCAmelCase__     :int											=				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, 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 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Tuple											=				self.get_config()
         lowerCAmelCase__     :str											=				self.get_dummy_retriever()
         lowerCAmelCase__     :Tuple											=				retriever.tokenizer

         lowerCAmelCase__     :List[str]											=				np.array([0, 3, 5]       ,					dtype='long' )
         lowerCAmelCase__     :Optional[Any]											=				tokenizer(['Test question'] ).input_ids
         lowerCAmelCase__     :List[Any]											=				tokenizer(
             ['the fourth', 'longer longer']       ,					add_special_tokens=snake_case_       ,					return_token_type_ids=snake_case_       ,					return_attention_mask=snake_case_       ,					).input_ids
         lowerCAmelCase__     :Optional[Any]											=				config.reader_seq_len

         lowerCAmelCase__     :Tuple											=				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 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Tuple											=				self.get_dummy_retriever()
         retriever.save_pretrained(os.path.join(self.tmpdirname       ,					'realm_block_records' ) )

         # Test local path
         lowerCAmelCase__     :Any											=				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__     :Union[str, Any]											=				os.path.join(
               os.path.join(self.tmpdirname       ,					'realm_block_records' )       ,					_REALM_BLOCK_RECORDS_FILENAME )
           lowerCAmelCase__     :Any											=				RealmRetriever.from_pretrained('google/realm-cc-news-pretrained-openqa' )

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





'''simple docstring'''





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


__a  			=						{
    "configuration_gpt_bigcode": ["GPT_BIGCODE_PRETRAINED_CONFIG_ARCHIVE_MAP", "GPTBigCodeConfig"],
}

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __a  			=						[
           "GPT_BIGCODE_PRETRAINED_MODEL_ARCHIVE_LIST",
           "GPTBigCodeForSequenceClassification",
           "GPTBigCodeForTokenClassification",
           "GPTBigCodeForCausalLM",
           "GPTBigCodeModel",
           "GPTBigCodePreTrainedModel",
       ]

if TYPE_CHECKING:
       from .configuration_gpt_bigcode import GPT_BIGCODE_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTBigCodeConfig

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_gpt_bigcode import (
                  GPT_BIGCODE_PRETRAINED_MODEL_ARCHIVE_LIST,
                  GPTBigCodeForCausalLM,
                  GPTBigCodeForSequenceClassification,
                  GPTBigCodeForTokenClassification,
                  GPTBigCodeModel,
                  GPTBigCodePreTrainedModel,
              )


else:
       import sys

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