





'''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   :Dict		, lowercase_   :int		, lowercase_   :Tuple=2		, lowercase_   :int=8		, lowercase_   :Tuple=True		, lowercase_   :List[Any]=False		, lowercase_   :List[Any]=False		, lowercase_   :Dict=False		, lowercase_   :Union[str, Any]=False		, lowercase_   :Optional[Any]=99		, lowercase_   :List[str]=32		, lowercase_   :List[str]=3		, lowercase_   :Dict=4		, lowercase_   :Union[str, Any]=37		, lowercase_   :int=32		, lowercase_   :int=1_28		, lowercase_   :Any=8		, lowercase_   :Tuple=8		, lowercase_   :str=8		, lowercase_   :Optional[int]=1.0		, lowercase_   :int=True		, )->							Tuple:
     A__       =    parent
     A__       =    batch_size
     A__       =    seq_length
     A__       =    is_training
     A__       =    use_token_type_ids
     A__       =    use_input_mask
     A__       =    use_labels
     A__       =    use_mc_token_ids
     A__       =    vocab_size
     A__       =    hidden_size
     A__       =    num_hidden_layers
     A__       =    num_attention_heads
     A__       =    intermediate_size
     A__       =    num_buckets
     A__       =    max_distance
     A__       =    prompt_length
     A__       =    prompt_types
     A__       =    segment_types
     A__       =    init_std
     A__       =    return_dict
 def   UpperCAmelCase_							(       self   :Tuple							)->							List[Any]:
     A__       =    {}
     A__       =    ids_tensor([self.batch_size, self.seq_length]		, self.vocab_size							).type(torch.intaa							)
     A__       =    False

     A__       =    self.get_config()

     return (config, input_ids)
 def   UpperCAmelCase_							(       self   :Dict							)->							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=lowercase_		, init_std=self.init_std		, return_dict=self.return_dict		, )
 def   UpperCAmelCase_							(       self   :Tuple		, lowercase_   :Union[str, Any]		, lowercase_   :List[Any]		, *lowercase_   :str							)->							Any:
     A__       =    CpmAntModel(config=lowercase_							)
     model.to(lowercase_							)
     model.eval()

     A__       =    model(**lowercase_							).last_hidden_state

     self.parent.assertEqual(hidden_states.shape		, (self.batch_size, self.seq_length, config.hidden_size)							)
 def   UpperCAmelCase_							(       self   :List[Any]		, lowercase_   :str		, lowercase_   :Tuple		, *lowercase_   :str							)->							Optional[Any]:
     A__       =    CpmAntForCausalLM(lowercase_							)
     model.to(lowercase_							)
     A__       =    input_ids["input_ids"].to(lowercase_							)
     model.eval()

     A__       =    model(**lowercase_							)
     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   :str							)->							Tuple:
     A__,			A__       =    self.prepare_config_and_inputs()
     return config, inputs_dict






@require_torch
class      UpperCAmelCase			( UpperCamelCase__    , UpperCamelCase__    , 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   UpperCAmelCase_							(       self   :Dict							)->							int:
     A__       =    CpmAntModelTester(self							)
     A__       =    ConfigTester(self		, config_class=lowercase_							)
 def   UpperCAmelCase_							(       self   :Optional[Any]							)->							Dict:
     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]							)->							Optional[int]:
     unittest.skip("CPMAnt doesn't support input_embeds."							)(self.test_inputs_embeds							)
 def   UpperCAmelCase_							(       self   :Union[str, Any]							)->							Tuple:
     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   :Union[str, Any]							)->							List[Any]:
     A__,			A__       =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_cpmant_model(lowercase_		, lowercase_							)



 def   UpperCAmelCase_							(       self   :Optional[int]							)->							int:
     A__,			A__       =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_lm_head_model(lowercase_		, lowercase_							)






@require_torch
class      UpperCAmelCase			( unittest.TestCase       ):
 @tooslow
 def   UpperCAmelCase_							(       self   :Union[str, Any]							)->							int:
     A__       =    "今天天气真好！"
     A__       =    "openbmb/cpm-ant-10b"
     A__       =    CpmAntModel.from_pretrained(lowercase_							)
     A__       =    CpmAntTokenizer.from_pretrained(lowercase_							)
     A__       =    tokenizer(lowercase_		, return_tensors="pt"							)
     A__       =    model(**lowercase_							).last_hidden_state

     A__       =    torch.tensor(
         [[[6.1_7_0_8, 5.9_2_4_4, 1.0_8_3_5], [6.5_2_0_7, 6.2_8_9_3, -1_1.3_3_2_4], [-1.0_1_0_7, -0.0_5_7_6, -5.9_5_7_7]]]		, )
     self.assertTrue(torch.allclose(hidden_states[:, :3, :3]		, lowercase_		, atol=1E-2							)							)






@require_torch
class      UpperCAmelCase			( unittest.TestCase       ):
 @tooslow
 def   UpperCAmelCase_							(       self   :List[Any]							)->							int:
     A__       =    "今天天气真好！"
     A__       =    "openbmb/cpm-ant-10b"
     A__       =    CpmAntForCausalLM.from_pretrained(lowercase_							)
     A__       =    CpmAntTokenizer.from_pretrained(lowercase_							)
     A__       =    tokenizer(lowercase_		, return_tensors="pt"							)
     A__       =    model(**lowercase_							).logits

     A__       =    torch.tensor(
         [[[-6.4_2_6_7, -6.4_0_8_3, -6.3_9_5_8], [-5.8_8_0_2, -5.9_4_4_7, -5.7_8_1_1], [-5.3_8_9_6, -5.4_8_2_0, -5.4_2_9_5]]]		, )
     self.assertTrue(torch.allclose(hidden_states[:, :3, :3]		, lowercase_		, atol=1E-2							)							)
 @tooslow
 def   UpperCAmelCase_							(       self   :Tuple							)->							List[Any]:
     A__       =    "openbmb/cpm-ant-10b"
     A__       =    CpmAntForCausalLM.from_pretrained(lowercase_							)
     A__       =    CpmAntTokenizer.from_pretrained(lowercase_							)
     A__       =    "今天天气不错，"
     A__       =    "今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的"
     A__       =    tokenizer(lowercase_		, return_tensors="pt"							)
     A__       =    model.generate(**lowercase_							)
     A__       =    tokenizer.batch_decode(lowercase_							)
     self.assertEqual(lowercase_		, lowercase_							)



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










'''simple docstring'''

import torch

from transformers import AutoTokenizer






class      UpperCAmelCase			( UpperCamelCase__       ):
 def __init__(       self   :Any		, lowercase_   :Dict							)->							Any:
     A__       =    AutoTokenizer.from_pretrained(lowercase_							)
 def   UpperCAmelCase_							(       self   :Tuple		, lowercase_   :int							)->							Optional[int]:

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

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

     elif isinstance(lowercase_		, lowercase_							) and all(isinstance(_x		, lowercase_							) for _x in x							):
         A__       =    self.tokenizer(
             lowercase_		, padding="max_length"		, max_length=3_84		, truncation=lowercase_		, 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   UpperCAmelCase_							(       self   :Union[str, Any]		, lowercase_   :List[str]		, lowercase_   :str		, lowercase_   :str		, lowercase_   :List[Any]		, lowercase_   :Dict=0.7_0							)->							int:
     A__       =    []
     for idx in range(len(W_query["input_ids"]							)							):
         A__       =    A__       =    range(p_start.shape[1]							)

         A__       =    []
         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 lowercase_							: (tup[2] * tup[3])		, reverse=lowercase_							)
         A__       =    []
         for k in range(len(lowercase_							)							):
             if output[k][2] * output[k][3] >= thresh:
                 A__,			A__       =    output[k][0], output[k][1]
                 A__       =    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(lowercase_							)

     return final_outputs



