






'''simple docstring'''







import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaConfig, GPTaTokenizer, 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, floats_tensor, 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.gpta.modeling_flax_gpta import FlaxGPTaLMHeadModel, FlaxGPTaModel

if is_torch_available():
       import torch






class 					a__    :
     '''simple docstring'''







     def __init__(self    :	List[Any]				,				lowercase__    :	List[str]				,				lowercase__    :	List[Any]=14				,				lowercase__    :	Dict=7				,				lowercase__    :	str=True				,				lowercase__    :	Dict=True				,				lowercase__    :	List[Any]=False				,				lowercase__    :	Any=True				,				lowercase__    :	List[str]=99				,				lowercase__    :	Union[str, Any]=32				,				lowercase__    :	List[Any]=5				,				lowercase__    :	Union[str, Any]=4				,				lowercase__    :	Optional[int]=37				,				lowercase__    :	str="gelu"				,				lowercase__    :	List[str]=0.1				,				lowercase__    :	Dict=0.1				,				lowercase__    :	str=512				,				lowercase__    :	List[Any]=0.02				,				)  ->       List[str]:
           __A		   =      parent
           __A		   =      batch_size
           __A		   =      seq_length
           __A		   =      is_training
           __A		   =      use_input_mask
           __A		   =      use_token_type_ids
           __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		   =      initializer_range
           __A		   =      None
           __A		   =      vocab_size - 1
           __A		   =      vocab_size - 1
           __A		   =      vocab_size - 1






     def     _UpperCamelCase   (self    :	Optional[int]					)  ->       str:
           __A		   =      ids_tensor([self.batch_size, self.seq_length]				,				self.vocab_size					)

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

           __A		   =      GPTaConfig(
               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=lowercase__				,				bos_token_id=self.bos_token_id				,				eos_token_id=self.eos_token_id				,				pad_token_id=self.pad_token_id				,				)

           return (config, input_ids, input_mask)






     def     _UpperCamelCase   (self    :	List[str]					)  ->       int:
           __A		   =      self.prepare_config_and_inputs()
           __A   ,				__A   ,				__A		   =      config_and_inputs
           __A		   =      {"input_ids": input_ids, "attention_mask": attention_mask}
           return config, inputs_dict






     def     _UpperCamelCase   (self    :	List[str]					)  ->       Tuple:
           __A   ,				__A   ,				__A		   =      self.prepare_config_and_inputs()

           __A		   =      floats_tensor([self.batch_size, self.seq_length, self.hidden_size]					)
           __A		   =      ids_tensor([self.batch_size, self.seq_length]				,				vocab_size=2					)

           return (
               config,
               input_ids,
               attention_mask,
               encoder_hidden_states,
               encoder_attention_mask,
           )






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

           __A		   =      model.init_cache(input_ids.shape[0]				,				lowercase__					)
           __A		   =      jnp.ones((input_ids.shape[0], max_decoder_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    :	Union[str, Any]				,				lowercase__    :	List[Any]				,				lowercase__    :	List[Any]				,				lowercase__    :	List[Any]				,				lowercase__    :	Dict					)  ->       Tuple:
           __A		   =      20
           __A		   =      model_class_name(lowercase__					)

           __A		   =      jnp.concatenate(
               [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_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 ,						lowerCamelCase ,						unittest.TestCase						):
     '''simple docstring'''







     _snake_case             =  (FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
     _snake_case             =  (FlaxGPTaLMHeadModel,) if is_flax_available() else ()






     def     _UpperCamelCase   (self    :	List[str]					)  ->       List[str]:
           __A		   =      FlaxGPTaModelTester(self					)






     def     _UpperCamelCase   (self    :	str					)  ->       Tuple:
           for model_class_name in self.all_model_classes:
                 __A   ,				__A   ,				__A		   =      self.model_tester.prepare_config_and_inputs()
                 self.model_tester.check_use_cache_forward(lowercase__				,				lowercase__				,				lowercase__				,				lowercase__					)






     def     _UpperCamelCase   (self    :	List[Any]					)  ->       Optional[Any]:
           for model_class_name in self.all_model_classes:
                 __A   ,				__A   ,				__A		   =      self.model_tester.prepare_config_and_inputs()
                 self.model_tester.check_use_cache_forward_with_attn_mask(
                     lowercase__				,				lowercase__				,				lowercase__				,				lowercase__					)






     @slow
     def     _UpperCamelCase   (self    :	List[Any]					)  ->       str:
           __A		   =      GPTaTokenizer.from_pretrained("gpt2"				,				pad_token="</s>"				,				padding_side="left"					)
           __A		   =      tokenizer(["Hello this is a long string", "Hey"]				,				return_tensors="np"				,				padding=lowercase__				,				truncation=lowercase__					)

           __A		   =      FlaxGPTaLMHeadModel.from_pretrained("gpt2"					)
           __A		   =      False
           __A		   =      model.config.eos_token_id

           __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__					)

           __A		   =      [
               "Hello this is a long string of words. I'm going to start with the first one.\n",
               "Hey, I'm not sure if I'm going to be able to do",
           ]

           self.assertListEqual(lowercase__				,				lowercase__					)






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

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

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

                       __A		   =      convert_pytorch_state_dict_to_flax(pt_model.state_dict()				,				lowercase__					)
                       __A		   =      fx_state

                       with torch.no_grad():
                             __A		   =      pt_model(**lowercase__					).to_tuple()

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

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

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






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

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

                       __A		   =      pt_model_class(lowercase__					).eval()
                       __A		   =      model_class(lowercase__				,				dtype=jnp.floataa					)

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

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

                       with torch.no_grad():
                             __A		   =      pt_model(**lowercase__					).to_tuple()

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

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

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

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






     @slow
     def     _UpperCamelCase   (self    :	Union[str, Any]					)  ->       Optional[Any]:
           for model_class_name in self.all_model_classes:
                 __A		   =      model_class_name.from_pretrained("gpt2"				,				from_pt=lowercase__					)
                 __A		   =      model(np.ones((1, 1)					)					)
                 self.assertIsNotNone(lowercase__					)










'''simple docstring'''







from typing import List

import datasets
from datasets.tasks import ImageClassification

from ..folder_based_builder import folder_based_builder


__A      :       Tuple      =	datasets.utils.logging.get_logger(__name__)






class 					a__    (			folder_based_builder.FolderBasedBuilderConfig						):
     '''simple docstring'''







     _snake_case             =  None
     _snake_case             =  None







class 					a__    (			folder_based_builder.FolderBasedBuilder						):
     '''simple docstring'''







     _snake_case             =  datasets.Image()
     _snake_case             =  'image'
     _snake_case             =  ImageFolderConfig
     _snake_case             =  42  # definition at the bottom of the script
     _snake_case             =  ImageClassification(image_column='image' ,						label_column='label'						)







__A      :       Any      =	[
    '.blp',
    '.bmp',
    '.dib',
    '.bufr',
    '.cur',
    '.pcx',
    '.dcx',
    '.dds',
    '.ps',
    '.eps',
    '.fit',
    '.fits',
    '.fli',
    '.flc',
    '.ftc',
    '.ftu',
    '.gbr',
    '.gif',
    '.grib',
    '.h5',
    '.hdf',
    '.png',
    '.apng',
    '.jp2',
    '.j2k',
    '.jpc',
    '.jpf',
    '.jpx',
    '.j2c',
    '.icns',
    '.ico',
    '.im',
    '.iim',
    '.tif',
    '.tiff',
    '.jfif',
    '.jpe',
    '.jpg',
    '.jpeg',
    '.mpg',
    '.mpeg',
    '.msp',
    '.pcd',
    '.pxr',
    '.pbm',
    '.pgm',
    '.ppm',
    '.pnm',
    '.psd',
    '.bw',
    '.rgb',
    '.rgba',
    '.sgi',
    '.ras',
    '.tga',
    '.icb',
    '.vda',
    '.vst',
    '.webp',
    '.wmf',
    '.emf',
    '.xbm',
    '.xpm',
]
__A      :       List[Any]      =	IMAGE_EXTENSIONS


