

"""simple docstring"""




import argparse
import json
from pathlib import Path

import torch
import torchaudio
from datasets import load_dataset
from huggingface_hub import hf_hub_download

from transformers import ASTConfig, ASTFeatureExtractor, ASTForAudioClassification
from transformers.utils import logging


logging.set_verbosity_info()
SCREAMING_SNAKE_CASE_     :			int	  =			logging.get_logger(__name__)



def 	UpperCAmelCase__   (				A__   )      ->      Optional[int]:




   """simple docstring"""
   lowerCamelCase__						=			ASTConfig()

   if "10-10" in model_name:
      pass
   elif "speech-commands" in model_name:
      lowerCamelCase__						=			128
   elif "12-12" in model_name:
      lowerCamelCase__						=			12
      lowerCamelCase__						=			12
   elif "14-14" in model_name:
      lowerCamelCase__						=			14
      lowerCamelCase__						=			14
   elif "16-16" in model_name:
      lowerCamelCase__						=			16
      lowerCamelCase__						=			16
   else:
      raise ValueError("Model not supported"   )

   lowerCamelCase__						=			"huggingface/label-files"
   if "speech-commands" in model_name:
      lowerCamelCase__						=			35
      lowerCamelCase__						=			"speech-commands-v2-id2label.json"
   else:
      lowerCamelCase__						=			527
      lowerCamelCase__						=			"audioset-id2label.json"

   lowerCamelCase__						=			json.load(open(hf_hub_download(A__	,    A__	,    repo_type="dataset"   )	,    "r"   )   )
   lowerCamelCase__						=			{int(A__   ): v for k, v in idalabel.items()}
   lowerCamelCase__						=			idalabel
   lowerCamelCase__						=			{v: k for k, v in idalabel.items()}

   return config



def 	UpperCAmelCase__   (				A__   )      ->      str:




   """simple docstring"""
   if "module.v" in name:
      lowerCamelCase__						=			name.replace("module.v"	,    "audio_spectrogram_transformer"   )
   if "cls_token" in name:
      lowerCamelCase__						=			name.replace("cls_token"	,    "embeddings.cls_token"   )
   if "dist_token" in name:
      lowerCamelCase__						=			name.replace("dist_token"	,    "embeddings.distillation_token"   )
   if "pos_embed" in name:
      lowerCamelCase__						=			name.replace("pos_embed"	,    "embeddings.position_embeddings"   )
   if "patch_embed.proj" in name:
      lowerCamelCase__						=			name.replace("patch_embed.proj"	,    "embeddings.patch_embeddings.projection"   )
   # transformer blocks
   if "blocks" in name:
      lowerCamelCase__						=			name.replace("blocks"	,    "encoder.layer"   )
   if "attn.proj" in name:
      lowerCamelCase__						=			name.replace("attn.proj"	,    "attention.output.dense"   )
   if "attn" in name:
      lowerCamelCase__						=			name.replace("attn"	,    "attention.self"   )
   if "norm1" in name:
      lowerCamelCase__						=			name.replace("norm1"	,    "layernorm_before"   )
   if "norm2" in name:
      lowerCamelCase__						=			name.replace("norm2"	,    "layernorm_after"   )
   if "mlp.fc1" in name:
      lowerCamelCase__						=			name.replace("mlp.fc1"	,    "intermediate.dense"   )
   if "mlp.fc2" in name:
      lowerCamelCase__						=			name.replace("mlp.fc2"	,    "output.dense"   )
   # final layernorm
   if "audio_spectrogram_transformer.norm" in name:
      lowerCamelCase__						=			name.replace("audio_spectrogram_transformer.norm"	,    "audio_spectrogram_transformer.layernorm"   )
   # classifier head
   if "module.mlp_head.0" in name:
      lowerCamelCase__						=			name.replace("module.mlp_head.0"	,    "classifier.layernorm"   )
   if "module.mlp_head.1" in name:
      lowerCamelCase__						=			name.replace("module.mlp_head.1"	,    "classifier.dense"   )

   return name



