



import argparse
import collections

import torch
from flax import traverse_util
from tax import checkpoints

from transformers import TaConfig, TaEncoderModel, TaForConditionalGeneration
from transformers.utils import logging


logging.set_verbosity_info()



def       lowercase       (   SCREAMING_SNAKE_CASE__       :	str				,       SCREAMING_SNAKE_CASE__       :	Optional[int]				,       SCREAMING_SNAKE_CASE__       :	Optional[Any]				,       SCREAMING_SNAKE_CASE__       :	int="attention"				)  ->     List[Any]:
 _snake_case    :  Dict					=					params[F'''{prefix}/layers_{i}/{layer_name}/key/kernel''']
 _snake_case    :  Optional[int]					=					params[F'''{prefix}/layers_{i}/{layer_name}/out/kernel''']
 _snake_case    :  List[str]					=					params[F'''{prefix}/layers_{i}/{layer_name}/query/kernel''']
 _snake_case    :  List[Any]					=					params[F'''{prefix}/layers_{i}/{layer_name}/value/kernel''']
 return k, o, q, v



def       lowercase       (   SCREAMING_SNAKE_CASE__       :	Union[str, Any]				,       SCREAMING_SNAKE_CASE__       :	List[str]				,       SCREAMING_SNAKE_CASE__       :	int				,       SCREAMING_SNAKE_CASE__       :	str=False				)  ->     str:
 if split_mlp_wi:
  _snake_case    :  List[Any]					=					params[F'''{prefix}/layers_{i}/mlp/wi_0/kernel''']
  _snake_case    :  int					=					params[F'''{prefix}/layers_{i}/mlp/wi_1/kernel''']
  _snake_case    :  Dict					=					(wi_a, wi_a)
 else:
  _snake_case    :  int					=					params[F'''{prefix}/layers_{i}/mlp/wi/kernel''']

 _snake_case    :  Tuple					=					params[F'''{prefix}/layers_{i}/mlp/wo/kernel''']
 return wi, wo



def       lowercase       (   SCREAMING_SNAKE_CASE__       :	Union[str, Any]				,       SCREAMING_SNAKE_CASE__       :	Dict				,       SCREAMING_SNAKE_CASE__       :	Optional[int]				,       SCREAMING_SNAKE_CASE__       :	Optional[Any]				)  ->     Dict:
 return params[F'''{prefix}/layers_{i}/{layer_name}/scale''']



