



'''simple docstring'''



from __future__ import annotations

import os
import tempfile
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import is_tensorflow_text_available, is_tf_available
from transformers.testing_utils import require_tensorflow_text, require_tf, slow

from ..test_modeling_tf_common import floats_tensor
from .test_framework_agnostic import GenerationIntegrationTestsMixin


if is_tf_available():
	import tensorflow as tf

	from transformers import (
	    AutoTokenizer,
	    TFAutoModelForCausalLM,
	    TFAutoModelForSeqaSeqLM,
	    TFAutoModelForSpeechSeqaSeq,
	    TFAutoModelForVisionaSeq,
	    TFBartForConditionalGeneration,
	    TFLogitsProcessorList,
	    TFMinLengthLogitsProcessor,
	    tf_top_k_top_p_filtering,
	)

if is_tensorflow_text_available():
	import tensorflow_text as text






@require_tf
class 					UpperCAmelCase__	(		unittest.TestCase ):



					def        UpperCAmelCase_ (    self     )     ->			int:
						__lowerCAmelCase												=  tf.convert_to_tensor(
						    [
						        [
						            8.2_22_09_91,  # 3rd highest value; idx. 0
						            -0.5_62_00_44,
						            5.23_22_97_52,
						            4.0_38_63_93,
						            -6.8_79_83_78,
						            -0.54_78_58_02,
						            -3.2_01_21_53,
						            2.92_77_71_76,
						            1.88_17_19_53,
						            7.35_34_12_76,  # 5th highest value; idx. 9
						            8.43_20_78_33,  # 2nd highest value; idx. 10
						            -9.85_71_18_36,
						            -5.96_20_92_36,
						            -1.13_03_91_61,
						            -7.1_11_52_94,
						            -0.8_36_96_33,
						            -5.3_18_64_08,
						            7.06_42_74_07,
						            0.81_36_93_44,
						            -0.82_02_38_17,
						            -5.9_17_97_96,
						            0.58_81_34_43,
						            -6.99_77_84_38,
						            4.71_55_11_89,
						            -0.18_77_16_37,
						            7.44_02_07_59,  # 4th highest value; idx. 25
						            9.38_45_09_87,  # 1st highest value; idx. 26
						            2.12_66_29_41,
						            -9.32_56_20_38,
						            2.35_65_25_22,
						        ],  # cummulative prob of 5 highest values <= 0.6
						        [
						            0.58_42_55_18,
						            4.53_13_92_38,
						            -5.57_51_04_64,
						            -6.28_03_06_99,
						            -7.19_52_95_03,
						            -4.02_12_25_51,
						            1.39_33_70_37,
						            -6.06_70_70_57,
						            1.59_48_05_17,
						            -9.64_31_19,
						            0.03_90_77_99,
						            0.67_23_17_62,
						            -8.88_20_67_26,
						            6.27_11_59_22,  # 4th highest value; idx. 13
						            2.28_52_07_23,
						            4.82_76_75_06,
						            4.30_42_13_68,
						            8.8_27_53_13,  # 2nd highest value; idx. 17
						            5.44_02_99_58,  # 5th highest value; idx. 18
						            -4.4_73_57_94,
						            7.38_57_95_36,  # 3rd highest value; idx. 20
						            -2.91_05_16_63,
						            2.61_94_60_77,
						            -2.5_67_47_62,
						            -9.48_95_93_02,
						            -4.02_92_26_45,
						            -1.35_41_69_18,
						            9.67_70_23_23,  # 1st highest value; idx. 27
						            -5.89_47_85_53,
						            1.85_37_04_67,
						        ],  # cummulative prob of 5 highest values <= 0.6
						    ]			,				dtype=tf.floataa			,				)

						__lowerCAmelCase												=  tf.convert_to_tensor(
						    [[0, 0], [0, 9], [0, 10], [0, 25], [0, 26], [1, 13], [1, 17], [1, 18], [1, 20], [1, 27]]			,				dtype=tf.intaa			,				)  # expected non filtered idx as noted above

						__lowerCAmelCase												=  tf.convert_to_tensor(
						    [8.22_20_99, 7.3_53_41_26, 8.43_20_78, 7.4_40_20_75, 9.3_84_51, 6.27_11_59, 8.82_75_31, 5.4_40_29_95, 7.3_85_79_56, 9.67_70_23]			,				dtype=tf.floataa			,				)  # expected non filtered values as noted above

						__lowerCAmelCase												=  tf_top_k_top_p_filtering(UpperCamelCase			,				top_k=10			,				top_p=0.6			,				min_tokens_to_keep=4     )

						__lowerCAmelCase												=  output[output != -float("inf"     )]
						__lowerCAmelCase												=  tf.cast(
						    tf.where(tf.not_equal(UpperCamelCase			,				tf.constant(-float("inf"     )			,				dtype=tf.floataa     )     )     )			,				dtype=tf.intaa			,				)

						tf.debugging.assert_near(UpperCamelCase			,				UpperCamelCase			,				rtol=1E-12     )
						tf.debugging.assert_equal(UpperCamelCase			,				UpperCamelCase     )






