code: 
  
  
  import argparse
  import json
  
  import torch
  from huggingface_hub import hf_hub_download
  from PIL import Image
  from timm.models import create_model
  
  from transformers import (
      BeitImageProcessor,
      DataaVecVisionConfig,
      DataaVecVisionForImageClassification,
      DataaVecVisionModel,
  )
  
  
  
  
  def       a__						(		snake_case__				:       str	,      snake_case__				:       Any=False	,      snake_case__				:       Union[str, Any]=False	,      snake_case__				:       Dict="data2vec."    ):
  		_UpperCAmelCase		:       Tuple       =    """backbone.""" if is_semantic else """"""
  
  		_UpperCAmelCase		:       Tuple       =    []
  		for i in range(config.num_hidden_layers    ):
  				# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
  				rename_keys.append(
  				    (f'''{prefix}blocks.{i}.norm1.weight''', f'''{hf_prefix}encoder.layer.{i}.layernorm_before.weight''')    )
  				rename_keys.append((f'''{prefix}blocks.{i}.norm1.bias''', f'''{hf_prefix}encoder.layer.{i}.layernorm_before.bias''')    )
  				rename_keys.append(
  				    (f'''{prefix}blocks.{i}.attn.proj.weight''', f'''{hf_prefix}encoder.layer.{i}.attention.output.dense.weight''')    )
  				rename_keys.append(
  				    (f'''{prefix}blocks.{i}.attn.proj.bias''', f'''{hf_prefix}encoder.layer.{i}.attention.output.dense.bias''')    )
  				rename_keys.append(
  				    (f'''{prefix}blocks.{i}.norm2.weight''', f'''{hf_prefix}encoder.layer.{i}.layernorm_after.weight''')    )
  				rename_keys.append((f'''{prefix}blocks.{i}.norm2.bias''', f'''{hf_prefix}encoder.layer.{i}.layernorm_after.bias''')    )
  				rename_keys.append(
  				    (f'''{prefix}blocks.{i}.mlp.fc1.weight''', f'''{hf_prefix}encoder.layer.{i}.intermediate.dense.weight''')    )
  				rename_keys.append(
  				    (f'''{prefix}blocks.{i}.mlp.fc1.bias''', f'''{hf_prefix}encoder.layer.{i}.intermediate.dense.bias''')    )
  				rename_keys.append((f'''{prefix}blocks.{i}.mlp.fc2.weight''', f'''{hf_prefix}encoder.layer.{i}.output.dense.weight''')    )
  				rename_keys.append((f'''{prefix}blocks.{i}.mlp.fc2.bias''', f'''{hf_prefix}encoder.layer.{i}.output.dense.bias''')    )
  
  		# projection layer + position embeddings
  		rename_keys.extend(
  		    [
  		        (f'''{prefix}cls_token''', f'''{hf_prefix}embeddings.cls_token'''),
  		        (f'''{prefix}patch_embed.proj.weight''', f'''{hf_prefix}embeddings.patch_embeddings.projection.weight'''),
  		        (f'''{prefix}patch_embed.proj.bias''', f'''{hf_prefix}embeddings.patch_embeddings.projection.bias'''),
  		    ]    )
  
  		if has_lm_head:
  				# mask token + shared relative position bias + layernorm
  				rename_keys.extend(
  				    [
  				        ("""mask_token""", f'''{hf_prefix}embeddings.mask_token'''),
  				        (
  				            """rel_pos_bias.relative_position_bias_table""",
  				            f'''{hf_prefix}encoder.relative_position_bias.relative_position_bias_table''',
  				        ),
  				        (
  				            """rel_pos_bias.relative_position_index""",
  				            f'''{hf_prefix}encoder.relative_position_bias.relative_position_index''',
  				        ),
  				        ("""norm.weight""", """layernorm.weight"""),
  				        ("""norm.bias""", """layernorm.bias"""),
  				    ]    )
  		elif is_semantic:
  				# semantic segmentation classification heads
  				rename_keys.extend(
  				    [
  				        ("""decode_head.conv_seg.weight""", """decode_head.classifier.weight"""),
  				        ("""decode_head.conv_seg.bias""", """decode_head.classifier.bias"""),
  				        ("""auxiliary_head.conv_seg.weight""", """auxiliary_head.classifier.weight"""),
  				        ("""auxiliary_head.conv_seg.bias""", """auxiliary_head.classifier.bias"""),
  				    ]    )
  		else:
  				# layernorm + classification head
  				rename_keys.extend(
  				    [
  				        ("""fc_norm.weight""", f'''{hf_prefix}pooler.layernorm.weight'''),
  				        ("""fc_norm.bias""", f'''{hf_prefix}pooler.layernorm.bias'''),
  				        ("""head.weight""", """classifier.weight"""),
  				        ("""head.bias""", """classifier.bias"""),
  				    ]    )
  
  		return rename_keys
  
  
  
  
  def       a__						(		snake_case__				:       Tuple	,      snake_case__				:       str	,      snake_case__				:       Dict=False	,      snake_case__				:       List[str]=False	,      snake_case__				:       Union[str, Any]="data2vec_vision."    ):
  		for i in range(config.num_hidden_layers    ):
  				_UpperCAmelCase		:       int       =    """backbone.""" if is_semantic else """"""
  				# queries, keys and values
  				_UpperCAmelCase		:       Any       =    state_dict.pop(f'''{prefix}blocks.{i}.attn.qkv.weight'''    )
  				_UpperCAmelCase		:       Union[str, Any]       =    state_dict.pop(f'''{prefix}blocks.{i}.attn.q_bias'''    )
  				_UpperCAmelCase		:       List[str]       =    state_dict.pop(f'''{prefix}blocks.{i}.attn.v_bias'''    )
  
  				_UpperCAmelCase		:       Optional[int]       =    in_proj_weight[
  				    : config.hidden_size, :
  				]
  				_UpperCAmelCase		:       Tuple       =    q_bias
  				_UpperCAmelCase		:       Union[str, Any]       =    in_proj_weight[
  				    config.hidden_size : config.hidden_size * 2, :
  				]
  				_UpperCAmelCase		:       Optional[Any]       =    in_proj_weight[
  				    -config.hidden_size :, :
  				]
  				_UpperCAmelCase		:       Any       =    v_bias
  
  				# gamma_1 and gamma_2
  				# we call them lambda because otherwise they are renamed when using .from_pretrained
  				_UpperCAmelCase		:       List[Any]       =    state_dict.pop(f'''{prefix}blocks.{i}.gamma_1'''    )
  				_UpperCAmelCase		:       Optional[Any]       =    state_dict.pop(f'''{prefix}blocks.{i}.gamma_2'''    )
  
  				_UpperCAmelCase		:       Any       =    gamma_a
  				_UpperCAmelCase		:       Optional[Any]       =    gamma_a
  
  				# relative_position bias table + index
  				if not has_lm_head:
  						# each layer has its own relative position bias
  						_UpperCAmelCase		:       List[Any]       =    state_dict.pop(f'''{prefix}blocks.{i}.attn.relative_position_bias_table'''    )
  						_UpperCAmelCase		:       Optional[Any]       =    state_dict.pop(f'''{prefix}blocks.{i}.attn.relative_position_index'''    )
  
  						_UpperCAmelCase		:       str       =    table
  						_UpperCAmelCase		:       str       =    index
  
  
  
  
  def       a__						(		):
  		_UpperCAmelCase		:       Union[str, Any]       =    argparse.ArgumentParser(
  		    """Convert Data2VecVision to HF for image classification and pretraining"""	,      add_help=snake_case__    )
  		parser.add_argument("""--hf_checkpoint_name"""	,      type=snake_case__    )
  		parser.add_argument("""--input_size"""	,      default=224	,      type=snake_case__	,      help="""images input size"""    )
  		parser.add_argument("""--beit_checkpoint"""	,      default=""""""	,      help="""beit checkpoint"""    )
  
  		return parser.parse_args()
  
  
  
  
  def       a__						(		snake_case__				:       List[str]	,      snake_case__				:       Any	,      snake_case__				:       Dict    ):
  		def load_state_dict(snake_case__				:       List[Any]	,      snake_case__				:       Any	,      snake_case__				:       Optional[int]=""	,      snake_case__				:       List[Any]="relative_position_index"    ):
  				_UpperCAmelCase		:       Optional[int]       =    []
  				_UpperCAmelCase		:       List[Any]       =    []
  				_UpperCAmelCase		:       Dict       =    []
  				# copy state_dict so _load_from_state_dict can modify it
  				_UpperCAmelCase		:       Optional[int]       =    getattr(snake_case__	,      """_metadata"""	,      snake_case__    )
  				_UpperCAmelCase		:       Dict       =    state_dict.copy()
  				if metadata is not None:
  						_UpperCAmelCase		:       Optional[int]       =    metadata
  
  				def load(snake_case__				:       Dict	,      snake_case__				:       List[Any]=""    ):
  						_UpperCAmelCase		:       Tuple       =    {} if metadata is None else metadata.get(prefix[:-1]	,      {}    )
  						module._load_from_state_dict(
  						    snake_case__	,      snake_case__	,      snake_case__	,      snake_case__	,      snake_case__	,      snake_case__	,      snake_case__    )
  						for name, child in module._modules.items():
  								if child is not None:
  										load(snake_case__	,      prefix + name + """."""    )
  
  				load(snake_case__	,      prefix=snake_case__    )
  
  				_UpperCAmelCase		:       Union[str, Any]       =    []
  				_UpperCAmelCase		:       Optional[Any]       =    []
  				for key in missing_keys:
  						_UpperCAmelCase		:       str       =    True
  						for ignore_key in ignore_missing.split("""|"""    ):
  								if ignore_key in key:
  										_UpperCAmelCase		:       Optional[int]       =    False
  										break
  						if keep_flag:
  								warn_missing_keys.append(snake_case__    )
  						else:
  								ignore_missing_keys.append(snake_case__    )
  
  				_UpperCAmelCase		:       str       =    warn_missing_keys
  
  				if len(snake_case__    ) > 0:
  						print(
  						    """Weights of {} not initialized from pretrained model: {}""".format(
  						        model.__class__.__name__	,      snake_case__    )    )
  				if len(snake_case__    ) > 0:
  						print("""Weights from pretrained model not used in {}: {}""".format(model.__class__.__name__	,      snake_case__    )    )
  				if len(snake_case__    ) > 0:
  						print(
  						    """Ignored weights of {} not initialized from pretrained model: {}""".format(
  						        model.__class__.__name__	,      snake_case__    )    )
  				if len(snake_case__    ) > 0:
  						print("""\n""".join(snake_case__    )    )
  
  		_UpperCAmelCase		:       List[Any]       =    {
  		    """pretrained""": False,
  		    """use_shared_rel_pos_bias""": True,
  		    """use_abs_pos_emb""": False,
  		    """init_values""": 0.1,
  		}
  
  		if is_finetuned:
  				model_kwargs.update(
  				    {
  				        """num_classes""": 1000,
  				        """use_mean_pooling""": True,
  				        """init_scale""": 0.001,
  				        """use_rel_pos_bias""": True,
  				    }    )
  
  		_UpperCAmelCase		:       Dict       =    create_model(
  		    """beit_large_patch16_224""" if is_large else """beit_base_patch16_224"""	,      **snake_case__	,      )
  		_UpperCAmelCase		:       Optional[int]       =    model.patch_embed.patch_size
  		_UpperCAmelCase		:       Optional[int]       =    (args.input_size // patch_size[0], args.input_size // patch_size[1])
  		_UpperCAmelCase		:       Union[str, Any]       =    torch.load(args.beit_checkpoint	,      map_location="""cpu"""    )
  
  		print(f'''Load ckpt from {args.beit_checkpoint}'''    )
  		_UpperCAmelCase		:       Any       =    None
  		for model_key in ("model", "module"):
  				if model_key in checkpoint:
  						_UpperCAmelCase		:       List[str]       =    checkpoint[model_key]
  						print(f'''Load state_dict by model_key = {model_key}'''    )
  						break
  
  		_UpperCAmelCase		:       List[Any]       =    list(checkpoint_model.keys()    )
  		for key in all_keys:
  				if "relative_position_index" in key:
  						checkpoint_model.pop(snake_case__    )
  
  				if "relative_position_bias_table" in key:
  						_UpperCAmelCase		:       List[Any]       =    checkpoint_model[key]
  						_UpperCAmelCase		:       Tuple       =    rel_pos_bias.size()
  						_UpperCAmelCase		:       List[str]       =    model.state_dict()[key].size()
  						_UpperCAmelCase		:       Tuple       =    model.patch_embed.patch_shape
  						if dst_patch_shape[0] != dst_patch_shape[1]:
  								raise NotImplementedError()
  
  		load_state_dict(snake_case__	,      snake_case__	,      prefix=""""""    )
  
  		return model
  
  
  
  
  def       a__						(		):
  		_UpperCAmelCase		:       List[Any]       =    get_args()
  
  		_UpperCAmelCase		:       List[Any]       =    """ft1k""" in args.hf_checkpoint_name
  		_UpperCAmelCase		:       List[str]       =    """large""" in args.hf_checkpoint_name
  
  		if is_finetuned:
  				# To convert Beit's data2vec_vision to HF you need to copy
  				# https://github.com/facebookresearch/data2vec_vision/blob/main/beit/modeling_finetune.py
  				# into this folder.
  				import modeling_finetune  # noqa: F401
  		else:
  				# To convert Beit's data2vec_vision to HF you need to copy
  				# https://github.com/facebookresearch/data2vec_vision/blob/main/beit/modeling_cyclical.py
  				# into this folder
  				# IMPORTANT: Note that for now we've only converted the down-stream
  				# model and not the full pretrained model. This means for the integration
  				# test you need to add a `return x` after the following line:
  				# https://github.com/facebookresearch/data2vec_vision/blob/af9a36349aaed59ae66e69b5dabeef2d62fdc5da/beit/modeling_cyclical.py#L197
  				# to make the integration test pass.
  				import modeling_cyclical  # noqa: F401
  
  		# 1. Create model config
  		_UpperCAmelCase		:       str       =    DataaVecVisionConfig()
  		if is_finetuned:
  				_UpperCAmelCase		:       Union[str, Any]       =    True
  				_UpperCAmelCase		:       Optional[Any]       =    False
  				_UpperCAmelCase		:       Dict       =    True
  				_UpperCAmelCase		:       List[Any]       =    1000
  
  				_UpperCAmelCase		:       str       =    """huggingface/label-files"""
  				_UpperCAmelCase		:       Any       =    """imagenet-1k-id2label.json"""
  				_UpperCAmelCase		:       str       =    json.load(open(hf_hub_download(snake_case__	,      snake_case__	,      repo_type="""dataset"""    )	,      """r"""    )    )
  				_UpperCAmelCase		:       Any       =    {int(snake_case__    ): v for k, v in idalabel.items()}
  				_UpperCAmelCase		:       Optional[int]       =    idalabel
  				_UpperCAmelCase		:       Any       =    {v: k for k, v in idalabel.items()}
  		else:
  				_UpperCAmelCase		:       int       =    False
  				_UpperCAmelCase		:       Dict       =    True
  				_UpperCAmelCase		:       List[Any]       =    False
  
  		if is_large:
  				_UpperCAmelCase		:       Dict       =    1024
  				_UpperCAmelCase		:       Optional[int]       =    4096
  				_UpperCAmelCase		:       Dict       =    24
  				_UpperCAmelCase		:       Optional[Any]       =    16
  
  		# 2. Load Beit model
  		_UpperCAmelCase		:       List[str]       =    load_beit_model(snake_case__	,      snake_case__	,      snake_case__    )
  		orig_model.eval()
  
  		# 3. Forward Beit model
  		_UpperCAmelCase		:       Any       =    BeitImageProcessor(size=config.image_size	,      do_center_crop=snake_case__    )
  		_UpperCAmelCase		:       Union[str, Any]       =    Image.open("""../../../../tests/fixtures/tests_samples/COCO/000000039769.png"""    )
  		_UpperCAmelCase		:       Union[str, Any]       =    image_processor(images=snake_case__	,      return_tensors="""pt"""    )
  		_UpperCAmelCase		:       List[Any]       =    encoding["""pixel_values"""]
  
  		_UpperCAmelCase		:       int       =    (pixel_values,) if is_finetuned else (pixel_values, None)
  		with torch.no_grad():
  				_UpperCAmelCase		:       List[str]       =    orig_model(*snake_case__    )
  
  		# 4. Load HF Data2VecVision model
  		if is_finetuned:
  				_UpperCAmelCase		:       Optional[Any]       =    DataaVecVisionForImageClassification(snake_case__    )
  				hf_model.eval()
  				_UpperCAmelCase		:       Dict       =    False
  				_UpperCAmelCase		:       Any       =    """data2vec_vision."""
  		else:
  				_UpperCAmelCase		:       Union[str, Any]       =    DataaVecVisionModel(snake_case__    )
  				hf_model.eval()
  				_UpperCAmelCase		:       List[str]       =    True
  				_UpperCAmelCase		:       Tuple       =    """"""
  
  		_UpperCAmelCase		:       Optional[Any]       =    create_rename_keys(snake_case__	,      hf_prefix=snake_case__	,      has_lm_head=snake_case__    )
  		_UpperCAmelCase		:       List[str]       =    orig_model.state_dict()
  		for src, dest in rename_keys:
  				_UpperCAmelCase		:       str       =    state_dict.pop(snake_case__    )
  				_UpperCAmelCase		:       List[Any]       =    val
  
  		read_in_q_k_v(snake_case__	,      snake_case__	,      hf_prefix=snake_case__	,      has_lm_head=snake_case__    )
  		_UpperCAmelCase		:       Union[str, Any]       =    hf_model.load_state_dict(snake_case__	,      strict=snake_case__    )
  		print("""HF missing"""	,      snake_case__    )
  		print("""HF unexpected_keys"""	,      snake_case__    )
  
  		# 5. Forward HF Data2VecVision model
  		with torch.no_grad():
  				_UpperCAmelCase		:       List[str]       =    hf_model(snake_case__    )
  
  		_UpperCAmelCase		:       List[Any]       =    hf_model_output.logits if is_finetuned else hf_model_output.last_hidden_state
  
  		# 6. Compare
  		_UpperCAmelCase		:       Union[str, Any]       =    torch.max(torch.abs(hf_output - orig_model_output    )    ).item()
  
  		print(f'''max_absolute_diff = {max_absolute_diff}'''    )
  		_UpperCAmelCase		:       Tuple       =    torch.allclose(snake_case__	,      snake_case__	,      atol=1e-3    )
  		print("""Do both models output the same tensors?"""	,      """🔥""" if success else """💩"""    )
  		if not success:
  				raise Exception("""Something went wRoNg"""    )
  
  		# 7. Save
  		print(f'''Saving to {args.hf_checkpoint_name}'''    )
  		hf_model.save_pretrained(args.hf_checkpoint_name    )
  		image_processor.save_pretrained(args.hf_checkpoint_name    )
  
  
  if __name__ == "__main__":
  	main()
  	# Run the following to convert checkpoints
  	#  python ./convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.py \
  	#          --beit_checkpoint ./pretrained_base.pt \
  	#          --hf_checkpoint_name "./data2vec-vision-base"
  	#  python ./convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.py \
  	#          --beit_checkpoint ./finetuned_base.pt \
  	#          --hf_checkpoint_name "./data2vec-vision-base-ft1k"
  	#  python ./convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.py \
  	#          --beit_checkpoint ./pretrained_large.pt \
  	#          --hf_checkpoint_name "./data2vec-vision-large"
  	#  python ./convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.py \
  	#          --beit_checkpoint ./finetuned_large.pt \
  	#          --hf_checkpoint_name "./data2vec-vision-large-ft1k"
  
  
  
  
  
  
style_context: 
  
  
  import logging
  
  import torch
  from torch import nn
  from torch.nn import CrossEntropyLoss, MSELoss
  
  from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
  from transformers.models.albert.modeling_albert import (
      ALBERT_INPUTS_DOCSTRING,
      ALBERT_START_DOCSTRING,
      AlbertModel,
      AlbertPreTrainedModel,
      AlbertTransformer,
  )
  
  
  SCREAMING_SNAKE_CASE__				:    Dict         =    logging.getLogger(__name__)
  
  
  class 		_SCREAMING_SNAKE_CASE    (	A    ):
  
  
       def 					__snake_case(     self       ,				A_       ,				A_       ,				A_=None       ,				A_=None						):
         if current_layer == 0:
           _UpperCAmelCase		:       Any       =    self.embedding_hidden_mapping_in(A_						)
         else:
           _UpperCAmelCase		:       Optional[int]       =    hidden_states[0]
  
         _UpperCAmelCase		:       Tuple       =    int(self.config.num_hidden_layers / self.config.num_hidden_groups						)
  
         # Index of the hidden group
         _UpperCAmelCase		:       Union[str, Any]       =    int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups)						)
  
         _UpperCAmelCase		:       int       =    self.albert_layer_groups[group_idx](
             A_       ,				A_       ,				head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]       ,				)
         _UpperCAmelCase		:       Dict       =    layer_group_output[0]
  
         return (hidden_states,)
  
  
  @add_start_docstrings(
      '''The bare ALBERT Model transformer with PABEE outputting raw hidden-states without any specific head on top.'''       ,  A       ,  )
  class 		_SCREAMING_SNAKE_CASE    (	A    ):
  
  
       def __init__(     self       ,				A_						):
         super().__init__(A_						)
  
         _UpperCAmelCase		:       List[Any]       =    AlbertTransformerWithPabee(A_						)
  
         self.init_weights()
         _UpperCAmelCase		:       Tuple       =    0
         _UpperCAmelCase		:       Optional[Any]       =    0
         _UpperCAmelCase		:       int       =    0
  
         _UpperCAmelCase		:       Dict       =    0
  
  
       def 					__snake_case(     self       ,				A_						):
         _UpperCAmelCase		:       str       =    threshold
  
  
       def 					__snake_case(     self       ,				A_						):
         _UpperCAmelCase		:       Optional[Any]       =    patience
  
  
       def 					__snake_case(     self						):
         _UpperCAmelCase		:       Tuple       =    0
         _UpperCAmelCase		:       Tuple       =    0
  
  
       def 					__snake_case(     self						):
         _UpperCAmelCase		:       Dict       =    self.inference_layers_num / self.inference_instances_num
         _UpperCAmelCase		:       Optional[int]       =    (
             F'''*** Patience = {self.patience} Avg. Inference Layers = {avg_inf_layers:.2f} Speed Up ='''
             F''' {1 - avg_inf_layers / self.config.num_hidden_layers:.2f} ***'''
         )
         print(A_						)
  
  
  
  
       @add_start_docstrings_to_model_forward(A_						)
       def 					__snake_case(     self       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=False       ,				):
  
         if input_ids is not None and inputs_embeds is not None:
           raise ValueError("""You cannot specify both input_ids and inputs_embeds at the same time"""						)
         elif input_ids is not None:
           _UpperCAmelCase		:       List[str]       =    input_ids.size()
         elif inputs_embeds is not None:
           _UpperCAmelCase		:       Optional[Any]       =    inputs_embeds.size()[:-1]
         else:
           raise ValueError("""You have to specify either input_ids or inputs_embeds"""						)
  
         _UpperCAmelCase		:       List[str]       =    input_ids.device if input_ids is not None else inputs_embeds.device
  
         if attention_mask is None:
           _UpperCAmelCase		:       Any       =    torch.ones(A_       ,				device=A_						)
         if token_type_ids is None:
           _UpperCAmelCase		:       List[str]       =    torch.zeros(A_       ,				dtype=torch.long       ,				device=A_						)
  
         _UpperCAmelCase		:       str       =    attention_mask.unsqueeze(1						).unsqueeze(2						)
         _UpperCAmelCase		:       Tuple       =    extended_attention_mask.to(dtype=self.dtype						)  # fp16 compatibility
         _UpperCAmelCase		:       Tuple       =    (1.0 - extended_attention_mask) * -1_0_0_0_0.0
         _UpperCAmelCase		:       Union[str, Any]       =    self.get_head_mask(A_       ,				self.config.num_hidden_layers						)
  
         _UpperCAmelCase		:       Dict       =    self.embeddings(
             A_       ,				position_ids=A_       ,				token_type_ids=A_       ,				inputs_embeds=A_						)
         _UpperCAmelCase		:       str       =    embedding_output
  
         if self.training:
           _UpperCAmelCase		:       Any       =    []
           for i in range(self.config.num_hidden_layers						):
             _UpperCAmelCase		:       Tuple       =    self.encoder.adaptive_forward(
                 A_       ,				current_layer=A_       ,				attention_mask=A_       ,				head_mask=A_       ,				)
  
             _UpperCAmelCase		:       List[Any]       =    self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]						)						)
             _UpperCAmelCase		:       List[str]       =    output_layers[i](output_dropout(A_						)						)
             res.append(A_						)
         elif self.patience == 0:  # Use all layers for inference
           _UpperCAmelCase		:       Optional[int]       =    self.encoder(A_       ,				A_       ,				head_mask=A_						)
           _UpperCAmelCase		:       Optional[int]       =    self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]						)						)
           _UpperCAmelCase		:       List[Any]       =    [output_layers[self.config.num_hidden_layers - 1](A_						)]
         else:
           _UpperCAmelCase		:       str       =    0
           _UpperCAmelCase		:       str       =    None
           _UpperCAmelCase		:       str       =    0
           for i in range(self.config.num_hidden_layers						):
             calculated_layer_num += 1
             _UpperCAmelCase		:       int       =    self.encoder.adaptive_forward(
                 A_       ,				current_layer=A_       ,				attention_mask=A_       ,				head_mask=A_       ,				)
  
             _UpperCAmelCase		:       Tuple       =    self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]						)						)
             _UpperCAmelCase		:       List[str]       =    output_layers[i](A_						)
             if regression:
               _UpperCAmelCase		:       Optional[Any]       =    logits.detach()
               if patient_result is not None:
                 _UpperCAmelCase		:       str       =    patient_result.detach()
               if (patient_result is not None) and torch.abs(patient_result - labels						) < self.regression_threshold:
                 patient_counter += 1
               else:
                 _UpperCAmelCase		:       Union[str, Any]       =    0
             else:
               _UpperCAmelCase		:       Dict       =    logits.detach().argmax(dim=1						)
               if patient_result is not None:
                 _UpperCAmelCase		:       Optional[Any]       =    patient_result.detach().argmax(dim=1						)
               if (patient_result is not None) and torch.all(labels.eq(A_						)						):
                 patient_counter += 1
               else:
                 _UpperCAmelCase		:       Dict       =    0
  
             _UpperCAmelCase		:       int       =    logits
             if patient_counter == self.patience:
               break
           _UpperCAmelCase		:       Tuple       =    [patient_result]
           self.inference_layers_num += calculated_layer_num
           self.inference_instances_num += 1
  
         return res
  
  
  @add_start_docstrings(
      '''Albert Model transformer with PABEE and a sequence classification/regression head on top (a linear layer on top of
      the pooled output) e.g. for GLUE tasks. '''       ,  A       ,  )
  class 		_SCREAMING_SNAKE_CASE    (	A    ):
  
  
       def __init__(     self       ,				A_						):
         super().__init__(A_						)
         _UpperCAmelCase		:       List[str]       =    config.num_labels
  
         _UpperCAmelCase		:       List[Any]       =    AlbertModelWithPabee(A_						)
         _UpperCAmelCase		:       Optional[int]       =    nn.Dropout(config.classifier_dropout_prob						)
         _UpperCAmelCase		:       Tuple       =    nn.ModuleList(
             [nn.Linear(config.hidden_size       ,				self.config.num_labels						) for _ in range(config.num_hidden_layers						)]						)
  
         self.init_weights()
  
  
  
  
       @add_start_docstrings_to_model_forward(A_						)
       def 					__snake_case(     self       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				):
         _UpperCAmelCase		:       Optional[Any]       =    self.albert(
             input_ids=A_       ,				attention_mask=A_       ,				token_type_ids=A_       ,				position_ids=A_       ,				head_mask=A_       ,				inputs_embeds=A_       ,				output_dropout=self.dropout       ,				output_layers=self.classifiers       ,				regression=self.num_labels == 1       ,				)
  
         _UpperCAmelCase		:       str       =    (logits[-1],)
  
         if labels is not None:
           _UpperCAmelCase		:       List[Any]       =    None
           _UpperCAmelCase		:       Optional[int]       =    0
           for ix, logits_item in enumerate(A_						):
             if self.num_labels == 1:
               #  We are doing regression
               _UpperCAmelCase		:       int       =    MSELoss()
               _UpperCAmelCase		:       Union[str, Any]       =    loss_fct(logits_item.view(-1						)       ,				labels.view(-1						)						)
             else:
               _UpperCAmelCase		:       List[Any]       =    CrossEntropyLoss()
               _UpperCAmelCase		:       List[Any]       =    loss_fct(logits_item.view(-1       ,				self.num_labels						)       ,				labels.view(-1						)						)
             if total_loss is None:
               _UpperCAmelCase		:       List[str]       =    loss
             else:
               total_loss += loss * (ix + 1)
             total_weights += ix + 1
           _UpperCAmelCase		:       List[str]       =    (total_loss / total_weights,) + outputs
  
         return outputs
  
  
  
  
  
  