
"""simple docstring"""






import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTJConfig, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, tooslow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
 import jax
 import jax.numpy as jnp

 from transformers.modeling_flax_pytorch_utils import (
     convert_pytorch_state_dict_to_flax,
     load_flax_weights_in_pytorch_model,
 )
 from transformers.models.gptj.modeling_flax_gptj import FlaxGPTJForCausalLM, FlaxGPTJModel

if is_torch_available():
 import torch





class 				a__					:



     def __init__(							self :Tuple      ,  _lowerCamelCase :Tuple      ,  _lowerCamelCase :Dict=14      ,  _lowerCamelCase :int=7      ,  _lowerCamelCase :str=True      ,  _lowerCamelCase :Any=True      ,  _lowerCamelCase :List[str]=False      ,  _lowerCamelCase :List[str]=True      ,  _lowerCamelCase :Any=99      ,  _lowerCamelCase :Optional[int]=32      ,  _lowerCamelCase :Any=4      ,  _lowerCamelCase :Optional[Any]=4      ,  _lowerCamelCase :Tuple=4      ,  _lowerCamelCase :List[Any]=37      ,  _lowerCamelCase :str="gelu"      ,  _lowerCamelCase :str=0.1      ,  _lowerCamelCase :Optional[int]=0.1      ,  _lowerCamelCase :Any=512      ,  _lowerCamelCase :str=0.02      ,  ):







      '''simple docstring'''



      UpperCamelCase_     :					Union[str, Any]   	=parent
      UpperCamelCase_     :					Any   	=batch_size
      UpperCamelCase_     :					int   	=seq_length
      UpperCamelCase_     :					int   	=is_training
      UpperCamelCase_     :					str   	=use_input_mask
      UpperCamelCase_     :					Optional[Any]   	=use_token_type_ids
      UpperCamelCase_     :					Dict   	=use_labels
      UpperCamelCase_     :					Union[str, Any]   	=vocab_size
      UpperCamelCase_     :					Tuple   	=hidden_size
      UpperCamelCase_     :					Optional[Any]   	=rotary_dim
      UpperCamelCase_     :					str   	=num_hidden_layers
      UpperCamelCase_     :					Union[str, Any]   	=num_attention_heads
      UpperCamelCase_     :					Tuple   	=intermediate_size
      UpperCamelCase_     :					Dict   	=hidden_act
      UpperCamelCase_     :					Tuple   	=hidden_dropout_prob
      UpperCamelCase_     :					Any   	=attention_probs_dropout_prob
      UpperCamelCase_     :					Dict   	=max_position_embeddings
      UpperCamelCase_     :					str   	=initializer_range
      UpperCamelCase_     :					Tuple   	=None
      UpperCamelCase_     :					Optional[Any]   	=vocab_size - 1
      UpperCamelCase_     :					Dict   	=vocab_size - 1
      UpperCamelCase_     :					Tuple   	=vocab_size - 1



     def  lowerCamelCase_						(							self :Optional[Any]   ):







      '''simple docstring'''



      UpperCamelCase_     :					Union[str, Any]   	=ids_tensor([self.batch_size, self.seq_length]      ,  self.vocab_size   )

      UpperCamelCase_     :					Dict   	=None
      if self.use_input_mask:
       UpperCamelCase_     :					int   	=random_attention_mask([self.batch_size, self.seq_length]   )

      UpperCamelCase_     :					Any   	=GPTJConfig(
          vocab_size=self.vocab_size      ,  n_embd=self.hidden_size      ,  n_layer=self.num_hidden_layers      ,  n_head=self.num_attention_heads      ,  n_positions=self.max_position_embeddings      ,  use_cache=_lowerCamelCase      ,  bos_token_id=self.bos_token_id      ,  eos_token_id=self.eos_token_id      ,  pad_token_id=self.pad_token_id      ,  rotary_dim=self.rotary_dim      ,  )

      return (config, input_ids, input_mask)



     def  lowerCamelCase_						(							self :List[str]   ):







      '''simple docstring'''



      UpperCamelCase_     :					List[Any]   	=self.prepare_config_and_inputs()
      UpperCamelCase_	,     UpperCamelCase_	,     UpperCamelCase_     :					int   	=config_and_inputs
      UpperCamelCase_     :					Any   	={'input_ids': input_ids, 'attention_mask': attention_mask}
      return config, inputs_dict



     def  lowerCamelCase_						(							self :Union[str, Any]      ,  _lowerCamelCase :Tuple      ,  _lowerCamelCase :int      ,  _lowerCamelCase :Union[str, Any]      ,  _lowerCamelCase :List[Any]   ):







      '''simple docstring'''



      UpperCamelCase_     :					Any   	=20
      UpperCamelCase_     :					int   	=model_class_name(_lowerCamelCase   )

      UpperCamelCase_     :					List[str]   	=model.init_cache(input_ids.shape[0]      ,  _lowerCamelCase   )
      UpperCamelCase_     :					int   	=jnp.ones((input_ids.shape[0], max_decoder_length)      ,  dtype='i4'   )

      UpperCamelCase_     :					Tuple   	=jnp.broadcast_to(
          jnp.arange(input_ids.shape[-1] - 1   )[None, :]      ,  (input_ids.shape[0], input_ids.shape[-1] - 1)   )
      UpperCamelCase_     :					str   	=model(
          input_ids[:, :-1]      ,  attention_mask=_lowerCamelCase      ,  past_key_values=_lowerCamelCase      ,  position_ids=_lowerCamelCase      ,  )

      UpperCamelCase_     :					Tuple   	=jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]      ,  dtype='i4'   )
      UpperCamelCase_     :					str   	=model(
          input_ids[:, -1:]      ,  attention_mask=_lowerCamelCase      ,  past_key_values=outputs_cache.past_key_values      ,  position_ids=_lowerCamelCase      ,  )

      UpperCamelCase_     :					Union[str, Any]   	=model(_lowerCamelCase   )

      UpperCamelCase_     :					Any   	=np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])   )   )
      self.parent.assertTrue(diff < 1E-3      ,  msg=f'''Max diff is {diff}'''   )




     def  lowerCamelCase_						(							self :Tuple      ,  _lowerCamelCase :Union[str, Any]      ,  _lowerCamelCase :List[str]      ,  _lowerCamelCase :Optional[int]      ,  _lowerCamelCase :Union[str, Any]   ):







      '''simple docstring'''



      UpperCamelCase_     :					Any   	=20
      UpperCamelCase_     :					Dict   	=model_class_name(_lowerCamelCase   )

      UpperCamelCase_     :					Tuple   	=jnp.concatenate(
          [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])   )]      ,  axis=-1      ,  )

      UpperCamelCase_     :					Dict   	=model.init_cache(input_ids.shape[0]      ,  _lowerCamelCase   )
      UpperCamelCase_     :					Dict   	=jnp.broadcast_to(
          jnp.arange(input_ids.shape[-1] - 1   )[None, :]      ,  (input_ids.shape[0], input_ids.shape[-1] - 1)   )

      UpperCamelCase_     :					Optional[int]   	=model(
          input_ids[:, :-1]      ,  attention_mask=_lowerCamelCase      ,  past_key_values=_lowerCamelCase      ,  position_ids=_lowerCamelCase      ,  )
      UpperCamelCase_     :					Optional[int]   	=jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]      ,  dtype='i4'   )
      UpperCamelCase_     :					Union[str, Any]   	=model(
          input_ids[:, -1:]      ,  past_key_values=outputs_cache.past_key_values      ,  attention_mask=_lowerCamelCase      ,  position_ids=_lowerCamelCase      ,  )

      UpperCamelCase_     :					str   	=model(_lowerCamelCase      ,  attention_mask=_lowerCamelCase   )

      UpperCamelCase_     :					List[Any]   	=np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])   )   )
      self.parent.assertTrue(diff < 1E-3      ,  msg=f'''Max diff is {diff}'''   )