def       lowercase       (   SCREAMING_SNAKE_CASE__       :	dict				,       *, SCREAMING_SNAKE_CASE__       :	int				,       SCREAMING_SNAKE_CASE__       :	bool				)  ->     int:
 _snake_case    :  Any					=					traverse_util.flatten_dict(variables["""target"""]				)
 _snake_case    :  str					=					{"""/""".join(SCREAMING_SNAKE_CASE__				): v for k, v in old.items()}

 # v1.1 models have a gated GeLU with wi_0 and wi_1 instead of wi
 _snake_case    :  Optional[int]					=					"""encoder/layers_0/mlp/wi_0/kernel""" in old
 print("""Split MLP:"""				,       SCREAMING_SNAKE_CASE__				)

 _snake_case    :  Any					=					collections.OrderedDict()

 # Shared embeddings.
 _snake_case    :  Union[str, Any]					=					old["""token_embedder/embedding"""]

 # Encoder.
 for i in range(SCREAMING_SNAKE_CASE__				):
  # Block i, layer 0 (Self Attention).
  _snake_case    :  Tuple					=					tax_layer_norm_lookup(SCREAMING_SNAKE_CASE__				,       SCREAMING_SNAKE_CASE__				,       """encoder"""				,       """pre_attention_layer_norm"""				)
  _snake_case    :  int					=					tax_attention_lookup(SCREAMING_SNAKE_CASE__				,       SCREAMING_SNAKE_CASE__				,       """encoder"""				,       """attention"""				)
  _snake_case    :  Optional[int]					=					layer_norm
  _snake_case    :  int					=					k.T
  _snake_case    :  str					=					o.T
  _snake_case    :  List[str]					=					q.T
  _snake_case    :  Union[str, Any]					=					v.T

  # Block i, layer 1 (MLP).
  _snake_case    :  List[Any]					=					tax_layer_norm_lookup(SCREAMING_SNAKE_CASE__				,       SCREAMING_SNAKE_CASE__				,       """encoder"""				,       """pre_mlp_layer_norm"""				)
  _snake_case    :  Any					=					tax_mlp_lookup(SCREAMING_SNAKE_CASE__				,       SCREAMING_SNAKE_CASE__				,       """encoder"""				,       SCREAMING_SNAKE_CASE__				)
  _snake_case    :  Optional[int]					=					layer_norm
  if split_mlp_wi:
   _snake_case    :  Optional[int]					=					wi[0].T
   _snake_case    :  Any					=					wi[1].T
  else:
   _snake_case    :  List[str]					=					wi.T
  _snake_case    :  Any					=					wo.T

 _snake_case    :  Optional[Any]					=					old[
     """encoder/relpos_bias/rel_embedding"""
 ].T
 _snake_case    :  Optional[Any]					=					old["""encoder/encoder_norm/scale"""]

 if not is_encoder_only:
  # Decoder.
  for i in range(SCREAMING_SNAKE_CASE__				):
   # Block i, layer 0 (Self Attention).
   _snake_case    :  Tuple					=					tax_layer_norm_lookup(SCREAMING_SNAKE_CASE__				,       SCREAMING_SNAKE_CASE__				,       """decoder"""				,       """pre_self_attention_layer_norm"""				)
   _snake_case    :  Dict					=					tax_attention_lookup(SCREAMING_SNAKE_CASE__				,       SCREAMING_SNAKE_CASE__				,       """decoder"""				,       """self_attention"""				)
   _snake_case    :  List[Any]					=					layer_norm
   _snake_case    :  List[str]					=					k.T
   _snake_case    :  Optional[int]					=					o.T
   _snake_case    :  Dict					=					q.T
   _snake_case    :  Union[str, Any]					=					v.T

   # Block i, layer 1 (Cross Attention).
   _snake_case    :  Dict					=					tax_layer_norm_lookup(SCREAMING_SNAKE_CASE__				,       SCREAMING_SNAKE_CASE__				,       """decoder"""				,       """pre_cross_attention_layer_norm"""				)
   _snake_case    :  Optional[int]					=					tax_attention_lookup(SCREAMING_SNAKE_CASE__				,       SCREAMING_SNAKE_CASE__				,       """decoder"""				,       """encoder_decoder_attention"""				)
   _snake_case    :  Optional[int]					=					layer_norm
   _snake_case    :  int					=					k.T
   _snake_case    :  Any					=					o.T
   _snake_case    :  int					=					q.T
   _snake_case    :  int					=					v.T

   # Block i, layer 2 (MLP).
   _snake_case    :  Optional[int]					=					tax_layer_norm_lookup(SCREAMING_SNAKE_CASE__				,       SCREAMING_SNAKE_CASE__				,       """decoder"""				,       """pre_mlp_layer_norm"""				)
   _snake_case    :  Tuple					=					tax_mlp_lookup(SCREAMING_SNAKE_CASE__				,       SCREAMING_SNAKE_CASE__				,       """decoder"""				,       SCREAMING_SNAKE_CASE__				)
   _snake_case    :  Dict					=					layer_norm
   if split_mlp_wi:
    _snake_case    :  str					=					wi[0].T
    _snake_case    :  List[Any]					=					wi[1].T
   else:
    _snake_case    :  int					=					wi.T
   _snake_case    :  int					=					wo.T

  _snake_case    :  List[str]					=					old["""decoder/decoder_norm/scale"""]
  _snake_case    :  int					=					old[
      """decoder/relpos_bias/rel_embedding"""
  ].T

  # LM Head (only in v1.1 checkpoints, in v1.0 embeddings are used instead)
  if "decoder/logits_dense/kernel" in old:
   _snake_case    :  Dict					=					old["""decoder/logits_dense/kernel"""].T

 return new



def       lowercase       (   SCREAMING_SNAKE_CASE__       :	List[Any]				,       SCREAMING_SNAKE_CASE__       :	bool				)  ->     Optional[int]:
 _snake_case    :  Optional[Any]					=					collections.OrderedDict([(k, torch.from_numpy(v.copy()				)) for (k, v) in converted_params.items()]				)

 # Add what is missing.
 if "encoder.embed_tokens.weight" not in state_dict:
  _snake_case    :  Dict					=					state_dict["""shared.weight"""]

 if not is_encoder_only:
  if "decoder.embed_tokens.weight" not in state_dict:
   _snake_case    :  Tuple					=					state_dict["""shared.weight"""]

  if "lm_head.weight" not in state_dict:  # For old 1.0 models.
   print("""Using shared word embeddings as lm_head."""				)
   _snake_case    :  str					=					state_dict["""shared.weight"""]

 return state_dict



