



"""simple docstring"""


import gc
import unittest

from transformers import MODEL_FOR_MASKED_LM_MAPPING, TF_MODEL_FOR_MASKED_LM_MAPPING, FillMaskPipeline, pipeline
from transformers.pipelines import PipelineException
from transformers.testing_utils import (
    is_pipeline_test,
    is_torch_available,
    nested_simplify,
    require_tf,
    require_torch,
    require_torch_gpu,
    slow,
)

from .test_pipelines_common import ANY







@is_pipeline_test
class       A_			(     unittest.TestCase				):



						"""simple docstring"""
						__UpperCamelCase						=							MODEL_FOR_MASKED_LM_MAPPING
						__UpperCamelCase						=							TF_MODEL_FOR_MASKED_LM_MAPPING






						def 	UpperCAmelCase__     (					self :Tuple					)				->			Dict:
												super().tearDown()
												# clean-up as much as possible GPU memory occupied by PyTorch
												gc.collect()
												if is_torch_available():
																		import torch

																		torch.cuda.empty_cache()






						@require_tf
						def 	UpperCAmelCase__     (					self :Optional[int]					)				->			str:
												UpperCAmelCase             = pipeline(task='fill-mask'   ,  model='sshleifer/tiny-distilroberta-base'   ,  top_k=2   ,  framework='tf'					)
												UpperCAmelCase             = unmasker('My name is <mask>'					)
												self.assertEqual(
												    nested_simplify(__UpperCAmelCase   ,  decimals=6					)   ,  [
												        {'sequence': 'My name is grouped', 'score': 2.1E-05, 'token': 3_80_15, 'token_str': ' grouped'},
												        {'sequence': 'My name is accuser', 'score': 2.1E-05, 'token': 2_55_06, 'token_str': ' accuser'},
												    ]   ,  )

												UpperCAmelCase             = unmasker('The largest city in France is <mask>'					)
												self.assertEqual(
												    nested_simplify(__UpperCAmelCase   ,  decimals=6					)   ,  [
												        {
												            'sequence': 'The largest city in France is grouped',
												            'score': 2.1E-05,
												            'token': 3_80_15,
												            'token_str': ' grouped',
												        },
												        {
												            'sequence': 'The largest city in France is accuser',
												            'score': 2.1E-05,
												            'token': 2_55_06,
												            'token_str': ' accuser',
												        },
												    ]   ,  )

												UpperCAmelCase             = unmasker('My name is <mask>'   ,  targets=[' Patrick', ' Clara', ' Teven']   ,  top_k=3					)
												self.assertEqual(
												    nested_simplify(__UpperCAmelCase   ,  decimals=6					)   ,  [
												        {'sequence': 'My name is Clara', 'score': 2E-05, 'token': 1_36_06, 'token_str': ' Clara'},
												        {'sequence': 'My name is Patrick', 'score': 2E-05, 'token': 34_99, 'token_str': ' Patrick'},
												        {'sequence': 'My name is Te', 'score': 1.9E-05, 'token': 29_41, 'token_str': ' Te'},
												    ]   ,  )






						@require_torch
						def 	UpperCAmelCase__     (					self :Optional[int]					)				->			int:
												UpperCAmelCase             = pipeline(task='fill-mask'   ,  model='sshleifer/tiny-distilroberta-base'   ,  top_k=2   ,  framework='pt'					)

												UpperCAmelCase             = unmasker('My name is <mask>'					)
												self.assertEqual(
												    nested_simplify(__UpperCAmelCase   ,  decimals=6					)   ,  [
												        {'sequence': 'My name is Maul', 'score': 2.2E-05, 'token': 3_56_76, 'token_str': ' Maul'},
												        {'sequence': 'My name isELS', 'score': 2.2E-05, 'token': 1_64_16, 'token_str': 'ELS'},
												    ]   ,  )

												UpperCAmelCase             = unmasker('The largest city in France is <mask>'					)
												self.assertEqual(
												    nested_simplify(__UpperCAmelCase   ,  decimals=6					)   ,  [
												        {
												            'sequence': 'The largest city in France is Maul',
												            'score': 2.2E-05,
												            'token': 3_56_76,
												            'token_str': ' Maul',
												        },
												        {'sequence': 'The largest city in France isELS', 'score': 2.2E-05, 'token': 1_64_16, 'token_str': 'ELS'},
												    ]   ,  )

												UpperCAmelCase             = unmasker('My name is <mask>'   ,  targets=[' Patrick', ' Clara', ' Teven']   ,  top_k=3					)
												self.assertEqual(
												    nested_simplify(__UpperCAmelCase   ,  decimals=6					)   ,  [
												        {'sequence': 'My name is Patrick', 'score': 2.1E-05, 'token': 34_99, 'token_str': ' Patrick'},
												        {'sequence': 'My name is Te', 'score': 2E-05, 'token': 29_41, 'token_str': ' Te'},
												        {'sequence': 'My name is Clara', 'score': 2E-05, 'token': 1_36_06, 'token_str': ' Clara'},
												    ]   ,  )

												UpperCAmelCase             = unmasker('My name is <mask> <mask>'   ,  top_k=2					)

												self.assertEqual(
												    nested_simplify(__UpperCAmelCase   ,  decimals=6					)   ,  [
												        [
												            {
												                'score': 2.2E-05,
												                'token': 3_56_76,
												                'token_str': ' Maul',
												                'sequence': '<s>My name is Maul<mask></s>',
												            },
												            {'score': 2.2E-05, 'token': 1_64_16, 'token_str': 'ELS', 'sequence': '<s>My name isELS<mask></s>'},
												        ],
												        [
												            {
												                'score': 2.2E-05,
												                'token': 3_56_76,
												                'token_str': ' Maul',
												                'sequence': '<s>My name is<mask> Maul</s>',
												            },
												            {'score': 2.2E-05, 'token': 1_64_16, 'token_str': 'ELS', 'sequence': '<s>My name is<mask>ELS</s>'},
												        ],
												    ]   ,  )






						@require_torch_gpu
						def 	UpperCAmelCase__     (					self :int					)				->			Optional[Any]:
												UpperCAmelCase             = pipeline('fill-mask'   ,  model='hf-internal-testing/tiny-random-distilbert'   ,  device=0   ,  framework='pt'					)

												# convert model to fp16
												pipe.model.half()

												UpperCAmelCase             = pipe('Paris is the [MASK] of France.'					)
												# We actually don't care about the result, we just want to make sure
												# it works, meaning the float16 tensor got casted back to float32
												# for postprocessing.
												self.assertIsInstance(__UpperCAmelCase   ,  __UpperCAmelCase					)






						@slow
						@require_torch
						def 	UpperCAmelCase__     (					self :Dict					)				->			str:
												UpperCAmelCase             = pipeline(task='fill-mask'   ,  model='distilroberta-base'   ,  top_k=2   ,  framework='pt'					)
												self.run_large_test(__UpperCAmelCase					)






						@slow
						@require_tf
						def 	UpperCAmelCase__     (					self :Optional[int]					)				->			Any:
												UpperCAmelCase             = pipeline(task='fill-mask'   ,  model='distilroberta-base'   ,  top_k=2   ,  framework='tf'					)
												self.run_large_test(__UpperCAmelCase					)






						def 	UpperCAmelCase__     (					self :int   ,  lowercase_ :Optional[int]					)				->			Dict:
												UpperCAmelCase             = unmasker('My name is <mask>'					)
												self.assertEqual(
												    nested_simplify(__UpperCAmelCase					)   ,  [
												        {'sequence': 'My name is John', 'score': 0.008, 'token': 6_10, 'token_str': ' John'},
												        {'sequence': 'My name is Chris', 'score': 0.007, 'token': 15_73, 'token_str': ' Chris'},
												    ]   ,  )
												UpperCAmelCase             = unmasker('The largest city in France is <mask>'					)
												self.assertEqual(
												    nested_simplify(__UpperCAmelCase					)   ,  [
												        {
												            'sequence': 'The largest city in France is Paris',
												            'score': 0.251,
												            'token': 22_01,
												            'token_str': ' Paris',
												        },
												        {
												            'sequence': 'The largest city in France is Lyon',
												            'score': 0.214,
												            'token': 1_27_90,
												            'token_str': ' Lyon',
												        },
												    ]   ,  )

												UpperCAmelCase             = unmasker('My name is <mask>'   ,  targets=[' Patrick', ' Clara', ' Teven']   ,  top_k=3					)
												self.assertEqual(
												    nested_simplify(__UpperCAmelCase					)   ,  [
												        {'sequence': 'My name is Patrick', 'score': 0.005, 'token': 34_99, 'token_str': ' Patrick'},
												        {'sequence': 'My name is Clara', 'score': 0.000, 'token': 1_36_06, 'token_str': ' Clara'},
												        {'sequence': 'My name is Te', 'score': 0.000, 'token': 29_41, 'token_str': ' Te'},
												    ]   ,  )






						@require_torch
						def 	UpperCAmelCase__     (					self :Union[str, Any]					)				->			List[str]:
												UpperCAmelCase             = pipeline(task='fill-mask'   ,  model='sshleifer/tiny-distilroberta-base'   ,  framework='pt'					)
												UpperCAmelCase             = None
												UpperCAmelCase             = None
												self.run_pipeline_test(__UpperCAmelCase   ,  []					)






						@require_tf
						def 	UpperCAmelCase__     (					self :Tuple					)				->			List[str]:
												UpperCAmelCase             = pipeline(task='fill-mask'   ,  model='sshleifer/tiny-distilroberta-base'   ,  framework='tf'					)
												UpperCAmelCase             = None
												UpperCAmelCase             = None
												self.run_pipeline_test(__UpperCAmelCase   ,  []					)






						def 	UpperCAmelCase__     (					self :Dict   ,  lowercase_ :Tuple   ,  lowercase_ :Tuple   ,  lowercase_ :Optional[Any]					)				->			str:
												if tokenizer is None or tokenizer.mask_token_id is None:
																		self.skipTest('The provided tokenizer has no mask token, (probably reformer or wav2vec2)'					)

												UpperCAmelCase             = FillMaskPipeline(model=__UpperCAmelCase   ,  tokenizer=__UpperCAmelCase					)
												UpperCAmelCase             = [
												    f"""This is another {tokenizer.mask_token} test""",
												]
												return fill_masker, examples






						def 	UpperCAmelCase__     (					self :Optional[int]   ,  lowercase_ :Tuple   ,  lowercase_ :Optional[Any]					)				->			List[Any]:
												UpperCAmelCase             = fill_masker.tokenizer
												UpperCAmelCase             = fill_masker.model

												UpperCAmelCase             = fill_masker(
												    f"""This is a {tokenizer.mask_token}"""   ,  )
												self.assertEqual(
												    __UpperCAmelCase   ,  [
												        {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												        {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												        {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												        {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												        {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												    ]   ,  )

												UpperCAmelCase             = fill_masker([f"""This is a {tokenizer.mask_token}"""]					)
												self.assertEqual(
												    __UpperCAmelCase   ,  [
												        {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												        {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												        {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												        {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												        {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												    ]   ,  )

												UpperCAmelCase             = fill_masker([f"""This is a {tokenizer.mask_token}""", f"""Another {tokenizer.mask_token} great test."""]					)
												self.assertEqual(
												    __UpperCAmelCase   ,  [
												        [
												            {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												            {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												            {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												            {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												            {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												        ],
												        [
												            {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												            {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												            {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												            {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												            {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												        ],
												    ]   ,  )

												with self.assertRaises(__UpperCAmelCase					):
																		fill_masker([None]					)
												# No mask_token is not supported
												with self.assertRaises(__UpperCAmelCase					):
																		fill_masker('This is'					)

												self.run_test_top_k(__UpperCAmelCase   ,  __UpperCAmelCase					)
												self.run_test_targets(__UpperCAmelCase   ,  __UpperCAmelCase					)
												self.run_test_top_k_targets(__UpperCAmelCase   ,  __UpperCAmelCase					)
												self.fill_mask_with_duplicate_targets_and_top_k(__UpperCAmelCase   ,  __UpperCAmelCase					)
												self.fill_mask_with_multiple_masks(__UpperCAmelCase   ,  __UpperCAmelCase					)






						def 	UpperCAmelCase__     (					self :Union[str, Any]   ,  lowercase_ :str   ,  lowercase_ :int					)				->			Tuple:
												UpperCAmelCase             = tokenizer.get_vocab()
												UpperCAmelCase             = sorted(vocab.keys()					)[:2]
												# Pipeline argument
												UpperCAmelCase             = FillMaskPipeline(model=__UpperCAmelCase   ,  tokenizer=__UpperCAmelCase   ,  targets=__UpperCAmelCase					)
												UpperCAmelCase             = fill_masker(f"""This is a {tokenizer.mask_token}"""					)
												self.assertEqual(
												    __UpperCAmelCase   ,  [
												        {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												        {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												    ]   ,  )
												UpperCAmelCase             = {vocab[el] for el in targets}
												self.assertEqual({el['token'] for el in outputs}   ,  __UpperCAmelCase					)
												UpperCAmelCase             = [tokenizer.decode([x]					) for x in target_ids]
												self.assertEqual({el['token_str'] for el in outputs}   ,  set(__UpperCAmelCase					)					)

												# Call argument
												UpperCAmelCase             = FillMaskPipeline(model=__UpperCAmelCase   ,  tokenizer=__UpperCAmelCase					)
												UpperCAmelCase             = fill_masker(f"""This is a {tokenizer.mask_token}"""   ,  targets=__UpperCAmelCase					)
												self.assertEqual(
												    __UpperCAmelCase   ,  [
												        {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												        {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												    ]   ,  )
												UpperCAmelCase             = {vocab[el] for el in targets}
												self.assertEqual({el['token'] for el in outputs}   ,  __UpperCAmelCase					)
												UpperCAmelCase             = [tokenizer.decode([x]					) for x in target_ids]
												self.assertEqual({el['token_str'] for el in outputs}   ,  set(__UpperCAmelCase					)					)

												# Score equivalence
												UpperCAmelCase             = fill_masker(f"""This is a {tokenizer.mask_token}"""   ,  targets=__UpperCAmelCase					)
												UpperCAmelCase             = [top_mask["token_str"] for top_mask in outputs]
												UpperCAmelCase             = [top_mask["score"] for top_mask in outputs]

												# For some BPE tokenizers, `</w>` is removed during decoding, so `token_str` won't be the same as in `targets`.
												if set(__UpperCAmelCase					) == set(__UpperCAmelCase					):
																		UpperCAmelCase             = fill_masker(f"""This is a {tokenizer.mask_token}"""   ,  targets=__UpperCAmelCase					)
																		UpperCAmelCase             = [top_mask["score"] for top_mask in unmasked_targets]
																		self.assertEqual(nested_simplify(__UpperCAmelCase					)   ,  nested_simplify(__UpperCAmelCase					)					)

												# Raises with invalid
												with self.assertRaises(__UpperCAmelCase					):
																		UpperCAmelCase             = fill_masker(f"""This is a {tokenizer.mask_token}"""   ,  targets=[]					)
												# For some tokenizers, `""` is actually in the vocabulary and the expected error won't raised
												if "" not in tokenizer.get_vocab():
																		with self.assertRaises(__UpperCAmelCase					):
																								UpperCAmelCase             = fill_masker(f"""This is a {tokenizer.mask_token}"""   ,  targets=['']					)
																		with self.assertRaises(__UpperCAmelCase					):
																								UpperCAmelCase             = fill_masker(f"""This is a {tokenizer.mask_token}"""   ,  targets=''					)






						def 	UpperCAmelCase__     (					self :int   ,  lowercase_ :Optional[Any]   ,  lowercase_ :str					)				->			Optional[int]:
												UpperCAmelCase             = FillMaskPipeline(model=__UpperCAmelCase   ,  tokenizer=__UpperCAmelCase   ,  top_k=2					)
												UpperCAmelCase             = fill_masker(f"""This is a {tokenizer.mask_token}"""					)
												self.assertEqual(
												    __UpperCAmelCase   ,  [
												        {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												        {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												    ]   ,  )

												UpperCAmelCase             = FillMaskPipeline(model=__UpperCAmelCase   ,  tokenizer=__UpperCAmelCase					)
												UpperCAmelCase             = fill_masker(f"""This is a {tokenizer.mask_token}"""   ,  top_k=2					)
												self.assertEqual(
												    __UpperCAmelCase   ,  [
												        {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												        {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												    ]   ,  )
												self.assertEqual(nested_simplify(__UpperCAmelCase					)   ,  nested_simplify(__UpperCAmelCase					)					)






						def 	UpperCAmelCase__     (					self :Any   ,  lowercase_ :Optional[int]   ,  lowercase_ :str					)				->			Any:
												UpperCAmelCase             = tokenizer.get_vocab()
												UpperCAmelCase             = FillMaskPipeline(model=__UpperCAmelCase   ,  tokenizer=__UpperCAmelCase					)

												# top_k=2, ntargets=3
												UpperCAmelCase             = sorted(vocab.keys()					)[:3]
												UpperCAmelCase             = fill_masker(f"""This is a {tokenizer.mask_token}"""   ,  top_k=2   ,  targets=__UpperCAmelCase					)

												# If we use the most probably targets, and filter differently, we should still
												# have the same results
												UpperCAmelCase             = [el["token_str"] for el in sorted(__UpperCAmelCase   ,  key=lambda lowercase_					: x["score"]   ,  reverse=__UpperCAmelCase					)]
												# For some BPE tokenizers, `</w>` is removed during decoding, so `token_str` won't be the same as in `targets`.
												if set(__UpperCAmelCase					).issubset(__UpperCAmelCase					):
																		UpperCAmelCase             = fill_masker(f"""This is a {tokenizer.mask_token}"""   ,  top_k=3   ,  targets=__UpperCAmelCase					)
																		# They should yield exactly the same result
																		self.assertEqual(nested_simplify(__UpperCAmelCase					)   ,  nested_simplify(__UpperCAmelCase					)					)






						def 	UpperCAmelCase__     (					self :int   ,  lowercase_ :Any   ,  lowercase_ :Any					)				->			List[str]:
												UpperCAmelCase             = FillMaskPipeline(model=__UpperCAmelCase   ,  tokenizer=__UpperCAmelCase					)
												UpperCAmelCase             = tokenizer.get_vocab()
												# String duplicates + id duplicates
												UpperCAmelCase             = sorted(vocab.keys()					)[:3]
												UpperCAmelCase             = [targets[0], targets[1], targets[0], targets[2], targets[1]]
												UpperCAmelCase             = fill_masker(f"""My name is {tokenizer.mask_token}"""   ,  targets=__UpperCAmelCase   ,  top_k=10					)

												# The target list contains duplicates, so we can't output more
												# than them
												self.assertEqual(len(__UpperCAmelCase					)   ,  3					)






						def 	UpperCAmelCase__     (					self :Optional[int]   ,  lowercase_ :Union[str, Any]   ,  lowercase_ :List[Any]					)				->			Tuple:
												UpperCAmelCase             = FillMaskPipeline(model=__UpperCAmelCase   ,  tokenizer=__UpperCAmelCase					)

												UpperCAmelCase             = fill_masker(
												    f"""This is a {tokenizer.mask_token} {tokenizer.mask_token} {tokenizer.mask_token}"""   ,  top_k=2					)
												self.assertEqual(
												    __UpperCAmelCase   ,  [
												        [
												            {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												            {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												        ],
												        [
												            {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												            {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												        ],
												        [
												            {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												            {'sequence': ANY(__UpperCAmelCase					), 'score': ANY(__UpperCAmelCase					), 'token': ANY(__UpperCAmelCase					), 'token_str': ANY(__UpperCAmelCase					)},
												        ],
												    ]   ,  )









"""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()
__lowercase    					=      logging.get_logger(__name__)


def      lowercase			(				A_       )->      Dict:
  '''simple docstring'''







  a						:		str      							=			ASTConfig()

  if "10-10" in model_name:
    pass
  elif "speech-commands" in model_name:
    a						:		Union[str, Any]      							=			128
  elif "12-12" in model_name:
    a						:		List[Any]      							=			12
    a						:		str      							=			12
  elif "14-14" in model_name:
    a						:		List[Any]      							=			14
    a						:		Optional[int]      							=			14
  elif "16-16" in model_name:
    a						:		Any      							=			16
    a						:		List[Any]      							=			16
  else:
    raise ValueError("Model not supported"       )

  a						:		Optional[int]      							=			"huggingface/label-files"
  if "speech-commands" in model_name:
    a						:		Optional[int]      							=			35
    a						:		List[str]      							=			"speech-commands-v2-id2label.json"
  else:
    a						:		Optional[Any]      							=			527
    a						:		Tuple      							=			"audioset-id2label.json"

  a						:		List[str]      							=			json.load(open(hf_hub_download(A_  ,   A_  ,   repo_type="dataset"       )  ,   "r"       )       )
  a						:		Union[str, Any]      							=			{int(A_       ): v for k, v in idalabel.items()}
  a						:		Any      							=			idalabel
  a						:		str      							=			{v: k for k, v in idalabel.items()}

  return config


def      lowercase			(				A_       )->      Tuple:
  '''simple docstring'''







  if "module.v" in name:
    a						:		Union[str, Any]      							=			name.replace("module.v"  ,   "audio_spectrogram_transformer"       )
  if "cls_token" in name:
    a						:		List[Any]      							=			name.replace("cls_token"  ,   "embeddings.cls_token"       )
  if "dist_token" in name:
    a						:		Union[str, Any]      							=			name.replace("dist_token"  ,   "embeddings.distillation_token"       )
  if "pos_embed" in name:
    a						:		str      							=			name.replace("pos_embed"  ,   "embeddings.position_embeddings"       )
  if "patch_embed.proj" in name:
    a						:		Union[str, Any]      							=			name.replace("patch_embed.proj"  ,   "embeddings.patch_embeddings.projection"       )
  # transformer blocks
  if "blocks" in name:
    a						:		Union[str, Any]      							=			name.replace("blocks"  ,   "encoder.layer"       )
  if "attn.proj" in name:
    a						:		str      							=			name.replace("attn.proj"  ,   "attention.output.dense"       )
  if "attn" in name:
    a						:		Tuple      							=			name.replace("attn"  ,   "attention.self"       )
  if "norm1" in name:
    a						:		int      							=			name.replace("norm1"  ,   "layernorm_before"       )
  if "norm2" in name:
    a						:		Union[str, Any]      							=			name.replace("norm2"  ,   "layernorm_after"       )
  if "mlp.fc1" in name:
    a						:		Union[str, Any]      							=			name.replace("mlp.fc1"  ,   "intermediate.dense"       )
  if "mlp.fc2" in name:
    a						:		Optional[Any]      							=			name.replace("mlp.fc2"  ,   "output.dense"       )
  # final layernorm
  if "audio_spectrogram_transformer.norm" in name:
    a						:		Tuple      							=			name.replace("audio_spectrogram_transformer.norm"  ,   "audio_spectrogram_transformer.layernorm"       )
  # classifier head
  if "module.mlp_head.0" in name:
    a						:		List[str]      							=			name.replace("module.mlp_head.0"  ,   "classifier.layernorm"       )
  if "module.mlp_head.1" in name:
    a						:		Optional[int]      							=			name.replace("module.mlp_head.1"  ,   "classifier.dense"       )

  return name


def      lowercase			(				A_  ,   A_       )->      Any:
  '''simple docstring'''







  for key in orig_state_dict.copy().keys():
    a						:		str      							=			orig_state_dict.pop(A_       )

    if "qkv" in key:
      a						:		int      							=			key.split("."       )
      a						:		Optional[int]      							=			int(key_split[3]       )
      a						:		int      							=			config.hidden_size
      if "weight" in key:
        a						:		List[str]      							=			val[:dim, :]
        a						:		Any      							=			val[dim : dim * 2, :]
        a						:		int      							=			val[-dim:, :]
      else:
        a						:		Optional[Any]      							=			val[:dim]
        a						:		Union[str, Any]      							=			val[dim : dim * 2]
        a						:		str      							=			val[-dim:]
    else:
      a						:		str      							=			val

  return orig_state_dict


def      lowercase			(				A_       )->      Dict:
  '''simple docstring'''







  a						:		Union[str, Any]      							=			[
      "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      lowercase			(				A_  ,   A_  ,   A_=False       )->      Optional[int]:
  '''simple docstring'''







  a						:		Optional[int]      							=			get_audio_spectrogram_transformer_config(A_       )

  a						:		Dict      							=			{
      "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
  a						:		Any      							=			model_name_to_url[model_name]
  a						:		List[Any]      							=			torch.hub.load_state_dict_from_url(A_  ,   map_location="cpu"       )
  # remove some keys
  remove_keys(A_       )
  # rename some keys
  a						:		Union[str, Any]      							=			convert_state_dict(A_  ,   A_       )

  # load 🤗 model
  a						:		List[str]      							=			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
  a						:		Tuple      							=			-4.2_6_7_7_3_9_3 if "speech-commands" not in model_name else -6.8_4_5_9_7_8
  a						:		Union[str, Any]      							=			4.5_6_8_9_9_7_4 if "speech-commands" not in model_name else 5.5_6_5_4_5_2_6
  a						:		str      							=			1_024 if "speech-commands" not in model_name else 128
  a						:		List[Any]      							=			ASTFeatureExtractor(mean=A_  ,   std=A_  ,   max_length=A_       )

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

    a							,		a						:		Tuple      							=			torchaudio.load(A_       )
    a						:		Optional[Any]      							=			waveform.squeeze().numpy()

  a						:		Union[str, Any]      							=			feature_extractor(A_  ,   sampling_rate=16_000  ,   return_tensors="pt"       )

  # forward pass
  a						:		Optional[Any]      							=			model(**A_       )
  a						:		List[str]      							=			outputs.logits

  if model_name == "ast-finetuned-audioset-10-10-0.4593":
    a						:		Any      							=			torch.tensor([-0.8_7_6_0, -7.0_0_4_2, -8.6_6_0_2]       )
  elif model_name == "ast-finetuned-audioset-10-10-0.450":
    a						:		Optional[int]      							=			torch.tensor([-1.1_9_8_6, -7.0_9_0_3, -8.2_7_1_8]       )
  elif model_name == "ast-finetuned-audioset-10-10-0.448":
    a						:		List[str]      							=			torch.tensor([-2.6_1_2_8, -8.0_0_8_0, -9.4_3_4_4]       )
  elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
    a						:		Tuple      							=			torch.tensor([-1.5_0_8_0, -7.4_5_3_4, -8.8_9_1_7]       )
  elif model_name == "ast-finetuned-audioset-12-12-0.447":
    a						:		int      							=			torch.tensor([-0.5_0_5_0, -6.5_8_3_3, -8.0_8_4_3]       )
  elif model_name == "ast-finetuned-audioset-14-14-0.443":
    a						:		Any      							=			torch.tensor([-0.3_8_2_6, -7.0_3_3_6, -8.2_4_1_3]       )
  elif model_name == "ast-finetuned-audioset-16-16-0.442":
    a						:		Dict      							=			torch.tensor([-1.2_1_1_3, -6.9_1_0_1, -8.3_4_7_0]       )
  elif model_name == "ast-finetuned-speech-commands-v2":
    a						:		Union[str, Any]      							=			torch.tensor([6.1_5_8_9, -8.0_5_6_6, -8.7_9_8_4]       )
  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__":
 __lowercase    					=      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."""
 )

 __lowercase    					=      parser.parse_args()
 convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)






