

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


     def __init__(							self      :		Any      ,  UpperCAmelCase      :		Any      ,  UpperCAmelCase      :		Any=2      ,  UpperCAmelCase      :		List[str]=8      ,  UpperCAmelCase      :		List[str]=True      ,  UpperCAmelCase      :		int=False      ,  UpperCAmelCase      :		Tuple=False      ,  UpperCAmelCase      :		Dict=False      ,  UpperCAmelCase      :		List[str]=False      ,  UpperCAmelCase      :		Optional[Any]=9_9      ,  UpperCAmelCase      :		List[str]=3_2      ,  UpperCAmelCase      :		Any=3      ,  UpperCAmelCase      :		Dict=4      ,  UpperCAmelCase      :		List[str]=3_7      ,  UpperCAmelCase      :		List[str]=3_2      ,  UpperCAmelCase      :		Union[str, Any]=1_2_8      ,  UpperCAmelCase      :		Any=8      ,  UpperCAmelCase      :		Union[str, Any]=8      ,  UpperCAmelCase      :		Optional[int]=8      ,  UpperCAmelCase      :		Optional[Any]=1.0      ,  UpperCAmelCase      :		Any=True      ,  )							->     Dict:
      __lowerCAmelCase:       Union[str, Any]            =				parent
      __lowerCAmelCase:       Tuple            =				batch_size
      __lowerCAmelCase:       Dict            =				seq_length
      __lowerCAmelCase:       Optional[Any]            =				is_training
      __lowerCAmelCase:       Optional[int]            =				use_token_type_ids
      __lowerCAmelCase:       List[str]            =				use_input_mask
      __lowerCAmelCase:       str            =				use_labels
      __lowerCAmelCase:       Tuple            =				use_mc_token_ids
      __lowerCAmelCase:       Dict            =				vocab_size
      __lowerCAmelCase:       Any            =				hidden_size
      __lowerCAmelCase:       Union[str, Any]            =				num_hidden_layers
      __lowerCAmelCase:       List[Any]            =				num_attention_heads
      __lowerCAmelCase:       Union[str, Any]            =				intermediate_size
      __lowerCAmelCase:       Tuple            =				num_buckets
      __lowerCAmelCase:       Optional[Any]            =				max_distance
      __lowerCAmelCase:       List[Any]            =				prompt_length
      __lowerCAmelCase:       Optional[int]            =				prompt_types
      __lowerCAmelCase:       List[Any]            =				segment_types
      __lowerCAmelCase:       List[str]            =				init_std
      __lowerCAmelCase:       int            =				return_dict


     def   UpperCAmelCase  (							self      :		Optional[Any] )							->     Union[str, Any]:
      __lowerCAmelCase:       Optional[Any]            =				{}
      __lowerCAmelCase:       Optional[int]            =				ids_tensor([self.batch_size, self.seq_length]      ,  self.vocab_size ).type(torch.intaa )
      __lowerCAmelCase:       Any            =				False

      __lowerCAmelCase:       Any            =				self.get_config()

      return (config, input_ids)


     def   UpperCAmelCase  (							self      :		int )							->     Optional[Any]:
      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=UpperCAmelCase      ,  init_std=self.init_std      ,  return_dict=self.return_dict      ,  )


     def   UpperCAmelCase  (							self      :		Tuple      ,  UpperCAmelCase      :		Dict      ,  UpperCAmelCase      :		str      ,  *UpperCAmelCase      :		Union[str, Any] )							->     str:
      __lowerCAmelCase:       Union[str, Any]            =				CpmAntModel(config=UpperCAmelCase )
      model.to(UpperCAmelCase )
      model.eval()

      __lowerCAmelCase:       Any            =				model(**UpperCAmelCase ).last_hidden_state

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


     def   UpperCAmelCase  (							self      :		Tuple      ,  UpperCAmelCase      :		Union[str, Any]      ,  UpperCAmelCase      :		Any      ,  *UpperCAmelCase      :		Tuple )							->     int:
      __lowerCAmelCase:       Dict            =				CpmAntForCausalLM(UpperCAmelCase )
      model.to(UpperCAmelCase )
      __lowerCAmelCase:       List[Any]            =				input_ids['input_ids'].to(UpperCAmelCase )
      model.eval()

      __lowerCAmelCase:       Any            =				model(**UpperCAmelCase )
      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      :		Dict )							->     Optional[int]:
      __lowerCAmelCase     ,							__lowerCAmelCase:       Optional[Any]            =				self.prepare_config_and_inputs()
      return config, inputs_dict







