




"""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


_lowercase	      =    5_03_58
_lowercase	      =    5_03_62

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


				'''simple docstring'''


				def       _lowerCamelCase					(   self   :     Tuple      ):
										snake_case__       :					Union[str, Any]										=					'''openai/whisper-small.en'''
										snake_case__       :					Dict										=					tempfile.mkdtemp()

				def       _lowerCamelCase					(   self   :     int      ,**SCREAMING_SNAKE_CASE_   :     Any      ):
										return WhisperTokenizer.from_pretrained(self.checkpoint      ,**SCREAMING_SNAKE_CASE_      )

				def       _lowerCamelCase					(   self   :     Optional[Any]      ,**SCREAMING_SNAKE_CASE_   :     Tuple      ):
										return WhisperFeatureExtractor.from_pretrained(self.checkpoint      ,**SCREAMING_SNAKE_CASE_      )

				def       _lowerCamelCase					(   self   :     Any      ):
										shutil.rmtree(self.tmpdirname      )

				def       _lowerCamelCase					(   self   :     Union[str, Any]      ):
										snake_case__       :					str										=					self.get_tokenizer()
										snake_case__       :					Dict										=					self.get_feature_extractor()

										snake_case__       :					Union[str, Any]										=					WhisperProcessor(tokenizer=SCREAMING_SNAKE_CASE_      ,feature_extractor=SCREAMING_SNAKE_CASE_      )

										processor.save_pretrained(self.tmpdirname      )
										snake_case__       :					List[Any]										=					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       _lowerCamelCase					(   self   :     Tuple      ):
										snake_case__       :					Optional[Any]										=					WhisperProcessor(tokenizer=self.get_tokenizer()      ,feature_extractor=self.get_feature_extractor()      )
										processor.save_pretrained(self.tmpdirname      )

										snake_case__       :					Any										=					self.get_tokenizer(bos_token='''(BOS)'''      ,eos_token='''(EOS)'''      )
										snake_case__       :					Union[str, Any]										=					self.get_feature_extractor(do_normalize=SCREAMING_SNAKE_CASE_      ,padding_value=1.0      )

										snake_case__       :					Tuple										=					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       _lowerCamelCase					(   self   :     Union[str, Any]      ):
										snake_case__       :					Tuple										=					self.get_feature_extractor()
										snake_case__       :					Optional[int]										=					self.get_tokenizer()

										snake_case__       :					Union[str, Any]										=					WhisperProcessor(tokenizer=SCREAMING_SNAKE_CASE_      ,feature_extractor=SCREAMING_SNAKE_CASE_      )

										snake_case__       :					str										=					floats_list((3, 1_0_0_0)      )

										snake_case__       :					Any										=					feature_extractor(SCREAMING_SNAKE_CASE_      ,return_tensors='''np'''      )
										snake_case__       :					Optional[int]										=					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       _lowerCamelCase					(   self   :     Tuple      ):
										snake_case__       :					int										=					self.get_feature_extractor()
										snake_case__       :					Union[str, Any]										=					self.get_tokenizer()

										snake_case__       :					List[str]										=					WhisperProcessor(tokenizer=SCREAMING_SNAKE_CASE_      ,feature_extractor=SCREAMING_SNAKE_CASE_      )

										snake_case__       :					List[str]										=					'''This is a test string'''

										snake_case__       :					Any										=					processor(text=SCREAMING_SNAKE_CASE_      )

										snake_case__       :					List[Any]										=					tokenizer(SCREAMING_SNAKE_CASE_      )

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

				def       _lowerCamelCase					(   self   :     Optional[Any]      ):
										snake_case__       :					Any										=					self.get_feature_extractor()
										snake_case__       :					str										=					self.get_tokenizer()

										snake_case__       :					List[str]										=					WhisperProcessor(tokenizer=SCREAMING_SNAKE_CASE_      ,feature_extractor=SCREAMING_SNAKE_CASE_      )

										snake_case__       :					int										=					[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

										snake_case__       :					Optional[int]										=					processor.batch_decode(SCREAMING_SNAKE_CASE_      )
										snake_case__       :					Dict										=					tokenizer.batch_decode(SCREAMING_SNAKE_CASE_      )

										self.assertListEqual(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_      )

				def       _lowerCamelCase					(   self   :     Any      ):
										snake_case__       :					Tuple										=					self.get_feature_extractor()
										snake_case__       :					Union[str, Any]										=					self.get_tokenizer()

										snake_case__       :					Union[str, Any]										=					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       _lowerCamelCase					(   self   :     Any      ):
										snake_case__       :					List[Any]										=					self.get_feature_extractor()
										snake_case__       :					Optional[int]										=					self.get_tokenizer()

										snake_case__       :					Any										=					WhisperProcessor(tokenizer=SCREAMING_SNAKE_CASE_      ,feature_extractor=SCREAMING_SNAKE_CASE_      )
										snake_case__       :					Optional[int]										=					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__       :					List[Any]										=					[TRANSCRIBE, NOTIMESTAMPS]
										self.assertListEqual([ids[-1] for ids in forced_decoder_ids]      ,SCREAMING_SNAKE_CASE_      )

				def       _lowerCamelCase					(   self   :     Optional[int]      ):
										snake_case__       :					List[str]										=					WhisperProcessor(tokenizer=self.get_tokenizer()      ,feature_extractor=self.get_feature_extractor()      )
										snake_case__       :					Tuple										=					processor.get_prompt_ids('''Mr. Quilter'''      )
										snake_case__       :					int										=					processor.tokenizer.decode(SCREAMING_SNAKE_CASE_      )

										self.assertListEqual(prompt_ids.tolist()      ,[5_0_3_6_0, 1_7_7_0, 1_3, 2_2_6_4, 3_4_6, 3_5_3]      )
										self.assertEqual(SCREAMING_SNAKE_CASE_      ,'''<|startofprev|> Mr. Quilter'''      )

				def       _lowerCamelCase					(   self   :     List[Any]      ):
										snake_case__       :					Any										=					WhisperProcessor(tokenizer=self.get_tokenizer()      ,feature_extractor=self.get_feature_extractor()      )
										snake_case__       :					str										=					processor.get_prompt_ids(''''''      )
										snake_case__       :					List[str]										=					processor.tokenizer.decode(SCREAMING_SNAKE_CASE_      )

										self.assertListEqual(prompt_ids.tolist()      ,[5_0_3_6_0, 2_2_0]      )
										self.assertEqual(SCREAMING_SNAKE_CASE_      ,'''<|startofprev|> '''      )





				def       _lowerCamelCase					(   self   :     List[Any]      ):
										snake_case__       :					Optional[int]										=					WhisperProcessor(tokenizer=self.get_tokenizer()      ,feature_extractor=self.get_feature_extractor()      )

										def _test_prompt_error_raised_helper(SCREAMING_SNAKE_CASE_   :     Optional[Any]      ,SCREAMING_SNAKE_CASE_   :     Tuple      ):
																with pytest.raises(SCREAMING_SNAKE_CASE_      ) as excinfo:
																						processor.get_prompt_ids(SCREAMING_SNAKE_CASE_      )
																snake_case__       :					Dict										=					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 argparse
import os
import sys
import urllib.request
import zipfile


_lowercase	      =    ["""CoLA""", """SST""", """MRPC""", """QQP""", """STS""", """MNLI""", """SNLI""", """QNLI""", """RTE""", """WNLI""", """diagnostic"""]
_lowercase	      =    {
    """CoLA""": """https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FCoLA.zip?alt=media&token=46d5e637-3411-4188-bc44-5809b5bfb5f4""",
    """SST""": """https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FSST-2.zip?alt=media&token=aabc5f6b-e466-44a2-b9b4-cf6337f84ac8""",
    """MRPC""": """https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2Fmrpc_dev_ids.tsv?alt=media&token=ec5c0836-31d5-48f4-b431-7480817f1adc""",
    """QQP""": """https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FQQP.zip?alt=media&token=700c6acf-160d-4d89-81d1-de4191d02cb5""",
    """STS""": """https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FSTS-B.zip?alt=media&token=bddb94a7-8706-4e0d-a694-1109e12273b5""",
    """MNLI""": """https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FMNLI.zip?alt=media&token=50329ea1-e339-40e2-809c-10c40afff3ce""",
    """SNLI""": """https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FSNLI.zip?alt=media&token=4afcfbb2-ff0c-4b2d-a09a-dbf07926f4df""",
    """QNLI""": """https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FQNLIv2.zip?alt=media&token=6fdcf570-0fc5-4631-8456-9505272d1601""",
    """RTE""": """https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FRTE.zip?alt=media&token=5efa7e85-a0bb-4f19-8ea2-9e1840f077fb""",
    """WNLI""": """https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FWNLI.zip?alt=media&token=068ad0a0-ded7-4bd7-99a5-5e00222e0faf""",
    """diagnostic""": """https://storage.googleapis.com/mtl-sentence-representations.appspot.com/tsvsWithoutLabels%2FAX.tsv?GoogleAccessId=firebase-adminsdk-0khhl@mtl-sentence-representations.iam.gserviceaccount.com&Expires=2498860800&Signature=DuQ2CSPt2Yfre0C%2BiISrVYrIFaZH1Lc7hBVZDD4ZyR7fZYOMNOUGpi8QxBmTNOrNPjR3z1cggo7WXFfrgECP6FBJSsURv8Ybrue8Ypt%2FTPxbuJ0Xc2FhDi%2BarnecCBFO77RSbfuz%2Bs95hRrYhTnByqu3U%2FYZPaj3tZt5QdfpH2IUROY8LiBXoXS46LE%2FgOQc%2FKN%2BA9SoscRDYsnxHfG0IjXGwHN%2Bf88q6hOmAxeNPx6moDulUF6XMUAaXCSFU%2BnRO2RDL9CapWxj%2BDl7syNyHhB7987hZ80B%2FwFkQ3MEs8auvt5XW1%2Bd4aCU7ytgM69r8JDCwibfhZxpaa4gd50QXQ%3D%3D""",
}

