






"""simple docstring"""



from __future__ import annotations

_UpperCamelCase						:    Optional[int]		=						[2, 1, 4, 5, 6, 0, 8, 9, 1, 2, 0, 6, 4, 2, 0, 6, 5, 3, 2, 7]

class 			a	:







    def __init__(     self				,						_lowerCamelCase			):
      lowercase				       =			-1
      lowercase				       =			-1
      lowercase				       =			[-1] * length
      lowercase				       =			None
      lowercase				       =			None







    def __repr__(     self			):
      return F'Node(min_value={self.minn} max_value={self.maxx})'


def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    list[int]       ):





  '''simple docstring'''





  lowercase				       =			Node(len(__snake_case       )       )
  lowercase       ,	lowercase				       =			min(__snake_case       ), max(__snake_case       )
  # Leaf node case where the node contains only one unique value
  if root.minn == root.maxx:
    return root
  lowercase				       =			(root.minn + root.maxx) // 2

  lowercase				       =			[]
  lowercase				       =			[]

  for index, num in enumerate(__snake_case       ):
    if num <= pivot:
      left_arr.append(__snake_case       )
    else:
      right_arr.append(__snake_case       )
    lowercase				       =			len(__snake_case       )
  lowercase				       =			build_tree(__snake_case       )
  lowercase				       =			build_tree(__snake_case       )
  return root
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    Node | None						,   __snake_case     :    int						,   __snake_case     :    int       ):





  '''simple docstring'''





  if index < 0 or node is None:
    return 0
  # Leaf node cases
  if node.minn == node.maxx:
    return index + 1 if node.minn == num else 0
  lowercase				       =			(node.minn + node.maxx) // 2
  if num <= pivot:
    # go the left subtree and map index to the left subtree
    return rank_till_index(node.left						,   __snake_case						,   node.map_left[index] - 1       )
  else:
    # go to the right subtree and map index to the right subtree
    return rank_till_index(node.right						,   __snake_case						,   index - node.map_left[index]       )
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    Node | None						,   __snake_case     :    int						,   __snake_case     :    int						,   __snake_case     :    int       ):





  '''simple docstring'''





  if start > end:
    return 0
  lowercase				       =			rank_till_index(__snake_case						,   __snake_case						,   __snake_case       )
  lowercase				       =			rank_till_index(__snake_case						,   __snake_case						,   start - 1       )
  return rank_till_end - rank_before_start
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    Node | None						,   __snake_case     :    int						,   __snake_case     :    int						,   __snake_case     :    int       ):





  '''simple docstring'''





  if index > (end - start) or start > end or node is None:
    return -1
  # Leaf node case
  if node.minn == node.maxx:
    return node.minn
  # Number of elements in the left subtree in interval [start, end]
  lowercase				       =			node.map_left[end] - (
      node.map_left[start - 1] if start else 0
  )
  if num_elements_in_left_tree > index:
    return quantile(
        node.left						,   __snake_case						,   (node.map_left[start - 1] if start else 0)						,   node.map_left[end] - 1						,   )
  else:
    return quantile(
        node.right						,   index - num_elements_in_left_tree						,   start - (node.map_left[start - 1] if start else 0)						,   end - node.map_left[end]						,   )
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    Node | None						,   __snake_case     :    int						,   __snake_case     :    int						,   __snake_case     :    int						,   __snake_case     :    int       ):





  '''simple docstring'''





  if (
      start > end
      or node is None
      or start_num > end_num
      or node.minn > end_num
      or node.maxx < start_num
  ):
    return 0
  if start_num <= node.minn and node.maxx <= end_num:
    return end - start + 1
  lowercase				       =			range_counting(
      node.left						,   (node.map_left[start - 1] if start else 0)						,   node.map_left[end] - 1						,   __snake_case						,   __snake_case						,   )
  lowercase				       =			range_counting(
      node.right						,   start - (node.map_left[start - 1] if start else 0)						,   end - node.map_left[end]						,   __snake_case						,   __snake_case						,   )
  return left + right


