

import unittest

from transformers.testing_utils import is_torch_available, require_torch, tooslow

from ...generation.test_utils import torch_device
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch

       from transformers import (
           CpmAntConfig,
           CpmAntForCausalLM,
           CpmAntModel,
           CpmAntTokenizer,
       )







@require_torch
class 		__magic_name__     :

   '''simple docstring'''




   def __init__(							self,       lowercase_,       lowercase_=2,       lowercase_=8,       lowercase_=True,       lowercase_=False,       lowercase_=False,       lowercase_=False,       lowercase_=False,       lowercase_=99,       lowercase_=32,       lowercase_=3,       lowercase_=4,       lowercase_=37,       lowercase_=32,       lowercase_=128,       lowercase_=8,       lowercase_=8,       lowercase_=8,       lowercase_=1.0,       lowercase_=True,       )				->       Tuple:






     """simple docstring"""






     a__    =parent
     a__    =batch_size
     a__    =seq_length
     a__    =is_training
     a__    =use_token_type_ids
     a__    =use_input_mask
     a__    =use_labels
     a__    =use_mc_token_ids
     a__    =vocab_size
     a__    =hidden_size
     a__    =num_hidden_layers
     a__    =num_attention_heads
     a__    =intermediate_size
     a__    =num_buckets
     a__    =max_distance
     a__    =prompt_length
     a__    =prompt_types
     a__    =segment_types
     a__    =init_std
     a__    =return_dict



   def 					_UpperCAmelCase					(							self			)				->       int:






     """simple docstring"""






     a__    ={}
     a__    =ids_tensor([self.batch_size, self.seq_length],       self.vocab_size			).type(torch.intaa			)
     a__    =False

     a__    =self.get_config()

     return (config, input_ids)



   def 					_UpperCAmelCase					(							self			)				->       Dict:






     """simple docstring"""






     return CpmAntConfig(
         vocab_size=self.vocab_size,       hidden_size=self.hidden_size,       num_hidden_layers=self.num_hidden_layers,       num_attention_heads=self.num_attention_heads,       dim_ff=self.intermediate_size,       position_bias_num_buckets=self.num_buckets,       position_bias_max_distance=self.max_distance,       prompt_types=self.prompt_types,       prompt_length=self.prompt_length,       segment_types=self.segment_types,       use_cache=lowercase_,       init_std=self.init_std,       return_dict=self.return_dict,       )



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       *lowercase_			)				->       int:






     """simple docstring"""






     a__    =CpmAntModel(config=lowercase_			)
     model.to(lowercase_			)
     model.eval()

     a__    =model(**lowercase_			).last_hidden_state

     self.parent.assertEqual(hidden_states.shape,       (self.batch_size, self.seq_length, config.hidden_size)			)



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       *lowercase_			)				->       Tuple:






     """simple docstring"""






     a__    =CpmAntForCausalLM(lowercase_			)
     model.to(lowercase_			)
     a__    =input_ids['input_ids'].to(lowercase_			)
     model.eval()

     a__    =model(**lowercase_			)
     self.parent.assertEqual(
         model_output.logits.shape,       (self.batch_size, self.seq_length, config.vocab_size + config.prompt_types * config.prompt_length),       )



   def 					_UpperCAmelCase					(							self			)				->       List[str]:






     """simple docstring"""






     a__    =self.prepare_config_and_inputs()
     return config, inputs_dict







@require_torch
class 		__magic_name__     (				snake_case__						,			snake_case__						,			unittest.TestCase				):

   '''simple docstring'''




   lowerCamelCase__					:      List[str]        =  (CpmAntModel, CpmAntForCausalLM) if is_torch_available() else ()
   lowerCamelCase__					:      Optional[int]        =  (
       {'feature-extraction': CpmAntModel, 'text-generation': CpmAntForCausalLM} if is_torch_available() else {}
   )

   lowerCamelCase__					:      Optional[int]        =  False
   lowerCamelCase__					:      Tuple        =  False
   lowerCamelCase__					:      Dict        =  False
   lowerCamelCase__					:      str        =  False
   lowerCamelCase__					:      str        =  False



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






     """simple docstring"""






     a__    =CpmAntModelTester(self			)
     a__    =ConfigTester(self,       config_class=lowercase_			)



   def 					_UpperCAmelCase					(							self			)				->       List[str]:






     """simple docstring"""






     self.config_tester.create_and_test_config_common_properties()
     self.config_tester.create_and_test_config_to_json_string()
     self.config_tester.create_and_test_config_to_json_file()
     self.config_tester.create_and_test_config_from_and_save_pretrained()
     self.config_tester.check_config_can_be_init_without_params()
     self.config_tester.check_config_arguments_init()



   def 					_UpperCAmelCase					(							self			)				->       Dict:






     """simple docstring"""






     unittest.skip('''CPMAnt doesn\'t support input_embeds.'''			)(self.test_inputs_embeds			)



   def 					_UpperCAmelCase					(							self			)				->       Tuple:






     """simple docstring"""






     unittest.skip(
         '''CPMAnt doesn\'t support retain grad in hidden_states or attentions, because prompt management will peel off the output.hidden_states from graph.\
                 So is attentions. We strongly recommand you use loss to tune model.'''			)(self.test_retain_grad_hidden_states_attentions			)



   def 					_UpperCAmelCase					(							self			)				->       List[str]:






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_cpmant_model(lowercase_,       lowercase_			)



   def 					_UpperCAmelCase					(							self			)				->       int:






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_lm_head_model(lowercase_,       lowercase_			)