_lowercase	      =    """https://dl.fbaipublicfiles.com/senteval/senteval_data/msr_paraphrase_train.txt"""
_lowercase	      =    """https://dl.fbaipublicfiles.com/senteval/senteval_data/msr_paraphrase_test.txt"""

def 		lowerCAmelCase							( lowerCamelCase      ,			lowerCamelCase		):
						'''simple docstring'''


						print(f'''Downloading and extracting {task}...'''		)
						snake_case__       :					Optional[Any]										=					f'''{task}.zip'''
						urllib.request.urlretrieve(TASKaPATH[task]      ,			lowerCamelCase		)
						with zipfile.ZipFile(lowerCamelCase		) as zip_ref:
												zip_ref.extractall(lowerCamelCase		)
						os.remove(lowerCamelCase		)
						print('''\tCompleted!'''		)

def 		lowerCAmelCase							( lowerCamelCase      ,			lowerCamelCase		):
						'''simple docstring'''


						print('''Processing MRPC...'''		)
						snake_case__       :					Union[str, Any]										=					os.path.join(lowerCamelCase      ,			'''MRPC'''		)
						if not os.path.isdir(lowerCamelCase		):
												os.mkdir(lowerCamelCase		)
						if path_to_data:
												snake_case__       :					List[Any]										=					os.path.join(lowerCamelCase      ,			'''msr_paraphrase_train.txt'''		)
												snake_case__       :					Any										=					os.path.join(lowerCamelCase      ,			'''msr_paraphrase_test.txt'''		)
						else:
												print('''Local MRPC data not specified, downloading data from %s''' % MRPC_TRAIN		)
												snake_case__       :					Tuple										=					os.path.join(lowerCamelCase      ,			'''msr_paraphrase_train.txt'''		)
												snake_case__       :					Optional[int]										=					os.path.join(lowerCamelCase      ,			'''msr_paraphrase_test.txt'''		)
												urllib.request.urlretrieve(lowerCamelCase      ,			lowerCamelCase		)
												urllib.request.urlretrieve(lowerCamelCase      ,			lowerCamelCase		)
						if not os.path.isfile(lowerCamelCase		):
												raise ValueError(f'''Train data not found at {mrpc_train_file}'''		)
						if not os.path.isfile(lowerCamelCase		):
												raise ValueError(f'''Test data not found at {mrpc_test_file}'''		)
						urllib.request.urlretrieve(TASKaPATH['''MRPC''']      ,			os.path.join(lowerCamelCase      ,			'''dev_ids.tsv'''		)		)

						snake_case__       :					Optional[Any]										=					[]
						with open(os.path.join(lowerCamelCase      ,			'''dev_ids.tsv'''		)      ,			encoding='''utf8'''		) as ids_fh:
												for row in ids_fh:
																		dev_ids.append(row.strip().split('''\t'''		)		)

						with open(lowerCamelCase      ,			encoding='''utf8'''		) as data_fh, open(
						    os.path.join(lowerCamelCase      ,			'''train.tsv'''		)      ,			'''w'''      ,			encoding='''utf8'''		) as train_fh, open(os.path.join(lowerCamelCase      ,			'''dev.tsv'''		)      ,			'''w'''      ,			encoding='''utf8'''		) as dev_fh:
												snake_case__       :					Any										=					data_fh.readline()
												train_fh.write(lowerCamelCase		)
												dev_fh.write(lowerCamelCase		)
												for row in data_fh:
																		snake_case__			,	snake_case__			,	snake_case__			,	snake_case__			,	snake_case__       :					Dict										=					row.strip().split('''\t'''		)
																		if [ida, ida] in dev_ids:
																								dev_fh.write('''%s\t%s\t%s\t%s\t%s\n''' % (label, ida, ida, sa, sa)		)
																		else:
																								train_fh.write('''%s\t%s\t%s\t%s\t%s\n''' % (label, ida, ida, sa, sa)		)

						with open(lowerCamelCase      ,			encoding='''utf8'''		) as data_fh, open(
						    os.path.join(lowerCamelCase      ,			'''test.tsv'''		)      ,			'''w'''      ,			encoding='''utf8'''		) as test_fh:
												snake_case__       :					Tuple										=					data_fh.readline()
												test_fh.write('''index\t#1 ID\t#2 ID\t#1 String\t#2 String\n'''		)
												for idx, row in enumerate(lowerCamelCase		):
																		snake_case__			,	snake_case__			,	snake_case__			,	snake_case__			,	snake_case__       :					Union[str, Any]										=					row.strip().split('''\t'''		)
																		test_fh.write('''%d\t%s\t%s\t%s\t%s\n''' % (idx, ida, ida, sa, sa)		)
						print('''\tCompleted!'''		)

