






"""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
UpperCAmelCase_     :				List[str]         =				"""src/transformers"""


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


# Regexes that match TF/Flax/PT model names.
UpperCAmelCase_     :				Union[str, Any]         =				re.compile(r"""TF(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)""")
UpperCAmelCase_     :				Dict         =				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.
UpperCAmelCase_     :				List[str]         =				re.compile(r"""(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)""")


# Fill this with tuples (pipeline_tag, model_mapping, auto_model)
UpperCAmelCase_     :				Any         =				[
    ("""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       _A    (__a							)				-> str:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	re.finditer('''.+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)'''					,       __a							)
    return [m.group(0							) for m in matches]



def       _A    ()				-> Optional[Any]:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_   :						List[Any]							=	transformers_module.models.auto.configuration_auto.CONFIG_MAPPING_NAMES
    SCREAMING_SNAKE_CASE_   :						List[Any]							=	{
        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.
    SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	collections.defaultdict(__a							)
    SCREAMING_SNAKE_CASE_   :						int							=	collections.defaultdict(__a							)
    SCREAMING_SNAKE_CASE_   :						List[str]							=	collections.defaultdict(__a							)

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

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

    SCREAMING_SNAKE_CASE_   :						List[Any]							=	set(list(pt_models.keys()							) + list(tf_models.keys()							) + list(flax_models.keys()							)							)
    SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	list(__a							)
    all_models.sort()

    SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	{'''model_type''': all_models}
    SCREAMING_SNAKE_CASE_   :						str							=	[pt_models[t] for t in all_models]
    SCREAMING_SNAKE_CASE_   :						List[Any]							=	[tf_models[t] for t in all_models]
    SCREAMING_SNAKE_CASE_   :						str							=	[flax_models[t] for t in all_models]

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

    SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	[processors[t] for t in all_models]

    return pd.DataFrame(__a							)



def       _A    (__a							)				-> List[Any]:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_   :						Any							=	[
        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:
        SCREAMING_SNAKE_CASE_   :						List[Any]							=	[model_mapping, f'TF_{model_mapping}', f'FLAX_{model_mapping}']
        SCREAMING_SNAKE_CASE_   :						List[Any]							=	[auto_class, f'TF_{auto_class}', f'Flax_{auto_class}']
        # Loop through all three frameworks
        for module, cls, mapping in zip(__a					,       __a					,       __a							):
            # The type of pipeline may not exist in this framework
            if not hasattr(__a					,       __a							):
                continue
            # First extract all model_names
            SCREAMING_SNAKE_CASE_   :						Tuple							=	[]
            for name in getattr(__a					,       __a							).values():
                if isinstance(__a					,       __a							):
                    model_names.append(__a							)
                else:
                    model_names.extend(list(__a							)							)

            # 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       _A    (__a					,       __a							)				-> int:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	get_frameworks_table()
    SCREAMING_SNAKE_CASE_   :						Any							=	Dataset.from_pandas(__a							)

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

    # Sort the model classes to avoid some nondeterministic updates to create false update commits.
    SCREAMING_SNAKE_CASE_   :						List[Any]							=	sorted(table.keys()							)
    SCREAMING_SNAKE_CASE_   :						Tuple							=	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],
        }							)
    SCREAMING_SNAKE_CASE_   :						List[str]							=	Dataset.from_pandas(__a							)

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

        if commit_sha is not None:
            SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	(
                f'Update with commit {commit_sha}\n\nSee: '
                f'https://github.com/huggingface/transformers/commit/{commit_sha}'
            )
        else:
            SCREAMING_SNAKE_CASE_   :						int							=	'''Update'''

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



def       _A    ()				-> Dict:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_   :						Dict							=	{tag: cls for tag, _, cls in PIPELINE_TAGS_AND_AUTO_MODELS}
    SCREAMING_SNAKE_CASE_   :						Tuple							=	transformers_module.pipelines.SUPPORTED_TASKS
    SCREAMING_SNAKE_CASE_   :						List[str]							=	[]
    for key in pipeline_tasks:
        if key not in in_table:
            SCREAMING_SNAKE_CASE_   :						int							=	pipeline_tasks[key]['''pt''']
            if isinstance(__a					,       (list, tuple)							):
                SCREAMING_SNAKE_CASE_   :						str							=	model[0]
            SCREAMING_SNAKE_CASE_   :						List[str]							=	model.__name__
            if model not in in_table.values():
                missing.append(__a							)

    if len(__a							) > 0:
        SCREAMING_SNAKE_CASE_   :						Optional[int]							=	''', '''.join(__a							)
        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__":
     UpperCAmelCase_     :				List[str]         =				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.""")
     UpperCAmelCase_     :				Optional[Any]         =				parser.parse_args()

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










"""simple docstring"""






import multiprocessing
import time

from arguments import PretokenizationArguments
from datasets import load_dataset

from transformers import AutoTokenizer, HfArgumentParser



def       _A    (__a							)				-> Optional[int]:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	{}
    SCREAMING_SNAKE_CASE_   :						List[Any]							=	tokenizer(example['''content''']					,       truncation=__a							)['''input_ids''']
    SCREAMING_SNAKE_CASE_   :						List[Any]							=	len(example['''content''']							) / len(output['''input_ids''']							)
    return output


UpperCAmelCase_     :				Tuple         =				HfArgumentParser(PretokenizationArguments)
UpperCAmelCase_     :				Tuple         =				parser.parse_args()
if args.num_workers is None:
  UpperCAmelCase_     :				Tuple         =				multiprocessing.cpu_count()
UpperCAmelCase_     :				str         =				AutoTokenizer.from_pretrained(args.tokenizer_dir)

UpperCAmelCase_     :				Tuple         =				time.time()
UpperCAmelCase_     :				List[Any]         =				load_dataset(args.dataset_name, split="""train""")
print(f'''Dataset loaded in {time.time()-t_start:.2f}s''')

UpperCAmelCase_     :				Dict         =				time.time()
UpperCAmelCase_     :				Dict         =				ds.map(
    tokenize,
    num_proc=args.num_workers,
    remove_columns=[
        """repo_name""",
        """path""",
        """copies""",
        """size""",
        """content""",
        """license""",
        """hash""",
        """line_mean""",
        """line_max""",
        """alpha_frac""",
        """autogenerated""",
    ],
)
print(f'''Dataset tokenized in {time.time()-t_start:.2f}s''')

UpperCAmelCase_     :				str         =				time.time()
ds.push_to_hub(args.tokenized_data_repo)





print(f'''Data pushed to the hub in {time.time()-t_start:.2f}s''')