def       lowercase       (   SCREAMING_SNAKE_CASE__       :	str				,       SCREAMING_SNAKE_CASE__       :	int				,       SCREAMING_SNAKE_CASE__       :	Optional[int]				,       SCREAMING_SNAKE_CASE__       :	Union[str, Any]				)  ->     Optional[int]:
 _snake_case    :  Optional[int]					=					checkpoints.load_tax_checkpoint(SCREAMING_SNAKE_CASE__				)
 _snake_case    :  str					=					convert_tax_to_pytorch(SCREAMING_SNAKE_CASE__				,       num_layers=config.num_layers				,       is_encoder_only=SCREAMING_SNAKE_CASE__				)
 _snake_case    :  Dict					=					make_state_dict(SCREAMING_SNAKE_CASE__				,       SCREAMING_SNAKE_CASE__				)
 model.load_state_dict(SCREAMING_SNAKE_CASE__				,       strict=SCREAMING_SNAKE_CASE__				)



def       lowercase       (   SCREAMING_SNAKE_CASE__       :	Optional[Any]				,       SCREAMING_SNAKE_CASE__       :	int				,       SCREAMING_SNAKE_CASE__       :	int				,       SCREAMING_SNAKE_CASE__       :	bool = False				)  ->     Optional[int]:
 _snake_case    :  List[Any]					=					TaConfig.from_json_file(SCREAMING_SNAKE_CASE__				)
 print(F'''Building PyTorch model from configuration: {config}'''				)
 # Non-v1.1 checkpoints could also use T5Model, but this works for all.
 # The v1.0 checkpoints will simply have an LM head that is the word embeddings.
 if is_encoder_only:
  _snake_case    :  Dict					=					TaEncoderModel(SCREAMING_SNAKE_CASE__				)
 else:
  _snake_case    :  Optional[int]					=					TaForConditionalGeneration(SCREAMING_SNAKE_CASE__				)

 # Load weights from tf checkpoint
 load_tax_weights_in_ta(SCREAMING_SNAKE_CASE__				,       SCREAMING_SNAKE_CASE__				,       SCREAMING_SNAKE_CASE__				,       SCREAMING_SNAKE_CASE__				)

 # Save pytorch-model
 print(F'''Save PyTorch model to {pytorch_dump_path}'''				)
 model.save_pretrained(SCREAMING_SNAKE_CASE__				)

 # Verify that we can load the checkpoint.
 model.from_pretrained(SCREAMING_SNAKE_CASE__				)
 print("""Done"""				)


