'''simple docstring'''




import warnings

from transformers import AutoTokenizer
from transformers.utils import is_torch_available
from transformers.utils.generic import ExplicitEnum

from ...processing_utils import ProcessorMixin


if is_torch_available():
  import torch




class    a__				(       UpperCAmelCase__			):
     lowerCamelCase  :					str							="char"
     lowerCamelCase  :					Optional[Any]							="bpe"
     lowerCamelCase  :					Union[str, Any]							="wp"
__UpperCAmelCase												=(DecodeType.CHARACTER, DecodeType.BPE, DecodeType.WORDPIECE)




class    a__				(       UpperCAmelCase__			):
     lowerCamelCase  :					int							=["image_processor", "char_tokenizer"]
     lowerCamelCase  :					str							="ViTImageProcessor"
     lowerCamelCase  :					Optional[Any]							="MgpstrTokenizer"






     def __init__(			self		:		str  ,    a		:		Union[str, Any]=None  ,    a		:		List[Any]=None  ,    **a		:		Any   ):





         """simple docstring"""



         __lowerCamelCase          =  None
         if "feature_extractor" in kwargs:
             warnings.warn(
                 '''The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`'''
                 ''' instead.'''  ,    a  ,    )
             __lowerCamelCase          =  kwargs.pop('''feature_extractor'''   )

         __lowerCamelCase          =  image_processor if image_processor is not None else feature_extractor
         if image_processor is None:
             raise ValueError('''You need to specify an `image_processor`.'''   )
         if tokenizer is None:
             raise ValueError('''You need to specify a `tokenizer`.'''   )

         __lowerCamelCase          =  tokenizer
         __lowerCamelCase          =  AutoTokenizer.from_pretrained('''gpt2'''   )
         __lowerCamelCase          =  AutoTokenizer.from_pretrained('''bert-base-uncased'''   )

         super().__init__(a  ,    a   )






     def __call__(			self		:		Optional[int]  ,    a		:		Any=None  ,    a		:		Tuple=None  ,    a		:		Any=None  ,    **a		:		Optional[Any]   ):





         """simple docstring"""



         if images is None and text is None:
             raise ValueError('''You need to specify either an `images` or `text` input to process.'''   )

         if images is not None:
             __lowerCamelCase          =  self.image_processor(a  ,    return_tensors=a  ,    **a   )
         if text is not None:
             __lowerCamelCase          =  self.char_tokenizer(a  ,    return_tensors=a  ,    **a   )

         if text is None:
             return inputs
         elif images is None:
             return encodings
         else:
             __lowerCamelCase          =  encodings['''input_ids''']
             return inputs






     def        SCREAMING_SNAKE_CASE__       (			self		:		List[str]  ,    a		:		List[Any]   ):





         """simple docstring"""



         __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase          =  sequences
         __lowerCamelCase          =  char_preds.size(0   )

         __lowerCamelCase     ,  __lowerCamelCase          =  self._decode_helper(a  ,    '''char'''   )
         __lowerCamelCase     ,  __lowerCamelCase          =  self._decode_helper(a  ,    '''bpe'''   )
         __lowerCamelCase     ,  __lowerCamelCase          =  self._decode_helper(a  ,    '''wp'''   )

         __lowerCamelCase          =  []
         __lowerCamelCase          =  []
         for i in range(a   ):
             __lowerCamelCase          =  [char_scores[i], bpe_scores[i], wp_scores[i]]
             __lowerCamelCase          =  [char_strs[i], bpe_strs[i], wp_strs[i]]
             __lowerCamelCase          =  scores.index(max(a   )   )
             final_strs.append(strs[max_score_index]   )
             final_scores.append(scores[max_score_index]   )

         __lowerCamelCase          =  {}
         __lowerCamelCase          =  final_strs
         __lowerCamelCase          =  final_scores
         __lowerCamelCase          =  char_strs
         __lowerCamelCase          =  bpe_strs
         __lowerCamelCase          =  wp_strs
         return out






     def        SCREAMING_SNAKE_CASE__       (			self		:		int  ,    a		:		Any  ,    a		:		Tuple   ):





         """simple docstring"""



         if format == DecodeType.CHARACTER:
             __lowerCamelCase          =  self.char_decode
             __lowerCamelCase          =  1
             __lowerCamelCase          =  '''[s]'''
         elif format == DecodeType.BPE:
             __lowerCamelCase          =  self.bpe_decode
             __lowerCamelCase          =  2
             __lowerCamelCase          =  '''#'''
         elif format == DecodeType.WORDPIECE:
             __lowerCamelCase          =  self.wp_decode
             __lowerCamelCase          =  1_02
             __lowerCamelCase          =  '''[SEP]'''
         else:
             raise ValueError(f"""Format {format} is not supported."""   )

         __lowerCamelCase     ,  __lowerCamelCase          =  [], []
         __lowerCamelCase          =  pred_logits.size(0   )
         __lowerCamelCase          =  pred_logits.size(1   )
         __lowerCamelCase     ,  __lowerCamelCase          =  pred_logits.topk(1  ,    dim=-1  ,    largest=a  ,    sorted=a   )
         __lowerCamelCase          =  preds_index.view(-1  ,    a   )[:, 1:]
         __lowerCamelCase          =  decoder(a   )
         __lowerCamelCase     ,  __lowerCamelCase          =  torch.nn.functional.softmax(a  ,    dim=2   ).max(dim=2   )
         __lowerCamelCase          =  preds_max_prob[:, 1:]

         for index in range(a   ):
             __lowerCamelCase          =  preds_str[index].find(a   )
             __lowerCamelCase          =  preds_str[index][:pred_eos]
             __lowerCamelCase          =  preds_index[index].cpu().tolist()
             __lowerCamelCase          =  pred_index.index(a   ) if eos_token in pred_index else -1
             __lowerCamelCase          =  preds_max_prob[index][: pred_eos_index + 1]
             __lowerCamelCase          =  pred_max_prob.cumprod(dim=0   )[-1] if pred_max_prob.nelement() != 0 else 0.0
             dec_strs.append(a   )
             conf_scores.append(a   )

         return dec_strs, conf_scores






     def        SCREAMING_SNAKE_CASE__       (			self		:		int  ,    a		:		Union[str, Any]   ):





         """simple docstring"""



         __lowerCamelCase          =  [seq.replace(''' '''  ,    ''''''   ) for seq in self.char_tokenizer.batch_decode(a   )]
         return decode_strs






     def        SCREAMING_SNAKE_CASE__       (			self		:		int  ,    a		:		Optional[Any]   ):





         """simple docstring"""



         return self.bpe_tokenizer.batch_decode(a   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		Dict  ,    a		:		Optional[Any]   ):





         """simple docstring"""



         __lowerCamelCase          =  [seq.replace(''' '''  ,    ''''''   ) for seq in self.wp_tokenizer.batch_decode(a   )]
         return decode_strs

