




"""simple docstring"""


import json
import os
from dataclasses import dataclass
from functools import partial
from typing import Callable

import flax.linen as nn
import jax
import jax.numpy as jnp
import joblib
import optax
import wandb
from flax import jax_utils, struct, traverse_util
from flax.serialization import from_bytes, to_bytes
from flax.training import train_state
from flax.training.common_utils import shard
from tqdm.auto import tqdm

from transformers import BigBirdConfig, FlaxBigBirdForQuestionAnswering
from transformers.models.big_bird.modeling_flax_big_bird import FlaxBigBirdForQuestionAnsweringModule






class       lowercase(  UpperCamelCase__				):







	'''simple docstring'''





	lowercase__				      =      42
	lowercase__				      =      jnp.floataa
	lowercase__				      =      True


	def 				UpperCamelCase_     (       self:	Optional[int]	):




			'''simple docstring'''



			super().setup()
			_snake_case							:	Optional[Any]			=						nn.Dense(5,   dtype=self.dtype	)


	def __call__(       self:	Union[str, Any],   *a_:	Optional[Any],   **a_:	Optional[Any]	):




			'''simple docstring'''



			_snake_case							:	Union[str, Any]			=						super().__call__(*__A,   **__A	)
			_snake_case							:	Tuple			=						self.cls(outputs[2]	)
			return outputs[:2] + (cls_out,)







class       lowercase(  UpperCamelCase__				):







	'''simple docstring'''





	lowercase__				      =      FlaxBigBirdForNaturalQuestionsModule







def       UpperCAmelCase__ (snake_case__							:      List[str]	,     snake_case__							:      List[str]	,     snake_case__							:      Tuple	,     snake_case__							:      List[str]	,     snake_case__							:      Optional[int]	,     snake_case__							:      Optional[Any]					):
		"""simple docstring"""

		def cross_entropy(snake_case__							:      List[Any]	,     snake_case__							:      Tuple	,     snake_case__							:      Union[str, Any]=None					):
				_snake_case							:	str			=						logits.shape[-1]
				_snake_case							:	Dict			=						(labels[..., None] == jnp.arange(UpperCamelCase__					)[None]).astype("""f4"""					)
				_snake_case							:	Optional[int]			=						jax.nn.log_softmax(UpperCamelCase__	,     axis=-1					)
				_snake_case							:	List[str]			=						-jnp.sum(labels * logits	,     axis=-1					)
				if reduction is not None:
						_snake_case							:	List[str]			=						reduction(UpperCamelCase__					)
				return loss

		_snake_case							:	Dict			=						partial(UpperCamelCase__	,     reduction=jnp.mean					)
		_snake_case							:	List[Any]			=						cross_entropy(UpperCamelCase__	,     UpperCamelCase__					)
		_snake_case							:	Optional[int]			=						cross_entropy(UpperCamelCase__	,     UpperCamelCase__					)
		_snake_case							:	Optional[Any]			=						cross_entropy(UpperCamelCase__	,     UpperCamelCase__					)
		return (start_loss + end_loss + pooled_loss) / 3






@dataclass
class       lowercase:







	'''simple docstring'''





	lowercase__				      =      "google/bigbird-roberta-base"
	lowercase__				      =      30_00
	lowercase__				      =      1_05_00

	lowercase__				      =      1_28
	lowercase__				      =      3

	lowercase__				      =      1
	lowercase__				      =      5

	# tx_args
	lowercase__				      =      3e-5
	lowercase__				      =      0.0
	lowercase__				      =      2_00_00
	lowercase__				      =      0.0095

	lowercase__				      =      "bigbird-roberta-natural-questions"
	lowercase__				      =      "training-expt"
	lowercase__				      =      "data/nq-training.jsonl"
	lowercase__				      =      "data/nq-validation.jsonl"


	def 				UpperCamelCase_     (       self:	int	):




			'''simple docstring'''



			os.makedirs(self.base_dir,   exist_ok=__A	)
			_snake_case							:	str			=						os.path.join(self.base_dir,   self.save_dir	)
			_snake_case							:	List[str]			=						self.batch_size_per_device * jax.device_count()







