


import unittest

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

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 (
        RemBertConfig,
        RemBertForCausalLM,
        RemBertForMaskedLM,
        RemBertForMultipleChoice,
        RemBertForQuestionAnswering,
        RemBertForSequenceClassification,
        RemBertForTokenClassification,
        RemBertModel,
    )
    from transformers.models.rembert.modeling_rembert import REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST


class     __lowerCAmelCase     :






  """simple docstring"""





  def __init__(     self       :       Optional[Any]    ,				_lowerCAmelCase       :       Dict    ,				_lowerCAmelCase       :       Optional[Any]=1_3    ,				_lowerCAmelCase       :       Any=7    ,				_lowerCAmelCase       :       int=True    ,				_lowerCAmelCase       :       Dict=True    ,				_lowerCAmelCase       :       Union[str, Any]=True    ,				_lowerCAmelCase       :       Union[str, Any]=True    ,				_lowerCAmelCase       :       Union[str, Any]=9_9    ,				_lowerCAmelCase       :       str=3_2    ,				_lowerCAmelCase       :       Tuple=1_8    ,				_lowerCAmelCase       :       int=4_3    ,				_lowerCAmelCase       :       Dict=5    ,				_lowerCAmelCase       :       List[Any]=4    ,				_lowerCAmelCase       :       Any=3_7    ,				_lowerCAmelCase       :       str="gelu"    ,				_lowerCAmelCase       :       int=0.1    ,				_lowerCAmelCase       :       Optional[int]=0.1    ,				_lowerCAmelCase       :       List[str]=5_1_2    ,				_lowerCAmelCase       :       List[Any]=1_6    ,				_lowerCAmelCase       :       Tuple=2    ,				_lowerCAmelCase       :       List[Any]=0.02    ,				_lowerCAmelCase       :       List[Any]=3    ,				_lowerCAmelCase       :       Optional[Any]=4    ,				_lowerCAmelCase       :       int=None    ,				)   -> Dict:




     """simple docstring"""
     snake_case_					    =					parent
     snake_case_					    =					batch_size
     snake_case_					    =					seq_length
     snake_case_					    =					is_training
     snake_case_					    =					use_input_mask
     snake_case_					    =					use_token_type_ids
     snake_case_					    =					use_labels
     snake_case_					    =					vocab_size
     snake_case_					    =					hidden_size
     snake_case_					    =					input_embedding_size
     snake_case_					    =					output_embedding_size
     snake_case_					    =					num_hidden_layers
     snake_case_					    =					num_attention_heads
     snake_case_					    =					intermediate_size
     snake_case_					    =					hidden_act
     snake_case_					    =					hidden_dropout_prob
     snake_case_					    =					attention_probs_dropout_prob
     snake_case_					    =					max_position_embeddings
     snake_case_					    =					type_vocab_size
     snake_case_					    =					type_sequence_label_size
     snake_case_					    =					initializer_range
     snake_case_					    =					num_labels
     snake_case_					    =					num_choices
     snake_case_					    =					scope



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




     """simple docstring"""
     snake_case_					    =					ids_tensor([self.batch_size, self.seq_length]    ,				self.vocab_size  )

     snake_case_					    =					None
     if self.use_input_mask:
        snake_case_					    =					random_attention_mask([self.batch_size, self.seq_length]  )

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

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

     snake_case_					    =					RemBertConfig(
         vocab_size=self.vocab_size    ,				hidden_size=self.hidden_size    ,				input_embedding_size=self.input_embedding_size    ,				output_embedding_size=self.output_embedding_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    ,				)

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



  def 					lowerCAmelCase__  (     self       :       int  )   -> Optional[int]:




     """simple docstring"""
     (
         (
         snake_case_    
     )    ,    (
         snake_case_    
     )    ,    (
         snake_case_    
     )    ,    (
         snake_case_    
     )    ,    (
         snake_case_    
     )    ,    (
         snake_case_    
     )    ,    (
         snake_case_    
     )    ,        
     )					    =					self.prepare_config_and_inputs()

     snake_case_					    =					True
     snake_case_					    =					floats_tensor([self.batch_size, self.seq_length, self.hidden_size]  )
     snake_case_					    =					ids_tensor([self.batch_size, self.seq_length]    ,				vocab_size=2  )

     return (
         config,
         input_ids,
         token_type_ids,
         input_mask,
         sequence_labels,
         token_labels,
         choice_labels,
         encoder_hidden_states,
         encoder_attention_mask,
     )



  def 					lowerCAmelCase__  (     self       :       Dict    ,				_lowerCAmelCase       :       Optional[int]    ,				_lowerCAmelCase       :       Any    ,				_lowerCAmelCase       :       str    ,				_lowerCAmelCase       :       Tuple    ,				_lowerCAmelCase       :       Any    ,				_lowerCAmelCase       :       Any    ,				_lowerCAmelCase       :       str  )   -> int:




     """simple docstring"""
     snake_case_					    =					RemBertModel(config=_lowerCAmelCase  )
     model.to(_lowerCAmelCase  )
     model.eval()
     snake_case_					    =					model(_lowerCAmelCase    ,				attention_mask=_lowerCAmelCase    ,				token_type_ids=_lowerCAmelCase  )
     snake_case_					    =					model(_lowerCAmelCase    ,				token_type_ids=_lowerCAmelCase  )
     snake_case_					    =					model(_lowerCAmelCase  )
     self.parent.assertEqual(result.last_hidden_state.shape    ,				(self.batch_size, self.seq_length, self.hidden_size)  )



  def 					lowerCAmelCase__  (     self       :       Union[str, Any]    ,				_lowerCAmelCase       :       List[Any]    ,				_lowerCAmelCase       :       int    ,				_lowerCAmelCase       :       Any    ,				_lowerCAmelCase       :       Any    ,				_lowerCAmelCase       :       Optional[Any]    ,				_lowerCAmelCase       :       str    ,				_lowerCAmelCase       :       List[Any]    ,				_lowerCAmelCase       :       Union[str, Any]    ,				_lowerCAmelCase       :       List[str]    ,				)   -> Tuple:




     """simple docstring"""
     snake_case_					    =					True
     snake_case_					    =					RemBertModel(_lowerCAmelCase  )
     model.to(_lowerCAmelCase  )
     model.eval()
     snake_case_					    =					model(
         _lowerCAmelCase    ,				attention_mask=_lowerCAmelCase    ,				token_type_ids=_lowerCAmelCase    ,				encoder_hidden_states=_lowerCAmelCase    ,				encoder_attention_mask=_lowerCAmelCase    ,				)
     snake_case_					    =					model(
         _lowerCAmelCase    ,				attention_mask=_lowerCAmelCase    ,				token_type_ids=_lowerCAmelCase    ,				encoder_hidden_states=_lowerCAmelCase    ,				)
     snake_case_					    =					model(_lowerCAmelCase    ,				attention_mask=_lowerCAmelCase    ,				token_type_ids=_lowerCAmelCase  )
     self.parent.assertEqual(result.last_hidden_state.shape    ,				(self.batch_size, self.seq_length, self.hidden_size)  )



  def 					lowerCAmelCase__  (     self       :       Any    ,				_lowerCAmelCase       :       Any    ,				_lowerCAmelCase       :       List[str]    ,				_lowerCAmelCase       :       Union[str, Any]    ,				_lowerCAmelCase       :       str    ,				_lowerCAmelCase       :       str    ,				_lowerCAmelCase       :       List[Any]    ,				_lowerCAmelCase       :       Optional[Any]    ,				_lowerCAmelCase       :       List[Any]    ,				_lowerCAmelCase       :       str    ,				)   -> Union[str, Any]:




     """simple docstring"""
     snake_case_					    =					RemBertForCausalLM(config=_lowerCAmelCase  )
     model.to(_lowerCAmelCase  )
     model.eval()
     snake_case_					    =					model(_lowerCAmelCase    ,				attention_mask=_lowerCAmelCase    ,				token_type_ids=_lowerCAmelCase    ,				labels=_lowerCAmelCase  )
     self.parent.assertEqual(result.logits.shape    ,				(self.batch_size, self.seq_length, self.vocab_size)  )



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




     """simple docstring"""
     snake_case_					    =					RemBertForMaskedLM(config=_lowerCAmelCase  )
     model.to(_lowerCAmelCase  )
     model.eval()
     snake_case_					    =					model(_lowerCAmelCase    ,				attention_mask=_lowerCAmelCase    ,				token_type_ids=_lowerCAmelCase    ,				labels=_lowerCAmelCase  )
     self.parent.assertEqual(result.logits.shape    ,				(self.batch_size, self.seq_length, self.vocab_size)  )



  def 					lowerCAmelCase__  (     self       :       List[str]    ,				_lowerCAmelCase       :       Dict    ,				_lowerCAmelCase       :       int    ,				_lowerCAmelCase       :       Optional[int]    ,				_lowerCAmelCase       :       Tuple    ,				_lowerCAmelCase       :       List[Any]    ,				_lowerCAmelCase       :       List[str]    ,				_lowerCAmelCase       :       Optional[Any]    ,				_lowerCAmelCase       :       Optional[Any]    ,				_lowerCAmelCase       :       List[str]    ,				)   -> List[Any]:




     """simple docstring"""
     snake_case_					    =					True
     snake_case_					    =					True
     snake_case_					    =					RemBertForCausalLM(config=_lowerCAmelCase  )
     model.to(_lowerCAmelCase  )
     model.eval()

     # first forward pass
     snake_case_					    =					model(
         _lowerCAmelCase    ,				attention_mask=_lowerCAmelCase    ,				encoder_hidden_states=_lowerCAmelCase    ,				encoder_attention_mask=_lowerCAmelCase    ,				use_cache=_lowerCAmelCase    ,				)
     snake_case_					    =					outputs.past_key_values

     # create hypothetical multiple next token and extent to next_input_ids
     snake_case_					    =					ids_tensor((self.batch_size, 3)    ,				config.vocab_size  )
     snake_case_					    =					ids_tensor((self.batch_size, 3)    ,				vocab_size=2  )

     # append to next input_ids and
     snake_case_					    =					torch.cat([input_ids, next_tokens]    ,				dim=-1  )
     snake_case_					    =					torch.cat([input_mask, next_mask]    ,				dim=-1  )

     snake_case_					    =					model(
         _lowerCAmelCase    ,				attention_mask=_lowerCAmelCase    ,				encoder_hidden_states=_lowerCAmelCase    ,				encoder_attention_mask=_lowerCAmelCase    ,				output_hidden_states=_lowerCAmelCase    ,				)["hidden_states"][0]
     snake_case_					    =					model(
         _lowerCAmelCase    ,				attention_mask=_lowerCAmelCase    ,				encoder_hidden_states=_lowerCAmelCase    ,				encoder_attention_mask=_lowerCAmelCase    ,				past_key_values=_lowerCAmelCase    ,				output_hidden_states=_lowerCAmelCase    ,				)["hidden_states"][0]

     # select random slice
     snake_case_					    =					ids_tensor((1,)    ,				output_from_past.shape[-1]  ).item()
     snake_case_					    =					output_from_no_past[:, -3:, random_slice_idx].detach()
     snake_case_					    =					output_from_past[:, :, random_slice_idx].detach()

     self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]  )

     # test that outputs are equal for slice
     self.parent.assertTrue(torch.allclose(_lowerCAmelCase    ,				_lowerCAmelCase    ,				atol=1e-3  )  )



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




     """simple docstring"""
     snake_case_					    =					RemBertForQuestionAnswering(config=_lowerCAmelCase  )
     model.to(_lowerCAmelCase  )
     model.eval()
     snake_case_					    =					model(
         _lowerCAmelCase    ,				attention_mask=_lowerCAmelCase    ,				token_type_ids=_lowerCAmelCase    ,				start_positions=_lowerCAmelCase    ,				end_positions=_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)  )



  def 					lowerCAmelCase__  (     self       :       List[Any]    ,				_lowerCAmelCase       :       Optional[Any]    ,				_lowerCAmelCase       :       Dict    ,				_lowerCAmelCase       :       List[str]    ,				_lowerCAmelCase       :       Optional[Any]    ,				_lowerCAmelCase       :       int    ,				_lowerCAmelCase       :       List[str]    ,				_lowerCAmelCase       :       str  )   -> str:




     """simple docstring"""
     snake_case_					    =					self.num_labels
     snake_case_					    =					RemBertForSequenceClassification(_lowerCAmelCase  )
     model.to(_lowerCAmelCase  )
     model.eval()
     snake_case_					    =					model(_lowerCAmelCase    ,				attention_mask=_lowerCAmelCase    ,				token_type_ids=_lowerCAmelCase    ,				labels=_lowerCAmelCase  )
     self.parent.assertEqual(result.logits.shape    ,				(self.batch_size, self.num_labels)  )



  def 					lowerCAmelCase__  (     self       :       Any    ,				_lowerCAmelCase       :       Any    ,				_lowerCAmelCase       :       Optional[int]    ,				_lowerCAmelCase       :       Any    ,				_lowerCAmelCase       :       Dict    ,				_lowerCAmelCase       :       List[Any]    ,				_lowerCAmelCase       :       Tuple    ,				_lowerCAmelCase       :       Optional[Any]  )   -> Any:




     """simple docstring"""
     snake_case_					    =					self.num_labels
     snake_case_					    =					RemBertForTokenClassification(config=_lowerCAmelCase  )
     model.to(_lowerCAmelCase  )
     model.eval()
     snake_case_					    =					model(_lowerCAmelCase    ,				attention_mask=_lowerCAmelCase    ,				token_type_ids=_lowerCAmelCase    ,				labels=_lowerCAmelCase  )
     self.parent.assertEqual(result.logits.shape    ,				(self.batch_size, self.seq_length, self.num_labels)  )



  def 					lowerCAmelCase__  (     self       :       Tuple    ,				_lowerCAmelCase       :       List[str]    ,				_lowerCAmelCase       :       Optional[int]    ,				_lowerCAmelCase       :       List[Any]    ,				_lowerCAmelCase       :       str    ,				_lowerCAmelCase       :       Any    ,				_lowerCAmelCase       :       Optional[Any]    ,				_lowerCAmelCase       :       Tuple  )   -> str:




     """simple docstring"""
     snake_case_					    =					self.num_choices
     snake_case_					    =					RemBertForMultipleChoice(config=_lowerCAmelCase  )
     model.to(_lowerCAmelCase  )
     model.eval()
     snake_case_					    =					input_ids.unsqueeze(1  ).expand(-1    ,				self.num_choices    ,				-1  ).contiguous()
     snake_case_					    =					token_type_ids.unsqueeze(1  ).expand(-1    ,				self.num_choices    ,				-1  ).contiguous()
     snake_case_					    =					input_mask.unsqueeze(1  ).expand(-1    ,				self.num_choices    ,				-1  ).contiguous()
     snake_case_					    =					model(
         _lowerCAmelCase    ,				attention_mask=_lowerCAmelCase    ,				token_type_ids=_lowerCAmelCase    ,				labels=_lowerCAmelCase    ,				)
     self.parent.assertEqual(result.logits.shape    ,				(self.batch_size, self.num_choices)  )



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




     """simple docstring"""
     snake_case_					    =					self.prepare_config_and_inputs()
     (
         (
         snake_case_    
     )    ,    (
         snake_case_    
     )    ,    (
         snake_case_    
     )    ,    (
         snake_case_    
     )    ,    (
         snake_case_    
     )    ,    (
         snake_case_    
     )    ,    (
         snake_case_    
     )    ,        
     )					    =					config_and_inputs
     snake_case_					    =					{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
     return config, inputs_dict






@require_torch
class     __lowerCAmelCase     ( a ,  a ,  unittest.TestCase					):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		    =       (
      (
          RemBertModel,
          RemBertForMaskedLM,
          RemBertForCausalLM,
          RemBertForMultipleChoice,
          RemBertForQuestionAnswering,
          RemBertForSequenceClassification,
          RemBertForTokenClassification,
      )
      if is_torch_available()
      else ()
  )
  _SCREAMING_SNAKE_CASE		    =       (RemBertForCausalLM,) if is_torch_available() else ()
  _SCREAMING_SNAKE_CASE		    =       (
      {
          'feature-extraction': RemBertModel,
          'fill-mask': RemBertForMaskedLM,
          'question-answering': RemBertForQuestionAnswering,
          'text-classification': RemBertForSequenceClassification,
          'text-generation': RemBertForCausalLM,
          'token-classification': RemBertForTokenClassification,
          'zero-shot': RemBertForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )



  def 					lowerCAmelCase__  (     self       :       int  )   -> int:




     """simple docstring"""
     snake_case_					    =					RemBertModelTester(self  )
     snake_case_					    =					ConfigTester(self    ,				config_class=_lowerCAmelCase    ,				hidden_size=3_7  )



  def 					lowerCAmelCase__  (     self       :       int  )   -> Tuple:




     """simple docstring"""
     self.config_tester.run_common_tests()



  def 					lowerCAmelCase__  (     self       :       Optional[int]  )   -> Optional[Any]:




     """simple docstring"""
     snake_case_					    =					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*_lowerCAmelCase  )



  def 					lowerCAmelCase__  (     self       :       Tuple  )   -> int:




     """simple docstring"""
     snake_case_					    =					self.model_tester.prepare_config_and_inputs()
     for type in ["absolute", "relative_key", "relative_key_query"]:
        snake_case_					    =					type
        self.model_tester.create_and_check_model(*_lowerCAmelCase  )



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




     """simple docstring"""
     snake_case_					    =					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_lm(*_lowerCAmelCase  )



  def 					lowerCAmelCase__  (     self       :       Tuple  )   -> Any:




     """simple docstring"""
     snake_case_					    =					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_multiple_choice(*_lowerCAmelCase  )



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




     """simple docstring"""
     snake_case_					    =					self.model_tester.prepare_config_and_inputs_for_decoder()
     self.model_tester.create_and_check_decoder_model_past_large_inputs(*_lowerCAmelCase  )



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




     """simple docstring"""
     snake_case_					    =					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_question_answering(*_lowerCAmelCase  )



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




     """simple docstring"""
     snake_case_					    =					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_sequence_classification(*_lowerCAmelCase  )



  def 					lowerCAmelCase__  (     self       :       Optional[int]  )   -> Optional[Any]:




     """simple docstring"""
     snake_case_					    =					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_token_classification(*_lowerCAmelCase  )



  def 					lowerCAmelCase__  (     self       :       Any  )   -> Tuple:




     """simple docstring"""
     snake_case_					    =					self.model_tester.prepare_config_and_inputs_for_decoder()
     self.model_tester.create_and_check_model_as_decoder(*_lowerCAmelCase  )



  def 					lowerCAmelCase__  (     self       :       Union[str, Any]  )   -> List[str]:




     """simple docstring"""
     (
         (
         snake_case_    
     )    ,    (
         snake_case_    
     )    ,    (
         snake_case_    
     )    ,    (
         snake_case_    
     )    ,    (
         snake_case_    
     )    ,    (
         snake_case_    
     )    ,    (
         snake_case_    
     )    ,    (
         snake_case_    
     )    ,    (
         snake_case_    
     )    ,        
     )					    =					self.model_tester.prepare_config_and_inputs_for_decoder()

     snake_case_					    =					None

     self.model_tester.create_and_check_model_as_decoder(
         _lowerCAmelCase    ,				_lowerCAmelCase    ,				_lowerCAmelCase    ,				_lowerCAmelCase    ,				_lowerCAmelCase    ,				_lowerCAmelCase    ,				_lowerCAmelCase    ,				_lowerCAmelCase    ,				_lowerCAmelCase    ,				)



  @slow
  def 					lowerCAmelCase__  (     self       :       Union[str, Any]  )   -> str:




     """simple docstring"""
     for model_name in REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        snake_case_					    =					RemBertModel.from_pretrained(_lowerCAmelCase  )
        self.assertIsNotNone(_lowerCAmelCase  )






@require_torch
class     __lowerCAmelCase     ( unittest.TestCase					):






  """simple docstring"""





  @slow
  def 					lowerCAmelCase__  (     self       :       int  )   -> Tuple:




     """simple docstring"""
     snake_case_					    =					RemBertModel.from_pretrained("google/rembert"  )
     snake_case_					    =					torch.tensor([[3_1_2, 5_6_4_9_8, 3_1_3, 2_1_2_5, 3_1_3]]  )
     snake_case_					    =					torch.tensor([[0, 0, 0, 1, 1]]  )
     with torch.no_grad():
        snake_case_					    =					model(_lowerCAmelCase    ,				token_type_ids=_lowerCAmelCase    ,				output_hidden_states=_lowerCAmelCase  )

     snake_case_					    =					1_1_5_2

     snake_case_					    =					torch.Size((1, 5, hidden_size)  )
     self.assertEqual(output["last_hidden_state"].shape    ,				_lowerCAmelCase  )

     snake_case_					    =					torch.tensor(
         [
             [
                 [0.0_754, -0.2_022, 0.1_904],
                 [-0.3_354, -0.3_692, -0.4_791],
                 [-0.2_314, -0.6_729, -0.0_749],
                 [-0.0_396, -0.3_105, -0.4_234],
                 [-0.1_571, -0.0_525, 0.5_353],
             ]
         ]  )

     # Running on the original tf implementation gives slightly different results here.
     # Not clear why this variations is present
     # TODO: Find reason for discrepancy
     # expected_original_implementation = [[
     #     [0.07630594074726105, -0.20146065950393677, 0.19107051193714142],
     #     [-0.3405614495277405, -0.36971670389175415, -0.4808273911476135],
     #     [-0.22587086260318756, -0.6656315922737122, -0.07844287157058716],
     #     [-0.04145475849509239, -0.3077218234539032, -0.42316967248916626],
     #     [-0.15887849032878876, -0.054529931396245956, 0.5356100797653198]
     # ]]

     self.assertTrue(torch.allclose(output["last_hidden_state"][:, :, :3]    ,				_lowerCAmelCase    ,				atol=1e-4  )  )








import csv
import os

import pytest

from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
from datasets.io.csv import CsvDatasetReader, CsvDatasetWriter

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases



def 			_lowerCAmelCase						(							lowerCAmelCase_							:Optional[Any]		,						lowerCAmelCase_							:List[str]							)->Dict:
   '''simple docstring'''






   assert isinstance(lowerCAmelCase_		,						lowerCAmelCase_							)
   assert dataset.num_rows == 4
   assert dataset.num_columns == 3
   assert dataset.column_names == ["col_1", "col_2", "col_3"]
   for feature, expected_dtype in expected_features.items():
      assert dataset.features[feature].dtype == expected_dtype



@pytest.mark.parametrize("keep_in_memory"		,						[False, True]							)
def 			_lowerCAmelCase						(							lowerCAmelCase_							:Dict		,						lowerCAmelCase_							:List[Any]		,						lowerCAmelCase_							:List[str]							)->int:
   '''simple docstring'''






   snake_case_					    =					tmp_path / "cache"
   snake_case_					    =					{"col_1": "int64", "col_2": "int64", "col_3": "float64"}
   with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
      snake_case_					    =					CsvDatasetReader(lowerCAmelCase_		,						cache_dir=lowerCAmelCase_		,						keep_in_memory=lowerCAmelCase_							).read()
   _check_csv_dataset(lowerCAmelCase_		,						lowerCAmelCase_							)



@pytest.mark.parametrize(
    "features"		,						[
        None,
        {"col_1": "string", "col_2": "int64", "col_3": "float64"},
        {"col_1": "string", "col_2": "string", "col_3": "string"},
        {"col_1": "int32", "col_2": "int32", "col_3": "int32"},
        {"col_1": "float32", "col_2": "float32", "col_3": "float32"},
    ]		,						)
def 			_lowerCAmelCase						(							lowerCAmelCase_							:List[str]		,						lowerCAmelCase_							:Optional[Any]		,						lowerCAmelCase_							:Optional[Any]							)->Optional[int]:
   '''simple docstring'''






   snake_case_					    =					tmp_path / "cache"
   # CSV file loses col_1 string dtype information: default now is "int64" instead of "string"
   snake_case_					    =					{"col_1": "int64", "col_2": "int64", "col_3": "float64"}
   snake_case_					    =					features.copy() if features else default_expected_features
   snake_case_					    =					(
       Features({feature: Value(lowerCAmelCase_							) for feature, dtype in features.items()}							) if features is not None else None
   )
   snake_case_					    =					CsvDatasetReader(lowerCAmelCase_		,						features=lowerCAmelCase_		,						cache_dir=lowerCAmelCase_							).read()
   _check_csv_dataset(lowerCAmelCase_		,						lowerCAmelCase_							)



@pytest.mark.parametrize("split"		,						[None, NamedSplit("train"							), "train", "test"]							)
def 			_lowerCAmelCase						(							lowerCAmelCase_							:Optional[int]		,						lowerCAmelCase_							:Tuple		,						lowerCAmelCase_							:List[Any]							)->Optional[int]:
   '''simple docstring'''






   snake_case_					    =					tmp_path / "cache"
   snake_case_					    =					{"col_1": "int64", "col_2": "int64", "col_3": "float64"}
   snake_case_					    =					CsvDatasetReader(lowerCAmelCase_		,						cache_dir=lowerCAmelCase_		,						split=lowerCAmelCase_							).read()
   _check_csv_dataset(lowerCAmelCase_		,						lowerCAmelCase_							)
   assert dataset.split == split if split else "train"



@pytest.mark.parametrize("path_type"		,						[str, list]							)
def 			_lowerCAmelCase						(							lowerCAmelCase_							:Any		,						lowerCAmelCase_							:Dict		,						lowerCAmelCase_							:List[Any]							)->List[Any]:
   '''simple docstring'''






   if issubclass(lowerCAmelCase_		,						lowerCAmelCase_							):
      snake_case_					    =					csv_path
   elif issubclass(lowerCAmelCase_		,						lowerCAmelCase_							):
      snake_case_					    =					[csv_path]
   snake_case_					    =					tmp_path / "cache"
   snake_case_					    =					{"col_1": "int64", "col_2": "int64", "col_3": "float64"}
   snake_case_					    =					CsvDatasetReader(lowerCAmelCase_		,						cache_dir=lowerCAmelCase_							).read()
   _check_csv_dataset(lowerCAmelCase_		,						lowerCAmelCase_							)



def 			_lowerCAmelCase						(							lowerCAmelCase_							:List[str]		,						lowerCAmelCase_							:Dict		,						lowerCAmelCase_							:List[str]=("train",)							)->Optional[Any]:
   '''simple docstring'''






   assert isinstance(lowerCAmelCase_		,						lowerCAmelCase_							)
   for split in splits:
      snake_case_					    =					dataset_dict[split]
      assert dataset.num_rows == 4
      assert dataset.num_columns == 3
      assert dataset.column_names == ["col_1", "col_2", "col_3"]
      for feature, expected_dtype in expected_features.items():
         assert dataset.features[feature].dtype == expected_dtype



@pytest.mark.parametrize("keep_in_memory"		,						[False, True]							)
def 			_lowerCAmelCase						(							lowerCAmelCase_							:Union[str, Any]		,						lowerCAmelCase_							:Dict		,						lowerCAmelCase_							:List[Any]							)->Any:
   '''simple docstring'''






   snake_case_					    =					tmp_path / "cache"
   snake_case_					    =					{"col_1": "int64", "col_2": "int64", "col_3": "float64"}
   with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
      snake_case_					    =					CsvDatasetReader({"train": csv_path}		,						cache_dir=lowerCAmelCase_		,						keep_in_memory=lowerCAmelCase_							).read()
   _check_csv_datasetdict(lowerCAmelCase_		,						lowerCAmelCase_							)



@pytest.mark.parametrize(
    "features"		,						[
        None,
        {"col_1": "string", "col_2": "int64", "col_3": "float64"},
        {"col_1": "string", "col_2": "string", "col_3": "string"},
        {"col_1": "int32", "col_2": "int32", "col_3": "int32"},
        {"col_1": "float32", "col_2": "float32", "col_3": "float32"},
    ]		,						)
def 			_lowerCAmelCase						(							lowerCAmelCase_							:List[Any]		,						lowerCAmelCase_							:Optional[Any]		,						lowerCAmelCase_							:Any							)->str:
   '''simple docstring'''






   snake_case_					    =					tmp_path / "cache"
   # CSV file loses col_1 string dtype information: default now is "int64" instead of "string"
   snake_case_					    =					{"col_1": "int64", "col_2": "int64", "col_3": "float64"}
   snake_case_					    =					features.copy() if features else default_expected_features
   snake_case_					    =					(
       Features({feature: Value(lowerCAmelCase_							) for feature, dtype in features.items()}							) if features is not None else None
   )
   snake_case_					    =					CsvDatasetReader({"train": csv_path}		,						features=lowerCAmelCase_		,						cache_dir=lowerCAmelCase_							).read()
   _check_csv_datasetdict(lowerCAmelCase_		,						lowerCAmelCase_							)



@pytest.mark.parametrize("split"		,						[None, NamedSplit("train"							), "train", "test"]							)
def 			_lowerCAmelCase						(							lowerCAmelCase_							:str		,						lowerCAmelCase_							:Tuple		,						lowerCAmelCase_							:Any							)->str:
   '''simple docstring'''






   if split:
      snake_case_					    =					{split: csv_path}
   else:
      snake_case_					    =					{"train": csv_path, "test": csv_path}
   snake_case_					    =					tmp_path / "cache"
   snake_case_					    =					{"col_1": "int64", "col_2": "int64", "col_3": "float64"}
   snake_case_					    =					CsvDatasetReader(lowerCAmelCase_		,						cache_dir=lowerCAmelCase_							).read()
   _check_csv_datasetdict(lowerCAmelCase_		,						lowerCAmelCase_		,						splits=list(path.keys()							)							)
   assert all(dataset[split].split == split for split in path.keys()							)



def 			_lowerCAmelCase						(							lowerCAmelCase_							:Tuple							)->str:
   '''simple docstring'''






   with open(lowerCAmelCase_		,						encoding="utf-8"							) as csvfile:
      yield from csv.reader(lowerCAmelCase_							)



def 			_lowerCAmelCase						(							lowerCAmelCase_							:Tuple		,						lowerCAmelCase_							:Dict							)->Any:
   '''simple docstring'''






   snake_case_					    =					tmp_path / "cache"
   snake_case_					    =					os.path.join(lowerCAmelCase_		,						"tmp.csv"							)
   snake_case_					    =					CsvDatasetReader({"train": csv_path}		,						cache_dir=lowerCAmelCase_							).read()
   CsvDatasetWriter(dataset["train"]		,						lowerCAmelCase_		,						num_proc=1							).write()

   snake_case_					    =					iter_csv_file(lowerCAmelCase_							)
   snake_case_					    =					iter_csv_file(lowerCAmelCase_							)

   for rowa, rowa in zip(lowerCAmelCase_		,						lowerCAmelCase_							):
      assert rowa == rowa



def 			_lowerCAmelCase						(							lowerCAmelCase_							:str		,						lowerCAmelCase_							:Any							)->Union[str, Any]:
   '''simple docstring'''






   snake_case_					    =					tmp_path / "cache"
   snake_case_					    =					os.path.join(lowerCAmelCase_		,						"tmp.csv"							)
   snake_case_					    =					CsvDatasetReader({"train": csv_path}		,						cache_dir=lowerCAmelCase_							).read()
   CsvDatasetWriter(dataset["train"]		,						lowerCAmelCase_		,						num_proc=2							).write()

   snake_case_					    =					iter_csv_file(lowerCAmelCase_							)
   snake_case_					    =					iter_csv_file(lowerCAmelCase_							)

   for rowa, rowa in zip(lowerCAmelCase_		,						lowerCAmelCase_							):
      assert rowa == rowa



def 			_lowerCAmelCase						(							lowerCAmelCase_							:Any		,						lowerCAmelCase_							:Optional[int]							)->int:
   '''simple docstring'''






   snake_case_					    =					tmp_path / "cache"
   snake_case_					    =					os.path.join(lowerCAmelCase_		,						"tmp.csv"							)
   snake_case_					    =					CsvDatasetReader({"train": csv_path}		,						cache_dir=lowerCAmelCase_							).read()
   with pytest.raises(lowerCAmelCase_							):
      CsvDatasetWriter(dataset["train"]		,						lowerCAmelCase_		,						num_proc=0							)




