'''simple docstring'''


import pytest

from datasets.parallel import ParallelBackendConfig, parallel_backend
from datasets.utils.py_utils import map_nested

from .utils import require_dill_gt_0_3_2, require_joblibspark, require_not_windows





def       lowerCAmelCase					(		UpperCamelCase__				:   Tuple       ):  # picklable for multiprocessing


 """simple docstring"""
 return i + 1





@require_dill_gt_0_3_2
@require_joblibspark
@require_not_windows
def       lowerCAmelCase					(		):


 """simple docstring"""
 with parallel_backend('''spark'''       ):
  assert ParallelBackendConfig.backend_name == "spark"

 __UpperCAmelCase							=			[1, 2, 3]
 with pytest.raises(lowerCAmelCase__       ):
  with parallel_backend('''unsupported backend'''       ):
   map_nested(lowerCAmelCase__    ,					lowerCAmelCase__    ,					num_proc=2       )

 with pytest.raises(lowerCAmelCase__       ):
  with parallel_backend('''unsupported backend'''       ):
   map_nested(lowerCAmelCase__    ,					lowerCAmelCase__    ,					num_proc=-1       )





@require_dill_gt_0_3_2
@require_joblibspark
@require_not_windows
@pytest.mark.parametrize('''num_proc'''    ,					[2, -1]       )
def       lowerCAmelCase					(		UpperCamelCase__				:   Optional[Any]       ):


 """simple docstring"""
 __UpperCAmelCase							=			[1, 2]
 __UpperCAmelCase							=			{'''a''': 1, '''b''': 2}
 __UpperCAmelCase							=			{'''a''': [1, 2], '''b''': [3, 4]}
 __UpperCAmelCase							=			{'''a''': {'''1''': 1}, '''b''': 2}
 __UpperCAmelCase							=			{'''a''': 1, '''b''': 2, '''c''': 3, '''d''': 4}
 __UpperCAmelCase							=			[2, 3]
 __UpperCAmelCase							=			{'''a''': 2, '''b''': 3}
 __UpperCAmelCase							=			{'''a''': [2, 3], '''b''': [4, 5]}
 __UpperCAmelCase							=			{'''a''': {'''1''': 2}, '''b''': 3}
 __UpperCAmelCase							=			{'''a''': 2, '''b''': 3, '''c''': 4, '''d''': 5}

 with parallel_backend('''spark'''       ):
  assert map_nested(lowerCAmelCase__    ,					lowerCAmelCase__    ,					num_proc=lowerCAmelCase__       ) == expected_map_nested_sa
  assert map_nested(lowerCAmelCase__    ,					lowerCAmelCase__    ,					num_proc=lowerCAmelCase__       ) == expected_map_nested_sa
  assert map_nested(lowerCAmelCase__    ,					lowerCAmelCase__    ,					num_proc=lowerCAmelCase__       ) == expected_map_nested_sa
  assert map_nested(lowerCAmelCase__    ,					lowerCAmelCase__    ,					num_proc=lowerCAmelCase__       ) == expected_map_nested_sa
  assert map_nested(lowerCAmelCase__    ,					lowerCAmelCase__    ,					num_proc=lowerCAmelCase__       ) == expected_map_nested_sa

'''simple docstring'''


import argparse
import collections
import os
import re
import tempfile

import pandas as pd
from datasets import Dataset
from huggingface_hub import hf_hub_download, upload_folder

from transformers.utils import direct_transformers_import


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/update_metadata.py
__lowerCAmelCase	:       List[Any]						       =	"src/transformers"


# This is to make sure the transformers module imported is the one in the repo.
__lowerCAmelCase	:       str						       =	direct_transformers_import(TRANSFORMERS_PATH)


# Regexes that match TF/Flax/PT model names.
__lowerCAmelCase	:       int						       =	re.compile(r"TF(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)")
__lowerCAmelCase	:       List[Any]						       =	re.compile(r"Flax(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)")
# Will match any TF or Flax model too so need to be in an else branch afterthe two previous regexes.
__lowerCAmelCase	:       List[str]						       =	re.compile(r"(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)")