def 		lowerCAmelCase							( lowerCamelCase		):
						'''simple docstring'''


						print('''Downloading and extracting diagnostic...'''		)
						if not os.path.isdir(os.path.join(lowerCamelCase      ,			'''diagnostic'''		)		):
												os.mkdir(os.path.join(lowerCamelCase      ,			'''diagnostic'''		)		)
						snake_case__       :					Any										=					os.path.join(lowerCamelCase      ,			'''diagnostic'''      ,			'''diagnostic.tsv'''		)
						urllib.request.urlretrieve(TASKaPATH['''diagnostic''']      ,			lowerCamelCase		)
						print('''\tCompleted!'''		)
						return

def 		lowerCAmelCase							( lowerCamelCase		):
						'''simple docstring'''


						snake_case__       :					Any										=					task_names.split(''','''		)
						if "all" in task_names:
												snake_case__       :					str										=					TASKS
						else:
												snake_case__       :					Union[str, Any]										=					[]
												for task_name in task_names:
																		if task_name not in TASKS:
																								raise ValueError(f'''Task {task_name} not found!'''		)
																		tasks.append(lowerCamelCase		)
						return tasks

def 		lowerCAmelCase							( lowerCamelCase		):
						'''simple docstring'''


						snake_case__       :					int										=					argparse.ArgumentParser()
						parser.add_argument('''--data_dir'''      ,			help='''directory to save data to'''      ,			type=lowerCamelCase      ,			default='''glue_data'''		)
						parser.add_argument(
						    '''--tasks'''      ,			help='''tasks to download data for as a comma separated string'''      ,			type=lowerCamelCase      ,			default='''all'''		)
						parser.add_argument(
						    '''--path_to_mrpc'''      ,			help='''path to directory containing extracted MRPC data, msr_paraphrase_train.txt and msr_paraphrase_text.txt'''      ,			type=lowerCamelCase      ,			default=''''''      ,			)
						snake_case__       :					Dict										=					parser.parse_args(lowerCamelCase		)

						if not os.path.isdir(args.data_dir		):
												os.mkdir(args.data_dir		)
						snake_case__       :					Union[str, Any]										=					get_tasks(args.tasks		)

						for task in tasks:
												if task == "MRPC":
																		format_mrpc(args.data_dir      ,			args.path_to_mrpc		)
												elif task == "diagnostic":
																		download_diagnostic(args.data_dir		)
												else:
																		download_and_extract(lowerCamelCase      ,			args.data_dir		)


if __name__ == "__main__":
							sys.exit(main(sys.argv[1:]))


