

'''simple docstring'''






from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import ResNetConfig
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 TFResNetForImageClassification, TFResNetModel
 from transformers.models.resnet.modeling_tf_resnet import TF_RESNET_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]					,lowercase__   :     Tuple					,lowercase__   :     List[str]=3					,lowercase__   :     Optional[int]=3_2					,lowercase__   :     Dict=3					,lowercase__   :     List[str]=1_0					,lowercase__   :     Dict=[1_0, 2_0, 3_0, 4_0]					,lowercase__   :     Optional[Any]=[1, 1, 2, 1]					,lowercase__   :     Tuple=True					,lowercase__   :     int=True					,lowercase__   :     Tuple="relu"					,lowercase__   :     List[Any]=3					,lowercase__   :     str=None					,):
   __lowercase							  = parent
   __lowercase							  = batch_size
   __lowercase							  = image_size
   __lowercase							  = num_channels
   __lowercase							  = embeddings_size
   __lowercase							  = hidden_sizes
   __lowercase							  = depths
   __lowercase							  = is_training
   __lowercase							  = use_labels
   __lowercase							  = hidden_act
   __lowercase							  = num_labels
   __lowercase							  = scope
   __lowercase							  = len(lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[Any]			):
   __lowercase							  = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]			)

   __lowercase							  = None
   if self.use_labels:
    __lowercase							  = ids_tensor([self.batch_size]					,self.num_labels			)

   __lowercase							  = self.get_config()

   return config, pixel_values, labels






  def 					SCREAMING_SNAKE_CASE      (      self   :     Tuple			):
   return ResNetConfig(
       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					,image_size=self.image_size					,)






  def 					SCREAMING_SNAKE_CASE      (      self   :     int					,lowercase__   :     Any					,lowercase__   :     List[str]					,lowercase__   :     Optional[Any]			):
   __lowercase							  = TFResNetModel(config=lowercase__			)
   __lowercase							  = model(lowercase__			)
   # 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 // 3_2, self.image_size // 3_2)					,)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Union[str, Any]					,lowercase__   :     Dict					,lowercase__   :     List[Any]					,lowercase__   :     Optional[Any]			):
   __lowercase							  = self.num_labels
   __lowercase							  = TFResNetForImageClassification(lowercase__			)
   __lowercase							  = model(lowercase__					,labels=lowercase__			)
   self.parent.assertEqual(result.logits.shape					,(self.batch_size, self.num_labels)			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[Any]			):
   __lowercase							  = self.prepare_config_and_inputs()
   __lowercase  ,	__lowercase  ,	__lowercase							  = config_and_inputs
   __lowercase							  = {'''pixel_values''': pixel_values}
   return config, inputs_dict






