




'''simple docstring'''





import inspect
import unittest

from transformers import DecisionTransformerConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import DecisionTransformerModel
		from transformers.models.decision_transformer.modeling_decision_transformer import (
		    DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
		)


class 							__UpperCamelCase			:





				'''simple docstring'''



				def __init__( self       , UpperCAmelCase_       , UpperCAmelCase_=13       , UpperCAmelCase_=7       , UpperCAmelCase_=6       , UpperCAmelCase_=17       , UpperCAmelCase_=23       , UpperCAmelCase_=11       , UpperCAmelCase_=True       , ):
											lowerCAmelCase   =		parent
											lowerCAmelCase   =		batch_size
											lowerCAmelCase   =		seq_length
											lowerCAmelCase   =		act_dim
											lowerCAmelCase   =		state_dim
											lowerCAmelCase   =		hidden_size
											lowerCAmelCase   =		max_length
											lowerCAmelCase   =		is_training



				def       __snake_case		( self						):
											lowerCAmelCase   =		floats_tensor((self.batch_size, self.seq_length, self.state_dim)						)
											lowerCAmelCase   =		floats_tensor((self.batch_size, self.seq_length, self.act_dim)						)
											lowerCAmelCase   =		floats_tensor((self.batch_size, self.seq_length, 1)						)
											lowerCAmelCase   =		floats_tensor((self.batch_size, self.seq_length, 1)						)
											lowerCAmelCase   =		ids_tensor((self.batch_size, self.seq_length)       , vocab_size=10_00						)
											lowerCAmelCase   =		random_attention_mask((self.batch_size, self.seq_length)						)

											lowerCAmelCase   =		self.get_config()

											return (
											    config,
											    states,
											    actions,
											    rewards,
											    returns_to_go,
											    timesteps,
											    attention_mask,
											)



				def       __snake_case		( self						):
											return DecisionTransformerConfig(
											    batch_size=self.batch_size       , seq_length=self.seq_length       , act_dim=self.act_dim       , state_dim=self.state_dim       , hidden_size=self.hidden_size       , max_length=self.max_length       , )



				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , ):
											lowerCAmelCase   =		DecisionTransformerModel(config=__a						)
											model.to(__a						)
											model.eval()
											lowerCAmelCase   =		model(__a       , __a       , __a       , __a       , __a       , __a						)

											self.parent.assertEqual(result.state_preds.shape       , states.shape						)
											self.parent.assertEqual(result.action_preds.shape       , actions.shape						)
											self.parent.assertEqual(result.return_preds.shape       , returns_to_go.shape						)
											self.parent.assertEqual(
											    result.last_hidden_state.shape       , (self.batch_size, self.seq_length * 3, self.hidden_size)						)  # seq length *3 as there are 3 modelities: states, returns and actions






				def       __snake_case		( self						):
											lowerCAmelCase   =		self.prepare_config_and_inputs()
											(
											    (
											    lowerCAmelCase						
											)   ,     (
											    lowerCAmelCase						
											)   ,     (
											    lowerCAmelCase						
											)   ,     (
											    lowerCAmelCase						
											)   ,     (
											    lowerCAmelCase						
											)   ,     (
											    lowerCAmelCase						
											)   ,     (
											    lowerCAmelCase						
											)   ,     						
											)   =		config_and_inputs
											lowerCAmelCase   =		{
											    '''states''': states,
											    '''actions''': actions,
											    '''rewards''': rewards,
											    '''returns_to_go''': returns_to_go,
											    '''timesteps''': timesteps,
											    '''attention_mask''': attention_mask,
											}
											return config, inputs_dict




