


'''simple docstring'''
import argparse
import logging
import os
from pathlib import Path
from typing import Any, Dict

import pytorch_lightning as pl
from pytorch_lightning.utilities import rank_zero_info

from transformers import (
    AdamW,
    AutoConfig,
    AutoModel,
    AutoModelForPreTraining,
    AutoModelForQuestionAnswering,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification,
    AutoModelWithLMHead,
    AutoTokenizer,
    PretrainedConfig,
    PreTrainedTokenizer,
)
from transformers.optimization import (
    Adafactor,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.utils.versions import require_version


__snake_case							      =     logging.getLogger(__name__)

require_version('''pytorch_lightning>=1.0.4''')

__snake_case							      =     {
    '''base''': AutoModel,
    '''sequence-classification''': AutoModelForSequenceClassification,
    '''question-answering''': AutoModelForQuestionAnswering,
    '''pretraining''': AutoModelForPreTraining,
    '''token-classification''': AutoModelForTokenClassification,
    '''language-modeling''': AutoModelWithLMHead,
    '''summarization''': AutoModelForSeqaSeqLM,
    '''translation''': AutoModelForSeqaSeqLM,
}


# update this and the import above to support new schedulers from transformers.optimization
__snake_case							      =     {
    '''linear''': get_linear_schedule_with_warmup,
    '''cosine''': get_cosine_schedule_with_warmup,
    '''cosine_w_restarts''': get_cosine_with_hard_restarts_schedule_with_warmup,
    '''polynomial''': get_polynomial_decay_schedule_with_warmup,
    # '': get_constant_schedule,             # not supported for now
    # '': get_constant_schedule_with_warmup, # not supported for now
}
__snake_case							      =     sorted(arg_to_scheduler.keys())
__snake_case							      =     '''{''' + ''', '''.join(arg_to_scheduler_choices) + '''}'''
class    lowercase	(							pl.LightningModule					):

		"""simple docstring"""

		def __init__(		self    ,   UpperCamelCase_    ,   UpperCamelCase_=None    ,   UpperCamelCase_="base"    ,   UpperCamelCase_=None    ,   UpperCamelCase_=None    ,   UpperCamelCase_=None    ,   **UpperCamelCase_    ,   ):

							'''simple docstring'''


							super().__init__()
							# TODO: move to self.save_hyperparameters()
							# self.save_hyperparameters()
							# can also expand arguments into trainer signature for easier reading

							self.save_hyperparameters(UpperCamelCase_							)
							UpperCamelCase__   :Any												= 0
							UpperCamelCase__   :str												= Path(self.hparams.output_dir							)
							UpperCamelCase__   :List[str]												= self.hparams.cache_dir if self.hparams.cache_dir else None
							if config is None:
												UpperCamelCase__   :Optional[int]												= AutoConfig.from_pretrained(
												    self.hparams.config_name if self.hparams.config_name else self.hparams.model_name_or_path    ,   **({'''num_labels''': num_labels} if num_labels is not None else {})    ,   cache_dir=UpperCamelCase_    ,   **UpperCamelCase_    ,   )
							else:
												UpperCamelCase__   :PretrainedConfig												= config

							UpperCamelCase__   :int												= ('''encoder_layerdrop''', '''decoder_layerdrop''', '''dropout''', '''attention_dropout''')
							for p in extra_model_params:
												if getattr(self.hparams    ,   UpperCamelCase_    ,   UpperCamelCase_							):
																	assert hasattr(self.config    ,   UpperCamelCase_							), F'''model config doesn\'t have a `{p}` attribute'''
																	setattr(self.config    ,   UpperCamelCase_    ,   getattr(self.hparams    ,   UpperCamelCase_							)							)

							if tokenizer is None:
												UpperCamelCase__   :Dict												= AutoTokenizer.from_pretrained(
												    self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path    ,   cache_dir=UpperCamelCase_    ,   )
							else:
												UpperCamelCase__   :PreTrainedTokenizer												= tokenizer
							UpperCamelCase__   :Optional[int]												= MODEL_MODES[mode]
							if model is None:
												UpperCamelCase__   :str												= self.model_type.from_pretrained(
												    self.hparams.model_name_or_path    ,   from_tf=bool('''.ckpt''' in self.hparams.model_name_or_path							)    ,   config=self.config    ,   cache_dir=UpperCamelCase_    ,   )
							else:
												UpperCamelCase__   :List[Any]												= model




		def     lowerCAmelCase__				(		self    ,   *UpperCamelCase_    ,   **UpperCamelCase_							):

							'''simple docstring'''


							UpperCamelCase__   :Tuple												= self.model_type.from_pretrained(*UpperCamelCase_    ,   **UpperCamelCase_							)




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							UpperCamelCase__   :int												= arg_to_scheduler[self.hparams.lr_scheduler]
							UpperCamelCase__   :Dict												= get_schedule_func(
							    self.opt    ,   num_warmup_steps=self.hparams.warmup_steps    ,   num_training_steps=self.total_steps()							)
							UpperCamelCase__   :Any												= {'''scheduler''': scheduler, '''interval''': '''step''', '''frequency''': 1}
							return scheduler




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							UpperCamelCase__   :List[str]												= self.model
							UpperCamelCase__   :Optional[Any]												= ['''bias''', '''LayerNorm.weight''']
							UpperCamelCase__   :Optional[Any]												= [
							    {
							        '''params''': [
							            p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay							)
							        ],  # check this named paramters
							        '''weight_decay''': self.hparams.weight_decay,
							    },
							    {
							        '''params''': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay							)],
							        '''weight_decay''': 0.0,
							    },
							]
							if self.hparams.adafactor:
												UpperCamelCase__   :List[str]												= Adafactor(
												    UpperCamelCase_    ,   lr=self.hparams.learning_rate    ,   scale_parameter=UpperCamelCase_    ,   relative_step=UpperCamelCase_							)

							else:
												UpperCamelCase__   :int												= AdamW(
												    UpperCamelCase_    ,   lr=self.hparams.learning_rate    ,   eps=self.hparams.adam_epsilon							)
							UpperCamelCase__   :str												= optimizer

							UpperCamelCase__   :int												= self.get_lr_scheduler()

							return [optimizer], [scheduler]




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_    ,   UpperCamelCase_							):

							'''simple docstring'''


							return self.validation_step(UpperCamelCase_    ,   UpperCamelCase_							)




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_							):

							'''simple docstring'''


							return self.validation_end(UpperCamelCase_							)




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							UpperCamelCase__   :int												= max(1    ,   self.hparams.gpus							)  # TODO: consider num_tpu_cores
							UpperCamelCase__   :List[str]												= self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
							return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_							):

							'''simple docstring'''


							if stage == "test":
												UpperCamelCase__   :Union[str, Any]												= len(self.test_dataloader().dataset							)
							else:
												UpperCamelCase__   :List[Any]												= self.get_dataloader('''train'''    ,   self.hparams.train_batch_size    ,   shuffle=UpperCamelCase_							)
												UpperCamelCase__   :List[Any]												= len(self.train_dataloader().dataset							)




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_ = False							):

							'''simple docstring'''


							raise NotImplementedError('''You must implement this for your task'''							)




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							return self.train_loader




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							return self.get_dataloader('''dev'''    ,   self.hparams.eval_batch_size    ,   shuffle=UpperCamelCase_							)




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							return self.get_dataloader('''test'''    ,   self.hparams.eval_batch_size    ,   shuffle=UpperCamelCase_							)




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_							):

							'''simple docstring'''


							return os.path.join(
							    self.hparams.data_dir    ,   '''cached_{}_{}_{}'''.format(
							        UpperCamelCase_    ,   list(filter(UpperCamelCase_    ,   self.hparams.model_name_or_path.split('''/'''							)							)							).pop()    ,   str(self.hparams.max_seq_length							)    ,   )    ,   )




		@pl.utilities.rank_zero_only
		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_							):

							'''simple docstring'''


							UpperCamelCase__   :Optional[Any]												= self.output_dir.joinpath('''best_tfmr'''							)
							UpperCamelCase__   :int												= self.step_count
							self.model.save_pretrained(UpperCamelCase_							)
							self.tokenizer.save_pretrained(UpperCamelCase_							)




		@staticmethod
		def     lowerCAmelCase__				(		UpperCamelCase_    ,   UpperCamelCase_							):

							'''simple docstring'''


							parser.add_argument(
							    '''--model_name_or_path'''    ,   default=UpperCamelCase_    ,   type=UpperCamelCase_    ,   required=UpperCamelCase_    ,   help='''Path to pretrained model or model identifier from huggingface.co/models'''    ,   )
							parser.add_argument(
							    '''--config_name'''    ,   default=''''''    ,   type=UpperCamelCase_    ,   help='''Pretrained config name or path if not the same as model_name'''							)
							parser.add_argument(
							    '''--tokenizer_name'''    ,   default=UpperCamelCase_    ,   type=UpperCamelCase_    ,   help='''Pretrained tokenizer name or path if not the same as model_name'''    ,   )
							parser.add_argument(
							    '''--cache_dir'''    ,   default=str(Path(UpperCamelCase_							).parent / '''test_run''' / '''cache'''							)    ,   type=UpperCamelCase_    ,   help='''Where do you want to store the pre-trained models downloaded from huggingface.co'''    ,   )
							parser.add_argument(
							    '''--encoder_layerdrop'''    ,   type=UpperCamelCase_    ,   help='''Encoder layer dropout probability (Optional). Goes into model.config'''    ,   )
							parser.add_argument(
							    '''--decoder_layerdrop'''    ,   type=UpperCamelCase_    ,   help='''Decoder layer dropout probability (Optional). Goes into model.config'''    ,   )
							parser.add_argument(
							    '''--dropout'''    ,   type=UpperCamelCase_    ,   help='''Dropout probability (Optional). Goes into model.config'''    ,   )
							parser.add_argument(
							    '''--attention_dropout'''    ,   type=UpperCamelCase_    ,   help='''Attention dropout probability (Optional). Goes into model.config'''    ,   )
							parser.add_argument('''--learning_rate'''    ,   default=5e-5    ,   type=UpperCamelCase_    ,   help='''The initial learning rate for Adam.'''							)
							parser.add_argument(
							    '''--lr_scheduler'''    ,   default='''linear'''    ,   choices=UpperCamelCase_    ,   metavar=UpperCamelCase_    ,   type=UpperCamelCase_    ,   help='''Learning rate scheduler'''    ,   )
							parser.add_argument('''--weight_decay'''    ,   default=0.0    ,   type=UpperCamelCase_    ,   help='''Weight decay if we apply some.'''							)
							parser.add_argument('''--adam_epsilon'''    ,   default=1e-8    ,   type=UpperCamelCase_    ,   help='''Epsilon for Adam optimizer.'''							)
							parser.add_argument('''--warmup_steps'''    ,   default=0    ,   type=UpperCamelCase_    ,   help='''Linear warmup over warmup_steps.'''							)
							parser.add_argument('''--num_workers'''    ,   default=4    ,   type=UpperCamelCase_    ,   help='''kwarg passed to DataLoader'''							)
							parser.add_argument('''--num_train_epochs'''    ,   dest='''max_epochs'''    ,   default=3    ,   type=UpperCamelCase_							)
							parser.add_argument('''--train_batch_size'''    ,   default=32    ,   type=UpperCamelCase_							)
							parser.add_argument('''--eval_batch_size'''    ,   default=32    ,   type=UpperCamelCase_							)
							parser.add_argument('''--adafactor'''    ,   action='''store_true'''							)



