

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    ,    _UpperCamelCase    ,    _UpperCamelCase=2    ,    _UpperCamelCase=8    ,    _UpperCamelCase=True    ,    _UpperCamelCase=False    ,    _UpperCamelCase=False    ,    _UpperCamelCase=False    ,    _UpperCamelCase=False    ,    _UpperCamelCase=99    ,    _UpperCamelCase=32    ,    _UpperCamelCase=3    ,    _UpperCamelCase=4    ,    _UpperCamelCase=37    ,    _UpperCamelCase=32    ,    _UpperCamelCase=128    ,    _UpperCamelCase=8    ,    _UpperCamelCase=8    ,    _UpperCamelCase=8    ,    _UpperCamelCase=1.0    ,    _UpperCamelCase=True    ,    )    ->				Any:
							lowerCAmelCase_      =     parent
							lowerCAmelCase_      =     batch_size
							lowerCAmelCase_      =     seq_length
							lowerCAmelCase_      =     is_training
							lowerCAmelCase_      =     use_token_type_ids
							lowerCAmelCase_      =     use_input_mask
							lowerCAmelCase_      =     use_labels
							lowerCAmelCase_      =     use_mc_token_ids
							lowerCAmelCase_      =     vocab_size
							lowerCAmelCase_      =     hidden_size
							lowerCAmelCase_      =     num_hidden_layers
							lowerCAmelCase_      =     num_attention_heads
							lowerCAmelCase_      =     intermediate_size
							lowerCAmelCase_      =     num_buckets
							lowerCAmelCase_      =     max_distance
							lowerCAmelCase_      =     prompt_length
							lowerCAmelCase_      =     prompt_types
							lowerCAmelCase_      =     segment_types
							lowerCAmelCase_      =     init_std
							lowerCAmelCase_      =     return_dict




				def 					__a						( self					)    ->				Tuple:
							lowerCAmelCase_      =     {}
							lowerCAmelCase_      =     ids_tensor([self.batch_size, self.seq_length]    ,    self.vocab_size					).type(torch.intaa					)
							lowerCAmelCase_      =     False

							lowerCAmelCase_      =     self.get_config()

							return (config, input_ids)




				def 					__a						( self					)    ->				Dict:
							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 					__a						( self    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    *_UpperCamelCase					)    ->				Dict:
							lowerCAmelCase_      =     CpmAntModel(config=UpperCamelCase__					)
							model.to(UpperCamelCase__					)
							model.eval()

							lowerCAmelCase_      =     model(**UpperCamelCase__					).last_hidden_state

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




				def 					__a						( self    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    *_UpperCamelCase					)    ->				int:
							lowerCAmelCase_      =     CpmAntForCausalLM(UpperCamelCase__					)
							model.to(UpperCamelCase__					)
							lowerCAmelCase_      =     input_ids["input_ids"].to(UpperCamelCase__					)
							model.eval()

							lowerCAmelCase_      =     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 					__a						( self					)    ->				Union[str, Any]:
							lowerCAmelCase_     ,  lowerCAmelCase_      =     self.prepare_config_and_inputs()
							return config, inputs_dict





