code: 
  
  
  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   snake_case__ :
  
  
  
  
  
  
  
         def __init__( self:   Union[str, Any],A_:   Tuple,A_:   List[str]=2,A_:   str=8,A_:   Any=True,A_:   Union[str, Any]=False,A_:   Union[str, Any]=False,A_:   Union[str, Any]=False,A_:   List[Any]=False,A_:   Optional[Any]=99,A_:   Any=32,A_:   Optional[int]=3,A_:   Any=4,A_:   Union[str, Any]=37,A_:   str=32,A_:   Optional[int]=128,A_:   int=8,A_:   Any=8,A_:   Any=8,A_:   str=1.0,A_:   List[Any]=True,):
  
  
  
  
  
  
              '''simple docstring'''
  
  
  
  
              __UpperCamelCase   = parent
              __UpperCamelCase   = batch_size
              __UpperCamelCase   = seq_length
              __UpperCamelCase   = is_training
              __UpperCamelCase   = use_token_type_ids
              __UpperCamelCase   = use_input_mask
              __UpperCamelCase   = use_labels
              __UpperCamelCase   = use_mc_token_ids
              __UpperCamelCase   = vocab_size
              __UpperCamelCase   = hidden_size
              __UpperCamelCase   = num_hidden_layers
              __UpperCamelCase   = num_attention_heads
              __UpperCamelCase   = intermediate_size
              __UpperCamelCase   = num_buckets
              __UpperCamelCase   = max_distance
              __UpperCamelCase   = prompt_length
              __UpperCamelCase   = prompt_types
              __UpperCamelCase   = segment_types
              __UpperCamelCase   = init_std
              __UpperCamelCase   = return_dict
  
  
  
  
  
  
  
         def 		snake_case_  ( self:   List[str]				):
  
  
  
  
  
  
              '''simple docstring'''
  
  
  
  
              __UpperCamelCase   = {}
              __UpperCamelCase   = ids_tensor([self.batch_size, self.seq_length],self.vocab_size				).type(torch.intaa				)
              __UpperCamelCase   = False
  
              __UpperCamelCase   = self.get_config()
  
              return (config, input_ids)
  
  
  
  
  
  
  
         def 		snake_case_  ( self:   List[Any]				):
  
  
  
  
  
  
              '''simple docstring'''
  
  
  
  
              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 		snake_case_  ( self:   List[str],A_:   str,A_:   List[str],*A_:   Optional[int]				):
  
  
  
  
  
  
              '''simple docstring'''
  
  
  
  
              __UpperCamelCase   = CpmAntModel(config=A_				)
              model.to(A_				)
              model.eval()
  
              __UpperCamelCase   = model(**A_				).last_hidden_state
  
              self.parent.assertEqual(hidden_states.shape,(self.batch_size, self.seq_length, config.hidden_size)				)
  
  
  
  
  
  
  
         def 		snake_case_  ( self:   Tuple,A_:   int,A_:   Optional[int],*A_:   Union[str, Any]				):
  
  
  
  
  
  
              '''simple docstring'''
  
  
  
  
              __UpperCamelCase   = CpmAntForCausalLM(A_				)
              model.to(A_				)
              __UpperCamelCase   = input_ids['input_ids'].to(A_				)
              model.eval()
  
              __UpperCamelCase   = 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 		snake_case_  ( self:   str				):
  
  
  
  
  
  
              '''simple docstring'''
  
  
  
  
              __UpperCamelCase,						__UpperCamelCase   = self.prepare_config_and_inputs()
              return config, inputs_dict
  
  
  
  @require_torch
  class   snake_case__ (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 		snake_case_  ( self:   Union[str, Any]				):
  
  
  
  
  
  
              '''simple docstring'''
  
  
  
  
              __UpperCamelCase   = CpmAntModelTester(self				)
              __UpperCamelCase   = ConfigTester(self,config_class=A_				)
  
  
  
  
  
  
  
         def 		snake_case_  ( self:   Optional[int]				):
  
  
  
  
  
  
              '''simple docstring'''
  
  
  
  
              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 		snake_case_  ( self:   Dict				):
  
  
  
  
  
  
              '''simple docstring'''
  
  
  
  
              unittest.skip('CPMAnt doesn\'t support input_embeds.'				)(self.test_inputs_embeds				)
  
  
  
  
  
  
  
         def 		snake_case_  ( self:   str				):
  
  
  
  
  
  
              '''simple docstring'''
  
  
  
  
              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 		snake_case_  ( self:   Dict				):
  
  
  
  
  
  
              '''simple docstring'''
  
  
  
  
              __UpperCamelCase,						__UpperCamelCase   = self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_cpmant_model(A_,A_				)
  
  
  
  
  
  
  
         def 		snake_case_  ( self:   Optional[Any]				):
  
  
  
  
  
  
              '''simple docstring'''
  
  
  
  
              __UpperCamelCase,						__UpperCamelCase   = self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_lm_head_model(A_,A_				)
  
  
  
  @require_torch
  class   snake_case__ (unittest.TestCase	):
  
  
  
  
  
  
  
         @tooslow
         def 		snake_case_  ( self:   Dict				):
  
  
  
  
  
  
              '''simple docstring'''
  
  
  
  
              __UpperCamelCase   = '今天天气真好！'
              __UpperCamelCase   = 'openbmb/cpm-ant-10b'
              __UpperCamelCase   = CpmAntModel.from_pretrained(A_				)
              __UpperCamelCase   = CpmAntTokenizer.from_pretrained(A_				)
              __UpperCamelCase   = tokenizer(A_,return_tensors='pt'				)
              __UpperCamelCase   = model(**A_				).last_hidden_state
  
              __UpperCamelCase   = 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],A_,atol=1E-2				)				)
  
  
  
  @require_torch
  class   snake_case__ (unittest.TestCase	):
  
  
  
  
  
  
  
         @tooslow
         def 		snake_case_  ( self:   Any				):
  
  
  
  
  
  
              '''simple docstring'''
  
  
  
  
              __UpperCamelCase   = '今天天气真好！'
              __UpperCamelCase   = 'openbmb/cpm-ant-10b'
              __UpperCamelCase   = CpmAntForCausalLM.from_pretrained(A_				)
              __UpperCamelCase   = CpmAntTokenizer.from_pretrained(A_				)
              __UpperCamelCase   = tokenizer(A_,return_tensors='pt'				)
              __UpperCamelCase   = model(**A_				).logits
  
              __UpperCamelCase   = 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],A_,atol=1E-2				)				)
  
  
  
  
  
  
  
         @tooslow
         def 		snake_case_  ( self:   Any				):
  
  
  
  
  
  
              '''simple docstring'''
  
  
  
  
              __UpperCamelCase   = 'openbmb/cpm-ant-10b'
              __UpperCamelCase   = CpmAntForCausalLM.from_pretrained(A_				)
              __UpperCamelCase   = CpmAntTokenizer.from_pretrained(A_				)
              __UpperCamelCase   = '今天天气不错，'
              __UpperCamelCase   = '今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的'
              __UpperCamelCase   = tokenizer(A_,return_tensors='pt'				)
              __UpperCamelCase   = model.generate(**A_				)
              __UpperCamelCase   = tokenizer.batch_decode(A_				)
              self.assertEqual(A_,A_				)
  
  
  
  
  
  
  
         @tooslow
         def 		snake_case_  ( self:   int				):
  
  
  
  
  
  
              '''simple docstring'''
  
  
  
  
              __UpperCamelCase   = 'openbmb/cpm-ant-10b'
              __UpperCamelCase   = CpmAntForCausalLM.from_pretrained(A_				)
              __UpperCamelCase   = CpmAntTokenizer.from_pretrained(A_				)
              __UpperCamelCase   = ['今天天气不错，', '新年快乐，万事如意！']
              __UpperCamelCase   = [
                  '今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的',
                  '新年快乐，万事如意！在这辞旧迎新的美好时刻，我谨代表《农村新技术》杂志社全体同仁，向一直以来关心、支持《农村新技术》杂志发展的各级领导、各界朋友和广大读者致以最诚挚的',
              ]
              __UpperCamelCase   = tokenizer(A_,return_tensors='pt',padding=A_				)
              __UpperCamelCase   = model.generate(**A_				)
              __UpperCamelCase   = tokenizer.batch_decode(A_				)
              self.assertEqual(A_,A_				)
  
  
  
