





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       __lowerCAmelCase (    unittest.TestCase  ):






			def     lowerCamelCase						(self					)  ->     Tuple:


							'''simple docstring'''



							snake_case_		:					int			    =						'''openai/whisper-small.en'''
							snake_case_		:					Optional[Any]			    =						tempfile.mkdtemp()






			def     lowerCamelCase						(self      ,							**__magic_name__					)  ->     Optional[Any]:


							'''simple docstring'''



							return WhisperTokenizer.from_pretrained(self.checkpoint      ,							**__magic_name__					)






			def     lowerCamelCase						(self      ,							**__magic_name__					)  ->     str:


							'''simple docstring'''



							return WhisperFeatureExtractor.from_pretrained(self.checkpoint      ,							**__magic_name__					)






			def     lowerCamelCase						(self					)  ->     Union[str, Any]:


							'''simple docstring'''



							shutil.rmtree(self.tmpdirname					)






			def     lowerCamelCase						(self					)  ->     Union[str, Any]:


							'''simple docstring'''



							snake_case_		:					List[str]			    =						self.get_tokenizer()
							snake_case_		:					Any			    =						self.get_feature_extractor()

							snake_case_		:					Dict			    =						WhisperProcessor(tokenizer=__magic_name__      ,							feature_extractor=__magic_name__					)

							processor.save_pretrained(self.tmpdirname					)
							snake_case_		:					Tuple			    =						WhisperProcessor.from_pretrained(self.tmpdirname					)

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

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






			def     lowerCamelCase						(self					)  ->     Any:


							'''simple docstring'''



							snake_case_		:					Dict			    =						WhisperProcessor(tokenizer=self.get_tokenizer()      ,							feature_extractor=self.get_feature_extractor()					)
							processor.save_pretrained(self.tmpdirname					)

							snake_case_		:					Dict			    =						self.get_tokenizer(bos_token='''(BOS)'''      ,							eos_token='''(EOS)'''					)
							snake_case_		:					List[str]			    =						self.get_feature_extractor(do_normalize=__magic_name__      ,							padding_value=1.0					)

							snake_case_		:					List[Any]			    =						WhisperProcessor.from_pretrained(
							    self.tmpdirname      ,							bos_token='''(BOS)'''      ,							eos_token='''(EOS)'''      ,							do_normalize=__magic_name__      ,							padding_value=1.0					)

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

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






			def     lowerCamelCase						(self					)  ->     Optional[int]:


							'''simple docstring'''



							snake_case_		:					Any			    =						self.get_feature_extractor()
							snake_case_		:					str			    =						self.get_tokenizer()

							snake_case_		:					Union[str, Any]			    =						WhisperProcessor(tokenizer=__magic_name__      ,							feature_extractor=__magic_name__					)

							snake_case_		:					Tuple			    =						floats_list((3, 1000)					)

							snake_case_		:					int			    =						feature_extractor(__magic_name__      ,							return_tensors='''np'''					)
							snake_case_		:					int			    =						processor(__magic_name__      ,							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					)  ->     str:


							'''simple docstring'''



							snake_case_		:					Optional[int]			    =						self.get_feature_extractor()
							snake_case_		:					int			    =						self.get_tokenizer()

							snake_case_		:					Any			    =						WhisperProcessor(tokenizer=__magic_name__      ,							feature_extractor=__magic_name__					)

							snake_case_		:					Tuple			    =						'''This is a test string'''

							snake_case_		:					Dict			    =						processor(text=__magic_name__					)

							snake_case_		:					str			    =						tokenizer(__magic_name__					)

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






			def     lowerCamelCase						(self					)  ->     str:


							'''simple docstring'''



							snake_case_		:					Optional[Any]			    =						self.get_feature_extractor()
							snake_case_		:					Tuple			    =						self.get_tokenizer()

							snake_case_		:					Dict			    =						WhisperProcessor(tokenizer=__magic_name__      ,							feature_extractor=__magic_name__					)

							snake_case_		:					Optional[Any]			    =						[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

							snake_case_		:					str			    =						processor.batch_decode(__magic_name__					)
							snake_case_		:					Tuple			    =						tokenizer.batch_decode(__magic_name__					)

							self.assertListEqual(__magic_name__      ,							__magic_name__					)






			def     lowerCamelCase						(self					)  ->     List[Any]:


							'''simple docstring'''



							snake_case_		:					List[Any]			    =						self.get_feature_extractor()
							snake_case_		:					Dict			    =						self.get_tokenizer()

							snake_case_		:					Tuple			    =						WhisperProcessor(tokenizer=__magic_name__      ,							feature_extractor=__magic_name__					)

							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					)  ->     List[str]:


							'''simple docstring'''



							snake_case_		:					int			    =						self.get_feature_extractor()
							snake_case_		:					Any			    =						self.get_tokenizer()

							snake_case_		:					Optional[int]			    =						WhisperProcessor(tokenizer=__magic_name__      ,							feature_extractor=__magic_name__					)
							snake_case_		:					Tuple			    =						processor.get_decoder_prompt_ids(task='''transcribe'''      ,							no_timestamps=__magic_name__					)

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

							snake_case_		:					Optional[int]			    =						[TRANSCRIBE, NOTIMESTAMPS]
							self.assertListEqual([ids[-1] for ids in forced_decoder_ids]      ,							__magic_name__					)






			def     lowerCamelCase						(self					)  ->     Union[str, Any]:


							'''simple docstring'''



							snake_case_		:					Optional[Any]			    =						WhisperProcessor(tokenizer=self.get_tokenizer()      ,							feature_extractor=self.get_feature_extractor()					)
							snake_case_		:					Union[str, Any]			    =						processor.get_prompt_ids('''Mr. Quilter'''					)
							snake_case_		:					List[Any]			    =						processor.tokenizer.decode(__magic_name__					)

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






			def     lowerCamelCase						(self					)  ->     int:


							'''simple docstring'''



							snake_case_		:					Any			    =						WhisperProcessor(tokenizer=self.get_tokenizer()      ,							feature_extractor=self.get_feature_extractor()					)
							snake_case_		:					List[str]			    =						processor.get_prompt_ids(''''''					)
							snake_case_		:					int			    =						processor.tokenizer.decode(__magic_name__					)

							self.assertListEqual(prompt_ids.tolist()      ,							[5_0360, 220]					)
							self.assertEqual(__magic_name__      ,							'''<|startofprev|> '''					)






			def     lowerCamelCase						(self					)  ->     Tuple:


							'''simple docstring'''



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

							def _test_prompt_error_raised_helper(__magic_name__      ,							__magic_name__					):
											with pytest.raises(__magic_name__					) as excinfo:
															processor.get_prompt_ids(__magic_name__					)
											snake_case_		:					str			    =						F'''Encountered text in the prompt corresponding to disallowed special token: {special_token}.'''
											self.assertEqual(__magic_name__      ,							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 typing import Optional, Tuple

import flax.linen as nn
import jax
import jax.numpy as jnp
from configuration_hybrid_clip import HybridCLIPConfig
from flax.core.frozen_dict import FrozenDict

from transformers import FLAX_MODEL_MAPPING, FlaxCLIPVisionModel
from transformers.modeling_flax_utils import FlaxPreTrainedModel
from transformers.models.clip.modeling_flax_clip import FlaxCLIPOutput
from transformers.utils import logging


lowerCAmelCase_												=       logging.get_logger(__name__)



class       __lowerCAmelCase (    nn.Module  ):
			lowerCamelCase_  :  HybridCLIPConfig
			lowerCamelCase_  :  jnp.dtype     =   jnp.floataa






			def     lowerCamelCase						(self					)  ->     int:


							'''simple docstring'''



							snake_case_		:					Any			    =						self.config.text_config
							snake_case_		:					Any			    =						self.config.vision_config

							snake_case_		:					Union[str, Any]			    =						self.config.projection_dim
							snake_case_		:					Dict			    =						text_config.hidden_size
							snake_case_		:					Union[str, Any]			    =						vision_config.hidden_size

							snake_case_		:					Union[str, Any]			    =						FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
							snake_case_		:					Dict			    =						FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__      ,							__magic_name__					).module_class

							snake_case_		:					Optional[Any]			    =						text_module(__magic_name__      ,							dtype=self.dtype					)
							snake_case_		:					Dict			    =						vision_module(__magic_name__      ,							dtype=self.dtype					)

							snake_case_		:					Dict			    =						nn.Dense(
							    self.projection_dim      ,							dtype=self.dtype      ,							kernel_init=jax.nn.initializers.normal(0.02					)      ,							use_bias=__magic_name__      ,							)
							snake_case_		:					Optional[Any]			    =						nn.Dense(
							    self.projection_dim      ,							dtype=self.dtype      ,							kernel_init=jax.nn.initializers.normal(0.02					)      ,							use_bias=__magic_name__      ,							)
							snake_case_		:					Optional[Any]			    =						self.param('''logit_scale'''      ,							jax.nn.initializers.ones      ,							[]					)






			def __call__(self      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__ = True      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__=None      ,							)  ->     Any:


							'''simple docstring'''



							snake_case_		:					Tuple			    =						return_dict if return_dict is not None else self.config.return_dict

							snake_case_		:					Tuple			    =						self.vision_model(
							    pixel_values=__magic_name__      ,							deterministic=__magic_name__      ,							output_attentions=__magic_name__      ,							output_hidden_states=__magic_name__      ,							return_dict=__magic_name__      ,							)

							snake_case_		:					Tuple			    =						self.text_model(
							    input_ids=__magic_name__      ,							attention_mask=__magic_name__      ,							token_type_ids=__magic_name__      ,							position_ids=__magic_name__      ,							deterministic=__magic_name__      ,							output_attentions=__magic_name__      ,							output_hidden_states=__magic_name__      ,							return_dict=__magic_name__      ,							)

							snake_case_		:					Dict			    =						vision_outputs[1]
							snake_case_		:					Optional[Any]			    =						self.visual_projection(__magic_name__					)

							snake_case_		:					List[str]			    =						text_outputs[1]
							snake_case_		:					Union[str, Any]			    =						self.text_projection(__magic_name__					)

							# normalized features
							snake_case_		:					List[str]			    =						image_embeds / jnp.linalg.norm(__magic_name__      ,							axis=-1      ,							keepdims=__magic_name__					)
							snake_case_		:					Any			    =						text_embeds / jnp.linalg.norm(__magic_name__      ,							axis=-1      ,							keepdims=__magic_name__					)

							# cosine similarity as logits
							snake_case_		:					Optional[Any]			    =						jnp.exp(self.logit_scale					)
							snake_case_		:					str			    =						jnp.matmul(__magic_name__      ,							image_embeds.T					) * logit_scale
							snake_case_		:					List[str]			    =						logits_per_text.T

							if not return_dict:
											return (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs)

							return FlaxCLIPOutput(
							    logits_per_image=__magic_name__      ,							logits_per_text=__magic_name__      ,							text_embeds=__magic_name__      ,							image_embeds=__magic_name__      ,							text_model_output=__magic_name__      ,							vision_model_output=__magic_name__      ,							)




