



import shutil
import tempfile
import unittest
from unittest.mock import patch

from transformers import (
    DefaultFlowCallback,
    IntervalStrategy,
    PrinterCallback,
    ProgressCallback,
    Trainer,
    TrainerCallback,
    TrainingArguments,
    is_torch_available,
)
from transformers.testing_utils import require_torch


if is_torch_available():
		from transformers.trainer import DEFAULT_CALLBACKS

		from .test_trainer import RegressionDataset, RegressionModelConfig, RegressionPreTrainedModel






class      UpperCamelCase    (						_a     ):



	def __init__(       self		):
							A__							=       []



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	**UpperCAmelCase__		):
							self.events.append("on_init_end"		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	**UpperCAmelCase__		):
							self.events.append("on_train_begin"		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	**UpperCAmelCase__		):
							self.events.append("on_train_end"		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	**UpperCAmelCase__		):
							self.events.append("on_epoch_begin"		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	**UpperCAmelCase__		):
							self.events.append("on_epoch_end"		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	**UpperCAmelCase__		):
							self.events.append("on_step_begin"		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	**UpperCAmelCase__		):
							self.events.append("on_step_end"		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	**UpperCAmelCase__		):
							self.events.append("on_evaluate"		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	**UpperCAmelCase__		):
							self.events.append("on_predict"		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	**UpperCAmelCase__		):
							self.events.append("on_save"		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	**UpperCAmelCase__		):
							self.events.append("on_log"		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	**UpperCAmelCase__		):
							self.events.append("on_prediction_step"		)






