



import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTNeoConfig, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, slow

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.gpt_neo.modeling_flax_gpt_neo import FlaxGPTNeoForCausalLM, FlaxGPTNeoModel

if is_torch_available():
      import torch

class    _SCREAMING_SNAKE_CASE						:





     def __init__(							self			:				Optional[Any]  ,	UpperCamelCase			:				Union[str, Any]  ,	UpperCamelCase			:				str=1_4  ,	UpperCamelCase			:				str=7  ,	UpperCamelCase			:				str=True  ,	UpperCamelCase			:				Union[str, Any]=True  ,	UpperCamelCase			:				Dict=False  ,	UpperCamelCase			:				Optional[Any]=True  ,	UpperCamelCase			:				Any=9_9  ,	UpperCamelCase			:				Optional[Any]=3_2  ,	UpperCamelCase			:				Union[str, Any]=4  ,	UpperCamelCase			:				List[Any]=4  ,	UpperCamelCase			:				Tuple=[[["global", "local"], 2]]  ,	UpperCamelCase			:				str=3_7  ,	UpperCamelCase			:				int="gelu"  ,	UpperCamelCase			:				Union[str, Any]=0.1  ,	UpperCamelCase			:				Union[str, Any]=0.1  ,	UpperCamelCase			:				Dict=5_1_2  ,	UpperCamelCase			:				Union[str, Any]=7  ,	UpperCamelCase			:				Optional[Any]=0.0_2  ,	)->Union[str, Any]:
      __SCREAMING_SNAKE_CASE							:     Dict							=						parent
      __SCREAMING_SNAKE_CASE							:     Dict							=						batch_size
      __SCREAMING_SNAKE_CASE							:     Optional[int]							=						seq_length
      __SCREAMING_SNAKE_CASE							:     List[Any]							=						is_training
      __SCREAMING_SNAKE_CASE							:     int							=						use_input_mask
      __SCREAMING_SNAKE_CASE							:     int							=						use_token_type_ids
      __SCREAMING_SNAKE_CASE							:     Any							=						use_labels
      __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						vocab_size
      __SCREAMING_SNAKE_CASE							:     List[str]							=						hidden_size
      __SCREAMING_SNAKE_CASE							:     int							=						num_hidden_layers
      __SCREAMING_SNAKE_CASE							:     Optional[int]							=						num_attention_heads
      __SCREAMING_SNAKE_CASE							:     str							=						attention_types
      __SCREAMING_SNAKE_CASE							:     Dict							=						intermediate_size
      __SCREAMING_SNAKE_CASE							:     Tuple							=						hidden_act
      __SCREAMING_SNAKE_CASE							:     Tuple							=						hidden_dropout_prob
      __SCREAMING_SNAKE_CASE							:     Dict							=						attention_probs_dropout_prob
      __SCREAMING_SNAKE_CASE							:     Any							=						max_position_embeddings
      __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						window_size
      __SCREAMING_SNAKE_CASE							:     int							=						initializer_range
      __SCREAMING_SNAKE_CASE							:     List[str]							=						None
      __SCREAMING_SNAKE_CASE							:     Tuple							=						vocab_size - 1
      __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						vocab_size - 1
      __SCREAMING_SNAKE_CASE							:     str							=						vocab_size - 1





     def 						__snake_case     (							self			:				Dict					)->Optional[Any]:
      __SCREAMING_SNAKE_CASE							:     str							=						ids_tensor([self.batch_size, self.seq_length]  ,	self.vocab_size					)

      __SCREAMING_SNAKE_CASE							:     List[str]							=						None
      if self.use_input_mask:
       __SCREAMING_SNAKE_CASE							:     List[str]							=						random_attention_mask([self.batch_size, self.seq_length]					)

      __SCREAMING_SNAKE_CASE							:     Tuple							=						GPTNeoConfig(
          vocab_size=self.vocab_size  ,	hidden_size=self.hidden_size  ,	num_layers=self.num_hidden_layers  ,	num_heads=self.num_attention_heads  ,	max_position_embeddings=self.max_position_embeddings  ,	use_cache=UpperCamelCase  ,	bos_token_id=self.bos_token_id  ,	eos_token_id=self.eos_token_id  ,	pad_token_id=self.pad_token_id  ,	window_size=self.window_size  ,	attention_types=self.attention_types  ,	)

      return (config, input_ids, input_mask)





     def 						__snake_case     (							self			:				List[str]					)->Union[str, Any]:
      __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						self.prepare_config_and_inputs()
      __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE							:     List[Any]							=						config_and_inputs
      __SCREAMING_SNAKE_CASE							:     int							=						{"input_ids": input_ids, "attention_mask": attention_mask}
      return config, inputs_dict





     def 						__snake_case     (							self			:				Optional[int]  ,	UpperCamelCase			:				List[str]  ,	UpperCamelCase			:				int  ,	UpperCamelCase			:				str  ,	UpperCamelCase			:				Tuple					)->int:
      __SCREAMING_SNAKE_CASE							:     List[Any]							=						2_0
      __SCREAMING_SNAKE_CASE							:     Tuple							=						model_class_name(UpperCamelCase					)

      __SCREAMING_SNAKE_CASE							:     List[Any]							=						model.init_cache(input_ids.shape[0]  ,	UpperCamelCase					)
      __SCREAMING_SNAKE_CASE							:     List[Any]							=						jnp.ones((input_ids.shape[0], max_decoder_length)  ,	dtype="i4"					)

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

      __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]  ,	dtype="i4"					)
      __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						model(
          input_ids[:, -1:]  ,	attention_mask=UpperCamelCase  ,	past_key_values=outputs_cache.past_key_values  ,	position_ids=UpperCamelCase  ,	)

      __SCREAMING_SNAKE_CASE							:     Dict							=						model(UpperCamelCase					)

      __SCREAMING_SNAKE_CASE							:     str							=						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			:				str  ,	UpperCamelCase			:				Optional[int]  ,	UpperCamelCase			:				int  ,	UpperCamelCase			:				Dict  ,	UpperCamelCase			:				Optional[Any]					)->List[str]:
      __SCREAMING_SNAKE_CASE							:     Tuple							=						2_0
      __SCREAMING_SNAKE_CASE							:     Any							=						model_class_name(UpperCamelCase					)

      __SCREAMING_SNAKE_CASE							:     List[Any]							=						jnp.concatenate(
          [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])					)]  ,	axis=-1  ,	)

      __SCREAMING_SNAKE_CASE							:     List[str]							=						model.init_cache(input_ids.shape[0]  ,	UpperCamelCase					)
      __SCREAMING_SNAKE_CASE							:     str							=						jnp.broadcast_to(
          jnp.arange(input_ids.shape[-1] - 1					)[None, :]  ,	(input_ids.shape[0], input_ids.shape[-1] - 1)					)

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

      __SCREAMING_SNAKE_CASE							:     Any							=						model(UpperCamelCase  ,	attention_mask=UpperCamelCase					)

      __SCREAMING_SNAKE_CASE							:     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    _SCREAMING_SNAKE_CASE						(UpperCamelCase      ,     UpperCamelCase      ,     unittest.TestCase			):
     lowerCAmelCase       =      (FlaxGPTNeoModel, FlaxGPTNeoForCausalLM) if is_flax_available() else ()
     lowerCAmelCase       =      (FlaxGPTNeoForCausalLM,) if is_flax_available() else ()





     def 						__snake_case     (							self			:				Any					)->List[Any]:
      __SCREAMING_SNAKE_CASE							:     Any							=						FlaxGPTNeoModelTester(self					)





     def 						__snake_case     (							self			:				List[str]					)->Optional[int]:
      for model_class_name in self.all_model_classes:
       __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE							:     List[str]							=						self.model_tester.prepare_config_and_inputs()
       self.model_tester.check_use_cache_forward(UpperCamelCase  ,	UpperCamelCase  ,	UpperCamelCase  ,	UpperCamelCase					)





     def 						__snake_case     (							self			:				List[Any]					)->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(
           UpperCamelCase  ,	UpperCamelCase  ,	UpperCamelCase  ,	UpperCamelCase					)





     @slow
     def 						__snake_case     (							self			:				str					)->Optional[int]:
      __SCREAMING_SNAKE_CASE							:     str							=						GPTaTokenizer.from_pretrained("gpt2"  ,	pad_token="<|endoftext|>"  ,	padding_side="left"					)
      __SCREAMING_SNAKE_CASE							:     str							=						tokenizer(["Hello this is a long string", "Hey"]  ,	return_tensors="np"  ,	padding=UpperCamelCase  ,	truncation=UpperCamelCase					)

      __SCREAMING_SNAKE_CASE							:     str							=						FlaxGPTNeoForCausalLM.from_pretrained("EleutherAI/gpt-neo-125M"					)
      __SCREAMING_SNAKE_CASE							:     Tuple							=						False
      __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						model.config.eos_token_id

      __SCREAMING_SNAKE_CASE							:     Tuple							=						jax.jit(model.generate					)

      __SCREAMING_SNAKE_CASE							:     List[Any]							=						jit_generate(
          inputs["input_ids"]  ,	attention_mask=inputs["attention_mask"]  ,	pad_token_id=tokenizer.pad_token_id					).sequences

      __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						tokenizer.batch_decode(UpperCamelCase  ,	skip_special_tokens=UpperCamelCase					)

      __SCREAMING_SNAKE_CASE							:     List[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(UpperCamelCase  ,	UpperCamelCase					)





     @is_pt_flax_cross_test
     def 						__snake_case     (							self			:				Optional[int]					)->Optional[int]:
      __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							:     str							=						self._prepare_for_class(UpperCamelCase  ,	UpperCamelCase					)
        __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						{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							:     Union[str, Any]							=						getattr(UpperCamelCase  ,	UpperCamelCase					)

        __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE							:     Tuple							=						pt_inputs["input_ids"].shape
        __SCREAMING_SNAKE_CASE							:     Optional[int]							=						np.random.randint(0  ,	seq_length - 1  ,	size=(batch_size,)					)
        for batch_idx, start_index in enumerate(UpperCamelCase					):
         __SCREAMING_SNAKE_CASE							:     int							=						0
         __SCREAMING_SNAKE_CASE							:     str							=						1
         __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						0
         __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						1
        __SCREAMING_SNAKE_CASE							:     int							=						pt_model_class(UpperCamelCase					).eval()
        __SCREAMING_SNAKE_CASE							:     int							=						model_class(UpperCamelCase  ,	dtype=jnp.floataa					)

        __SCREAMING_SNAKE_CASE							:     List[str]							=						convert_pytorch_state_dict_to_flax(pt_model.state_dict()  ,	UpperCamelCase					)
        __SCREAMING_SNAKE_CASE							:     List[str]							=						fx_state

        with torch.no_grad():
         __SCREAMING_SNAKE_CASE							:     int							=						pt_model(**UpperCamelCase					).to_tuple()

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

        with tempfile.TemporaryDirectory() as tmpdirname:
         pt_model.save_pretrained(UpperCamelCase					)
         __SCREAMING_SNAKE_CASE							:     List[Any]							=						model_class.from_pretrained(UpperCamelCase  ,	from_pt=UpperCamelCase					)

        __SCREAMING_SNAKE_CASE							:     List[Any]							=						fx_model_loaded(**UpperCamelCase					).to_tuple()
        self.assertEqual(
            len(UpperCamelCase					)  ,	len(UpperCamelCase					)  ,	"Output lengths differ between Flax and PyTorch"					)
        for fx_output_loaded, pt_output in zip(UpperCamelCase  ,	UpperCamelCase					):
         self.assert_almost_equals(fx_output_loaded[:, -1]  ,	pt_output[:, -1].numpy()  ,	4E-2					)





     @is_pt_flax_cross_test
     def 						__snake_case     (							self			:				str					)->Dict:
      __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[str]							=						self._prepare_for_class(UpperCamelCase  ,	UpperCamelCase					)
        __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						{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							:     str							=						getattr(UpperCamelCase  ,	UpperCamelCase					)

        __SCREAMING_SNAKE_CASE							:     Any							=						pt_model_class(UpperCamelCase					).eval()
        __SCREAMING_SNAKE_CASE							:     List[str]							=						model_class(UpperCamelCase  ,	dtype=jnp.floataa					)

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

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

        with torch.no_grad():
         __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						pt_model(**UpperCamelCase					).to_tuple()

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

        with tempfile.TemporaryDirectory() as tmpdirname:
         fx_model.save_pretrained(UpperCamelCase					)
         __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						pt_model_class.from_pretrained(UpperCamelCase  ,	from_flax=UpperCamelCase					)

        with torch.no_grad():
         __SCREAMING_SNAKE_CASE							:     Dict							=						pt_model_loaded(**UpperCamelCase					).to_tuple()

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






     @slow
     def 						__snake_case     (							self			:				Dict					)->Any:
      for model_class_name in self.all_model_classes:
       __SCREAMING_SNAKE_CASE							:     int							=						model_class_name.from_pretrained("EleutherAI/gpt-neo-125M"					)
       __SCREAMING_SNAKE_CASE							:     List[str]							=						model(np.ones((1, 1)					)					)
       self.assertIsNotNone(UpperCamelCase					)








from __future__ import annotations

from collections.abc import Iterable
from dataclasses import dataclass
from typing import Any

class 		__magic_name__				:


  """simple docstring"""





  _UpperCamelCase						 = ("_front", "_back", "_len")

  @dataclass
  class 		__magic_name__				:


    """simple docstring"""





    _UpperCamelCase						 = None
    _UpperCamelCase						 = None
    _UpperCamelCase						 = None







  class 		__magic_name__				:


    """simple docstring"""





    _UpperCamelCase						 = ("_cur",)


    def __init__(					self	,		a__		):
      _lowerCamelCase      =						cur


    def __iter__(					self		):
      return self


    def     _UpperCAmelCase  (					self		):
      if self._cur is None:
        # finished iterating
        raise StopIteration
      _lowerCamelCase      =						self._cur.val
      _lowerCamelCase      =						self._cur.next_node

      return val







  def __init__(					self	,		a__ = None		):
    _lowerCamelCase      =						None
    _lowerCamelCase      =						None
    _lowerCamelCase      =						0

    if iterable is not None:
      # append every value to the deque
      for val in iterable:
        self.append(a__		)


  def     _UpperCAmelCase  (					self	,		a__		):
    _lowerCamelCase      =						self._Node(a__	,		a__	,		a__		)
    if self.is_empty():
      # front = back
      _lowerCamelCase      =						_lowerCamelCase      =						node
      _lowerCamelCase      =						1
    else:
      # connect nodes
      _lowerCamelCase      =						node
      _lowerCamelCase      =						self._back
      _lowerCamelCase      =						node  # assign new back to the new node

      self._len += 1

      # make sure there were no errors
      assert not self.is_empty(), "Error on appending value."


  def     _UpperCAmelCase  (					self	,		a__		):
    _lowerCamelCase      =						self._Node(a__	,		a__	,		a__		)
    if self.is_empty():
      # front = back
      _lowerCamelCase      =						_lowerCamelCase      =						node
      _lowerCamelCase      =						1
    else:
      # connect nodes
      _lowerCamelCase      =						self._front
      _lowerCamelCase      =						node
      _lowerCamelCase      =						node  # assign new front to the new node

      self._len += 1

      # make sure there were no errors
      assert not self.is_empty(), "Error on appending value."


  def     _UpperCAmelCase  (					self	,		a__		):
    for val in iterable:
      self.append(a__		)


  def     _UpperCAmelCase  (					self	,		a__		):
    for val in iterable:
      self.appendleft(a__		)


  def     _UpperCAmelCase  (					self		):
    assert not self.is_empty(), "Deque is empty."

    _lowerCamelCase      =						self._back
    _lowerCamelCase      =						self._back.prev_node  # set new back
    # drop the last node - python will deallocate memory automatically
    _lowerCamelCase      =						None

    self._len -= 1

    return topop.val


  def     _UpperCAmelCase  (					self		):
    assert not self.is_empty(), "Deque is empty."

    _lowerCamelCase      =						self._front
    _lowerCamelCase      =						self._front.next_node  # set new front and drop the first node
    _lowerCamelCase      =						None

    self._len -= 1

    return topop.val


  def     _UpperCAmelCase  (					self		):
    return self._front is None


  def __len__(					self		):
    return self._len


  def __eq__(					self	,		a__		):

    if not isinstance(a__	,		a__		):
      return NotImplemented

    _lowerCamelCase      =						self._front
    _lowerCamelCase      =						other._front

    # if the length of the dequeues are not the same, they are not equal
    if len(self		) != len(a__		):
      return False

    while me is not None and oth is not None:
      # compare every value
      if me.val != oth.val:
        return False
      _lowerCamelCase      =						me.next_node
      _lowerCamelCase      =						oth.next_node

    return True


  def __iter__(					self		):
    return Deque._Iterator(self._front		)


  def __repr__(					self		):
    _lowerCamelCase      =						[]
    _lowerCamelCase      =						self._front
    while aux is not None:
      # append the values in a list to display
      values_list.append(aux.val		)
      _lowerCamelCase      =						aux.next_node

    return f'''[{', '.join(repr(a__		) for val in values_list		)}]'''


if __name__ == "__main__":
  import doctest

  doctest.testmod()