style_context: 
  
  
  import jax
  import jax.numpy as jnp
  from bigbird_flax import FlaxBigBirdForNaturalQuestions
  from datasets import load_from_disk
  
  from transformers import BigBirdTokenizerFast
  
  
  __snake_case             =	{0: '''null''', 1: '''short''', 2: '''long''', 3: '''yes''', 4: '''no'''}
  __snake_case             =	set(''''''.join(['''‘''', '''’''', '''´''', '''`''', '''.''', ''',''', '''-''', '''"''']))
  
  
  
  
  
  
  def 		_A		(		_lowercase				,     _lowercase=0				,     _lowercase=None   )				->   Union[str, Any]:
  
  
       """simple docstring"""
  
  
       return [" ".join(x.split(' '   )[begin:end]   ) for x in answers if len(x.split(' '   )   ) > 1]
  
  
  
  
  
  
  def 		_A		(		_lowercase				,     _lowercase=False   )				->   Tuple:
  
  
       """simple docstring"""
  
  
       if make_sub_answers:
            # if answers are longer than one word, make sure a predictions is correct if it coresponds to the complete 1: or :-1 sub word
            # *e.g.* if the correct answer contains a prefix such as "the", or "a"
            __UpperCamelCase   = (
                given_answers + get_sub_answers(_lowercase				,     begin=1   ) + get_sub_answers(_lowercase				,     end=-1   )
            )
       __UpperCamelCase   = []
       for answer in given_answers:
            __UpperCamelCase   = answer.replace('_'				,     ' '   ).lower()
            __UpperCamelCase   = ''.join(c if c not in PUNCTUATION_SET_TO_EXCLUDE else ' ' for c in alias   )
            answers.append(' '.join(alias.split()   ).strip()   )
       return set(_lowercase   )
  
  
  
  
  
  
  def 		_A		(		_lowercase				,     _lowercase				,     _lowercase=1				,     _lowercase=1_00   )				->   Any:
  
  
       """simple docstring"""
  
  
       __UpperCamelCase,						__UpperCamelCase   = jax.lax.top_k(_lowercase				,     _lowercase   )
       __UpperCamelCase,						__UpperCamelCase   = jax.lax.top_k(_lowercase				,     _lowercase   )
  
       __UpperCamelCase   = best_end_idx[:, None] - best_start_idx[None, :]
       __UpperCamelCase   = jnp.logical_or(widths < 0				,     widths > max_size   )
       __UpperCamelCase   = (best_end_scores[:, None] + best_start_scores[None, :]) - (1e8 * mask)
       __UpperCamelCase   = jnp.argmax(_lowercase   ).item()
  
       return best_start_idx[best_score % top_k], best_end_idx[best_score // top_k]
  
  
  
  
  
  
  def 		_A		(		_lowercase   )				->   List[str]:
  
  
       """simple docstring"""
  
  
       __UpperCamelCase   = sample['question']['text']
       __UpperCamelCase   = sample['document']['tokens']['token']
       __UpperCamelCase   = sample['document']['tokens']['is_html']
       __UpperCamelCase   = sample['annotations']['long_answer']
       __UpperCamelCase   = sample['annotations']['short_answers']
  
       __UpperCamelCase   = ' '.join([context[i] for i in range(len(_lowercase   )   ) if not is_html[i]]   )
  
       # 0 - No ; 1 - Yes
       for answer in sample["annotations"]["yes_no_answer"]:
            if answer == 0 or answer == 1:
                 return {
                     "question": question,
                     "context": context_string,
                     "short": [],
                     "long": [],
                     "category": "no" if answer == 0 else "yes",
                 }
  
       __UpperCamelCase   = []
       for s in short_answers:
            short_targets.extend(s['text']   )
       __UpperCamelCase   = list(set(_lowercase   )   )
  
       __UpperCamelCase   = []
       for s in long_answers:
            if s["start_token"] == -1:
                 continue
            __UpperCamelCase   = context[s['start_token'] : s['end_token']]
            __UpperCamelCase   = is_html[s['start_token'] : s['end_token']]
            __UpperCamelCase   = ' '.join([answer[i] for i in range(len(_lowercase   )   ) if not html[i]]   )
            if new_answer not in long_targets:
                 long_targets.append(_lowercase   )
  
       __UpperCamelCase   = 'long_short' if len(short_targets + long_targets   ) > 0 else 'null'
  
       return {
           "question": question,
           "context": context_string,
           "short": short_targets,
           "long": long_targets,
           "category": category,
       }
  
  
  
  
  
  
  def 		_A		(		)				->   Optional[int]:
  
  
       """simple docstring"""
  
  
       __UpperCamelCase   = load_from_disk('natural-questions-validation'   )
       __UpperCamelCase   = dataset.map(_lowercase   ).remove_columns(['annotations', 'document', 'id']   )
       print(_lowercase   )
  
       __UpperCamelCase   = dataset.filter(lambda _lowercase   : (len(x['question']   ) + len(x['context']   )) < 4 * 40_96   )
       __UpperCamelCase   = short_validation_dataset.filter(lambda _lowercase   : x["category"] != "null"   )
       short_validation_dataset
  
       __UpperCamelCase   = 'vasudevgupta/flax-bigbird-natural-questions'
       __UpperCamelCase   = FlaxBigBirdForNaturalQuestions.from_pretrained(_lowercase   )
       __UpperCamelCase   = BigBirdTokenizerFast.from_pretrained(_lowercase   )
  
       @jax.jit
       def forward(*_lowercase				,     **_lowercase   ):
            __UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase   = model(*_lowercase				,     **_lowercase   )
            return start_logits, end_logits, jnp.argmax(_lowercase				,     axis=-1   )
  
       def evaluate(_lowercase   ):
            # encode question and context so that they are separated by a tokenizer.sep_token and cut at max_length
            __UpperCamelCase   = tokenizer(
                example['question']				,     example['context']				,     return_tensors='np'				,     max_length=40_96				,     padding='max_length'				,     truncation=_lowercase				,     )
  
            __UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase   = forward(**_lowercase   )
  
            __UpperCamelCase   = CATEGORY_MAPPING[category.item()]
  
            __UpperCamelCase   = example['long'] + example['short']
            if example["category"] in ["yes", "no", "null"]:
                 __UpperCamelCase   = [example['category']]
            __UpperCamelCase   = example['category'] != 'null'
            # Now target can be: "yes", "no", "null", "list of long & short answers"
  
            if predicted_category in ["yes", "no", "null"]:
                 __UpperCamelCase   = [predicted_category]
                 __UpperCamelCase   = example['output'] == example['targets']
                 __UpperCamelCase   = predicted_category != 'null'
                 return example
  
            __UpperCamelCase   = 38 if predicted_category == 'short' else 10_24
            __UpperCamelCase,						__UpperCamelCase   = get_best_valid_start_end_idx(
                start_scores[0]				,     end_scores[0]				,     top_k=8				,     max_size=_lowercase   )
  
            __UpperCamelCase   = inputs['input_ids'][0].tolist()
            __UpperCamelCase   = [tokenizer.decode(input_ids[start_score : end_score + 1]   )]
  
            __UpperCamelCase   = expand_to_aliases(example['targets']				,     make_sub_answers=_lowercase   )
            __UpperCamelCase   = expand_to_aliases(example['output']   )
  
            # some preprocessing to both prediction and answer
            __UpperCamelCase   = {''.join(a.split()   ) for a in answers}
            __UpperCamelCase   = {''.join(p.split()   ) for p in predictions}
            __UpperCamelCase   = {s for s in predictions if s not in ['``', '\'\'', '`', '\'']}
  
            # if there is a common element, it's a exact match
            __UpperCamelCase   = len(list(answers & predictions   )   ) > 0
            __UpperCamelCase   = predicted_category != 'null' and len(_lowercase   ) > 0
  
            return example
  
       __UpperCamelCase   = short_validation_dataset.map(_lowercase   )
  
       __UpperCamelCase   = len(_lowercase   )
       __UpperCamelCase   = len(short_validation_dataset.filter(lambda _lowercase   : x["match"] == 1   )   )
       print('EM score:'				,     (matched / total) * 1_00				,     '%'   )
  
  
  if __name__ == "__main__":
      main()
  
  
  