





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        SCREAMING_SNAKE_CASE     :


   def __init__(					self    :       Optional[Any]    ,  a    :       Optional[Any]    ,  a    :       str=13    ,  a    :       Any=7    ,  a    :       Union[str, Any]=6    ,  a    :       Optional[Any]=17    ,  a    :       Optional[Any]=23    ,  a    :       Optional[int]=11    ,  a    :       Any=True    ,  )->     Dict:





     """simple docstring"""







     lowercase__			    =						parent
     lowercase__			    =						batch_size
     lowercase__			    =						seq_length
     lowercase__			    =						act_dim
     lowercase__			    =						state_dim
     lowercase__			    =						hidden_size
     lowercase__			    =						max_length
     lowercase__			    =						is_training


   def  SCREAMING_SNAKE_CASE_							(					self    :       Any )->     Tuple:





     """simple docstring"""







     lowercase__			    =						floats_tensor((self.batch_size, self.seq_length, self.state_dim) )
     lowercase__			    =						floats_tensor((self.batch_size, self.seq_length, self.act_dim) )
     lowercase__			    =						floats_tensor((self.batch_size, self.seq_length, 1) )
     lowercase__			    =						floats_tensor((self.batch_size, self.seq_length, 1) )
     lowercase__			    =						ids_tensor((self.batch_size, self.seq_length)    ,  vocab_size=1_000 )
     lowercase__			    =						random_attention_mask((self.batch_size, self.seq_length) )

     lowercase__			    =						self.get_config()

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


   def  SCREAMING_SNAKE_CASE_							(					self    :       Any )->     Optional[Any]:





     """simple docstring"""







     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  SCREAMING_SNAKE_CASE_							(					self    :       List[Any]    ,  a    :       Any    ,  a    :       str    ,  a    :       Any    ,  a    :       Tuple    ,  a    :       List[str]    ,  a    :       str    ,  a    :       List[str]    ,  )->     Optional[Any]:





     """simple docstring"""







     lowercase__			    =						DecisionTransformerModel(config=a )
     model.to(a )
     model.eval()
     lowercase__			    =						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  SCREAMING_SNAKE_CASE_							(					self    :       str )->     List[str]:





     """simple docstring"""







     lowercase__			    =						self.prepare_config_and_inputs()
     (
         (
         lowercase__
     )    ,					(
         lowercase__
     )    ,					(
         lowercase__
     )    ,					(
         lowercase__
     )    ,					(
         lowercase__
     )    ,					(
         lowercase__
     )    ,					(
         lowercase__
     )    ,					
     )			    =						config_and_inputs
     lowercase__			    =						{
         '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        SCREAMING_SNAKE_CASE     (UpperCAmelCase    ,			UpperCAmelCase    ,			UpperCAmelCase    ,			unittest.TestCase   ):
   _UpperCamelCase					:	str				      = (DecisionTransformerModel,) if is_torch_available() else ()
   _UpperCamelCase					:	Optional[int]				      = ()
   _UpperCamelCase					:	int				      = {'feature-extraction': DecisionTransformerModel} if is_torch_available() else {}

   # Ignoring of a failing test from GenerationTesterMixin, as the model does not use inputs_ids
   _UpperCamelCase					:	List[Any]				      = False

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


   def  SCREAMING_SNAKE_CASE_							(					self    :       Dict )->     str:





     """simple docstring"""







     lowercase__			    =						DecisionTransformerModelTester(self )
     lowercase__			    =						ConfigTester(self    ,  config_class=a    ,  hidden_size=37 )


   def  SCREAMING_SNAKE_CASE_							(					self    :       List[str] )->     Any:





     """simple docstring"""







     self.config_tester.run_common_tests()


   def  SCREAMING_SNAKE_CASE_							(					self    :       Dict )->     Dict:





     """simple docstring"""







     lowercase__			    =						self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*a )


   @slow
   def  SCREAMING_SNAKE_CASE_							(					self    :       str )->     List[str]:





     """simple docstring"""







     for model_name in DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       lowercase__			    =						DecisionTransformerModel.from_pretrained(a )
       self.assertIsNotNone(a )


   def  SCREAMING_SNAKE_CASE_							(					self    :       Dict )->     Union[str, Any]:





     """simple docstring"""







     lowercase__    ,					lowercase__			    =						self.model_tester.prepare_config_and_inputs_for_common()

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

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

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