@require_torch
class 							A_			(		snake_case__							,     snake_case__							,     unittest.TestCase ):
     _lowercase      :					Any       =     (CpmAntModel, CpmAntForCausalLM) if is_torch_available() else ()
     _lowercase      :					List[Any]       =     (
         {'feature-extraction': CpmAntModel, 'text-generation': CpmAntForCausalLM} if is_torch_available() else {}
     )

     _lowercase      :					int       =     False
     _lowercase      :					Union[str, Any]       =     False
     _lowercase      :					Optional[Any]       =     False
     _lowercase      :					Any       =     False
     _lowercase      :					str       =     False


     def   UpperCAmelCase  (							self      :		Dict )							->     List[str]:
      __lowerCAmelCase:       Dict            =				CpmAntModelTester(self )
      __lowerCAmelCase:       List[str]            =				ConfigTester(self      ,  config_class=UpperCAmelCase )


     def   UpperCAmelCase  (							self      :		int )							->     Any:
      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 )							->     List[Any]:
      unittest.skip('CPMAnt doesn\'t support input_embeds.' )(self.test_inputs_embeds )


     def   UpperCAmelCase  (							self      :		int )							->     Optional[Any]:
      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[Any] )							->     Dict:
      __lowerCAmelCase     ,							__lowerCAmelCase:       Union[str, Any]            =				self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_cpmant_model(UpperCAmelCase      ,  UpperCAmelCase )


     def   UpperCAmelCase  (							self      :		str )							->     int:
      __lowerCAmelCase     ,							__lowerCAmelCase:       Any            =				self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_lm_head_model(UpperCAmelCase      ,  UpperCAmelCase )







@require_torch
class 							A_			(		unittest.TestCase ):


     @tooslow
     def   UpperCAmelCase  (							self      :		List[str] )							->     List[Any]:
      __lowerCAmelCase:       Optional[int]            =				'今天天气真好！'
      __lowerCAmelCase:       List[str]            =				'openbmb/cpm-ant-10b'
      __lowerCAmelCase:       Optional[Any]            =				CpmAntModel.from_pretrained(UpperCAmelCase )
      __lowerCAmelCase:       Optional[Any]            =				CpmAntTokenizer.from_pretrained(UpperCAmelCase )
      __lowerCAmelCase:       List[Any]            =				tokenizer(UpperCAmelCase      ,  return_tensors='pt' )
      __lowerCAmelCase:       str            =				model(**UpperCAmelCase ).last_hidden_state

      __lowerCAmelCase:       Union[str, Any]            =				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]      ,  UpperCAmelCase      ,  atol=1E-2 ) )







@require_torch
class 							A_			(		unittest.TestCase ):


     @tooslow
     def   UpperCAmelCase  (							self      :		Tuple )							->     Optional[int]:
      __lowerCAmelCase:       Optional[int]            =				'今天天气真好！'
      __lowerCAmelCase:       Any            =				'openbmb/cpm-ant-10b'
      __lowerCAmelCase:       int            =				CpmAntForCausalLM.from_pretrained(UpperCAmelCase )
      __lowerCAmelCase:       Optional[Any]            =				CpmAntTokenizer.from_pretrained(UpperCAmelCase )
      __lowerCAmelCase:       Dict            =				tokenizer(UpperCAmelCase      ,  return_tensors='pt' )
      __lowerCAmelCase:       Dict            =				model(**UpperCAmelCase ).logits

      __lowerCAmelCase:       List[Any]            =				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]      ,  UpperCAmelCase      ,  atol=1E-2 ) )


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


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








from __future__ import annotations