@require_tf
class 					UpperCAmelCase__	(		unittest.TestCase							,       UpperCamelCase__ ):
					# setting framework_dependent_parameters needs to be gated, just like its contents' imports
					if is_tf_available():
										a      :     Dict        =							{
										    """AutoModelForCausalLM""": TFAutoModelForCausalLM,
										    """AutoModelForSpeechSeq2Seq""": TFAutoModelForSpeechSeqaSeq,
										    """AutoModelForSeq2SeqLM""": TFAutoModelForSeqaSeqLM,
										    """AutoModelForVision2Seq""": TFAutoModelForVisionaSeq,
										    """LogitsProcessorList""": TFLogitsProcessorList,
										    """MinLengthLogitsProcessor""": TFMinLengthLogitsProcessor,
										    """create_tensor_fn""": tf.convert_to_tensor,
										    """floats_tensor""": floats_tensor,
										    """return_tensors""": """tf""",
										}



					@slow
					def        UpperCAmelCase_ (    self     )     ->			Any:
						# TF-only test: tf.saved_model export
						__lowerCAmelCase												=  TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2"     )
						__lowerCAmelCase												=  2
						__lowerCAmelCase												=  2






						class 					UpperCAmelCase__	(		tf.Module ):



											def __init__(    self			,				UpperCamelCase     )     ->			Tuple:
												super(UpperCamelCase			,				self     ).__init__()
												__lowerCAmelCase												=  model



											@tf.function(
											    input_signature=(
											        tf.TensorSpec((None, input_length)			,				tf.intaa			,				name="input_ids"     ),
											        tf.TensorSpec((None, input_length)			,				tf.intaa			,				name="attention_mask"     ),
											    )			,				jit_compile=UpperCamelCase			,				)
											def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase     )     ->			Optional[Any]:
												__lowerCAmelCase												=  self.model.generate(
												    input_ids=UpperCamelCase			,				attention_mask=UpperCamelCase			,				max_new_tokens=UpperCamelCase			,				return_dict_in_generate=UpperCamelCase			,				)
												return {"sequences": outputs["sequences"]}


						__lowerCAmelCase												=  [[2, 0], [102, 103]]
						__lowerCAmelCase												=  [[1, 0], [1, 1]]
						__lowerCAmelCase												=  DummyModel(model=UpperCamelCase     )
						with tempfile.TemporaryDirectory() as tmp_dir:
							tf.saved_model.save(UpperCamelCase			,				UpperCamelCase			,				signatures={"serving_default": dummy_model.serving}     )
							__lowerCAmelCase												=  tf.saved_model.load(UpperCamelCase     ).signatures["serving_default"]
							for batch_size in range(1			,				len(UpperCamelCase     ) + 1     ):
								__lowerCAmelCase												=  {
								    "input_ids": tf.constant(dummy_input_ids[:batch_size]     ),
								    "attention_mask": tf.constant(dummy_attention_masks[:batch_size]     ),
								}
								__lowerCAmelCase												=  serving_func(**UpperCamelCase     )["sequences"]
								__lowerCAmelCase												=  test_model.generate(**UpperCamelCase			,				max_new_tokens=UpperCamelCase     )
								tf.debugging.assert_equal(UpperCamelCase			,				UpperCamelCase     )



					@slow
					def        UpperCAmelCase_ (    self     )     ->			Optional[Any]:
						# TF-only test: tf.saved_model export
						__lowerCAmelCase												=  TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2"     )
						__lowerCAmelCase												=  1
						__lowerCAmelCase												=  2






						class 					UpperCAmelCase__	(		tf.Module ):



											def __init__(    self			,				UpperCamelCase     )     ->			Optional[int]:
												super(UpperCamelCase			,				self     ).__init__()
												__lowerCAmelCase												=  model



											@tf.function(
											    input_signature=(
											        tf.TensorSpec((batch_size, None)			,				tf.intaa			,				name="input_ids"     ),
											        tf.TensorSpec((batch_size, None)			,				tf.intaa			,				name="attention_mask"     ),
											    )			,				jit_compile=UpperCamelCase			,				)
											def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase     )     ->			str:
												__lowerCAmelCase												=  self.model.generate(
												    input_ids=UpperCamelCase			,				attention_mask=UpperCamelCase			,				max_new_tokens=UpperCamelCase			,				return_dict_in_generate=UpperCamelCase			,				)
												return {"sequences": outputs["sequences"]}


						__lowerCAmelCase												=  [[2], [102, 103]]
						__lowerCAmelCase												=  [[1], [1, 1]]
						__lowerCAmelCase												=  DummyModel(model=UpperCamelCase     )
						with tempfile.TemporaryDirectory() as tmp_dir:
							tf.saved_model.save(UpperCamelCase			,				UpperCamelCase			,				signatures={"serving_default": dummy_model.serving}     )
							__lowerCAmelCase												=  tf.saved_model.load(UpperCamelCase     ).signatures["serving_default"]
							for input_row in range(len(UpperCamelCase     )     ):
								__lowerCAmelCase												=  {
								    "input_ids": tf.constant([dummy_input_ids[input_row]]     ),
								    "attention_mask": tf.constant([dummy_attention_masks[input_row]]     ),
								}
								__lowerCAmelCase												=  serving_func(**UpperCamelCase     )["sequences"]
								__lowerCAmelCase												=  test_model.generate(**UpperCamelCase			,				max_new_tokens=UpperCamelCase     )
								tf.debugging.assert_equal(UpperCamelCase			,				UpperCamelCase     )



					@slow
					@require_tensorflow_text
					def        UpperCAmelCase_ (    self     )     ->			Union[str, Any]:
						# TF-only test: tf.saved_model export
						with tempfile.TemporaryDirectory() as tmp_dir:
							# file needed to load the TF tokenizer
							hf_hub_download(repo_id="google/flan-t5-small"			,				filename="spiece.model"			,				local_dir=UpperCamelCase     )






							class 					UpperCAmelCase__	(		tf.keras.layers.Layer ):



												def __init__(    self     )     ->			List[Any]:
													super().__init__()
													__lowerCAmelCase												=  text.SentencepieceTokenizer(
													    model=tf.io.gfile.GFile(os.path.join(UpperCamelCase			,				"spiece.model"     )			,				"rb"     ).read()     )
													__lowerCAmelCase												=  TFAutoModelForSeqaSeqLM.from_pretrained("hf-internal-testing/tiny-random-t5"     )



												def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				*UpperCamelCase			,				**UpperCamelCase     )     ->			Union[str, Any]:
													__lowerCAmelCase												=  self.tokenizer.tokenize(UpperCamelCase     )
													__lowerCAmelCase							,       __lowerCAmelCase												=  text.pad_model_inputs(
													    UpperCamelCase			,				max_seq_length=64			,				pad_value=self.model.config.pad_token_id     )
													__lowerCAmelCase												=  self.model.generate(input_ids=UpperCamelCase			,				attention_mask=UpperCamelCase     )
													return self.tokenizer.detokenize(UpperCamelCase     )


							__lowerCAmelCase												=  CompleteSentenceTransformer()
							__lowerCAmelCase												=  tf.keras.layers.Input(shape=(1,)			,				dtype=tf.string			,				name="inputs"     )
							__lowerCAmelCase												=  complete_model(UpperCamelCase     )
							__lowerCAmelCase												=  tf.keras.Model(UpperCamelCase			,				UpperCamelCase     )
							keras_model.save(UpperCamelCase     )



					def        UpperCAmelCase_ (    self     )     ->			Optional[int]:
						# Has PT equivalent: this test relies on random sampling
						__lowerCAmelCase												=  {
						    "do_sample": True,
						    "num_beams": 1,
						    "top_p": 0.7,
						    "top_k": 10,
						    "temperature": 0.7,
						}
						__lowerCAmelCase												=  14

						__lowerCAmelCase												=  AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2"     )
						__lowerCAmelCase												=  "Hello, my dog is cute and"
						__lowerCAmelCase												=  tokenizer(UpperCamelCase			,				return_tensors="tf"     )
						__lowerCAmelCase												=  TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2"     )

						__lowerCAmelCase												=  638
						# forces the generation to happen on CPU, to avoid GPU-related quirks
						with tf.device(":/CPU:0"     ):
							tf.random.set_seed(0     )
							__lowerCAmelCase												=  model.generate(**UpperCamelCase			,				eos_token_id=UpperCamelCase			,				**UpperCamelCase     )
						self.assertTrue(expectation == len(generated_tokens[0]     )     )

						__lowerCAmelCase												=  [638, 198]
						with tf.device(":/CPU:0"     ):
							tf.random.set_seed(0     )
							__lowerCAmelCase												=  model.generate(**UpperCamelCase			,				eos_token_id=UpperCamelCase			,				**UpperCamelCase     )
						self.assertTrue(expectation == len(generated_tokens[0]     )     )



					def        UpperCAmelCase_ (    self     )     ->			Any:
						# Has PT equivalent: ample use of framework-specific code
						__lowerCAmelCase												=  AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-bart"     )
						__lowerCAmelCase												=  "Hugging Face is a technology company based in New York and Paris."
						__lowerCAmelCase												=  bart_tokenizer(UpperCamelCase			,				return_tensors="tf"     ).input_ids
						__lowerCAmelCase												=  TFBartForConditionalGeneration.from_pretrained("hf-internal-testing/tiny-random-bart"     )
						__lowerCAmelCase												=  bart_model.generate(UpperCamelCase     ).numpy()






						class 					UpperCAmelCase__	(		UpperCamelCase__ ):



											def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase=None			,				**UpperCamelCase     )     ->			Any:
												return super().call(UpperCamelCase			,				**UpperCamelCase     )


						__lowerCAmelCase												=  FakeBart.from_pretrained("hf-internal-testing/tiny-random-bart"     )
						__lowerCAmelCase												=  bart_model.generate(UpperCamelCase			,				foo="bar"     ).numpy()
						self.assertTrue(np.array_equal(UpperCamelCase			,				UpperCamelCase     )     )






						class 					UpperCAmelCase__	(		bart_model.model.encoder.__class__ ):



											def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				**UpperCamelCase     )     ->			List[str]:
												return super().call(UpperCamelCase			,				**UpperCamelCase     )


						__lowerCAmelCase												=  FakeEncoder(bart_model.config			,				bart_model.model.shared     )
						__lowerCAmelCase												=  fake_encoder

						# Normal generation still works (the output will be different because the encoder weights are different)
						__lowerCAmelCase												=  bart_model.generate(UpperCamelCase     ).numpy()
						with self.assertRaises(UpperCamelCase     ):
							# FakeEncoder.call() accepts **kwargs -> no filtering -> value error due to unexpected input "foo"
							bart_model.generate(UpperCamelCase			,				foo="bar"     )




