



import gc
import importlib.metadata
import tempfile
import unittest

from packaging import version

from transformers import (
    AutoModel,
    AutoModelForCausalLM,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    BitsAndBytesConfig,
    pipeline,
)
from transformers.testing_utils import (
    is_torch_available,
    require_accelerate,
    require_bitsandbytes,
    require_torch,
    require_torch_gpu,
    require_torch_multi_gpu,
    slow,
)

def 			A_      (	__a     :				Optional[Any]		):

				"""simple docstring"""
				if model.config.model_type == "gpt2":
								return model.transformer.h[0].mlp.c_fc
				return model.transformer.h[0].mlp.dense_ah_to_h


if is_torch_available():
	import torch
	import torch.nn as nn

	class  __snake_case      (   nn.Module):







								'''simple docstring'''






								def __init__(							self  ,  a_  ,  a_		):
												super().__init__()
												a__				      =		module
												a__				      =		nn.Sequential(
												    nn.Linear(module.in_features  ,  a_  ,  bias=a_		)  ,  nn.Linear(a_  ,  module.out_features  ,  bias=a_		)  ,  )
												a__				      =		(2.0 / (5 * min(module.in_features  ,  module.out_features		))) ** 0.5
												nn.init.normal_(self.adapter[0].weight  ,  std=a_		)
												nn.init.zeros_(self.adapter[1].weight		)
												self.adapter.to(module.weight.device		)







								def _a      (							self  ,  a_  ,  *a_  ,  **a_		):
												return self.module(a_  ,  *a_  ,  **a_		) + self.adapter(a_		)

@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class  __snake_case      (   unittest.TestCase):







							'''simple docstring'''






							UpperCamelCase__				:      Optional[Any]         =					"""bigscience/bloom-1b7"""

							# Constant values
							UpperCamelCase__				:      List[Any]         =					2.109659552692574

							UpperCamelCase__				:      str         =					"""Hello my name is"""
							UpperCamelCase__				:      Any         =					set()
							EXPECTED_OUTPUTS.add("""Hello my name is John and I am a professional photographer. I""")
							EXPECTED_OUTPUTS.add("""Hello my name is John.\nI am a friend of your father.\n""")
							EXPECTED_OUTPUTS.add("""Hello my name is John Doe, I am a student at the University""")
							UpperCamelCase__				:      List[str]         =					1_0


							def _a      (							self		):
											# Models and tokenizer
											a__				      =		AutoTokenizer.from_pretrained(self.model_name		)