@require_torch
class 		__magic_name__     (				unittest.TestCase				):

   '''simple docstring'''




   @tooslow
   def 					_UpperCAmelCase					(							self			)				->       Any:






     """simple docstring"""






     a__    ='今天天气真好！'
     a__    ='openbmb/cpm-ant-10b'
     a__    =CpmAntModel.from_pretrained(lowercase_			)
     a__    =CpmAntTokenizer.from_pretrained(lowercase_			)
     a__    =tokenizer(lowercase_,       return_tensors='''pt'''			)
     a__    =model(**lowercase_			).last_hidden_state

     a__    =torch.tensor(
         [[[6.1708, 5.9244, 1.0835], [6.5207, 6.2893, -11.3324], [-1.0107, -0.0576, -5.9577]]],       )
     self.assertTrue(torch.allclose(hidden_states[:, :3, :3],       lowercase_,       atol=1E-2			)			)







@require_torch
class 		__magic_name__     (				unittest.TestCase				):

   '''simple docstring'''




   @tooslow
   def 					_UpperCAmelCase					(							self			)				->       int:






     """simple docstring"""






     a__    ='今天天气真好！'
     a__    ='openbmb/cpm-ant-10b'
     a__    =CpmAntForCausalLM.from_pretrained(lowercase_			)
     a__    =CpmAntTokenizer.from_pretrained(lowercase_			)
     a__    =tokenizer(lowercase_,       return_tensors='''pt'''			)
     a__    =model(**lowercase_			).logits

     a__    =torch.tensor(
         [[[-6.4267, -6.4083, -6.3958], [-5.8802, -5.9447, -5.7811], [-5.3896, -5.4820, -5.4295]]],       )
     self.assertTrue(torch.allclose(hidden_states[:, :3, :3],       lowercase_,       atol=1E-2			)			)



   @tooslow
   def 					_UpperCAmelCase					(							self			)				->       Any:






     """simple docstring"""






     a__    ='openbmb/cpm-ant-10b'
     a__    =CpmAntForCausalLM.from_pretrained(lowercase_			)
     a__    =CpmAntTokenizer.from_pretrained(lowercase_			)
     a__    ='今天天气不错，'
     a__    ='今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的'
     a__    =tokenizer(lowercase_,       return_tensors='''pt'''			)
     a__    =model.generate(**lowercase_			)
     a__    =tokenizer.batch_decode(lowercase_			)
     self.assertEqual(lowercase_,       lowercase_			)



   @tooslow
   def 					_UpperCAmelCase					(							self			)				->       Any:






     """simple docstring"""






     a__    ='openbmb/cpm-ant-10b'
     a__    =CpmAntForCausalLM.from_pretrained(lowercase_			)
     a__    =CpmAntTokenizer.from_pretrained(lowercase_			)
     a__    =['今天天气不错，', '新年快乐，万事如意！']
     a__    =[
         '今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的',
         '新年快乐，万事如意！在这辞旧迎新的美好时刻，我谨代表《农村新技术》杂志社全体同仁，向一直以来关心、支持《农村新技术》杂志发展的各级领导、各界朋友和广大读者致以最诚挚的',
     ]
     a__    =tokenizer(lowercase_,       return_tensors='''pt''',       padding=lowercase_			)
     a__    =model.generate(**lowercase_			)
     a__    =tokenizer.batch_decode(lowercase_			)
     self.assertEqual(lowercase_,       lowercase_			)



from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import TensorType, is_torch_available, logging


_a									=  logging.get_logger(__name__)

_a									=  {
    '''facebook/m2m100_418M''': '''https://huggingface.co/facebook/m2m100_418M/resolve/main/config.json''',
    # See all M2M100 models at https://huggingface.co/models?filter=m2m_100
}