class    lowercase	(							pl.Callback					):

		"""simple docstring"""

		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_    ,   UpperCamelCase_							):

							'''simple docstring'''


							if (
							    trainer.is_global_zero and trainer.global_rank == 0
							):  # we initialize the retriever only on master worker with RAY. In new pytorch-lightning accelorators are removed.
												pl_module.model.rag.retriever.init_retrieval()  # better to use hook functions.



class    lowercase	(							pl.Callback					):

		"""simple docstring"""

		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_    ,   UpperCamelCase_							):

							'''simple docstring'''


							for name, param in pl_module.model.rag.named_parameters():
												if param.grad is None:
																	print(UpperCamelCase_							)



class    lowercase	(							pl.Callback					):

		"""simple docstring"""

		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_    ,   UpperCamelCase_							):

							'''simple docstring'''


							UpperCamelCase__   :Union[str, Any]												= trainer.lr_schedulers[0]['''scheduler''']
							UpperCamelCase__   :Dict												= {F'''lr_group_{i}''': lr for i, lr in enumerate(lr_scheduler.get_lr()							)}
							pl_module.logger.log_metrics(UpperCamelCase_							)




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_    ,   UpperCamelCase_							):

							'''simple docstring'''


							rank_zero_info('''***** Validation results *****'''							)
							UpperCamelCase__   :Dict												= trainer.callback_metrics
							# Log results
							for key in sorted(UpperCamelCase_							):
												if key not in ["log", "progress_bar"]:
																	rank_zero_info('''{} = {}\n'''.format(UpperCamelCase_    ,   str(metrics[key]							)							)							)




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_    ,   UpperCamelCase_							):

							'''simple docstring'''


							rank_zero_info('''***** Test results *****'''							)
							UpperCamelCase__   :int												= trainer.callback_metrics
							# Log and save results to file
							UpperCamelCase__   :Any												= os.path.join(pl_module.hparams.output_dir    ,   '''test_results.txt'''							)
							with open(UpperCamelCase_    ,   '''w'''							) as writer:
												for key in sorted(UpperCamelCase_							):
																	if key not in ["log", "progress_bar"]:
																						rank_zero_info('''{} = {}\n'''.format(UpperCamelCase_    ,   str(metrics[key]							)							)							)
																						writer.write('''{} = {}\n'''.format(UpperCamelCase_    ,   str(metrics[key]							)							)							)



