






'''simple docstring'''







import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import OPTConfig, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, slow

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


if is_flax_available():
       import os

       # The slow tests are often failing with OOM error on GPU
       # This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
       # but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
       __A   :							Any     	=						'platform'

       import jax
       import jax.numpy as jnp

       from transformers import FlaxOPTForCausalLM, FlaxOPTModel, GPTaTokenizer


def     _UpperCamelCase				(  snake_case__				,      snake_case__				,      snake_case__=None				,      snake_case__=None				):
      if attention_mask is None:
            __A		 =   np.where(input_ids != config.pad_token_id				,      1				,      0				)
      return {
          "input_ids": input_ids,
          "attention_mask": attention_mask,
      }






@require_flax
class        a__     :
     '''simple docstring'''







     def __init__(self     :						str					,					lowercase__     :						Optional[Any]					,					lowercase__     :						Optional[Any]=13					,					lowercase__     :						Union[str, Any]=7					,					lowercase__     :						int=True					,					lowercase__     :						List[str]=False					,					lowercase__     :						str=99					,					lowercase__     :						int=16					,					lowercase__     :						Union[str, Any]=2					,					lowercase__     :						Dict=4					,					lowercase__     :						Any=4					,					lowercase__     :						Tuple="gelu"					,					lowercase__     :						int=0.1					,					lowercase__     :						str=0.1					,					lowercase__     :						Union[str, Any]=20					,					lowercase__     :						Tuple=2					,					lowercase__     :						str=1					,					lowercase__     :						List[str]=0					,					lowercase__     :						Any=16					,					lowercase__     :						Any=16					,					lowercase__     :						Dict=0.02					,					)	->							List[Any]:
           __A		 =   parent
           __A		 =   batch_size
           __A		 =   seq_length
           __A		 =   is_training
           __A		 =   use_labels
           __A		 =   vocab_size
           __A		 =   hidden_size
           __A		 =   num_hidden_layers
           __A		 =   num_attention_heads
           __A		 =   intermediate_size
           __A		 =   hidden_act
           __A		 =   hidden_dropout_prob
           __A		 =   attention_probs_dropout_prob
           __A		 =   max_position_embeddings
           __A		 =   eos_token_id
           __A		 =   pad_token_id
           __A		 =   bos_token_id
           __A		 =   embed_dim
           __A		 =   word_embed_proj_dim
           __A		 =   initializer_range
           __A		 =   False






     def      _UpperCamelCase (self     :						Optional[Any]       )	->							Union[str, Any]:
           __A		 =   np.clip(ids_tensor([self.batch_size, self.seq_length - 1]					,					self.vocab_size       )					,					3					,					self.vocab_size       )
           __A		 =   np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)					,					dtype=np.intaa       ))					,					-1       )

           __A		 =   OPTConfig(
               vocab_size=self.vocab_size					,					hidden_size=self.hidden_size					,					num_hidden_layers=self.num_hidden_layers					,					num_attention_heads=self.num_attention_heads					,					ffn_dim=self.intermediate_size					,					dropout=self.hidden_dropout_prob					,					attention_dropout=self.attention_probs_dropout_prob					,					max_position_embeddings=self.max_position_embeddings					,					eos_token_id=self.eos_token_id					,					bos_token_id=self.bos_token_id					,					pad_token_id=self.pad_token_id					,					embed_dim=self.embed_dim					,					is_encoder_decoder=lowercase__					,					word_embed_proj_dim=self.word_embed_proj_dim					,					initializer_range=self.initializer_range					,					use_cache=lowercase__					,					)
           __A		 =   prepare_opt_inputs_dict(lowercase__					,					lowercase__       )
           return config, inputs_dict






     def      _UpperCamelCase (self     :						Dict       )	->							List[Any]:
           __A ,					__A		 =   self.prepare_config_and_inputs()
           return config, inputs_dict






     def      _UpperCamelCase (self     :						Optional[int]					,					lowercase__     :						Tuple					,					lowercase__     :						Optional[int]					,					lowercase__     :						int       )	->							Optional[int]:
           __A		 =   20
           __A		 =   model_class_name(lowercase__       )

           __A		 =   inputs_dict["input_ids"]
           __A		 =   inputs_dict["attention_mask"]

           __A		 =   model.init_cache(input_ids.shape[0]					,					lowercase__       )
           __A		 =   jnp.ones((input_ids.shape[0], max_length)					,					dtype="i4"       )

           __A		 =   jnp.broadcast_to(
               jnp.arange(input_ids.shape[-1] - 1       )[None, :]					,					(input_ids.shape[0], input_ids.shape[-1] - 1)					,					)
           __A		 =   model(
               input_ids[:, :-1]					,					attention_mask=lowercase__					,					past_key_values=lowercase__					,					position_ids=lowercase__					,					)

           __A		 =   jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]					,					dtype="i4"       )
           __A		 =   model(
               input_ids[:, -1:]					,					attention_mask=lowercase__					,					past_key_values=outputs_cache.past_key_values					,					position_ids=lowercase__					,					)

           __A		 =   model(lowercase__       )

           __A		 =   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      _UpperCamelCase (self     :						Any					,					lowercase__     :						int					,					lowercase__     :						Optional[Any]					,					lowercase__     :						Optional[int]       )	->							Dict:
           __A		 =   20
           __A		 =   model_class_name(lowercase__       )

           __A ,					__A		 =   (
               inputs_dict["input_ids"],
               inputs_dict["attention_mask"],
           )

           __A		 =   jnp.concatenate(
               [
                   attention_mask,
                   jnp.zeros((attention_mask.shape[0], max_length - attention_mask.shape[1])       ),
               ]					,					axis=-1					,					)

           __A		 =   model.init_cache(input_ids.shape[0]					,					lowercase__       )
           __A		 =   jnp.broadcast_to(
               jnp.arange(input_ids.shape[-1] - 1       )[None, :]					,					(input_ids.shape[0], input_ids.shape[-1] - 1)					,					)

           __A		 =   model(
               input_ids[:, :-1]					,					attention_mask=lowercase__					,					past_key_values=lowercase__					,					position_ids=lowercase__					,					)
           __A		 =   jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]					,					dtype="i4"       )
           __A		 =   model(
               input_ids[:, -1:]					,					past_key_values=outputs_cache.past_key_values					,					attention_mask=lowercase__					,					position_ids=lowercase__					,					)

           __A		 =   model(lowercase__					,					attention_mask=lowercase__       )

           __A		 =   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__     (      lowerCamelCase			,    unittest.TestCase			,    lowerCamelCase    ):
     '''simple docstring'''







     _snake_case							   =	(FlaxOPTModel, FlaxOPTForCausalLM) if is_flax_available() else ()
     _snake_case							   =	() if is_flax_available() else ()






     def      _UpperCamelCase (self     :						Tuple       )	->							str:
           __A		 =   FlaxOPTModelTester(self       )






     def      _UpperCamelCase (self     :						List[str]       )	->							Optional[int]:
           __A ,					__A		 =   self.model_tester.prepare_config_and_inputs()
           for model_class in self.all_model_classes:
                 self.model_tester.check_use_cache_forward(lowercase__					,					lowercase__					,					lowercase__       )






     def      _UpperCamelCase (self     :						Union[str, Any]       )	->							Any:
           __A ,					__A		 =   self.model_tester.prepare_config_and_inputs()
           for model_class in self.all_model_classes:
                 self.model_tester.check_use_cache_forward_with_attn_mask(lowercase__					,					lowercase__					,					lowercase__       )






     @slow
     def      _UpperCamelCase (self     :						Any       )	->							Union[str, Any]:
           for model_class_name in self.all_model_classes:
                 __A		 =   model_class_name.from_pretrained("facebook/opt-125m"       )
                 __A		 =   np.ones((1, 1)       ) * model.config.eos_token_id
                 __A		 =   model(lowercase__       )
                 self.assertIsNotNone(lowercase__       )