@require_torch
class      UpperCamelCase    (						unittest.TestCase     ):



	def 			__A							(       self		):
							A__							=       tempfile.mkdtemp()



	def 			__A							(       self		):
							shutil.rmtree(self.output_dir		)



	def 			__A							(       self			,	UpperCAmelCase__=0			,	UpperCAmelCase__=0			,	UpperCAmelCase__=64			,	UpperCAmelCase__=64			,	UpperCAmelCase__=None			,	UpperCAmelCase__=False			,	**UpperCAmelCase__		):
							# disable_tqdm in TrainingArguments has a flaky default since it depends on the level of logging. We make sure
							# its set to False since the tests later on depend on its value.
							A__							=       RegressionDataset(length=snake_case_		)
							A__							=       RegressionDataset(length=snake_case_		)
							A__							=       RegressionModelConfig(a=snake_case_			,	b=snake_case_		)
							A__							=       RegressionPreTrainedModel(snake_case_		)

							A__							=       TrainingArguments(self.output_dir			,	disable_tqdm=snake_case_			,	report_to=[]			,	**snake_case_		)
							return Trainer(
							    snake_case_			,	snake_case_			,	train_dataset=snake_case_			,	eval_dataset=snake_case_			,	callbacks=snake_case_			,	)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__		):
							self.assertEqual(len(snake_case_		)			,	len(snake_case_		)		)

							# Order doesn't matter
							A__							=       sorted(snake_case_			,	key=lambda UpperCAmelCase__		: cb.__name__ if isinstance(snake_case_			,	snake_case_		) else cb.__class__.__name__		)
							A__							=       sorted(snake_case_			,	key=lambda UpperCAmelCase__		: cb.__name__ if isinstance(snake_case_			,	snake_case_		) else cb.__class__.__name__		)

							for cba, cba in zip(snake_case_			,	snake_case_		):
													if isinstance(snake_case_			,	snake_case_		) and isinstance(snake_case_			,	snake_case_		):
																			self.assertEqual(snake_case_			,	snake_case_		)
													elif isinstance(snake_case_			,	snake_case_		) and not isinstance(snake_case_			,	snake_case_		):
																			self.assertEqual(snake_case_			,	cba.__class__		)
													elif not isinstance(snake_case_			,	snake_case_		) and isinstance(snake_case_			,	snake_case_		):
																			self.assertEqual(cba.__class__			,	snake_case_		)
													else:
																			self.assertEqual(snake_case_			,	snake_case_		)



	def 			__A							(       self			,	UpperCAmelCase__		):
							A__							=       ["""on_init_end""", """on_train_begin"""]
							A__							=       0
							A__							=       len(trainer.get_eval_dataloader()		)
							A__							=       ["""on_prediction_step"""] * len(trainer.get_eval_dataloader()		) + ["""on_log""", """on_evaluate"""]
							for _ in range(trainer.state.num_train_epochs		):
													expected_events.append("on_epoch_begin"		)
													for _ in range(snake_case_		):
																			step += 1
																			expected_events += ["on_step_begin", "on_step_end"]
																			if step % trainer.args.logging_steps == 0:
																									expected_events.append("on_log"		)
																			if trainer.args.evaluation_strategy == IntervalStrategy.STEPS and step % trainer.args.eval_steps == 0:
																									expected_events += evaluation_events.copy()
																			if step % trainer.args.save_steps == 0:
																									expected_events.append("on_save"		)
													expected_events.append("on_epoch_end"		)
													if trainer.args.evaluation_strategy == IntervalStrategy.EPOCH:
																			expected_events += evaluation_events.copy()
							expected_events += ["on_log", "on_train_end"]
							return expected_events



	def 			__A							(       self		):
							A__							=       self.get_trainer()
							A__							=       DEFAULT_CALLBACKS.copy() + [ProgressCallback]
							self.check_callbacks_equality(trainer.callback_handler.callbacks			,	snake_case_		)

							# Callbacks passed at init are added to the default callbacks
							A__							=       self.get_trainer(callbacks=[MyTestTrainerCallback]		)
							expected_callbacks.append(snake_case_		)
							self.check_callbacks_equality(trainer.callback_handler.callbacks			,	snake_case_		)

							# TrainingArguments.disable_tqdm controls if use ProgressCallback or PrinterCallback
							A__							=       self.get_trainer(disable_tqdm=snake_case_		)
							A__							=       DEFAULT_CALLBACKS.copy() + [PrinterCallback]
							self.check_callbacks_equality(trainer.callback_handler.callbacks			,	snake_case_		)



	def 			__A							(       self		):
							A__							=       DEFAULT_CALLBACKS.copy() + [ProgressCallback]
							A__							=       self.get_trainer()

							# We can add, pop, or remove by class name
							trainer.remove_callback(snake_case_		)
							expected_callbacks.remove(snake_case_		)
							self.check_callbacks_equality(trainer.callback_handler.callbacks			,	snake_case_		)

							A__							=       self.get_trainer()
							A__							=       trainer.pop_callback(snake_case_		)
							self.assertEqual(cb.__class__			,	snake_case_		)
							self.check_callbacks_equality(trainer.callback_handler.callbacks			,	snake_case_		)

							trainer.add_callback(snake_case_		)
							expected_callbacks.insert(0			,	snake_case_		)
							self.check_callbacks_equality(trainer.callback_handler.callbacks			,	snake_case_		)

							# We can also add, pop, or remove by instance
							A__							=       self.get_trainer()
							A__							=       trainer.callback_handler.callbacks[0]
							trainer.remove_callback(snake_case_		)
							expected_callbacks.remove(snake_case_		)
							self.check_callbacks_equality(trainer.callback_handler.callbacks			,	snake_case_		)

							A__							=       self.get_trainer()
							A__							=       trainer.callback_handler.callbacks[0]
							A__							=       trainer.pop_callback(snake_case_		)
							self.assertEqual(snake_case_			,	snake_case_		)
							self.check_callbacks_equality(trainer.callback_handler.callbacks			,	snake_case_		)

							trainer.add_callback(snake_case_		)
							expected_callbacks.insert(0			,	snake_case_		)
							self.check_callbacks_equality(trainer.callback_handler.callbacks			,	snake_case_		)



	def 			__A							(       self		):
							import warnings

							# XXX: for now ignore scatter_gather warnings in this test since it's not relevant to what's being tested
							warnings.simplefilter(action="ignore"			,	category=snake_case_		)

							A__							=       self.get_trainer(callbacks=[MyTestTrainerCallback]		)
							trainer.train()
							A__							=       trainer.callback_handler.callbacks[-2].events
							self.assertEqual(snake_case_			,	self.get_expected_events(snake_case_		)		)

							# Independent log/save/eval
							A__							=       self.get_trainer(callbacks=[MyTestTrainerCallback]			,	logging_steps=5		)
							trainer.train()
							A__							=       trainer.callback_handler.callbacks[-2].events
							self.assertEqual(snake_case_			,	self.get_expected_events(snake_case_		)		)

							A__							=       self.get_trainer(callbacks=[MyTestTrainerCallback]			,	save_steps=5		)
							trainer.train()
							A__							=       trainer.callback_handler.callbacks[-2].events
							self.assertEqual(snake_case_			,	self.get_expected_events(snake_case_		)		)

							A__							=       self.get_trainer(callbacks=[MyTestTrainerCallback]			,	eval_steps=5			,	evaluation_strategy="steps"		)
							trainer.train()
							A__							=       trainer.callback_handler.callbacks[-2].events
							self.assertEqual(snake_case_			,	self.get_expected_events(snake_case_		)		)

							A__							=       self.get_trainer(callbacks=[MyTestTrainerCallback]			,	evaluation_strategy="epoch"		)
							trainer.train()
							A__							=       trainer.callback_handler.callbacks[-2].events
							self.assertEqual(snake_case_			,	self.get_expected_events(snake_case_		)		)

							# A bit of everything
							A__							=       self.get_trainer(
							    callbacks=[MyTestTrainerCallback]			,	logging_steps=3			,	save_steps=10			,	eval_steps=5			,	evaluation_strategy="steps"			,	)
							trainer.train()
							A__							=       trainer.callback_handler.callbacks[-2].events
							self.assertEqual(snake_case_			,	self.get_expected_events(snake_case_		)		)

							# warning should be emitted for duplicated callbacks
							with patch("transformers.trainer_callback.logger.warning"		) as warn_mock:
													A__							=       self.get_trainer(
													    callbacks=[MyTestTrainerCallback, MyTestTrainerCallback]			,	)
													assert str(snake_case_		) in warn_mock.call_args[0][0]










