






"""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        _SCREAMING_SNAKE_CASE						(  __UpperCAmelCase							):





							"""simple docstring"""





							def  lowercase		(   self:     List[str]    ):

												'''simple docstring'''

												a__          =  tempfile.mkdtemp()
												a__          =  5

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

												a__          =  os.path.join(self.tmpdirname				,		'''realm_block_records'''    )
												os.makedirs(_lowerCamelCase				,		exist_ok=_lowerCamelCase    )




							def  lowercase		(   self:     str    ):

												'''simple docstring'''

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




							def  lowercase		(   self:     List[str]    ):

												'''simple docstring'''

												shutil.rmtree(self.tmpdirname    )




							def  lowercase		(   self:     Tuple    ):

												'''simple docstring'''

												a__          =  RealmConfig(num_block_records=self.num_block_records    )
												return config




							def  lowercase		(   self:     Optional[Any]    ):

												'''simple docstring'''

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




							def  lowercase		(   self:     List[Any]    ):

												'''simple docstring'''

												a__          =  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:     str    ):

												'''simple docstring'''

												a__          =  RealmRetriever(
												    block_records=self.get_dummy_block_records()				,		tokenizer=self.get_tokenizer()				,		)
												return retriever




							def  lowercase		(   self:     Any    ):

												'''simple docstring'''

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

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

												a__       ,a__       ,a__       ,a__          =  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, 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]    ):

												'''simple docstring'''

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

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

												a__       ,a__       ,a__       ,a__          =  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]    ):

												'''simple docstring'''

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

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

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










"""simple docstring"""







import flax.linen as nn
import jax.numpy as jnp

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD







class        _SCREAMING_SNAKE_CASE						(  nn.Module							):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =0.0
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =True
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =jnp.floataa




							def  lowercase		(   self:     Union[str, Any]    ):

												'''simple docstring'''

												a__          =  []
												a__          =  []

												for i in range(self.num_layers    ):
																	a__          =  self.in_channels if i == 0 else self.out_channels

																	a__          =  FlaxResnetBlockaD(
																	    in_channels=__A				,		out_channels=self.out_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
																	resnets.append(__A    )

																	a__          =  FlaxTransformeraDModel(
																	    in_channels=self.out_channels				,		n_heads=self.num_attention_heads				,		d_head=self.out_channels // self.num_attention_heads				,		depth=1				,		use_linear_projection=self.use_linear_projection				,		only_cross_attention=self.only_cross_attention				,		use_memory_efficient_attention=self.use_memory_efficient_attention				,		dtype=self.dtype				,		)
																	attentions.append(__A    )

												a__          =  resnets
												a__          =  attentions

												if self.add_downsample:
																	a__          =  FlaxDownsampleaD(self.out_channels				,		dtype=self.dtype    )





							def __call__(   self:     Optional[int]				,		__A:     Union[str, Any]				,		__A:     str				,		__A:     Optional[Any]				,		__A:     Any=True    ):

												'''simple docstring'''

												a__          =  ()

												for resnet, attn in zip(self.resnets				,		self.attentions    ):
																	a__          =  resnet(__A				,		__A				,		deterministic=__A    )
																	a__          =  attn(__A				,		__A				,		deterministic=__A    )
																	output_states += (hidden_states,)

												if self.add_downsample:
																	a__          =  self.downsamplers_a(__A    )
																	output_states += (hidden_states,)

												return hidden_states, output_states







class        _SCREAMING_SNAKE_CASE						(  nn.Module							):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =0.0
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =True
							_SCREAMING_SNAKE_CASE    =jnp.floataa




							def  lowercase		(   self:     Union[str, Any]    ):

												'''simple docstring'''

												a__          =  []

												for i in range(self.num_layers    ):
																	a__          =  self.in_channels if i == 0 else self.out_channels

																	a__          =  FlaxResnetBlockaD(
																	    in_channels=__A				,		out_channels=self.out_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
																	resnets.append(__A    )
												a__          =  resnets

												if self.add_downsample:
																	a__          =  FlaxDownsampleaD(self.out_channels				,		dtype=self.dtype    )





							def __call__(   self:     Dict				,		__A:     int				,		__A:     Dict				,		__A:     Optional[Any]=True    ):

												'''simple docstring'''

												a__          =  ()

												for resnet in self.resnets:
																	a__          =  resnet(__A				,		__A				,		deterministic=__A    )
																	output_states += (hidden_states,)

												if self.add_downsample:
																	a__          =  self.downsamplers_a(__A    )
																	output_states += (hidden_states,)

												return hidden_states, output_states







class        _SCREAMING_SNAKE_CASE						(  nn.Module							):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =0.0
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =True
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =jnp.floataa




							def  lowercase		(   self:     Optional[int]    ):

												'''simple docstring'''

												a__          =  []
												a__          =  []

												for i in range(self.num_layers    ):
																	a__          =  self.in_channels if (i == self.num_layers - 1) else self.out_channels
																	a__          =  self.prev_output_channel if i == 0 else self.out_channels

																	a__          =  FlaxResnetBlockaD(
																	    in_channels=resnet_in_channels + res_skip_channels				,		out_channels=self.out_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
																	resnets.append(__A    )

																	a__          =  FlaxTransformeraDModel(
																	    in_channels=self.out_channels				,		n_heads=self.num_attention_heads				,		d_head=self.out_channels // self.num_attention_heads				,		depth=1				,		use_linear_projection=self.use_linear_projection				,		only_cross_attention=self.only_cross_attention				,		use_memory_efficient_attention=self.use_memory_efficient_attention				,		dtype=self.dtype				,		)
																	attentions.append(__A    )

												a__          =  resnets
												a__          =  attentions

												if self.add_upsample:
																	a__          =  FlaxUpsampleaD(self.out_channels				,		dtype=self.dtype    )





							def __call__(   self:     Any				,		__A:     Optional[int]				,		__A:     List[Any]				,		__A:     List[str]				,		__A:     Optional[Any]				,		__A:     Any=True    ):

												'''simple docstring'''

												for resnet, attn in zip(self.resnets				,		self.attentions    ):
																	# pop res hidden states
																	a__          =  res_hidden_states_tuple[-1]
																	a__          =  res_hidden_states_tuple[:-1]
																	a__          =  jnp.concatenate((hidden_states, res_hidden_states)				,		axis=-1    )

																	a__          =  resnet(__A				,		__A				,		deterministic=__A    )
																	a__          =  attn(__A				,		__A				,		deterministic=__A    )

												if self.add_upsample:
																	a__          =  self.upsamplers_a(__A    )

												return hidden_states







class        _SCREAMING_SNAKE_CASE						(  nn.Module							):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =0.0
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =True
							_SCREAMING_SNAKE_CASE    =jnp.floataa




							def  lowercase		(   self:     str    ):

												'''simple docstring'''

												a__          =  []

												for i in range(self.num_layers    ):
																	a__          =  self.in_channels if (i == self.num_layers - 1) else self.out_channels
																	a__          =  self.prev_output_channel if i == 0 else self.out_channels

																	a__          =  FlaxResnetBlockaD(
																	    in_channels=resnet_in_channels + res_skip_channels				,		out_channels=self.out_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
																	resnets.append(__A    )

												a__          =  resnets

												if self.add_upsample:
																	a__          =  FlaxUpsampleaD(self.out_channels				,		dtype=self.dtype    )





							def __call__(   self:     Tuple				,		__A:     Optional[Any]				,		__A:     Optional[Any]				,		__A:     Union[str, Any]				,		__A:     Dict=True    ):

												'''simple docstring'''

												for resnet in self.resnets:
																	# pop res hidden states
																	a__          =  res_hidden_states_tuple[-1]
																	a__          =  res_hidden_states_tuple[:-1]
																	a__          =  jnp.concatenate((hidden_states, res_hidden_states)				,		axis=-1    )

																	a__          =  resnet(__A				,		__A				,		deterministic=__A    )

												if self.add_upsample:
																	a__          =  self.upsamplers_a(__A    )

												return hidden_states







class        _SCREAMING_SNAKE_CASE						(  nn.Module							):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =0.0
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =jnp.floataa




							def  lowercase		(   self:     Tuple    ):

												'''simple docstring'''

												a__          =  [
												    FlaxResnetBlockaD(
												        in_channels=self.in_channels				,		out_channels=self.in_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
												]

												a__          =  []

												for _ in range(self.num_layers    ):
																	a__          =  FlaxTransformeraDModel(
																	    in_channels=self.in_channels				,		n_heads=self.num_attention_heads				,		d_head=self.in_channels // self.num_attention_heads				,		depth=1				,		use_linear_projection=self.use_linear_projection				,		use_memory_efficient_attention=self.use_memory_efficient_attention				,		dtype=self.dtype				,		)
																	attentions.append(__A    )

																	a__          =  FlaxResnetBlockaD(
																	    in_channels=self.in_channels				,		out_channels=self.in_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
																	resnets.append(__A    )

												a__          =  resnets
												a__          =  attentions





							def __call__(   self:     Any				,		__A:     Optional[int]				,		__A:     int				,		__A:     Tuple				,		__A:     str=True    ):

												'''simple docstring'''

												a__          =  self.resnets[0](__A				,		__A    )
												for attn, resnet in zip(self.attentions				,		self.resnets[1:]    ):
																	a__          =  attn(__A				,		__A				,		deterministic=__A    )
																	a__          =  resnet(__A				,		__A				,		deterministic=__A    )

												return hidden_states