@require_torch
class      _lowerCAmelCase (     _A ,				_A ,				unittest.TestCase	):
				_lowercase      		=(CpmAntModel, CpmAntForCausalLM) if is_torch_available() else ()
				_lowercase      		=(
				    {'''feature-extraction''': CpmAntModel, '''text-generation''': CpmAntForCausalLM} if is_torch_available() else {}
				)

				_lowercase      		=False
				_lowercase      		=False
				_lowercase      		=False
				_lowercase      		=False
				_lowercase      		=False




				def 					__a						( self					)    ->				Any:
							lowerCAmelCase_      =     CpmAntModelTester(self					)
							lowerCAmelCase_      =     ConfigTester(self    ,    config_class=UpperCamelCase__					)




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




				def 					__a						( self					)    ->				Union[str, 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.\\n                 So is attentions. We strongly recommand you use loss to tune model."					)(self.test_retain_grad_hidden_states_attentions					)




				def 					__a						( self					)    ->				Tuple:
							lowerCAmelCase_     ,  lowerCAmelCase_      =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_cpmant_model(UpperCamelCase__    ,    UpperCamelCase__					)




				def 					__a						( self					)    ->				List[Any]:
							lowerCAmelCase_     ,  lowerCAmelCase_      =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_lm_head_model(UpperCamelCase__    ,    UpperCamelCase__					)





@require_torch
class      _lowerCAmelCase (     unittest.TestCase	):




				@tooslow
				def 					__a						( self					)    ->				int:
							lowerCAmelCase_      =     "今天天气真好！"
							lowerCAmelCase_      =     "openbmb/cpm-ant-10b"
							lowerCAmelCase_      =     CpmAntModel.from_pretrained(UpperCamelCase__					)
							lowerCAmelCase_      =     CpmAntTokenizer.from_pretrained(UpperCamelCase__					)
							lowerCAmelCase_      =     tokenizer(UpperCamelCase__    ,    return_tensors="pt"					)
							lowerCAmelCase_      =     model(**UpperCamelCase__					).last_hidden_state

							lowerCAmelCase_      =     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      _lowerCAmelCase (     unittest.TestCase	):




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

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




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





from __future__ import annotations






def        lowerCamelCase__    (    __lowerCAmelCase  :  int   ,							__lowerCAmelCase  :  int   ,							__lowerCAmelCase  :  int    ):




			"""simple docstring"""


			assert n > 1 and a > 0 and greatest_common_divisor(__lowerCAmelCase   ,							__lowerCAmelCase    ) == 1
			((lowerCAmelCase_)     ,  (lowerCAmelCase_)     ,  (lowerCAmelCase_))      =     extended_gcd(__lowerCAmelCase   ,							__lowerCAmelCase    )  # Implemented below
			lowerCAmelCase_      =     (b * s) % n
			return x






def        lowerCamelCase__    (    __lowerCAmelCase  :  int   ,							__lowerCAmelCase  :  int    ):




			"""simple docstring"""


			((lowerCAmelCase_)     ,  (lowerCAmelCase_))      =     extended_euclid(__lowerCAmelCase   ,							__lowerCAmelCase    )  # Implemented below
			if b < 0:
						lowerCAmelCase_      =     (b % n + n) % n
			return b






def        lowerCamelCase__    (    __lowerCAmelCase  :  int   ,							__lowerCAmelCase  :  int   ,							__lowerCAmelCase  :  int    ):




			"""simple docstring"""


			lowerCAmelCase_      =     invert_modulo(__lowerCAmelCase   ,							__lowerCAmelCase    )
			lowerCAmelCase_      =     (b * s) % n
			return x






def        lowerCamelCase__    (    __lowerCAmelCase  :  int   ,							__lowerCAmelCase  :  int    ):




			"""simple docstring"""


			assert a >= 0 and b >= 0

			if b == 0:
						lowerCAmelCase_     ,  lowerCAmelCase_     ,  lowerCAmelCase_      =     a, 1, 0
			else:
						((lowerCAmelCase_)     ,  (lowerCAmelCase_)     ,  (lowerCAmelCase_))      =     extended_gcd(__lowerCAmelCase   ,							a % b    )
						lowerCAmelCase_      =     q
						lowerCAmelCase_      =     p - q * (a // b)

			assert a % d == 0 and b % d == 0
			assert d == a * x + b * y

			return (d, x, y)






def        lowerCamelCase__    (    __lowerCAmelCase  :  int   ,							__lowerCAmelCase  :  int    ):




			"""simple docstring"""


			if b == 0:
						return (1, 0)
			((lowerCAmelCase_)     ,  (lowerCAmelCase_))      =     extended_euclid(__lowerCAmelCase   ,							a % b    )
			lowerCAmelCase_      =     a // b
			return (y, x - k * y)






def        lowerCamelCase__    (    __lowerCAmelCase  :  int   ,							__lowerCAmelCase  :  int    ):




			"""simple docstring"""


			if a < b:
						lowerCAmelCase_     ,  lowerCAmelCase_      =     b, a

			while a % b != 0:
						lowerCAmelCase_     ,  lowerCAmelCase_      =     b, a % b

			return b


if __name__ == "__main__":
			from doctest import testmod

			testmod(name="modular_division", verbose=True)
			testmod(name="modular_division2", verbose=True)
			testmod(name="invert_modulo", verbose=True)
			testmod(name="extended_gcd", verbose=True)
			testmod(name="extended_euclid", verbose=True)
			testmod(name="greatest_common_divisor", verbose=True)