import multiprocessing
import os
from typing import BinaryIO, Optional, Union

import fsspec

from .. import Dataset, Features, NamedSplit, config
from ..formatting import query_table
from ..packaged_modules.json.json import Json
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader






class      UpperCamelCase    (						_UpperCAmelCase     ):



	def __init__(       self			,	UpperCAmelCase__			,	UpperCAmelCase__ = None			,	UpperCAmelCase__ = None			,	UpperCAmelCase__ = None			,	UpperCAmelCase__ = False			,	UpperCAmelCase__ = False			,	UpperCAmelCase__ = None			,	UpperCAmelCase__ = None			,	**UpperCAmelCase__			,	):
							super().__init__(
							    UpperCAmelCase__			,	split=UpperCAmelCase__			,	features=UpperCAmelCase__			,	cache_dir=UpperCAmelCase__			,	keep_in_memory=UpperCAmelCase__			,	streaming=UpperCAmelCase__			,	num_proc=UpperCAmelCase__			,	**UpperCAmelCase__			,	)
							A__							=       field
							A__							=       path_or_paths if isinstance(UpperCAmelCase__			,	UpperCAmelCase__		) else {self.split: path_or_paths}
							A__							=       Json(
							    cache_dir=UpperCAmelCase__			,	data_files=UpperCAmelCase__			,	features=UpperCAmelCase__			,	field=UpperCAmelCase__			,	**UpperCAmelCase__			,	)



	def 			__A							(       self		):
							# Build iterable dataset
							if self.streaming:
													A__							=       self.builder.as_streaming_dataset(split=self.split		)
							# Build regular (map-style) dataset
							else:
													A__							=       None
													A__							=       None
													A__							=       None
													A__							=       None

													self.builder.download_and_prepare(
													    download_config=UpperCAmelCase__			,	download_mode=UpperCAmelCase__			,	verification_mode=UpperCAmelCase__			,	base_path=UpperCAmelCase__			,	num_proc=self.num_proc			,	)
													A__							=       self.builder.as_dataset(
													    split=self.split			,	verification_mode=UpperCAmelCase__			,	in_memory=self.keep_in_memory		)
							return dataset






