
'''simple docstring'''

import logging
import os
import sys
from dataclasses import dataclass, field
from typing import Optional

from seqaseq_trainer import SeqaSeqTrainer
from seqaseq_training_args import SeqaSeqTrainingArguments

import transformers
from transformers import (
    AutoConfig,
    AutoModelForSeqaSeqLM,
    AutoTokenizer,
    HfArgumentParser,
    MBartTokenizer,
    MBartTokenizerFast,
    set_seed,
)
from transformers.trainer_utils import EvaluationStrategy, is_main_process
from transformers.training_args import ParallelMode
from utils import (
    SeqaSeqDataCollator,
    SeqaSeqDataset,
    assert_all_frozen,
    build_compute_metrics_fn,
    check_output_dir,
    freeze_embeds,
    freeze_params,
    lmap,
    save_json,
    use_task_specific_params,
    write_txt_file,
)


__snake_case:	str										=      logging.getLogger(__name__)


@dataclass
class 			_UpperCAmelCase			:

     """simple docstring"""

     a_     						=       field(
         metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}     )
     a_     						=       field(
         default=__A						,metadata={"help": "Pretrained config name or path if not the same as model_name"}     )
     a_     						=       field(
         default=__A						,metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}     )
     a_     						=       field(
         default=__A						,metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"}						,)
     a_     						=       field(default=__A						,metadata={"help": "Whether tp freeze the encoder."}     )
     a_     						=       field(default=__A						,metadata={"help": "Whether  to freeze the embeddings."}     )







@dataclass
class 			_UpperCAmelCase			:

     """simple docstring"""

     a_     						=       field(
         metadata={"help": "The input data dir. Should contain the .tsv files (or other data files) for the task."}     )
     a_     						=       field(
         default="summarization"						,metadata={"help": "Task name, summarization (or summarization_{dataset} for pegasus) or translation"}						,)
     a_     						=       field(
         default=1_024						,metadata={
             "help": (
                 "The maximum total input sequence length after tokenization. Sequences longer "
                 "than this will be truncated, sequences shorter will be padded."
             )
         }						,)
     a_     						=       field(
         default=128						,metadata={
             "help": (
                 "The maximum total sequence length for target text after tokenization. Sequences longer "
                 "than this will be truncated, sequences shorter will be padded."
             )
         }						,)
     a_     						=       field(
         default=142						,metadata={
             "help": (
                 "The maximum total sequence length for validation target text after tokenization. Sequences longer "
                 "than this will be truncated, sequences shorter will be padded. "
                 "This argument is also used to override the ``max_length`` param of ``model.generate``, which is used "
                 "during ``evaluate`` and ``predict``."
             )
         }						,)
     a_     						=       field(
         default=142						,metadata={
             "help": (
                 "The maximum total sequence length for test target text after tokenization. Sequences longer "
                 "than this will be truncated, sequences shorter will be padded."
             )
         }						,)
     a_     						=       field(default=-1						,metadata={"help": "# training examples. -1 means use all."}     )
     a_     						=       field(default=-1						,metadata={"help": "# validation examples. -1 means use all."}     )
     a_     						=       field(default=-1						,metadata={"help": "# test examples. -1 means use all."}     )
     a_     						=       field(default=__A						,metadata={"help": "Source language id for translation."}     )
     a_     						=       field(default=__A						,metadata={"help": "Target language id for translation."}     )
     a_     						=       field(default=__A						,metadata={"help": "# num_beams to use for evaluation."}     )
     a_     						=       field(
         default=__A						,metadata={"help": "If only pad tokens should be ignored. This assumes that `config.pad_token_id` is defined."}						,)







def 						_snake_case    (   A_							:     str  ,		A_							:     Optional[Any]  ,		A_							:     List[str] ):
  """simple docstring"""







  logger.info(f'''***** {split} metrics *****''' )
  for key in sorted(metrics.keys() ):
    logger.info(f'''  {key} = {metrics[key]}''' )
  save_json(A_  ,		os.path.join(A_  ,		f'''{split}_results.json''' ) )






def 						_snake_case    (   ):
  """simple docstring"""







  a_   :						Optional[int]   =				HfArgumentParser((ModelArguments, DataTrainingArguments, SeqaSeqTrainingArguments) )

  if len(sys.argv ) == 2 and sys.argv[1].endswith(""".json""" ):
    # If we pass only one argument to the script and it's the path to a json file,
    # let's parse it to get our arguments.
    a_   :						Optional[Any]   =				parser.parse_json_file(json_file=os.path.abspath(sys.argv[1] ) )
  else:
    a_   :						Tuple   =				parser.parse_args_into_dataclasses()

  check_output_dir(A_ )

  # Setup logging
  logging.basicConfig(
      format="""%(asctime)s - %(levelname)s - %(name)s - %(message)s"""  ,		datefmt="""%m/%d/%Y %H:%M:%S"""  ,		level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN  ,		)
  logger.warning(
      """Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s"""  ,		training_args.local_rank  ,		training_args.device  ,		training_args.n_gpu  ,		bool(training_args.parallel_mode == ParallelMode.DISTRIBUTED )  ,		training_args.fpaa  ,		)
  transformers.utils.logging.enable_default_handler()
  transformers.utils.logging.enable_explicit_format()
  # Set the verbosity to info of the Transformers logger (on main process only):
  if is_main_process(training_args.local_rank ):
    transformers.utils.logging.set_verbosity_info()
  logger.info("""Training/evaluation parameters %s"""  ,		A_ )

  # Set seed
  set_seed(training_args.seed )

  # Load pretrained model and tokenizer
  #
  # Distributed training:
  # The .from_pretrained methods guarantee that only one local process can concurrently
  # download model & vocab.

  a_   :						Dict   =				AutoConfig.from_pretrained(
      model_args.config_name if model_args.config_name else model_args.model_name_or_path  ,		cache_dir=model_args.cache_dir  ,		)

  a_   :						int   =				("""encoder_layerdrop""", """decoder_layerdrop""", """dropout""", """attention_dropout""")
  for p in extra_model_params:
    if getattr(A_  ,		A_  ,		A_ ):
      assert hasattr(A_  ,		A_ ), f'''({config.__class__.__name__}) doesn\'t have a `{p}` attribute'''
      setattr(A_  ,		A_  ,		getattr(A_  ,		A_ ) )

  a_   :						int   =				AutoTokenizer.from_pretrained(
      model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path  ,		cache_dir=model_args.cache_dir  ,		)
  a_   :						int   =				AutoModelForSeqaSeqLM.from_pretrained(
      model_args.model_name_or_path  ,		from_tf=""".ckpt""" in model_args.model_name_or_path  ,		config=A_  ,		cache_dir=model_args.cache_dir  ,		)

  # use task specific params
  use_task_specific_params(A_  ,		data_args.task )

  # set num_beams for evaluation
  if data_args.eval_beams is None:
    a_   :						int   =				model.config.num_beams

  # set decoder_start_token_id for MBart
  if model.config.decoder_start_token_id is None and isinstance(A_  ,		(MBartTokenizer, MBartTokenizerFast) ):
    assert (
        data_args.tgt_lang is not None and data_args.src_lang is not None
    ), "mBart requires --tgt_lang and --src_lang"
    if isinstance(A_  ,		A_ ):
      a_   :						Optional[int]   =				tokenizer.lang_code_to_id[data_args.tgt_lang]
    else:
      a_   :						Union[str, Any]   =				tokenizer.convert_tokens_to_ids(data_args.tgt_lang )

  if model_args.freeze_embeds:
    freeze_embeds(A_ )
  if model_args.freeze_encoder:
    freeze_params(model.get_encoder() )
    assert_all_frozen(model.get_encoder() )

  a_   :						List[Any]   =				SeqaSeqDataset

  # Get datasets
  a_   :						Optional[int]   =				(
      dataset_class(
          A_  ,		type_path="""train"""  ,		data_dir=data_args.data_dir  ,		n_obs=data_args.n_train  ,		max_target_length=data_args.max_target_length  ,		max_source_length=data_args.max_source_length  ,		prefix=model.config.prefix or """"""  ,		)
      if training_args.do_train
      else None
  )
  a_   :						Any   =				(
      dataset_class(
          A_  ,		type_path="""val"""  ,		data_dir=data_args.data_dir  ,		n_obs=data_args.n_val  ,		max_target_length=data_args.val_max_target_length  ,		max_source_length=data_args.max_source_length  ,		prefix=model.config.prefix or """"""  ,		)
      if training_args.do_eval or training_args.evaluation_strategy != EvaluationStrategy.NO
      else None
  )
  a_   :						Optional[Any]   =				(
      dataset_class(
          A_  ,		type_path="""test"""  ,		data_dir=data_args.data_dir  ,		n_obs=data_args.n_test  ,		max_target_length=data_args.test_max_target_length  ,		max_source_length=data_args.max_source_length  ,		prefix=model.config.prefix or """"""  ,		)
      if training_args.do_predict
      else None
  )

  # Initialize our Trainer
  a_   :						Union[str, Any]   =				(
      build_compute_metrics_fn(data_args.task  ,		A_ ) if training_args.predict_with_generate else None
  )
  a_   :						List[Any]   =				SeqaSeqTrainer(
      model=A_  ,		args=A_  ,		data_args=A_  ,		train_dataset=A_  ,		eval_dataset=A_  ,		data_collator=SeqaSeqDataCollator(
          A_  ,		A_  ,		model.config.decoder_start_token_id  ,		training_args.tpu_num_cores )  ,		compute_metrics=A_  ,		tokenizer=A_  ,		)

  a_   :						Optional[Any]   =				{}
  # Training
  if training_args.do_train:
    logger.info("""*** Train ***""" )

    a_   :						int   =				trainer.train(
        model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path ) else None )
    a_   :						Dict   =				train_result.metrics
    a_   :						Tuple   =				data_args.n_train

    trainer.save_model()  # this also saves the tokenizer

    if trainer.is_world_process_zero():
      handle_metrics("""train"""  ,		A_  ,		training_args.output_dir )
      all_metrics.update(A_ )

      # Need to save the state, since Trainer.save_model saves only the tokenizer with the model
      trainer.state.save_to_json(os.path.join(training_args.output_dir  ,		"""trainer_state.json""" ) )

      # For convenience, we also re-save the tokenizer to the same directory,
      # so that you can share your model easily on huggingface.co/models =)
      tokenizer.save_pretrained(training_args.output_dir )

    # Evaluation
  if training_args.do_eval:
    logger.info("""*** Evaluate ***""" )

    a_   :						Any   =				trainer.evaluate(metric_key_prefix="""val""" )
    a_   :						Tuple   =				data_args.n_val
    a_   :						Union[str, Any]   =				round(metrics["""val_loss"""]  ,		4 )

    if trainer.is_world_process_zero():
      handle_metrics("""val"""  ,		A_  ,		training_args.output_dir )
      all_metrics.update(A_ )

  if training_args.do_predict:
    logger.info("""*** Predict ***""" )

    a_   :						str   =				trainer.predict(test_dataset=A_  ,		metric_key_prefix="""test""" )
    a_   :						int   =				test_output.metrics
    a_   :						List[str]   =				data_args.n_test

    if trainer.is_world_process_zero():
      a_   :						Optional[Any]   =				round(metrics["""test_loss"""]  ,		4 )
      handle_metrics("""test"""  ,		A_  ,		training_args.output_dir )
      all_metrics.update(A_ )

      if training_args.predict_with_generate:
        a_   :						Union[str, Any]   =				tokenizer.batch_decode(
            test_output.predictions  ,		skip_special_tokens=A_  ,		clean_up_tokenization_spaces=A_ )
        a_   :						Optional[Any]   =				lmap(str.strip  ,		A_ )
        write_txt_file(A_  ,		os.path.join(training_args.output_dir  ,		"""test_generations.txt""" ) )

  if trainer.is_world_process_zero():
    save_json(A_  ,		os.path.join(training_args.output_dir  ,		"""all_results.json""" ) )

  return all_metrics