class  __snake_case      (   SCREAMING_SNAKE_CASE):







							'''simple docstring'''






							def _a      (							self		):
											super().setUp()

											# Models and tokenizer
											a__				      =		AutoModelForCausalLM.from_pretrained(
											    self.model_name  ,  torch_dtype=torch.floataa  ,  device_map="""auto"""		)
											a__				      =		AutoModelForCausalLM.from_pretrained(self.model_name  ,  load_in_abit=a_  ,  device_map="""auto"""		)


							def _a      (							self		):
											del self.model_fpaa
											del self.model_abit

											gc.collect()
											torch.cuda.empty_cache()


							def _a      (							self		):
											a__				      =		self.model_abit.config

											self.assertTrue(hasattr(a_  ,  """quantization_config"""		)		)

											a__				      =		config.to_dict()
											a__				      =		config.to_diff_dict()

											a__				      =		config.to_json_string()


							def _a      (							self		):
											from bitsandbytes.nn import Paramsabit

											a__				      =		self.model_fpaa.get_memory_footprint()
											a__				      =		self.model_abit.get_memory_footprint()

											self.assertAlmostEqual(mem_fpaa / mem_abit  ,  self.EXPECTED_RELATIVE_DIFFERENCE		)
											a__				      =		get_some_linear_layer(self.model_abit		)
											self.assertTrue(linear.weight.__class__ == Paramsabit		)


							def _a      (							self		):
											from transformers import TaPreTrainedModel

											self.model_fpaa.get_memory_footprint()
											self.model_abit.get_memory_footprint()

											for name, module in self.model_abit.named_modules():
															if isinstance(a_  ,  torch.nn.Linear		):
																			if name not in ["lm_head"] + TaPreTrainedModel._keep_in_fpaa_modules:
																							# 4-bit parameters are packed in uint8 variables
																							self.assertTrue(module.weight.dtype == torch.uinta		)


							def _a      (							self		):
											a__				      =		self.tokenizer(self.input_text  ,  return_tensors="""pt"""		)
											a__				      =		self.model_abit.generate(input_ids=encoded_input["""input_ids"""].to(0		)  ,  max_new_tokens=10		)

											self.assertIn(self.tokenizer.decode(output_sequences[0]  ,  skip_special_tokens=a_		)  ,  self.EXPECTED_OUTPUTS		)


							def _a      (							self		):
											a__				      =		BitsAndBytesConfig()
											a__				      =		True

											a__				      =		AutoModelForCausalLM.from_pretrained(
											    self.model_name  ,  quantization_config=a_  ,  device_map="""auto"""		)

											a__				      =		self.tokenizer(self.input_text  ,  return_tensors="""pt"""		)
											a__				      =		model_abit_from_config.generate(
											    input_ids=encoded_input["""input_ids"""].to(0		)  ,  max_new_tokens=10		)

											self.assertIn(self.tokenizer.decode(output_sequences[0]  ,  skip_special_tokens=a_		)  ,  self.EXPECTED_OUTPUTS		)


							def _a      (							self		):
											with self.assertRaises(a_		), tempfile.TemporaryDirectory() as tmpdirname:
															self.model_abit.save_pretrained(a_		)


							def _a      (							self		):
											a__				      =		BitsAndBytesConfig()

											with self.assertRaises(a_		):
															a__				      =		AutoModelForCausalLM.from_pretrained(
															    self.model_name  ,  quantization_config=a_  ,  load_in_abit=a_  ,  device_map="""auto"""  ,  bnb_abit_quant_type="""nf4"""  ,  )


							def _a      (							self		):
											with self.assertRaises(a_		):
															# Tries with `str`
															self.model_abit.to("""cpu"""		)

											with self.assertRaises(a_		):
															# Tries with a `dtype``
															self.model_abit.to(torch.floataa		)

											with self.assertRaises(a_		):
															# Tries with a `device`
															self.model_abit.to(torch.device("""cuda:0"""		)		)

											with self.assertRaises(a_		):
															# Tries with a `device`
															self.model_abit.float()

											with self.assertRaises(a_		):
															# Tries with a `device`
															self.model_abit.half()

											# Test if we did not break anything
											a__				      =		self.tokenizer(self.input_text  ,  return_tensors="""pt"""		)

											a__				      =		self.model_fpaa.to(torch.floataa		)
											a__				      =		self.model_fpaa.generate(input_ids=encoded_input["""input_ids"""].to(0		)  ,  max_new_tokens=10		)

											# Check this does not throw an error
											a__				      =		self.model_fpaa.to("""cpu"""		)

											# Check this does not throw an error
											a__				      =		self.model_fpaa.half()

											# Check this does not throw an error
											a__				      =		self.model_fpaa.float()







							def _a      (							self		):
											a__				      =		AutoModelForSeqaSeqLM.from_pretrained("""t5-small"""  ,  load_in_abit=a_  ,  device_map="""auto"""		)
											self.assertTrue(model.decoder.block[0].layer[2].DenseReluDense.wo.weight.dtype == torch.floataa		)