@require_flax
class 				a__					(						A__   ,							A__   ,							unittest.TestCase       ):
     UpperCAmelCase__         =      (FlaxGPTJModel, FlaxGPTJForCausalLM) if is_flax_available() else ()
     UpperCAmelCase__         =      (FlaxGPTJForCausalLM,) if is_flax_available() else ()



     def  lowerCamelCase_						(							self :List[str]   ):







      '''simple docstring'''



      UpperCamelCase_     :					Tuple   	=FlaxGPTJModelTester(self   )



     def  lowerCamelCase_						(							self :Union[str, Any]   ):







      '''simple docstring'''



      for model_class_name in self.all_model_classes:
       UpperCamelCase_	,     UpperCamelCase_	,     UpperCamelCase_     :					List[str]   	=self.model_tester.prepare_config_and_inputs()
       self.model_tester.check_use_cache_forward(_lowerCamelCase      ,  _lowerCamelCase      ,  _lowerCamelCase      ,  _lowerCamelCase   )



     def  lowerCamelCase_						(							self :Dict   ):







      '''simple docstring'''



      for model_class_name in self.all_model_classes:
       UpperCamelCase_	,     UpperCamelCase_	,     UpperCamelCase_     :					Tuple   	=self.model_tester.prepare_config_and_inputs()
       self.model_tester.check_use_cache_forward_with_attn_mask(
           _lowerCamelCase      ,  _lowerCamelCase      ,  _lowerCamelCase      ,  _lowerCamelCase   )



     @tooslow
     def  lowerCamelCase_						(							self :Union[str, Any]   ):







      '''simple docstring'''



      UpperCamelCase_     :					Any   	=GPTaTokenizer.from_pretrained('gpt2'      ,  pad_token='<|endoftext|>'      ,  padding_side='left'   )
      UpperCamelCase_     :					Optional[int]   	=tokenizer(['Hello this is a long string', 'Hey']      ,  return_tensors='np'      ,  padding=_lowerCamelCase      ,  truncation=_lowerCamelCase   )

      UpperCamelCase_     :					Any   	=FlaxGPTJForCausalLM.from_pretrained('EleutherAI/gpt-j-6B'   )
      UpperCamelCase_     :					Dict   	=False
      UpperCamelCase_     :					str   	=model.config.eos_token_id

      UpperCamelCase_     :					Dict   	=jax.jit(model.generate   )

      UpperCamelCase_     :					int   	=jit_generate(
          inputs['input_ids']      ,  attention_mask=inputs['attention_mask']      ,  pad_token_id=tokenizer.pad_token_id   ).sequences

      UpperCamelCase_     :					str   	=tokenizer.batch_decode(_lowerCamelCase      ,  skip_special_tokens=_lowerCamelCase   )

      UpperCamelCase_     :					Dict   	=[
          'Hello this is a long string of text.\n\nI\'m trying to get the text of the',
          'Hey, I\'m a little late to the party. I\'m going to',
      ]

      self.assertListEqual(_lowerCamelCase      ,  _lowerCamelCase   )



     @is_pt_flax_cross_test
     def  lowerCamelCase_						(							self :Tuple   ):







      '''simple docstring'''



      UpperCamelCase_	,     UpperCamelCase_     :					List[Any]   	=self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
       with self.subTest(model_class.__name__   ):
        # prepare inputs
        UpperCamelCase_     :					Optional[Any]   	=self._prepare_for_class(_lowerCamelCase      ,  _lowerCamelCase   )
        UpperCamelCase_     :					Tuple   	={k: torch.tensor(v.tolist()   ) for k, v in prepared_inputs_dict.items()}

        # load corresponding PyTorch class
        UpperCamelCase_     :					Tuple   	=model_class.__name__[4:]  # Skip the "Flax" at the beginning
        UpperCamelCase_     :					Tuple   	=getattr(_lowerCamelCase      ,  _lowerCamelCase   )

        UpperCamelCase_	,     UpperCamelCase_     :					List[str]   	=pt_inputs['input_ids'].shape
        UpperCamelCase_     :					Union[str, Any]   	=np.random.randint(0      ,  seq_length - 1      ,  size=(batch_size,)   )
        for batch_idx, start_index in enumerate(_lowerCamelCase   ):
         UpperCamelCase_     :					Any   	=0
         UpperCamelCase_     :					Union[str, Any]   	=1
         UpperCamelCase_     :					List[str]   	=0
         UpperCamelCase_     :					Any   	=1
        UpperCamelCase_     :					List[str]   	=pt_model_class(_lowerCamelCase   ).eval()
        UpperCamelCase_     :					Any   	=model_class(_lowerCamelCase      ,  dtype=jnp.floataa   )

        UpperCamelCase_     :					Optional[int]   	=convert_pytorch_state_dict_to_flax(pt_model.state_dict()      ,  _lowerCamelCase   )
        UpperCamelCase_     :					Dict   	=fx_state

        with torch.no_grad():
         UpperCamelCase_     :					List[Any]   	=pt_model(**_lowerCamelCase   ).to_tuple()

        UpperCamelCase_     :					str   	=fx_model(**_lowerCamelCase   ).to_tuple()
        self.assertEqual(len(_lowerCamelCase   )      ,  len(_lowerCamelCase   )      ,  'Output lengths differ between Flax and PyTorch'   )
        for fx_output, pt_output in zip(_lowerCamelCase      ,  _lowerCamelCase   ):
         self.assert_almost_equals(fx_output[:, -1]      ,  pt_output[:, -1].numpy()      ,  4E-2   )

        with tempfile.TemporaryDirectory() as tmpdirname:
         pt_model.save_pretrained(_lowerCamelCase   )
         UpperCamelCase_     :					Tuple   	=model_class.from_pretrained(_lowerCamelCase      ,  from_pt=_lowerCamelCase   )

        UpperCamelCase_     :					Optional[int]   	=fx_model_loaded(**_lowerCamelCase   ).to_tuple()
        self.assertEqual(
            len(_lowerCamelCase   )      ,  len(_lowerCamelCase   )      ,  'Output lengths differ between Flax and PyTorch'   )
        for fx_output_loaded, pt_output in zip(_lowerCamelCase      ,  _lowerCamelCase   ):
         self.assert_almost_equals(fx_output_loaded[:, -1]      ,  pt_output[:, -1].numpy()      ,  4E-2   )



     @is_pt_flax_cross_test
     def  lowerCamelCase_						(							self :Optional[Any]   ):







      '''simple docstring'''



      UpperCamelCase_	,     UpperCamelCase_     :					Any   	=self.model_tester.prepare_config_and_inputs_for_common()
      for model_class in self.all_model_classes:
       with self.subTest(model_class.__name__   ):
        # prepare inputs
        UpperCamelCase_     :					int   	=self._prepare_for_class(_lowerCamelCase      ,  _lowerCamelCase   )
        UpperCamelCase_     :					Dict   	={k: torch.tensor(v.tolist()   ) for k, v in prepared_inputs_dict.items()}

        # load corresponding PyTorch class
        UpperCamelCase_     :					Tuple   	=model_class.__name__[4:]  # Skip the "Flax" at the beginning
        UpperCamelCase_     :					Optional[int]   	=getattr(_lowerCamelCase      ,  _lowerCamelCase   )

        UpperCamelCase_     :					Tuple   	=pt_model_class(_lowerCamelCase   ).eval()
        UpperCamelCase_     :					str   	=model_class(_lowerCamelCase      ,  dtype=jnp.floataa   )

        UpperCamelCase_     :					Optional[int]   	=load_flax_weights_in_pytorch_model(_lowerCamelCase      ,  fx_model.params   )
        UpperCamelCase_	,     UpperCamelCase_     :					Union[str, Any]   	=pt_inputs['input_ids'].shape
        UpperCamelCase_     :					str   	=np.random.randint(0      ,  seq_length - 1      ,  size=(batch_size,)   )
        for batch_idx, start_index in enumerate(_lowerCamelCase   ):
         UpperCamelCase_     :					Union[str, Any]   	=0
         UpperCamelCase_     :					str   	=1
         UpperCamelCase_     :					List[Any]   	=0
         UpperCamelCase_     :					Any   	=1

        # make sure weights are tied in PyTorch
        pt_model.tie_weights()

        with torch.no_grad():
         UpperCamelCase_     :					Dict   	=pt_model(**_lowerCamelCase   ).to_tuple()

        UpperCamelCase_     :					str   	=fx_model(**_lowerCamelCase   ).to_tuple()
        self.assertEqual(len(_lowerCamelCase   )      ,  len(_lowerCamelCase   )      ,  'Output lengths differ between Flax and PyTorch'   )
        for fx_output, pt_output in zip(_lowerCamelCase      ,  _lowerCamelCase   ):
         self.assert_almost_equals(fx_output[:, -1]      ,  pt_output[:, -1].numpy()      ,  4E-2   )

        with tempfile.TemporaryDirectory() as tmpdirname:
         fx_model.save_pretrained(_lowerCamelCase   )
         UpperCamelCase_     :					str   	=pt_model_class.from_pretrained(_lowerCamelCase      ,  from_flax=_lowerCamelCase   )

        with torch.no_grad():
         UpperCamelCase_     :					Tuple   	=pt_model_loaded(**_lowerCamelCase   ).to_tuple()

        self.assertEqual(
            len(_lowerCamelCase   )      ,  len(_lowerCamelCase   )      ,  'Output lengths differ between Flax and PyTorch'   )
        for fx_output, pt_output in zip(_lowerCamelCase      ,  _lowerCamelCase   ):
         self.assert_almost_equals(fx_output[:, -1]      ,  pt_output[:, -1].numpy()      ,  4E-2   )




     @tooslow
     def  lowerCamelCase_						(							self :Optional[Any]   ):







      '''simple docstring'''



      for model_class_name in self.all_model_classes:
       UpperCamelCase_     :					Tuple   	=model_class_name.from_pretrained('EleutherAI/gpt-j-6B'   )
       UpperCamelCase_     :					Optional[Any]   	=model(np.ones((1, 1)   )   )
       self.assertIsNotNone(_lowerCamelCase   )


