


import itertools
import os
import subprocess
from os.path import dirname

from parameterized import parameterized

from tests.trainer.test_trainer import TrainerIntegrationCommon  # noqa
from transformers import is_torch_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
    get_tests_dir,
    require_deepspeed,
    require_torch_gpu,
    slow,
)
from transformers.trainer_utils import set_seed


if is_torch_available():
						from tests.trainer.test_trainer import (  # noqa
						    RegressionModelConfig,
						    RegressionPreTrainedModel,
						    get_regression_trainer,
						)


set_seed(42)

snake_case			:      Optional[Any] 							=   get_tests_dir('fixtures')
snake_case			:      str 							=   os.path.join(dirname(get_tests_dir()))
snake_case			:      Union[str, Any] 							=   dirname(os.path.abspath(__file__))

# default torch.distributed port
snake_case			:      Union[str, Any] 							=   '10999'

snake_case			:      Optional[Any] 							=   't5-small'

# *** Working Models ***
snake_case			:      Tuple 							=   'hf-internal-testing/tiny-albert'
snake_case			:      Dict 							=   'sshleifer/bart-tiny-random'
snake_case			:      Union[str, Any] 							=   'hf-internal-testing/tiny-bert'
snake_case			:      str 							=   'hf-internal-testing/tiny-random-bigbird_pegasus'
snake_case			:      Optional[Any] 							=   'hf-internal-testing/tiny-random-big_bird'
snake_case			:      Optional[Any] 							=   'hf-internal-testing/tiny-random-blenderbot'
snake_case			:      List[str] 							=   'bigscience/bigscience-small-testing'
snake_case			:      List[str] 							=   'hf-internal-testing/tiny-random-deberta'
snake_case			:      int 							=   'hf-internal-testing/tiny-random-deberta-v2'
snake_case			:      Union[str, Any] 							=   'sshleifer/tiny-distilbert-base-cased'
snake_case			:      List[str] 							=   'hf-internal-testing/tiny-electra'
snake_case			:      Tuple 							=   'hf-internal-testing/tiny-random-flaubert'
snake_case			:      Any 							=   'stas/tiny-wmt19-en-de'
snake_case			:      Optional[Any] 							=   'hf-internal-testing/tiny-random-funnel'
snake_case			:      Optional[Any] 							=   'sshleifer/tiny-gpt2'
snake_case			:      str 							=   'hf-internal-testing/tiny-random-gptj'
snake_case			:      str 							=   'hf-internal-testing/tiny-random-gpt_neo'
snake_case			:      List[Any] 							=   'hf-internal-testing/tiny-layoutlm'
snake_case			:      Optional[Any] 							=   'hf-internal-testing/tiny-random-led'
snake_case			:      Optional[Any] 							=   'hf-internal-testing/tiny-random-longformer'
snake_case			:      str 							=   'stas/tiny-m2m_100'  # hf tiny model is unsuitable
snake_case			:      Dict 							=   'sshleifer/tiny-marian-en-de'
snake_case			:      int 							=   'sshleifer/tiny-mbart'
snake_case			:      Dict 							=   'hf-internal-testing/tiny-random-mobilebert'
snake_case			:      Tuple 							=   'hf-internal-testing/tiny-random-mpnet'
snake_case			:      Any 							=   'stas/pegasus-cnn_dailymail-tiny-random'
snake_case			:      List[str] 							=   'hf-internal-testing/tiny-random-prophetnet'
snake_case			:      str 							=   'sshleifer/tiny-distilroberta-base'
snake_case			:      List[Any] 							=   'hf-internal-testing/tiny-random-squeezebert'
snake_case			:      List[str] 							=   'patrickvonplaten/t5-tiny-random'
snake_case			:      List[Any] 							=   'hf-internal-testing/tiny-random-t5-v1.1'
snake_case			:      List[str] 							=   'hf-internal-testing/tiny-random-vit'
snake_case			:      Union[str, Any] 							=   'hf-internal-testing/tiny-xlm-roberta'
snake_case			:      Any 							=   'sshleifer/tiny-xlnet-base-cased'