# Fill this with tuples (pipeline_tag, model_mapping, auto_model)
__lowerCAmelCase	:       Optional[int]						       =	[
    ("pretraining", "MODEL_FOR_PRETRAINING_MAPPING_NAMES", "AutoModelForPreTraining"),
    ("feature-extraction", "MODEL_MAPPING_NAMES", "AutoModel"),
    ("audio-classification", "MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES", "AutoModelForAudioClassification"),
    ("text-generation", "MODEL_FOR_CAUSAL_LM_MAPPING_NAMES", "AutoModelForCausalLM"),
    ("automatic-speech-recognition", "MODEL_FOR_CTC_MAPPING_NAMES", "AutoModelForCTC"),
    ("image-classification", "MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES", "AutoModelForImageClassification"),
    ("image-segmentation", "MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES", "AutoModelForImageSegmentation"),
    ("fill-mask", "MODEL_FOR_MASKED_LM_MAPPING_NAMES", "AutoModelForMaskedLM"),
    ("object-detection", "MODEL_FOR_OBJECT_DETECTION_MAPPING_NAMES", "AutoModelForObjectDetection"),
    (
        "zero-shot-object-detection",
        "MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING_NAMES",
        "AutoModelForZeroShotObjectDetection",
    ),
    ("question-answering", "MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES", "AutoModelForQuestionAnswering"),
    ("text2text-generation", "MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES", "AutoModelForSeq2SeqLM"),
    ("text-classification", "MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES", "AutoModelForSequenceClassification"),
    ("automatic-speech-recognition", "MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES", "AutoModelForSpeechSeq2Seq"),
    (
        "table-question-answering",
        "MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING_NAMES",
        "AutoModelForTableQuestionAnswering",
    ),
    ("token-classification", "MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES", "AutoModelForTokenClassification"),
    ("multiple-choice", "MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES", "AutoModelForMultipleChoice"),
    (
        "next-sentence-prediction",
        "MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING_NAMES",
        "AutoModelForNextSentencePrediction",
    ),
    (
        "audio-frame-classification",
        "MODEL_FOR_AUDIO_FRAME_CLASSIFICATION_MAPPING_NAMES",
        "AutoModelForAudioFrameClassification",
    ),
    ("audio-xvector", "MODEL_FOR_AUDIO_XVECTOR_MAPPING_NAMES", "AutoModelForAudioXVector"),
    (
        "document-question-answering",
        "MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING_NAMES",
        "AutoModelForDocumentQuestionAnswering",
    ),
    (
        "visual-question-answering",
        "MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING_NAMES",
        "AutoModelForVisualQuestionAnswering",
    ),
    ("image-to-text", "MODEL_FOR_FOR_VISION_2_SEQ_MAPPING_NAMES", "AutoModelForVision2Seq"),
    (
        "zero-shot-image-classification",
        "MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING_NAMES",
        "AutoModelForZeroShotImageClassification",
    ),
    ("depth-estimation", "MODEL_FOR_DEPTH_ESTIMATION_MAPPING_NAMES", "AutoModelForDepthEstimation"),
    ("video-classification", "MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING_NAMES", "AutoModelForVideoClassification"),
    ("mask-generation", "MODEL_FOR_MASK_GENERATION_MAPPING_NAMES", "AutoModelForMaskGeneration"),
]





def       lowerCAmelCase					(		UpperCamelCase__				:   int       ):


	"""simple docstring"""
	__UpperCAmelCase							=			re.finditer('''.+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)'''    ,					UpperCamelCase__       )
	return [m.group(0       ) for m in matches]