@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class  __snake_case      (   unittest.TestCase):







							'''simple docstring'''






							@classmethod
							def _a      (							cls		):
											a__				      =		"""t5-small"""
											a__				      =		"""google/flan-t5-small"""  # flan-t5 uses dense-act instead of dense-relu-dense
											a__				      =		AutoTokenizer.from_pretrained(cls.model_name		)
											a__				      =		"""Translate in German: Hello, my dog is cute"""


							def _a      (							self		):
											gc.collect()
											torch.cuda.empty_cache()


							def _a      (							self		):
											from transformers import TaForConditionalGeneration

											a__				      =		TaForConditionalGeneration._keep_in_fpaa_modules
											a__				      =		None

											# test with `t5-small`
											a__				      =		TaForConditionalGeneration.from_pretrained(self.model_name  ,  load_in_abit=a_  ,  device_map="""auto"""		)
											a__				      =		self.tokenizer(self.input_text  ,  return_tensors="""pt"""		).to(0		)
											a__				      =		model.generate(**a_		)

											# test with `flan-t5-small`
											a__				      =		TaForConditionalGeneration.from_pretrained(
											    self.dense_act_model_name  ,  load_in_abit=a_  ,  device_map="""auto"""		)
											a__				      =		self.tokenizer(self.input_text  ,  return_tensors="""pt"""		).to(0		)
											a__				      =		model.generate(**a_		)
											a__				      =		modules







							def _a      (							self		):
											import bitsandbytes as bnb

											from transformers import TaForConditionalGeneration

											# test with `t5-small`
											a__				      =		TaForConditionalGeneration.from_pretrained(self.model_name  ,  load_in_abit=a_  ,  device_map="""auto"""		)

											# there was a bug with decoders - this test checks that it is fixed
											self.assertTrue(isinstance(model.decoder.block[0].layer[0].SelfAttention.q  ,  bnb.nn.Linearabit		)		)

											a__				      =		self.tokenizer(self.input_text  ,  return_tensors="""pt"""		).to(0		)
											a__				      =		model.generate(**a_		)

											# test with `flan-t5-small`
											a__				      =		TaForConditionalGeneration.from_pretrained(
											    self.dense_act_model_name  ,  load_in_abit=a_  ,  device_map="""auto"""		)
											a__				      =		self.tokenizer(self.input_text  ,  return_tensors="""pt"""		).to(0		)
											a__				      =		model.generate(**a_		)




class  __snake_case      (   SCREAMING_SNAKE_CASE):







							'''simple docstring'''






							def _a      (							self		):
											super().setUp()
											# model_name
											a__				      =		"""bigscience/bloom-560m"""
											a__				      =		"""t5-small"""

											# Different types of model

											a__				      =		AutoModel.from_pretrained(self.model_name  ,  load_in_abit=a_  ,  device_map="""auto"""		)
											# Sequence classification model
											a__				      =		AutoModelForSequenceClassification.from_pretrained(
											    self.model_name  ,  load_in_abit=a_  ,  device_map="""auto"""		)
											# CausalLM model
											a__				      =		AutoModelForCausalLM.from_pretrained(self.model_name  ,  load_in_abit=a_  ,  device_map="""auto"""		)
											# Seq2seq model
											a__				      =		AutoModelForSeqaSeqLM.from_pretrained(
											    self.seq_to_seq_name  ,  load_in_abit=a_  ,  device_map="""auto"""		)


							def _a      (							self		):
											del self.base_model
											del self.sequence_model
											del self.model_abit
											del self.seq_to_seq_model

											gc.collect()
											torch.cuda.empty_cache()







							def _a      (							self		):
											from bitsandbytes.nn import Paramsabit

											self.assertTrue(self.base_model.h[-1].mlp.dense_ah_to_h.weight.__class__ == Paramsabit		)

											# Other heads should be nn.Parameter
											self.assertTrue(self.model_abit.lm_head.weight.__class__ == torch.nn.Parameter		)
											self.assertTrue(self.sequence_model.score.weight.__class__ == torch.nn.Parameter		)
											self.assertTrue(self.seq_to_seq_model.lm_head.weight.__class__ == torch.nn.Parameter		)




class  __snake_case      (   SCREAMING_SNAKE_CASE):







							'''simple docstring'''






							def _a      (							self		):
											super().setUp()


							def _a      (							self		):
											del self.pipe

											gc.collect()
											torch.cuda.empty_cache()







							def _a      (							self		):
											a__				      =		pipeline(
											    """text-generation"""  ,  model=self.model_name  ,  model_kwargs={"""device_map""": """auto""", """load_in_4bit""": True, """torch_dtype""": torch.floataa}  ,  max_new_tokens=self.MAX_NEW_TOKENS  ,  )

											# Real second forward pass
											a__				      =		self.pipe(self.input_text		)
											self.assertIn(pipeline_output[0]["""generated_text"""]  ,  self.EXPECTED_OUTPUTS		)




@require_torch_multi_gpu
class  __snake_case      (   SCREAMING_SNAKE_CASE):







							'''simple docstring'''






							def _a      (							self		):
											super().setUp()







							def _a      (							self		):
											a__				      =		AutoModelForCausalLM.from_pretrained(
											    self.model_name  ,  load_in_abit=a_  ,  device_map="""balanced"""		)

											# Check correct device map
											self.assertEqual(set(model_parallel.hf_device_map.values()		)  ,  {0, 1}		)

											# Check that inference pass works on the model
											a__				      =		self.tokenizer(self.input_text  ,  return_tensors="""pt"""		)

											# Second real batch
											a__				      =		model_parallel.generate(input_ids=encoded_input["""input_ids"""].to(0		)  ,  max_new_tokens=10		)
											self.assertIn(self.tokenizer.decode(output_parallel[0]  ,  skip_special_tokens=a_		)  ,  self.EXPECTED_OUTPUTS		)