# *** To Fix ***


# *** tiny model issues ***
# missing model files:
snake_case			:      Tuple 							=   'hf-internal-testing/tiny-random-mt5'
snake_case			:      Optional[int] 							=   'hf-internal-testing/tiny-random-camembert'
snake_case			:      Optional[int] 							=   'hf-internal-testing/tiny-random-openai-gpt'

# missing tokenizer files
snake_case			:      Union[str, Any] 							=   'hf-internal-testing/tiny-random-convbert'
snake_case			:      str 							=   'hf-internal-testing/tiny-random-layoutlmv2'
snake_case			:      List[str] 							=   'hf-internal-testing/tiny-random-hubert'

# issues with tokenizer
snake_case			:      Tuple 							=   'hf-internal-testing/tiny-random-ctrl'
snake_case			:      List[str] 							=   'hf-internal-testing/tiny-random-transfo-xl'  # same as ctrl

# other issues with tiny models
snake_case			:      int 							=   'hf-internal-testing/tiny-random-ibert'  # multiple issues with either mlm/qa/clas
snake_case			:      Union[str, Any] 							=   'hf-internal-testing/tiny-random-reformer'  # multiple issues with either mlm/qa/clas

# *** Lacking official examples to test with ***
# or not working with examples
snake_case			:      Any 							=   'hf-internal-testing/tiny-random-dpr'
# - "dpr"  examples/research_projects/rag-end2end-retriever/
snake_case			:      Optional[Any] 							=   'hf-internal-testing/tiny-random-rag'
# - "rag" research_projects
snake_case			:      List[Any] 							=   ''
# - "luke" Entities classes - no plan to make such example
snake_case			:      Optional[int] 							=   'hf-internal-testing/tiny-random-lxmert'
# - "lxmert" doesn't work with run_qa.py
snake_case			:      Tuple 							=   'hf-internal-testing/tiny-random-clip'
# - "clip" nothing under pytorch examples - XXX: Suraj is working on adding some - check by end of Sep
snake_case			:      Optional[Any] 							=   'hf-internal-testing/tiny-random-speech_to_text'
# - "speech_to_text", nothing under pytorch examples


# *** Reactive mode ***
# models with low usage, unstable API, things about to change - do nothing about the following until someone runs into a problem
snake_case			:      str 							=   'hf-internal-testing/tiny-random-tapas'


def 				SCREAMING_SNAKE_CASE (     UpperCAmelCase__=False							):

				"""simple docstring"""






				_SCREAMING_SNAKE_CASE      = min(2   ,get_gpu_count()							) if distributed else 1
				_SCREAMING_SNAKE_CASE      = os.environ.get('DS_TEST_PORT'   ,UpperCAmelCase__							)
				return f'''deepspeed --num_nodes 1 --num_gpus {num_gpus} --master_port {master_port}'''.split()