def 			a      (			__a ,					__a   )  ->	None:



					'''simple docstring'''






					parser.add_argument(
					    '''--output_dir''' ,					default=str(Path(__a   ).parent / '''test_run''' / '''model_checkpoints'''   ) ,					type=__a ,					help='''The output directory where the model predictions and checkpoints will be written.''' ,					)
					parser.add_argument(
					    '''--fp16''' ,					action='''store_true''' ,					help='''Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit''' ,					)

					parser.add_argument(
					    '''--fp16_opt_level''' ,					type=__a ,					default='''O2''' ,					help=(
					        '''For fp16: Apex AMP optimization level selected in [\'O0\', \'O1\', \'O2\', and \'O3\'].'''
					        '''See details at https://nvidia.github.io/apex/amp.html'''
					    ) ,					)
					parser.add_argument('''--n_tpu_cores''' ,					dest='''tpu_cores''' ,					type=__a   )
					parser.add_argument('''--max_grad_norm''' ,					dest='''gradient_clip_val''' ,					default=1.0 ,					type=__a ,					help='''Max gradient norm'''   )
					parser.add_argument('''--do_train''' ,					action='''store_true''' ,					help='''Whether to run training.'''   )
					parser.add_argument('''--do_predict''' ,					action='''store_true''' ,					help='''Whether to run predictions on the test set.'''   )
					parser.add_argument(
					    '''--gradient_accumulation_steps''' ,					dest='''accumulate_grad_batches''' ,					type=__a ,					default=1 ,					help='''Number of updates steps to accumulate before performing a backward/update pass.''' ,					)
					parser.add_argument('''--seed''' ,					type=__a ,					default=42 ,					help='''random seed for initialization'''   )
					parser.add_argument(
					    '''--data_dir''' ,					default=str(Path(__a   ).parent / '''test_run''' / '''dummy-train-data'''   ) ,					type=__a ,					help='''The input data dir. Should contain the training files for the CoNLL-2003 NER task.''' ,					)







