





import math
from typing import Optional

import torch
import torch.utils.checkpoint as checkpoint
from torch import nn

from ....modeling_utils import PreTrainedModel
from ....utils import add_start_docstrings, logging
from ...bert.modeling_bert import BertModel
from .configuration_retribert import RetriBertConfig


_UpperCamelCase   :							Union[str, Any]     =	logging.get_logger(__name__)

_UpperCamelCase   :							List[str]     =	[
    """yjernite/retribert-base-uncased""",
    # See all RetriBert models at https://huggingface.co/models?filter=retribert
]







class 						_lowerCAmelCase(						_a):

  """simple docstring"""







  lowerCamelCase__      					=   RetriBertConfig
  lowerCamelCase__      					=   None
  lowerCamelCase__      					=   '''retribert'''

  def 		SCREAMING_SNAKE_CASE__	( self  ,    UpperCAmelCase     )-> Optional[Any]:
     if isinstance(UpperCAmelCase  ,    nn.Linear     ):
        module.weight.data.normal_(mean=0.0  ,    std=self.config.initializer_range     )
        if module.bias is not None:
           module.bias.data.zero_()
     elif isinstance(UpperCAmelCase  ,    nn.Embedding     ):
        module.weight.data.normal_(mean=0.0  ,    std=self.config.initializer_range     )
        if module.padding_idx is not None:
           module.weight.data[module.padding_idx].zero_()
     elif isinstance(UpperCAmelCase  ,    nn.LayerNorm     ):
        module.bias.data.zero_()
        module.weight.data.fill_(1.0     )
