

# XXX: we want transformers master here - in the absense of conftest manipulating sys.path:
# hack it in for now:
import sys
from pathlib import Path


UpperCAmelCase_						=Path(__file__).resolve().parents[3] / """src"""
sys.path.insert(1, str(git_repo_path))

import dataclasses  # noqa
import io  # noqa
import itertools  # noqa
import json  # noqa
import os  # noqa
import unittest  # noqa
from copy import deepcopy  # noqa

from parameterized import parameterized  # noqa
from transformers import TrainingArguments, is_torch_available  # noqa
from transformers.deepspeed import is_deepspeed_available  # noqa
from transformers.file_utils import WEIGHTS_NAME  # noqa
from transformers.testing_utils import (  # noqa
    CaptureLogger,
    ExtendSysPath,
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
    mockenv_context,
    require_deepspeed,
    require_torch_gpu,
    require_torch_multi_gpu,
    slow,
)
from transformers.trainer_utils import set_seed  # noqa


set_seed(42)

UpperCAmelCase_						={"""base""": """patrickvonplaten/wav2vec2_tiny_random""", """robust""": """patrickvonplaten/wav2vec2_tiny_random_robust"""}

UpperCAmelCase_						="""zero2"""
UpperCAmelCase_						="""zero3"""
UpperCAmelCase_						=[ZEROa, ZEROa]



def       UpperCAmelCase       (     _snake_case						, _snake_case						, _snake_case							):
							lowerCAmelCase   =		parameterized.to_safe_name('''_'''.join(str(__a							) for x in param.args							)							)
							return F"""{func.__name__}_{param_based_name}"""


# Cartesian-product of zero stages with models to test
UpperCAmelCase_						=list(itertools.product(stages, models.keys()))