def       lowerCAmelCase					(		):


	"""simple docstring"""
	__UpperCAmelCase							=			transformers_module.models.auto.configuration_auto.CONFIG_MAPPING_NAMES
	__UpperCAmelCase							=			{
	    config.replace('''Config'''    ,					''''''       ): model_type for model_type, config in config_maping_names.items()
	}

	# Dictionaries flagging if each model prefix has a backend in PT/TF/Flax.
	__UpperCAmelCase							=			collections.defaultdict(UpperCamelCase__       )
	__UpperCAmelCase							=			collections.defaultdict(UpperCamelCase__       )
	__UpperCAmelCase							=			collections.defaultdict(UpperCamelCase__       )

	# Let's lookup through all transformers object (once) and find if models are supported by a given backend.
	for attr_name in dir(UpperCamelCase__       ):
		__UpperCAmelCase							=			None
		if _re_tf_models.match(UpperCamelCase__       ) is not None:
			__UpperCAmelCase							=			tf_models
			__UpperCAmelCase							=			_re_tf_models.match(UpperCamelCase__       ).groups()[0]
		elif _re_flax_models.match(UpperCamelCase__       ) is not None:
			__UpperCAmelCase							=			flax_models
			__UpperCAmelCase							=			_re_flax_models.match(UpperCamelCase__       ).groups()[0]
		elif _re_pt_models.match(UpperCamelCase__       ) is not None:
			__UpperCAmelCase							=			pt_models
			__UpperCAmelCase							=			_re_pt_models.match(UpperCamelCase__       ).groups()[0]

		if lookup_dict is not None:
			while len(UpperCamelCase__       ) > 0:
				if attr_name in model_prefix_to_model_type:
					__UpperCAmelCase							=			True
					break
				# Try again after removing the last word in the name
				__UpperCAmelCase							=			''''''.join(camel_case_split(UpperCamelCase__       )[:-1]       )

	__UpperCAmelCase							=			set(list(pt_models.keys()       ) + list(tf_models.keys()       ) + list(flax_models.keys()       )       )
	__UpperCAmelCase							=			list(UpperCamelCase__       )
	all_models.sort()

	__UpperCAmelCase							=			{'''model_type''': all_models}
	__UpperCAmelCase							=			[pt_models[t] for t in all_models]
	__UpperCAmelCase							=			[tf_models[t] for t in all_models]
	__UpperCAmelCase							=			[flax_models[t] for t in all_models]

	# Now let's use the auto-mapping names to make sure
	__UpperCAmelCase							=			{}
	for t in all_models:
		if t in transformers_module.models.auto.processing_auto.PROCESSOR_MAPPING_NAMES:
			__UpperCAmelCase							=			'''AutoProcessor'''
		elif t in transformers_module.models.auto.tokenization_auto.TOKENIZER_MAPPING_NAMES:
			__UpperCAmelCase							=			'''AutoTokenizer'''
		elif t in transformers_module.models.auto.feature_extraction_auto.FEATURE_EXTRACTOR_MAPPING_NAMES:
			__UpperCAmelCase							=			'''AutoFeatureExtractor'''
		else:
			# Default to AutoTokenizer if a model has nothing, for backward compatibility.
			__UpperCAmelCase							=			'''AutoTokenizer'''

	__UpperCAmelCase							=			[processors[t] for t in all_models]

	return pd.DataFrame(UpperCamelCase__       )





def       lowerCAmelCase					(		UpperCamelCase__				:   List[str]       ):


	"""simple docstring"""
	__UpperCAmelCase							=			[
	    transformers_module.models.auto.modeling_auto,
	    transformers_module.models.auto.modeling_tf_auto,
	    transformers_module.models.auto.modeling_flax_auto,
	]
	for pipeline_tag, model_mapping, auto_class in PIPELINE_TAGS_AND_AUTO_MODELS:
		__UpperCAmelCase							=			[model_mapping, f"""TF_{model_mapping}""", f"""FLAX_{model_mapping}"""]
		__UpperCAmelCase							=			[auto_class, f"""TF_{auto_class}""", f"""Flax_{auto_class}"""]
		# Loop through all three frameworks
		for module, cls, mapping in zip(UpperCamelCase__    ,					UpperCamelCase__    ,					UpperCamelCase__       ):
			# The type of pipeline may not exist in this framework
			if not hasattr(UpperCamelCase__    ,					UpperCamelCase__       ):
				continue
			# First extract all model_names
			__UpperCAmelCase							=			[]
			for name in getattr(UpperCamelCase__    ,					UpperCamelCase__       ).values():
				if isinstance(UpperCamelCase__    ,					UpperCamelCase__       ):
					model_names.append(UpperCamelCase__       )
				else:
					model_names.extend(list(UpperCamelCase__       )       )

            # Add pipeline tag and auto model class for those models
			table.update({model_name: (pipeline_tag, cls) for model_name in model_names}       )

	return table





