





"""simple docstring"""






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






       def __init__( self     ,	a__     ,	a__=2     ,	a__=8     ,	a__=True     ,	a__=False     ,	a__=False     ,	a__=False     ,	a__=False     ,	a__=99     ,	a__=32     ,	a__=3     ,	a__=4     ,	a__=37     ,	a__=32     ,	a__=1_28     ,	a__=8     ,	a__=8     ,	a__=8     ,	a__=1.0     ,	a__=True     ,	):
              _UpperCAmelCase			     =  parent
              _UpperCAmelCase			     =  batch_size
              _UpperCAmelCase			     =  seq_length
              _UpperCAmelCase			     =  is_training
              _UpperCAmelCase			     =  use_token_type_ids
              _UpperCAmelCase			     =  use_input_mask
              _UpperCAmelCase			     =  use_labels
              _UpperCAmelCase			     =  use_mc_token_ids
              _UpperCAmelCase			     =  vocab_size
              _UpperCAmelCase			     =  hidden_size
              _UpperCAmelCase			     =  num_hidden_layers
              _UpperCAmelCase			     =  num_attention_heads
              _UpperCAmelCase			     =  intermediate_size
              _UpperCAmelCase			     =  num_buckets
              _UpperCAmelCase			     =  max_distance
              _UpperCAmelCase			     =  prompt_length
              _UpperCAmelCase			     =  prompt_types
              _UpperCAmelCase			     =  segment_types
              _UpperCAmelCase			     =  init_std
              _UpperCAmelCase			     =  return_dict






       def 				__A			( self  ):
              _UpperCAmelCase			     =  {}
              _UpperCAmelCase			     =  ids_tensor([self.batch_size, self.seq_length]     ,	self.vocab_size  ).type(torch.intaa  )
              _UpperCAmelCase			     =  False

              _UpperCAmelCase			     =  self.get_config()

              return (config, input_ids)






       def 				__A			( self  ):
              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=a__     ,	init_std=self.init_std     ,	return_dict=self.return_dict     ,	)






       def 				__A			( self     ,	a__     ,	a__     ,	*a__  ):
              _UpperCAmelCase			     =  CpmAntModel(config=a__  )
              model.to(a__  )
              model.eval()

              _UpperCAmelCase			     =  model(**a__  ).last_hidden_state

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






       def 				__A			( self     ,	a__     ,	a__     ,	*a__  ):
              _UpperCAmelCase			     =  CpmAntForCausalLM(a__  )
              model.to(a__  )
              _UpperCAmelCase			     =  input_ids["""input_ids"""].to(a__  )
              model.eval()

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






       def 				__A			( self  ):
              _UpperCAmelCase			     =  self.prepare_config_and_inputs()
              return config, inputs_dict



@require_torch
class 						lowerCAmelCase       (   snake_case      ,						snake_case      ,						unittest.TestCase				):
       lowerCAmelCase__    			=	(CpmAntModel, CpmAntForCausalLM) if is_torch_available() else ()
       lowerCAmelCase__    			=	(
           {"""feature-extraction""": CpmAntModel, """text-generation""": CpmAntForCausalLM} if is_torch_available() else {}
       )

       lowerCAmelCase__    			=	False
       lowerCAmelCase__    			=	False
       lowerCAmelCase__    			=	False
       lowerCAmelCase__    			=	False
       lowerCAmelCase__    			=	False






       def 				__A			( self  ):
              _UpperCAmelCase			     =  CpmAntModelTester(self  )
              _UpperCAmelCase			     =  ConfigTester(self     ,	config_class=a__  )






       def 				__A			( self  ):
              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 				__A			( self  ):
              unittest.skip('CPMAnt doesn\'t support input_embeds.'  )(self.test_inputs_embeds  )






       def 				__A			( self  ):
              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.\\n                 So is attentions. We strongly recommand you use loss to tune model.'  )(self.test_retain_grad_hidden_states_attentions  )






       def 				__A			( self  ):
              _UpperCAmelCase			     =  self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_cpmant_model(a__     ,	a__  )






       def 				__A			( self  ):
              _UpperCAmelCase			     =  self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_lm_head_model(a__     ,	a__  )



@require_torch
class 						lowerCAmelCase       (   unittest.TestCase				):






       @tooslow
       def 				__A			( self  ):
              _UpperCAmelCase			     =  """今天天气真好！"""
              _UpperCAmelCase			     =  """openbmb/cpm-ant-10b"""
              _UpperCAmelCase			     =  CpmAntModel.from_pretrained(a__  )
              _UpperCAmelCase			     =  CpmAntTokenizer.from_pretrained(a__  )
              _UpperCAmelCase			     =  tokenizer(a__     ,	return_tensors='pt'  )
              _UpperCAmelCase			     =  model(**a__  ).last_hidden_state

              _UpperCAmelCase			     =  torch.tensor(
                  [[[6.1_708, 5.9_244, 1.0_835], [6.5_207, 6.2_893, -11.3_324], [-1.0_107, -0.0_576, -5.9_577]]]     ,	)
              self.assertTrue(torch.allclose(hidden_states[:, :3, :3]     ,	a__     ,	atol=1E-2  )  )