def 	UpperCAmelCase__   (				A__	,    A__   )      ->      Union[str, Any]:




   """simple docstring"""
   for key in orig_state_dict.copy().keys():
      lowerCamelCase__						=			orig_state_dict.pop(A__   )

      if "qkv" in key:
         lowerCamelCase__						=			key.split("."   )
         lowerCamelCase__						=			int(key_split[3]   )
         lowerCamelCase__						=			config.hidden_size
         if "weight" in key:
            lowerCamelCase__						=			val[:dim, :]
            lowerCamelCase__						=			val[dim : dim * 2, :]
            lowerCamelCase__						=			val[-dim:, :]
         else:
            lowerCamelCase__						=			val[:dim]
            lowerCamelCase__						=			val[dim : dim * 2]
            lowerCamelCase__						=			val[-dim:]
      else:
         lowerCamelCase__						=			val

   return orig_state_dict



def 	UpperCAmelCase__   (				A__   )      ->      Dict:




   """simple docstring"""
   lowerCamelCase__						=			[
       "module.v.head.weight",
       "module.v.head.bias",
       "module.v.head_dist.weight",
       "module.v.head_dist.bias",
   ]
   for k in ignore_keys:
      state_dict.pop(A__	,    A__   )



@torch.no_grad()
def 	UpperCAmelCase__   (				A__	,    A__	,    A__=False   )      ->      str:




   """simple docstring"""
   lowerCamelCase__						=			get_audio_spectrogram_transformer_config(A__   )

   lowerCamelCase__						=			{
       "ast-finetuned-audioset-10-10-0.4593": (
           "https://www.dropbox.com/s/ca0b1v2nlxzyeb4/audioset_10_10_0.4593.pth?dl=1"
       ),
       "ast-finetuned-audioset-10-10-0.450": (
           "https://www.dropbox.com/s/1tv0hovue1bxupk/audioset_10_10_0.4495.pth?dl=1"
       ),
       "ast-finetuned-audioset-10-10-0.448": (
           "https://www.dropbox.com/s/6u5sikl4b9wo4u5/audioset_10_10_0.4483.pth?dl=1"
       ),
       "ast-finetuned-audioset-10-10-0.448-v2": (
           "https://www.dropbox.com/s/kt6i0v9fvfm1mbq/audioset_10_10_0.4475.pth?dl=1"
       ),
       "ast-finetuned-audioset-12-12-0.447": (
           "https://www.dropbox.com/s/snfhx3tizr4nuc8/audioset_12_12_0.4467.pth?dl=1"
       ),
       "ast-finetuned-audioset-14-14-0.443": (
           "https://www.dropbox.com/s/z18s6pemtnxm4k7/audioset_14_14_0.4431.pth?dl=1"
       ),
       "ast-finetuned-audioset-16-16-0.442": (
           "https://www.dropbox.com/s/mdsa4t1xmcimia6/audioset_16_16_0.4422.pth?dl=1"
       ),
       "ast-finetuned-speech-commands-v2": (
           "https://www.dropbox.com/s/q0tbqpwv44pquwy/speechcommands_10_10_0.9812.pth?dl=1"
       ),
   }

   # load original state_dict
   lowerCamelCase__						=			model_name_to_url[model_name]
   lowerCamelCase__						=			torch.hub.load_state_dict_from_url(A__	,    map_location="cpu"   )
   # remove some keys
   remove_keys(A__   )
   # rename some keys
   lowerCamelCase__						=			convert_state_dict(A__	,    A__   )

   # load 🤗 model
   lowerCamelCase__						=			ASTForAudioClassification(A__   )
   model.eval()

   model.load_state_dict(A__   )

   # verify outputs on dummy input
   # source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
   lowerCamelCase__						=			-4.2677393 if "speech-commands" not in model_name else -6.845978
   lowerCamelCase__						=			4.5689974 if "speech-commands" not in model_name else 5.5654526
   lowerCamelCase__						=			1024 if "speech-commands" not in model_name else 128
   lowerCamelCase__						=			ASTFeatureExtractor(mean=A__	,    std=A__	,    max_length=A__   )

   if "speech-commands" in model_name:
      lowerCamelCase__						=			load_dataset("speech_commands"	,    "v0.02"	,    split="validation"   )
      lowerCamelCase__						=			dataset[0]["audio"]["array"]
   else:
      lowerCamelCase__						=			hf_hub_download(
          repo_id="nielsr/audio-spectogram-transformer-checkpoint"	,    filename="sample_audio.flac"	,    repo_type="dataset"	,    )

      lowerCamelCase__					,       lowerCamelCase__						=			torchaudio.load(A__   )
      lowerCamelCase__						=			waveform.squeeze().numpy()

   lowerCamelCase__						=			feature_extractor(A__	,    sampling_rate=1_6000	,    return_tensors="pt"   )

   # forward pass
   lowerCamelCase__						=			model(**A__   )
   lowerCamelCase__						=			outputs.logits

   if model_name == "ast-finetuned-audioset-10-10-0.4593":
      lowerCamelCase__						=			torch.tensor([-0.8760, -7.0042, -8.6602]   )
   elif model_name == "ast-finetuned-audioset-10-10-0.450":
      lowerCamelCase__						=			torch.tensor([-1.1986, -7.0903, -8.2718]   )
   elif model_name == "ast-finetuned-audioset-10-10-0.448":
      lowerCamelCase__						=			torch.tensor([-2.6128, -8.0080, -9.4344]   )
   elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
      lowerCamelCase__						=			torch.tensor([-1.5080, -7.4534, -8.8917]   )
   elif model_name == "ast-finetuned-audioset-12-12-0.447":
      lowerCamelCase__						=			torch.tensor([-0.5050, -6.5833, -8.0843]   )
   elif model_name == "ast-finetuned-audioset-14-14-0.443":
      lowerCamelCase__						=			torch.tensor([-0.3826, -7.0336, -8.2413]   )
   elif model_name == "ast-finetuned-audioset-16-16-0.442":
      lowerCamelCase__						=			torch.tensor([-1.2113, -6.9101, -8.3470]   )
   elif model_name == "ast-finetuned-speech-commands-v2":
      lowerCamelCase__						=			torch.tensor([6.1589, -8.0566, -8.7984]   )
   else:
      raise ValueError("Unknown model name"   )
   if not torch.allclose(logits[0, :3]	,    A__	,    atol=1E-4   ):
      raise ValueError("Logits don't match"   )
   print("Looks ok!"   )

   if pytorch_dump_folder_path is not None:
      Path(A__   ).mkdir(exist_ok=A__   )
      print(f'Saving model {model_name} to {pytorch_dump_folder_path}'   )
      model.save_pretrained(A__   )
      print(f'Saving feature extractor to {pytorch_dump_folder_path}'   )
      feature_extractor.save_pretrained(A__   )

   if push_to_hub:
      print("Pushing model and feature extractor to the hub..."   )
      model.push_to_hub(f'MIT/{model_name}'   )
      feature_extractor.push_to_hub(f'MIT/{model_name}'   )