def 						_snake_case    (   A_							:     int ):
  """simple docstring"""







  main()


if __name__ == "__main__":
    main()



"""simple docstring"""



from __future__ import annotations

import inspect
import unittest
from typing import List, Tuple

from transformers import RegNetConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
				import tensorflow as tf

				from transformers import TF_REGNET_PRETRAINED_MODEL_ARCHIVE_LIST, TFRegNetForImageClassification, TFRegNetModel


if is_vision_available():
				from PIL import Image

				from transformers import AutoImageProcessor
class 						__UpperCAmelCase    :







			'''simple docstring'''





			def __init__(			self   ,      snake_case_   ,      snake_case_=3   ,      snake_case_=32   ,      snake_case_=3   ,      snake_case_=10   ,      snake_case_=[10, 20, 30, 40]   ,      snake_case_=[1, 1, 2, 1]   ,      snake_case_=True   ,      snake_case_=True   ,      snake_case_="relu"   ,      snake_case_=3   ,      snake_case_=None   ,      ):







										'''simple docstring'''


										A__       :     Tuple								=			parent
										A__       :     int								=			batch_size
										A__       :     List[Any]								=			image_size
										A__       :     Any								=			num_channels
										A__       :     Tuple								=			embeddings_size
										A__       :     Union[str, Any]								=			hidden_sizes
										A__       :     Tuple								=			depths
										A__       :     Union[str, Any]								=			is_training
										A__       :     Any								=			use_labels
										A__       :     Optional[Any]								=			hidden_act
										A__       :     str								=			num_labels
										A__       :     Optional[Any]								=			scope
										A__       :     List[Any]								=			len(snake_case_							)





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     List[Any]								=			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

										A__       :     str								=			None
										if self.use_labels:
																	A__       :     Tuple								=			ids_tensor([self.batch_size]   ,      self.num_labels							)

										A__       :     Dict								=			self.get_config()
										return config, pixel_values, labels





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										return RegNetConfig(
										    num_channels=self.num_channels   ,      embeddings_size=self.embeddings_size   ,      hidden_sizes=self.hidden_sizes   ,      depths=self.depths   ,      hidden_act=self.hidden_act   ,      num_labels=self.num_labels   ,      )





			def 		lowerCamelCase					(			self   ,      snake_case_   ,      snake_case_   ,      snake_case_							):







										'''simple docstring'''


										A__       :     Union[str, Any]								=			TFRegNetModel(config=snake_case_							)
										A__       :     Optional[int]								=			model(snake_case_   ,      training=snake_case_							)
										# 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 		lowerCamelCase					(			self   ,      snake_case_   ,      snake_case_   ,      snake_case_							):







										'''simple docstring'''


										A__       :     Optional[int]								=			self.num_labels
										A__       :     Tuple								=			TFRegNetForImageClassification(snake_case_							)
										A__       :     Optional[int]								=			model(snake_case_   ,      labels=snake_case_   ,      training=snake_case_							)
										self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.num_labels)							)





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


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