@require_tf
class 			lowercase_	(lowerCamelCase__				,       lowerCamelCase__				,       unittest.TestCase      ):




  """simple docstring"""




  SCREAMING_SNAKE_CASE    :  Any							=			(TFResNetModel, TFResNetForImageClassification) if is_tf_available() else ()
  SCREAMING_SNAKE_CASE    :  int							=			(
      {'feature-extraction': TFResNetModel, 'image-classification': TFResNetForImageClassification}
      if is_tf_available()
      else {}
  )

  SCREAMING_SNAKE_CASE    :  int							=			False
  SCREAMING_SNAKE_CASE    :  Dict							=			False
  SCREAMING_SNAKE_CASE    :  str							=			False
  SCREAMING_SNAKE_CASE    :  Any							=			False
  SCREAMING_SNAKE_CASE    :  List[Any]							=			False






  def 					SCREAMING_SNAKE_CASE      (      self   :     List[Any]			):
   __lowercase							  = TFResNetModelTester(self			)
   __lowercase							  = ConfigTester(self					,config_class=lowercase__					,has_text_modality=lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Any			):
   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 					SCREAMING_SNAKE_CASE      (      self   :     Dict			):
   return






  @unittest.skip(reason='''ResNet does not use inputs_embeds'''			)
  def 					SCREAMING_SNAKE_CASE      (      self   :     int			):
   pass






  @unittest.skip(reason='''ResNet does not support input and output embeddings'''			)
  def 					SCREAMING_SNAKE_CASE      (      self   :     Tuple			):
   pass






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[int]			):
   __lowercase  ,	__lowercase							  = self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    __lowercase							  = model_class(lowercase__			)
    __lowercase							  = inspect.signature(model.call			)
    # signature.parameters is an OrderedDict => so arg_names order is deterministic
    __lowercase							  = [*signature.parameters.keys()]

    __lowercase							  = ['''pixel_values''']
    self.assertListEqual(arg_names[:1]					,lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     List[str]			):
   __lowercase							  = self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Tuple			):
   def check_hidden_states_output(lowercase__   :     Optional[int]					,lowercase__   :     Any					,lowercase__   :     str			):
    __lowercase							  = model_class(lowercase__			)
    __lowercase							  = model(**self._prepare_for_class(lowercase__					,lowercase__			)			)

    __lowercase							  = outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

    __lowercase							  = self.model_tester.num_stages
    self.assertEqual(len(lowercase__			)					,expected_num_stages + 1			)

    # ResNet'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]					,)

   __lowercase  ,	__lowercase							  = self.model_tester.prepare_config_and_inputs_for_common()
   __lowercase							  = ['''basic''', '''bottleneck''']
   for model_class in self.all_model_classes:
    for layer_type in layers_type:
     __lowercase							  = layer_type
     __lowercase							  = True
     check_hidden_states_output(lowercase__					,lowercase__					,lowercase__			)

     # check that output_hidden_states also work using config
     del inputs_dict["output_hidden_states"]
     __lowercase							  = True

     check_hidden_states_output(lowercase__					,lowercase__					,lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     List[str]			):
   __lowercase							  = self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_image_classification(*lowercase__			)






  @slow
  def 					SCREAMING_SNAKE_CASE      (      self   :     Dict			):
   for model_name in TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
    __lowercase							  = TFResNetModel.from_pretrained(lowercase__			)
    self.assertIsNotNone(lowercase__			)


def 		_A     (    ):






 """simple docstring"""



 __lowercase							  = Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''      )
 return image






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




  """simple docstring"""




  @cached_property
  def 					SCREAMING_SNAKE_CASE      (      self   :     List[str]			):
   return (
       AutoImageProcessor.from_pretrained(TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]			)
       if is_vision_available()
       else None
   )






  @slow
  def 					SCREAMING_SNAKE_CASE      (      self   :     int			):
   __lowercase							  = TFResNetForImageClassification.from_pretrained(TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]			)

   __lowercase							  = self.default_image_processor
   __lowercase							  = prepare_img()
   __lowercase							  = image_processor(images=lowercase__					,return_tensors='''tf'''			)

   # forward pass
   __lowercase							  = model(**lowercase__			)

   # verify the logits
   __lowercase							  = tf.TensorShape((1, 1_0_0_0)			)
   self.assertEqual(outputs.logits.shape					,lowercase__			)

   __lowercase							  = tf.constant([-1_1.1_0_6_9, -9.7_8_7_7, -8.3_7_7_7]			)

   self.assertTrue(np.allclose(outputs.logits[0, :3].numpy()					,lowercase__					,atol=1e-4			)			)




'''simple docstring'''






import torch
from torch import nn