class  __snake_case      (   SCREAMING_SNAKE_CASE):







							'''simple docstring'''






							def _a      (							self		):
											a__				      =		"""facebook/opt-350m"""
											super().setUp()







							def _a      (							self		):
											if version.parse(importlib.metadata.version("""bitsandbytes"""		)		) < version.parse("""0.37.0"""		):
															return

											# Step 1: freeze all parameters
											a__				      =		AutoModelForCausalLM.from_pretrained(self.model_name  ,  load_in_abit=a_		)

											self.assertEqual(set(model.hf_device_map.values()		)  ,  {torch.cuda.current_device()}		)

											for param in model.parameters():
															a__				      =		False  # freeze the model - train adapters later
															if param.ndim == 1:
																			# cast the small parameters (e.g. layernorm) to fp32 for stability
																			a__				      =		param.data.to(torch.floataa		)

        # Step 2: add adapters
											for _, module in model.named_modules():
															if "OPTAttention" in repr(type(a_		)		):
																			a__				      =		LoRALayer(module.q_proj  ,  rank=16		)
																			a__				      =		LoRALayer(module.k_proj  ,  rank=16		)
																			a__				      =		LoRALayer(module.v_proj  ,  rank=16		)

        # Step 3: dummy batch
											a__				      =		self.tokenizer("""Test batch """  ,  return_tensors="""pt"""		).to(0		)

											# Step 4: Check if the gradient is not None
											with torch.cuda.amp.autocast():
															a__				      =		model.forward(**a_		)
															out.logits.norm().backward()

											for module in model.modules():
															if isinstance(a_  ,  a_		):
																			self.assertTrue(module.adapter[1].weight.grad is not None		)
																			self.assertTrue(module.adapter[1].weight.grad.norm().item() > 0		)
															elif isinstance(a_  ,  nn.Embedding		):
																			self.assertTrue(module.weight.grad is None		)




class  __snake_case      (   SCREAMING_SNAKE_CASE):







							'''simple docstring'''






							UpperCamelCase__				:      Any         =					"""gpt2-xl"""
							UpperCamelCase__				:      Any         =					3.3191854854152187






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     :				List[str] ,		__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     :				Tuple ,		__a     :				Dict ,		__a     :				List[str]		):

				"""simple docstring"""
				a__				      =		tmp_path / """cache"""
				a__				      =		{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
				with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
								a__				      =		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     :				int ,		__a     :				Dict ,		__a     :				List[Any]		):

				"""simple docstring"""
				a__				      =		tmp_path / """cache"""
				a__				      =		{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
				a__				      =		features.copy() if features else default_expected_features
				a__				      =		(
				    Features({feature: Value(__a		) for feature, dtype in features.items()}		) if features is not None else None
				)
				a__				      =		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     :				Tuple ,		__a     :				Tuple ,		__a     :				Tuple		):

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

				"""simple docstring"""
				# jsonl_312_path features are {"col_3": "float64", "col_1": "string", "col_2": "int64"}
				a__				      =		{"""col_2""": """int64""", """col_3""": """float64""", """col_1""": """string"""}
				a__				      =		features.copy()
				a__				      =		(
				    Features({feature: Value(__a		) for feature, dtype in features.items()}		) if features is not None else None
				)
				a__				      =		tmp_path / """cache"""
				a__				      =		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     :				List[Any] ,		__a     :				int ,		__a     :				List[str]		):

				"""simple docstring"""
				a__				      =		tmp_path / """cache"""
				a__				      =		{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
				a__				      =		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     :				Dict ,		__a     :				Tuple ,		__a     :				List[str]		):

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

