






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(	lowerCAmelCase__ ):





       """simple docstring"""







       def       _a   ( self							):
         UpperCamelCase_: Dict   				=    tempfile.mkdtemp()
         UpperCamelCase_: Optional[Any]   				=    5

         # Realm tok
         UpperCamelCase_: Optional[int]   				=    [
             '[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_: Any   				=    os.path.join(self.tmpdirname  ,      'realm_tokenizer'							)
         os.makedirs(_SCREAMING_SNAKE_CASE  ,      exist_ok=_SCREAMING_SNAKE_CASE							)
         UpperCamelCase_: Optional[int]   				=    os.path.join(_SCREAMING_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]							)							)

         UpperCamelCase_: str   				=    os.path.join(self.tmpdirname  ,      'realm_block_records'							)
         os.makedirs(_SCREAMING_SNAKE_CASE  ,      exist_ok=_SCREAMING_SNAKE_CASE							)





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





       def       _a   ( self							):
         shutil.rmtree(self.tmpdirname							)





       def       _a   ( self							):
         UpperCamelCase_: Optional[int]   				=    RealmConfig(num_block_records=self.num_block_records							)
         return config





       def       _a   ( self							):
         UpperCamelCase_: Union[str, Any]   				=    Dataset.from_dict(
             {
                 'id': ['0', '1'],
                 'question': ['foo', 'bar'],
                 'answers': [['Foo', 'Bar'], ['Bar']],
             }							)
         return dataset





       def       _a   ( self							):
         UpperCamelCase_: Union[str, 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=_SCREAMING_SNAKE_CASE  ,      )
         return block_records





       def       _a   ( self							):
         UpperCamelCase_: Optional[Any]   				=    RealmRetriever(
             block_records=self.get_dummy_block_records()  ,      tokenizer=self.get_tokenizer()  ,      )
         return retriever





       def       _a   ( self							):
         UpperCamelCase_: Optional[int]   				=    self.get_config()
         UpperCamelCase_: List[Any]   				=    self.get_dummy_retriever()
         UpperCamelCase_: Any   				=    retriever.tokenizer

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

         UpperCamelCase_				,UpperCamelCase_				,UpperCamelCase_				,UpperCamelCase_: Any   				=    retriever(
             _SCREAMING_SNAKE_CASE  ,      _SCREAMING_SNAKE_CASE  ,      answer_ids=_SCREAMING_SNAKE_CASE  ,      max_length=_SCREAMING_SNAKE_CASE  ,      return_tensors='np'							)

         self.assertEqual(len(_SCREAMING_SNAKE_CASE							)  ,      2							)
         self.assertEqual(len(_SCREAMING_SNAKE_CASE							)  ,      2							)
         self.assertEqual(len(_SCREAMING_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       _a   ( self							):
         UpperCamelCase_: str   				=    self.get_config()
         UpperCamelCase_: Optional[Any]   				=    self.get_dummy_retriever()
         UpperCamelCase_: Tuple   				=    retriever.tokenizer

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

         UpperCamelCase_				,UpperCamelCase_				,UpperCamelCase_				,UpperCamelCase_: int   				=    retriever(
             _SCREAMING_SNAKE_CASE  ,      _SCREAMING_SNAKE_CASE  ,      answer_ids=_SCREAMING_SNAKE_CASE  ,      max_length=_SCREAMING_SNAKE_CASE  ,      return_tensors='np'							)

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





       def       _a   ( self							):
         UpperCamelCase_: Any   				=    self.get_dummy_retriever()
         retriever.save_pretrained(os.path.join(self.tmpdirname  ,      'realm_block_records'							)							)

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

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







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 (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_DEFAULT_MEAN,
    IMAGENET_DEFAULT_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    is_batched,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, logging


A_		:		int		     =			logging.get_logger(__name__)




class     _lowerCAmelCase(	UpperCAmelCase_ ):





       """simple docstring"""







       a    :		str									=['''pixel_values''']





       def __init__( self  ,      _lowerCamelCase = True  ,      _lowerCamelCase = None  ,      _lowerCamelCase = PILImageResampling.BICUBIC  ,      _lowerCamelCase = True  ,      _lowerCamelCase = True  ,      _lowerCamelCase = 1 / 2_5_5  ,      _lowerCamelCase = None  ,      _lowerCamelCase = True  ,      _lowerCamelCase = None  ,      _lowerCamelCase = None  ,      **_lowerCamelCase  ,      ):
         super().__init__(**_lowerCamelCase							)
         UpperCamelCase_: Any   				=    size if size is not None else {'height': 2_2_4, 'width': 2_2_4}
         UpperCamelCase_: Any   				=    get_size_dict(_lowerCamelCase							)
         UpperCamelCase_: Any   				=    crop_size if crop_size is not None else {'height': 2_2_4, 'width': 2_2_4}
         UpperCamelCase_: List[Any]   				=    get_size_dict(_lowerCamelCase  ,      default_to_square=_lowerCamelCase  ,      param_name='crop_size'							)

         UpperCamelCase_: Optional[int]   				=    do_resize
         UpperCamelCase_: Tuple   				=    do_rescale
         UpperCamelCase_: Dict   				=    do_normalize
         UpperCamelCase_: Optional[int]   				=    do_center_crop
         UpperCamelCase_: Tuple   				=    crop_size
         UpperCamelCase_: Optional[int]   				=    size
         UpperCamelCase_: Dict   				=    resample
         UpperCamelCase_: Tuple   				=    rescale_factor
         UpperCamelCase_: Tuple   				=    image_mean if image_mean is not None else IMAGENET_DEFAULT_MEAN
         UpperCamelCase_: Dict   				=    image_std if image_std is not None else IMAGENET_DEFAULT_STD





       def       _a   ( self  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase = PILImageResampling.BILINEAR  ,      _lowerCamelCase = None  ,      **_lowerCamelCase  ,      ):
         UpperCamelCase_: Optional[Any]   				=    get_size_dict(_lowerCamelCase							)

         if "shortest_edge" in size:
           UpperCamelCase_: str   				=    get_resize_output_image_size(_lowerCamelCase  ,      size=size['shortest_edge']  ,      default_to_square=_lowerCamelCase							)
           # size = get_resize_output_image_size(image, size["shortest_edge"], size["longest_edge"])
         elif "height" in size and "width" in size:
           UpperCamelCase_: Any   				=    (size['height'], size['width'])
         else:
           raise ValueError(f'''Size must contain \'height\' and \'width\' keys or \'shortest_edge\' key. Got {size.keys()}'''							)
         return resize(_lowerCamelCase  ,      size=_lowerCamelCase  ,      resample=_lowerCamelCase  ,      data_format=_lowerCamelCase  ,      **_lowerCamelCase							)





       def       _a   ( self  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase = None  ,      **_lowerCamelCase  ,      ):
         UpperCamelCase_: List[str]   				=    get_size_dict(_lowerCamelCase							)
         if "height" not in size or "width" not in size:
           raise ValueError(f'''The `size` parameter must contain the keys (height, width). Got {size.keys()}'''							)
         return center_crop(_lowerCamelCase  ,      size=(size['height'], size['width'])  ,      data_format=_lowerCamelCase  ,      **_lowerCamelCase							)





       def       _a   ( self  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase = None  ,      **_lowerCamelCase							):
         return rescale(_lowerCamelCase  ,      scale=_lowerCamelCase  ,      data_format=_lowerCamelCase  ,      **_lowerCamelCase							)





       def       _a   ( self  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase = None  ,      **_lowerCamelCase  ,      ):
         return normalize(_lowerCamelCase  ,      mean=_lowerCamelCase  ,      std=_lowerCamelCase  ,      data_format=_lowerCamelCase  ,      **_lowerCamelCase							)





       def       _a   ( self  ,      _lowerCamelCase  ,      _lowerCamelCase = None  ,      _lowerCamelCase = None  ,      _lowerCamelCase = None  ,      _lowerCamelCase = None  ,      _lowerCamelCase = None  ,      _lowerCamelCase = None  ,      _lowerCamelCase = None  ,      _lowerCamelCase = None  ,      _lowerCamelCase = None  ,      _lowerCamelCase = None  ,      _lowerCamelCase = None  ,      _lowerCamelCase = ChannelDimension.FIRST  ,      **_lowerCamelCase  ,      ):
         UpperCamelCase_: List[str]   				=    do_resize if do_resize is not None else self.do_resize
         UpperCamelCase_: str   				=    do_rescale if do_rescale is not None else self.do_rescale
         UpperCamelCase_: Dict   				=    do_normalize if do_normalize is not None else self.do_normalize
         UpperCamelCase_: List[str]   				=    do_center_crop if do_center_crop is not None else self.do_center_crop
         UpperCamelCase_: Optional[Any]   				=    crop_size if crop_size is not None else self.crop_size
         UpperCamelCase_: List[str]   				=    get_size_dict(_lowerCamelCase  ,      param_name='crop_size'  ,      default_to_square=_lowerCamelCase							)
         UpperCamelCase_: Tuple   				=    resample if resample is not None else self.resample
         UpperCamelCase_: str   				=    rescale_factor if rescale_factor is not None else self.rescale_factor
         UpperCamelCase_: Dict   				=    image_mean if image_mean is not None else self.image_mean
         UpperCamelCase_: Dict   				=    image_std if image_std is not None else self.image_std

         UpperCamelCase_: Any   				=    size if size is not None else self.size
         UpperCamelCase_: Optional[Any]   				=    get_size_dict(_lowerCamelCase							)

         if not is_batched(_lowerCamelCase							):
           UpperCamelCase_: Dict   				=    [images]

         if not valid_images(_lowerCamelCase							):
           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:
           raise ValueError('Size must be specified if do_resize is True.'							)

         if do_center_crop and crop_size is None:
           raise ValueError('Crop size must be specified if do_center_crop is True.'							)

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

         # All transformations expect numpy arrays.
         UpperCamelCase_: List[str]   				=    [to_numpy_array(_lowerCamelCase							) for image in images]

         if do_resize:
           UpperCamelCase_: List[str]   				=    [self.resize(image=_lowerCamelCase  ,      size=_lowerCamelCase  ,      resample=_lowerCamelCase							) for image in images]

         if do_center_crop:
           UpperCamelCase_: List[str]   				=    [self.center_crop(image=_lowerCamelCase  ,      size=_lowerCamelCase							) for image in images]

         if do_rescale:
           UpperCamelCase_: int   				=    [self.rescale(image=_lowerCamelCase  ,      scale=_lowerCamelCase							) for image in images]

         if do_normalize:
           UpperCamelCase_: Optional[Any]   				=    [self.normalize(image=_lowerCamelCase  ,      mean=_lowerCamelCase  ,      std=_lowerCamelCase							) for image in images]

         UpperCamelCase_: List[str]   				=    [to_channel_dimension_format(_lowerCamelCase  ,      _lowerCamelCase							) for image in images]

         UpperCamelCase_: Optional[int]   				=    {'pixel_values': images}
         return BatchFeature(data=_lowerCamelCase  ,      tensor_type=_lowerCamelCase							)