def 				SCREAMING_SNAKE_CASE (     ):

				"""simple docstring"""






				_SCREAMING_SNAKE_CASE      = f'''{FIXTURE_DIRECTORY}/tests_samples'''
				_SCREAMING_SNAKE_CASE      = f'''{data_dir_samples}/wmt_en_ro'''
				_SCREAMING_SNAKE_CASE      = f'''{data_dir_samples}/xsum'''
				_SCREAMING_SNAKE_CASE      = '\n        --do_train\n        --max_train_samples 4\n        --per_device_train_batch_size 2\n        --num_train_epochs 1\n        --fp16\n        --report_to none\n        --overwrite_output_dir\n        '.split()

				# try to cover as many models as possible once (it's enough to run on one task per model)
				# but need a tiny model for each
				#
				# should have "{model_type.upper()}_TINY" corresponding vars defined, e.g., T5_TINY, etc.
				_SCREAMING_SNAKE_CASE      = {
				    'trans': [
				        'bart',
				        'fsmt',
				        'm2m_100',
				        'marian',
				        'mbart',
				        't5',
				        't5_v1',
				        # "mt5", missing model files
				    ],
				    'sum': [
				        'pegasus',
				    ],
				    'clm': [
				        'big_bird',
				        'bigbird_pegasus',
				        'blenderbot',
				        'bloom',
				        'gpt2',
				        'gpt_neo',
				        'gptj',
				        'xlm-roberta',
				        'prophetnet',
				        # "camembert", missing model files
				    ],
				    'mlm': [
				        'albert',
				        'deberta',
				        'deberta-v2',
				        'distilbert',
				        'electra',
				        'flaubert',
				        'funnel',
				        'layoutlm',
				        # "reformer", # multiple issues with either mlm/qa/clas
				    ],
				    'qa': [
				        'led',
				        'longformer',
				        'mobilebert',
				        'mpnet',
				        'roberta',
				        'squeezebert',
				        # "convbert", # missing tokenizer files
				        # "layoutlmv2", missing model files
				    ],
				    'clas': [
				        'bert',
				        'xlnet',
				        # "hubert", # missing tokenizer files
				        # "ibert", # multiple issues with either mlm/qa/clas
				        # "transfo-xl", # tokenizer issues as ctrl
				        # "ctrl", # tokenizer issues
				        # "openai-gpt", missing model files
				        # "tapas", multiple issues
				    ],
				    'img_clas': [
				        'vit',
				    ],
				}

				_SCREAMING_SNAKE_CASE      = f'''{ROOT_DIRECTORY}/examples/pytorch'''

				_SCREAMING_SNAKE_CASE      = {
				    'trans': f'''
        {scripts_dir}/translation/run_translation.py
        --train_file {data_dir_wmt}/train.json
        --source_lang en
        --target_lang ro
        ''',
				    'sum': f'''
        {scripts_dir}/summarization/run_summarization.py
        --train_file {data_dir_xsum}/sample.json
        --max_source_length 12
        --max_target_length 12
        --lang en
        ''',
				    'clm': f'''
        {scripts_dir}/language-modeling/run_clm.py
        --train_file {FIXTURE_DIRECTORY}/sample_text.txt
        --block_size 8
        ''',
				    'mlm': f'''
        {scripts_dir}/language-modeling/run_mlm.py
        --train_file {FIXTURE_DIRECTORY}/sample_text.txt
        ''',
				    'qa': f'''
        {scripts_dir}/question-answering/run_qa.py
        --train_file {data_dir_samples}/SQUAD/sample.json
        ''',
				    'clas': f'''
        {scripts_dir}/text-classification/run_glue.py
        --train_file {data_dir_samples}/MRPC/train.csv
        --max_seq_length 12
        --task_name MRPC
        ''',
				    'img_clas': f'''
        {scripts_dir}/image-classification/run_image_classification.py
            --dataset_name hf-internal-testing/cats_vs_dogs_sample
            --remove_unused_columns False
            --max_steps 10
            --image_processor_name {DS_TESTS_DIRECTORY}/vit_feature_extractor.json
        ''',
				}

				_SCREAMING_SNAKE_CASE      = get_launcher(distributed=UpperCAmelCase__							)

				_SCREAMING_SNAKE_CASE      = {}
				for task, args in tasks.items():
								_SCREAMING_SNAKE_CASE      = args.split()
								for model in tasksamodels[task]:
												_SCREAMING_SNAKE_CASE      = globals()[f'''{model.upper().replace("-"   ,"_"							)}_TINY''']
												_SCREAMING_SNAKE_CASE      = f'''--model_name_or_path {model_name}'''.split()
												_SCREAMING_SNAKE_CASE      = launcher + args + args_model + args_main

												# # generation special case
												# if task == "gen":
												#     launcher = f"deepspeed --num_nodes 1 --num_gpus 1".split()
												#     args_model += f"--model_type {model}".split()
												#     cmds[f"{task}_{model}"] = launcher + args + args_model
												# else:

				return cmds


snake_case			:      Union[str, Any] 							=   make_task_cmds()

snake_case			:      Tuple 							=   'zero2'
snake_case			:      str 							=   'zero3'

snake_case			:      str 							=   [ZEROa, ZEROa]


