
'''simple docstring'''







from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import convert_to_rgb, normalize, rescale, resize, to_channel_dimension_format
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
							import PIL


_lowercase :  Any    							=       logging.get_logger(__name__)






class  lowerCAmelCase__   (  lowerCamelCase_     ):
					lowerCAmelCase_        =     ['''pixel_values''']






					def __init__(   self				,							__SCREAMING_SNAKE_CASE = True				,							__SCREAMING_SNAKE_CASE = None				,							__SCREAMING_SNAKE_CASE = PILImageResampling.BICUBIC				,							__SCREAMING_SNAKE_CASE = True				,							__SCREAMING_SNAKE_CASE = 1 / 2_55				,							__SCREAMING_SNAKE_CASE = True				,							__SCREAMING_SNAKE_CASE = None				,							__SCREAMING_SNAKE_CASE = None				,							__SCREAMING_SNAKE_CASE = True				,							**__SCREAMING_SNAKE_CASE				,							):


											"""simple docstring"""





											super().__init__(**__SCREAMING_SNAKE_CASE		)
											lowercase_	:     str													=       size if size is not None else {'''height''': 3_84, '''width''': 3_84}
											lowercase_	:     Union[str, Any]													=       get_size_dict(__SCREAMING_SNAKE_CASE				,							default_to_square=__SCREAMING_SNAKE_CASE		)

											lowercase_	:     Optional[Any]													=       do_resize
											lowercase_	:     Dict													=       size
											lowercase_	:     Optional[Any]													=       resample
											lowercase_	:     Tuple													=       do_rescale
											lowercase_	:     str													=       rescale_factor
											lowercase_	:     Union[str, Any]													=       do_normalize
											lowercase_	:     Dict													=       image_mean if image_mean is not None else OPENAI_CLIP_MEAN
											lowercase_	:     Union[str, Any]													=       image_std if image_std is not None else OPENAI_CLIP_STD
											lowercase_	:     Dict													=       do_convert_rgb






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE = PILImageResampling.BICUBIC				,							__SCREAMING_SNAKE_CASE = None				,							**__SCREAMING_SNAKE_CASE				,							):


											"""simple docstring"""





											lowercase_	:     List[str]													=       get_size_dict(__SCREAMING_SNAKE_CASE				,							default_to_square=__SCREAMING_SNAKE_CASE		)
											if "height" not in size or "width" not in size:
																	raise ValueError(F'''The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}'''		)
											lowercase_	:     Any													=       (size['''height'''], size['''width'''])
											return resize(__SCREAMING_SNAKE_CASE				,							size=__SCREAMING_SNAKE_CASE				,							resample=__SCREAMING_SNAKE_CASE				,							data_format=__SCREAMING_SNAKE_CASE				,							**__SCREAMING_SNAKE_CASE		)






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE = None				,							**__SCREAMING_SNAKE_CASE				,							):


											"""simple docstring"""





											return rescale(__SCREAMING_SNAKE_CASE				,							scale=__SCREAMING_SNAKE_CASE				,							data_format=__SCREAMING_SNAKE_CASE				,							**__SCREAMING_SNAKE_CASE		)






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE = None				,							**__SCREAMING_SNAKE_CASE				,							):


											"""simple docstring"""





											return normalize(__SCREAMING_SNAKE_CASE				,							mean=__SCREAMING_SNAKE_CASE				,							std=__SCREAMING_SNAKE_CASE				,							data_format=__SCREAMING_SNAKE_CASE				,							**__SCREAMING_SNAKE_CASE		)






					def 	_snake_case    (   self				,							__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE = None				,							__SCREAMING_SNAKE_CASE = None				,							__SCREAMING_SNAKE_CASE = None				,							__SCREAMING_SNAKE_CASE = None				,							__SCREAMING_SNAKE_CASE = None				,							__SCREAMING_SNAKE_CASE = None				,							__SCREAMING_SNAKE_CASE = None				,							__SCREAMING_SNAKE_CASE = None				,							__SCREAMING_SNAKE_CASE = None				,							__SCREAMING_SNAKE_CASE = None				,							__SCREAMING_SNAKE_CASE = ChannelDimension.FIRST				,							**__SCREAMING_SNAKE_CASE				,							):


											"""simple docstring"""





											lowercase_	:     List[Any]													=       do_resize if do_resize is not None else self.do_resize
											lowercase_	:     Dict													=       resample if resample is not None else self.resample
											lowercase_	:     List[str]													=       do_rescale if do_rescale is not None else self.do_rescale
											lowercase_	:     Dict													=       rescale_factor if rescale_factor is not None else self.rescale_factor
											lowercase_	:     str													=       do_normalize if do_normalize is not None else self.do_normalize
											lowercase_	:     Union[str, Any]													=       image_mean if image_mean is not None else self.image_mean
											lowercase_	:     List[Any]													=       image_std if image_std is not None else self.image_std
											lowercase_	:     Optional[Any]													=       do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

											lowercase_	:     List[Any]													=       size if size is not None else self.size
											lowercase_	:     Dict													=       get_size_dict(__SCREAMING_SNAKE_CASE				,							default_to_square=__SCREAMING_SNAKE_CASE		)

											lowercase_	:     Optional[int]													=       make_list_of_images(__SCREAMING_SNAKE_CASE		)

											if not valid_images(__SCREAMING_SNAKE_CASE		):
																	raise ValueError(
																	    '''Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '''
																	    '''torch.Tensor, tf.Tensor or jax.ndarray.'''		)

											if do_resize and size is None or resample is None:
																	raise ValueError('''Size and resample must be specified if do_resize is True.'''		)

											if do_rescale and rescale_factor is None:
																	raise ValueError('''Rescale factor must be specified if do_rescale is True.'''		)

											if do_normalize and (image_mean is None or image_std is None):
																	raise ValueError('''Image mean and std must be specified if do_normalize is True.'''		)

											# PIL RGBA images are converted to RGB
											if do_convert_rgb:
																	lowercase_	:     int													=       [convert_to_rgb(__SCREAMING_SNAKE_CASE		) for image in images]

											# All transformations expect numpy arrays.
											lowercase_	:     Union[str, Any]													=       [to_numpy_array(__SCREAMING_SNAKE_CASE		) for image in images]

											if do_resize:
																	lowercase_	:     List[Any]													=       [self.resize(image=__SCREAMING_SNAKE_CASE				,							size=__SCREAMING_SNAKE_CASE				,							resample=__SCREAMING_SNAKE_CASE		) for image in images]

											if do_rescale:
																	lowercase_	:     Any													=       [self.rescale(image=__SCREAMING_SNAKE_CASE				,							scale=__SCREAMING_SNAKE_CASE		) for image in images]

											if do_normalize:
																	lowercase_	:     str													=       [self.normalize(image=__SCREAMING_SNAKE_CASE				,							mean=__SCREAMING_SNAKE_CASE				,							std=__SCREAMING_SNAKE_CASE		) for image in images]

											lowercase_	:     List[str]													=       [to_channel_dimension_format(__SCREAMING_SNAKE_CASE				,							__SCREAMING_SNAKE_CASE		) for image in images]

											lowercase_	:     Dict													=       BatchFeature(data={'''pixel_values''': images}				,							tensor_type=__SCREAMING_SNAKE_CASE		)

											return encoded_outputs