"""simple docstring"""






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__SCREAMING_SNAKE_CASE							=     {
    'configuration_whisper': ['WHISPER_PRETRAINED_CONFIG_ARCHIVE_MAP', 'WhisperConfig', 'WhisperOnnxConfig'],
    'feature_extraction_whisper': ['WhisperFeatureExtractor'],
    'processing_whisper': ['WhisperProcessor'],
    'tokenization_whisper': ['WhisperTokenizer'],
}

try:
 if not is_tokenizers_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __SCREAMING_SNAKE_CASE							=     ['WhisperTokenizerFast']

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __SCREAMING_SNAKE_CASE							=     [
     'WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST',
     'WhisperForConditionalGeneration',
     'WhisperModel',
     'WhisperPreTrainedModel',
     'WhisperForAudioClassification',
 ]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __SCREAMING_SNAKE_CASE							=     [
     'TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST',
     'TFWhisperForConditionalGeneration',
     'TFWhisperModel',
     'TFWhisperPreTrainedModel',
 ]

try:
 if not is_flax_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __SCREAMING_SNAKE_CASE							=     [
     'FlaxWhisperForConditionalGeneration',
     'FlaxWhisperModel',
     'FlaxWhisperPreTrainedModel',
     'FlaxWhisperForAudioClassification',
 ]


if TYPE_CHECKING:
 from .configuration_whisper import WHISPER_PRETRAINED_CONFIG_ARCHIVE_MAP, WhisperConfig, WhisperOnnxConfig
 from .feature_extraction_whisper import WhisperFeatureExtractor
 from .processing_whisper import WhisperProcessor
 from .tokenization_whisper import WhisperTokenizer

 try:
  if not is_tokenizers_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_whisper_fast import WhisperTokenizerFast

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_whisper import (
      WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST,
      WhisperForAudioClassification,
      WhisperForConditionalGeneration,
      WhisperModel,
      WhisperPreTrainedModel,
  )

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_whisper import (
      TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFWhisperForConditionalGeneration,
      TFWhisperModel,
      TFWhisperPreTrainedModel,
  )

 try:
  if not is_flax_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_flax_whisper import (
      FlaxWhisperForAudioClassification,
      FlaxWhisperForConditionalGeneration,
      FlaxWhisperModel,
      FlaxWhisperPreTrainedModel,
  )

else:
 import sys

 __SCREAMING_SNAKE_CASE							=     _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)