@slow
@require_deepspeed
@require_torch_gpu
class 							__UpperCamelCase			( __lowerCamelCase  ):





				'''simple docstring'''



				@parameterized.expand(SCREAMING_SNAKE_CASE_       , name_func=SCREAMING_SNAKE_CASE_						)
				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_						):
											self.run_and_check(
											    stage=SCREAMING_SNAKE_CASE_       , model=SCREAMING_SNAKE_CASE_       , distributed=SCREAMING_SNAKE_CASE_       , fpaa=SCREAMING_SNAKE_CASE_       , )



				@require_torch_multi_gpu
				@parameterized.expand(SCREAMING_SNAKE_CASE_       , name_func=SCREAMING_SNAKE_CASE_						)
				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_						):
											self.run_and_check(
											    stage=SCREAMING_SNAKE_CASE_       , model=SCREAMING_SNAKE_CASE_       , distributed=SCREAMING_SNAKE_CASE_       , fpaa=SCREAMING_SNAKE_CASE_       , )



				@parameterized.expand(SCREAMING_SNAKE_CASE_       , name_func=SCREAMING_SNAKE_CASE_						)
				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_						):
											self.run_and_check(
											    stage=SCREAMING_SNAKE_CASE_       , model=SCREAMING_SNAKE_CASE_       , distributed=SCREAMING_SNAKE_CASE_       , fpaa=SCREAMING_SNAKE_CASE_       , )



				@require_torch_multi_gpu
				@parameterized.expand(SCREAMING_SNAKE_CASE_       , name_func=SCREAMING_SNAKE_CASE_						)
				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_						):
											self.run_and_check(
											    stage=SCREAMING_SNAKE_CASE_       , model=SCREAMING_SNAKE_CASE_       , distributed=SCREAMING_SNAKE_CASE_       , fpaa=SCREAMING_SNAKE_CASE_       , )



				def       __snake_case		( self       , UpperCAmelCase_						):
											# XXX: run_asr is premature and doesn't save any results
											# so all we check for now is that the process didn't fail
											pass



				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_ = 10       , UpperCAmelCase_ = True       , UpperCAmelCase_ = True       , UpperCAmelCase_ = True       , ):
											lowerCAmelCase   =		models[model]

											lowerCAmelCase   =		self.run_trainer(
											    stage=SCREAMING_SNAKE_CASE_       , model_name=SCREAMING_SNAKE_CASE_       , eval_steps=SCREAMING_SNAKE_CASE_       , num_train_epochs=1       , distributed=SCREAMING_SNAKE_CASE_       , fpaa=SCREAMING_SNAKE_CASE_       , )

											self.do_checks(SCREAMING_SNAKE_CASE_						)

											return output_dir



				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_ = 10       , UpperCAmelCase_ = 1       , UpperCAmelCase_ = True       , UpperCAmelCase_ = True       , ):
											lowerCAmelCase   =		self.get_auto_remove_tmp_dir('''./xxx'''       , after=SCREAMING_SNAKE_CASE_						)
											lowerCAmelCase   =		F"""\n            --model_name_or_path {model_name}\n            --dataset_name hf-internal-testing/librispeech_asr_dummy\n            --dataset_config_name clean\n            --train_split_name validation\n            --validation_split_name validation\n            --output_dir {output_dir}\n            --num_train_epochs {str(SCREAMING_SNAKE_CASE_						)}\n            --per_device_train_batch_size 2\n            --per_device_eval_batch_size 2\n            --evaluation_strategy steps\n            --learning_rate 5e-4\n            --warmup_steps 8\n            --orthography timit\n            --preprocessing_num_workers 1\n            --group_by_length\n            --freeze_feature_extractor\n            --report_to none\n            --save_steps 0\n            --eval_steps {eval_steps}\n            --report_to none\n        """.split()

											if fpaa:
																		args.extend(['''--fp16''']						)

											# currently ds_config_wav2vec2_zero.json requires "zero_optimization.find_unused_parameters": true,
											# hence the separate config files
											lowerCAmelCase   =		F"""--deepspeed {self.test_file_dir_str}/ds_config_wav2vec2_{stage}.json""".split()
											lowerCAmelCase   =		[F"""{self.examples_dir_str}/research_projects/wav2vec2/run_asr.py"""]
											lowerCAmelCase   =		self.get_launcher(SCREAMING_SNAKE_CASE_						)

											lowerCAmelCase   =		launcher + script + args + ds_args
											# keep for quick debug
											# print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
											execute_subprocess_async(SCREAMING_SNAKE_CASE_       , env=self.get_env()						)

											return output_dir






				def       __snake_case		( self       , UpperCAmelCase_=False						):
											# 1. explicitly set --num_nodes=1 just in case these tests end up run on a multi-node setup
											# - it won't be able to handle that
											# 2. for now testing with just 2 gpus max (since some quality tests may give different
											# results with mode gpus because we use very little data)
											lowerCAmelCase   =		min(2       , get_gpu_count()						) if distributed else 1
											return F"""deepspeed --num_nodes 1 --num_gpus {num_gpus}""".split()









from collections import Counter
from pathlib import Path
from typing import Optional, Tuple

import yaml


class 							__UpperCamelCase			( yaml.SafeLoader  ):





				'''simple docstring'''



				def       __snake_case		( self       , UpperCAmelCase_						):
											lowerCAmelCase   =		[self.constructed_objects[key_node] for key_node, _ in node.value]
											lowerCAmelCase   =		[tuple(UpperCAmelCase_						) if isinstance(UpperCAmelCase_       , UpperCAmelCase_						) else key for key in keys]
											lowerCAmelCase   =		Counter(UpperCAmelCase_						)
											lowerCAmelCase   =		[key for key in counter if counter[key] > 1]
											if duplicate_keys:
																		raise TypeError(F"""Got duplicate yaml keys: {duplicate_keys}"""						)






				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_=False						):
											lowerCAmelCase   =		super().construct_mapping(UpperCAmelCase_       , deep=UpperCAmelCase_						)
											self._check_no_duplicates_on_constructed_node(UpperCAmelCase_						)
											return mapping




def       UpperCAmelCase       (     _snake_case							):
							lowerCAmelCase   =		list(readme_content.splitlines()							)
							if full_content and full_content[0] == "---" and "---" in full_content[1:]:
														lowerCAmelCase   =		full_content[1:].index('''---'''							) + 1
														lowerCAmelCase   =		'''\n'''.join(full_content[1:sep_idx]							)
														return yamlblock, "\n".join(full_content[sep_idx + 1 :]							)

							return None, "\n".join(_snake_case							)


