
"""simple docstring"""






import itertools
import random
import unittest

import numpy as np

from transformers import ASTFeatureExtractor
from transformers.testing_utils import require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


__SCREAMING_SNAKE_CASE							=     random.Random()

if is_torch_available():
 import torch




def       A_  (				__lowercase						,	__lowercase=1.0						,	__lowercase=None						,	__lowercase=None	):
 if rng is None:
  UpperCamelCase_     :					str   	=global_rng

 UpperCamelCase_     :					Tuple   	=[]
 for batch_idx in range(shape[0]	):
  values.append([]	)
  for _ in range(shape[1]	):
   values[-1].append(rng.random() * scale	)

 return values





class 				a__					(						unittest.TestCase       ):



     def __init__(							self :Any      ,  _lowerCamelCase :List[str]      ,  _lowerCamelCase :Dict=7      ,  _lowerCamelCase :Tuple=400      ,  _lowerCamelCase :List[str]=2_000      ,  _lowerCamelCase :int=1      ,  _lowerCamelCase :str=0.0      ,  _lowerCamelCase :str=16_000      ,  _lowerCamelCase :Union[str, Any]=True      ,  _lowerCamelCase :Dict=True      ,  ):







      '''simple docstring'''



      UpperCamelCase_     :					Optional[Any]   	=parent
      UpperCamelCase_     :					Optional[Any]   	=batch_size
      UpperCamelCase_     :					Dict   	=min_seq_length
      UpperCamelCase_     :					Union[str, Any]   	=max_seq_length
      UpperCamelCase_     :					Any   	=(self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
      UpperCamelCase_     :					List[Any]   	=feature_size
      UpperCamelCase_     :					Optional[int]   	=padding_value
      UpperCamelCase_     :					Union[str, Any]   	=sampling_rate
      UpperCamelCase_     :					List[str]   	=return_attention_mask
      UpperCamelCase_     :					Optional[Any]   	=do_normalize



     def  lowerCamelCase_						(							self :Union[str, Any]   ):







      '''simple docstring'''



      return {
          "feature_size": self.feature_size,
          "padding_value": self.padding_value,
          "sampling_rate": self.sampling_rate,
          "return_attention_mask": self.return_attention_mask,
          "do_normalize": self.do_normalize,
      }




     def  lowerCamelCase_						(							self :int      ,  _lowerCamelCase :Dict=False      ,  _lowerCamelCase :int=False   ):







      '''simple docstring'''



      def _flatten(_lowerCamelCase :str   ):
       return list(itertools.chain(*_lowerCamelCase   )   )

      if equal_length:
       UpperCamelCase_     :					Optional[int]   	=floats_list((self.batch_size, self.max_seq_length)   )
      else:
       # make sure that inputs increase in size
       UpperCamelCase_     :					Dict   	=[
           _flatten(floats_list((x, self.feature_size)   )   )
           for x in range(self.min_seq_length      ,  self.max_seq_length      ,  self.seq_length_diff   )
       ]

      if numpify:
       UpperCamelCase_     :					int   	=[np.asarray(_lowerCamelCase   ) for x in speech_inputs]

      return speech_inputs





@require_torch
@require_torchaudio
class 				a__					(						A__   ,							unittest.TestCase       ):
     UpperCAmelCase__         =      ASTFeatureExtractor



     def  lowerCamelCase_						(							self :List[Any]   ):







      '''simple docstring'''



      UpperCamelCase_     :					List[str]   	=ASTFeatureExtractionTester(self   )



     def  lowerCamelCase_						(							self :Tuple   ):







      '''simple docstring'''



      UpperCamelCase_     :					List[str]   	=self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
      # create three inputs of length 800, 1000, and 1200
      UpperCamelCase_     :					List[str]   	=[floats_list((1, x)   )[0] for x in range(800      ,  1_400      ,  200   )]
      UpperCamelCase_     :					int   	=[np.asarray(_lowerCamelCase   ) for speech_input in speech_inputs]

      # Test not batched input
      UpperCamelCase_     :					Dict   	=feat_extract(speech_inputs[0]      ,  return_tensors='np'   ).input_values
      UpperCamelCase_     :					Optional[int]   	=feat_extract(np_speech_inputs[0]      ,  return_tensors='np'   ).input_values
      self.assertTrue(np.allclose(_lowerCamelCase      ,  _lowerCamelCase      ,  atol=1E-3   )   )

      # Test batched
      UpperCamelCase_     :					Dict   	=feat_extract(_lowerCamelCase      ,  padding=_lowerCamelCase      ,  return_tensors='np'   ).input_values
      UpperCamelCase_     :					List[Any]   	=feat_extract(_lowerCamelCase      ,  padding=_lowerCamelCase      ,  return_tensors='np'   ).input_values
      for enc_seq_a, enc_seq_a in zip(_lowerCamelCase      ,  _lowerCamelCase   ):
       self.assertTrue(np.allclose(_lowerCamelCase      ,  _lowerCamelCase      ,  atol=1E-3   )   )

      # Test 2-D numpy arrays are batched.
      UpperCamelCase_     :					Tuple   	=[floats_list((1, x)   )[0] for x in (800, 800, 800)]
      UpperCamelCase_     :					Optional[int]   	=np.asarray(_lowerCamelCase   )
      UpperCamelCase_     :					Union[str, Any]   	=feat_extract(_lowerCamelCase      ,  return_tensors='np'   ).input_values
      UpperCamelCase_     :					Optional[Any]   	=feat_extract(_lowerCamelCase      ,  return_tensors='np'   ).input_values
      for enc_seq_a, enc_seq_a in zip(_lowerCamelCase      ,  _lowerCamelCase   ):
       self.assertTrue(np.allclose(_lowerCamelCase      ,  _lowerCamelCase      ,  atol=1E-3   )   )



     @require_torch
     def  lowerCamelCase_						(							self :Optional[Any]   ):







      '''simple docstring'''



      import torch

      UpperCamelCase_     :					Tuple   	=self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
      UpperCamelCase_     :					List[Any]   	=np.random.rand(100   ).astype(np.floataa   )
      UpperCamelCase_     :					Tuple   	=np_speech_inputs.tolist()

      for inputs in [py_speech_inputs, np_speech_inputs]:
       UpperCamelCase_     :					Dict   	=feature_extractor.pad([{'input_values': inputs}]      ,  return_tensors='np'   )
       self.assertTrue(np_processed.input_values.dtype == np.floataa   )
       UpperCamelCase_     :					Union[str, Any]   	=feature_extractor.pad([{'input_values': inputs}]      ,  return_tensors='pt'   )
       self.assertTrue(pt_processed.input_values.dtype == torch.floataa   )



     def  lowerCamelCase_						(							self :str      ,  _lowerCamelCase :Optional[int]   ):







      '''simple docstring'''



      from datasets import load_dataset

      UpperCamelCase_     :					List[str]   	=load_dataset('hf-internal-testing/librispeech_asr_dummy'      ,  'clean'      ,  split='validation'   )
      # automatic decoding with librispeech
      UpperCamelCase_     :					Union[str, Any]   	=ds.sort('id'   ).select(range(_lowerCamelCase   )   )[:num_samples]['audio']

      return [x["array"] for x in speech_samples]




     @require_torch
     def  lowerCamelCase_						(							self :List[str]   ):







      '''simple docstring'''



      UpperCamelCase_     :					Union[str, Any]   	=torch.tensor(
          [-0.9894, -1.2776, -0.9066, -1.2776, -0.9349, -1.2609, -1.0386, -1.2776,
           -1.1561, -1.2776, -1.2052, -1.2723, -1.2190, -1.2132, -1.2776, -1.1133,
           -1.1953, -1.1343, -1.1584, -1.2203, -1.1770, -1.2474, -1.2381, -1.1936,
           -0.9270, -0.8317, -0.8049, -0.7706, -0.7565, -0.7869]   )
      # fmt: on

      UpperCamelCase_     :					Optional[Any]   	=self._load_datasamples(1   )
      UpperCamelCase_     :					Dict   	=ASTFeatureExtractor()
      UpperCamelCase_     :					List[Any]   	=feature_extractor(_lowerCamelCase      ,  return_tensors='pt'   ).input_values
      self.assertEquals(input_values.shape      ,  (1, 1_024, 128)   )
      self.assertTrue(torch.allclose(input_values[0, 0, :30]      ,  _lowerCamelCase      ,  atol=1E-4   )   )


"""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  lowerCamelCase_						(							self :List[str]   ):







      '''simple docstring'''



      UpperCamelCase_     :					Optional[Any]   	=tempfile.mkdtemp()
      UpperCamelCase_     :					str   	=5

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



     def  lowerCamelCase_						(							self :Union[str, Any]   ):







      '''simple docstring'''



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



     def  lowerCamelCase_						(							self :Dict   ):







      '''simple docstring'''



      shutil.rmtree(self.tmpdirname   )



     def  lowerCamelCase_						(							self :Dict   ):







      '''simple docstring'''



      UpperCamelCase_     :					List[str]   	=RealmConfig(num_block_records=self.num_block_records   )
      return config



     def  lowerCamelCase_						(							self :Optional[int]   ):







      '''simple docstring'''



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



     def  lowerCamelCase_						(							self :Optional[Any]   ):







      '''simple docstring'''



      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=_lowerCamelCase      ,  )
      return block_records



     def  lowerCamelCase_						(							self :Any   ):







      '''simple docstring'''



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



     def  lowerCamelCase_						(							self :int   ):







      '''simple docstring'''



      UpperCamelCase_     :					Any   	=self.get_config()
      UpperCamelCase_     :					Tuple   	=self.get_dummy_retriever()
      UpperCamelCase_     :					Tuple   	=retriever.tokenizer

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

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

      self.assertEqual(len(_lowerCamelCase   )      ,  2   )
      self.assertEqual(len(_lowerCamelCase   )      ,  2   )
      self.assertEqual(len(_lowerCamelCase   )      ,  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  lowerCamelCase_						(							self :Optional[int]   ):







      '''simple docstring'''



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

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

      UpperCamelCase_	,     UpperCamelCase_	,     UpperCamelCase_	,     UpperCamelCase_     :					Union[str, Any]   	=retriever(
          _lowerCamelCase      ,  _lowerCamelCase      ,  answer_ids=_lowerCamelCase      ,  max_length=_lowerCamelCase      ,  return_tensors='np'   )

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




     def  lowerCamelCase_						(							self :List[Any]   ):







      '''simple docstring'''



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

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

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