






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



      """simple docstring"""
      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   ,    )    ->				Dict:
          UpperCamelCase	:Union[str, Any]     							=       parent
          UpperCamelCase	:Tuple     							=       batch_size
          UpperCamelCase	:Dict     							=       seq_length
          UpperCamelCase	:Optional[Any]     							=       is_training
          UpperCamelCase	:Optional[int]     							=       use_token_type_ids
          UpperCamelCase	:List[str]     							=       use_input_mask
          UpperCamelCase	:str     							=       use_labels
          UpperCamelCase	:Tuple     							=       use_mc_token_ids
          UpperCamelCase	:Dict     							=       vocab_size
          UpperCamelCase	:Any     							=       hidden_size
          UpperCamelCase	:Union[str, Any]     							=       num_hidden_layers
          UpperCamelCase	:List[Any]     							=       num_attention_heads
          UpperCamelCase	:Union[str, Any]     							=       intermediate_size
          UpperCamelCase	:Tuple     							=       num_buckets
          UpperCamelCase	:Optional[Any]     							=       max_distance
          UpperCamelCase	:List[Any]     							=       prompt_length
          UpperCamelCase	:Optional[int]     							=       prompt_types
          UpperCamelCase	:List[Any]     							=       segment_types
          UpperCamelCase	:List[str]     							=       init_std
          UpperCamelCase	:int     							=       return_dict
      def    UpperCAmelCase						(    self			)    ->				Union[str, Any]:
          UpperCamelCase	:Optional[Any]     							=       {}
          UpperCamelCase	:Optional[int]     							=       ids_tensor([self.batch_size, self.seq_length]   ,    self.vocab_size			).type(torch.intaa			)
          UpperCamelCase	:Any     							=       False

          UpperCamelCase	:Any     							=       self.get_config()

          return (config, input_ids)
      def    UpperCAmelCase						(    self			)    ->				Optional[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=SCREAMING_SNAKE_CASE_   ,    init_std=self.init_std   ,    return_dict=self.return_dict   ,    )
      def    UpperCAmelCase						(    self   ,    SCREAMING_SNAKE_CASE_   ,    SCREAMING_SNAKE_CASE_   ,    *SCREAMING_SNAKE_CASE_			)    ->				str:
          UpperCamelCase	:Union[str, Any]     							=       CpmAntModel(config=SCREAMING_SNAKE_CASE_			)
          model.to(SCREAMING_SNAKE_CASE_			)
          model.eval()

          UpperCamelCase	:Any     							=       model(**SCREAMING_SNAKE_CASE_			).last_hidden_state

          self.parent.assertEqual(hidden_states.shape   ,    (self.batch_size, self.seq_length, config.hidden_size)			)
      def    UpperCAmelCase						(    self   ,    SCREAMING_SNAKE_CASE_   ,    SCREAMING_SNAKE_CASE_   ,    *SCREAMING_SNAKE_CASE_			)    ->				int:
          UpperCamelCase	:Dict     							=       CpmAntForCausalLM(SCREAMING_SNAKE_CASE_			)
          model.to(SCREAMING_SNAKE_CASE_			)
          UpperCamelCase	:List[Any]     							=       input_ids['input_ids'].to(SCREAMING_SNAKE_CASE_			)
          model.eval()

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







