




'''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 			_a  (	__a ):
       def     A      (	self						:	List[Any]				):



           '''simple docstring'''

           UpperCAmelCase			       =		tempfile.mkdtemp()
           UpperCAmelCase			       =		5

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

           UpperCAmelCase			       =		os.path.join(self.tmpdirname   ,				'''realm_block_records'''				)
           os.makedirs(lowercase   ,				exist_ok=lowercase				)
       def     A      (	self						:	List[Any]				):



           '''simple docstring'''

           return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname   ,				'''realm_tokenizer'''				)				)
       def     A      (	self						:	int				):



           '''simple docstring'''

           shutil.rmtree(self.tmpdirname				)
       def     A      (	self						:	Optional[int]				):



           '''simple docstring'''

           UpperCAmelCase			       =		RealmConfig(num_block_records=self.num_block_records				)
           return config
       def     A      (	self						:	List[str]				):



           '''simple docstring'''

           UpperCAmelCase			       =		Dataset.from_dict(
               {
                   '''id''': ['''0''', '''1'''],
                   '''question''': ['''foo''', '''bar'''],
                   '''answers''': [['''Foo''', '''Bar'''], ['''Bar''']],
               }				)
           return dataset
       def     A      (	self						:	Optional[int]				):



           '''simple docstring'''

           UpperCAmelCase			       =		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=lowercase   ,				)
           return block_records
       def     A      (	self						:	Optional[Any]				):



           '''simple docstring'''

           UpperCAmelCase			       =		RealmRetriever(
               block_records=self.get_dummy_block_records()   ,				tokenizer=self.get_tokenizer()   ,				)
           return retriever
       def     A      (	self						:	str				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.get_config()
           UpperCAmelCase			       =		self.get_dummy_retriever()
           UpperCAmelCase			       =		retriever.tokenizer

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

           UpperCAmelCase       ,     UpperCAmelCase       ,     UpperCAmelCase       ,     UpperCAmelCase			       =		retriever(
               lowercase   ,				lowercase   ,				answer_ids=lowercase   ,				max_length=lowercase   ,				return_tensors='''np'''				)

           self.assertEqual(len(lowercase				)   ,				2				)
           self.assertEqual(len(lowercase				)   ,				2				)
           self.assertEqual(len(lowercase				)   ,				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     A      (	self						:	List[str]				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.get_config()
           UpperCAmelCase			       =		self.get_dummy_retriever()
           UpperCAmelCase			       =		retriever.tokenizer

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

           UpperCAmelCase       ,     UpperCAmelCase       ,     UpperCAmelCase       ,     UpperCAmelCase			       =		retriever(
               lowercase   ,				lowercase   ,				answer_ids=lowercase   ,				max_length=lowercase   ,				return_tensors='''np'''				)

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







       def     A      (	self						:	Dict				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.get_dummy_retriever()
           retriever.save_pretrained(os.path.join(self.tmpdirname   ,				'''realm_block_records'''				)				)

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

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






from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer, TensorType, is_torch_available
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast
from ...utils import logging


lowercase  :					List[str]						    =							logging.get_logger(__name__)

lowercase  :					Optional[Any]						    =							{
    """EleutherAI/gpt-neo-1.3B""": """https://huggingface.co/EleutherAI/gpt-neo-1.3B/resolve/main/config.json""",
    # See all GPTNeo models at https://huggingface.co/models?filter=gpt_neo
}

class __snake_case		(      lowerCAmelCase     ):
    _a		:						str=				"gpt_neo"
    _a		:						Optional[int]=				["past_key_values"]
    _a		:						Dict=				{"num_attention_heads": "num_heads", "num_hidden_layers": "num_layers"}
    def __init__(  self		,snake_case=50257		,snake_case=2048		,snake_case=2048		,snake_case=24		,snake_case=[[["global", "local"], 12]]		,snake_case=16		,snake_case=None		,snake_case=256		,snake_case="gelu_new"		,snake_case=0.0		,snake_case=0.0		,snake_case=0.0		,snake_case=0.1		,snake_case=1e-5		,snake_case=0.02		,snake_case=True		,snake_case=50256		,snake_case=50256		,**snake_case		,):



           '''simple docstring'''







           lowercase							:     int       	=      vocab_size
           lowercase							:     Union[str, Any]       	=      max_position_embeddings
           lowercase							:     Dict       	=      hidden_size
           lowercase							:     Union[str, Any]       	=      num_layers
           lowercase							:     Union[str, Any]       	=      num_heads
           lowercase							:     Optional[int]       	=      intermediate_size
           lowercase							:     List[str]       	=      window_size
           lowercase							:     Optional[int]       	=      activation_function
           lowercase							:     List[str]       	=      resid_dropout
           lowercase							:     int       	=      embed_dropout
           lowercase							:     Optional[int]       	=      attention_dropout
           lowercase							:     Tuple       	=      classifier_dropout
           lowercase							:     Optional[int]       	=      layer_norm_epsilon
           lowercase							:     Dict       	=      initializer_range
           lowercase							:     List[str]       	=      use_cache

           lowercase							:     Optional[int]       	=      bos_token_id
           lowercase							:     int       	=      eos_token_id

           lowercase							:     Union[str, Any]       	=      attention_types
           lowercase							:     Dict       	=      self.expand_attention_types_params(snake_case      )

           if len(self.attention_layers      ) != self.num_layers:
                  raise ValueError(
                      """Configuration for convolutional module is incorrect. """
                      """It is required that `len(config.attention_layers)` == `config.num_layers` """
                      f"but is `len(config.attention_layers) = {len(self.attention_layers      )}`, "
                      f"`config.num_layers = {self.num_layers}`. "
                      """`config.attention_layers` is prepared using `config.attention_types`. """
                      """Please verify the value of `config.attention_types` argument."""      )

           super().__init__(bos_token_id=snake_case		,eos_token_id=snake_case		,**snake_case      )
    @staticmethod
    def       _SCREAMING_SNAKE_CASE       (  snake_case      ):



           '''simple docstring'''







           lowercase							:     List[Any]       	=      []
           for item in attention_types:
                  for _ in range(item[1]      ):
                         attentions.extend(item[0]      )
           return attentions







def 						_snake_case(  SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				) ->    Tuple:
       import torch

       lowercase							:     Tuple       	=      input.size()
       lowercase							:     int       	=      len(SCREAMING_SNAKE_CASE__				)
       lowercase							:     Union[str, Any]       	=      shape[dimension]

       lowercase							:     int       	=      torch.arange(0					,	SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				)
       lowercase							:     Optional[Any]       	=      torch.div(sizedim - size					,	SCREAMING_SNAKE_CASE__					,	rounding_mode="""floor"""				) + 1
       lowercase							:     Dict       	=      torch.arange(SCREAMING_SNAKE_CASE__				) + low_indices[:min_length][:, None]

       lowercase							:     Union[str, Any]       	=      [slice(SCREAMING_SNAKE_CASE__				)] * rank
       lowercase							:     Optional[Any]       	=      indices
       lowercase							:     List[str]       	=      input[s]

       lowercase							:     Optional[int]       	=      list(range(0					,	rank + 1				)				)
       perm.append(perm.pop(dimension + 1				)				)

       return sliced.permute(SCREAMING_SNAKE_CASE__				)







def 						_snake_case(  SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				) ->    Union[str, Any]:
       import torch

       lowercase							:     Union[str, Any]       	=      torch.arange(1					,	SCREAMING_SNAKE_CASE__				)
       lowercase							:     Optional[Any]       	=      torch.remainder(SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				)
       lowercase							:     List[Any]       	=      remainders == 0
       lowercase							:     Optional[int]       	=      candidates[divisor_indices]
       lowercase							:     List[Any]       	=      torch.max(SCREAMING_SNAKE_CASE__				)
       return largest_divisor, torch.div(SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__					,	rounding_mode="""floor"""				)




class __snake_case		(      lowerCAmelCase     ):
    @property
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           lowercase							:     Dict       	=      OrderedDict({"""input_ids""": {0: """batch""", 1: """sequence"""}}      )
           if self.use_past:
                  self.fill_with_past_key_values_(snake_case		,direction="""inputs"""      )
                  lowercase							:     Dict       	=      {0: """batch""", 1: """past_sequence + sequence"""}
           else:
                  lowercase							:     List[str]       	=      {0: """batch""", 1: """sequence"""}

           return common_inputs
    @property
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           return self._config.num_heads
    def       _SCREAMING_SNAKE_CASE       (  self		,snake_case		,snake_case = -1		,snake_case = -1		,snake_case = False		,snake_case = None		,):



           '''simple docstring'''







           lowercase							:     Any       	=      super(snake_case		,self      ).generate_dummy_inputs(
               snake_case		,batch_size=snake_case		,seq_length=snake_case		,is_pair=snake_case		,framework=snake_case      )

           # We need to order the input in the way they appears in the forward()
           lowercase							:     List[str]       	=      OrderedDict({"""input_ids""": common_inputs["""input_ids"""]}      )

           # Need to add the past_keys
           if self.use_past:
                  if not is_torch_available():
                         raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""      )
                  else:
                         import torch

                         lowercase	,			lowercase							:     List[Any]       	=      common_inputs["""input_ids"""].shape
                         # Not using the same length for past_key_values
                         lowercase							:     Optional[int]       	=      seqlen + 2
                         lowercase							:     int       	=      (
                             batch,
                             self.num_attention_heads,
                             past_key_values_length,
                             self._config.hidden_size // self.num_attention_heads,
                         )
                         lowercase							:     Optional[Any]       	=      [
                             (torch.zeros(snake_case      ), torch.zeros(snake_case      )) for _ in range(self.num_layers      )
                         ]

           lowercase							:     Optional[Any]       	=      common_inputs["""attention_mask"""]
           if self.use_past:
                  lowercase							:     Any       	=      ordered_inputs["""attention_mask"""].dtype
                  lowercase							:     Union[str, Any]       	=      torch.cat(
                      [ordered_inputs["""attention_mask"""], torch.ones(snake_case		,snake_case		,dtype=snake_case      )]		,dim=1      )

           return ordered_inputs
    @property
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           return 13