@require_sentencepiece
@require_flax
class        a__     (      unittest.TestCase    ):
     '''simple docstring'''







     @slow
     def      _UpperCamelCase (self     :						Any       )	->							str:
           __A		 =   FlaxOPTModel.from_pretrained("facebook/opt-350m"       )
           __A		 =   jnp.array([[0, 31414, 232, 328, 740, 1140, 12695, 69, 46078, 1588, 2]]       )
           __A		 =   model(input_ids=lowercase__       ).last_hidden_state
           __A		 =   (1, 11, 512)
           self.assertEqual(output.shape					,					lowercase__       )
           __A		 =   jnp.array(
               [[-0.28_67, -1.92_56, -0.30_62], [-1.27_11, -0.13_37, -0.18_97], [0.41_09, 0.11_87, -1.31_42]]       )
           self.assertTrue(jnp.allclose(output[:, :3, :3]					,					lowercase__					,					atol=4e-2       )       )







@require_flax
@slow
class        a__     (      unittest.TestCase    ):
     '''simple docstring'''







     def      _UpperCamelCase (self     :						Dict       )	->							List[Any]:
           super().setUp()
           __A		 =   "facebook/opt-350m"






     def      _UpperCamelCase (self     :						str       )	->							List[str]:
           __A		 =   FlaxOPTForCausalLM.from_pretrained(self.path_model       )
           __A		 =   GPTaTokenizer.from_pretrained(self.path_model       )

           __A		 =   [
               "Today is a beautiful day and I want to",
               "In the city of",
               "Paris is the capital of France and",
               "Computers and mobile phones have taken",
           ]
           # verify that prompt without BOS token is identical to Metaseq -> add_special_tokens=False
           __A		 =   tokenizer(lowercase__					,					return_tensors="jax"					,					padding=lowercase__					,					add_special_tokens=lowercase__       )
           __A		 =   model(inputs.input_ids					,					attention_mask=inputs.attention_mask       )[0].mean(axis=-1       )
           __A		 =   jnp.array(
               [
                   [1.38_51, -13.89_23, -10.52_29, -10.75_33, -0.23_09, -10.23_84, -0.53_65, -9.09_47, -5.16_70],
                   [-4.70_73, -10.62_76, -3.94_15, -21.52_42, -0.28_22, -0.28_22, -0.28_22, -0.28_22, -0.28_22],
                   [0.62_47, -3.42_29, -8.91_79, -1.42_97, -14.16_50, 1.41_46, -9.02_18, -0.27_03, -0.27_03],
                   [6.47_83, -1.99_13, -10.79_26, -2.33_36, 1.50_92, -0.99_74, -6.82_13, 1.34_77, 1.34_77],
               ]       )
           self.assertTrue(jnp.allclose(lowercase__					,					lowercase__					,					atol=4e-2       )       )

           __A		 =   jax.jit(lowercase__       )
           __A		 =   model(inputs.input_ids					,					attention_mask=inputs.attention_mask       )[0].mean(axis=-1       )
           self.assertTrue(jnp.allclose(lowercase__					,					lowercase__					,					atol=4e-2       )       )







