
'''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__ :



  '''simple docstring'''






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

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

        snake_case_    							=				self.get_config()

        return (config, input_ids)

  def 		_lowerCamelCase						(self	:			List[Any]     ) ->							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=lowerCAmelCase_						,			init_std=self.init_std						,			return_dict=self.return_dict						,			)

  def 		_lowerCamelCase						(self	:			str						,			lowerCAmelCase_	:			str						,			lowerCAmelCase_	:			Tuple						,			*lowerCAmelCase_	:			List[str]     ) ->							List[str]:
        snake_case_    							=				CpmAntModel(config=lowerCAmelCase_     )
        model.to(lowerCAmelCase_     )
        model.eval()

        snake_case_    							=				model(**lowerCAmelCase_     ).last_hidden_state

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

  def 		_lowerCamelCase						(self	:			str						,			lowerCAmelCase_	:			Tuple						,			lowerCAmelCase_	:			Union[str, Any]						,			*lowerCAmelCase_	:			str     ) ->							Optional[int]:
        snake_case_    							=				CpmAntForCausalLM(lowerCAmelCase_     )
        model.to(lowerCAmelCase_     )
        snake_case_    							=				input_ids["input_ids"].to(lowerCAmelCase_     )
        model.eval()

        snake_case_    							=				model(**lowerCAmelCase_     )
        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	:			Tuple     ) ->							Optional[Any]:
        snake_case_							,						snake_case_    							=				self.prepare_config_and_inputs()
        return config, inputs_dict



@require_torch
class 			lowerCAmelCase__ (							UpperCamelCase       ,				UpperCamelCase       ,				unittest.TestCase							):



  '''simple docstring'''






  _UpperCamelCase							=			(CpmAntModel, CpmAntForCausalLM) if is_torch_available() else ()
  _UpperCamelCase							=			(
      {'feature-extraction': CpmAntModel, 'text-generation': CpmAntForCausalLM} if is_torch_available() else {}
  )

  _UpperCamelCase							=			False
  _UpperCamelCase							=			False
  _UpperCamelCase							=			False
  _UpperCamelCase							=			False
  _UpperCamelCase							=			False

  def 		_lowerCamelCase						(self	:			str     ) ->							int:
        snake_case_    							=				CpmAntModelTester(self     )
        snake_case_    							=				ConfigTester(self						,			config_class=lowerCAmelCase_     )

  def 		_lowerCamelCase						(self	:			List[str]     ) ->							Optional[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 		_lowerCamelCase						(self	:			str     ) ->							List[Any]:
        unittest.skip("CPMAnt doesn't support input_embeds."     )(self.test_inputs_embeds     )

  def 		_lowerCamelCase						(self	:			Any     ) ->							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 		_lowerCamelCase						(self	:			List[Any]     ) ->							Any:
        snake_case_							,						snake_case_    							=				self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_cpmant_model(lowerCAmelCase_						,			lowerCAmelCase_     )


  def 		_lowerCamelCase						(self	:			Union[str, Any]     ) ->							Tuple:
        snake_case_							,						snake_case_    							=				self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_lm_head_model(lowerCAmelCase_						,			lowerCAmelCase_     )



@require_torch
class 			lowerCAmelCase__ (							unittest.TestCase							):



  '''simple docstring'''






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

        snake_case_    							=				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]						,			lowerCAmelCase_						,			atol=1E-2     )     )



@require_torch
class 			lowerCAmelCase__ (							unittest.TestCase							):



  '''simple docstring'''






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

        snake_case_    							=				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]						,			lowerCAmelCase_						,			atol=1E-2     )     )

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


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





"""simple docstring"""






import torch

from transformers import AutoTokenizer


class       lowerCamelCase_		(      lowercase ):





				"""simple docstring"""







				def __init__(    self							,		UpperCAmelCase__  ):
								SCREAMING_SNAKE_CASE__  				=					AutoTokenizer.from_pretrained(UpperCAmelCase__  )



				def 	lowerCAmelCase__		(    self							,		UpperCAmelCase__  ):

								if isinstance(UpperCAmelCase__							,		UpperCAmelCase__  ) and all(isinstance(_x							,		UpperCAmelCase__  ) for _x in x  ):
												SCREAMING_SNAKE_CASE__  				=					None
												for l in x:
																SCREAMING_SNAKE_CASE__  				=					self.tokenizer(
																    UpperCAmelCase__							,		padding="max_length"							,		max_length=384							,		truncation=UpperCAmelCase__							,		return_tensors="pt"							,		)
																SCREAMING_SNAKE_CASE__  				=					torch.tensor([len(UpperCAmelCase__  )]  )
																if d is not None:
																				for k in d.keys():
																								SCREAMING_SNAKE_CASE__  				=					torch.cat((d[k], t[k])							,		0  )
																else:
																				SCREAMING_SNAKE_CASE__  				=					t

												SCREAMING_SNAKE_CASE__  				=					torch.tensor(self.tokenizer.convert_tokens_to_ids("[E]"  )  )
												SCREAMING_SNAKE_CASE__  				=					torch.tensor(self.tokenizer.convert_tokens_to_ids("[/E]"  )  )

								elif isinstance(UpperCAmelCase__							,		UpperCAmelCase__  ) and all(isinstance(_x							,		UpperCAmelCase__  ) for _x in x  ):
												SCREAMING_SNAKE_CASE__  				=					self.tokenizer(
												    UpperCAmelCase__							,		padding="max_length"							,		max_length=384							,		truncation=UpperCAmelCase__							,		return_tensors="pt"							,		)

								else:
												raise Exception(
												    "Type of parameter x was not recognized! Only `list of strings` for query or `list of lists of"
												    " strings` for supports are supported."  )

								return d






				def 	lowerCAmelCase__		(    self							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__=0.70  ):
								SCREAMING_SNAKE_CASE__  				=					[]
								for idx in range(len(W_query["input_ids"]  )  ):
												SCREAMING_SNAKE_CASE__  				=					SCREAMING_SNAKE_CASE__  				=					range(p_start.shape[1]  )

												SCREAMING_SNAKE_CASE__  				=					[]
												for start_id in start_indexes:
																for end_id in end_indexes:
																				if start_id < end_id:
																								output.append(
																								    (
																								        start_id,
																								        end_id,
																								        p_start[idx][start_id].item(),
																								        p_end[idx][end_id].item(),
																								    )  )

												output.sort(key=lambda UpperCAmelCase__  : (tup[2] * tup[3])							,		reverse=UpperCAmelCase__  )
												SCREAMING_SNAKE_CASE__  				=					[]
												for k in range(len(UpperCAmelCase__  )  ):
																if output[k][2] * output[k][3] >= thresh:
																				SCREAMING_SNAKE_CASE__				,   SCREAMING_SNAKE_CASE__  				=					output[k][0], output[k][1]
																				SCREAMING_SNAKE_CASE__  				=					self.tokenizer.decode(W_query["input_ids"][idx][c_start_pos:c_end_pos]  )
																				temp.append((decoded, output[k][2] * output[k][3])  )

												final_outputs.append(UpperCAmelCase__  )

								return final_outputs