class       __lowerCAmelCase (    _a  ):
			lowerCamelCase_  :  Dict     =   HybridCLIPConfig
			lowerCamelCase_  :  int     =   FlaxHybridCLIPModule






			def __init__(self      ,							__magic_name__      ,							__magic_name__ = None      ,							__magic_name__ = 0      ,							__magic_name__ = jnp.floataa      ,							**__magic_name__      ,							)  ->     str:


							'''simple docstring'''



							if input_shape is None:
											snake_case_		:					int			    =						((1, 1), (1, config.vision_config.image_size, config.vision_config.image_size, 3))

							snake_case_		:					Union[str, Any]			    =						self.module_class(config=__magic_name__      ,							dtype=__magic_name__      ,							**__magic_name__					)
							super().__init__(__magic_name__      ,							__magic_name__      ,							input_shape=__magic_name__      ,							seed=__magic_name__      ,							dtype=__magic_name__					)






			def     lowerCamelCase						(self      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__ = None					)  ->     FrozenDict:


							'''simple docstring'''



							snake_case_		:					List[Any]			    =						jnp.zeros(input_shape[0]      ,							dtype='''i4'''					)
							snake_case_		:					Optional[Any]			    =						jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__magic_name__					).shape[-1]					)      ,							input_shape[0]					)
							snake_case_		:					Union[str, Any]			    =						jnp.ones_like(__magic_name__					)
							snake_case_		:					List[Any]			    =						jnp.ones_like(__magic_name__					)

							snake_case_		:					Union[str, Any]			    =						jax.random.normal(__magic_name__      ,							input_shape[1]					)

							snake_case_    ,							snake_case_		:					Optional[int]			    =						jax.random.split(__magic_name__					)
							snake_case_		:					int			    =						{'''params''': params_rng, '''dropout''': dropout_rng}

							return self.module.init(__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__					)["params"]






			def __call__(self      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__ = None      ,							__magic_name__ = None      ,							__magic_name__ = False      ,							__magic_name__ = None      ,							__magic_name__ = None      ,							__magic_name__ = None      ,							)  ->     Any:


							'''simple docstring'''



							snake_case_		:					Optional[int]			    =						output_attentions if output_attentions is not None else self.config.output_attentions
							snake_case_		:					Tuple			    =						(
							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
							)
							snake_case_		:					List[Any]			    =						return_dict if return_dict is not None else self.config.return_dict

							if position_ids is None:
											snake_case_		:					Tuple			    =						jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__magic_name__					).shape[-1]					)      ,							input_ids.shape					)

							if token_type_ids is None:
											snake_case_		:					Tuple			    =						jnp.zeros_like(__magic_name__					)

							if attention_mask is None:
											snake_case_		:					Optional[Any]			    =						jnp.ones_like(__magic_name__					)

							# Handle any PRNG if needed
							snake_case_		:					List[str]			    =						{}
							if dropout_rng is not None:
											snake_case_		:					Tuple			    =						dropout_rng

							return self.module.apply(
							    {'''params''': params or self.params}      ,							jnp.array(__magic_name__      ,							dtype='''i4'''					)      ,							jnp.array(__magic_name__      ,							dtype=jnp.floataa					)      ,							jnp.array(__magic_name__      ,							dtype='''i4'''					)      ,							jnp.array(__magic_name__      ,							dtype='''i4'''					)      ,							jnp.array(__magic_name__      ,							dtype='''i4'''					)      ,							not train      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							rngs=__magic_name__      ,							)






			def     lowerCamelCase						(self      ,							__magic_name__      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__=None      ,							__magic_name__ = None      ,							__magic_name__ = None      ,							__magic_name__=False      ,							)  ->     Union[str, Any]:


							'''simple docstring'''



							if position_ids is None:
											snake_case_		:					Tuple			    =						jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__magic_name__					).shape[-1]					)      ,							input_ids.shape					)

							if token_type_ids is None:
											snake_case_		:					List[str]			    =						jnp.zeros_like(__magic_name__					)

							if attention_mask is None:
											snake_case_		:					Union[str, Any]			    =						jnp.ones_like(__magic_name__					)

							# Handle any PRNG if needed
							snake_case_		:					Optional[int]			    =						{}
							if dropout_rng is not None:
											snake_case_		:					int			    =						dropout_rng

							def _get_features(__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__					):
											snake_case_		:					int			    =						module.text_model(
											    input_ids=__magic_name__      ,							attention_mask=__magic_name__      ,							position_ids=__magic_name__      ,							token_type_ids=__magic_name__      ,							deterministic=__magic_name__      ,							)
											snake_case_		:					Optional[int]			    =						text_outputs[1]
											snake_case_		:					int			    =						module.text_projection(__magic_name__					)
											return text_features

							return self.module.apply(
							    {'''params''': params or self.params}      ,							jnp.array(__magic_name__      ,							dtype='''i4'''					)      ,							jnp.array(__magic_name__      ,							dtype='''i4'''					)      ,							jnp.array(__magic_name__      ,							dtype='''i4'''					)      ,							jnp.array(__magic_name__      ,							dtype='''i4'''					)      ,							not train      ,							method=_get_features      ,							rngs=__magic_name__      ,							)






			def     lowerCamelCase						(self      ,							__magic_name__      ,							__magic_name__ = None      ,							__magic_name__ = None      ,							__magic_name__=False					)  ->     Dict:


							'''simple docstring'''



							snake_case_		:					Any			    =						{}
							if dropout_rng is not None:
											snake_case_		:					Optional[int]			    =						dropout_rng

							def _get_features(__magic_name__      ,							__magic_name__      ,							__magic_name__					):
											snake_case_		:					Dict			    =						module.vision_model(pixel_values=__magic_name__      ,							deterministic=__magic_name__					)
											snake_case_		:					str			    =						vision_outputs[1]  # pooled_output
											snake_case_		:					str			    =						module.visual_projection(__magic_name__					)
											return image_features

							return self.module.apply(
							    {'''params''': params or self.params}      ,							jnp.array(__magic_name__      ,							dtype=jnp.floataa					)      ,							not train      ,							method=_get_features      ,							rngs=__magic_name__      ,							)






			@classmethod
			def     lowerCamelCase						(cls      ,							__magic_name__ = None      ,							__magic_name__ = None      ,							*__magic_name__      ,							**__magic_name__      ,							)  ->     FlaxPreTrainedModel:


							'''simple docstring'''



							snake_case_		:					Optional[Any]			    =						{
							    argument[len('''text_'''					) :]: value for argument, value in kwargs.items() if argument.startswith('''text_'''					)
							}

							snake_case_		:					Dict			    =						{
							    argument[len('''vision_'''					) :]: value for argument, value in kwargs.items() if argument.startswith('''vision_'''					)
							}

							# remove text, vision kwargs from kwargs
							for key in kwargs_text.keys():
											del kwargs["text_" + key]
							for key in kwargs_vision.keys():
											del kwargs["vision_" + key]

							# Load and initialize the text and vision model
							snake_case_		:					str			    =						kwargs_text.pop('''model'''      ,							__magic_name__					)
							if text_model is None:
											assert (
											    text_model_name_or_path is not None
											), "If `model` is not defined as an argument, a `text_model_name_or_path` has to be defined"
											from transformers import FlaxAutoModel

											if "config" not in kwargs_text:
															from transformers import AutoConfig

															snake_case_		:					Dict			    =						AutoConfig.from_pretrained(__magic_name__					)
															snake_case_		:					int			    =						text_config

											snake_case_		:					Optional[int]			    =						FlaxAutoModel.from_pretrained(__magic_name__      ,							*__magic_name__      ,							**__magic_name__					)

							snake_case_		:					Any			    =						kwargs_vision.pop('''model'''      ,							__magic_name__					)
							if vision_model is None:
											assert (
											    vision_model_name_or_path is not None
											), "If `model` is not defined as an argument, a `vision_model_name_or_path` has to be defined"
											from transformers import FlaxAutoModel

											if "config" not in kwargs_vision:
															from transformers import AutoConfig

															snake_case_		:					Optional[int]			    =						AutoConfig.from_pretrained(__magic_name__					)
															snake_case_		:					List[str]			    =						vision_config

											snake_case_		:					List[Any]			    =						FlaxAutoModel.from_pretrained(__magic_name__      ,							*__magic_name__      ,							**__magic_name__					)

							# instantiate config with corresponding kwargs
							snake_case_		:					Dict			    =						kwargs.pop('''dtype'''      ,							jnp.floataa					)
							snake_case_		:					str			    =						HybridCLIPConfig.from_text_vision_configs(text_model.config      ,							vision_model.config      ,							**__magic_name__					)

							# init model
							snake_case_		:					int			    =						cls(__magic_name__      ,							*__magic_name__      ,							dtype=__magic_name__      ,							**__magic_name__					)

							if vision_config.model_type == "clip":
											snake_case_		:					int			    =						vision_model.params['''vision_model''']
											snake_case_		:					Any			    =						vision_model.params['''visual_projection''']['''kernel''']
							else:
											snake_case_		:					List[Any]			    =						vision_model.params

							snake_case_		:					Union[str, Any]			    =						text_model.params

							return model