class 							__UpperCamelCase			( __UpperCAmelCase  ):





				'''simple docstring'''



				__a							: Any										={"""train_eval_index"""}  # train-eval-index in the YAML metadata



				@classmethod
				def       __snake_case		( cls       , UpperCAmelCase_						):
											with open(UpperCAmelCase_       , encoding='''utf-8'''						) as readme_file:
																		lowerCAmelCase   ,     lowerCAmelCase   =		_split_yaml_from_readme(readme_file.read()						)
											if yaml_string is not None:
																		return cls.from_yaml_string(UpperCAmelCase_						)
											else:
																		return cls()



				def       __snake_case		( self       , UpperCAmelCase_						):
											if path.exists():
																		with open(UpperCAmelCase_       , encoding='''utf-8'''						) as readme_file:
																									lowerCAmelCase   =		readme_file.read()
											else:
																		lowerCAmelCase   =		None
											lowerCAmelCase   =		self._to_readme(UpperCAmelCase_						)
											with open(UpperCAmelCase_       , '''w'''       , encoding='''utf-8'''						) as readme_file:
																		readme_file.write(UpperCAmelCase_						)



				def       __snake_case		( self       , UpperCAmelCase_ = None						):
											if readme_content is not None:
																		lowerCAmelCase   ,     lowerCAmelCase   =		_split_yaml_from_readme(UpperCAmelCase_						)
																		lowerCAmelCase   =		'''---\n''' + self.to_yaml_string() + '''---\n''' + content
											else:
																		lowerCAmelCase   =		'''---\n''' + self.to_yaml_string() + '''---\n'''
											return full_content



				@classmethod
				def       __snake_case		( cls       , UpperCAmelCase_						):
											lowerCAmelCase   =		yaml.load(UpperCAmelCase_       , Loader=_NoDuplicateSafeLoader						) or {}

											# Convert the YAML keys to DatasetMetadata fields
											lowerCAmelCase   =		{
											    (key.replace('''-'''       , '''_'''						) if key.replace('''-'''       , '''_'''						) in cls._FIELDS_WITH_DASHES else key): value
											    for key, value in metadata_dict.items()
											}
											return cls(**UpperCAmelCase_						)






				def       __snake_case		( self						):
											return yaml.safe_dump(
											    {
											        (key.replace('''_'''       , '''-'''						) if key in self._FIELDS_WITH_DASHES else key): value
											        for key, value in self.items()
											    }       , sort_keys=UpperCAmelCase_       , allow_unicode=UpperCAmelCase_       , encoding='''utf-8'''       , ).decode('''utf-8'''						)




UpperCAmelCase_						={
    """image-classification""": [],
    """translation""": [],
    """image-segmentation""": [],
    """fill-mask""": [],
    """automatic-speech-recognition""": [],
    """token-classification""": [],
    """sentence-similarity""": [],
    """audio-classification""": [],
    """question-answering""": [],
    """summarization""": [],
    """zero-shot-classification""": [],
    """table-to-text""": [],
    """feature-extraction""": [],
    """other""": [],
    """multiple-choice""": [],
    """text-classification""": [],
    """text-to-image""": [],
    """text2text-generation""": [],
    """zero-shot-image-classification""": [],
    """tabular-classification""": [],
    """tabular-regression""": [],
    """image-to-image""": [],
    """tabular-to-text""": [],
    """unconditional-image-generation""": [],
    """text-retrieval""": [],
    """text-to-speech""": [],
    """object-detection""": [],
    """audio-to-audio""": [],
    """text-generation""": [],
    """conversational""": [],
    """table-question-answering""": [],
    """visual-question-answering""": [],
    """image-to-text""": [],
    """reinforcement-learning""": [],
    """voice-activity-detection""": [],
    """time-series-forecasting""": [],
    """document-question-answering""": [],
}


if __name__ == "__main__":
		from argparse import ArgumentParser

		UpperCAmelCase_						=ArgumentParser(usage="""Validate the yaml metadata block of a README.md file.""")
		ap.add_argument("""readme_filepath""")
		UpperCAmelCase_						=ap.parse_args()

		UpperCAmelCase_						=Path(args.readme_filepath)
		UpperCAmelCase_						=DatasetMetadata.from_readme(readme_filepath)
		print(dataset_metadata)
		dataset_metadata.to_readme(readme_filepath)