def 				SCREAMING_SNAKE_CASE (     UpperCAmelCase__   ,UpperCAmelCase__   ,UpperCAmelCase__							):

				"""simple docstring"""






				_SCREAMING_SNAKE_CASE      = parameterized.to_safe_name('_'.join(str(UpperCAmelCase__							) for x in param.args							)							)
				return f'''{func.__name__}_{param_based_name}'''


# Cartesian-product of zero stages with models to test
snake_case			:      List[Any] 							=   list(itertools.product(stages, task_cmds.keys()))



@slow
@require_deepspeed
@require_torch_gpu
class    __lowercase  (       UpperCamelCase							):

							"""simple docstring"""


							def 		__magic_name__			(    self				, A_				, A_	)->	Optional[int]:
											# return a ready to run train cmd
											if task not in task_cmds:
															raise ValueError(F'''don\'t know of task {task}, have {task_cmds.keys()}'''	)

											_SCREAMING_SNAKE_CASE      = task_cmds[task]
											_SCREAMING_SNAKE_CASE      = F'''--deepspeed {self.test_file_dir_str}/ds_config_{stage}.json'''.split()

											_SCREAMING_SNAKE_CASE      = self.get_auto_remove_tmp_dir()
											_SCREAMING_SNAKE_CASE      = F'''--output_dir {output_dir}'''.split()

											cmd += args_ds + args_out

											return cmd, output_dir





							@parameterized.expand(A_				, name_func=A_	)
							def 		__magic_name__			(    self				, A_				, A_	)->	Tuple:
											# testing the ability to do a run followed by recovery of full fp32 weights

											_SCREAMING_SNAKE_CASE      ,      _SCREAMING_SNAKE_CASE      = self.get_task_cmd(A_				, A_	)

											# 1. generate the checkpoint
											cmd += "--save_steps 1".split()
											# keep for quick debug
											# print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] + cmd)); die
											execute_subprocess_async(A_				, env=self.get_env()	)

											# 2. test that the fp32 weights get reconsolidated
											_SCREAMING_SNAKE_CASE      = F'''{output_dir}/checkpoint-1'''
											_SCREAMING_SNAKE_CASE      = F'''{chkpt_dir}/out.bin'''
											_SCREAMING_SNAKE_CASE      = F'''{chkpt_dir}/zero_to_fp32.py {chkpt_dir} {recovered_model_path}'''
											# keep for quick debug
											# print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
											subprocess.check_call(A_				, shell=A_	)
											assert os.path.exists(A_	), F'''{recovered_model_path} was not found'''

											# possibly could also test that the resulting saved model is usable but given that we use
											# random models we won't know if it's any good




import sys
from dataclasses import dataclass
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union

import pandas as pd
import pyarrow as pa

import datasets
import datasets.config
from datasets.features.features import require_storage_cast
from datasets.table import table_cast


if TYPE_CHECKING:
							import sqlitea

							import sqlalchemy


snake_case			:      List[str] 							=   datasets.utils.logging.get_logger(__name__)