@require_torch
class 							__UpperCamelCase			( __UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					unittest.TestCase  ):





				'''simple docstring'''



				__a							: List[Any]										=(DecisionTransformerModel,) if is_torch_available() else ()
				__a							: Any										=()
				__a							: List[Any]										={"""feature-extraction""": DecisionTransformerModel} if is_torch_available() else {}

				# Ignoring of a failing test from GenerationTesterMixin, as the model does not use inputs_ids
				__a							: Dict										=False

				# Ignoring of a failing tests from ModelTesterMixin, as the model does not implement these features
				__a							: int										=False
				__a							: List[Any]										=False
				__a							: List[Any]										=False
				__a							: int										=False
				__a							: Dict										=False
				__a							: Any										=False
				__a							: Dict										=False
				__a							: List[str]										=False
				__a							: str										=False



				def       __snake_case		( self						):
											lowerCAmelCase   =		DecisionTransformerModelTester(self						)
											lowerCAmelCase   =		ConfigTester(self       , config_class=__a       , hidden_size=37						)



				def       __snake_case		( self						):
											self.config_tester.run_common_tests()



				def       __snake_case		( self						):
											lowerCAmelCase   =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*__a						)



				@slow
				def       __snake_case		( self						):
											for model_name in DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																		lowerCAmelCase   =		DecisionTransformerModel.from_pretrained(__a						)
																		self.assertIsNotNone(__a						)






				def       __snake_case		( self						):
											lowerCAmelCase   ,     lowerCAmelCase   =		self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																		lowerCAmelCase   =		model_class(__a						)
																		lowerCAmelCase   =		inspect.signature(model.forward						)
																		# signature.parameters is an OrderedDict => so arg_names order is deterministic
																		lowerCAmelCase   =		[*signature.parameters.keys()]

																		lowerCAmelCase   =		[
																		    '''states''',
																		    '''actions''',
																		    '''rewards''',
																		    '''returns_to_go''',
																		    '''timesteps''',
																		    '''attention_mask''',
																		]

																		self.assertListEqual(arg_names[: len(__a						)]       , __a						)




@require_torch
class 							__UpperCamelCase			( unittest.TestCase  ):





				'''simple docstring'''



				@slow
				def       __snake_case		( self						):
											lowerCAmelCase   =		2  # number of steps of autoregressive prediction we will perform
											lowerCAmelCase   =		10  # defined by the RL environment, may be normalized
											lowerCAmelCase   =		DecisionTransformerModel.from_pretrained('''edbeeching/decision-transformer-gym-hopper-expert'''						)
											lowerCAmelCase   =		model.to(__a						)
											lowerCAmelCase   =		model.config
											torch.manual_seed(0						)
											lowerCAmelCase   =		torch.randn(1       , 1       , config.state_dim						).to(device=__a       , dtype=torch.floataa						)  # env.reset()

											lowerCAmelCase   =		torch.tensor(
											    [[0.242793, -0.28693074, 0.8742613], [0.67815274, -0.08101085, -0.12952147]]       , device=__a						)

											lowerCAmelCase   =		torch.tensor(__a       , device=__a       , dtype=torch.floataa						).reshape(1       , 1       , 1						)
											lowerCAmelCase   =		state
											lowerCAmelCase   =		torch.zeros(1       , 0       , config.act_dim       , device=__a       , dtype=torch.floataa						)
											lowerCAmelCase   =		torch.zeros(1       , 0       , device=__a       , dtype=torch.floataa						)
											lowerCAmelCase   =		torch.tensor(0       , device=__a       , dtype=torch.long						).reshape(1       , 1						)

											for step in range(__a						):
																		lowerCAmelCase   =		torch.cat([actions, torch.zeros(1       , 1       , config.act_dim       , device=__a						)]       , dim=1						)
																		lowerCAmelCase   =		torch.cat([rewards, torch.zeros(1       , 1       , device=__a						)]       , dim=1						)

																		lowerCAmelCase   =		torch.ones(1       , states.shape[1]						).to(dtype=torch.long       , device=states.device						)

																		with torch.no_grad():
																									lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   =		model(
																									    states=__a       , actions=__a       , rewards=__a       , returns_to_go=__a       , timesteps=__a       , attention_mask=__a       , return_dict=__a       , )

																		self.assertEqual(action_pred.shape       , actions.shape						)
																		self.assertTrue(torch.allclose(action_pred[0, -1]       , expected_outputs[step]       , atol=1E-4						)						)
																		lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   =		(  # env.step(action)
																		    torch.randn(1       , 1       , config.state_dim						).to(device=__a       , dtype=torch.floataa						),
																		    1.0,
																		    False,
																		    {},
																		)

																		lowerCAmelCase   =		action_pred[0, -1]
																		lowerCAmelCase   =		torch.cat([states, state]       , dim=1						)
																		lowerCAmelCase   =		returns_to_go[0, -1] - reward
																		lowerCAmelCase   =		torch.cat([returns_to_go, pred_return.reshape(1       , 1       , 1						)]       , dim=1						)
																		lowerCAmelCase   =		torch.cat(
																		    [timesteps, torch.ones((1, 1)       , device=__a       , dtype=torch.long						) * (step + 1)]       , dim=1						)