@require_torch
class       UpperCAmelCase_						( snake_case__,     snake_case__,     unittest.TestCase			):



      """simple docstring"""
      UpperCamelCase_    : Any									=(CpmAntModel, CpmAntForCausalLM) if is_torch_available() else ()
      UpperCamelCase_    : List[Any]									=(
          {'feature-extraction': CpmAntModel, 'text-generation': CpmAntForCausalLM} if is_torch_available() else {}
      )

      UpperCamelCase_    : int									=False
      UpperCamelCase_    : Union[str, Any]									=False
      UpperCamelCase_    : Optional[Any]									=False
      UpperCamelCase_    : Any									=False
      UpperCamelCase_    : str									=False
      def    UpperCAmelCase						(    self			)    ->				List[str]:
          UpperCamelCase	:Dict     							=       CpmAntModelTester(self			)
          UpperCamelCase	:List[str]     							=       ConfigTester(self   ,    config_class=SCREAMING_SNAKE_CASE_			)
      def    UpperCAmelCase						(    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    UpperCAmelCase						(    self			)    ->				List[Any]:
          unittest.skip('''CPMAnt doesn\'t support input_embeds.'''			)(self.test_inputs_embeds			)
      def    UpperCAmelCase						(    self			)    ->				Optional[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    UpperCAmelCase						(    self			)    ->				Dict:
          UpperCamelCase	:Union[str, Any]     							=       self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_cpmant_model(SCREAMING_SNAKE_CASE_   ,    SCREAMING_SNAKE_CASE_			)




      def    UpperCAmelCase						(    self			)    ->				int:
          UpperCamelCase	: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       UpperCAmelCase_						( unittest.TestCase			):



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

          UpperCamelCase	:Union[str, Any]     							=       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       UpperCAmelCase_						( unittest.TestCase			):



      """simple docstring"""
      @tooslow
      def    UpperCAmelCase						(    self			)    ->				Optional[int]:
          UpperCamelCase	:Optional[int]     							=       '今天天气真好！'
          UpperCamelCase	:Any     							=       'openbmb/cpm-ant-10b'
          UpperCamelCase	:int     							=       CpmAntForCausalLM.from_pretrained(SCREAMING_SNAKE_CASE_			)
          UpperCamelCase	:Optional[Any]     							=       CpmAntTokenizer.from_pretrained(SCREAMING_SNAKE_CASE_			)
          UpperCamelCase	:Dict     							=       tokenizer(SCREAMING_SNAKE_CASE_   ,    return_tensors='''pt'''			)
          UpperCamelCase	:Dict     							=       model(**SCREAMING_SNAKE_CASE_			).logits

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




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













from transformers import (
    DPRContextEncoderTokenizer,
    DPRContextEncoderTokenizerFast,
    DPRQuestionEncoderTokenizer,
    DPRQuestionEncoderTokenizerFast,
    DPRReaderOutput,
    DPRReaderTokenizer,
    DPRReaderTokenizerFast,
)
from transformers.testing_utils import require_tokenizers, slow
from transformers.tokenization_utils_base import BatchEncoding

from ..bert.test_tokenization_bert import BertTokenizationTest







@require_tokenizers
class       UpperCAmelCase_						( lowercase			):



      """simple docstring"""
      UpperCamelCase_    : Optional[int]									=DPRContextEncoderTokenizer
      UpperCamelCase_    : Any									=DPRContextEncoderTokenizerFast
      UpperCamelCase_    : Optional[int]									=True







@require_tokenizers
class       UpperCAmelCase_						( lowercase			):



      """simple docstring"""
      UpperCamelCase_    : Tuple									=DPRQuestionEncoderTokenizer
      UpperCamelCase_    : Union[str, Any]									=DPRQuestionEncoderTokenizerFast
      UpperCamelCase_    : Any									=True







@require_tokenizers
class       UpperCAmelCase_						( lowercase			):



      """simple docstring"""
      UpperCamelCase_    : Tuple									=DPRReaderTokenizer
      UpperCamelCase_    : List[Any]									=DPRReaderTokenizerFast
      UpperCamelCase_    : Union[str, Any]									=True
      @slow
      def    UpperCAmelCase						(    self			)    ->				Dict:
          UpperCamelCase	:Optional[int]     							=       self.tokenizer_class.from_pretrained('''bert-base-uncased'''			)

          UpperCamelCase	:Optional[Any]     							=       tokenizer.encode('''question sequence'''   ,    add_special_tokens=SCREAMING_SNAKE_CASE_			)
          UpperCamelCase	:Optional[Any]     							=       tokenizer.encode('''title sequence'''   ,    add_special_tokens=SCREAMING_SNAKE_CASE_			)
          UpperCamelCase	:List[str]     							=       tokenizer.encode('''text sequence ''' * 4   ,    add_special_tokens=SCREAMING_SNAKE_CASE_			)
          UpperCamelCase	:Optional[int]     							=       [[101] + text_a + [102] + text_a + [102] + text_a]
          UpperCamelCase	:Tuple     							=       BatchEncoding({'''input_ids''': input_ids}			)

          UpperCamelCase	:Dict     							=       [[0] * len(input_ids[0]			)]
          UpperCamelCase	:List[Any]     							=       [[0] * len(input_ids[0]			)]
          UpperCamelCase	:Tuple     							=       [0]
          UpperCamelCase	:Union[str, Any]     							=       DPRReaderOutput(SCREAMING_SNAKE_CASE_   ,    SCREAMING_SNAKE_CASE_   ,    SCREAMING_SNAKE_CASE_			)

          UpperCamelCase							,       UpperCamelCase	:Union[str, Any]     							=       8, 9
          UpperCamelCase	:Any     							=       10
          UpperCamelCase	:Optional[Any]     							=       10
          UpperCamelCase	:Optional[int]     							=       tokenizer.decode_best_spans(SCREAMING_SNAKE_CASE_   ,    SCREAMING_SNAKE_CASE_			)
          self.assertEqual(predicted_spans[0].start_index   ,    SCREAMING_SNAKE_CASE_			)
          self.assertEqual(predicted_spans[0].end_index   ,    SCREAMING_SNAKE_CASE_			)
          self.assertEqual(predicted_spans[0].doc_id   ,    0			)




      @slow
      def    UpperCAmelCase						(    self			)    ->				Any:
          UpperCamelCase	:List[str]     							=       self.tokenizer_class.from_pretrained('''bert-base-uncased'''			)

          UpperCamelCase	:Any     							=       tokenizer.encode('''question sequence'''   ,    add_special_tokens=SCREAMING_SNAKE_CASE_			)
          UpperCamelCase	:str     							=       tokenizer.encode('''title sequence'''   ,    add_special_tokens=SCREAMING_SNAKE_CASE_			)
          UpperCamelCase	:Union[str, Any]     							=       tokenizer.encode('''text sequence'''   ,    add_special_tokens=SCREAMING_SNAKE_CASE_			)
          UpperCamelCase	:List[Any]     							=       [101] + text_a + [102] + text_a + [102] + text_a
          UpperCamelCase	:Any     							=       tokenizer(questions=['''question sequence''']   ,    titles=['''title sequence''']   ,    texts=['''text sequence''']			)
          self.assertIn('''input_ids'''   ,    SCREAMING_SNAKE_CASE_			)
          self.assertIn('''attention_mask'''   ,    SCREAMING_SNAKE_CASE_			)
          self.assertListEqual(encoded_input['''input_ids'''][0]   ,    SCREAMING_SNAKE_CASE_			)