class      UpperCamelCase    :



	def __init__(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__ = None			,	UpperCAmelCase__ = None			,	**UpperCAmelCase__			,	):
							if num_proc is not None and num_proc <= 0:
													raise ValueError(F"""num_proc {num_proc} must be an integer > 0."""		)

							A__							=       dataset
							A__							=       path_or_buf
							A__							=       batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
							A__							=       num_proc
							A__							=       "utf-8"
							A__							=       to_json_kwargs



	def 			__A							(       self		):
							A__							=       self.to_json_kwargs.pop("path_or_buf"			,	UpperCAmelCase__		)
							A__							=       self.to_json_kwargs.pop("orient"			,	"records"		)
							A__							=       self.to_json_kwargs.pop("lines"			,	True if orient == "records" else False		)
							A__							=       self.to_json_kwargs.pop("index"			,	False if orient in ["split", "table"] else True		)
							A__							=       self.to_json_kwargs.pop("compression"			,	UpperCAmelCase__		)

							if compression not in [None, "infer", "gzip", "bz2", "xz"]:
													raise NotImplementedError(F"""`datasets` currently does not support {compression} compression"""		)

							if isinstance(self.path_or_buf			,	(str, bytes, os.PathLike)		):
													with fsspec.open(self.path_or_buf			,	"wb"			,	compression=UpperCAmelCase__		) as buffer:
																			A__							=       self._write(file_obj=UpperCAmelCase__			,	orient=UpperCAmelCase__			,	lines=UpperCAmelCase__			,	index=UpperCAmelCase__			,	**self.to_json_kwargs		)
							else:
													if compression:
																			raise NotImplementedError(
																			    F"""The compression parameter is not supported when writing to a buffer, but compression={compression}"""
																			    " was passed. Please provide a local path instead."		)
													A__							=       self._write(
													    file_obj=self.path_or_buf			,	orient=UpperCAmelCase__			,	lines=UpperCAmelCase__			,	index=UpperCAmelCase__			,	**self.to_json_kwargs		)
							return written



	def 			__A							(       self			,	UpperCAmelCase__		):
							A__						,       A__						,       A__						,       A__						,       A__							=       args

							A__							=       query_table(
							    table=self.dataset.data			,	key=slice(UpperCAmelCase__			,	offset + self.batch_size		)			,	indices=self.dataset._indices			,	)
							A__							=       batch.to_pandas().to_json(
							    path_or_buf=UpperCAmelCase__			,	orient=UpperCAmelCase__			,	lines=UpperCAmelCase__			,	index=UpperCAmelCase__			,	**UpperCAmelCase__		)
							if not json_str.endswith("\n"		):
													json_str += "\n"
							return json_str.encode(self.encoding		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	**UpperCAmelCase__			,	):
							A__							=       0

							if self.num_proc is None or self.num_proc == 1:
													for offset in logging.tqdm(
													    range(0			,	len(self.dataset		)			,	self.batch_size		)			,	unit="ba"			,	disable=not logging.is_progress_bar_enabled()			,	desc="Creating json from Arrow format"			,	):
																			A__							=       self._batch_json((offset, orient, lines, index, to_json_kwargs)		)
																			written += file_obj.write(UpperCAmelCase__		)
							else:
													A__						,       A__							=       len(self.dataset		), self.batch_size
													with multiprocessing.Pool(self.num_proc		) as pool:
																			for json_str in logging.tqdm(
																			    pool.imap(
																			        self._batch_json			,	[(offset, orient, lines, index, to_json_kwargs) for offset in range(0			,	UpperCAmelCase__			,	UpperCAmelCase__		)]			,	)			,	total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size			,	unit="ba"			,	disable=not logging.is_progress_bar_enabled()			,	desc="Creating json from Arrow format"			,	):
																									written += file_obj.write(UpperCAmelCase__		)

							return written