@dataclass
class    __lowercase  (       datasets.BuilderConfig							):

							"""simple docstring"""


							SCREAMING_SNAKE_CASE   :					Union[str, "sqlalchemy.sql.Selectable"]        =						None
							SCREAMING_SNAKE_CASE   :					Union[str, "sqlalchemy.engine.Connection", "sqlalchemy.engine.Engine", "sqlite3.Connection"]        =						None
							SCREAMING_SNAKE_CASE   :					Optional[Union[str, List[str]]]        =						None
							SCREAMING_SNAKE_CASE   :					bool        =						True
							SCREAMING_SNAKE_CASE   :					Optional[Union[List, Tuple, Dict]]        =						None
							SCREAMING_SNAKE_CASE   :					Optional[Union[List, Dict]]        =						None
							SCREAMING_SNAKE_CASE   :					Optional[List[str]]        =						None
							SCREAMING_SNAKE_CASE   :					Optional[int]        =						1_00_00
							SCREAMING_SNAKE_CASE   :					Optional[datasets.Features]        =						None





							def 		__magic_name__			(    self	)->	int:
											if self.sql is None:
															raise ValueError('sql must be specified'	)
											if self.con is None:
															raise ValueError('con must be specified'	)





							def 		__magic_name__			(    self				, A_				, A_ = None				, )->	str:
											_SCREAMING_SNAKE_CASE      = config_kwargs.copy()
											# We need to stringify the Selectable object to make its hash deterministic

											# The process of stringifying is explained here: http://docs.sqlalchemy.org/en/latest/faq/sqlexpressions.html
											_SCREAMING_SNAKE_CASE      = config_kwargs['sql']
											if not isinstance(A_				, A_	):
															if datasets.config.SQLALCHEMY_AVAILABLE and "sqlalchemy" in sys.modules:
																			import sqlalchemy

																			if isinstance(A_				, sqlalchemy.sql.Selectable	):
																							_SCREAMING_SNAKE_CASE      = sqlalchemy.create_engine(config_kwargs['con'].split('://'	)[0] + '://'	)
																							_SCREAMING_SNAKE_CASE      = str(sql.compile(dialect=engine.dialect	)	)
																							_SCREAMING_SNAKE_CASE      = sql_str
																			else:
																							raise TypeError(
																							    F'''Supported types for \'sql\' are string and sqlalchemy.sql.Selectable but got {type(A_	)}: {sql}'''	)
															else:
																			raise TypeError(
																			    F'''Supported types for \'sql\' are string and sqlalchemy.sql.Selectable but got {type(A_	)}: {sql}'''	)
											_SCREAMING_SNAKE_CASE      = config_kwargs['con']
											if not isinstance(A_				, A_	):
															_SCREAMING_SNAKE_CASE      = id(A_	)
															logger.info(
															    F'''SQL connection \'con\' of type {type(A_	)} couldn\'t be hashed properly. To enable hashing, specify \'con\' as URI string instead.'''	)

											return super().create_config_id(A_				, custom_features=A_	)





							@property
							def 		__magic_name__			(    self	)->	Dict:
											_SCREAMING_SNAKE_CASE      = {
											    'index_col': self.index_col,
											    'columns': self.columns,
											    'params': self.params,
											    'coerce_float': self.coerce_float,
											    'parse_dates': self.parse_dates,
											}
											return pd_read_sql_kwargs







class    __lowercase  (       datasets.ArrowBasedBuilder							):

							"""simple docstring"""


							SCREAMING_SNAKE_CASE   :					Any        =						SqlConfig





							def 		__magic_name__			(    self	)->	Dict:
											return datasets.DatasetInfo(features=self.config.features	)





							def 		__magic_name__			(    self				, A_	)->	Optional[Any]:
											return [datasets.SplitGenerator(name=datasets.Split.TRAIN				, gen_kwargs={}	)]





							def 		__magic_name__			(    self				, A_	)->	pa.Table:
											if self.config.features is not None:
															_SCREAMING_SNAKE_CASE      = self.config.features.arrow_schema
															if all(not require_storage_cast(A_	) for feature in self.config.features.values()	):
																			# cheaper cast
																			_SCREAMING_SNAKE_CASE      = pa.Table.from_arrays([pa_table[field.name] for field in schema]				, schema=A_	)
															else:
																			# more expensive cast; allows str <-> int/float or str to Audio for example
																			_SCREAMING_SNAKE_CASE      = table_cast(A_				, A_	)
											return pa_table





							def 		__magic_name__			(    self	)->	Optional[int]:
											_SCREAMING_SNAKE_CASE      = self.config.chunksize
											_SCREAMING_SNAKE_CASE      = pd.read_sql(
											    self.config.sql				, self.config.con				, chunksize=A_				, **self.config.pd_read_sql_kwargs	)
											_SCREAMING_SNAKE_CASE      = [sql_reader] if chunksize is None else sql_reader
											for chunk_idx, df in enumerate(A_	):
															_SCREAMING_SNAKE_CASE      = pa.Table.from_pandas(A_	)
															yield chunk_idx, self._cast_table(A_	)