def _a			(	SCREAMING_SNAKE_CASE		:      list[int]		)   ->   list[int]:

 """simple docstring"""






 if not list_of_ints:
  return []

 if min(SCREAMING_SNAKE_CASE		) < 0:
  raise ValueError('All numbers must be positive'		)

 __lowerCAmelCase:       Tuple            =				max(len(bin(SCREAMING_SNAKE_CASE		)[2:]		) for x in list_of_ints		)
 return _msd_radix_sort(SCREAMING_SNAKE_CASE   , SCREAMING_SNAKE_CASE		)







def _a			(	SCREAMING_SNAKE_CASE		:      list[int]   , SCREAMING_SNAKE_CASE		:      int		)   ->   list[int]:

 """simple docstring"""






 if bit_position == 0 or len(SCREAMING_SNAKE_CASE		) in [0, 1]:
  return list_of_ints

 __lowerCAmelCase:       List[str]            =				[]
 __lowerCAmelCase:       Any            =				[]
 # Split numbers based on bit at bit_position from the right
 for number in list_of_ints:
  if (number >> (bit_position - 1)) & 1:
   # number has a one at bit bit_position
   ones.append(SCREAMING_SNAKE_CASE		)
  else:
   # number has a zero at bit bit_position
   zeros.append(SCREAMING_SNAKE_CASE		)

    # recursively split both lists further
 __lowerCAmelCase:       Dict            =				_msd_radix_sort(SCREAMING_SNAKE_CASE   , bit_position - 1		)
 __lowerCAmelCase:       Union[str, Any]            =				_msd_radix_sort(SCREAMING_SNAKE_CASE   , bit_position - 1		)

 # recombine lists
 __lowerCAmelCase:       List[str]            =				zeros
 res.extend(SCREAMING_SNAKE_CASE		)

 return res







def _a			(	SCREAMING_SNAKE_CASE		:      list[int]		)   ->   List[str]:

 """simple docstring"""






 __lowerCAmelCase:       Optional[int]            =				len(SCREAMING_SNAKE_CASE		)
 if not list_of_ints or length == 1:
  return

 if min(SCREAMING_SNAKE_CASE		) < 0:
  raise ValueError('All numbers must be positive'		)

 __lowerCAmelCase:       Optional[int]            =				max(len(bin(SCREAMING_SNAKE_CASE		)[2:]		) for x in list_of_ints		)
 _msd_radix_sort_inplace(SCREAMING_SNAKE_CASE   , SCREAMING_SNAKE_CASE   , 0   , SCREAMING_SNAKE_CASE		)







def _a			(	SCREAMING_SNAKE_CASE		:      list[int]   , SCREAMING_SNAKE_CASE		:      int   , SCREAMING_SNAKE_CASE		:      int   , SCREAMING_SNAKE_CASE		:      int		)   ->   str:

 """simple docstring"""






 if bit_position == 0 or end_index - begin_index <= 1:
  return

 bit_position -= 1

 __lowerCAmelCase:       Optional[int]            =				begin_index
 __lowerCAmelCase:       Optional[Any]            =				end_index - 1
 while i <= j:
  __lowerCAmelCase:       Any            =				False
  if not (list_of_ints[i] >> bit_position) & 1:
   # found zero at the beginning
   i += 1
   __lowerCAmelCase:       Optional[Any]            =				True
  if (list_of_ints[j] >> bit_position) & 1:
   # found one at the end
   j -= 1
   __lowerCAmelCase:       List[str]            =				True

  if changed:
   continue

  __lowerCAmelCase     ,							__lowerCAmelCase:       Union[str, Any]            =				list_of_ints[j], list_of_ints[i]
  j -= 1
  if j != i:
   i += 1

 _msd_radix_sort_inplace(SCREAMING_SNAKE_CASE   , SCREAMING_SNAKE_CASE   , SCREAMING_SNAKE_CASE   , SCREAMING_SNAKE_CASE		)
 _msd_radix_sort_inplace(SCREAMING_SNAKE_CASE   , SCREAMING_SNAKE_CASE   , SCREAMING_SNAKE_CASE   , SCREAMING_SNAKE_CASE		)


if __name__ == "__main__":
  import doctest

  doctest.testmod()