@dataclass
class       lowercase:







	'''simple docstring'''





	lowercase__				      =      42
	lowercase__				      =      40_96  # no dynamic padding on TPUs


	def __call__(       self:	str,   a_:	Union[str, Any]	):




			'''simple docstring'''



			_snake_case							:	str			=						self.collate_fn(__A	)
			_snake_case							:	Union[str, Any]			=						jax.tree_util.tree_map(__A,   __A	)
			return batch


	def 				UpperCamelCase_     (       self:	List[str],   a_:	Union[str, Any]	):




			'''simple docstring'''



			_snake_case	,      _snake_case							:	Optional[Any]			=						self.fetch_inputs(features["""input_ids"""]	)
			_snake_case							:	Optional[int]			=						{
			    """input_ids""": jnp.array(__A,   dtype=jnp.intaa	),
			    """attention_mask""": jnp.array(__A,   dtype=jnp.intaa	),
			    """start_labels""": jnp.array(features["""start_token"""],   dtype=jnp.intaa	),
			    """end_labels""": jnp.array(features["""end_token"""],   dtype=jnp.intaa	),
			    """pooled_labels""": jnp.array(features["""category"""],   dtype=jnp.intaa	),
			}
			return batch


	def 				UpperCamelCase_     (       self:	List[Any],   a_:	list	):




			'''simple docstring'''



			_snake_case							:	str			=						[self._fetch_inputs(__A	) for ids in input_ids]
			return zip(*__A	)


	def 				UpperCamelCase_     (       self:	Union[str, Any],   a_:	list	):




			'''simple docstring'''



			_snake_case							:	Optional[Any]			=						[1 for _ in range(len(__A	)	)]
			while len(__A	) < self.max_length:
					input_ids.append(self.pad_id	)
					attention_mask.append(0	)
			return input_ids, attention_mask