@require_flax
@slow
class        a__     (      unittest.TestCase    ):
     '''simple docstring'''







     @property
     def      _UpperCamelCase (self     :						Any       )	->							Union[str, Any]:
           return [
               "Today is a beautiful day and I want",
               "In the city of",
               "Paris is the capital of France and",
               "Computers and mobile phones have taken",
           ]






     def      _UpperCamelCase (self     :						List[str]       )	->							Tuple:
           __A		 =   "facebook/opt-125m"

           __A		 =   [
               "Today is a beautiful day and I want to",
               "In the city of New York, the city",
               "Paris is the capital of France and the capital",
               "Computers and mobile phones have taken over the",
           ]

           __A		 =   []

           __A		 =   FlaxOPTForCausalLM.from_pretrained(lowercase__       )
           __A		 =   GPTaTokenizer.from_pretrained(lowercase__       )

           for prompt in self.prompts:
                 __A		 =   tokenizer(lowercase__					,					return_tensors="jax"       ).input_ids

                 __A		 =   model.generate(lowercase__					,					max_length=10       )
                 __A		 =   generated_ids[0]

                 __A		 =   tokenizer.batch_decode(lowercase__					,					skip_special_tokens=lowercase__       )
                 predicted_outputs += generated_string

           self.assertListEqual(lowercase__					,					lowercase__       )






     def      _UpperCamelCase (self     :						List[Any]       )	->							Dict:
           __A		 =   "facebook/opt-350m"

           __A		 =   [
               "Today is a beautiful day and I want to",
               "In the city of San Francisco, the city",
               "Paris is the capital of France and the capital",
               "Computers and mobile phones have taken over the",
           ]

           __A		 =   []
           __A		 =   FlaxOPTForCausalLM.from_pretrained(lowercase__       )
           __A		 =   GPTaTokenizer.from_pretrained(lowercase__       )

           for prompt in self.prompts:
                 __A		 =   tokenizer(lowercase__					,					return_tensors="jax"       ).input_ids

                 __A		 =   model.generate(lowercase__					,					max_length=10       )
                 __A		 =   generated_ids[0]

                 __A		 =   tokenizer.batch_decode(lowercase__					,					skip_special_tokens=lowercase__       )
                 predicted_outputs += generated_string

           self.assertListEqual(lowercase__					,					lowercase__       )






     def      _UpperCamelCase (self     :						str       )	->							Optional[int]:
           __A		 =   "facebook/opt-125m"
           __A		 =   [
               "Today is a beautiful day and I want to thank",
               "In the city of Rome Canaver Canaver Canaver Canaver",
           ]
           __A		 =   FlaxOPTForCausalLM.from_pretrained(lowercase__       )
           __A		 =   GPTaTokenizer.from_pretrained(lowercase__       )
           __A		 =   tokenizer(
               [
                   "Today is a beautiful day and I want to",
                   "In the city of",
               ]					,					return_tensors="jax"					,					padding=lowercase__					,					)

           __A		 =   jax.jit(model.generate       )

           __A		 =   jit_generate(inputs["input_ids"]					,					attention_mask=inputs["attention_mask"]       ).sequences

           __A		 =   tokenizer.batch_decode(lowercase__					,					skip_special_tokens=lowercase__       )

           self.assertIsNotNone(lowercase__					,					lowercase__       )






     def      _UpperCamelCase (self     :						Optional[int]       )	->							List[Any]:
           __A		 =   "facebook/opt-350m"

           __A		 =   GPTaTokenizer.from_pretrained(lowercase__       )
           __A		 =   FlaxOPTForCausalLM.from_pretrained(lowercase__       )

           __A		 =   "left"

           # use different length sentences to test batching
           __A		 =   [
               "Hello, my dog is a little",
               "Today, I",
           ]

           __A		 =   tokenizer(lowercase__					,					return_tensors="jax"					,					padding=lowercase__       )
           __A		 =   inputs["input_ids"]

           __A		 =   model.generate(input_ids=lowercase__					,					attention_mask=inputs["attention_mask"]					,					trace=lowercase__       )

           __A		 =   tokenizer(sentences[0]					,					return_tensors="jax"       ).input_ids
           __A		 =   model.generate(input_ids=lowercase__       )

           __A		 =   inputs_non_padded.shape[-1] - inputs["attention_mask"][-1].sum()
           __A		 =   tokenizer(sentences[1]					,					return_tensors="jax"       ).input_ids
           __A		 =   model.generate(input_ids=lowercase__					,					max_length=model.config.max_length - num_paddings       )

           __A		 =   tokenizer.batch_decode(outputs[0]					,					skip_special_tokens=lowercase__       )
           __A		 =   tokenizer.decode(output_non_padded[0][0]					,					skip_special_tokens=lowercase__       )
           __A		 =   tokenizer.decode(output_padded[0][0]					,					skip_special_tokens=lowercase__       )

           __A		 =   [
               "Hello, my dog is a little bit of a dork.\nI'm a little bit",
               "Today, I was in the middle of a conversation with a friend about the",
           ]
           self.assertListEqual(lowercase__					,					lowercase__       )
           self.assertListEqual(lowercase__					,					[non_padded_sentence, padded_sentence]       )