@require_torch
class        SCREAMING_SNAKE_CASE     (unittest.TestCase   ):


   @slow
   def  SCREAMING_SNAKE_CASE_							(					self    :       Tuple )->     List[Any]:





     """simple docstring"""







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

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

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

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

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

       with torch.no_grad():
         lowercase__    ,					lowercase__    ,					lowercase__			    =						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 ) )
       lowercase__    ,					lowercase__    ,					lowercase__    ,					lowercase__			    =						(  # env.step(action)
           torch.randn(1    ,  1    ,  config.state_dim ).to(device=a    ,  dtype=torch.floataa ),
           1.0,
           False,
           {},
       )

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









import unittest

import numpy as np

from transformers.testing_utils import is_flaky, require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
      import torch

if is_vision_available():
      from PIL import Image

      from transformers import DonutImageProcessor
class        SCREAMING_SNAKE_CASE     (unittest.TestCase   ):


   def __init__(					self    :       Any    ,  a    :       str    ,  a    :       List[Any]=7    ,  a    :       int=3    ,  a    :       int=18    ,  a    :       Optional[Any]=30    ,  a    :       Optional[int]=400    ,  a    :       int=True    ,  a    :       Tuple=None    ,  a    :       Optional[Any]=True    ,  a    :       str=False    ,  a    :       str=True    ,  a    :       int=True    ,  a    :       Tuple=[0.5, 0.5, 0.5]    ,  a    :       Any=[0.5, 0.5, 0.5]    ,  )->     Optional[int]:





     """simple docstring"""







     lowercase__			    =						parent
     lowercase__			    =						batch_size
     lowercase__			    =						num_channels
     lowercase__			    =						image_size
     lowercase__			    =						min_resolution
     lowercase__			    =						max_resolution
     lowercase__			    =						do_resize
     lowercase__			    =						size if size is not None else {'height': 18, 'width': 20}
     lowercase__			    =						do_thumbnail
     lowercase__			    =						do_align_axis
     lowercase__			    =						do_pad
     lowercase__			    =						do_normalize
     lowercase__			    =						image_mean
     lowercase__			    =						image_std


   def  SCREAMING_SNAKE_CASE_							(					self    :       Optional[Any] )->     Tuple:





     """simple docstring"""







     return {
         "do_resize": self.do_resize,
         "size": self.size,
         "do_thumbnail": self.do_thumbnail,
         "do_align_long_axis": self.do_align_axis,
         "do_pad": self.do_pad,
         "do_normalize": self.do_normalize,
         "image_mean": self.image_mean,
         "image_std": self.image_std,
     }