'''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        SCREAMING_SNAKE_CASE__       (			self		:		Optional[Any]   ):





         """simple docstring"""



         __lowerCamelCase          =  tempfile.mkdtemp()
         __lowerCamelCase          =  5

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

         __lowerCamelCase          =  os.path.join(self.tmpdirname  ,    '''realm_block_records'''   )
         os.makedirs(a  ,    exist_ok=a   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[Any]   ):





         """simple docstring"""



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






     def        SCREAMING_SNAKE_CASE__       (			self		:		Any   ):





         """simple docstring"""



         shutil.rmtree(self.tmpdirname   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[Any]   ):





         """simple docstring"""



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






     def        SCREAMING_SNAKE_CASE__       (			self		:		int   ):





         """simple docstring"""



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






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[Any]   ):





         """simple docstring"""



         __lowerCamelCase          =  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        SCREAMING_SNAKE_CASE__       (			self		:		str   ):





         """simple docstring"""



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






     def        SCREAMING_SNAKE_CASE__       (			self		:		Any   ):





         """simple docstring"""



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

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

         __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase          =  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        SCREAMING_SNAKE_CASE__       (			self		:		Optional[Any]   ):





         """simple docstring"""



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

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

         __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase          =  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        SCREAMING_SNAKE_CASE__       (			self		:		int   ):





         """simple docstring"""



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

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

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