import io
import itertools
import json
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa
import pyarrow.json as paj

import datasets
from datasets.table import table_cast
from datasets.utils.file_utils import readline


UpperCAmelCase_						=datasets.utils.logging.get_logger(__name__)


@dataclass
class 							__UpperCamelCase			( datasets.BuilderConfig  ):





				'''simple docstring'''



				__a							: Optional[datasets.Features]										=None
				__a							: str										="utf-8"
				__a							: Optional[str]										=None
				__a							: Optional[str]										=None
				__a							: bool										=True  # deprecated
				__a							: Optional[int]										=None  # deprecated
				__a							: int										=1_0 << 2_0  # 10MB
				__a							: Optional[bool]										=None




class 							__UpperCamelCase			( datasets.ArrowBasedBuilder  ):





				'''simple docstring'''



				__a							: str										=JsonConfig



				def       __snake_case		( self						):
											if self.config.block_size is not None:
																		logger.warning('''The JSON loader parameter `block_size` is deprecated. Please use `chunksize` instead'''						)
																		lowerCAmelCase   =		self.config.block_size
											if self.config.use_threads is not True:
																		logger.warning(
																		    '''The JSON loader parameter `use_threads` is deprecated and doesn\'t have any effect anymore.'''						)
											if self.config.newlines_in_values is not None:
																		raise ValueError('''The JSON loader parameter `newlines_in_values` is no longer supported'''						)
											return datasets.DatasetInfo(features=self.config.features						)



				def       __snake_case		( self       , UpperCAmelCase_						):
											if not self.config.data_files:
																		raise ValueError(F"""At least one data file must be specified, but got data_files={self.config.data_files}"""						)
											lowerCAmelCase   =		dl_manager.download_and_extract(self.config.data_files						)
											if isinstance(UpperCAmelCase_       , (str, list, tuple)						):
																		lowerCAmelCase   =		data_files
																		if isinstance(UpperCAmelCase_       , UpperCAmelCase_						):
																									lowerCAmelCase   =		[files]
																		lowerCAmelCase   =		[dl_manager.iter_files(UpperCAmelCase_						) for file in files]
																		return [datasets.SplitGenerator(name=datasets.Split.TRAIN       , gen_kwargs={'''files''': files}						)]
											lowerCAmelCase   =		[]
											for split_name, files in data_files.items():
																		if isinstance(UpperCAmelCase_       , UpperCAmelCase_						):
																									lowerCAmelCase   =		[files]
																		lowerCAmelCase   =		[dl_manager.iter_files(UpperCAmelCase_						) for file in files]
																		splits.append(datasets.SplitGenerator(name=UpperCAmelCase_       , gen_kwargs={'''files''': files}						)						)
											return splits



				def       __snake_case		( self       , UpperCAmelCase_						):
											if self.config.features is not None:
																		# adding missing columns
																		for column_name in set(self.config.features						) - set(pa_table.column_names						):
																									lowerCAmelCase   =		self.config.features.arrow_schema.field(UpperCAmelCase_						).type
																									lowerCAmelCase   =		pa_table.append_column(UpperCAmelCase_       , pa.array([None] * len(UpperCAmelCase_						)       , type=UpperCAmelCase_						)						)
																		# more expensive cast to support nested structures with keys in a different order
																		# allows str <-> int/float or str to Audio for example
																		lowerCAmelCase   =		table_cast(UpperCAmelCase_       , self.config.features.arrow_schema						)
											return pa_table






				def       __snake_case		( self       , UpperCAmelCase_						):
											for file_idx, file in enumerate(itertools.chain.from_iterable(UpperCAmelCase_						)						):
																		# If the file is one json object and if we need to look at the list of items in one specific field
																		if self.config.field is not None:
																									with open(UpperCAmelCase_       , encoding=self.config.encoding       , errors=self.config.encoding_errors						) as f:
																																lowerCAmelCase   =		json.load(UpperCAmelCase_						)

																									# We keep only the field we are interested in
																									lowerCAmelCase   =		dataset[self.config.field]

																									# We accept two format: a list of dicts or a dict of lists
																									if isinstance(UpperCAmelCase_       , (list, tuple)						):
																																lowerCAmelCase   =		set().union(*[row.keys() for row in dataset]						)
																																lowerCAmelCase   =		{col: [row.get(UpperCAmelCase_						) for row in dataset] for col in keys}
																									else:
																																lowerCAmelCase   =		dataset
																									lowerCAmelCase   =		pa.Table.from_pydict(UpperCAmelCase_						)
																									yield file_idx, self._cast_table(UpperCAmelCase_						)

																		# If the file has one json object per line
																		else:
																									with open(UpperCAmelCase_       , '''rb'''						) as f:
																																lowerCAmelCase   =		0
																																# Use block_size equal to the chunk size divided by 32 to leverage multithreading
																																# Set a default minimum value of 16kB if the chunk size is really small
																																lowerCAmelCase   =		max(self.config.chunksize // 32       , 16 << 10						)
																																lowerCAmelCase   =		(
																																    self.config.encoding_errors if self.config.encoding_errors is not None else '''strict'''
																																)
																																while True:
																																							lowerCAmelCase   =		f.read(self.config.chunksize						)
																																							if not batch:
																																														break
																																							# Finish current line
																																							try:
																																														batch += f.readline()
																																							except (AttributeError, io.UnsupportedOperation):
																																														batch += readline(UpperCAmelCase_						)
																																							# PyArrow only accepts utf-8 encoded bytes
																																							if self.config.encoding != "utf-8":
																																														lowerCAmelCase   =		batch.decode(self.config.encoding       , errors=UpperCAmelCase_						).encode('''utf-8'''						)
																																							try:
																																														while True:
																																																					try:
																																																												lowerCAmelCase   =		paj.read_json(
																																																												    io.BytesIO(UpperCAmelCase_						)       , read_options=paj.ReadOptions(block_size=UpperCAmelCase_						)						)
																																																												break
																																																					except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as e:
																																																												if (
																																																												    isinstance(UpperCAmelCase_       , pa.ArrowInvalid						)
																																																												    and "straddling" not in str(UpperCAmelCase_						)
																																																												    or block_size > len(UpperCAmelCase_						)
																																																												):
																																																																			raise
																																																												else:
																																																																			# Increase the block size in case it was too small.
																																																																			# The block size will be reset for the next file.
																																																																			logger.debug(
																																																																			    F"""Batch of {len(UpperCAmelCase_						)} bytes couldn't be parsed with block_size={block_size}. Retrying with block_size={block_size * 2}."""						)
																																																																			block_size *= 2
																																							except pa.ArrowInvalid as e:
																																														try:
																																																					with open(
																																																					    UpperCAmelCase_       , encoding=self.config.encoding       , errors=self.config.encoding_errors						) as f:
																																																												lowerCAmelCase   =		json.load(UpperCAmelCase_						)
																																														except json.JSONDecodeError:
																																																					logger.error(F"""Failed to read file '{file}' with error {type(UpperCAmelCase_						)}: {e}"""						)
																																																					raise e
																																														# If possible, parse the file as a list of json objects and exit the loop
																																														if isinstance(UpperCAmelCase_       , UpperCAmelCase_						):  # list is the only sequence type supported in JSON
																																																					try:
																																																												lowerCAmelCase   =		set().union(*[row.keys() for row in dataset]						)
																																																												lowerCAmelCase   =		{col: [row.get(UpperCAmelCase_						) for row in dataset] for col in keys}
																																																												lowerCAmelCase   =		pa.Table.from_pydict(UpperCAmelCase_						)
																																																					except (pa.ArrowInvalid, AttributeError) as e:
																																																												logger.error(F"""Failed to read file '{file}' with error {type(UpperCAmelCase_						)}: {e}"""						)
																																																												raise ValueError(F"""Not able to read records in the JSON file at {file}."""						) from None
																																																					yield file_idx, self._cast_table(UpperCAmelCase_						)
																																																					break
																																														else:
																																																					logger.error(F"""Failed to read file '{file}' with error {type(UpperCAmelCase_						)}: {e}"""						)
																																																					raise ValueError(
																																																					    F"""Not able to read records in the JSON file at {file}. """
																																																					    F"""You should probably indicate the field of the JSON file containing your records. """
																																																					    F"""This JSON file contain the following fields: {str(list(dataset.keys()						)						)}. """
																																																					    F"""Select the correct one and provide it as `field='XXX'` to the dataset loading method. """						) from None
                        # Uncomment for debugging (will print the Arrow table size and elements)
                        # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                        # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
																																							yield (file_idx, batch_idx), self._cast_table(UpperCAmelCase_						)
																																							batch_idx += 1