if __name__ == "__main__":
    import doctest

    doctest.testmod()









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:			Optional[Any]       , _SCREAMING_SNAKE_CASE:			Any       , _SCREAMING_SNAKE_CASE:			List[Any]=2       , _SCREAMING_SNAKE_CASE:			List[str]=8       , _SCREAMING_SNAKE_CASE:			Optional[Any]=True       , _SCREAMING_SNAKE_CASE:			List[str]=False       , _SCREAMING_SNAKE_CASE:			int=False       , _SCREAMING_SNAKE_CASE:			Optional[Any]=False       , _SCREAMING_SNAKE_CASE:			List[str]=False       , _SCREAMING_SNAKE_CASE:			List[Any]=99       , _SCREAMING_SNAKE_CASE:			str=32       , _SCREAMING_SNAKE_CASE:			Tuple=3       , _SCREAMING_SNAKE_CASE:			Union[str, Any]=4       , _SCREAMING_SNAKE_CASE:			List[Any]=37       , _SCREAMING_SNAKE_CASE:			int=32       , _SCREAMING_SNAKE_CASE:			int=128       , _SCREAMING_SNAKE_CASE:			Dict=8       , _SCREAMING_SNAKE_CASE:			Optional[Any]=8       , _SCREAMING_SNAKE_CASE:			str=8       , _SCREAMING_SNAKE_CASE:			Dict=1.0       , _SCREAMING_SNAKE_CASE:			List[str]=True       , )							->			Optional[int]:




              """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 lowercase			(		self:			Dict		)							->			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 lowercase			(		self:			Union[str, Any]		)							->			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=_SCREAMING_SNAKE_CASE       , init_std=self.init_std       , return_dict=self.return_dict       , )

       def lowercase			(		self:			Dict       , _SCREAMING_SNAKE_CASE:			Union[str, Any]       , _SCREAMING_SNAKE_CASE:			List[str]       , *_SCREAMING_SNAKE_CASE:			Tuple		)							->			str:




              """simple docstring"""


              UpperCamelCase_ 							=	CpmAntModel(config=_SCREAMING_SNAKE_CASE		)
              model.to(_SCREAMING_SNAKE_CASE		)
              model.eval()

              UpperCamelCase_ 							=	model(**_SCREAMING_SNAKE_CASE		).last_hidden_state

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

       def lowercase			(		self:			Optional[Any]       , _SCREAMING_SNAKE_CASE:			Tuple       , _SCREAMING_SNAKE_CASE:			Optional[Any]       , *_SCREAMING_SNAKE_CASE:			Dict		)							->			Optional[int]:




              """simple docstring"""


              UpperCamelCase_ 							=	CpmAntForCausalLM(_SCREAMING_SNAKE_CASE		)
              model.to(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_ 							=	input_ids["input_ids"].to(_SCREAMING_SNAKE_CASE		)
              model.eval()

              UpperCamelCase_ 							=	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 lowercase			(		self:			Any		)							->			Dict:




              """simple docstring"""


              UpperCamelCase_							,			UpperCamelCase_ 							=	self.prepare_config_and_inputs()
              return config, inputs_dict



@require_torch
class 				_UpperCamelCase       (						lowerCAmelCase_	,    lowerCAmelCase_	,    unittest.TestCase  ):
       _UpperCamelCase					:       Dict       = (CpmAntModel, CpmAntForCausalLM) if is_torch_available() else ()
       _UpperCamelCase					:       Union[str, Any]       = (
           {'''feature-extraction''': CpmAntModel, '''text-generation''': CpmAntForCausalLM} if is_torch_available() else {}
       )

       _UpperCamelCase					:       Optional[Any]       = False
       _UpperCamelCase					:       Dict       = False
       _UpperCamelCase					:       Any       = False
       _UpperCamelCase					:       Optional[Any]       = False
       _UpperCamelCase					:       Any       = False

       def lowercase			(		self:			Union[str, Any]		)							->			str:




              """simple docstring"""


              UpperCamelCase_ 							=	CpmAntModelTester(self		)
              UpperCamelCase_ 							=	ConfigTester(self       , config_class=_SCREAMING_SNAKE_CASE		)

       def lowercase			(		self:			str		)							->			Any:




              """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 lowercase			(		self:			Any		)							->			List[str]:




              """simple docstring"""


              unittest.skip("CPMAnt doesn't support input_embeds."		)(self.test_inputs_embeds		)

       def lowercase			(		self:			List[str]		)							->			Optional[Any]:




              """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.\
                 So is attentions. We strongly recommand you use loss to tune model."		)(self.test_retain_grad_hidden_states_attentions		)

       def lowercase			(		self:			List[str]		)							->			Optional[Any]:




              """simple docstring"""


              UpperCamelCase_							,			UpperCamelCase_ 							=	self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_cpmant_model(_SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE		)



       def lowercase			(		self:			Union[str, Any]		)							->			List[str]:




              """simple docstring"""


              UpperCamelCase_							,			UpperCamelCase_ 							=	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  ):

       @tooslow
       def lowercase			(		self:			List[Any]		)							->			List[Any]:




              """simple docstring"""


              UpperCamelCase_ 							=	"今天天气真好！"
              UpperCamelCase_ 							=	"openbmb/cpm-ant-10b"
              UpperCamelCase_ 							=	CpmAntModel.from_pretrained(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_ 							=	CpmAntTokenizer.from_pretrained(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_ 							=	tokenizer(_SCREAMING_SNAKE_CASE       , return_tensors="pt"		)
              UpperCamelCase_ 							=	model(**_SCREAMING_SNAKE_CASE		).last_hidden_state

              UpperCamelCase_ 							=	torch.tensor(
                  [[[6.17_08, 5.92_44, 1.08_35], [6.52_07, 6.28_93, -11.33_24], [-1.01_07, -0.05_76, -5.95_77]]]       , )
              self.assertTrue(torch.allclose(hidden_states[:, :3, :3]       , _SCREAMING_SNAKE_CASE       , atol=1e-2		)		)



@require_torch
class 				_UpperCamelCase       (						unittest.TestCase  ):

       @tooslow
       def lowercase			(		self:			Optional[Any]		)							->			Union[str, Any]:




              """simple docstring"""


              UpperCamelCase_ 							=	"今天天气真好！"
              UpperCamelCase_ 							=	"openbmb/cpm-ant-10b"
              UpperCamelCase_ 							=	CpmAntForCausalLM.from_pretrained(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_ 							=	CpmAntTokenizer.from_pretrained(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_ 							=	tokenizer(_SCREAMING_SNAKE_CASE       , return_tensors="pt"		)
              UpperCamelCase_ 							=	model(**_SCREAMING_SNAKE_CASE		).logits

              UpperCamelCase_ 							=	torch.tensor(
                  [[[-6.42_67, -6.40_83, -6.39_58], [-5.88_02, -5.94_47, -5.78_11], [-5.38_96, -5.48_20, -5.42_95]]]       , )
              self.assertTrue(torch.allclose(hidden_states[:, :3, :3]       , _SCREAMING_SNAKE_CASE       , atol=1e-2		)		)

       @tooslow
       def lowercase			(		self:			Dict		)							->			Optional[int]:




              """simple docstring"""


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



       @tooslow
       def lowercase			(		self:			List[str]		)							->			Optional[Any]:




              """simple docstring"""


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

