




'''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 pickle import UnpicklingError

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict

from ..utils import logging


lowerCAmelCase       =							logging.get_logger(__name__)
def 			_a (      SCREAMING_SNAKE_CASE						,				SCREAMING_SNAKE_CASE     ):


		"""simple docstring"""




		try:
				with open(SCREAMING_SNAKE_CASE						,				'''rb'''     ) as flax_state_f:
						lowercase__            =      from_bytes(SCREAMING_SNAKE_CASE						,				flax_state_f.read()     )
		except UnpicklingError as e:
				try:
						with open(SCREAMING_SNAKE_CASE     ) as f:
								if f.read().startswith('''version'''     ):
										raise OSError(
										    '''You seem to have cloned a repository without having git-lfs installed. Please'''
										    ''' install git-lfs and run `git lfs install` followed by `git lfs pull` in the'''
										    ''' folder you cloned.'''     )
								else:
										raise ValueError from e
				except (UnicodeDecodeError, ValueError):
						raise EnvironmentError(f'Unable to convert {model_file} to Flax deserializable object. '     )

		return load_flax_weights_in_pytorch_model(SCREAMING_SNAKE_CASE						,				SCREAMING_SNAKE_CASE     )
def 			_a (      SCREAMING_SNAKE_CASE						,				SCREAMING_SNAKE_CASE     ):


		"""simple docstring"""




		try:
				import torch  # noqa: F401
		except ImportError:
				logger.error(
				    '''Loading Flax weights in PyTorch requires both PyTorch and Flax to be installed. Please see'''
				    ''' https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation'''
				    ''' instructions.'''     )
				raise

		# check if we have bf16 weights
		lowercase__            =      flatten_dict(jax.tree_util.tree_map(lambda SCREAMING_SNAKE_CASE     : x.dtype == jnp.bfloataa						,				SCREAMING_SNAKE_CASE     )     ).values()
		if any(SCREAMING_SNAKE_CASE     ):
				# convert all weights to fp32 if they are bf16 since torch.from_numpy can-not handle bf16

				# and bf16 is not fully supported in PT yet.
				logger.warning(
				    '''Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` '''
				    '''before loading those in PyTorch model.'''     )
				lowercase__            =      jax.tree_util.tree_map(
				    lambda SCREAMING_SNAKE_CASE     : params.astype(np.floataa     ) if params.dtype == jnp.bfloataa else params						,				SCREAMING_SNAKE_CASE     )

		lowercase__            =      ''''''

		lowercase__            =      flatten_dict(SCREAMING_SNAKE_CASE						,				sep='''.'''     )
		lowercase__            =      pt_model.state_dict()

		# keep track of unexpected & missing keys
		lowercase__            =      []
		lowercase__            =      set(pt_model_dict.keys()     )

		for flax_key_tuple, flax_tensor in flax_state_dict.items():
				lowercase__            =      flax_key_tuple.split('''.'''     )

				if flax_key_tuple_array[-1] == "kernel" and flax_tensor.ndim == 4:
						lowercase__            =      flax_key_tuple_array[:-1] + ['''weight''']
						lowercase__            =      jnp.transpose(SCREAMING_SNAKE_CASE						,				(3, 2, 0, 1)     )
				elif flax_key_tuple_array[-1] == "kernel":
						lowercase__            =      flax_key_tuple_array[:-1] + ['''weight''']
						lowercase__            =      flax_tensor.T
				elif flax_key_tuple_array[-1] == "scale":
						lowercase__            =      flax_key_tuple_array[:-1] + ['''weight''']

				if "time_embedding" not in flax_key_tuple_array:
						for i, flax_key_tuple_string in enumerate(SCREAMING_SNAKE_CASE     ):
								lowercase__            =      (
								    flax_key_tuple_string.replace('''_0'''						,				'''.0'''     )
								    .replace('''_1'''						,				'''.1'''     )
								    .replace('''_2'''						,				'''.2'''     )
								    .replace('''_3'''						,				'''.3'''     )
								    .replace('''_4'''						,				'''.4'''     )
								    .replace('''_5'''						,				'''.5'''     )
								    .replace('''_6'''						,				'''.6'''     )
								    .replace('''_7'''						,				'''.7'''     )
								    .replace('''_8'''						,				'''.8'''     )
								    .replace('''_9'''						,				'''.9'''     )
								)

				lowercase__            =      '''.'''.join(SCREAMING_SNAKE_CASE     )

				if flax_key in pt_model_dict:
						if flax_tensor.shape != pt_model_dict[flax_key].shape:
								raise ValueError(
								    f'Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected '
								    f'to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}.'     )
						else:
								# add weight to pytorch dict
								lowercase__            =      np.asarray(SCREAMING_SNAKE_CASE     ) if not isinstance(SCREAMING_SNAKE_CASE						,				np.ndarray     ) else flax_tensor
								lowercase__            =      torch.from_numpy(SCREAMING_SNAKE_CASE     )
								# remove from missing keys
								missing_keys.remove(SCREAMING_SNAKE_CASE     )
				else:
						# weight is not expected by PyTorch model
						unexpected_keys.append(SCREAMING_SNAKE_CASE     )

		pt_model.load_state_dict(SCREAMING_SNAKE_CASE     )

		# re-transform missing_keys to list
		lowercase__            =      list(SCREAMING_SNAKE_CASE     )

		if len(SCREAMING_SNAKE_CASE     ) > 0:
				logger.warning(
				    '''Some weights of the Flax model were not used when initializing the PyTorch model'''
				    f' {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing'
				    f' {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture'
				    ''' (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This'''
				    f' IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect'
				    ''' to be exactly identical (e.g. initializing a BertForSequenceClassification model from a'''
				    ''' FlaxBertForSequenceClassification model).'''     )
		if len(SCREAMING_SNAKE_CASE     ) > 0:
				logger.warning(
				    f'Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly'
				    f' initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to'
				    ''' use it for predictions and inference.'''     )

		return pt_model