@require_torch
@require_vision
class        SCREAMING_SNAKE_CASE     (UpperCAmelCase    ,			unittest.TestCase   ):
   _UpperCamelCase					:	Optional[Any]				      = DonutImageProcessor if is_vision_available() else None


   def  SCREAMING_SNAKE_CASE_							(					self    :       int )->     List[Any]:





     """simple docstring"""







     lowercase__			    =						DonutImageProcessingTester(self )


   @property
   def  SCREAMING_SNAKE_CASE_							(					self    :       Union[str, Any] )->     Optional[int]:





     """simple docstring"""







     return self.image_processor_tester.prepare_image_processor_dict()


   def  SCREAMING_SNAKE_CASE_							(					self    :       Any )->     int:





     """simple docstring"""







     lowercase__			    =						self.image_processing_class(**self.image_processor_dict )
     self.assertTrue(hasattr(a    ,  'do_resize' ) )
     self.assertTrue(hasattr(a    ,  'size' ) )
     self.assertTrue(hasattr(a    ,  'do_thumbnail' ) )
     self.assertTrue(hasattr(a    ,  'do_align_long_axis' ) )
     self.assertTrue(hasattr(a    ,  'do_pad' ) )
     self.assertTrue(hasattr(a    ,  'do_normalize' ) )
     self.assertTrue(hasattr(a    ,  'image_mean' ) )
     self.assertTrue(hasattr(a    ,  'image_std' ) )


   def  SCREAMING_SNAKE_CASE_							(					self    :       Optional[Any] )->     Dict:





     """simple docstring"""







     lowercase__			    =						self.image_processing_class.from_dict(self.image_processor_dict )
     self.assertEqual(image_processor.size    ,  {'height': 18, 'width': 20} )

     lowercase__			    =						self.image_processing_class.from_dict(self.image_processor_dict    ,  size=42 )
     self.assertEqual(image_processor.size    ,  {'height': 42, 'width': 42} )

     # Previous config had dimensions in (width, height) order
     lowercase__			    =						self.image_processing_class.from_dict(self.image_processor_dict    ,  size=(42, 84) )
     self.assertEqual(image_processor.size    ,  {'height': 84, 'width': 42} )


   def  SCREAMING_SNAKE_CASE_							(					self    :       Optional[Any] )->     Dict:





     """simple docstring"""







     pass


   @is_flaky()
   def  SCREAMING_SNAKE_CASE_							(					self    :       str )->     Optional[int]:





     """simple docstring"""







     lowercase__			    =						self.image_processing_class(**self.image_processor_dict )
     # create random PIL images
     lowercase__			    =						prepare_image_inputs(self.image_processor_tester    ,  equal_resolution=a )
     for image in image_inputs:
       self.assertIsInstance(a    ,  Image.Image )

     # Test not batched input
     lowercase__			    =						image_processing(image_inputs[0]    ,  return_tensors='pt' ).pixel_values
     self.assertEqual(
         encoded_images.shape    ,  (
             1,
             self.image_processor_tester.num_channels,
             self.image_processor_tester.size['height'],
             self.image_processor_tester.size['width'],
         )    ,  )

     # Test batched
     lowercase__			    =						image_processing(a    ,  return_tensors='pt' ).pixel_values
     self.assertEqual(
         encoded_images.shape    ,  (
             self.image_processor_tester.batch_size,
             self.image_processor_tester.num_channels,
             self.image_processor_tester.size['height'],
             self.image_processor_tester.size['width'],
         )    ,  )


   @is_flaky()
   def  SCREAMING_SNAKE_CASE_							(					self    :       Optional[Any] )->     Tuple:





     """simple docstring"""







     lowercase__			    =						self.image_processing_class(**self.image_processor_dict )
     # create random numpy tensors
     lowercase__			    =						prepare_image_inputs(self.image_processor_tester    ,  equal_resolution=a    ,  numpify=a )
     for image in image_inputs:
       self.assertIsInstance(a    ,  np.ndarray )

     # Test not batched input
     lowercase__			    =						image_processing(image_inputs[0]    ,  return_tensors='pt' ).pixel_values
     self.assertEqual(
         encoded_images.shape    ,  (
             1,
             self.image_processor_tester.num_channels,
             self.image_processor_tester.size['height'],
             self.image_processor_tester.size['width'],
         )    ,  )

     # Test batched
     lowercase__			    =						image_processing(a    ,  return_tensors='pt' ).pixel_values
     self.assertEqual(
         encoded_images.shape    ,  (
             self.image_processor_tester.batch_size,
             self.image_processor_tester.num_channels,
             self.image_processor_tester.size['height'],
             self.image_processor_tester.size['width'],
         )    ,  )


   @is_flaky()
   def  SCREAMING_SNAKE_CASE_							(					self    :       List[str] )->     Dict:





     """simple docstring"""







     lowercase__			    =						self.image_processing_class(**self.image_processor_dict )
     # create random PyTorch tensors
     lowercase__			    =						prepare_image_inputs(self.image_processor_tester    ,  equal_resolution=a    ,  torchify=a )
     for image in image_inputs:
       self.assertIsInstance(a    ,  torch.Tensor )

     # Test not batched input
     lowercase__			    =						image_processing(image_inputs[0]    ,  return_tensors='pt' ).pixel_values
     self.assertEqual(
         encoded_images.shape    ,  (
             1,
             self.image_processor_tester.num_channels,
             self.image_processor_tester.size['height'],
             self.image_processor_tester.size['width'],
         )    ,  )

     # Test batched
     lowercase__			    =						image_processing(a    ,  return_tensors='pt' ).pixel_values
     self.assertEqual(
         encoded_images.shape    ,  (
             self.image_processor_tester.batch_size,
             self.image_processor_tester.num_channels,
             self.image_processor_tester.size['height'],
             self.image_processor_tester.size['width'],
         )    ,  )