def 			A_      (	__a     :				str ,		__a     :				List[Any] ,		__a     :				Dict=("train",)		):

				"""simple docstring"""
				assert isinstance(__a ,		__a		)
				for split in splits:
								a__				      =		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     :				int ,		__a     :				int ,		__a     :				Dict		):

				"""simple docstring"""
				a__				      =		tmp_path / """cache"""
				a__				      =		{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
				with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
								a__				      =		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     :				Tuple ,		__a     :				Dict ,		__a     :				Optional[Any]		):

				"""simple docstring"""
				a__				      =		tmp_path / """cache"""
				a__				      =		{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
				a__				      =		features.copy() if features else default_expected_features
				a__				      =		(
				    Features({feature: Value(__a		) for feature, dtype in features.items()}		) if features is not None else None
				)
				a__				      =		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     :				Dict ,		__a     :				Any ,		__a     :				Tuple		):

				"""simple docstring"""
				if split:
								a__				      =		{split: jsonl_path}
				else:
								a__				      =		"""train"""
								a__				      =		{"""train""": jsonl_path, """test""": jsonl_path}
				a__				      =		tmp_path / """cache"""
				a__				      =		{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
				a__				      =		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     :				Any		):

				"""simple docstring"""
				return json.load(__a		)

def 			A_      (	__a     :				Tuple		):

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




class  __snake_case      :







							'''simple docstring'''






							@pytest.mark.parametrize("""lines, load_json_function"""  ,  [(True, load_json_lines), (False, load_json)]		)
							def _a      (							self  ,  a_  ,  a_  ,  a_		):
											with io.BytesIO() as buffer:
															JsonDatasetWriter(a_  ,  a_  ,  lines=a_		).write()
															buffer.seek(0		)
															a__				      =		load_json_function(a_		)
											assert isinstance(a_  ,  a_		)
											assert isinstance(exported_content[0]  ,  a_		)
											assert len(a_		) == 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 _a      (							self  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_		):
											with io.BytesIO() as buffer:
															JsonDatasetWriter(a_  ,  a_  ,  lines=a_  ,  orient=a_		).write()
															buffer.seek(0		)
															a__				      =		load_json(a_		)
											assert isinstance(a_  ,  a_		)
											if keys:
															if container is dict:
																			assert exported_content.keys() == keys
															else:
																			assert exported_content[0].keys() == keys
											else:
															assert not hasattr(a_  ,  """keys"""		) and not hasattr(exported_content[0]  ,  """keys"""		)
											if len_at:
															assert len(exported_content[len_at]		) == 10
											else:
															assert len(a_		) == 10


							@pytest.mark.parametrize("""lines, load_json_function"""  ,  [(True, load_json_lines), (False, load_json)]		)
							def _a      (							self  ,  a_  ,  a_  ,  a_		):
											with io.BytesIO() as buffer:
															JsonDatasetWriter(a_  ,  a_  ,  lines=a_  ,  num_proc=2		).write()
															buffer.seek(0		)
															a__				      =		load_json_function(a_		)
											assert isinstance(a_  ,  a_		)
											assert isinstance(exported_content[0]  ,  a_		)
											assert len(a_		) == 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 _a      (							self  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_		):
											with io.BytesIO() as buffer:
															JsonDatasetWriter(a_  ,  a_  ,  lines=a_  ,  orient=a_  ,  num_proc=2		).write()
															buffer.seek(0		)
															a__				      =		load_json(a_		)
											assert isinstance(a_  ,  a_		)
											if keys:
															if container is dict:
																			assert exported_content.keys() == keys
															else:
																			assert exported_content[0].keys() == keys
											else:
															assert not hasattr(a_  ,  """keys"""		) and not hasattr(exported_content[0]  ,  """keys"""		)
											if len_at:
															assert len(exported_content[len_at]		) == 10
											else:
															assert len(a_		) == 10


							def _a      (							self  ,  a_		):
											with pytest.raises(a_		):
															with io.BytesIO() as buffer:
																			JsonDatasetWriter(a_  ,  a_  ,  num_proc=0		)







							@pytest.mark.parametrize("""compression, extension"""  ,  [("""gzip""", """gz"""), ("""bz2""", """bz2"""), ("""xz""", """xz""")]		)
							def _a      (							self  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_		):
											a__				      =		tmp_path_factory.mktemp("""data"""		) / F'''test.json.{extension}'''
											a__				      =		str(shared_datadir / F'''test_file.json.{extension}'''		)
											JsonDatasetWriter(a_  ,  a_  ,  compression=a_		).write()

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

