
"""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   _UpperCAmelCase				(		SCREAMING_SNAKE_CASE_				):




      def      a_   (				self				)    ->       Tuple:
          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				)    ->       RealmTokenizer:
          return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname ,							'realm_tokenizer'				)				)




      def      a_   (				self				)    ->       Optional[int]:
          shutil.rmtree(self.tmpdirname				)




      def      a_   (				self				)    ->       Optional[int]:
          UpperCAmelCase          =       RealmConfig(num_block_records=self.num_block_records				)
          return config




      def      a_   (				self				)    ->       Optional[Any]:
          UpperCAmelCase          =       Dataset.from_dict(
              {
                  'id': ['0', '1'],
                  'question': ['foo', 'bar'],
                  'answers': [['Foo', 'Bar'], ['Bar']],
              }				)
          return dataset




      def      a_   (				self				)    ->       Tuple:
          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				)    ->       str:
          UpperCAmelCase          =       RealmRetriever(
              block_records=self.get_dummy_block_records() ,							tokenizer=self.get_tokenizer() ,							)
          return retriever




      def      a_   (				self				)    ->       Optional[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=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, 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				)    ->       str:
          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:
          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'				)





"""simple docstring"""

import functools
import operator

from ...configuration_utils import PretrainedConfig
from ...utils import logging


SCREAMING_SNAKE_CASE_            =    logging.get_logger(__name__)

SCREAMING_SNAKE_CASE_            =    {
    '''asapp/sew-d-tiny-100k''': '''https://huggingface.co/asapp/sew-d-tiny-100k/resolve/main/config.json''',
    # See all SEW-D models at https://huggingface.co/models?filter=sew-d
}
class   _UpperCAmelCase				(		SCREAMING_SNAKE_CASE_				):
      __SCREAMING_SNAKE_CASE :				List[str]				   =						"sew-d"




      def __init__(				self ,							lowercase_=3_2 ,							lowercase_=7_6_8 ,							lowercase_=1_2 ,							lowercase_=1_2 ,							lowercase_=3_0_7_2 ,							lowercase_=2 ,							lowercase_=5_1_2 ,							lowercase_=2_5_6 ,							lowercase_=True ,							lowercase_=True ,							lowercase_=("p2c", "c2p") ,							lowercase_="layer_norm" ,							lowercase_="gelu_python" ,							lowercase_=0.1 ,							lowercase_=0.1 ,							lowercase_=0.1 ,							lowercase_=0.0 ,							lowercase_=0.1 ,							lowercase_=0.0_2 ,							lowercase_=1E-7 ,							lowercase_=1E-5 ,							lowercase_="group" ,							lowercase_="gelu" ,							lowercase_=(6_4, 1_2_8, 1_2_8, 1_2_8, 1_2_8, 2_5_6, 2_5_6, 2_5_6, 2_5_6, 5_1_2, 5_1_2, 5_1_2, 5_1_2) ,							lowercase_=(5, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1) ,							lowercase_=(1_0, 3, 1, 3, 1, 3, 1, 3, 1, 2, 1, 2, 1) ,							lowercase_=False ,							lowercase_=1_2_8 ,							lowercase_=1_6 ,							lowercase_=True ,							lowercase_=0.0_5 ,							lowercase_=1_0 ,							lowercase_=2 ,							lowercase_=0.0 ,							lowercase_=1_0 ,							lowercase_=0 ,							lowercase_="mean" ,							lowercase_=False ,							lowercase_=False ,							lowercase_=2_5_6 ,							lowercase_=0 ,							lowercase_=1 ,							lowercase_=2 ,							**lowercase_ ,							)    ->       Dict:
          super().__init__(**lowercase_ ,							pad_token_id=lowercase_ ,							bos_token_id=lowercase_ ,							eos_token_id=lowercase_				)
          UpperCAmelCase          =       hidden_size
          UpperCAmelCase          =       feat_extract_norm
          UpperCAmelCase          =       feat_extract_activation
          UpperCAmelCase          =       list(lowercase_				)
          UpperCAmelCase          =       list(lowercase_				)
          UpperCAmelCase          =       list(lowercase_				)
          UpperCAmelCase          =       conv_bias
          UpperCAmelCase          =       num_conv_pos_embeddings
          UpperCAmelCase          =       num_conv_pos_embedding_groups
          UpperCAmelCase          =       len(self.conv_dim				)
          UpperCAmelCase          =       num_hidden_layers
          UpperCAmelCase          =       intermediate_size
          UpperCAmelCase          =       squeeze_factor
          UpperCAmelCase          =       max_position_embeddings
          UpperCAmelCase          =       position_buckets
          UpperCAmelCase          =       share_att_key
          UpperCAmelCase          =       relative_attention
          UpperCAmelCase          =       norm_rel_ebd
          UpperCAmelCase          =       list(lowercase_				)
          UpperCAmelCase          =       hidden_act
          UpperCAmelCase          =       num_attention_heads
          UpperCAmelCase          =       hidden_dropout
          UpperCAmelCase          =       attention_dropout
          UpperCAmelCase          =       activation_dropout
          UpperCAmelCase          =       feat_proj_dropout
          UpperCAmelCase          =       final_dropout
          UpperCAmelCase          =       layer_norm_eps
          UpperCAmelCase          =       feature_layer_norm_eps
          UpperCAmelCase          =       initializer_range
          UpperCAmelCase          =       vocab_size

          if (
              (len(self.conv_stride				) != self.num_feat_extract_layers)
              or (len(self.conv_kernel				) != self.num_feat_extract_layers)
              or (len(self.conv_dim				) != self.num_feat_extract_layers)
          ):
              raise ValueError(
                  'Configuration for convolutional layers is incorrect.'
                  'It is required that `len(config.conv_dim)` == `len(config.conv_stride)` == `len(config.conv_kernel)`,'
                  F"but is `len(config.conv_dim) = {len(self.conv_dim				)}`, `len(config.conv_stride)"
                  F"= {len(self.conv_stride				)}`, `len(config.conv_kernel) = {len(self.conv_kernel				)}`."				)

          # fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
          UpperCAmelCase          =       apply_spec_augment
          UpperCAmelCase          =       mask_time_prob
          UpperCAmelCase          =       mask_time_length
          UpperCAmelCase          =       mask_time_min_masks
          UpperCAmelCase          =       mask_feature_prob
          UpperCAmelCase          =       mask_feature_length
          UpperCAmelCase          =       mask_feature_min_masks

          # ctc loss
          UpperCAmelCase          =       ctc_loss_reduction
          UpperCAmelCase          =       ctc_zero_infinity

          # sequence classification
          UpperCAmelCase          =       use_weighted_layer_sum
          UpperCAmelCase          =       classifier_proj_size




      @property
      def      a_   (				self				)    ->       Tuple:
          return functools.reduce(operator.mul ,							self.conv_stride ,							1				)