def       lowerCAmelCase					(		UpperCamelCase__				:   List[str]    ,					UpperCamelCase__				:   Dict       ):


	"""simple docstring"""
	__UpperCAmelCase							=			get_frameworks_table()
	__UpperCAmelCase							=			Dataset.from_pandas(UpperCamelCase__       )

	__UpperCAmelCase							=			hf_hub_download(
	    '''huggingface/transformers-metadata'''    ,					'''pipeline_tags.json'''    ,					repo_type='''dataset'''    ,					token=UpperCamelCase__       )
	__UpperCAmelCase							=			Dataset.from_json(UpperCamelCase__       )
	__UpperCAmelCase							=			{
	    tags_dataset[i]['''model_class''']: (tags_dataset[i]['''pipeline_tag'''], tags_dataset[i]['''auto_class'''])
	    for i in range(len(UpperCamelCase__       )       )
	}
	__UpperCAmelCase							=			update_pipeline_and_auto_class_table(UpperCamelCase__       )

	# Sort the model classes to avoid some nondeterministic updates to create false update commits.
	__UpperCAmelCase							=			sorted(table.keys()       )
	__UpperCAmelCase							=			pd.DataFrame(
	    {
	        '''model_class''': model_classes,
	        '''pipeline_tag''': [table[m][0] for m in model_classes],
	        '''auto_class''': [table[m][1] for m in model_classes],
	    }       )
	__UpperCAmelCase							=			Dataset.from_pandas(UpperCamelCase__       )

	with tempfile.TemporaryDirectory() as tmp_dir:
		frameworks_dataset.to_json(os.path.join(UpperCamelCase__    ,					'''frameworks.json'''       )       )
		tags_dataset.to_json(os.path.join(UpperCamelCase__    ,					'''pipeline_tags.json'''       )       )

		if commit_sha is not None:
			__UpperCAmelCase							=			(
			    f"""Update with commit {commit_sha}\n\nSee: """
			    f"""https://github.com/huggingface/transformers/commit/{commit_sha}"""
			)
		else:
			__UpperCAmelCase							=			'''Update'''

		upload_folder(
		    repo_id='''huggingface/transformers-metadata'''    ,					folder_path=UpperCamelCase__    ,					repo_type='''dataset'''    ,					token=UpperCamelCase__    ,					commit_message=UpperCamelCase__    ,					)





def       lowerCAmelCase					(		):


	"""simple docstring"""
	__UpperCAmelCase							=			{tag: cls for tag, _, cls in PIPELINE_TAGS_AND_AUTO_MODELS}
	__UpperCAmelCase							=			transformers_module.pipelines.SUPPORTED_TASKS
	__UpperCAmelCase							=			[]
	for key in pipeline_tasks:
		if key not in in_table:
			__UpperCAmelCase							=			pipeline_tasks[key]['''pt''']
			if isinstance(UpperCamelCase__    ,					(list, tuple)       ):
				__UpperCAmelCase							=			model[0]
			__UpperCAmelCase							=			model.__name__
			if model not in in_table.values():
				missing.append(UpperCamelCase__       )

	if len(UpperCamelCase__       ) > 0:
		__UpperCAmelCase							=			''', '''.join(UpperCamelCase__       )
		raise ValueError(
		    '''The following pipeline tags are not present in the `PIPELINE_TAGS_AND_AUTO_MODELS` constant inside '''
		    f"""`utils/update_metadata.py`: {msg}. Please add them!"""       )


if __name__ == "__main__":
					__lowerCAmelCase	:       Union[str, Any]						       =	argparse.ArgumentParser()
					parser.add_argument("--token", type=str, help="The token to use to push to the transformers-metadata dataset.")
					parser.add_argument("--commit_sha", type=str, help="The sha of the commit going with this update.")
					parser.add_argument("--check-only", action="store_true", help="Activate to just check all pipelines are present.")
					__lowerCAmelCase	:       Tuple						       =	parser.parse_args()

					if args.check_only:
										check_pipeline_tags()
					else:
										update_metadata(args.token, args.commit_sha)
