





'''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 		UpperCAmelCase :


  """simple docstring"""







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


  def 	UpperCamelCase__   (					self     :       Dict	)  ->    Any:
    _UpperCamelCase							      =ids_tensor([self.batch_size, self.seq_length]				,						self.vocab_size	)

    _UpperCamelCase							      =None
    if self.use_input_mask:
      _UpperCamelCase							      =random_attention_mask([self.batch_size, self.seq_length]	)

    _UpperCamelCase							      =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=UpperCamelCase__				,						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 	UpperCamelCase__   (					self     :       Optional[Any]	)  ->    Optional[Any]:
    _UpperCamelCase							      =self.prepare_config_and_inputs()
    _UpperCamelCase      ,  _UpperCamelCase      ,  _UpperCamelCase							      =config_and_inputs
    _UpperCamelCase							      ={'''input_ids''': input_ids, '''attention_mask''': attention_mask}
    return config, inputs_dict


  def 	UpperCamelCase__   (					self     :       str				,						UpperCamelCase__     :       str				,						UpperCamelCase__     :       int				,						UpperCamelCase__     :       List[str]				,						UpperCamelCase__     :       Tuple	)  ->    Dict:
    _UpperCamelCase							      =20
    _UpperCamelCase							      =model_class_name(UpperCamelCase__	)

    _UpperCamelCase							      =model.init_cache(input_ids.shape[0]				,						UpperCamelCase__	)
    _UpperCamelCase							      =jnp.ones((input_ids.shape[0], max_decoder_length)				,						dtype='''i4'''	)

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

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

    _UpperCamelCase							      =model(UpperCamelCase__	)

    _UpperCamelCase							      =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     :       Union[str, Any]				,						UpperCamelCase__     :       Dict				,						UpperCamelCase__     :       List[str]				,						UpperCamelCase__     :       str				,						UpperCamelCase__     :       Tuple	)  ->    List[str]:
    _UpperCamelCase							      =20
    _UpperCamelCase							      =model_class_name(UpperCamelCase__	)

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

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

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

    _UpperCamelCase							      =model(UpperCamelCase__				,						attention_mask=UpperCamelCase__	)

    _UpperCamelCase							      =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 		UpperCAmelCase (					lowercase_ ,  lowercase_ ,  unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_											=					(FlaxGPTJModel, FlaxGPTJForCausalLM) if is_flax_available() else ()
  lowerCAmelCase_											=					(FlaxGPTJForCausalLM,) if is_flax_available() else ()


  def 	UpperCamelCase__   (					self     :       Any	)  ->    List[Any]:
    _UpperCamelCase							      =FlaxGPTJModelTester(self	)


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






    for model_class_name in self.all_model_classes:
      _UpperCamelCase      ,  _UpperCamelCase      ,  _UpperCamelCase							      =self.model_tester.prepare_config_and_inputs()
      self.model_tester.check_use_cache_forward(UpperCamelCase__				,						UpperCamelCase__				,						UpperCamelCase__				,						UpperCamelCase__	)


  def 	UpperCamelCase__   (					self     :       Dict	)  ->    List[str]:






    for model_class_name in self.all_model_classes:
      _UpperCamelCase      ,  _UpperCamelCase      ,  _UpperCamelCase							      =self.model_tester.prepare_config_and_inputs()
      self.model_tester.check_use_cache_forward_with_attn_mask(
          UpperCamelCase__				,						UpperCamelCase__				,						UpperCamelCase__				,						UpperCamelCase__	)


  @tooslow
  def 	UpperCamelCase__   (					self     :       int	)  ->    str:
    _UpperCamelCase							      =GPTaTokenizer.from_pretrained('''gpt2'''				,						pad_token='''<|endoftext|>'''				,						padding_side='''left'''	)
    _UpperCamelCase							      =tokenizer(['''Hello this is a long string''', '''Hey''']				,						return_tensors='''np'''				,						padding=UpperCamelCase__				,						truncation=UpperCamelCase__	)

    _UpperCamelCase							      =FlaxGPTJForCausalLM.from_pretrained('''EleutherAI/gpt-j-6B'''	)
    _UpperCamelCase							      =False
    _UpperCamelCase							      =model.config.eos_token_id

    _UpperCamelCase							      =jax.jit(model.generate	)

    _UpperCamelCase							      =jit_generate(
        inputs['''input_ids''']				,						attention_mask=inputs['''attention_mask''']				,						pad_token_id=tokenizer.pad_token_id	).sequences

    _UpperCamelCase							      =tokenizer.batch_decode(UpperCamelCase__				,						skip_special_tokens=UpperCamelCase__	)

    _UpperCamelCase							      =[
        '''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 	UpperCamelCase__   (					self     :       Dict	)  ->    List[Any]:
    _UpperCamelCase      ,  _UpperCamelCase							      =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							      =self._prepare_for_class(UpperCamelCase__				,						UpperCamelCase__	)
        _UpperCamelCase							      ={k: torch.tensor(v.tolist()	) for k, v in prepared_inputs_dict.items()}

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

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

        _UpperCamelCase							      =convert_pytorch_state_dict_to_flax(pt_model.state_dict()				,						UpperCamelCase__	)
        _UpperCamelCase							      =fx_state

        with torch.no_grad():
          _UpperCamelCase							      =pt_model(**UpperCamelCase__	).to_tuple()

        _UpperCamelCase							      =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__	)
          _UpperCamelCase							      =model_class.from_pretrained(UpperCamelCase__				,						from_pt=UpperCamelCase__	)

        _UpperCamelCase							      =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 	UpperCamelCase__   (					self     :       List[str]	)  ->    Tuple:
    _UpperCamelCase      ,  _UpperCamelCase							      =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							      =self._prepare_for_class(UpperCamelCase__				,						UpperCamelCase__	)
        _UpperCamelCase							      ={k: torch.tensor(v.tolist()	) for k, v in prepared_inputs_dict.items()}

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

        _UpperCamelCase							      =pt_model_class(UpperCamelCase__	).eval()
        _UpperCamelCase							      =model_class(UpperCamelCase__				,						dtype=jnp.floataa	)

        _UpperCamelCase							      =load_flax_weights_in_pytorch_model(UpperCamelCase__				,						fx_model.params	)
        _UpperCamelCase      ,  _UpperCamelCase							      =pt_inputs['''input_ids'''].shape
        _UpperCamelCase							      =np.random.randint(0				,						seq_length - 1				,						size=(batch_size,)	)
        for batch_idx, start_index in enumerate(UpperCamelCase__	):
          _UpperCamelCase							      =0
          _UpperCamelCase							      =1
          _UpperCamelCase							      =0
          _UpperCamelCase							      =1

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

        with torch.no_grad():
          _UpperCamelCase							      =pt_model(**UpperCamelCase__	).to_tuple()

        _UpperCamelCase							      =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__	)
          _UpperCamelCase							      =pt_model_class.from_pretrained(UpperCamelCase__				,						from_flax=UpperCamelCase__	)

        with torch.no_grad():
          _UpperCamelCase							      =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	)




  @tooslow
  def 	UpperCamelCase__   (					self     :       Optional[int]	)  ->    List[str]:






    for model_class_name in self.all_model_classes:
      _UpperCamelCase							      =model_class_name.from_pretrained('''EleutherAI/gpt-j-6B'''	)
      _UpperCamelCase							      =model(np.ones((1, 1)	)	)
      self.assertIsNotNone(UpperCamelCase__	)