class 				A_					(						snake_case__   ):
     _lowercase     :	List[Any]        =      'm2m_100'
     _lowercase     :	str        =      ['past_key_values']
     _lowercase     :	List[Any]        =      {'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}


     def __init__(				self     :						int     ,    UpperCAmelCase     :						Any=1_2_8_1_1_2     ,    UpperCAmelCase     :						Optional[int]=1_0_2_4     ,    UpperCAmelCase     :						Tuple=1_2     ,    UpperCAmelCase     :						List[str]=4_0_9_6     ,    UpperCAmelCase     :						Dict=1_6     ,    UpperCAmelCase     :						int=1_2     ,    UpperCAmelCase     :						int=4_0_9_6     ,    UpperCAmelCase     :						List[str]=1_6     ,    UpperCAmelCase     :						Optional[Any]=0.05     ,    UpperCAmelCase     :						Dict=0.05     ,    UpperCAmelCase     :						Union[str, Any]=True     ,    UpperCAmelCase     :						Dict=True     ,    UpperCAmelCase     :						List[str]="relu"     ,    UpperCAmelCase     :						Dict=1_0_2_4     ,    UpperCAmelCase     :						str=0.1     ,    UpperCAmelCase     :						Dict=0.1     ,    UpperCAmelCase     :						List[str]=0.0     ,    UpperCAmelCase     :						Union[str, Any]=0.02     ,    UpperCAmelCase     :						Optional[Any]=2     ,    UpperCAmelCase     :						str=True     ,    UpperCAmelCase     :						Optional[int]=1     ,    UpperCAmelCase     :						List[Any]=0     ,    UpperCAmelCase     :						Union[str, Any]=2     ,    **UpperCAmelCase     :						Union[str, Any]     ,    )				->      Tuple:
      __lowerCAmelCase:  Optional[Any]        = vocab_size
      __lowerCAmelCase:  Dict        = max_position_embeddings
      __lowerCAmelCase:  Tuple        = d_model
      __lowerCAmelCase:  List[str]        = encoder_ffn_dim
      __lowerCAmelCase:  Any        = encoder_layers
      __lowerCAmelCase:  int        = encoder_attention_heads
      __lowerCAmelCase:  str        = decoder_ffn_dim
      __lowerCAmelCase:  Tuple        = decoder_layers
      __lowerCAmelCase:  Tuple        = decoder_attention_heads
      __lowerCAmelCase:  Any        = dropout
      __lowerCAmelCase:  Optional[int]        = attention_dropout
      __lowerCAmelCase:  Tuple        = activation_dropout
      __lowerCAmelCase:  Optional[Any]        = activation_function
      __lowerCAmelCase:  Union[str, Any]        = init_std
      __lowerCAmelCase:  List[str]        = encoder_layerdrop
      __lowerCAmelCase:  Optional[Any]        = decoder_layerdrop
      __lowerCAmelCase:  Dict        = use_cache
      __lowerCAmelCase:  Any        = encoder_layers
      __lowerCAmelCase:  Tuple        = scale_embedding  # scale factor will be sqrt(d_model) if True

      super().__init__(
          pad_token_id=UpperCAmelCase     ,    bos_token_id=UpperCAmelCase     ,    eos_token_id=UpperCAmelCase     ,    is_encoder_decoder=UpperCAmelCase     ,    decoder_start_token_id=UpperCAmelCase     ,    **UpperCAmelCase     ,    )







class 				A_					(						snake_case__   ):


     @property
     def     UpperCAmelCase    (				self     :						List[Any]   )				->      Mapping[str, Mapping[int, str]]:
      __lowerCAmelCase:  Union[str, Any]        = OrderedDict(
          [
              ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
              ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
          ]   )

      if self.use_past:
       __lowerCAmelCase:  str        = {0: 'batch'}
       __lowerCAmelCase:  Union[str, Any]        = {0: 'batch', 1: 'past_decoder_sequence + sequence'}
      else:
       __lowerCAmelCase:  Any        = {0: 'batch', 1: 'decoder_sequence'}
       __lowerCAmelCase:  Dict        = {0: 'batch', 1: 'decoder_sequence'}

      if self.use_past:
       self.fill_with_past_key_values_(UpperCAmelCase     ,    direction='inputs'   )
      return common_inputs


     def     UpperCAmelCase    (				self     :						Tuple     ,    UpperCAmelCase     :						PreTrainedTokenizer     ,    UpperCAmelCase     :						int = -1     ,    UpperCAmelCase     :						int = -1     ,    UpperCAmelCase     :						bool = False     ,    UpperCAmelCase     :						Optional[TensorType] = None     ,    )				->      Mapping[str, Any]:
      # Copied from OnnxConfig.generate_dummy_inputs
      # Did not use super(OnnxConfigWithPast, self).generate_dummy_inputs for code clarity.
      # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
      __lowerCAmelCase:  List[str]        = compute_effective_axis_dimension(
          UpperCAmelCase     ,    fixed_dimension=OnnxConfig.default_fixed_batch     ,    num_token_to_add=0   )

      # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
      __lowerCAmelCase:  Tuple        = tokenizer.num_special_tokens_to_add(UpperCAmelCase   )
      __lowerCAmelCase:  List[str]        = compute_effective_axis_dimension(
          UpperCAmelCase     ,    fixed_dimension=OnnxConfig.default_fixed_sequence     ,    num_token_to_add=UpperCAmelCase   )

      # Generate dummy inputs according to compute batch and sequence
      __lowerCAmelCase:  int        = [' '.join([tokenizer.unk_token]   ) * seq_length] * batch_size
      __lowerCAmelCase:  Optional[int]        = dict(tokenizer(UpperCAmelCase     ,    return_tensors=UpperCAmelCase   )   )
      return common_inputs


     def     UpperCAmelCase    (				self     :						str     ,    UpperCAmelCase     :						PreTrainedTokenizer     ,    UpperCAmelCase     :						int = -1     ,    UpperCAmelCase     :						int = -1     ,    UpperCAmelCase     :						bool = False     ,    UpperCAmelCase     :						Optional[TensorType] = None     ,    )				->      Mapping[str, Any]:
      __lowerCAmelCase:  Optional[int]        = self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
          UpperCAmelCase     ,    UpperCAmelCase     ,    UpperCAmelCase     ,    UpperCAmelCase     ,    UpperCAmelCase   )

      # Generate decoder inputs
      __lowerCAmelCase:  Any        = seq_length if not self.use_past else 1
      __lowerCAmelCase:  Union[str, Any]        = self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
          UpperCAmelCase     ,    UpperCAmelCase     ,    UpperCAmelCase     ,    UpperCAmelCase     ,    UpperCAmelCase   )
      __lowerCAmelCase:  List[Any]        = {F'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
      __lowerCAmelCase:  List[str]        = dict(**UpperCAmelCase     ,    **UpperCAmelCase   )

      if self.use_past:
       if not is_torch_available():
        raise ValueError('Cannot generate dummy past_keys inputs without PyTorch installed.'   )
       else:
        import torch
       __lowerCAmelCase      ,				__lowerCAmelCase:  List[Any]        = common_inputs['input_ids'].shape
       __lowerCAmelCase:  int        = common_inputs['decoder_input_ids'].shape[1]
       __lowerCAmelCase      ,				__lowerCAmelCase:  str        = self.num_attention_heads
       __lowerCAmelCase:  str        = (
           batch,
           num_encoder_attention_heads,
           encoder_seq_length,
           self._config.hidden_size // num_encoder_attention_heads,
       )
       __lowerCAmelCase:  Optional[int]        = decoder_seq_length + 3
       __lowerCAmelCase:  List[str]        = (
           batch,
           num_decoder_attention_heads,
           decoder_past_length,
           self._config.hidden_size // num_decoder_attention_heads,
       )

       __lowerCAmelCase:  Dict        = torch.cat(
           [common_inputs['decoder_attention_mask'], torch.ones(UpperCAmelCase     ,    UpperCAmelCase   )]     ,    dim=1   )

       __lowerCAmelCase:  str        = []
       # If the number of encoder and decoder layers are present in the model configuration, both are considered
       __lowerCAmelCase      ,				__lowerCAmelCase:  str        = self.num_layers
       __lowerCAmelCase:  Tuple        = min(UpperCAmelCase     ,    UpperCAmelCase   )
       __lowerCAmelCase:  List[Any]        = max(UpperCAmelCase     ,    UpperCAmelCase   ) - min_num_layers
       __lowerCAmelCase:  Any        = 'encoder' if num_encoder_layers > num_decoder_layers else 'decoder'

       for _ in range(UpperCAmelCase   ):
        common_inputs["past_key_values"].append(
            (
                torch.zeros(UpperCAmelCase   ),
                torch.zeros(UpperCAmelCase   ),
                torch.zeros(UpperCAmelCase   ),
                torch.zeros(UpperCAmelCase   ),
            )   )
       # TODO: test this.
       __lowerCAmelCase:  List[Any]        = encoder_shape if remaining_side_name == 'encoder' else decoder_shape
       for _ in range(UpperCAmelCase     ,    UpperCAmelCase   ):
        common_inputs["past_key_values"].append((torch.zeros(UpperCAmelCase   ), torch.zeros(UpperCAmelCase   ))   )
      return common_inputs


     _lowercase     :	int        =      _generate_dummy_inputs_for_default_and_seqaseq_lm





