



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_	       =   50_358
UpperCamelCase_	       =   50_362







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

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

						def 	__a       (    self	:Tuple   ,						**_lowercase	:int)		->					Union[str, Any]:
												return WhisperTokenizer.from_pretrained(self.checkpoint   ,						**_lowercase)

						def 	__a       (    self	:Any   ,						**_lowercase	:Union[str, Any])		->					Dict:
												return WhisperFeatureExtractor.from_pretrained(self.checkpoint   ,						**_lowercase)

						def 	__a       (    self	:Optional[int])		->					List[Any]:
												shutil.rmtree(self.tmpdirname)

						def 	__a       (    self	:Optional[int])		->					Any:
												UpperCAmelCase_     							=     self.get_tokenizer()
												UpperCAmelCase_     							=     self.get_feature_extractor()

												UpperCAmelCase_     							=     WhisperProcessor(tokenizer=_lowercase   ,						feature_extractor=_lowercase)

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

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

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

						def 	__a       (    self	:str)		->					Optional[int]:
												UpperCAmelCase_     							=     WhisperProcessor(tokenizer=self.get_tokenizer()   ,						feature_extractor=self.get_feature_extractor())
												processor.save_pretrained(self.tmpdirname)

												UpperCAmelCase_     							=     self.get_tokenizer(bos_token='''(BOS)'''   ,						eos_token='''(EOS)''')
												UpperCAmelCase_     							=     self.get_feature_extractor(do_normalize=_lowercase   ,						padding_value=1.0)

												UpperCAmelCase_     							=     WhisperProcessor.from_pretrained(
												    self.tmpdirname   ,						bos_token='''(BOS)'''   ,						eos_token='''(EOS)'''   ,						do_normalize=_lowercase   ,						padding_value=1.0)

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

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

						def 	__a       (    self	:Any)		->					Dict:
												UpperCAmelCase_     							=     self.get_feature_extractor()
												UpperCAmelCase_     							=     self.get_tokenizer()

												UpperCAmelCase_     							=     WhisperProcessor(tokenizer=_lowercase   ,						feature_extractor=_lowercase)

												UpperCAmelCase_     							=     floats_list((3, 1000))

												UpperCAmelCase_     							=     feature_extractor(_lowercase   ,						return_tensors='''np''')
												UpperCAmelCase_     							=     processor(_lowercase   ,						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	:Optional[Any])		->					Optional[int]:
												UpperCAmelCase_     							=     self.get_feature_extractor()
												UpperCAmelCase_     							=     self.get_tokenizer()

												UpperCAmelCase_     							=     WhisperProcessor(tokenizer=_lowercase   ,						feature_extractor=_lowercase)

												UpperCAmelCase_     							=     '''This is a test string'''

												UpperCAmelCase_     							=     processor(text=_lowercase)

												UpperCAmelCase_     							=     tokenizer(_lowercase)

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

						def 	__a       (    self	:Optional[Any])		->					List[str]:
												UpperCAmelCase_     							=     self.get_feature_extractor()
												UpperCAmelCase_     							=     self.get_tokenizer()

												UpperCAmelCase_     							=     WhisperProcessor(tokenizer=_lowercase   ,						feature_extractor=_lowercase)

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

												UpperCAmelCase_     							=     processor.batch_decode(_lowercase)
												UpperCAmelCase_     							=     tokenizer.batch_decode(_lowercase)

												self.assertListEqual(_lowercase   ,						_lowercase)

						def 	__a       (    self	:Optional[Any])		->					Union[str, Any]:
												UpperCAmelCase_     							=     self.get_feature_extractor()
												UpperCAmelCase_     							=     self.get_tokenizer()

												UpperCAmelCase_     							=     WhisperProcessor(tokenizer=_lowercase   ,						feature_extractor=_lowercase)

												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	:List[str])		->					Dict:
												UpperCAmelCase_     							=     self.get_feature_extractor()
												UpperCAmelCase_     							=     self.get_tokenizer()

												UpperCAmelCase_     							=     WhisperProcessor(tokenizer=_lowercase   ,						feature_extractor=_lowercase)
												UpperCAmelCase_     							=     processor.get_decoder_prompt_ids(task='''transcribe'''   ,						no_timestamps=_lowercase)

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

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

						def 	__a       (    self	:Optional[Any])		->					Optional[int]:
												UpperCAmelCase_     							=     WhisperProcessor(tokenizer=self.get_tokenizer()   ,						feature_extractor=self.get_feature_extractor())
												UpperCAmelCase_     							=     processor.get_prompt_ids('''Mr. Quilter''')
												UpperCAmelCase_     							=     processor.tokenizer.decode(_lowercase)

												self.assertListEqual(prompt_ids.tolist()   ,						[50360, 1770, 13, 2264, 346, 353])
												self.assertEqual(_lowercase   ,						'''<|startofprev|> Mr. Quilter''')

						def 	__a       (    self	:Dict)		->					str:
												UpperCAmelCase_     							=     WhisperProcessor(tokenizer=self.get_tokenizer()   ,						feature_extractor=self.get_feature_extractor())
												UpperCAmelCase_     							=     processor.get_prompt_ids('''''')
												UpperCAmelCase_     							=     processor.tokenizer.decode(_lowercase)

												self.assertListEqual(prompt_ids.tolist()   ,						[50360, 220])
												self.assertEqual(_lowercase   ,						'''<|startofprev|> ''')




						def 	__a       (    self	:str)		->					Union[str, Any]:
												UpperCAmelCase_     							=     WhisperProcessor(tokenizer=self.get_tokenizer()   ,						feature_extractor=self.get_feature_extractor())

												def _test_prompt_error_raised_helper(_lowercase	:int   ,						_lowercase	:Optional[Any]):
																		with pytest.raises(_lowercase) as excinfo:
																								processor.get_prompt_ids(_lowercase)
																		UpperCAmelCase_     							=     f"Encountered text in the prompt corresponding to disallowed special token: {special_token}."
																		self.assertEqual(_lowercase   ,						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|>''')










from .chunk_utils import chunk_layer
from .data_transforms import make_atomaa_masks
from .feats import atomaa_to_atomaa, frames_and_literature_positions_to_atomaa_pos, torsion_angles_to_frames
from .loss import compute_predicted_aligned_error, compute_tm
from .protein import Protein as OFProtein
from .protein import to_pdb
from .rigid_utils import Rigid, Rotation
from .tensor_utils import dict_multimap, flatten_final_dims, permute_final_dims





