"""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 					UpperCAmelCase_     :


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


      def 		_UpperCAmelCase	(       self							)	-> Optional[Any]:
             lowercase__      :		str          =     {}
             lowercase__      :		List[Any]          =     ids_tensor([self.batch_size, self.seq_length]		,      self.vocab_size							).type(torch.intaa							)
             lowercase__      :		str          =     False

             lowercase__      :		List[Any]          =     self.get_config()

             return (config, input_ids)


      def 		_UpperCAmelCase	(       self							)	-> Optional[int]:
             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 		_UpperCAmelCase	(       self		,      a		,      a		,      *a							)	-> Union[str, Any]:
             lowercase__      :		str          =     CpmAntModel(config=a							)
             model.to(a							)
             model.eval()

             lowercase__      :		Dict          =     model(**a							).last_hidden_state

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


      def 		_UpperCAmelCase	(       self		,      a		,      a		,      *a							)	-> Tuple:
             lowercase__      :		Optional[int]          =     CpmAntForCausalLM(a							)
             model.to(a							)
             lowercase__      :		Any          =     input_ids["""input_ids"""].to(a							)
             model.eval()

             lowercase__      :		List[Any]          =     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 		_UpperCAmelCase	(       self							)	-> Optional[int]:
             lowercase__      :		Dict          =     self.prepare_config_and_inputs()
             return config, inputs_dict





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

      lowerCamelCase__		:   Tuple			  =  False
      lowerCamelCase__		:   List[Any]			  =  False
      lowerCamelCase__		:   Optional[int]			  =  False
      lowerCamelCase__		:   Dict			  =  False
      lowerCamelCase__		:   Union[str, Any]			  =  False


      def 		_UpperCAmelCase	(       self							)	-> Any:
             lowercase__      :		str          =     CpmAntModelTester(self							)
             lowercase__      :		List[Any]          =     ConfigTester(self		,      config_class=a							)


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


      def 		_UpperCAmelCase	(       self							)	-> int:
             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 		_UpperCAmelCase	(       self							)	-> List[str]:
             lowercase__      :		Union[str, Any]          =     self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_cpmant_model(a		,      a							)


      def 		_UpperCAmelCase	(       self							)	-> Optional[int]:
             lowercase__      :		List[Any]          =     self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_lm_head_model(a		,      a							)





@require_torch
class 					UpperCAmelCase_     (   unittest.TestCase):


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

             lowercase__      :		Union[str, Any]          =     torch.tensor(
                 [[[6.1_708, 5.9_244, 1.0_835], [6.5_207, 6.2_893, -1_1.3_3_2_4], [-1.0_107, -0.0_576, -5.9_577]]]		,      )
             self.assertTrue(torch.allclose(hidden_states[:, :3, :3]		,      a		,      atol=1e-2							)							)





@require_torch
class 					UpperCAmelCase_     (   unittest.TestCase):


      @tooslow
      def 		_UpperCAmelCase	(       self							)	-> Dict:
             lowercase__      :		str          =     """今天天气真好！"""
             lowercase__      :		List[str]          =     """openbmb/cpm-ant-10b"""
             lowercase__      :		List[Any]          =     CpmAntForCausalLM.from_pretrained(a							)
             lowercase__      :		Dict          =     CpmAntTokenizer.from_pretrained(a							)
             lowercase__      :		str          =     tokenizer(a		,      return_tensors='pt'							)
             lowercase__      :		Any          =     model(**a							).logits

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


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


"""simple docstring"""





from typing import TYPE_CHECKING

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


_UpperCamelCase       :						int     =   {
    "configuration_retribert": ["RETRIBERT_PRETRAINED_CONFIG_ARCHIVE_MAP", "RetriBertConfig"],
    "tokenization_retribert": ["RetriBertTokenizer"],
}

try:
       if not is_tokenizers_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       _UpperCamelCase       :						Tuple     =   ["RetriBertTokenizerFast"]

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       _UpperCamelCase       :						int     =   [
           "RETRIBERT_PRETRAINED_MODEL_ARCHIVE_LIST",
           "RetriBertModel",
           "RetriBertPreTrainedModel",
       ]


if TYPE_CHECKING:
       from .configuration_retribert import RETRIBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, RetriBertConfig
       from .tokenization_retribert import RetriBertTokenizer

       try:
              if not is_tokenizers_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .tokenization_retribert_fast import RetriBertTokenizerFast

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_retribert import (
                  RETRIBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
                  RetriBertModel,
                  RetriBertPreTrainedModel,
              )

else:
       import sys

       _UpperCamelCase       :						Tuple     =   _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)

