




'''simple docstring'''





import shutil
import tempfile
import unittest

import pytest

from transformers import WhisperTokenizer, is_speech_available
from transformers.testing_utils import require_sentencepiece, require_torch, require_torchaudio

from .test_feature_extraction_whisper import floats_list


if is_speech_available():
   from transformers import WhisperFeatureExtractor, WhisperProcessor


UpperCamelCase_		:    List[str]									=		50358
UpperCamelCase_		:    Any									=		50362

@require_torch
@require_torchaudio
@require_sentencepiece
class      _a       (		unittest.TestCase							):






     def        _lowercase			(				self					)							->    int:
          _snake_case							=   "openai/whisper-small.en"
          _snake_case							=   tempfile.mkdtemp()






     def        _lowercase			(				self      ,**_SCREAMING_SNAKE_CASE					)							->    Any:
          return WhisperTokenizer.from_pretrained(self.checkpoint      ,**_SCREAMING_SNAKE_CASE					)






     def        _lowercase			(				self      ,**_SCREAMING_SNAKE_CASE					)							->    List[str]:
          return WhisperFeatureExtractor.from_pretrained(self.checkpoint      ,**_SCREAMING_SNAKE_CASE					)






     def        _lowercase			(				self					)							->    Tuple:
          shutil.rmtree(self.tmpdirname					)






     def        _lowercase			(				self					)							->    Tuple:
          _snake_case							=   self.get_tokenizer()
          _snake_case							=   self.get_feature_extractor()

          _snake_case							=   WhisperProcessor(tokenizer=_SCREAMING_SNAKE_CASE      ,feature_extractor=_SCREAMING_SNAKE_CASE					)

          processor.save_pretrained(self.tmpdirname					)
          _snake_case							=   WhisperProcessor.from_pretrained(self.tmpdirname					)

          self.assertEqual(processor.tokenizer.get_vocab()      ,tokenizer.get_vocab()					)
          self.assertIsInstance(processor.tokenizer      ,_SCREAMING_SNAKE_CASE					)

          self.assertEqual(processor.feature_extractor.to_json_string()      ,feature_extractor.to_json_string()					)
          self.assertIsInstance(processor.feature_extractor      ,_SCREAMING_SNAKE_CASE					)






     def        _lowercase			(				self					)							->    Dict:
          _snake_case							=   WhisperProcessor(tokenizer=self.get_tokenizer()      ,feature_extractor=self.get_feature_extractor()					)
          processor.save_pretrained(self.tmpdirname					)

          _snake_case							=   self.get_tokenizer(bos_token="(BOS)"      ,eos_token="(EOS)"					)
          _snake_case							=   self.get_feature_extractor(do_normalize=_SCREAMING_SNAKE_CASE      ,padding_value=1.0					)

          _snake_case							=   WhisperProcessor.from_pretrained(
              self.tmpdirname      ,bos_token="(BOS)"      ,eos_token="(EOS)"      ,do_normalize=_SCREAMING_SNAKE_CASE      ,padding_value=1.0					)

          self.assertEqual(processor.tokenizer.get_vocab()      ,tokenizer_add_kwargs.get_vocab()					)
          self.assertIsInstance(processor.tokenizer      ,_SCREAMING_SNAKE_CASE					)

          self.assertEqual(processor.feature_extractor.to_json_string()      ,feature_extractor_add_kwargs.to_json_string()					)
          self.assertIsInstance(processor.feature_extractor      ,_SCREAMING_SNAKE_CASE					)






     def        _lowercase			(				self					)							->    Union[str, Any]:
          _snake_case							=   self.get_feature_extractor()
          _snake_case							=   self.get_tokenizer()

          _snake_case							=   WhisperProcessor(tokenizer=_SCREAMING_SNAKE_CASE      ,feature_extractor=_SCREAMING_SNAKE_CASE					)

          _snake_case							=   floats_list((3, 1_000)					)

          _snake_case							=   feature_extractor(_SCREAMING_SNAKE_CASE      ,return_tensors="np"					)
          _snake_case							=   processor(_SCREAMING_SNAKE_CASE      ,return_tensors="np"					)

          for key in input_feat_extract.keys():
               self.assertAlmostEqual(input_feat_extract[key].sum()      ,input_processor[key].sum()      ,delta=1e-2					)






     def        _lowercase			(				self					)							->    List[str]:
          _snake_case							=   self.get_feature_extractor()
          _snake_case							=   self.get_tokenizer()

          _snake_case							=   WhisperProcessor(tokenizer=_SCREAMING_SNAKE_CASE      ,feature_extractor=_SCREAMING_SNAKE_CASE					)

          _snake_case							=   "This is a test string"

          _snake_case							=   processor(text=_SCREAMING_SNAKE_CASE					)

          _snake_case							=   tokenizer(_SCREAMING_SNAKE_CASE					)

          for key in encoded_tok.keys():
               self.assertListEqual(encoded_tok[key]      ,encoded_processor[key]					)






     def        _lowercase			(				self					)							->    Tuple:
          _snake_case							=   self.get_feature_extractor()
          _snake_case							=   self.get_tokenizer()

          _snake_case							=   WhisperProcessor(tokenizer=_SCREAMING_SNAKE_CASE      ,feature_extractor=_SCREAMING_SNAKE_CASE					)

          _snake_case							=   [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

          _snake_case							=   processor.batch_decode(_SCREAMING_SNAKE_CASE					)
          _snake_case							=   tokenizer.batch_decode(_SCREAMING_SNAKE_CASE					)

          self.assertListEqual(_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE					)






     def        _lowercase			(				self					)							->    Optional[Any]:
          _snake_case							=   self.get_feature_extractor()
          _snake_case							=   self.get_tokenizer()

          _snake_case							=   WhisperProcessor(tokenizer=_SCREAMING_SNAKE_CASE      ,feature_extractor=_SCREAMING_SNAKE_CASE					)

          self.assertListEqual(
              processor.model_input_names      ,feature_extractor.model_input_names      ,msg="`processor` and `feature_extractor` model input names do not match"      ,)






     def        _lowercase			(				self					)							->    Optional[int]:
          _snake_case							=   self.get_feature_extractor()
          _snake_case							=   self.get_tokenizer()

          _snake_case							=   WhisperProcessor(tokenizer=_SCREAMING_SNAKE_CASE      ,feature_extractor=_SCREAMING_SNAKE_CASE					)
          _snake_case							=   processor.get_decoder_prompt_ids(task="transcribe"      ,no_timestamps=_SCREAMING_SNAKE_CASE					)

          self.assertIsInstance(_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE					)
          for ids in forced_decoder_ids:
               self.assertIsInstance(_SCREAMING_SNAKE_CASE      ,(list, tuple)					)

          _snake_case							=   [TRANSCRIBE, NOTIMESTAMPS]
          self.assertListEqual([ids[-1] for ids in forced_decoder_ids]      ,_SCREAMING_SNAKE_CASE					)






     def        _lowercase			(				self					)							->    Dict:
          _snake_case							=   WhisperProcessor(tokenizer=self.get_tokenizer()      ,feature_extractor=self.get_feature_extractor()					)
          _snake_case							=   processor.get_prompt_ids("Mr. Quilter"					)
          _snake_case							=   processor.tokenizer.decode(_SCREAMING_SNAKE_CASE					)

          self.assertListEqual(prompt_ids.tolist()      ,[50_360, 1_770, 13, 2_264, 346, 353]					)
          self.assertEqual(_SCREAMING_SNAKE_CASE      ,"<|startofprev|> Mr. Quilter"					)






     def        _lowercase			(				self					)							->    Tuple:
          _snake_case							=   WhisperProcessor(tokenizer=self.get_tokenizer()      ,feature_extractor=self.get_feature_extractor()					)
          _snake_case							=   processor.get_prompt_ids(""					)
          _snake_case							=   processor.tokenizer.decode(_SCREAMING_SNAKE_CASE					)

          self.assertListEqual(prompt_ids.tolist()      ,[50_360, 220]					)
          self.assertEqual(_SCREAMING_SNAKE_CASE      ,"<|startofprev|> "					)






     def        _lowercase			(				self					)							->    Dict:
          _snake_case							=   WhisperProcessor(tokenizer=self.get_tokenizer()      ,feature_extractor=self.get_feature_extractor()					)

          def _test_prompt_error_raised_helper(_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE					):
               with pytest.raises(_SCREAMING_SNAKE_CASE					) as excinfo:
                    processor.get_prompt_ids(_SCREAMING_SNAKE_CASE					)
               _snake_case							=   f"""Encountered text in the prompt corresponding to disallowed special token: {special_token}."""
               self.assertEqual(_SCREAMING_SNAKE_CASE      ,str(excinfo.value					)					)

          _test_prompt_error_raised_helper("<|startofprev|> test"      ,"<|startofprev|>"					)
          _test_prompt_error_raised_helper("test <|notimestamps|>"      ,"<|notimestamps|>"					)
          _test_prompt_error_raised_helper("test <|zh|> test <|transcribe|>"      ,"<|zh|>"					)








'''simple docstring'''





import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


UpperCamelCase_		:    List[str]									=		500000

UpperCamelCase_			,			UpperCamelCase_		:    Dict									=		os.path.split(__file__)
UpperCamelCase_		:    int									=		os.path.join(RESULTS_BASEPATH, '''results''', RESULTS_FILENAME.replace('''.py''', '''.json'''))





@get_duration
def 			__a			(     _UpperCamelCase:					datasets.Dataset					)		->				List[Any]:

     """simple docstring"""




     _snake_case							=   dataset.select(range(0      ,		len(_UpperCamelCase					)      ,		2					)					)





@get_duration
def 			__a			(     _UpperCamelCase:					datasets.Dataset					)		->				List[Any]:

     """simple docstring"""




     _snake_case							=   dataset.sort("numbers"					)





@get_duration
def 			__a			(     _UpperCamelCase:					datasets.Dataset					)		->				List[Any]:

     """simple docstring"""




     _snake_case							=   dataset.shuffle()





@get_duration
def 			__a			(     _UpperCamelCase:					datasets.Dataset					)		->				Any:

     """simple docstring"""




     _snake_case							=   dataset.train_test_split(0.1					)





@get_duration
def 			__a			(     _UpperCamelCase:					datasets.Dataset      ,		_UpperCamelCase:					int=10					)		->				List[str]:

     """simple docstring"""




     for shard_id in range(_UpperCamelCase					):
          _snake_case							=   dataset.shard(_UpperCamelCase      ,		_UpperCamelCase					)





def 			__a			(     )		->				int:

     """simple docstring"""




     _snake_case							=   {"num examples": SPEED_TEST_N_EXAMPLES}
     _snake_case							=   (select, sort, shuffle, train_test_split, shard)
     with tempfile.TemporaryDirectory() as tmp_dir:
          print("generating dataset"					)
          _snake_case							=   datasets.Features({"text": datasets.Value("string"					), "numbers": datasets.Value("float32"					)}					)
          _snake_case							=   generate_example_dataset(
              os.path.join(_UpperCamelCase      ,		"dataset.arrow"					)      ,		_UpperCamelCase      ,		num_examples=_UpperCamelCase					)
          print("Functions"					)
          for func in functions:
               print(func.__name__					)
               _snake_case							=   func(_UpperCamelCase					)

     with open(_UpperCamelCase      ,		"wb"					) as f:
          f.write(json.dumps(_UpperCamelCase					).encode("utf-8"					)					)


if __name__ == "__main__":  # useful to run the profiler
   benchmark_indices_mapping()