@require_tf
class 						__UpperCAmelCase    (__A							,   __A							,   unittest.TestCase     ):







			'''simple docstring'''
			_UpperCamelCase  :						str					     =				(TFRegNetModel, TFRegNetForImageClassification) if is_tf_available() else ()
			_UpperCamelCase  :						Tuple					     =				(
			    {'feature-extraction': TFRegNetModel, 'image-classification': TFRegNetForImageClassification}
			    if is_tf_available()
			    else {}
			)

			_UpperCamelCase  :						str					     =				False
			_UpperCamelCase  :						int					     =				False
			_UpperCamelCase  :						Tuple					     =				False
			_UpperCamelCase  :						Any					     =				False
			_UpperCamelCase  :						Union[str, Any]					     =				False





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     Dict								=			TFRegNetModelTester(self							)
										A__       :     str								=			ConfigTester(self   ,      config_class=snake_case_   ,      has_text_modality=snake_case_							)





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										return





			@unittest.skip(reason="""RegNet does not use inputs_embeds"""							)
			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										pass





			@unittest.skipIf(
			    not is_tf_available() or len(tf.config.list_physical_devices("""GPU"""							)							) == 0   ,      reason="""TF does not support backprop for grouped convolutions on CPU."""   ,      )
			@slow
			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										super().test_keras_fit()





			@unittest.skip(reason="""RegNet does not support input and output embeddings"""							)
			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										pass





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__				,							A__       :     int								=			self.model_tester.prepare_config_and_inputs_for_common()

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

																	A__       :     Tuple								=			["""pixel_values"""]
																	self.assertListEqual(arg_names[:1]   ,      snake_case_							)





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     Union[str, Any]								=			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*snake_case_							)





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										def check_hidden_states_output(snake_case_   ,      snake_case_   ,      snake_case_							):
																	A__       :     str								=			model_class(snake_case_							)
																	A__       :     Union[str, Any]								=			model(**self._prepare_for_class(snake_case_   ,      snake_case_							)   ,      training=snake_case_							)

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

																	A__       :     Union[str, Any]								=			self.model_tester.num_stages
																	self.assertEqual(len(snake_case_							)   ,      expected_num_stages + 1							)

																	# RegNet'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 // 2, self.model_tester.image_size // 2]   ,      )

										A__				,							A__       :     int								=			self.model_tester.prepare_config_and_inputs_for_common()
										A__       :     Optional[int]								=			["""basic""", """bottleneck"""]
										for model_class in self.all_model_classes:
																	for layer_type in layers_type:
																								A__       :     Optional[Any]								=			layer_type
																								A__       :     Dict								=			True
																								check_hidden_states_output(snake_case_   ,      snake_case_   ,      snake_case_							)

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

																								check_hidden_states_output(snake_case_   ,      snake_case_   ,      snake_case_							)





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


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

										def check_equivalence(snake_case_   ,      snake_case_   ,      snake_case_   ,      snake_case_={}							):
																	A__       :     List[str]								=			model(snake_case_   ,      return_dict=snake_case_   ,      **snake_case_							)
																	A__       :     Union[str, Any]								=			model(snake_case_   ,      return_dict=snake_case_   ,      **snake_case_							).to_tuple()

																	def recursive_check(snake_case_   ,      snake_case_							):
																								if isinstance(snake_case_   ,      (List, Tuple)							):
																															for tuple_iterable_value, dict_iterable_value in zip(snake_case_   ,      snake_case_							):
																																						recursive_check(snake_case_   ,      snake_case_							)
																								elif tuple_object is None:
																															return
																								else:
																															self.assertTrue(
																															    all(tf.equal(snake_case_   ,      snake_case_							)							)   ,      msg=(
																															        """Tuple and dict output are not equal. Difference:"""
																															        F''' {tf.math.reduce_max(tf.abs(tuple_object - dict_object							)							)}'''
																															    )   ,      )

																								recursive_check(snake_case_   ,      snake_case_							)

										for model_class in self.all_model_classes:
																	A__       :     Any								=			model_class(snake_case_							)

																	A__       :     Dict								=			self._prepare_for_class(snake_case_   ,      snake_case_							)
																	A__       :     List[Any]								=			self._prepare_for_class(snake_case_   ,      snake_case_							)
																	check_equivalence(snake_case_   ,      snake_case_   ,      snake_case_							)

																	A__       :     Tuple								=			self._prepare_for_class(snake_case_   ,      snake_case_   ,      return_labels=snake_case_							)
																	A__       :     List[str]								=			self._prepare_for_class(snake_case_   ,      snake_case_   ,      return_labels=snake_case_							)
																	check_equivalence(snake_case_   ,      snake_case_   ,      snake_case_							)

																	A__       :     Optional[int]								=			self._prepare_for_class(snake_case_   ,      snake_case_							)
																	A__       :     List[str]								=			self._prepare_for_class(snake_case_   ,      snake_case_							)
																	check_equivalence(snake_case_   ,      snake_case_   ,      snake_case_   ,      {"""output_hidden_states""": True}							)

																	A__       :     Dict								=			self._prepare_for_class(snake_case_   ,      snake_case_   ,      return_labels=snake_case_							)
																	A__       :     Union[str, Any]								=			self._prepare_for_class(snake_case_   ,      snake_case_   ,      return_labels=snake_case_							)
																	check_equivalence(snake_case_   ,      snake_case_   ,      snake_case_   ,      {"""output_hidden_states""": True}							)





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     Tuple								=			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*snake_case_							)





			@slow
			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										for model_name in TF_REGNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	A__       :     Any								=			TFRegNetModel.from_pretrained(snake_case_							)
																	self.assertIsNotNone(snake_case_							)





def  _A(		):
							A__       :     Optional[int]								=			Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""							)
							return image



@require_tf
@require_vision
class 						__UpperCAmelCase    (unittest.TestCase     ):







			'''simple docstring'''





			@cached_property
			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										return (
										    AutoImageProcessor.from_pretrained(TF_REGNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]							)
										    if is_vision_available()
										    else None
										)





			@slow
			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     int								=			TFRegNetForImageClassification.from_pretrained(TF_REGNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]							)

										A__       :     str								=			self.default_image_processor
										A__       :     Any								=			prepare_img()
										A__       :     Tuple								=			image_processor(images=snake_case_   ,      return_tensors="""tf"""							)

										# forward pass
										A__       :     List[Any]								=			model(**snake_case_   ,      training=snake_case_							)

										# verify the logits
										A__       :     Tuple								=			tf.TensorShape((1, 1_000)							)
										self.assertEqual(outputs.logits.shape   ,      snake_case_							)

										A__       :     Tuple								=			tf.constant([-0.41_80, -1.50_51, -3.48_36]							)

										tf.debugging.assert_near(outputs.logits[0, :3]   ,      snake_case_   ,      atol=1E-4							)