class 			lowercase_	(nn.Module      ):




  """simple docstring"""




  def __init__(      self   :     Optional[int]					,lowercase__   :     List[str]					,lowercase__   :     Any					,lowercase__   :     Union[str, Any]					,lowercase__   :     Optional[int]					,lowercase__   :     Optional[int]=1					,lowercase__   :     Optional[Any]=False			):
   super().__init__()

   __lowercase							  = n_token
   __lowercase							  = d_embed
   __lowercase							  = d_proj

   __lowercase							  = cutoffs + [n_token]
   __lowercase							  = [0] + self.cutoffs
   __lowercase							  = div_val

   __lowercase							  = self.cutoffs[0]
   __lowercase							  = len(self.cutoffs			) - 1
   __lowercase							  = self.shortlist_size + self.n_clusters

   if self.n_clusters > 0:
    __lowercase							  = nn.Parameter(torch.zeros(self.n_clusters					,self.d_embed			)			)
    __lowercase							  = nn.Parameter(torch.zeros(self.n_clusters			)			)

   __lowercase							  = nn.ModuleList()
   __lowercase							  = nn.ParameterList()

   if div_val == 1:
    for i in range(len(self.cutoffs			)			):
     if d_proj != d_embed:
      self.out_projs.append(nn.Parameter(torch.FloatTensor(lowercase__					,lowercase__			)			)			)
     else:
      self.out_projs.append(lowercase__			)

    self.out_layers.append(nn.Linear(lowercase__					,lowercase__			)			)
   else:
    for i in range(len(self.cutoffs			)			):
     __lowercase  ,	__lowercase							  = self.cutoff_ends[i], self.cutoff_ends[i + 1]
     __lowercase							  = d_embed // (div_val**i)

     self.out_projs.append(nn.Parameter(torch.FloatTensor(lowercase__					,lowercase__			)			)			)

     self.out_layers.append(nn.Linear(lowercase__					,r_idx - l_idx			)			)

   __lowercase							  = keep_order






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[int]					,lowercase__   :     List[Any]					,lowercase__   :     Dict					,lowercase__   :     List[Any]					,lowercase__   :     Any			):
   if proj is None:
    __lowercase							  = nn.functional.linear(lowercase__					,lowercase__					,bias=lowercase__			)
   else:
    # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
    __lowercase							  = nn.functional.linear(lowercase__					,proj.t().contiguous()			)
    __lowercase							  = nn.functional.linear(lowercase__					,lowercase__					,bias=lowercase__			)
    # else:
    #     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
    #     if bias is not None:
    #         logit = logit + bias

   return logit






  def 					SCREAMING_SNAKE_CASE      (      self   :     Union[str, Any]					,lowercase__   :     Optional[Any]					,lowercase__   :     Any=None					,lowercase__   :     List[str]=False			):

   if labels is not None:
    # Shift so that tokens < n predict n
    __lowercase							  = hidden[..., :-1, :].contiguous()
    __lowercase							  = labels[..., 1:].contiguous()
    __lowercase							  = hidden.view(-1					,hidden.size(-1			)			)
    __lowercase							  = labels.view(-1			)
    if hidden.size(0			) != labels.size(0			):
     raise RuntimeError('''Input and labels should have the same size in the batch dimension.'''			)
   else:
    __lowercase							  = hidden.view(-1					,hidden.size(-1			)			)

   if self.n_clusters == 0:
    __lowercase							  = self._compute_logit(lowercase__					,self.out_layers[0].weight					,self.out_layers[0].bias					,self.out_projs[0]			)
    if labels is not None:
     __lowercase							  = labels != -1_0_0
     __lowercase							  = torch.zeros_like(lowercase__					,dtype=hidden.dtype					,device=hidden.device			)
     __lowercase							  = (
         -nn.functional.log_softmax(lowercase__					,dim=-1			)[mask].gather(1					,labels[mask].unsqueeze(1			)			).squeeze(1			)
     )
    else:
     __lowercase							  = nn.functional.log_softmax(lowercase__					,dim=-1			)
   else:
    # construct weights and biases
    __lowercase  ,	__lowercase							  = [], []
    for i in range(len(self.cutoffs			)			):
     if self.div_val == 1:
      __lowercase  ,	__lowercase							  = self.cutoff_ends[i], self.cutoff_ends[i + 1]
      __lowercase							  = self.out_layers[0].weight[l_idx:r_idx]
      __lowercase							  = self.out_layers[0].bias[l_idx:r_idx]
     else:
      __lowercase							  = self.out_layers[i].weight
      __lowercase							  = self.out_layers[i].bias

     if i == 0:
      __lowercase							  = torch.cat([weight_i, self.cluster_weight]					,dim=0			)
      __lowercase							  = torch.cat([bias_i, self.cluster_bias]					,dim=0			)

     weights.append(lowercase__			)
     biases.append(lowercase__			)

    __lowercase  ,	__lowercase  ,	__lowercase							  = weights[0], biases[0], self.out_projs[0]

    __lowercase							  = self._compute_logit(lowercase__					,lowercase__					,lowercase__					,lowercase__			)
    __lowercase							  = nn.functional.log_softmax(lowercase__					,dim=1			)

    if labels is None:
     __lowercase							  = hidden.new_empty((head_logit.size(0			), self.n_token)			)
    else:
     __lowercase							  = torch.zeros_like(lowercase__					,dtype=hidden.dtype					,device=hidden.device			)

    __lowercase							  = 0
    __lowercase							  = [0] + self.cutoffs
    for i in range(len(lowercase__			) - 1			):
     __lowercase  ,	__lowercase							  = cutoff_values[i], cutoff_values[i + 1]

     if labels is not None:
      __lowercase							  = (labels >= l_idx) & (labels < r_idx)
      __lowercase							  = mask_i.nonzero().squeeze()

      if indices_i.numel() == 0:
       continue

      __lowercase							  = labels.index_select(0					,lowercase__			) - l_idx
      __lowercase							  = head_logprob.index_select(0					,lowercase__			)
      __lowercase							  = hidden.index_select(0					,lowercase__			)
     else:
      __lowercase							  = hidden

     if i == 0:
      if labels is not None:
       __lowercase							  = head_logprob_i.gather(1					,target_i[:, None]			).squeeze(1			)
      else:
       __lowercase							  = head_logprob[:, : self.cutoffs[0]]
     else:
      __lowercase  ,	__lowercase  ,	__lowercase							  = weights[i], biases[i], self.out_projs[i]

      __lowercase							  = self._compute_logit(lowercase__					,lowercase__					,lowercase__					,lowercase__			)
      __lowercase							  = nn.functional.log_softmax(lowercase__					,dim=1			)
      __lowercase							  = self.cutoffs[0] + i - 1  # No probability for the head cluster
      if labels is not None:
       __lowercase							  = head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
           1					,target_i[:, None]			).squeeze(1			)
      else:
       __lowercase							  = head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
       __lowercase							  = logprob_i

     if labels is not None:
      if (hasattr(self					,'''keep_order'''			) and self.keep_order) or keep_order:
       out.index_copy_(0					,lowercase__					,-logprob_i			)
      else:
       out[offset : offset + logprob_i.size(0			)].copy_(-logprob_i			)
      offset += logprob_i.size(0			)

   return out






  def 					SCREAMING_SNAKE_CASE      (      self   :     Any					,lowercase__   :     Union[str, Any]			):
   if self.n_clusters == 0:
    __lowercase							  = self._compute_logit(lowercase__					,self.out_layers[0].weight					,self.out_layers[0].bias					,self.out_projs[0]			)
    return nn.functional.log_softmax(lowercase__					,dim=-1			)
   else:
    # construct weights and biases
    __lowercase  ,	__lowercase							  = [], []
    for i in range(len(self.cutoffs			)			):
     if self.div_val == 1:
      __lowercase  ,	__lowercase							  = self.cutoff_ends[i], self.cutoff_ends[i + 1]
      __lowercase							  = self.out_layers[0].weight[l_idx:r_idx]
      __lowercase							  = self.out_layers[0].bias[l_idx:r_idx]
     else:
      __lowercase							  = self.out_layers[i].weight
      __lowercase							  = self.out_layers[i].bias

     if i == 0:
      __lowercase							  = torch.cat([weight_i, self.cluster_weight]					,dim=0			)
      __lowercase							  = torch.cat([bias_i, self.cluster_bias]					,dim=0			)

     weights.append(lowercase__			)
     biases.append(lowercase__			)

    __lowercase  ,	__lowercase  ,	__lowercase							  = weights[0], biases[0], self.out_projs[0]
    __lowercase							  = self._compute_logit(lowercase__					,lowercase__					,lowercase__					,lowercase__			)

    __lowercase							  = hidden.new_empty((head_logit.size(0			), self.n_token)			)
    __lowercase							  = nn.functional.log_softmax(lowercase__					,dim=1			)

    __lowercase							  = [0] + self.cutoffs
    for i in range(len(lowercase__			) - 1			):
     __lowercase  ,	__lowercase							  = cutoff_values[i], cutoff_values[i + 1]

     if i == 0:
      __lowercase							  = head_logprob[:, : self.cutoffs[0]]
     else:
      __lowercase  ,	__lowercase  ,	__lowercase							  = weights[i], biases[i], self.out_projs[i]

      __lowercase							  = self._compute_logit(lowercase__					,lowercase__					,lowercase__					,lowercase__			)
      __lowercase							  = nn.functional.log_softmax(lowercase__					,dim=1			)

      __lowercase							  = head_logprob[:, -i] + tail_logprob_i
      __lowercase							  = logprob_i

    return out

