






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		:







			def __init__(     self:  Optional[Any]				,     UpperCAmelCase_:  Any				,     UpperCAmelCase_:  List[Any]=13				,     UpperCAmelCase_:  Optional[int]=7				,     UpperCAmelCase_:  Optional[int]=6				,     UpperCAmelCase_:  Optional[Any]=17				,     UpperCAmelCase_:  Dict=23				,     UpperCAmelCase_:  Dict=11				,     UpperCAmelCase_:  Optional[int]=True				,     ):







									'''simple docstring'''




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







			def 							UpperCamelCase				(     self:  Any					):







									'''simple docstring'''




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

									_SCREAMING_SNAKE_CASE     				=  self.get_config()

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







			def 							UpperCamelCase				(     self:  List[str]					):







									'''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 							UpperCamelCase				(     self:  int				,     UpperCAmelCase_:  Dict				,     UpperCAmelCase_:  int				,     UpperCAmelCase_:  List[Any]				,     UpperCAmelCase_:  str				,     UpperCAmelCase_:  str				,     UpperCAmelCase_:  int				,     UpperCAmelCase_:  Dict				,     ):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  DecisionTransformerModel(config=__A					)
									model.to(__A					)
									model.eval()
									_SCREAMING_SNAKE_CASE     				=  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 							UpperCamelCase				(     self:  List[str]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.prepare_config_and_inputs()
									(
									    _SCREAMING_SNAKE_CASE	
									)     				=  config_and_inputs
									_SCREAMING_SNAKE_CASE     				=  {
									    """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		(A__	,A__	,A__	,unittest.TestCase   ):
			__snake_case    :				Any									=							(DecisionTransformerModel,) if is_torch_available() else ()
			__snake_case    :				Tuple									=							()
			__snake_case    :				Optional[int]									=							{"feature-extraction": DecisionTransformerModel} if is_torch_available() else {}

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

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







			def 							UpperCamelCase				(     self:  List[str]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  DecisionTransformerModelTester(self					)
									_SCREAMING_SNAKE_CASE     				=  ConfigTester(self				,     config_class=__A				,     hidden_size=37					)







			def 							UpperCamelCase				(     self:  List[Any]					):







									'''simple docstring'''




									self.config_tester.run_common_tests()







			def 							UpperCamelCase				(     self:  List[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*__A					)







			@slow
			def 							UpperCamelCase				(     self:  Tuple					):







									'''simple docstring'''




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







			def 							UpperCamelCase				(     self:  Tuple					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs_for_common()

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

															_SCREAMING_SNAKE_CASE     				=  [
															    """states""",
															    """actions""",
															    """rewards""",
															    """returns_to_go""",
															    """timesteps""",
															    """attention_mask""",
															]

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





@require_torch
class 				__UpperCAmelCase		(unittest.TestCase   ):







			@slow
			def 							UpperCamelCase				(     self:  Any					):







									'''simple docstring'''




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

									_SCREAMING_SNAKE_CASE     				=  torch.tensor(
									    [[0.24_27_93, -0.28_69_30_74, 0.8_74_26_13], [0.67_81_52_74, -0.08_10_10_85, -0.12_95_21_47]]				,     device=__A					)

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

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

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

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

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









from __future__ import annotations

import unittest

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

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
		import numpy
		import tensorflow as tf

		from transformers import (
		    TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    BertConfig,
		    DPRConfig,
		    TFDPRContextEncoder,
		    TFDPRQuestionEncoder,
		    TFDPRReader,
		)



class 				__UpperCAmelCase		:







			def __init__(     self:  Optional[int]				,     UpperCAmelCase_:  List[str]				,     UpperCAmelCase_:  int=13				,     UpperCAmelCase_:  List[str]=7				,     UpperCAmelCase_:  Union[str, Any]=True				,     UpperCAmelCase_:  List[str]=True				,     UpperCAmelCase_:  Dict=True				,     UpperCAmelCase_:  str=True				,     UpperCAmelCase_:  List[str]=99				,     UpperCAmelCase_:  int=32				,     UpperCAmelCase_:  List[Any]=2				,     UpperCAmelCase_:  Union[str, Any]=4				,     UpperCAmelCase_:  Dict=37				,     UpperCAmelCase_:  Tuple="gelu"				,     UpperCAmelCase_:  Any=0.1				,     UpperCAmelCase_:  List[str]=0.1				,     UpperCAmelCase_:  List[Any]=512				,     UpperCAmelCase_:  Optional[int]=16				,     UpperCAmelCase_:  Optional[Any]=2				,     UpperCAmelCase_:  Optional[int]=0.02				,     UpperCAmelCase_:  Union[str, Any]=3				,     UpperCAmelCase_:  Optional[int]=4				,     UpperCAmelCase_:  Tuple=None				,     UpperCAmelCase_:  Any=0				,     ):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  parent
									_SCREAMING_SNAKE_CASE     				=  batch_size
									_SCREAMING_SNAKE_CASE     				=  seq_length
									_SCREAMING_SNAKE_CASE     				=  is_training
									_SCREAMING_SNAKE_CASE     				=  use_input_mask
									_SCREAMING_SNAKE_CASE     				=  use_token_type_ids
									_SCREAMING_SNAKE_CASE     				=  use_labels
									_SCREAMING_SNAKE_CASE     				=  vocab_size
									_SCREAMING_SNAKE_CASE     				=  hidden_size
									_SCREAMING_SNAKE_CASE     				=  num_hidden_layers
									_SCREAMING_SNAKE_CASE     				=  num_attention_heads
									_SCREAMING_SNAKE_CASE     				=  intermediate_size
									_SCREAMING_SNAKE_CASE     				=  hidden_act
									_SCREAMING_SNAKE_CASE     				=  hidden_dropout_prob
									_SCREAMING_SNAKE_CASE     				=  attention_probs_dropout_prob
									_SCREAMING_SNAKE_CASE     				=  max_position_embeddings
									_SCREAMING_SNAKE_CASE     				=  type_vocab_size
									_SCREAMING_SNAKE_CASE     				=  type_sequence_label_size
									_SCREAMING_SNAKE_CASE     				=  initializer_range
									_SCREAMING_SNAKE_CASE     				=  num_labels
									_SCREAMING_SNAKE_CASE     				=  num_choices
									_SCREAMING_SNAKE_CASE     				=  scope
									_SCREAMING_SNAKE_CASE     				=  projection_dim







			def 							UpperCamelCase				(     self:  Any					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  ids_tensor([self.batch_size, self.seq_length]				,     self.vocab_size					)

									_SCREAMING_SNAKE_CASE     				=  None
									if self.use_input_mask:
															# follow test_modeling_tf_ctrl.py
															_SCREAMING_SNAKE_CASE     				=  random_attention_mask([self.batch_size, self.seq_length]					)

									_SCREAMING_SNAKE_CASE     				=  None
									if self.use_token_type_ids:
															_SCREAMING_SNAKE_CASE     				=  ids_tensor([self.batch_size, self.seq_length]				,     self.type_vocab_size					)

									_SCREAMING_SNAKE_CASE     				=  None
									_SCREAMING_SNAKE_CASE     				=  None
									_SCREAMING_SNAKE_CASE     				=  None
									if self.use_labels:
															_SCREAMING_SNAKE_CASE     				=  ids_tensor([self.batch_size]				,     self.type_sequence_label_size					)
															_SCREAMING_SNAKE_CASE     				=  ids_tensor([self.batch_size, self.seq_length]				,     self.num_labels					)
															_SCREAMING_SNAKE_CASE     				=  ids_tensor([self.batch_size]				,     self.num_choices					)

									_SCREAMING_SNAKE_CASE     				=  BertConfig(
									    vocab_size=self.vocab_size				,     hidden_size=self.hidden_size				,     num_hidden_layers=self.num_hidden_layers				,     num_attention_heads=self.num_attention_heads				,     intermediate_size=self.intermediate_size				,     hidden_act=self.hidden_act				,     hidden_dropout_prob=self.hidden_dropout_prob				,     attention_probs_dropout_prob=self.attention_probs_dropout_prob				,     max_position_embeddings=self.max_position_embeddings				,     type_vocab_size=self.type_vocab_size				,     is_decoder=UpperCAmelCase_				,     initializer_range=self.initializer_range				,     )
									_SCREAMING_SNAKE_CASE     				=  DPRConfig(projection_dim=self.projection_dim				,     **config.to_dict()					)

									return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels







			def 							UpperCamelCase				(     self:  Optional[int]				,     UpperCAmelCase_:  Optional[Any]				,     UpperCAmelCase_:  Any				,     UpperCAmelCase_:  int				,     UpperCAmelCase_:  List[str]				,     UpperCAmelCase_:  Dict				,     UpperCAmelCase_:  Dict				,     UpperCAmelCase_:  Any					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFDPRContextEncoder(config=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_				,     attention_mask=UpperCAmelCase_				,     token_type_ids=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_				,     token_type_ids=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)
									self.parent.assertEqual(result.pooler_output.shape				,     (self.batch_size, self.projection_dim or self.hidden_size)					)







			def 							UpperCamelCase				(     self:  List[Any]				,     UpperCAmelCase_:  Union[str, Any]				,     UpperCAmelCase_:  Any				,     UpperCAmelCase_:  Dict				,     UpperCAmelCase_:  Union[str, Any]				,     UpperCAmelCase_:  Union[str, Any]				,     UpperCAmelCase_:  List[Any]				,     UpperCAmelCase_:  Any					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFDPRQuestionEncoder(config=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_				,     attention_mask=UpperCAmelCase_				,     token_type_ids=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_				,     token_type_ids=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)
									self.parent.assertEqual(result.pooler_output.shape				,     (self.batch_size, self.projection_dim or self.hidden_size)					)







			def 							UpperCamelCase				(     self:  Optional[Any]				,     UpperCAmelCase_:  Dict				,     UpperCAmelCase_:  Optional[int]				,     UpperCAmelCase_:  List[Any]				,     UpperCAmelCase_:  str				,     UpperCAmelCase_:  Optional[Any]				,     UpperCAmelCase_:  Union[str, Any]				,     UpperCAmelCase_:  str					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFDPRReader(config=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_				,     attention_mask=UpperCAmelCase_					)

									self.parent.assertEqual(result.start_logits.shape				,     (self.batch_size, self.seq_length)					)
									self.parent.assertEqual(result.end_logits.shape				,     (self.batch_size, self.seq_length)					)
									self.parent.assertEqual(result.relevance_logits.shape				,     (self.batch_size,)					)







			def 							UpperCamelCase				(     self:  Optional[int]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.prepare_config_and_inputs()
									(
									    (
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					
									)     				=  config_and_inputs
									_SCREAMING_SNAKE_CASE     				=  {"""input_ids""": input_ids}
									return config, inputs_dict



@require_tf
class 				__UpperCAmelCase		(_UpperCAmelCase	,_UpperCAmelCase	,unittest.TestCase   ):
			__snake_case    :				Dict									=							(
			    (
			        TFDPRContextEncoder,
			        TFDPRQuestionEncoder,
			        TFDPRReader,
			    )
			    if is_tf_available()
			    else ()
			)
			__snake_case    :				Optional[Any]									=							{"feature-extraction": TFDPRQuestionEncoder} if is_tf_available() else {}

			__snake_case    :				str									=							False
			__snake_case    :				List[Any]									=							False
			__snake_case    :				Any									=							False
			__snake_case    :				List[Any]									=							False
			__snake_case    :				Tuple									=							False







			def 							UpperCamelCase				(     self:  int					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFDPRModelTester(self					)
									_SCREAMING_SNAKE_CASE     				=  ConfigTester(self				,     config_class=UpperCAmelCase_				,     hidden_size=37					)







			def 							UpperCamelCase				(     self:  Any					):







									'''simple docstring'''




									self.config_tester.run_common_tests()







			def 							UpperCamelCase				(     self:  List[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_dpr_context_encoder(*UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  Optional[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_dpr_question_encoder(*UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  int					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_dpr_reader(*UpperCAmelCase_					)







			@slow
			def 							UpperCamelCase				(     self:  List[Any]					):







									'''simple docstring'''




									for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															_SCREAMING_SNAKE_CASE     				=  TFDPRContextEncoder.from_pretrained(UpperCAmelCase_					)
															self.assertIsNotNone(UpperCAmelCase_					)

									for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															_SCREAMING_SNAKE_CASE     				=  TFDPRContextEncoder.from_pretrained(UpperCAmelCase_					)
															self.assertIsNotNone(UpperCAmelCase_					)

									for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															_SCREAMING_SNAKE_CASE     				=  TFDPRQuestionEncoder.from_pretrained(UpperCAmelCase_					)
															self.assertIsNotNone(UpperCAmelCase_					)

									for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															_SCREAMING_SNAKE_CASE     				=  TFDPRReader.from_pretrained(UpperCAmelCase_					)
															self.assertIsNotNone(UpperCAmelCase_					)



@require_tf
class 				__UpperCAmelCase		(unittest.TestCase   ):







			@slow
			def 							UpperCamelCase				(     self:  Union[str, Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFDPRQuestionEncoder.from_pretrained("""facebook/dpr-question_encoder-single-nq-base"""					)

									_SCREAMING_SNAKE_CASE     				=  tf.constant(
									    [[101, 7_592, 1_010, 2_003, 2_026, 3_899, 10_140, 1_029, 102]]					)  # [CLS] hello, is my dog cute? [SEP]
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)[0]  # embedding shape = (1, 768)
									# compare the actual values for a slice.
									_SCREAMING_SNAKE_CASE     				=  tf.constant(
									    [
									        [
									            0.03_23_62_53,
									            0.12_75_33_35,
									            0.16_81_85_09,
									            0.00_27_97_86,
									            0.3_89_69_33,
									            0.24_26_49_45,
									            0.2_17_89_71,
									            -0.02_33_52_27,
									            -0.08_48_19_59,
									            -0.14_32_41_17,
									        ]
									    ]					)
									self.assertTrue(numpy.allclose(output[:, :10].numpy()				,     expected_slice.numpy()				,     atol=1E-4					)					)