'''simple docstring'''







from typing import List, Optional, Union

import numpy as np
import torch
import torchaudio.compliance.kaldi as ta_kaldi

from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
from ...feature_extraction_utils import BatchFeature
from ...utils import TensorType, logging


__A   :							int     	=						logging.get_logger(__name__)






class        a__     (      lowerCamelCase    ):
     '''simple docstring'''







     _snake_case							   =	['input_values', 'attention_mask']






     def __init__(self     :						Optional[int]					,					lowercase__     :						Tuple=1					,					lowercase__     :						Optional[int]=16000					,					lowercase__     :						Any=128					,					lowercase__     :						Optional[Any]=1024					,					lowercase__     :						int=0.0					,					lowercase__     :						List[str]=True					,					lowercase__     :						Dict=-4.2_67_73_93					,					lowercase__     :						Tuple=4.5_68_99_74					,					lowercase__     :						Optional[int]=False					,					**lowercase__     :						Optional[int]					,					)	->							Tuple:
           super().__init__(feature_size=lowercase__					,					sampling_rate=lowercase__					,					padding_value=lowercase__					,					**lowercase__       )
           __A		 =   num_mel_bins
           __A		 =   max_length
           __A		 =   do_normalize
           __A		 =   mean
           __A		 =   std
           __A		 =   return_attention_mask






     def      _UpperCamelCase (self     :						List[Any]					,					lowercase__     :						np.ndarray					,					lowercase__     :						int					,					)	->							np.ndarray:
           __A		 =   torch.from_numpy(lowercase__       ).unsqueeze(0       )
           __A		 =   ta_kaldi.fbank(
               lowercase__					,					htk_compat=lowercase__					,					sample_frequency=self.sampling_rate					,					use_energy=lowercase__					,					window_type="hanning"					,					num_mel_bins=self.num_mel_bins					,					dither=0.0					,					frame_shift=10					,					)

           __A		 =   fbank.shape[0]
           __A		 =   max_length - n_frames

           # pad or truncate, depending on difference
           if difference > 0:
                 __A		 =   torch.nn.ZeroPadad((0, 0, 0, difference)       )
                 __A		 =   pad_module(lowercase__       )
           elif difference < 0:
                 __A		 =   fbank[0:max_length, :]

           __A		 =   fbank.numpy()

           return fbank






     def      _UpperCamelCase (self     :						int					,					lowercase__     :						np.ndarray       )	->							np.ndarray:
           return (input_values - (self.mean)) / (self.std * 2)






     def __call__(self     :						str					,					lowercase__     :						Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]]					,					lowercase__     :						Optional[int] = None					,					lowercase__     :						Optional[Union[str, TensorType]] = None					,					**lowercase__     :						List[Any]					,					)	->							BatchFeature:

           if sampling_rate is not None:
                 if sampling_rate != self.sampling_rate:
                       raise ValueError(
                           f"The model corresponding to this feature extractor: {self} was trained using a sampling rate of"
                           f" {self.sampling_rate}. Please make sure that the provided `raw_speech` input was sampled with"
                           f" {self.sampling_rate} and not {sampling_rate}."       )
           else:
                 logger.warning(
                     "It is strongly recommended to pass the `sampling_rate` argument to this function. "
                     "Failing to do so can result in silent errors that might be hard to debug."       )

           __A		 =   isinstance(lowercase__					,					np.ndarray       ) and len(raw_speech.shape       ) > 1
           if is_batched_numpy and len(raw_speech.shape       ) > 2:
                 raise ValueError(f"Only mono-channel audio is supported for input to {self}"       )
           __A		 =   is_batched_numpy or (
               isinstance(lowercase__					,					(list, tuple)       ) and (isinstance(raw_speech[0]					,					(np.ndarray, tuple, list)       ))
           )

           if is_batched:
                 __A		 =   [np.asarray(lowercase__					,					dtype=np.floataa       ) for speech in raw_speech]
           elif not is_batched and not isinstance(lowercase__					,					np.ndarray       ):
                 __A		 =   np.asarray(lowercase__					,					dtype=np.floataa       )
           elif isinstance(lowercase__					,					np.ndarray       ) and raw_speech.dtype is np.dtype(np.floataa       ):
                 __A		 =   raw_speech.astype(np.floataa       )

           # always return batch
           if not is_batched:
                 __A		 =   [raw_speech]

           # extract fbank features and pad/truncate to max_length
           __A		 =   [self._extract_fbank_features(lowercase__					,					max_length=self.max_length       ) for waveform in raw_speech]

           # convert into BatchFeature
           __A		 =   BatchFeature({"input_values": features}       )

           # make sure list is in array format
           __A		 =   padded_inputs.get("input_values"       )
           if isinstance(input_values[0]					,					lowercase__       ):
                 __A		 =   [np.asarray(lowercase__					,					dtype=np.floataa       ) for feature in input_values]

           # normalization
           if self.do_normalize:
                 __A		 =   [self.normalize(lowercase__       ) for feature in input_values]

           if return_tensors is not None:
                 __A		 =   padded_inputs.convert_to_tensors(lowercase__       )

           return padded_inputs


