


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


      def __init__(							self							,						SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__=2							,						SCREAMING_SNAKE_CASE__=8							,						SCREAMING_SNAKE_CASE__=True							,						SCREAMING_SNAKE_CASE__=False							,						SCREAMING_SNAKE_CASE__=False							,						SCREAMING_SNAKE_CASE__=False							,						SCREAMING_SNAKE_CASE__=False							,						SCREAMING_SNAKE_CASE__=99							,						SCREAMING_SNAKE_CASE__=32							,						SCREAMING_SNAKE_CASE__=3							,						SCREAMING_SNAKE_CASE__=4							,						SCREAMING_SNAKE_CASE__=37							,						SCREAMING_SNAKE_CASE__=32							,						SCREAMING_SNAKE_CASE__=128							,						SCREAMING_SNAKE_CASE__=8							,						SCREAMING_SNAKE_CASE__=8							,						SCREAMING_SNAKE_CASE__=8							,						SCREAMING_SNAKE_CASE__=1.0							,						SCREAMING_SNAKE_CASE__=True							,						):
            lowercase			:   Tuple					   =		parent
            lowercase			:   str					   =		batch_size
            lowercase			:   Dict					   =		seq_length
            lowercase			:   List[str]					   =		is_training
            lowercase			:   Union[str, Any]					   =		use_token_type_ids
            lowercase			:   List[str]					   =		use_input_mask
            lowercase			:   str					   =		use_labels
            lowercase			:   str					   =		use_mc_token_ids
            lowercase			:   List[str]					   =		vocab_size
            lowercase			:   str					   =		hidden_size
            lowercase			:   Union[str, Any]					   =		num_hidden_layers
            lowercase			:   List[str]					   =		num_attention_heads
            lowercase			:   Optional[int]					   =		intermediate_size
            lowercase			:   List[Any]					   =		num_buckets
            lowercase			:   Any					   =		max_distance
            lowercase			:   Tuple					   =		prompt_length
            lowercase			:   List[str]					   =		prompt_types
            lowercase			:   int					   =		segment_types
            lowercase			:   Tuple					   =		init_std
            lowercase			:   Any					   =		return_dict


      def        __lowerCamelCase			(							self		):
            lowercase			:   Optional[Any]					   =		{}
            lowercase			:   Tuple					   =		ids_tensor([self.batch_size, self.seq_length]							,						self.vocab_size		).type(torch.intaa		)
            lowercase			:   int					   =		False

            lowercase			:   Optional[int]					   =		self.get_config()

            return (config, input_ids)


      def        __lowerCamelCase			(							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=SCREAMING_SNAKE_CASE__							,						init_std=self.init_std							,						return_dict=self.return_dict							,						)


      def        __lowerCamelCase			(							self							,						SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__							,						*SCREAMING_SNAKE_CASE__		):
            lowercase			:   Union[str, Any]					   =		CpmAntModel(config=SCREAMING_SNAKE_CASE__		)
            model.to(SCREAMING_SNAKE_CASE__		)
            model.eval()

            lowercase			:   Tuple					   =		model(**SCREAMING_SNAKE_CASE__		).last_hidden_state

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


      def        __lowerCamelCase			(							self							,						SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__							,						*SCREAMING_SNAKE_CASE__		):
            lowercase			:   Dict					   =		CpmAntForCausalLM(SCREAMING_SNAKE_CASE__		)
            model.to(SCREAMING_SNAKE_CASE__		)
            lowercase			:   List[str]					   =		input_ids['''input_ids'''].to(SCREAMING_SNAKE_CASE__		)
            model.eval()

            lowercase			:   int					   =		model(**SCREAMING_SNAKE_CASE__		)
            self.parent.assertEqual(
                model_output.logits.shape							,						(self.batch_size, self.seq_length, config.vocab_size + config.prompt_types * config.prompt_length)							,						)






      def        __lowerCamelCase			(							self		):
            lowercase   ,			lowercase			:   Dict					   =		self.prepare_config_and_inputs()
            return config, inputs_dict







@require_torch
class 				__SCREAMING_SNAKE_CASE   (			A__  ,						A__  ,						unittest.TestCase   ):
      A				:							Union[str, Any]    							=     (CpmAntModel, CpmAntForCausalLM) if is_torch_available() else ()
      A				:							Dict    							=     (
          {'feature-extraction': CpmAntModel, 'text-generation': CpmAntForCausalLM} if is_torch_available() else {}
      )

      A				:							Optional[Any]    							=     False
      A				:							Dict    							=     False
      A				:							int    							=     False
      A				:							Optional[int]    							=     False
      A				:							Union[str, Any]    							=     False


      def        __lowerCamelCase			(							self		):
            lowercase			:   Optional[int]					   =		CpmAntModelTester(self		)
            lowercase			:   Dict					   =		ConfigTester(self							,						config_class=SCREAMING_SNAKE_CASE__		)


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


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


      def        __lowerCamelCase			(							self		):
            lowercase   ,			lowercase			:   Optional[Any]					   =		self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_cpmant_model(SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__		)






      def        __lowerCamelCase			(							self		):
            lowercase   ,			lowercase			:   Union[str, Any]					   =		self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_lm_head_model(SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__		)







@require_torch
class 				__SCREAMING_SNAKE_CASE   (			unittest.TestCase   ):


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

            lowercase			:   int					   =		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]							,						SCREAMING_SNAKE_CASE__							,						atol=1E-2		)		)







@require_torch
class 				__SCREAMING_SNAKE_CASE   (			unittest.TestCase   ):


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

            lowercase			:   Tuple					   =		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]							,						SCREAMING_SNAKE_CASE__							,						atol=1E-2		)		)


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






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





from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


__a   				=      {'''configuration_yoso''': ['''YOSO_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''YosoConfig''']}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __a   				=      [
      '''YOSO_PRETRAINED_MODEL_ARCHIVE_LIST''',
      '''YosoForMaskedLM''',
      '''YosoForMultipleChoice''',
      '''YosoForQuestionAnswering''',
      '''YosoForSequenceClassification''',
      '''YosoForTokenClassification''',
      '''YosoLayer''',
      '''YosoModel''',
      '''YosoPreTrainedModel''',
  ]


if TYPE_CHECKING:
  from .configuration_yoso import YOSO_PRETRAINED_CONFIG_ARCHIVE_MAP, YosoConfig

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_yoso import (
        YOSO_PRETRAINED_MODEL_ARCHIVE_LIST,
        YosoForMaskedLM,
        YosoForMultipleChoice,
        YosoForQuestionAnswering,
        YosoForSequenceClassification,
        YosoForTokenClassification,
        YosoLayer,
        YosoModel,
        YosoPreTrainedModel,
    )


else:
  import sys

  __a   				=      _LazyModule(__name__, globals()['''__file__'''], _import_structure)

