

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


lowerCAmelCase__             =		5_0_3_5_8
lowerCAmelCase__             =		5_0_3_6_2







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


		"""simple docstring"""




		def 		A__       (					self):
			_UpperCamelCase :		Any								=    'openai/whisper-small.en'
			_UpperCamelCase :		Optional[int]								=    tempfile.mkdtemp()


		def 		A__       (					self						,							**__snake_case):
			return WhisperTokenizer.from_pretrained(self.checkpoint						,							**__snake_case)


		def 		A__       (					self						,							**__snake_case):
			return WhisperFeatureExtractor.from_pretrained(self.checkpoint						,							**__snake_case)


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


		def 		A__       (					self):
			_UpperCamelCase :		int								=    self.get_tokenizer()
			_UpperCamelCase :		Optional[int]								=    self.get_feature_extractor()

			_UpperCamelCase :		str								=    WhisperProcessor(tokenizer=__snake_case						,							feature_extractor=__snake_case)

			processor.save_pretrained(self.tmpdirname)
			_UpperCamelCase :		Dict								=    WhisperProcessor.from_pretrained(self.tmpdirname)

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

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


		def 		A__       (					self):
			_UpperCamelCase :		List[str]								=    WhisperProcessor(tokenizer=self.get_tokenizer()						,							feature_extractor=self.get_feature_extractor())
			processor.save_pretrained(self.tmpdirname)

			_UpperCamelCase :		Dict								=    self.get_tokenizer(bos_token='(BOS)'						,							eos_token='(EOS)')
			_UpperCamelCase :		Tuple								=    self.get_feature_extractor(do_normalize=__snake_case						,							padding_value=1.0)

			_UpperCamelCase :		Optional[int]								=    WhisperProcessor.from_pretrained(
			    self.tmpdirname						,							bos_token='(BOS)'						,							eos_token='(EOS)'						,							do_normalize=__snake_case						,							padding_value=1.0)

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

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


		def 		A__       (					self):
			_UpperCamelCase :		List[Any]								=    self.get_feature_extractor()
			_UpperCamelCase :		int								=    self.get_tokenizer()

			_UpperCamelCase :		Dict								=    WhisperProcessor(tokenizer=__snake_case						,							feature_extractor=__snake_case)

			_UpperCamelCase :		Any								=    floats_list((3, 10_00))

			_UpperCamelCase :		Union[str, Any]								=    feature_extractor(__snake_case						,							return_tensors='np')
			_UpperCamelCase :		Union[str, Any]								=    processor(__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 		A__       (					self):
			_UpperCamelCase :		List[str]								=    self.get_feature_extractor()
			_UpperCamelCase :		Dict								=    self.get_tokenizer()

			_UpperCamelCase :		Tuple								=    WhisperProcessor(tokenizer=__snake_case						,							feature_extractor=__snake_case)

			_UpperCamelCase :		str								=    'This is a test string'

			_UpperCamelCase :		Optional[Any]								=    processor(text=__snake_case)

			_UpperCamelCase :		int								=    tokenizer(__snake_case)

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


		def 		A__       (					self):
			_UpperCamelCase :		Union[str, Any]								=    self.get_feature_extractor()
			_UpperCamelCase :		List[Any]								=    self.get_tokenizer()

			_UpperCamelCase :		int								=    WhisperProcessor(tokenizer=__snake_case						,							feature_extractor=__snake_case)

			_UpperCamelCase :		List[Any]								=    [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

			_UpperCamelCase :		List[str]								=    processor.batch_decode(__snake_case)
			_UpperCamelCase :		str								=    tokenizer.batch_decode(__snake_case)

			self.assertListEqual(__snake_case						,							__snake_case)


		def 		A__       (					self):
			_UpperCamelCase :		Union[str, Any]								=    self.get_feature_extractor()
			_UpperCamelCase :		Optional[Any]								=    self.get_tokenizer()

			_UpperCamelCase :		List[Any]								=    WhisperProcessor(tokenizer=__snake_case						,							feature_extractor=__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 		A__       (					self):
			_UpperCamelCase :		Tuple								=    self.get_feature_extractor()
			_UpperCamelCase :		List[str]								=    self.get_tokenizer()

			_UpperCamelCase :		Tuple								=    WhisperProcessor(tokenizer=__snake_case						,							feature_extractor=__snake_case)
			_UpperCamelCase :		Tuple								=    processor.get_decoder_prompt_ids(task='transcribe'						,							no_timestamps=__snake_case)

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

			_UpperCamelCase :		Tuple								=    [TRANSCRIBE, NOTIMESTAMPS]
			self.assertListEqual([ids[-1] for ids in forced_decoder_ids]						,							__snake_case)


		def 		A__       (					self):
			_UpperCamelCase :		Any								=    WhisperProcessor(tokenizer=self.get_tokenizer()						,							feature_extractor=self.get_feature_extractor())
			_UpperCamelCase :		Tuple								=    processor.get_prompt_ids('Mr. Quilter')
			_UpperCamelCase :		Dict								=    processor.tokenizer.decode(__snake_case)

			self.assertListEqual(prompt_ids.tolist()						,							[5_03_60, 17_70, 13, 22_64, 3_46, 3_53])
			self.assertEqual(__snake_case						,							'<|startofprev|> Mr. Quilter')


		def 		A__       (					self):
			_UpperCamelCase :		List[Any]								=    WhisperProcessor(tokenizer=self.get_tokenizer()						,							feature_extractor=self.get_feature_extractor())
			_UpperCamelCase :		Union[str, Any]								=    processor.get_prompt_ids('')
			_UpperCamelCase :		str								=    processor.tokenizer.decode(__snake_case)

			self.assertListEqual(prompt_ids.tolist()						,							[5_03_60, 2_20])
			self.assertEqual(__snake_case						,							'<|startofprev|> ')






		def 		A__       (					self):
			_UpperCamelCase :		Union[str, Any]								=    WhisperProcessor(tokenizer=self.get_tokenizer()						,							feature_extractor=self.get_feature_extractor())

			def _test_prompt_error_raised_helper(__snake_case						,							__snake_case):
				with pytest.raises(__snake_case) as excinfo:
					processor.get_prompt_ids(__snake_case)
				_UpperCamelCase :		Dict								=    f'''Encountered text in the prompt corresponding to disallowed special token: {special_token}.'''
				self.assertEqual(__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|>')







import os

from apache_beam.io.filesystems import FileSystems
from apache_beam.pipeline import Pipeline

from .logging import get_logger


lowerCAmelCase__             =		2 << 2_0  # 2mb
lowerCAmelCase__             =		get_logger(__name__)







class 						lowercase     (					_lowercase				):


		"""simple docstring"""




		def 		A__       (					self):
			_UpperCamelCase :		Optional[Any]								=    self._options.get_all_options().get('runner')
			return runner in [None, "DirectRunner", "PortableRunner"]






def        lowerCamelCase_  (				UpperCAmelCase_				:  Any				,							UpperCAmelCase_				:  Optional[Any]				,							UpperCAmelCase_				:  str=False						)    ->					str:




	'''simple docstring'''




	_UpperCamelCase :		Optional[int]								=    FileSystems
	if fs.exists(UpperCAmelCase_						):
		if force_upload:
			logger.info(F'''Remote path already exist: {remote_file_path}. Overwriting it as force_upload=True.'''						)
		else:
			logger.info(F'''Remote path already exist: {remote_file_path}. Skipping it as force_upload=False.'''						)
			return
	with fs.create(UpperCAmelCase_						) as remote_file:
		with open(UpperCAmelCase_				,							'rb'						) as local_file:
			_UpperCamelCase :		List[str]								=    local_file.read(UpperCAmelCase_						)
			while chunk:
				remote_file.write(UpperCAmelCase_						)
				_UpperCamelCase :		List[str]								=    local_file.read(UpperCAmelCase_						)







def        lowerCamelCase_  (				UpperCAmelCase_				:  Any				,							UpperCAmelCase_				:  List[Any]				,							UpperCAmelCase_				:  Optional[int]=False						)    ->					Any:




	'''simple docstring'''




	_UpperCamelCase :		List[str]								=    FileSystems
	if os.path.exists(UpperCAmelCase_						):
		if force_download:
			logger.info(F'''Local path already exist: {remote_file_path}. Overwriting it as force_upload=True.'''						)
		else:
			logger.info(F'''Local path already exist: {remote_file_path}. Skipping it as force_upload=False.'''						)
			return
	with fs.open(UpperCAmelCase_						) as remote_file:
		with open(UpperCAmelCase_				,							'wb'						) as local_file:
			_UpperCamelCase :		Union[str, Any]								=    remote_file.read(UpperCAmelCase_						)
			while chunk:
				local_file.write(UpperCAmelCase_						)
				_UpperCamelCase :		List[str]								=    remote_file.read(UpperCAmelCase_						)