if __name__ == "__main__":
    SCREAMING_SNAKE_CASE_     :			List[str]	  =			argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        '''--model_name''',
        default='''ast-finetuned-audioset-10-10-0.4593''',
        type=str,
        help='''Name of the Audio Spectrogram Transformer model you\'d like to convert.''',
    )
    parser.add_argument(
        '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model directory.'''
    )
    parser.add_argument(
        '''--push_to_hub''', action='''store_true''', help='''Whether or not to push the converted model to the 🤗 hub.'''
    )

    SCREAMING_SNAKE_CASE_     :			str	  =			parser.parse_args()
    convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)







"""simple docstring"""




import unittest

from transformers import AlbertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

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


if is_torch_available():
     import torch

     from transformers import (
         MODEL_FOR_PRETRAINING_MAPPING,
         AlbertForMaskedLM,
         AlbertForMultipleChoice,
         AlbertForPreTraining,
         AlbertForQuestionAnswering,
         AlbertForSequenceClassification,
         AlbertForTokenClassification,
         AlbertModel,
     )
     from transformers.models.albert.modeling_albert import ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST


class 		_A       :





      def __init__(  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__=13    ,						SCREAMING_SNAKE_CASE__=7    ,						SCREAMING_SNAKE_CASE__=True    ,						SCREAMING_SNAKE_CASE__=True    ,						SCREAMING_SNAKE_CASE__=True    ,						SCREAMING_SNAKE_CASE__=True    ,						SCREAMING_SNAKE_CASE__=99    ,						SCREAMING_SNAKE_CASE__=16    ,						SCREAMING_SNAKE_CASE__=36    ,						SCREAMING_SNAKE_CASE__=6    ,						SCREAMING_SNAKE_CASE__=6    ,						SCREAMING_SNAKE_CASE__=6    ,						SCREAMING_SNAKE_CASE__=37    ,						SCREAMING_SNAKE_CASE__="gelu"    ,						SCREAMING_SNAKE_CASE__=0.1    ,						SCREAMING_SNAKE_CASE__=0.1    ,						SCREAMING_SNAKE_CASE__=512    ,						SCREAMING_SNAKE_CASE__=16    ,						SCREAMING_SNAKE_CASE__=2    ,						SCREAMING_SNAKE_CASE__=0.02    ,						SCREAMING_SNAKE_CASE__=3    ,						SCREAMING_SNAKE_CASE__=4    ,						SCREAMING_SNAKE_CASE__=None    ,						)   ->					Optional[Any]:
         lowerCamelCase__						=			parent
         lowerCamelCase__						=			batch_size
         lowerCamelCase__						=			seq_length
         lowerCamelCase__						=			is_training
         lowerCamelCase__						=			use_input_mask
         lowerCamelCase__						=			use_token_type_ids
         lowerCamelCase__						=			use_labels
         lowerCamelCase__						=			vocab_size
         lowerCamelCase__						=			embedding_size
         lowerCamelCase__						=			hidden_size
         lowerCamelCase__						=			num_hidden_layers
         lowerCamelCase__						=			num_hidden_groups
         lowerCamelCase__						=			num_attention_heads
         lowerCamelCase__						=			intermediate_size
         lowerCamelCase__						=			hidden_act
         lowerCamelCase__						=			hidden_dropout_prob
         lowerCamelCase__						=			attention_probs_dropout_prob
         lowerCamelCase__						=			max_position_embeddings
         lowerCamelCase__						=			type_vocab_size
         lowerCamelCase__						=			type_sequence_label_size
         lowerCamelCase__						=			initializer_range
         lowerCamelCase__						=			num_labels
         lowerCamelCase__						=			num_choices
         lowerCamelCase__						=			scope





      def 					_lowerCamelCase  (  self		)   ->					Any:
         lowerCamelCase__						=			ids_tensor([self.batch_size, self.seq_length]    ,						self.vocab_size		)

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

         lowerCamelCase__						=			None
         if self.use_token_type_ids:
            lowerCamelCase__						=			ids_tensor([self.batch_size, self.seq_length]    ,						self.type_vocab_size		)

         lowerCamelCase__						=			None
         lowerCamelCase__						=			None
         lowerCamelCase__						=			None
         if self.use_labels:
            lowerCamelCase__						=			ids_tensor([self.batch_size]    ,						self.type_sequence_label_size		)
            lowerCamelCase__						=			ids_tensor([self.batch_size, self.seq_length]    ,						self.num_labels		)
            lowerCamelCase__						=			ids_tensor([self.batch_size]    ,						self.num_choices		)

         lowerCamelCase__						=			self.get_config()

         return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels





      def 					_lowerCamelCase  (  self		)   ->					List[str]:
         return AlbertConfig(
             vocab_size=self.vocab_size    ,						hidden_size=self.hidden_size    ,						num_hidden_layers=self.num_hidden_layers    ,						num_attention_heads=self.num_attention_heads    ,						intermediate_size=self.intermediate_size    ,						hidden_act=self.hidden_act    ,						hidden_dropout_prob=self.hidden_dropout_prob    ,						attention_probs_dropout_prob=self.attention_probs_dropout_prob    ,						max_position_embeddings=self.max_position_embeddings    ,						type_vocab_size=self.type_vocab_size    ,						initializer_range=self.initializer_range    ,						num_hidden_groups=self.num_hidden_groups    ,						)





      def 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)   ->					int:
         lowerCamelCase__						=			AlbertModel(config=SCREAMING_SNAKE_CASE__		)
         model.to(SCREAMING_SNAKE_CASE__		)
         model.eval()
         lowerCamelCase__						=			model(SCREAMING_SNAKE_CASE__    ,						attention_mask=SCREAMING_SNAKE_CASE__    ,						token_type_ids=SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			model(SCREAMING_SNAKE_CASE__    ,						token_type_ids=SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			model(SCREAMING_SNAKE_CASE__		)
         self.parent.assertEqual(result.last_hidden_state.shape    ,						(self.batch_size, self.seq_length, self.hidden_size)		)
         self.parent.assertEqual(result.pooler_output.shape    ,						(self.batch_size, self.hidden_size)		)





      def 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)   ->					str:
         lowerCamelCase__						=			AlbertForPreTraining(config=SCREAMING_SNAKE_CASE__		)
         model.to(SCREAMING_SNAKE_CASE__		)
         model.eval()
         lowerCamelCase__						=			model(
             SCREAMING_SNAKE_CASE__    ,						attention_mask=SCREAMING_SNAKE_CASE__    ,						token_type_ids=SCREAMING_SNAKE_CASE__    ,						labels=SCREAMING_SNAKE_CASE__    ,						sentence_order_label=SCREAMING_SNAKE_CASE__    ,						)
         self.parent.assertEqual(result.prediction_logits.shape    ,						(self.batch_size, self.seq_length, self.vocab_size)		)
         self.parent.assertEqual(result.sop_logits.shape    ,						(self.batch_size, config.num_labels)		)





      def 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)   ->					Any:
         lowerCamelCase__						=			AlbertForMaskedLM(config=SCREAMING_SNAKE_CASE__		)
         model.to(SCREAMING_SNAKE_CASE__		)
         model.eval()
         lowerCamelCase__						=			model(SCREAMING_SNAKE_CASE__    ,						attention_mask=SCREAMING_SNAKE_CASE__    ,						token_type_ids=SCREAMING_SNAKE_CASE__    ,						labels=SCREAMING_SNAKE_CASE__		)
         self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.seq_length, self.vocab_size)		)





      def 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)   ->					Any:
         lowerCamelCase__						=			AlbertForQuestionAnswering(config=SCREAMING_SNAKE_CASE__		)
         model.to(SCREAMING_SNAKE_CASE__		)
         model.eval()
         lowerCamelCase__						=			model(
             SCREAMING_SNAKE_CASE__    ,						attention_mask=SCREAMING_SNAKE_CASE__    ,						token_type_ids=SCREAMING_SNAKE_CASE__    ,						start_positions=SCREAMING_SNAKE_CASE__    ,						end_positions=SCREAMING_SNAKE_CASE__    ,						)
         self.parent.assertEqual(result.start_logits.shape    ,						(self.batch_size, self.seq_length)		)
         self.parent.assertEqual(result.end_logits.shape    ,						(self.batch_size, self.seq_length)		)





      def 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)   ->					Optional[int]:
         lowerCamelCase__						=			self.num_labels
         lowerCamelCase__						=			AlbertForSequenceClassification(SCREAMING_SNAKE_CASE__		)
         model.to(SCREAMING_SNAKE_CASE__		)
         model.eval()
         lowerCamelCase__						=			model(SCREAMING_SNAKE_CASE__    ,						attention_mask=SCREAMING_SNAKE_CASE__    ,						token_type_ids=SCREAMING_SNAKE_CASE__    ,						labels=SCREAMING_SNAKE_CASE__		)
         self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.num_labels)		)





      def 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)   ->					int:
         lowerCamelCase__						=			self.num_labels
         lowerCamelCase__						=			AlbertForTokenClassification(config=SCREAMING_SNAKE_CASE__		)
         model.to(SCREAMING_SNAKE_CASE__		)
         model.eval()
         lowerCamelCase__						=			model(SCREAMING_SNAKE_CASE__    ,						attention_mask=SCREAMING_SNAKE_CASE__    ,						token_type_ids=SCREAMING_SNAKE_CASE__    ,						labels=SCREAMING_SNAKE_CASE__		)
         self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.seq_length, self.num_labels)		)





      def 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)   ->					int:
         lowerCamelCase__						=			self.num_choices
         lowerCamelCase__						=			AlbertForMultipleChoice(config=SCREAMING_SNAKE_CASE__		)
         model.to(SCREAMING_SNAKE_CASE__		)
         model.eval()
         lowerCamelCase__						=			input_ids.unsqueeze(1		).expand(-1    ,						self.num_choices    ,						-1		).contiguous()
         lowerCamelCase__						=			token_type_ids.unsqueeze(1		).expand(-1    ,						self.num_choices    ,						-1		).contiguous()
         lowerCamelCase__						=			input_mask.unsqueeze(1		).expand(-1    ,						self.num_choices    ,						-1		).contiguous()
         lowerCamelCase__						=			model(
             SCREAMING_SNAKE_CASE__    ,						attention_mask=SCREAMING_SNAKE_CASE__    ,						token_type_ids=SCREAMING_SNAKE_CASE__    ,						labels=SCREAMING_SNAKE_CASE__    ,						)
         self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.num_choices)		)





      def 					_lowerCamelCase  (  self		)   ->					Union[str, Any]:
         lowerCamelCase__						=			self.prepare_config_and_inputs()
         (
             (
             lowerCamelCase__
         )					,       (
             lowerCamelCase__
         )					,       (
             lowerCamelCase__
         )					,       (
             lowerCamelCase__
         )					,       (
             lowerCamelCase__
         )					,       (
             lowerCamelCase__
         )					,       (
             lowerCamelCase__
         )					,       
         )						=			config_and_inputs
         lowerCamelCase__						=			{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
         return config, inputs_dict






@require_torch
class 		_A       (			__a				,       __a				,       unittest.TestCase					):
      __a								=	(
          (
              AlbertModel,
              AlbertForPreTraining,
              AlbertForMaskedLM,
              AlbertForMultipleChoice,
              AlbertForSequenceClassification,
              AlbertForTokenClassification,
              AlbertForQuestionAnswering,
          )
          if is_torch_available()
          else ()
      )
      __a								=	(
          {
              'feature-extraction': AlbertModel,
              'fill-mask': AlbertForMaskedLM,
              'question-answering': AlbertForQuestionAnswering,
              'text-classification': AlbertForSequenceClassification,
              'token-classification': AlbertForTokenClassification,
              'zero-shot': AlbertForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      __a								=	True





      def 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__=False		)   ->					Optional[Any]:
         lowerCamelCase__						=			super()._prepare_for_class(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						return_labels=SCREAMING_SNAKE_CASE__		)

         if return_labels:
            if model_class in get_values(SCREAMING_SNAKE_CASE__		):
               lowerCamelCase__						=			torch.zeros(
                   (self.model_tester.batch_size, self.model_tester.seq_length)    ,						dtype=torch.long    ,						device=SCREAMING_SNAKE_CASE__		)
               lowerCamelCase__						=			torch.zeros(
                   self.model_tester.batch_size    ,						dtype=torch.long    ,						device=SCREAMING_SNAKE_CASE__		)
         return inputs_dict





      def 					_lowerCamelCase  (  self		)   ->					Union[str, Any]:
         lowerCamelCase__						=			AlbertModelTester(self		)
         lowerCamelCase__						=			ConfigTester(self    ,						config_class=SCREAMING_SNAKE_CASE__    ,						hidden_size=37		)





      def 					_lowerCamelCase  (  self		)   ->					Any:
         self.config_tester.run_common_tests()





      def 					_lowerCamelCase  (  self		)   ->					int:
         lowerCamelCase__						=			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE__		)





      def 					_lowerCamelCase  (  self		)   ->					Optional[Any]:
         lowerCamelCase__						=			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_pretraining(*SCREAMING_SNAKE_CASE__		)





      def 					_lowerCamelCase  (  self		)   ->					Dict:
         lowerCamelCase__						=			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE__		)





      def 					_lowerCamelCase  (  self		)   ->					Dict:
         lowerCamelCase__						=			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_multiple_choice(*SCREAMING_SNAKE_CASE__		)





      def 					_lowerCamelCase  (  self		)   ->					int:
         lowerCamelCase__						=			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(*SCREAMING_SNAKE_CASE__		)





      def 					_lowerCamelCase  (  self		)   ->					Union[str, Any]:
         lowerCamelCase__						=			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_sequence_classification(*SCREAMING_SNAKE_CASE__		)





      def 					_lowerCamelCase  (  self		)   ->					Optional[int]:
         lowerCamelCase__						=			self.model_tester.prepare_config_and_inputs()
         for type in ["absolute", "relative_key", "relative_key_query"]:
            lowerCamelCase__						=			type
            self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE__		)





      @slow
      def 					_lowerCamelCase  (  self		)   ->					int:
         for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            lowerCamelCase__						=			AlbertModel.from_pretrained(SCREAMING_SNAKE_CASE__		)
            self.assertIsNotNone(SCREAMING_SNAKE_CASE__		)






@require_torch
class 		_A       (			unittest.TestCase					):





      @slow
      def 					_lowerCamelCase  (  self		)   ->					int:
         lowerCamelCase__						=			AlbertModel.from_pretrained("albert-base-v2"		)
         lowerCamelCase__						=			torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]		)
         lowerCamelCase__						=			torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]		)
         with torch.no_grad():
            lowerCamelCase__						=			model(SCREAMING_SNAKE_CASE__    ,						attention_mask=SCREAMING_SNAKE_CASE__		)[0]
         lowerCamelCase__						=			torch.Size((1, 11, 768)		)
         self.assertEqual(output.shape    ,						SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			torch.tensor(
             [[[-0.65_13, 1.50_35, -0.27_66], [-0.65_15, 1.50_46, -0.27_80], [-0.65_12, 1.50_49, -0.27_84]]]		)

         self.assertTrue(torch.allclose(output[:, 1:4, 1:4]    ,						SCREAMING_SNAKE_CASE__    ,						atol=1e-4		)		)