def       UpperCAmelCase__ (snake_case__							:      List[Any]	,     snake_case__							:      Any	,     snake_case__							:      int=None					):
		"""simple docstring"""

		if seed is not None:
				_snake_case							:	List[str]			=						dataset.shuffle(seed=UpperCamelCase__					)
		for i in range(len(UpperCamelCase__					) // batch_size					):
				_snake_case							:	Union[str, Any]			=						dataset[i * batch_size : (i + 1) * batch_size]
				yield dict(UpperCamelCase__					)


@partial(jax.pmap	,     axis_name="""batch"""					)
def       UpperCAmelCase__ (snake_case__							:      Optional[Any]	,     snake_case__							:      int	,     **snake_case__							:      Dict					):
		"""simple docstring"""

		def loss_fn(snake_case__							:      Any					):
				_snake_case							:	Optional[int]			=						model_inputs.pop("""start_labels"""					)
				_snake_case							:	Dict			=						model_inputs.pop("""end_labels"""					)
				_snake_case							:	Optional[int]			=						model_inputs.pop("""pooled_labels"""					)

				_snake_case							:	int			=						state.apply_fn(**UpperCamelCase__	,     params=UpperCamelCase__	,     dropout_rng=UpperCamelCase__	,     train=UpperCamelCase__					)
				_snake_case	,      _snake_case	,      _snake_case							:	Optional[int]			=						outputs

				return state.loss_fn(
				    UpperCamelCase__	,     UpperCamelCase__	,     UpperCamelCase__	,     UpperCamelCase__	,     UpperCamelCase__	,     UpperCamelCase__	,     )

		_snake_case	,      _snake_case							:	Dict			=						jax.random.split(UpperCamelCase__					)
		_snake_case							:	Optional[int]			=						jax.value_and_grad(UpperCamelCase__					)
		_snake_case	,      _snake_case							:	int			=						grad_fn(state.params					)
		_snake_case							:	List[str]			=						jax.lax.pmean({"""loss""": loss}	,     axis_name="""batch"""					)
		_snake_case							:	Dict			=						jax.lax.pmean(UpperCamelCase__	,     """batch"""					)

		_snake_case							:	Union[str, Any]			=						state.apply_gradients(grads=UpperCamelCase__					)
		return state, metrics, new_drp_rng


@partial(jax.pmap	,     axis_name="""batch"""					)
def       UpperCAmelCase__ (snake_case__							:      Optional[Any]	,     **snake_case__							:      Optional[Any]					):
		"""simple docstring"""

		_snake_case							:	int			=						model_inputs.pop("""start_labels"""					)
		_snake_case							:	List[Any]			=						model_inputs.pop("""end_labels"""					)
		_snake_case							:	Optional[Any]			=						model_inputs.pop("""pooled_labels"""					)

		_snake_case							:	Union[str, Any]			=						state.apply_fn(**UpperCamelCase__	,     params=state.params	,     train=UpperCamelCase__					)
		_snake_case	,      _snake_case	,      _snake_case							:	Optional[Any]			=						outputs

		_snake_case							:	int			=						state.loss_fn(UpperCamelCase__	,     UpperCamelCase__	,     UpperCamelCase__	,     UpperCamelCase__	,     UpperCamelCase__	,     UpperCamelCase__					)
		_snake_case							:	str			=						jax.lax.pmean({"""loss""": loss}	,     axis_name="""batch"""					)
		return metrics






class       lowercase(  train_state.TrainState				):







	'''simple docstring'''





	lowercase__				      =      struct.field(pytree_node=UpperCamelCase__				)







@dataclass
class       lowercase:







	'''simple docstring'''





	lowercase__				      =      42
	lowercase__				      =      42
	lowercase__				      =      42
	lowercase__				      =      42
	lowercase__				      =      42
	lowercase__				      =      42
	lowercase__				      =      None


	def 				UpperCamelCase_     (       self:	Dict,   a_:	int,   a_:	str,   a_:	Optional[Any],   a_:	Union[str, Any]=None	):




			'''simple docstring'''



			_snake_case							:	Dict			=						model.params
			_snake_case							:	str			=						TrainState.create(
			    apply_fn=model.__call__,   params=__A,   tx=__A,   loss_fn=__A,   )
			if ckpt_dir is not None:
					_snake_case	,      _snake_case	,      _snake_case	,      _snake_case	,      _snake_case							:	str			=						restore_checkpoint(__A,   __A	)
					_snake_case							:	Tuple			=						{
					    """lr""": args.lr,
					    """init_lr""": args.init_lr,
					    """warmup_steps""": args.warmup_steps,
					    """num_train_steps""": num_train_steps,
					    """weight_decay""": args.weight_decay,
					}
					_snake_case	,      _snake_case							:	Any			=						build_tx(**__A	)
					_snake_case							:	str			=						train_state.TrainState(
					    step=__A,   apply_fn=model.__call__,   params=__A,   tx=__A,   opt_state=__A,   )
					_snake_case							:	Any			=						args
					_snake_case							:	Tuple			=						data_collator
					_snake_case							:	Optional[int]			=						lr
					_snake_case							:	Optional[int]			=						params
			_snake_case							:	Dict			=						jax_utils.replicate(__A	)
			return state


	def 				UpperCamelCase_     (       self:	Union[str, Any],   a_:	Any,   a_:	Optional[Any],   a_:	Optional[Any]	):




			'''simple docstring'''



			_snake_case							:	str			=						self.args
			_snake_case							:	List[str]			=						len(__A	) // args.batch_size

			_snake_case							:	List[str]			=						jax.random.PRNGKey(0	)
			_snake_case							:	List[str]			=						jax.random.split(__A,   jax.device_count()	)
			for epoch in range(args.max_epochs	):
					_snake_case							:	str			=						jnp.array(0,   dtype=jnp.floataa	)
					_snake_case							:	List[Any]			=						get_batched_dataset(__A,   args.batch_size,   seed=__A	)
					_snake_case							:	Any			=						0
					for batch in tqdm(__A,   total=__A,   desc=f"Running EPOCH-{epoch}"	):
							_snake_case							:	Any			=						self.data_collator(__A	)
							_snake_case	,      _snake_case	,      _snake_case							:	List[Any]			=						self.train_step_fn(__A,   __A,   **__A	)
							running_loss += jax_utils.unreplicate(metrics["""loss"""]	)
							i += 1
							if i % args.logging_steps == 0:
									_snake_case							:	Optional[Any]			=						jax_utils.unreplicate(state.step	)
									_snake_case							:	List[str]			=						running_loss.item() / i
									_snake_case							:	Any			=						self.scheduler_fn(state_step - 1	)

									_snake_case							:	Optional[Any]			=						self.evaluate(__A,   __A	)
									_snake_case							:	List[Any]			=						{
									    """step""": state_step.item(),
									    """eval_loss""": eval_loss.item(),
									    """tr_loss""": tr_loss,
									    """lr""": lr.item(),
									}
									tqdm.write(str(__A	)	)
									self.logger.log(__A,   commit=__A	)

							if i % args.save_steps == 0:
									self.save_checkpoint(args.save_dir + f"-e{epoch}-s{i}",   state=__A	)


	def 				UpperCamelCase_     (       self:	str,   a_:	Union[str, Any],   a_:	Optional[int]	):




			'''simple docstring'''



			_snake_case							:	int			=						get_batched_dataset(__A,   self.args.batch_size	)
			_snake_case							:	Dict			=						len(__A	) // self.args.batch_size
			_snake_case							:	Tuple			=						jnp.array(0,   dtype=jnp.floataa	)
			_snake_case							:	List[Any]			=						0
			for batch in tqdm(__A,   total=__A,   desc="""Evaluating ... """	):
					_snake_case							:	List[Any]			=						self.data_collator(__A	)
					_snake_case							:	int			=						self.val_step_fn(__A,   **__A	)
					running_loss += jax_utils.unreplicate(metrics["""loss"""]	)
					i += 1
			return running_loss / i


	def 				UpperCamelCase_     (       self:	int,   a_:	int,   a_:	Tuple	):




			'''simple docstring'''



			_snake_case							:	int			=						jax_utils.unreplicate(__A	)
			print(f"SAVING CHECKPOINT IN {save_dir}",   end=""" ... """	)
			self.model_save_fn(__A,   params=state.params	)
			with open(os.path.join(__A,   """opt_state.msgpack"""	),   """wb"""	) as f:
					f.write(to_bytes(state.opt_state	)	)
			joblib.dump(self.args,   os.path.join(__A,   """args.joblib"""	)	)
			joblib.dump(self.data_collator,   os.path.join(__A,   """data_collator.joblib"""	)	)
			with open(os.path.join(__A,   """training_state.json"""	),   """w"""	) as f:
					json.dump({"""step""": state.step.item()},   __A	)
			print("""DONE"""	)







def       UpperCAmelCase__ (snake_case__							:      List[str]	,     snake_case__							:      str					):
		"""simple docstring"""

		print(F"RESTORING CHECKPOINT FROM {save_dir}"	,     end=""" ... """					)
		with open(os.path.join(UpperCamelCase__	,     """flax_model.msgpack"""					)	,     """rb"""					) as f:
				_snake_case							:	str			=						from_bytes(state.params	,     f.read()					)

		with open(os.path.join(UpperCamelCase__	,     """opt_state.msgpack"""					)	,     """rb"""					) as f:
				_snake_case							:	List[str]			=						from_bytes(state.opt_state	,     f.read()					)

		_snake_case							:	Optional[Any]			=						joblib.load(os.path.join(UpperCamelCase__	,     """args.joblib"""					)					)
		_snake_case							:	Optional[Any]			=						joblib.load(os.path.join(UpperCamelCase__	,     """data_collator.joblib"""					)					)

		with open(os.path.join(UpperCamelCase__	,     """training_state.json"""					)	,     """r"""					) as f:
				_snake_case							:	List[str]			=						json.load(UpperCamelCase__					)
		_snake_case							:	int			=						training_state["""step"""]

		print("""DONE"""					)
		return params, opt_state, step, args, data_collator


def       UpperCAmelCase__ (snake_case__							:      int	,     snake_case__							:      Optional[Any]	,     snake_case__							:      Optional[int]	,     snake_case__							:      Optional[int]					):
		"""simple docstring"""

		_snake_case							:	Tuple			=						num_train_steps - warmup_steps
		_snake_case							:	str			=						optax.linear_schedule(init_value=UpperCamelCase__	,     end_value=UpperCamelCase__	,     transition_steps=UpperCamelCase__					)
		_snake_case							:	Optional[Any]			=						optax.linear_schedule(init_value=UpperCamelCase__	,     end_value=1e-7	,     transition_steps=UpperCamelCase__					)
		_snake_case							:	str			=						optax.join_schedules(schedules=[warmup_fn, decay_fn]	,     boundaries=[warmup_steps]					)
		return lr






def       UpperCAmelCase__ (snake_case__							:      Optional[Any]	,     snake_case__							:      Optional[int]	,     snake_case__							:      List[str]	,     snake_case__							:      Optional[int]	,     snake_case__							:      Dict					):
		"""simple docstring"""

		def weight_decay_mask(snake_case__							:      Optional[Any]					):
				_snake_case							:	List[Any]			=						traverse_util.flatten_dict(UpperCamelCase__					)
				_snake_case							:	List[Any]			=						{k: (v[-1] != """bias""" and v[-2:] != ("""LayerNorm""", """scale""")) for k, v in params.items()}
				return traverse_util.unflatten_dict(UpperCamelCase__					)

		_snake_case							:	Any			=						scheduler_fn(UpperCamelCase__	,     UpperCamelCase__	,     UpperCamelCase__	,     UpperCamelCase__					)

		_snake_case							:	Union[str, Any]			=						optax.adamw(learning_rate=UpperCamelCase__	,     weight_decay=UpperCamelCase__	,     mask=UpperCamelCase__					)
		return tx, lr








"""simple docstring"""


import inspect
import unittest

from transformers import ConvNextVaConfig
from transformers.models.auto import get_values
from transformers.models.auto.modeling_auto import MODEL_FOR_BACKBONE_MAPPING_NAMES, MODEL_MAPPING_NAMES
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import ConvNextVaBackbone, ConvNextVaForImageClassification, ConvNextVaModel
     from transformers.models.convnextva.modeling_convnextva import CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
     from PIL import Image

     from transformers import AutoImageProcessor






class       lowercase:







 '''simple docstring'''





 def __init__(       self:	List[str],   a_:	List[Any],   a_:	str=13,   a_:	Dict=32,   a_:	Union[str, Any]=3,   a_:	Union[str, Any]=4,   a_:	Tuple=[10, 20, 30, 40],   a_:	Dict=[2, 2, 3, 2],   a_:	Tuple=True,   a_:	Optional[Any]=True,   a_:	Any=37,   a_:	Any="gelu",   a_:	int=10,   a_:	Tuple=0.02,   a_:	str=["stage2", "stage3", "stage4"],   a_:	List[str]=[2, 3, 4],   a_:	List[str]=None,   ):




   '''simple docstring'''



   _snake_case							:	int			=						parent
   _snake_case							:	int			=						batch_size
   _snake_case							:	List[Any]			=						image_size
   _snake_case							:	List[str]			=						num_channels
   _snake_case							:	Tuple			=						num_stages
   _snake_case							:	Union[str, Any]			=						hidden_sizes
   _snake_case							:	List[Any]			=						depths
   _snake_case							:	Tuple			=						is_training
   _snake_case							:	List[str]			=						use_labels
   _snake_case							:	Tuple			=						intermediate_size
   _snake_case							:	List[str]			=						hidden_act
   _snake_case							:	Optional[Any]			=						num_labels
   _snake_case							:	Tuple			=						initializer_range
   _snake_case							:	Tuple			=						out_features
   _snake_case							:	Tuple			=						out_indices
   _snake_case							:	Dict			=						scope


 def 				UpperCamelCase_     (       self:	List[str]	):




   '''simple docstring'''



   _snake_case							:	Optional[int]			=						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

   _snake_case							:	Any			=						None
   if self.use_labels:
     _snake_case							:	Dict			=						ids_tensor([self.batch_size],   self.num_labels	)

   _snake_case							:	Optional[Any]			=						self.get_config()

   return config, pixel_values, labels


 def 				UpperCamelCase_     (       self:	Tuple	):




   '''simple docstring'''



   return ConvNextVaConfig(
       num_channels=self.num_channels,   hidden_sizes=self.hidden_sizes,   depths=self.depths,   num_stages=self.num_stages,   hidden_act=self.hidden_act,   is_decoder=a_,   initializer_range=self.initializer_range,   out_features=self.out_features,   out_indices=self.out_indices,   num_labels=self.num_labels,   )


 def 				UpperCamelCase_     (       self:	int,   a_:	Tuple,   a_:	Any,   a_:	Dict	):




   '''simple docstring'''



   _snake_case							:	int			=						ConvNextVaModel(config=a_	)
   model.to(a_	)
   model.eval()
   _snake_case							:	Any			=						model(a_	)
   # expected last hidden states: B, C, H // 32, W // 32
   self.parent.assertEqual(
       result.last_hidden_state.shape,   (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32),   )


 def 				UpperCamelCase_     (       self:	Optional[int],   a_:	List[str],   a_:	Tuple,   a_:	Union[str, Any]	):




   '''simple docstring'''



   _snake_case							:	Optional[int]			=						ConvNextVaForImageClassification(a_	)
   model.to(a_	)
   model.eval()
   _snake_case							:	Optional[int]			=						model(a_,   labels=a_	)
   self.parent.assertEqual(result.logits.shape,   (self.batch_size, self.num_labels)	)


 def 				UpperCamelCase_     (       self:	Union[str, Any],   a_:	Tuple,   a_:	Tuple,   a_:	Tuple	):




   '''simple docstring'''



   _snake_case							:	List[str]			=						ConvNextVaBackbone(config=a_	)
   model.to(a_	)
   model.eval()
   _snake_case							:	int			=						model(a_	)

   # verify hidden states
   self.parent.assertEqual(len(result.feature_maps	),   len(config.out_features	)	)
   self.parent.assertListEqual(list(result.feature_maps[0].shape	),   [self.batch_size, self.hidden_sizes[1], 4, 4]	)

   # verify channels
   self.parent.assertEqual(len(model.channels	),   len(config.out_features	)	)
   self.parent.assertListEqual(model.channels,   config.hidden_sizes[1:]	)

   # verify backbone works with out_features=None
   _snake_case							:	Tuple			=						None
   _snake_case							:	Tuple			=						ConvNextVaBackbone(config=a_	)
   model.to(a_	)
   model.eval()
   _snake_case							:	List[Any]			=						model(a_	)

   # verify feature maps
   self.parent.assertEqual(len(result.feature_maps	),   1	)
   self.parent.assertListEqual(list(result.feature_maps[0].shape	),   [self.batch_size, self.hidden_sizes[-1], 1, 1]	)

   # verify channels
   self.parent.assertEqual(len(model.channels	),   1	)
   self.parent.assertListEqual(model.channels,   [config.hidden_sizes[-1]]	)


 def 				UpperCamelCase_     (       self:	Dict	):




   '''simple docstring'''



   _snake_case							:	Dict			=						self.prepare_config_and_inputs()
   _snake_case	,      _snake_case	,      _snake_case							:	Any			=						config_and_inputs
   _snake_case							:	str			=						{"""pixel_values""": pixel_values}
   return config, inputs_dict


 def 				UpperCamelCase_     (       self:	Optional[int]	):




   '''simple docstring'''



   _snake_case							:	List[Any]			=						self.prepare_config_and_inputs()
   _snake_case	,      _snake_case	,      _snake_case							:	List[str]			=						config_and_inputs
   _snake_case							:	Any			=						{"""pixel_values""": pixel_values, """labels""": labels}
   return config, inputs_dict







@require_torch
class       lowercase(  __a ,    __a ,    unittest.TestCase				):







 '''simple docstring'''





 lowercase__				      =      (
     (
         ConvNextVaModel,
         ConvNextVaForImageClassification,
         ConvNextVaBackbone,
     )
     if is_torch_available()
     else ()
 )
 lowercase__				      =      (
     {"feature-extraction": ConvNextVaModel, "image-classification": ConvNextVaForImageClassification}
     if is_torch_available()
     else {}
 )

 lowercase__				      =      False
 lowercase__				      =      False
 lowercase__				      =      False
 lowercase__				      =      False
 lowercase__				      =      False


 def 				UpperCamelCase_     (       self:	Tuple	):




   '''simple docstring'''



   _snake_case							:	Tuple			=						ConvNextVaModelTester(self	)
   _snake_case							:	int			=						ConfigTester(self,   config_class=a_,   has_text_modality=a_,   hidden_size=37	)


 def 				UpperCamelCase_     (       self:	List[str]	):




   '''simple docstring'''



   self.create_and_test_config_common_properties()
   self.config_tester.create_and_test_config_to_json_string()
   self.config_tester.create_and_test_config_to_json_file()
   self.config_tester.create_and_test_config_from_and_save_pretrained()
   self.config_tester.create_and_test_config_with_num_labels()
   self.config_tester.check_config_can_be_init_without_params()
   self.config_tester.check_config_arguments_init()


 def 				UpperCamelCase_     (       self:	Tuple	):




   '''simple docstring'''



   return


 @unittest.skip(reason="""ConvNextV2 does not use inputs_embeds"""	)
 def 				UpperCamelCase_     (       self:	int	):




   '''simple docstring'''



   pass


 @unittest.skip(reason="""ConvNextV2 does not support input and output embeddings"""	)
 def 				UpperCamelCase_     (       self:	Any	):




   '''simple docstring'''



   pass


 @unittest.skip(reason="""ConvNextV2 does not use feedforward chunking"""	)
 def 				UpperCamelCase_     (       self:	str	):




   '''simple docstring'''



   pass


 def 				UpperCamelCase_     (       self:	int	):




   '''simple docstring'''



   if not self.model_tester.is_training:
     return

   for model_class in self.all_model_classes:
     _snake_case	,      _snake_case							:	Union[str, Any]			=						self.model_tester.prepare_config_and_inputs_with_labels()
     _snake_case							:	List[Any]			=						True

     if model_class.__name__ in [
         *get_values(a_	),
         *get_values(a_	),
     ]:
       continue

     _snake_case							:	Tuple			=						model_class(a_	)
     model.to(a_	)
     model.train()
     _snake_case							:	Optional[Any]			=						self._prepare_for_class(a_,   a_,   return_labels=a_	)
     _snake_case							:	Any			=						model(**a_	).loss
     loss.backward()


 def 				UpperCamelCase_     (       self:	Optional[int]	):




   '''simple docstring'''



   if not self.model_tester.is_training:
     return

   for model_class in self.all_model_classes:
     _snake_case	,      _snake_case							:	str			=						self.model_tester.prepare_config_and_inputs_with_labels()
     _snake_case							:	Any			=						False
     _snake_case							:	List[Any]			=						True

     if (
         model_class.__name__
         in [*get_values(a_	), *get_values(a_	)]
         or not model_class.supports_gradient_checkpointing
     ):
       continue

     _snake_case							:	Dict			=						model_class(a_	)
     model.to(a_	)
     model.gradient_checkpointing_enable()
     model.train()
     _snake_case							:	str			=						self._prepare_for_class(a_,   a_,   return_labels=a_	)
     _snake_case							:	Optional[int]			=						model(**a_	).loss
     loss.backward()


 def 				UpperCamelCase_     (       self:	str	):




   '''simple docstring'''



   _snake_case	,      _snake_case							:	List[Any]			=						self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     _snake_case							:	List[str]			=						model_class(a_	)
     _snake_case							:	Union[str, Any]			=						inspect.signature(model.forward	)
     # signature.parameters is an OrderedDict => so arg_names order is deterministic
     _snake_case							:	int			=						[*signature.parameters.keys()]

     _snake_case							:	Union[str, Any]			=						["""pixel_values"""]
     self.assertListEqual(arg_names[:1],   a_	)


 def 				UpperCamelCase_     (       self:	int	):




   '''simple docstring'''



   _snake_case							:	Tuple			=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*a_	)


 def 				UpperCamelCase_     (       self:	Union[str, Any]	):




   '''simple docstring'''



   def check_hidden_states_output(a_:	str,   a_:	Tuple,   a_:	Tuple	):
     _snake_case							:	Optional[Any]			=						model_class(a_	)
     model.to(a_	)
     model.eval()

     with torch.no_grad():
       _snake_case							:	Any			=						model(**self._prepare_for_class(a_,   a_	)	)

     _snake_case							:	Optional[int]			=						outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

     _snake_case							:	Optional[int]			=						self.model_tester.num_stages
     self.assertEqual(len(a_	),   expected_num_stages + 1	)

     # ConvNextV2's feature maps are of shape (batch_size, num_channels, height, width)
     self.assertListEqual(
         list(hidden_states[0].shape[-2:]	),   [self.model_tester.image_size // 4, self.model_tester.image_size // 4],   )

   _snake_case	,      _snake_case							:	Optional[Any]			=						self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     _snake_case							:	Optional[Any]			=						True
     check_hidden_states_output(a_,   a_,   a_	)

     # check that output_hidden_states also work using config
     del inputs_dict["output_hidden_states"]
     _snake_case							:	List[str]			=						True

     check_hidden_states_output(a_,   a_,   a_	)


 def 				UpperCamelCase_     (       self:	Dict	):




   '''simple docstring'''



   _snake_case							:	List[str]			=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_image_classification(*a_	)


 @slow
 def 				UpperCamelCase_     (       self:	Dict	):




   '''simple docstring'''



   for model_name in CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _snake_case							:	str			=						ConvNextVaModel.from_pretrained(a_	)
     self.assertIsNotNone(a_	)







def       UpperCAmelCase__ ():
  """simple docstring"""

  _snake_case							:	List[Any]			=						Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""					)
  return image






@require_torch
@require_vision
class       lowercase(  unittest.TestCase				):







 '''simple docstring'''





 @cached_property
 def 				UpperCamelCase_     (       self:	Optional[Any]	):




   '''simple docstring'''



   return AutoImageProcessor.from_pretrained("""facebook/convnextv2-tiny-1k-224"""	) if is_vision_available() else None


 @slow
 def 				UpperCamelCase_     (       self:	int	):




   '''simple docstring'''



   _snake_case							:	Optional[Any]			=						ConvNextVaForImageClassification.from_pretrained("""facebook/convnextv2-tiny-1k-224"""	).to(a_	)

   _snake_case							:	Union[str, Any]			=						self.default_image_processor
   _snake_case							:	List[Any]			=						prepare_img()
   _snake_case							:	Optional[int]			=						preprocessor(images=a_,   return_tensors="""pt"""	).to(a_	)

   # forward pass
   with torch.no_grad():
     _snake_case							:	Optional[int]			=						model(**a_	)

   # verify the logits
   _snake_case							:	Optional[int]			=						torch.Size((1, 1_000)	)
   self.assertEqual(outputs.logits.shape,   a_	)

   _snake_case							:	Optional[int]			=						torch.tensor([0.9_996, 0.1_966, -0.4_386]	).to(a_	)
   self.assertTrue(torch.allclose(outputs.logits[0, :3],   a_,   atol=1E-4	)	)