def 			a      (			__a ,					__a ,					__a=None ,					__a=True ,					__a=[] ,					__a=None ,					__a=None ,					**__a ,					)  ->	str:



					'''simple docstring'''






					pl.seed_everything(args.seed   )

					# init model
					UpperCamelCase__   :Union[str, Any]												= Path(model.hparams.output_dir   )
					odir.mkdir(exist_ok=__a   )

					# add custom checkpoints
					if checkpoint_callback is None:
										UpperCamelCase__   :Any												= pl.callbacks.ModelCheckpoint(
										    filepath=args.output_dir ,					prefix='''checkpoint''' ,					monitor='''val_loss''' ,					mode='''min''' ,					save_top_k=1   )
					if early_stopping_callback:
										extra_callbacks.append(__a   )
					if logging_callback is None:
										UpperCamelCase__   :str												= LoggingCallback()

					UpperCamelCase__   :Optional[Any]												= {}

					if args.fpaa:
										UpperCamelCase__   :List[Any]												= 16

					if args.gpus > 1:
										UpperCamelCase__   :Union[str, Any]												= '''auto'''
										UpperCamelCase__   :Optional[Any]												= '''ddp'''

					UpperCamelCase__   :List[str]												= args.accumulate_grad_batches
					UpperCamelCase__   :Dict												= None
					UpperCamelCase__   :Optional[Any]												= '''auto'''

					UpperCamelCase__   :List[str]												= pl.Trainer.from_argparse_args(
					    __a ,					weights_summary=__a ,					callbacks=[logging_callback] + extra_callbacks + [InitCallback()] + [checkpoint_callback] ,					logger=__a ,					val_check_interval=1 ,					num_sanity_val_steps=2 ,					**__a ,					)

					if args.do_train:
										trainer.fit(__a   )

					else:
										print('''RAG modeling tests with new set functions successfuly executed!'''   )
					return trainer