@require_torch
class 						lowerCAmelCase       (   unittest.TestCase				):






       @tooslow
       def 				__A			( self  ):
              _UpperCAmelCase			     =  """今天天气真好！"""
              _UpperCAmelCase			     =  """openbmb/cpm-ant-10b"""
              _UpperCAmelCase			     =  CpmAntForCausalLM.from_pretrained(a__  )
              _UpperCAmelCase			     =  CpmAntTokenizer.from_pretrained(a__  )
              _UpperCAmelCase			     =  tokenizer(a__     ,	return_tensors='pt'  )
              _UpperCAmelCase			     =  model(**a__  ).logits

              _UpperCAmelCase			     =  torch.tensor(
                  [[[-6.4_267, -6.4_083, -6.3_958], [-5.8_802, -5.9_447, -5.7_811], [-5.3_896, -5.4_820, -5.4_295]]]     ,	)
              self.assertTrue(torch.allclose(hidden_states[:, :3, :3]     ,	a__     ,	atol=1E-2  )  )






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






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










'''simple docstring'''
from typing import Dict, Optional, Tuple

import torch



def 				lowerCAmelCase__	(       UpperCAmelCase     ):






       """simple docstring"""





       snake_case__					:	Optional[int]          =					boundaries[1] - boundaries[0]
       snake_case__					:	Optional[Any]          =					boundaries + step / 2
       snake_case__					:	str          =					torch.cat([bin_centers, (bin_centers[-1] + step).unsqueeze(-1     )]    ,       dim=0     )
       return bin_centers



def 				lowerCAmelCase__	(       UpperCAmelCase    ,       UpperCAmelCase    ,       ):






       """simple docstring"""





       snake_case__					:	str          =					_calculate_bin_centers(UpperCAmelCase     )
       return (
           torch.sum(aligned_distance_error_probs * bin_centers    ,       dim=-1     ),
           bin_centers[-1],
       )



def 				lowerCAmelCase__	(       UpperCAmelCase    ,       UpperCAmelCase = 31    ,       UpperCAmelCase = 64    ,       **UpperCAmelCase    ,       ):






       """simple docstring"""





       snake_case__					:	Dict          =					torch.linspace(0    ,       UpperCAmelCase    ,       steps=(no_bins - 1)    ,       device=logits.device     )

       snake_case__					:	str          =					torch.nn.functional.softmax(UpperCAmelCase    ,       dim=-1     )
       snake_case__       ,   snake_case__					:	List[str]          =					_calculate_expected_aligned_error(
           alignment_confidence_breaks=UpperCAmelCase    ,       aligned_distance_error_probs=UpperCAmelCase    ,       )

       return {
           "aligned_confidence_probs": aligned_confidence_probs,
           "predicted_aligned_error": predicted_aligned_error,
           "max_predicted_aligned_error": max_predicted_aligned_error,
       }






def 				lowerCAmelCase__	(       UpperCAmelCase    ,       UpperCAmelCase = None    ,       UpperCAmelCase = 31    ,       UpperCAmelCase = 64    ,       UpperCAmelCase = 1E-8    ,       **UpperCAmelCase    ,       ):






       """simple docstring"""





       if residue_weights is None:
              snake_case__					:	List[Any]          =					logits.new_ones(logits.shape[-2]     )

       snake_case__					:	Optional[Any]          =					torch.linspace(0    ,       UpperCAmelCase    ,       steps=(no_bins - 1)    ,       device=logits.device     )

       snake_case__					:	int          =					_calculate_bin_centers(UpperCAmelCase     )
       torch.sum(UpperCAmelCase     )
       snake_case__					:	Dict          =					logits.shape[-2]
       snake_case__					:	Optional[Any]          =					max(UpperCAmelCase    ,       19     )

       snake_case__					:	Optional[int]          =					1.2_4 * (clipped_n - 15) ** (1.0 / 3) - 1.8

       snake_case__					:	Tuple          =					torch.nn.functional.softmax(UpperCAmelCase    ,       dim=-1     )

       snake_case__					:	int          =					1.0 / (1 + (bin_centers**2) / (da**2))
       snake_case__					:	Optional[int]          =					torch.sum(probs * tm_per_bin    ,       dim=-1     )

       snake_case__					:	Optional[Any]          =					residue_weights / (eps + residue_weights.sum())
       snake_case__					:	Any          =					torch.sum(predicted_tm_term * normed_residue_mask    ,       dim=-1     )

       snake_case__					:	Dict          =					per_alignment * residue_weights

       snake_case__					:	List[str]          =					(weighted == torch.max(UpperCAmelCase     )).nonzero()[0]
       return per_alignment[tuple(UpperCAmelCase     )]