if __name__ == "__main__":
  a__    =	argparse.ArgumentParser(description="""Converts a native T5X checkpoint into a PyTorch checkpoint.""")
  # Required parameters
  parser.add_argument(
      """--t5x_checkpoint_path""", default=None, type=str, required=True, help="""Path to the T5X checkpoint."""
  )
  parser.add_argument(
      """--config_file""",
      default=None,
      type=str,
      required=True,
      help="""The config json file corresponding to the pre-trained T5 model.\nThis specifies the model architecture.""",
  )
  parser.add_argument(
      """--pytorch_dump_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
  )
  parser.add_argument(
      """--is_encoder_only""", action="""store_true""", help="""Check if the model is encoder-decoder model""", default=False
  )
  a__    =	parser.parse_args()
  convert_tax_checkpoint_to_pytorch(
      args.tax_checkpoint_path, args.config_file, args.pytorch_dump_path, args.is_encoder_only
  )





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



   '''simple docstring'''






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

    """simple docstring"""







    _snake_case    :  Dict					=					parent
    _snake_case    :  int					=					batch_size
    _snake_case    :  str					=					seq_length
    _snake_case    :  List[str]					=					is_training
    _snake_case    :  Tuple					=					use_input_mask
    _snake_case    :  Optional[Any]					=					use_token_type_ids
    _snake_case    :  Any					=					use_labels
    _snake_case    :  str					=					vocab_size
    _snake_case    :  List[str]					=					hidden_size
    _snake_case    :  List[str]					=					num_hidden_layers
    _snake_case    :  Dict					=					num_attention_heads
    _snake_case    :  int					=					intermediate_size
    _snake_case    :  Any					=					hidden_act
    _snake_case    :  List[Any]					=					hidden_dropout_prob
    _snake_case    :  Tuple					=					attention_probs_dropout_prob
    _snake_case    :  List[Any]					=					max_position_embeddings
    _snake_case    :  int					=					type_vocab_size
    _snake_case    :  Dict					=					type_sequence_label_size
    _snake_case    :  List[str]					=					initializer_range
    _snake_case    :  Dict					=					num_labels
    _snake_case    :  Optional[Any]					=					num_choices
    _snake_case    :  Dict					=					scope
    _snake_case    :  List[Any]					=					projection_dim

   def      UpperCamelCase_  (					self							:			Dict)				->				List[str]:

    """simple docstring"""







    _snake_case    :  Optional[int]					=					ids_tensor([self.batch_size, self.seq_length]				,		self.vocab_size)

    _snake_case    :  List[str]					=					None
    if self.use_input_mask:
     # follow test_modeling_tf_ctrl.py
     _snake_case    :  List[Any]					=					random_attention_mask([self.batch_size, self.seq_length])

    _snake_case    :  str					=					None
    if self.use_token_type_ids:
     _snake_case    :  Optional[Any]					=					ids_tensor([self.batch_size, self.seq_length]				,		self.type_vocab_size)

    _snake_case    :  Tuple					=					None
    _snake_case    :  Any					=					None
    _snake_case    :  str					=					None
    if self.use_labels:
     _snake_case    :  Optional[Any]					=					ids_tensor([self.batch_size]				,		self.type_sequence_label_size)
     _snake_case    :  Union[str, Any]					=					ids_tensor([self.batch_size, self.seq_length]				,		self.num_labels)
     _snake_case    :  Optional[Any]					=					ids_tensor([self.batch_size]				,		self.num_choices)

    _snake_case    :  Union[str, Any]					=					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=lowerCAmelCase				,		initializer_range=self.initializer_range				,		)
    _snake_case    :  Optional[int]					=					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]				,		lowerCAmelCase							:			Optional[Any]				,		lowerCAmelCase							:			Optional[Any]				,		lowerCAmelCase							:			List[str]				,		lowerCAmelCase							:			Optional[Any]				,		lowerCAmelCase							:			Tuple				,		lowerCAmelCase							:			Tuple				,		lowerCAmelCase							:			List[str])				->				List[Any]:

    """simple docstring"""







    _snake_case    :  Dict					=					TFDPRContextEncoder(config=lowerCAmelCase)
    _snake_case    :  Union[str, Any]					=					model(lowerCAmelCase				,		attention_mask=lowerCAmelCase				,		token_type_ids=lowerCAmelCase)
    _snake_case    :  Optional[int]					=					model(lowerCAmelCase				,		token_type_ids=lowerCAmelCase)
    _snake_case    :  List[str]					=					model(lowerCAmelCase)
    self.parent.assertEqual(result.pooler_output.shape				,		(self.batch_size, self.projection_dim or self.hidden_size))

   def      UpperCamelCase_  (					self							:			List[Any]				,		lowerCAmelCase							:			int				,		lowerCAmelCase							:			List[str]				,		lowerCAmelCase							:			Optional[Any]				,		lowerCAmelCase							:			Any				,		lowerCAmelCase							:			Any				,		lowerCAmelCase							:			Union[str, Any]				,		lowerCAmelCase							:			Optional[int])				->				int:

    """simple docstring"""







    _snake_case    :  Optional[int]					=					TFDPRQuestionEncoder(config=lowerCAmelCase)
    _snake_case    :  List[Any]					=					model(lowerCAmelCase				,		attention_mask=lowerCAmelCase				,		token_type_ids=lowerCAmelCase)
    _snake_case    :  List[Any]					=					model(lowerCAmelCase				,		token_type_ids=lowerCAmelCase)
    _snake_case    :  Tuple					=					model(lowerCAmelCase)
    self.parent.assertEqual(result.pooler_output.shape				,		(self.batch_size, self.projection_dim or self.hidden_size))

   def      UpperCamelCase_  (					self							:			List[str]				,		lowerCAmelCase							:			List[str]				,		lowerCAmelCase							:			Union[str, Any]				,		lowerCAmelCase							:			Dict				,		lowerCAmelCase							:			Optional[int]				,		lowerCAmelCase							:			str				,		lowerCAmelCase							:			Tuple				,		lowerCAmelCase							:			str)				->				Tuple:

    """simple docstring"""







    _snake_case    :  Optional[Any]					=					TFDPRReader(config=lowerCAmelCase)
    _snake_case    :  Any					=					model(lowerCAmelCase				,		attention_mask=lowerCAmelCase)

    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							:			str)				->				Any:

    """simple docstring"""







    _snake_case    :  List[Any]					=					self.prepare_config_and_inputs()
    (
        (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    
    )    :  Union[str, Any]					=					config_and_inputs
    _snake_case    :  Union[str, Any]					=					{"""input_ids""": input_ids}
    return config, inputs_dict







@require_tf
class 		snake_case						(	SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_					,unittest.TestCase  ):



   '''simple docstring'''






   snake_case_					: Any          =		(
       (
           TFDPRContextEncoder,
           TFDPRQuestionEncoder,
           TFDPRReader,
       )
       if is_tf_available()
       else ()
   )
   snake_case_					: List[str]          =		{"""feature-extraction""": TFDPRQuestionEncoder} if is_tf_available() else {}

   snake_case_					: Tuple          =		False
   snake_case_					: str          =		False
   snake_case_					: Tuple          =		False
   snake_case_					: int          =		False
   snake_case_					: List[Any]          =		False

   def      UpperCamelCase_  (					self							:			Tuple)				->				str:

    """simple docstring"""







    _snake_case    :  List[str]					=					TFDPRModelTester(self)
    _snake_case    :  Dict					=					ConfigTester(self				,		config_class=lowerCAmelCase				,		hidden_size=37)

   def      UpperCamelCase_  (					self							:			List[Any])				->				Optional[Any]:

    """simple docstring"""







    self.config_tester.run_common_tests()

   def      UpperCamelCase_  (					self							:			List[str])				->				Optional[int]:

    """simple docstring"""







    _snake_case    :  Any					=					self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_dpr_context_encoder(*lowerCAmelCase)

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

    """simple docstring"""







    _snake_case    :  str					=					self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_dpr_question_encoder(*lowerCAmelCase)

   def      UpperCamelCase_  (					self							:			List[str])				->				Tuple:

    """simple docstring"""







    _snake_case    :  str					=					self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_dpr_reader(*lowerCAmelCase)


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

    """simple docstring"""







    for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _snake_case    :  List[Any]					=					TFDPRContextEncoder.from_pretrained(lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)

    for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _snake_case    :  List[str]					=					TFDPRContextEncoder.from_pretrained(lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)

    for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _snake_case    :  Any					=					TFDPRQuestionEncoder.from_pretrained(lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)

    for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _snake_case    :  Union[str, Any]					=					TFDPRReader.from_pretrained(lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)







@require_tf
class 		snake_case						(	unittest.TestCase  ):



   '''simple docstring'''






   @slow
   def      UpperCamelCase_  (					self							:			List[str])				->				Optional[int]:

    """simple docstring"""







    _snake_case    :  List[str]					=					TFDPRQuestionEncoder.from_pretrained("""facebook/dpr-question_encoder-single-nq-base""")

    _snake_case    :  str					=					tf.constant(
        [[101, 7592, 1010, 2003, 2026, 3899, 1_0140, 1029, 102]])  # [CLS] hello, is my dog cute? [SEP]
    _snake_case    :  Optional[int]					=					model(lowerCAmelCase)[0]  # embedding shape = (1, 768)
    # compare the actual values for a slice.
    _snake_case    :  Union[str, Any]					=					tf.constant(
        [
            [
                0.03_236_253,
                0.12_753_335,
                0.16_818_509,
                0.00_279_786,
                0.3_896_933,
                0.24_264_945,
                0.2_178_971,
                -0.02_335_227,
                -0.08_481_959,
                -0.14_324_117,
            ]
        ])
    self.assertTrue(numpy.allclose(output[:, :10].numpy()				,		expected_slice.numpy()				,		atol=1E-4))