'''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      ( UpperCamelCase__	):



    def 					_lowercase						(self    :    Optional[Any] ):
         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(__a   ,			exist_ok=__a )
         UpperCAmelCase_           =   os.path.join(__a   ,			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(__a   ,			exist_ok=__a )



    def 					_lowercase						(self    :    Optional[Any] ):
         return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname   ,			"realm_tokenizer" ) )



    def 					_lowercase						(self    :    Any ):
         shutil.rmtree(self.tmpdirname )



    def 					_lowercase						(self    :    List[Any] ):
         UpperCAmelCase_           =   RealmConfig(num_block_records=self.num_block_records )
         return config



    def 					_lowercase						(self    :    List[str] ):
         UpperCAmelCase_           =   Dataset.from_dict(
             {
                 "id": ["0", "1"],
                 "question": ["foo", "bar"],
                 "answers": [["Foo", "Bar"], ["Bar"]],
             } )
         return dataset



    def 					_lowercase						(self    :    Any ):
         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=__a   ,			)
         return block_records



    def 					_lowercase						(self    :    Union[str, Any] ):
         UpperCAmelCase_           =   RealmRetriever(
             block_records=self.get_dummy_block_records()   ,			tokenizer=self.get_tokenizer()   ,			)
         return retriever



    def 					_lowercase						(self    :    int ):
         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=__a   ,			return_token_type_ids=__a   ,			return_attention_mask=__a   ,			).input_ids
         UpperCAmelCase_           =   config.reader_seq_len

         UpperCAmelCase_    , UpperCAmelCase_    , UpperCAmelCase_    , UpperCAmelCase_           =   retriever(
             __a   ,			__a   ,			answer_ids=__a   ,			max_length=__a   ,			return_tensors="np" )

         self.assertEqual(len(__a )   ,			2 )
         self.assertEqual(len(__a )   ,			2 )
         self.assertEqual(len(__a )   ,			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[Any] ):
         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=__a   ,			return_token_type_ids=__a   ,			return_attention_mask=__a   ,			).input_ids
         UpperCAmelCase_           =   config.reader_seq_len

         UpperCAmelCase_    , UpperCAmelCase_    , UpperCAmelCase_    , UpperCAmelCase_           =   retriever(
             __a   ,			__a   ,			answer_ids=__a   ,			max_length=__a   ,			return_tensors="np" )

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






    def 					_lowercase						(self    :    Optional[Any] ):
         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" )