'''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   __lowerCAmelCase					(							lowerCamelCase   :						str				,		lowerCamelCase   :						Union[str, Any]					):






	'''simple docstring'''




	assert isinstance(lowerCamelCase				,		lowerCamelCase					)
	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   __lowerCAmelCase					(							lowerCamelCase   :						Tuple				,		lowerCamelCase   :						int				,		lowerCamelCase   :						Optional[int]					):






	'''simple docstring'''




	__lowerCAmelCase												=  tmp_path / "cache"
	__lowerCAmelCase												=  {"col_1": "string", "col_2": "int64", "col_3": "float64"}
	with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
		__lowerCAmelCase												=  JsonDatasetReader(lowerCamelCase				,		cache_dir=lowerCamelCase				,		keep_in_memory=lowerCamelCase					).read()
	_check_json_dataset(lowerCamelCase				,		lowerCamelCase					)



@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   __lowerCAmelCase					(							lowerCamelCase   :						List[Any]				,		lowerCamelCase   :						int				,		lowerCamelCase   :						Dict					):






	'''simple docstring'''




	__lowerCAmelCase												=  tmp_path / "cache"
	__lowerCAmelCase												=  {"col_1": "string", "col_2": "int64", "col_3": "float64"}
	__lowerCAmelCase												=  features.copy() if features else default_expected_features
	__lowerCAmelCase												=  (
	    Features({feature: Value(lowerCamelCase					) for feature, dtype in features.items()}					) if features is not None else None
	)
	__lowerCAmelCase												=  JsonDatasetReader(lowerCamelCase				,		features=lowerCamelCase				,		cache_dir=lowerCamelCase					).read()
	_check_json_dataset(lowerCamelCase				,		lowerCamelCase					)



@pytest.mark.parametrize(
    "features"				,		[
        None,
        {"col_3": "float64", "col_1": "string", "col_2": "int64"},
    ]				,		)
def   __lowerCAmelCase					(							lowerCamelCase   :						List[Any]				,		lowerCamelCase   :						Optional[int]				,		lowerCamelCase   :						Any					):






	'''simple docstring'''




	__lowerCAmelCase												=  tmp_path / "cache"
	__lowerCAmelCase												=  {"col_3": "float64", "col_1": "string", "col_2": "int64"}
	__lowerCAmelCase												=  features.copy() if features else default_expected_features
	__lowerCAmelCase												=  (
	    Features({feature: Value(lowerCamelCase					) for feature, dtype in features.items()}					) if features is not None else None
	)
	__lowerCAmelCase												=  JsonDatasetReader(lowerCamelCase				,		features=lowerCamelCase				,		cache_dir=lowerCamelCase					).read()
	assert isinstance(lowerCamelCase				,		lowerCamelCase					)
	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   __lowerCAmelCase					(							lowerCamelCase   :						Any				,		lowerCamelCase   :						Optional[int]					):






	'''simple docstring'''




	__lowerCAmelCase												=  {"col_2": "int64", "col_3": "float64", "col_1": "string"}
	__lowerCAmelCase												=  features.copy()
	__lowerCAmelCase												=  (
	    Features({feature: Value(lowerCamelCase					) for feature, dtype in features.items()}					) if features is not None else None
	)
	__lowerCAmelCase												=  tmp_path / "cache"
	__lowerCAmelCase												=  JsonDatasetReader(lowerCamelCase				,		features=lowerCamelCase				,		cache_dir=lowerCamelCase					).read()
	assert isinstance(lowerCamelCase				,		lowerCamelCase					)
	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   __lowerCAmelCase					(							lowerCamelCase   :						List[Any]				,		lowerCamelCase   :						Optional[Any]				,		lowerCamelCase   :						Union[str, Any]					):






	'''simple docstring'''




	__lowerCAmelCase												=  tmp_path / "cache"
	__lowerCAmelCase												=  {"col_1": "string", "col_2": "int64", "col_3": "float64"}
	__lowerCAmelCase												=  JsonDatasetReader(lowerCamelCase				,		cache_dir=lowerCamelCase				,		split=lowerCamelCase					).read()
	_check_json_dataset(lowerCamelCase				,		lowerCamelCase					)
	assert dataset.split == split if split else "train"



@pytest.mark.parametrize("path_type"				,		[str, list]					)
def   __lowerCAmelCase					(							lowerCamelCase   :						Any				,		lowerCamelCase   :						List[str]				,		lowerCamelCase   :						Optional[int]					):






	'''simple docstring'''




	if issubclass(lowerCamelCase				,		lowerCamelCase					):
		__lowerCAmelCase												=  jsonl_path
	elif issubclass(lowerCamelCase				,		lowerCamelCase					):
		__lowerCAmelCase												=  [jsonl_path]
	__lowerCAmelCase												=  tmp_path / "cache"
	__lowerCAmelCase												=  {"col_1": "string", "col_2": "int64", "col_3": "float64"}
	__lowerCAmelCase												=  JsonDatasetReader(lowerCamelCase				,		cache_dir=lowerCamelCase					).read()
	_check_json_dataset(lowerCamelCase				,		lowerCamelCase					)



def   __lowerCAmelCase					(							lowerCamelCase   :						Optional[int]				,		lowerCamelCase   :						int				,		lowerCamelCase   :						str=("train",)					):






	'''simple docstring'''




	assert isinstance(lowerCamelCase				,		lowerCamelCase					)
	for split in splits:
		__lowerCAmelCase												=  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   __lowerCAmelCase					(							lowerCamelCase   :						Dict				,		lowerCamelCase   :						int				,		lowerCamelCase   :						List[str]					):






	'''simple docstring'''




	__lowerCAmelCase												=  tmp_path / "cache"
	__lowerCAmelCase												=  {"col_1": "string", "col_2": "int64", "col_3": "float64"}
	with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
		__lowerCAmelCase												=  JsonDatasetReader({"train": jsonl_path}				,		cache_dir=lowerCamelCase				,		keep_in_memory=lowerCamelCase					).read()
	_check_json_datasetdict(lowerCamelCase				,		lowerCamelCase					)



@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   __lowerCAmelCase					(							lowerCamelCase   :						int				,		lowerCamelCase   :						Any				,		lowerCamelCase   :						Dict					):






	'''simple docstring'''




	__lowerCAmelCase												=  tmp_path / "cache"
	__lowerCAmelCase												=  {"col_1": "string", "col_2": "int64", "col_3": "float64"}
	__lowerCAmelCase												=  features.copy() if features else default_expected_features
	__lowerCAmelCase												=  (
	    Features({feature: Value(lowerCamelCase					) for feature, dtype in features.items()}					) if features is not None else None
	)
	__lowerCAmelCase												=  JsonDatasetReader({"train": jsonl_path}				,		features=lowerCamelCase				,		cache_dir=lowerCamelCase					).read()
	_check_json_datasetdict(lowerCamelCase				,		lowerCamelCase					)



@pytest.mark.parametrize("split"				,		[None, NamedSplit("train"					), "train", "test"]					)
def   __lowerCAmelCase					(							lowerCamelCase   :						int				,		lowerCamelCase   :						Optional[int]				,		lowerCamelCase   :						int					):






	'''simple docstring'''




	if split:
		__lowerCAmelCase												=  {split: jsonl_path}
	else:
		__lowerCAmelCase												=  "train"
		__lowerCAmelCase												=  {"train": jsonl_path, "test": jsonl_path}
	__lowerCAmelCase												=  tmp_path / "cache"
	__lowerCAmelCase												=  {"col_1": "string", "col_2": "int64", "col_3": "float64"}
	__lowerCAmelCase												=  JsonDatasetReader(lowerCamelCase				,		cache_dir=lowerCamelCase					).read()
	_check_json_datasetdict(lowerCamelCase				,		lowerCamelCase				,		splits=list(path.keys()					)					)
	assert all(dataset[split].split == split for split in path.keys()					)



def   __lowerCAmelCase					(							lowerCamelCase   :						Optional[Any]					):






	'''simple docstring'''




	return json.load(lowerCamelCase					)



def   __lowerCAmelCase					(							lowerCamelCase   :						Tuple					):






	'''simple docstring'''




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






class 					UpperCAmelCase__	:



					@pytest.mark.parametrize("lines, load_json_function"			,				[(True, load_json_lines), (False, load_json)]     )
					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )     ->			int:
						with io.BytesIO() as buffer:
							JsonDatasetWriter(UpperCamelCase			,				UpperCamelCase			,				lines=UpperCamelCase     ).write()
							buffer.seek(0     )
							__lowerCAmelCase												=  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        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )     ->			int:
						with io.BytesIO() as buffer:
							JsonDatasetWriter(UpperCamelCase			,				UpperCamelCase			,				lines=UpperCamelCase			,				orient=UpperCamelCase     ).write()
							buffer.seek(0     )
							__lowerCAmelCase												=  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        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )     ->			str:
						with io.BytesIO() as buffer:
							JsonDatasetWriter(UpperCamelCase			,				UpperCamelCase			,				lines=UpperCamelCase			,				num_proc=2     ).write()
							buffer.seek(0     )
							__lowerCAmelCase												=  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        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )     ->			Union[str, Any]:
						with io.BytesIO() as buffer:
							JsonDatasetWriter(UpperCamelCase			,				UpperCamelCase			,				lines=UpperCamelCase			,				orient=UpperCamelCase			,				num_proc=2     ).write()
							buffer.seek(0     )
							__lowerCAmelCase												=  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        UpperCAmelCase_ (    self			,				UpperCamelCase     )     ->			Any:
						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        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )     ->			Tuple:
						__lowerCAmelCase												=  tmp_path_factory.mktemp("data"     ) / F'''test.json.{extension}'''
						__lowerCAmelCase												=  str(shared_datadir / F'''test_file.json.{extension}'''     )
						JsonDatasetWriter(UpperCamelCase			,				UpperCamelCase			,				compression=UpperCamelCase     ).write()

						with fsspec.open(UpperCamelCase			,				"rb"			,				compression="infer"     ) as f:
							__lowerCAmelCase												=  f.read()
						with fsspec.open(UpperCamelCase			,				"rb"			,				compression="infer"     ) as f:
							__lowerCAmelCase												=  f.read()
						assert exported_content == original_content