'''simple docstring'''
import io
import json

import fsspec
import pytest

from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
from datasets.io.json import JsonDatasetReader, JsonDatasetWriter

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases


def 			a      (			__a ,					__a   )  ->	Optional[int]:



					'''simple docstring'''






					assert isinstance(__a ,					__a   )
					assert dataset.num_rows == 4
					assert dataset.num_columns == 3
					assert dataset.column_names == ["col_1", "col_2", "col_3"]
					for feature, expected_dtype in expected_features.items():
										assert dataset.features[feature].dtype == expected_dtype


@pytest.mark.parametrize('''keep_in_memory''' ,					[False, True]   )
def 			a      (			__a ,					__a ,					__a   )  ->	Optional[Any]:



					'''simple docstring'''






					UpperCamelCase__   :Union[str, Any]												= tmp_path / '''cache'''
					UpperCamelCase__   :Dict												= {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
					with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
										UpperCamelCase__   :Tuple												= JsonDatasetReader(__a ,					cache_dir=__a ,					keep_in_memory=__a   ).read()
					_check_json_dataset(__a ,					__a   )


@pytest.mark.parametrize(
    '''features''' ,					[
        None,
        {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''},
        {'''col_1''': '''string''', '''col_2''': '''string''', '''col_3''': '''string'''},
        {'''col_1''': '''int32''', '''col_2''': '''int32''', '''col_3''': '''int32'''},
        {'''col_1''': '''float32''', '''col_2''': '''float32''', '''col_3''': '''float32'''},
    ] ,					)
def 			a      (			__a ,					__a ,					__a   )  ->	Any:



					'''simple docstring'''






					UpperCamelCase__   :Union[str, Any]												= tmp_path / '''cache'''
					UpperCamelCase__   :Optional[Any]												= {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
					UpperCamelCase__   :Optional[Any]												= features.copy() if features else default_expected_features
					UpperCamelCase__   :Tuple												= (
					    Features({feature: Value(__a   ) for feature, dtype in features.items()}   ) if features is not None else None
					)
					UpperCamelCase__   :int												= JsonDatasetReader(__a ,					features=__a ,					cache_dir=__a   ).read()
					_check_json_dataset(__a ,					__a   )


@pytest.mark.parametrize(
    '''features''' ,					[
        None,
        {'''col_3''': '''float64''', '''col_1''': '''string''', '''col_2''': '''int64'''},
    ] ,					)
def 			a      (			__a ,					__a ,					__a   )  ->	Tuple:



					'''simple docstring'''






					UpperCamelCase__   :int												= tmp_path / '''cache'''
					UpperCamelCase__   :str												= {'''col_3''': '''float64''', '''col_1''': '''string''', '''col_2''': '''int64'''}
					UpperCamelCase__   :Any												= features.copy() if features else default_expected_features
					UpperCamelCase__   :Union[str, Any]												= (
					    Features({feature: Value(__a   ) for feature, dtype in features.items()}   ) if features is not None else None
					)
					UpperCamelCase__   :Any												= JsonDatasetReader(__a ,					features=__a ,					cache_dir=__a   ).read()
					assert isinstance(__a ,					__a   )
					assert dataset.num_rows == 2
					assert dataset.num_columns == 3
					assert dataset.column_names == ["col_3", "col_1", "col_2"]
					for feature, expected_dtype in expected_features.items():
										assert dataset.features[feature].dtype == expected_dtype


def 			a      (			__a ,					__a   )  ->	List[Any]:



					'''simple docstring'''






					UpperCamelCase__   :Any												= {'''col_2''': '''int64''', '''col_3''': '''float64''', '''col_1''': '''string'''}
					UpperCamelCase__   :int												= features.copy()
					UpperCamelCase__   :List[Any]												= (
					    Features({feature: Value(__a   ) for feature, dtype in features.items()}   ) if features is not None else None
					)
					UpperCamelCase__   :Optional[int]												= tmp_path / '''cache'''
					UpperCamelCase__   :Dict												= JsonDatasetReader(__a ,					features=__a ,					cache_dir=__a   ).read()
					assert isinstance(__a ,					__a   )
					assert dataset.num_rows == 2
					assert dataset.num_columns == 3
					assert dataset.column_names == ["col_2", "col_3", "col_1"]
					for feature, expected_dtype in expected_features.items():
										assert dataset.features[feature].dtype == expected_dtype


@pytest.mark.parametrize('''split''' ,					[None, NamedSplit('''train'''   ), '''train''', '''test''']   )
def 			a      (			__a ,					__a ,					__a   )  ->	List[Any]:



					'''simple docstring'''






					UpperCamelCase__   :Union[str, Any]												= tmp_path / '''cache'''
					UpperCamelCase__   :Optional[Any]												= {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
					UpperCamelCase__   :List[Any]												= JsonDatasetReader(__a ,					cache_dir=__a ,					split=__a   ).read()
					_check_json_dataset(__a ,					__a   )
					assert dataset.split == split if split else "train"


@pytest.mark.parametrize('''path_type''' ,					[str, list]   )
def 			a      (			__a ,					__a ,					__a   )  ->	Any:



					'''simple docstring'''






					if issubclass(__a ,					__a   ):
										UpperCamelCase__   :Union[str, Any]												= jsonl_path
					elif issubclass(__a ,					__a   ):
										UpperCamelCase__   :int												= [jsonl_path]
					UpperCamelCase__   :Dict												= tmp_path / '''cache'''
					UpperCamelCase__   :Any												= {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
					UpperCamelCase__   :List[str]												= JsonDatasetReader(__a ,					cache_dir=__a   ).read()
					_check_json_dataset(__a ,					__a   )


def 			a      (			__a ,					__a ,					__a=("train",)   )  ->	Optional[Any]:



					'''simple docstring'''






					assert isinstance(__a ,					__a   )
					for split in splits:
										UpperCamelCase__   :Optional[int]												= dataset_dict[split]
										assert dataset.num_rows == 4
										assert dataset.num_columns == 3
										assert dataset.column_names == ["col_1", "col_2", "col_3"]
										for feature, expected_dtype in expected_features.items():
															assert dataset.features[feature].dtype == expected_dtype


@pytest.mark.parametrize('''keep_in_memory''' ,					[False, True]   )
def 			a      (			__a ,					__a ,					__a   )  ->	List[str]:



					'''simple docstring'''






					UpperCamelCase__   :List[str]												= tmp_path / '''cache'''
					UpperCamelCase__   :Dict												= {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
					with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
										UpperCamelCase__   :str												= JsonDatasetReader({'''train''': jsonl_path} ,					cache_dir=__a ,					keep_in_memory=__a   ).read()
					_check_json_datasetdict(__a ,					__a   )


@pytest.mark.parametrize(
    '''features''' ,					[
        None,
        {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''},
        {'''col_1''': '''string''', '''col_2''': '''string''', '''col_3''': '''string'''},
        {'''col_1''': '''int32''', '''col_2''': '''int32''', '''col_3''': '''int32'''},
        {'''col_1''': '''float32''', '''col_2''': '''float32''', '''col_3''': '''float32'''},
    ] ,					)
def 			a      (			__a ,					__a ,					__a   )  ->	int:



					'''simple docstring'''






					UpperCamelCase__   :Tuple												= tmp_path / '''cache'''
					UpperCamelCase__   :Any												= {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
					UpperCamelCase__   :Optional[int]												= features.copy() if features else default_expected_features
					UpperCamelCase__   :str												= (
					    Features({feature: Value(__a   ) for feature, dtype in features.items()}   ) if features is not None else None
					)
					UpperCamelCase__   :Dict												= JsonDatasetReader({'''train''': jsonl_path} ,					features=__a ,					cache_dir=__a   ).read()
					_check_json_datasetdict(__a ,					__a   )


@pytest.mark.parametrize('''split''' ,					[None, NamedSplit('''train'''   ), '''train''', '''test''']   )
def 			a      (			__a ,					__a ,					__a   )  ->	str:



					'''simple docstring'''






					if split:
										UpperCamelCase__   :List[str]												= {split: jsonl_path}
					else:
										UpperCamelCase__   :int												= '''train'''
										UpperCamelCase__   :int												= {'''train''': jsonl_path, '''test''': jsonl_path}
					UpperCamelCase__   :Any												= tmp_path / '''cache'''
					UpperCamelCase__   :Union[str, Any]												= {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
					UpperCamelCase__   :Any												= JsonDatasetReader(__a ,					cache_dir=__a   ).read()
					_check_json_datasetdict(__a ,					__a ,					splits=list(path.keys()   )   )
					assert all(dataset[split].split == split for split in path.keys()   )


def 			a      (			__a   )  ->	Union[str, Any]:



					'''simple docstring'''






					return json.load(__a   )


def 			a      (			__a   )  ->	int:



					'''simple docstring'''






					return [json.loads(__a   ) for line in buffer]







class    lowercase	:

		"""simple docstring"""

		@pytest.mark.parametrize('''lines, load_json_function'''    ,   [(True, load_json_lines), (False, load_json)]							)
		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_							):

							'''simple docstring'''


							with io.BytesIO() as buffer:
												JsonDatasetWriter(UpperCamelCase_    ,   UpperCamelCase_    ,   lines=UpperCamelCase_							).write()
												buffer.seek(0							)
												UpperCamelCase__   :List[Any]												= load_json_function(UpperCamelCase_							)
							assert isinstance(UpperCamelCase_    ,   UpperCamelCase_							)
							assert isinstance(exported_content[0]    ,   UpperCamelCase_							)
							assert len(UpperCamelCase_							) == 10




		@pytest.mark.parametrize(
		    '''orient, container, keys, len_at'''    ,   [
		        ('''records''', list, {'''tokens''', '''labels''', '''answers''', '''id'''}, None),
		        ('''split''', dict, {'''columns''', '''data'''}, '''data'''),
		        ('''index''', dict, set('''0123456789'''							), None),
		        ('''columns''', dict, {'''tokens''', '''labels''', '''answers''', '''id'''}, '''tokens'''),
		        ('''values''', list, None, None),
		        ('''table''', dict, {'''schema''', '''data'''}, '''data'''),
		    ]    ,   )
		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_							):

							'''simple docstring'''


							with io.BytesIO() as buffer:
												JsonDatasetWriter(UpperCamelCase_    ,   UpperCamelCase_    ,   lines=UpperCamelCase_    ,   orient=UpperCamelCase_							).write()
												buffer.seek(0							)
												UpperCamelCase__   :Optional[int]												= load_json(UpperCamelCase_							)
							assert isinstance(UpperCamelCase_    ,   UpperCamelCase_							)
							if keys:
												if container is dict:
																	assert exported_content.keys() == keys
												else:
																	assert exported_content[0].keys() == keys
							else:
												assert not hasattr(UpperCamelCase_    ,   '''keys'''							) and not hasattr(exported_content[0]    ,   '''keys'''							)
							if len_at:
												assert len(exported_content[len_at]							) == 10
							else:
												assert len(UpperCamelCase_							) == 10




		@pytest.mark.parametrize('''lines, load_json_function'''    ,   [(True, load_json_lines), (False, load_json)]							)
		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_							):

							'''simple docstring'''


							with io.BytesIO() as buffer:
												JsonDatasetWriter(UpperCamelCase_    ,   UpperCamelCase_    ,   lines=UpperCamelCase_    ,   num_proc=2							).write()
												buffer.seek(0							)
												UpperCamelCase__   :Union[str, Any]												= load_json_function(UpperCamelCase_							)
							assert isinstance(UpperCamelCase_    ,   UpperCamelCase_							)
							assert isinstance(exported_content[0]    ,   UpperCamelCase_							)
							assert len(UpperCamelCase_							) == 10




		@pytest.mark.parametrize(
		    '''orient, container, keys, len_at'''    ,   [
		        ('''records''', list, {'''tokens''', '''labels''', '''answers''', '''id'''}, None),
		        ('''split''', dict, {'''columns''', '''data'''}, '''data'''),
		        ('''index''', dict, set('''0123456789'''							), None),
		        ('''columns''', dict, {'''tokens''', '''labels''', '''answers''', '''id'''}, '''tokens'''),
		        ('''values''', list, None, None),
		        ('''table''', dict, {'''schema''', '''data'''}, '''data'''),
		    ]    ,   )
		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_							):

							'''simple docstring'''


							with io.BytesIO() as buffer:
												JsonDatasetWriter(UpperCamelCase_    ,   UpperCamelCase_    ,   lines=UpperCamelCase_    ,   orient=UpperCamelCase_    ,   num_proc=2							).write()
												buffer.seek(0							)
												UpperCamelCase__   :int												= load_json(UpperCamelCase_							)
							assert isinstance(UpperCamelCase_    ,   UpperCamelCase_							)
							if keys:
												if container is dict:
																	assert exported_content.keys() == keys
												else:
																	assert exported_content[0].keys() == keys
							else:
												assert not hasattr(UpperCamelCase_    ,   '''keys'''							) and not hasattr(exported_content[0]    ,   '''keys'''							)
							if len_at:
												assert len(exported_content[len_at]							) == 10
							else:
												assert len(UpperCamelCase_							) == 10




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_							):

							'''simple docstring'''


							with pytest.raises(UpperCamelCase_							):
												with io.BytesIO() as buffer:
																	JsonDatasetWriter(UpperCamelCase_    ,   UpperCamelCase_    ,   num_proc=0							)




		@pytest.mark.parametrize('''compression, extension'''    ,   [('''gzip''', '''gz'''), ('''bz2''', '''bz2'''), ('''xz''', '''xz''')]							)
		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_							):

							'''simple docstring'''


							UpperCamelCase__   :Tuple												= tmp_path_factory.mktemp('''data'''							) / F'''test.json.{extension}'''
							UpperCamelCase__   :Union[str, Any]												= str(shared_datadir / F'''test_file.json.{extension}'''							)
							JsonDatasetWriter(UpperCamelCase_    ,   UpperCamelCase_    ,   compression=UpperCamelCase_							).write()

							with fsspec.open(UpperCamelCase_    ,   '''rb'''    ,   compression='''infer'''							) as f:
												UpperCamelCase__   :Dict												= f.read()
							with fsspec.open(UpperCamelCase_    ,   '''rb'''    ,   compression='''infer'''							) as f:
												UpperCamelCase__   :int												= f.read()
							assert exported_content == original_content