import unittest

from transformers import load_tool

from .test_tools_common import ToolTesterMixin


lowercase_   							=  '\nHugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf originally as a company that developed a chatbot app targeted at teenagers.[2] After open-sourcing the model behind the chatbot, the company pivoted to focus on being a platform for machine learning.\n\nIn March 2021, Hugging Face raised $40 million in a Series B funding round.[3]\n\nOn April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model.[4] In 2022, the workshop concluded with the announcement of BLOOM, a multilingual large language model with 176 billion parameters.[5]\n'





class 					__lowerCAmelCase	(   unittest.TestCase     ,     SCREAMING_SNAKE_CASE				):




  def     A__	(							self   )    ->					Union[str, Any]:
        '''simple docstring'''






        _lowercase      			=load_tool('text-question-answering'   )
        self.tool.setup()
        _lowercase      			=load_tool('text-question-answering'				,  remote=lowerCAmelCase   )




  def     A__	(							self   )    ->					Dict:
        '''simple docstring'''






        _lowercase      			=self.tool(lowerCAmelCase				,  'What did Hugging Face do in April 2021?'   )
        self.assertEqual(lowerCAmelCase				,  'launched the BigScience Research Workshop'   )




  def     A__	(							self   )    ->					Optional[Any]:
        '''simple docstring'''






        _lowercase      			=self.remote_tool(lowerCAmelCase				,  'What did Hugging Face do in April 2021?'   )
        self.assertEqual(lowerCAmelCase				,  'launched the BigScience Research Workshop'   )




  def     A__	(							self   )    ->					Any:
        '''simple docstring'''






        _lowercase      			=self.tool(text=lowerCAmelCase				,  question='What did Hugging Face do in April 2021?'   )
        self.assertEqual(lowerCAmelCase				,  'launched the BigScience Research Workshop'   )




  def     A__	(							self   )    ->					Optional[int]:
        '''simple docstring'''






        _lowercase      			=self.remote_tool(text=lowerCAmelCase				,  question='What did Hugging Face do in April 2021?'   )
        self.assertEqual(lowerCAmelCase				,  'launched the BigScience Research Workshop'   )


