
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 			lowercase_		:


       def __init__(    self						,       __A						,       __A=14						,       __A=7						,       __A=True						,       __A=True						,       __A=False						,       __A=True						,       __A=99						,       __A=32						,       __A=4						,       __A=4						,       __A=4						,       __A=37						,       __A="gelu"						,       __A=0.1						,       __A=0.1						,       __A=512						,       __A=0.02						,       )						->    Optional[Any]:
              SCREAMING_SNAKE_CASE_							:				Tuple		=parent
              SCREAMING_SNAKE_CASE_							:				List[str]		=batch_size
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=seq_length
              SCREAMING_SNAKE_CASE_							:				int		=is_training
              SCREAMING_SNAKE_CASE_							:				List[str]		=use_input_mask
              SCREAMING_SNAKE_CASE_							:				Any		=use_token_type_ids
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=use_labels
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=vocab_size
              SCREAMING_SNAKE_CASE_							:				Any		=hidden_size
              SCREAMING_SNAKE_CASE_							:				int		=rotary_dim
              SCREAMING_SNAKE_CASE_							:				Dict		=num_hidden_layers
              SCREAMING_SNAKE_CASE_							:				List[str]		=num_attention_heads
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=intermediate_size
              SCREAMING_SNAKE_CASE_							:				Dict		=hidden_act
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=hidden_dropout_prob
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=attention_probs_dropout_prob
              SCREAMING_SNAKE_CASE_							:				str		=max_position_embeddings
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=initializer_range
              SCREAMING_SNAKE_CASE_							:				List[Any]		=None
              SCREAMING_SNAKE_CASE_							:				str		=vocab_size - 1
              SCREAMING_SNAKE_CASE_							:				Dict		=vocab_size - 1
              SCREAMING_SNAKE_CASE_							:				List[str]		=vocab_size - 1


       def 	_snake_case				(    self				)						->    List[Any]:
              SCREAMING_SNAKE_CASE_							:				str		=ids_tensor([self.batch_size, self.seq_length]						,       self.vocab_size				)

              SCREAMING_SNAKE_CASE_							:				List[Any]		=None
              if self.use_input_mask:
                     SCREAMING_SNAKE_CASE_							:				Dict		=random_attention_mask([self.batch_size, self.seq_length]				)

              SCREAMING_SNAKE_CASE_							:				str		=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=__A						,       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 	_snake_case				(    self				)						->    List[str]:
              SCREAMING_SNAKE_CASE_							:				Any		=self.prepare_config_and_inputs()
              SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_							:				int		=config_and_inputs
              SCREAMING_SNAKE_CASE_							:				Dict		={'''input_ids''': input_ids, '''attention_mask''': attention_mask}
              return config, inputs_dict


       def 	_snake_case				(    self						,       __A						,       __A						,       __A						,       __A				)						->    List[Any]:
              SCREAMING_SNAKE_CASE_							:				Tuple		=20
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=model_class_name(__A				)

              SCREAMING_SNAKE_CASE_							:				List[Any]		=model.init_cache(input_ids.shape[0]						,       __A				)
              SCREAMING_SNAKE_CASE_							:				Any		=jnp.ones((input_ids.shape[0], max_decoder_length)						,       dtype='''i4'''				)

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

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

              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=model(__A				)

              SCREAMING_SNAKE_CASE_							:				Optional[int]		=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 	_snake_case				(    self						,       __A						,       __A						,       __A						,       __A				)						->    Optional[Any]:
              SCREAMING_SNAKE_CASE_							:				str		=20
              SCREAMING_SNAKE_CASE_							:				List[str]		=model_class_name(__A				)

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

              SCREAMING_SNAKE_CASE_							:				List[Any]		=model.init_cache(input_ids.shape[0]						,       __A				)
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=jnp.broadcast_to(
                  jnp.arange(input_ids.shape[-1] - 1				)[None, :]						,       (input_ids.shape[0], input_ids.shape[-1] - 1)				)

              SCREAMING_SNAKE_CASE_							:				str		=model(
                  input_ids[:, :-1]						,       attention_mask=__A						,       past_key_values=__A						,       position_ids=__A						,       )
              SCREAMING_SNAKE_CASE_							:				Dict		=jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]						,       dtype='''i4'''				)
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=model(
                  input_ids[:, -1:]						,       past_key_values=outputs_cache.past_key_values						,       attention_mask=__A						,       position_ids=__A						,       )

              SCREAMING_SNAKE_CASE_							:				List[str]		=model(__A						,       attention_mask=__A				)

              SCREAMING_SNAKE_CASE_							:				Union[str, 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 			lowercase_		(      A  ,      A  ,      unittest.TestCase       ):
       __lowerCamelCase				       =						(FlaxGPTJModel, FlaxGPTJForCausalLM) if is_flax_available() else ()
       __lowerCamelCase				       =						(FlaxGPTJForCausalLM,) if is_flax_available() else ()


       def 	_snake_case				(    self				)						->    Tuple:
              SCREAMING_SNAKE_CASE_							:				List[str]		=FlaxGPTJModelTester(self				)


       def 	_snake_case				(    self				)						->    List[str]:
              for model_class_name in self.all_model_classes:
                     SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_							:				Dict		=self.model_tester.prepare_config_and_inputs()
                     self.model_tester.check_use_cache_forward(__A						,       __A						,       __A						,       __A				)


       def 	_snake_case				(    self				)						->    Optional[int]:
              for model_class_name in self.all_model_classes:
                     SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_							:				Optional[Any]		=self.model_tester.prepare_config_and_inputs()
                     self.model_tester.check_use_cache_forward_with_attn_mask(
                         __A						,       __A						,       __A						,       __A				)


       @tooslow
       def 	_snake_case				(    self				)						->    List[Any]:
              SCREAMING_SNAKE_CASE_							:				List[Any]		=GPTaTokenizer.from_pretrained('''gpt2'''						,       pad_token='''<|endoftext|>'''						,       padding_side='''left'''				)
              SCREAMING_SNAKE_CASE_							:				Tuple		=tokenizer(['''Hello this is a long string''', '''Hey''']						,       return_tensors='''np'''						,       padding=__A						,       truncation=__A				)

              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=FlaxGPTJForCausalLM.from_pretrained('''EleutherAI/gpt-j-6B'''				)
              SCREAMING_SNAKE_CASE_							:				Dict		=False
              SCREAMING_SNAKE_CASE_							:				Tuple		=model.config.eos_token_id

              SCREAMING_SNAKE_CASE_							:				Optional[int]		=jax.jit(model.generate				)

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

              SCREAMING_SNAKE_CASE_							:				Tuple		=tokenizer.batch_decode(__A						,       skip_special_tokens=__A				)

              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=[
                  '''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(__A						,       __A				)


       @is_pt_flax_cross_test
       def 	_snake_case				(    self				)						->    List[Any]:
              SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_							:				List[str]		=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
                            SCREAMING_SNAKE_CASE_							:				List[Any]		=self._prepare_for_class(__A						,       __A				)
                            SCREAMING_SNAKE_CASE_							:				str		={k: torch.tensor(v.tolist()				) for k, v in prepared_inputs_dict.items()}

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

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

                            SCREAMING_SNAKE_CASE_							:				Optional[Any]		=convert_pytorch_state_dict_to_flax(pt_model.state_dict()						,       __A				)
                            SCREAMING_SNAKE_CASE_							:				int		=fx_state

                            with torch.no_grad():
                                   SCREAMING_SNAKE_CASE_							:				Optional[Any]		=pt_model(**__A				).to_tuple()

                            SCREAMING_SNAKE_CASE_							:				Optional[Any]		=fx_model(**__A				).to_tuple()
                            self.assertEqual(len(__A				)						,       len(__A				)						,       '''Output lengths differ between Flax and PyTorch'''				)
                            for fx_output, pt_output in zip(__A						,       __A				):
                                   self.assert_almost_equals(fx_output[:, -1]						,       pt_output[:, -1].numpy()						,       4e-2				)

                            with tempfile.TemporaryDirectory() as tmpdirname:
                                   pt_model.save_pretrained(__A				)
                                   SCREAMING_SNAKE_CASE_							:				Any		=model_class.from_pretrained(__A						,       from_pt=__A				)

                            SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=fx_model_loaded(**__A				).to_tuple()
                            self.assertEqual(
                                len(__A				)						,       len(__A				)						,       '''Output lengths differ between Flax and PyTorch'''				)
                            for fx_output_loaded, pt_output in zip(__A						,       __A				):
                                   self.assert_almost_equals(fx_output_loaded[:, -1]						,       pt_output[:, -1].numpy()						,       4e-2				)


       @is_pt_flax_cross_test
       def 	_snake_case				(    self				)						->    Optional[Any]:
              SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_							:				int		=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
                            SCREAMING_SNAKE_CASE_							:				str		=self._prepare_for_class(__A						,       __A				)
                            SCREAMING_SNAKE_CASE_							:				Dict		={k: torch.tensor(v.tolist()				) for k, v in prepared_inputs_dict.items()}

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

                            SCREAMING_SNAKE_CASE_							:				List[str]		=pt_model_class(__A				).eval()
                            SCREAMING_SNAKE_CASE_							:				Optional[Any]		=model_class(__A						,       dtype=jnp.floataa				)

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

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

                            with torch.no_grad():
                                   SCREAMING_SNAKE_CASE_							:				List[Any]		=pt_model(**__A				).to_tuple()

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

                            with tempfile.TemporaryDirectory() as tmpdirname:
                                   fx_model.save_pretrained(__A				)
                                   SCREAMING_SNAKE_CASE_							:				Dict		=pt_model_class.from_pretrained(__A						,       from_flax=__A				)

                            with torch.no_grad():
                                   SCREAMING_SNAKE_CASE_							:				Dict		=pt_model_loaded(**__A				).to_tuple()

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





       @tooslow
       def 	_snake_case				(    self				)						->    Tuple:
              for model_class_name in self.all_model_classes:
                     SCREAMING_SNAKE_CASE_							:				Optional[Any]		=model_class_name.from_pretrained('''EleutherAI/gpt-j-6B'''				)
                     SCREAMING_SNAKE_CASE_							:				Dict		=model(np.ones((1, 1)				)				)
                     self.assertIsNotNone(__A				)


from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


_lowercase											=    {
    """configuration_falcon""": ["""FALCON_PRETRAINED_CONFIG_ARCHIVE_MAP""", """FalconConfig"""],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _lowercase											=    [
     """FALCON_PRETRAINED_MODEL_ARCHIVE_LIST""",
     """FalconForCausalLM""",
     """FalconModel""",
     """FalconPreTrainedModel""",
     """FalconForSequenceClassification""",
     """FalconForTokenClassification""",
     """FalconForQuestionAnswering""",
 ]


if TYPE_CHECKING:
 from .configuration_falcon import FALCON_PRETRAINED_CONFIG_ARCHIVE_MAP, FalconConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_falcon import (
      FALCON_PRETRAINED_MODEL_ARCHIVE_LIST,
      FalconForCausalLM,
      FalconForQuestionAnswering,
      FalconForSequenceClassification,
      FalconForTokenClassification,
      FalconModel,
      FalconPreTrainedModel,
  )


else:
 import sys

 _lowercase											=    _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