_UpperCamelCase   :							Optional[int]     =	r"""

    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
    etc.)

    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
    and behavior.

    Parameters:
        config ([`RetriBertConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""







@add_start_docstrings(
    '''Bert Based model to embed queries or document for document retrieval.'''		,   _a		,   )
class 						_lowerCAmelCase(						_a):

  """simple docstring"""







  def __init__( self  ,    UpperCAmelCase     )-> None:
     super().__init__(UpperCAmelCase     )
     __A				 =       config.projection_dim

     __A				 =       BertModel(UpperCAmelCase     )
     __A				 =       None if config.share_encoders else BertModel(UpperCAmelCase     )
     __A				 =       nn.Dropout(config.hidden_dropout_prob     )
     __A				 =       nn.Linear(config.hidden_size  ,    config.projection_dim  ,    bias=UpperCAmelCase     )
     __A				 =       nn.Linear(config.hidden_size  ,    config.projection_dim  ,    bias=UpperCAmelCase     )

     __A				 =       nn.CrossEntropyLoss(reduction='''mean'''     )

     # Initialize weights and apply final processing
     self.post_init()

  def 		SCREAMING_SNAKE_CASE__	( self  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase=-1  ,    )-> Optional[int]:
     # reproduces BERT forward pass with checkpointing
     if checkpoint_batch_size < 0 or input_ids.shape[0] < checkpoint_batch_size:
        return sent_encoder(UpperCAmelCase  ,    attention_mask=UpperCAmelCase     )[1]
     else:
        # prepare implicit variables
        __A				 =       input_ids.device
        __A				 =       input_ids.size()
        __A				 =       torch.zeros(UpperCAmelCase  ,    dtype=torch.long  ,    device=UpperCAmelCase     )
        __A				 =       [None] * sent_encoder.config.num_hidden_layers
        __A				 =       sent_encoder.get_extended_attention_mask(
            UpperCAmelCase  ,    UpperCAmelCase     )

        # define function for checkpointing
        def partial_encode(*UpperCAmelCase     ):
           __A				 =       sent_encoder.encoder(
               inputs[0]  ,    attention_mask=inputs[1]  ,    head_mask=UpperCAmelCase  ,    )
           __A				 =       encoder_outputs[0]
           __A				 =       sent_encoder.pooler(UpperCAmelCase     )
           return pooled_output

        # run embedding layer on everything at once
        __A				 =       sent_encoder.embeddings(
            input_ids=UpperCAmelCase  ,    position_ids=UpperCAmelCase  ,    token_type_ids=UpperCAmelCase  ,    inputs_embeds=UpperCAmelCase     )
        # run encoding and pooling on one mini-batch at a time
        __A				 =       []
        for b in range(math.ceil(input_ids.shape[0] / checkpoint_batch_size     )     ):
           __A				 =       embedding_output[b * checkpoint_batch_size : (b + 1) * checkpoint_batch_size]
           __A				 =       extended_attention_mask[b * checkpoint_batch_size : (b + 1) * checkpoint_batch_size]
           __A				 =       checkpoint.checkpoint(UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase     )
           pooled_output_list.append(UpperCAmelCase     )
        return torch.cat(UpperCAmelCase  ,    dim=0     )

  def 		SCREAMING_SNAKE_CASE__	( self  ,    UpperCAmelCase  ,    UpperCAmelCase=None  ,    UpperCAmelCase=-1  ,    )-> List[str]:
     __A				 =       self.embed_sentences_checkpointed(
         UpperCAmelCase  ,    UpperCAmelCase  ,    self.bert_query  ,    UpperCAmelCase  ,    )
     return self.project_query(UpperCAmelCase     )

  def 		SCREAMING_SNAKE_CASE__	( self  ,    UpperCAmelCase  ,    UpperCAmelCase=None  ,    UpperCAmelCase=-1  ,    )-> List[Any]:
     __A				 =       self.embed_sentences_checkpointed(
         UpperCAmelCase  ,    UpperCAmelCase  ,    self.bert_query if self.bert_doc is None else self.bert_doc  ,    UpperCAmelCase  ,    )
     return self.project_doc(UpperCAmelCase     )





  def 		SCREAMING_SNAKE_CASE__	( self  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase = -1  ,    )-> torch.FloatTensor:
     __A				 =       input_ids_query.device
     __A				 =       self.embed_questions(UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase     )
     __A				 =       self.embed_answers(UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase     )
     __A				 =       torch.mm(UpperCAmelCase  ,    a_reps.t()     )
     __A				 =       self.ce_loss(UpperCAmelCase  ,    torch.arange(compare_scores.shape[1]     ).to(UpperCAmelCase     )     )
     __A				 =       self.ce_loss(compare_scores.t()  ,    torch.arange(compare_scores.shape[0]     ).to(UpperCAmelCase     )     )
     __A				 =       (loss_qa + loss_aq) / 2
     return loss








import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import OPTConfig, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
    import os

    # The slow tests are often failing with OOM error on GPU
    # This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
    # but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
    _UpperCamelCase   :							Dict     =	"""platform"""

    import jax
    import jax.numpy as jnp

    from transformers import FlaxOPTForCausalLM, FlaxOPTModel, GPTaTokenizer
def 				__UpperCamelCase  (		snake_case    ,	snake_case    ,	snake_case=None    ,	snake_case=None     )  ->						str:


   '''simple docstring'''







   if attention_mask is None:
      __A				 =       np.where(input_ids != config.pad_token_id    ,	1    ,	0     )
   return {
       "input_ids": input_ids,
       "attention_mask": attention_mask,
   }







@require_flax
class 						_lowerCAmelCase:

  """simple docstring"""







  def __init__( self  ,    UpperCAmelCase  ,    UpperCAmelCase=13  ,    UpperCAmelCase=7  ,    UpperCAmelCase=True  ,    UpperCAmelCase=False  ,    UpperCAmelCase=99  ,    UpperCAmelCase=16  ,    UpperCAmelCase=2  ,    UpperCAmelCase=4  ,    UpperCAmelCase=4  ,    UpperCAmelCase="gelu"  ,    UpperCAmelCase=0.1  ,    UpperCAmelCase=0.1  ,    UpperCAmelCase=20  ,    UpperCAmelCase=2  ,    UpperCAmelCase=1  ,    UpperCAmelCase=0  ,    UpperCAmelCase=16  ,    UpperCAmelCase=16  ,    UpperCAmelCase=0.02  ,    )-> List[str]:
     __A				 =       parent
     __A				 =       batch_size
     __A				 =       seq_length
     __A				 =       is_training
     __A				 =       use_labels
     __A				 =       vocab_size
     __A				 =       hidden_size
     __A				 =       num_hidden_layers
     __A				 =       num_attention_heads
     __A				 =       intermediate_size
     __A				 =       hidden_act
     __A				 =       hidden_dropout_prob
     __A				 =       attention_probs_dropout_prob
     __A				 =       max_position_embeddings
     __A				 =       eos_token_id
     __A				 =       pad_token_id
     __A				 =       bos_token_id
     __A				 =       embed_dim
     __A				 =       word_embed_proj_dim
     __A				 =       initializer_range
     __A				 =       False

  def 		SCREAMING_SNAKE_CASE__	( self     )-> Any:
     __A				 =       np.clip(ids_tensor([self.batch_size, self.seq_length - 1]  ,    self.vocab_size     )  ,    3  ,    self.vocab_size     )
     __A				 =       np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)  ,    dtype=np.intaa     ))  ,    -1     )

     __A				 =       OPTConfig(
         vocab_size=self.vocab_size  ,    hidden_size=self.hidden_size  ,    num_hidden_layers=self.num_hidden_layers  ,    num_attention_heads=self.num_attention_heads  ,    ffn_dim=self.intermediate_size  ,    dropout=self.hidden_dropout_prob  ,    attention_dropout=self.attention_probs_dropout_prob  ,    max_position_embeddings=self.max_position_embeddings  ,    eos_token_id=self.eos_token_id  ,    bos_token_id=self.bos_token_id  ,    pad_token_id=self.pad_token_id  ,    embed_dim=self.embed_dim  ,    is_encoder_decoder=UpperCAmelCase  ,    word_embed_proj_dim=self.word_embed_proj_dim  ,    initializer_range=self.initializer_range  ,    use_cache=UpperCAmelCase  ,    )
     __A				 =       prepare_opt_inputs_dict(UpperCAmelCase  ,    UpperCAmelCase     )
     return config, inputs_dict

  def 		SCREAMING_SNAKE_CASE__	( self     )-> int:
     __A ,						__A				 =       self.prepare_config_and_inputs()
     return config, inputs_dict

  def 		SCREAMING_SNAKE_CASE__	( self  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase     )-> Any:
     __A				 =       20
     __A				 =       model_class_name(UpperCAmelCase     )

     __A				 =       inputs_dict['''input_ids''']
     __A				 =       inputs_dict['''attention_mask''']

     __A				 =       model.init_cache(input_ids.shape[0]  ,    UpperCAmelCase     )
     __A				 =       jnp.ones((input_ids.shape[0], max_length)  ,    dtype='''i4'''     )

     __A				 =       jnp.broadcast_to(
         jnp.arange(input_ids.shape[-1] - 1     )[None, :]  ,    (input_ids.shape[0], input_ids.shape[-1] - 1)  ,    )
     __A				 =       model(
         input_ids[:, :-1]  ,    attention_mask=UpperCAmelCase  ,    past_key_values=UpperCAmelCase  ,    position_ids=UpperCAmelCase  ,    )

     __A				 =       jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]  ,    dtype='''i4'''     )
     __A				 =       model(
         input_ids[:, -1:]  ,    attention_mask=UpperCAmelCase  ,    past_key_values=outputs_cache.past_key_values  ,    position_ids=UpperCAmelCase  ,    )

     __A				 =       model(UpperCAmelCase     )

     __A				 =       np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])     )     )
     self.parent.assertTrue(diff < 1e-3  ,    msg=f"Max diff is {diff}"     )





  def 		SCREAMING_SNAKE_CASE__	( self  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase     )-> str:
     __A				 =       20
     __A				 =       model_class_name(UpperCAmelCase     )

     __A ,						__A				 =       (
         inputs_dict['''input_ids'''],
         inputs_dict['''attention_mask'''],
     )

     __A				 =       jnp.concatenate(
         [
             attention_mask,
             jnp.zeros((attention_mask.shape[0], max_length - attention_mask.shape[1])     ),
         ]  ,    axis=-1  ,    )

     __A				 =       model.init_cache(input_ids.shape[0]  ,    UpperCAmelCase     )
     __A				 =       jnp.broadcast_to(
         jnp.arange(input_ids.shape[-1] - 1     )[None, :]  ,    (input_ids.shape[0], input_ids.shape[-1] - 1)  ,    )

     __A				 =       model(
         input_ids[:, :-1]  ,    attention_mask=UpperCAmelCase  ,    past_key_values=UpperCAmelCase  ,    position_ids=UpperCAmelCase  ,    )
     __A				 =       jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]  ,    dtype='''i4'''     )
     __A				 =       model(
         input_ids[:, -1:]  ,    past_key_values=outputs_cache.past_key_values  ,    attention_mask=UpperCAmelCase  ,    position_ids=UpperCAmelCase  ,    )

     __A				 =       model(UpperCAmelCase  ,    attention_mask=UpperCAmelCase     )

     __A				 =       np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])     )     )
     self.parent.assertTrue(diff < 1e-3  ,    msg=f"Max diff is {diff}"     )







@require_flax
class 						_lowerCAmelCase(						_a		,   unittest.TestCase		,   _a):

  """simple docstring"""







  lowerCamelCase__      					=   (FlaxOPTModel, FlaxOPTForCausalLM) if is_flax_available() else ()
  lowerCamelCase__      					=   () if is_flax_available() else ()

  def 		SCREAMING_SNAKE_CASE__	( self     )-> List[Any]:
     __A				 =       FlaxOPTModelTester(self     )

  def 		SCREAMING_SNAKE_CASE__	( self     )-> List[Any]:
     __A ,						__A				 =       self.model_tester.prepare_config_and_inputs()
     for model_class in self.all_model_classes:
        self.model_tester.check_use_cache_forward(UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase     )

  def 		SCREAMING_SNAKE_CASE__	( self     )-> int:
     __A ,						__A				 =       self.model_tester.prepare_config_and_inputs()
     for model_class in self.all_model_classes:
        self.model_tester.check_use_cache_forward_with_attn_mask(UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase     )





  @slow
  def 		SCREAMING_SNAKE_CASE__	( self     )-> Union[str, Any]:
     for model_class_name in self.all_model_classes:
        __A				 =       model_class_name.from_pretrained('''facebook/opt-125m'''     )
        __A				 =       np.ones((1, 1)     ) * model.config.eos_token_id
        __A				 =       model(UpperCAmelCase     )
        self.assertIsNotNone(UpperCAmelCase     )







@require_sentencepiece
@require_flax
class 						_lowerCAmelCase(						unittest.TestCase):

  """simple docstring"""







  @slow
  def 		SCREAMING_SNAKE_CASE__	( self     )-> Tuple:
     __A				 =       FlaxOPTModel.from_pretrained('''facebook/opt-350m'''     )
     __A				 =       jnp.array([[0, 3_14_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69, 4_60_78, 15_88, 2]]     )
     __A				 =       model(input_ids=UpperCAmelCase     ).last_hidden_state
     __A				 =       (1, 11, 5_12)
     self.assertEqual(output.shape  ,    UpperCAmelCase     )
     __A				 =       jnp.array(
         [[-0.2867, -1.9256, -0.3062], [-1.2711, -0.1337, -0.1897], [0.4109, 0.1187, -1.3142]]     )
     self.assertTrue(jnp.allclose(output[:, :3, :3]  ,    UpperCAmelCase  ,    atol=4e-2     )     )







@require_flax
@slow
class 						_lowerCAmelCase(						unittest.TestCase):

  """simple docstring"""







  def 		SCREAMING_SNAKE_CASE__	( self     )-> List[str]:
     super().setUp()
     __A				 =       '''facebook/opt-350m'''





  def 		SCREAMING_SNAKE_CASE__	( self     )-> Dict:
     __A				 =       FlaxOPTForCausalLM.from_pretrained(self.path_model     )
     __A				 =       GPTaTokenizer.from_pretrained(self.path_model     )

     __A				 =       [
         '''Today is a beautiful day and I want to''',
         '''In the city of''',
         '''Paris is the capital of France and''',
         '''Computers and mobile phones have taken''',
     ]
     # verify that prompt without BOS token is identical to Metaseq -> add_special_tokens=False
     __A				 =       tokenizer(UpperCAmelCase  ,    return_tensors='''jax'''  ,    padding=UpperCAmelCase  ,    add_special_tokens=UpperCAmelCase     )
     __A				 =       model(inputs.input_ids  ,    attention_mask=inputs.attention_mask     )[0].mean(axis=-1     )
     __A				 =       jnp.array(
         [
             [1.3851, -13.8923, -10.5229, -10.7533, -0.2309, -10.2384, -0.5365, -9.0947, -5.1670],
             [-4.7073, -10.6276, -3.9415, -21.5242, -0.2822, -0.2822, -0.2822, -0.2822, -0.2822],
             [0.6247, -3.4229, -8.9179, -1.4297, -14.1650, 1.4146, -9.0218, -0.2703, -0.2703],
             [6.4783, -1.9913, -10.7926, -2.3336, 1.5092, -0.9974, -6.8213, 1.3477, 1.3477],
         ]     )
     self.assertTrue(jnp.allclose(UpperCAmelCase  ,    UpperCAmelCase  ,    atol=4e-2     )     )

     __A				 =       jax.jit(UpperCAmelCase     )
     __A				 =       model(inputs.input_ids  ,    attention_mask=inputs.attention_mask     )[0].mean(axis=-1     )
     self.assertTrue(jnp.allclose(UpperCAmelCase  ,    UpperCAmelCase  ,    atol=4e-2     )     )







@require_flax
@slow
class 						_lowerCAmelCase(						unittest.TestCase):

  """simple docstring"""







  @property
  def 		SCREAMING_SNAKE_CASE__	( self     )-> Union[str, Any]:
     return [
         "Today is a beautiful day and I want",
         "In the city of",
         "Paris is the capital of France and",
         "Computers and mobile phones have taken",
     ]

  def 		SCREAMING_SNAKE_CASE__	( self     )-> str:
     __A				 =       '''facebook/opt-125m'''

     __A				 =       [
         '''Today is a beautiful day and I want to''',
         '''In the city of New York, the city''',
         '''Paris is the capital of France and the capital''',
         '''Computers and mobile phones have taken over the''',
     ]

     __A				 =       []

     __A				 =       FlaxOPTForCausalLM.from_pretrained(UpperCAmelCase     )
     __A				 =       GPTaTokenizer.from_pretrained(UpperCAmelCase     )

     for prompt in self.prompts:
        __A				 =       tokenizer(UpperCAmelCase  ,    return_tensors='''jax'''     ).input_ids

        __A				 =       model.generate(UpperCAmelCase  ,    max_length=10     )
        __A				 =       generated_ids[0]

        __A				 =       tokenizer.batch_decode(UpperCAmelCase  ,    skip_special_tokens=UpperCAmelCase     )
        predicted_outputs += generated_string

     self.assertListEqual(UpperCAmelCase  ,    UpperCAmelCase     )

  def 		SCREAMING_SNAKE_CASE__	( self     )-> List[Any]:
     __A				 =       '''facebook/opt-350m'''

     __A				 =       [
         '''Today is a beautiful day and I want to''',
         '''In the city of San Francisco, the city''',
         '''Paris is the capital of France and the capital''',
         '''Computers and mobile phones have taken over the''',
     ]

     __A				 =       []
     __A				 =       FlaxOPTForCausalLM.from_pretrained(UpperCAmelCase     )
     __A				 =       GPTaTokenizer.from_pretrained(UpperCAmelCase     )

     for prompt in self.prompts:
        __A				 =       tokenizer(UpperCAmelCase  ,    return_tensors='''jax'''     ).input_ids

        __A				 =       model.generate(UpperCAmelCase  ,    max_length=10     )
        __A				 =       generated_ids[0]

        __A				 =       tokenizer.batch_decode(UpperCAmelCase  ,    skip_special_tokens=UpperCAmelCase     )
        predicted_outputs += generated_string

     self.assertListEqual(UpperCAmelCase  ,    UpperCAmelCase     )

  def 		SCREAMING_SNAKE_CASE__	( self     )-> Dict:
     __A				 =       '''facebook/opt-125m'''
     __A				 =       [
         '''Today is a beautiful day and I want to thank''',
         '''In the city of Rome Canaver Canaver Canaver Canaver''',
     ]
     __A				 =       FlaxOPTForCausalLM.from_pretrained(UpperCAmelCase     )
     __A				 =       GPTaTokenizer.from_pretrained(UpperCAmelCase     )
     __A				 =       tokenizer(
         [
             '''Today is a beautiful day and I want to''',
             '''In the city of''',
         ]  ,    return_tensors='''jax'''  ,    padding=UpperCAmelCase  ,    )

     __A				 =       jax.jit(model.generate     )

     __A				 =       jit_generate(inputs['''input_ids''']  ,    attention_mask=inputs['''attention_mask''']     ).sequences

     __A				 =       tokenizer.batch_decode(UpperCAmelCase  ,    skip_special_tokens=UpperCAmelCase     )

     self.assertIsNotNone(UpperCAmelCase  ,    UpperCAmelCase     )





  def 		SCREAMING_SNAKE_CASE__	( self     )-> Any:
     __A				 =       '''facebook/opt-350m'''

     __A				 =       GPTaTokenizer.from_pretrained(UpperCAmelCase     )
     __A				 =       FlaxOPTForCausalLM.from_pretrained(UpperCAmelCase     )

     __A				 =       '''left'''

     # use different length sentences to test batching
     __A				 =       [
         '''Hello, my dog is a little''',
         '''Today, I''',
     ]

     __A				 =       tokenizer(UpperCAmelCase  ,    return_tensors='''jax'''  ,    padding=UpperCAmelCase     )
     __A				 =       inputs['''input_ids''']

     __A				 =       model.generate(input_ids=UpperCAmelCase  ,    attention_mask=inputs['''attention_mask''']  ,    trace=UpperCAmelCase     )

     __A				 =       tokenizer(sentences[0]  ,    return_tensors='''jax'''     ).input_ids
     __A				 =       model.generate(input_ids=UpperCAmelCase     )

     __A				 =       inputs_non_padded.shape[-1] - inputs['''attention_mask'''][-1].sum()
     __A				 =       tokenizer(sentences[1]  ,    return_tensors='''jax'''     ).input_ids
     __A				 =       model.generate(input_ids=UpperCAmelCase  ,    max_length=model.config.max_length - num_paddings     )

     __A				 =       tokenizer.batch_decode(outputs[0]  ,    skip_special_tokens=UpperCAmelCase     )
     __A				 =       tokenizer.decode(output_non_padded[0][0]  ,    skip_special_tokens=UpperCAmelCase     )
     __A				 =       tokenizer.decode(output_padded[0][0]  ,    skip_special_tokens=UpperCAmelCase     )

     __A				 =       [
         '''Hello, my dog is a little bit of a dork.\nI\'m a little bit''',
         '''Today, I was in the middle of a conversation with a friend about the''',
     ]
     self.assertListEqual(UpperCAmelCase  ,    UpperCAmelCase     )
     self.assertListEqual(UpperCAmelCase  ,    [non_padded_sentence, padded_sentence]     )

