



"""simple docstring"""
import unittest

from transformers import AutoTokenizer, FalconConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch

       from transformers import (
           FalconForCausalLM,
           FalconForQuestionAnswering,
           FalconForSequenceClassification,
           FalconForTokenClassification,
           FalconModel,
       )


class    SCREAMING_SNAKE_CASE__ :

    def __init__(				self			:				str ,  lowerCAmelCase_			:				Any ,  lowerCAmelCase_			:				Union[str, Any]=3 ,  lowerCAmelCase_			:				Dict=7 ,  lowerCAmelCase_			:				Any=True ,  lowerCAmelCase_			:				Union[str, Any]=True ,  lowerCAmelCase_			:				List[Any]=False ,  lowerCAmelCase_			:				str=True ,  lowerCAmelCase_			:				Dict=9_9 ,  lowerCAmelCase_			:				Optional[int]=3_2 ,  lowerCAmelCase_			:				List[str]=5 ,  lowerCAmelCase_			:				List[str]=4 ,  lowerCAmelCase_			:				Optional[int]=3_7 ,  lowerCAmelCase_			:				Any="gelu" ,  lowerCAmelCase_			:				List[str]=0.1 ,  lowerCAmelCase_			:				Optional[Any]=0.1 ,  lowerCAmelCase_			:				Optional[int]=5_1_2 ,  lowerCAmelCase_			:				Optional[Any]=1_6 ,  lowerCAmelCase_			:				Any=2 ,  lowerCAmelCase_			:				Optional[int]=0.02 ,  lowerCAmelCase_			:				int=3 ,  lowerCAmelCase_			:				Optional[int]=4 ,  lowerCAmelCase_			:				Optional[int]=None ,  ):







      """simple docstring"""







      lowercase_								= parent
      lowercase_								= batch_size
      lowercase_								= seq_length
      lowercase_								= is_training
      lowercase_								= use_input_mask
      lowercase_								= use_token_type_ids
      lowercase_								= use_labels
      lowercase_								= vocab_size
      lowercase_								= hidden_size
      lowercase_								= num_hidden_layers
      lowercase_								= num_attention_heads
      lowercase_								= intermediate_size
      lowercase_								= hidden_act
      lowercase_								= hidden_dropout_prob
      lowercase_								= attention_probs_dropout_prob
      lowercase_								= max_position_embeddings
      lowercase_								= type_vocab_size
      lowercase_								= type_sequence_label_size
      lowercase_								= initializer_range
      lowercase_								= num_labels
      lowercase_								= num_choices
      lowercase_								= scope

    def 					_UpperCAmelCase	(				self			:				str):







      """simple docstring"""







      lowercase_								= ids_tensor([self.batch_size, self.seq_length] ,  self.vocab_size)

      lowercase_								= None
      if self.use_input_mask:
        lowercase_								= random_attention_mask([self.batch_size, self.seq_length])

      lowercase_								= None

      lowercase_								= None
      lowercase_								= None
      lowercase_								= None
      if self.use_labels:
        lowercase_								= ids_tensor([self.batch_size] ,  self.type_sequence_label_size)
        lowercase_								= ids_tensor([self.batch_size, self.seq_length] ,  self.num_labels)
        lowercase_								= ids_tensor([self.batch_size] ,  self.num_choices)

      lowercase_								= self.get_config()

      return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels

    def 					_UpperCAmelCase	(				self			:				Tuple):







      """simple docstring"""







      return FalconConfig(
          vocab_size=self.vocab_size ,  hidden_size=self.hidden_size ,  num_hidden_layers=self.num_hidden_layers ,  num_attention_heads=self.num_attention_heads ,  intermediate_size=self.intermediate_size ,  hidden_act=self.hidden_act ,  hidden_dropout_prob=self.hidden_dropout_prob ,  attention_probs_dropout_prob=self.attention_probs_dropout_prob ,  max_position_embeddings=self.max_position_embeddings ,  type_vocab_size=self.type_vocab_size ,  is_decoder=lowerCAmelCase_ ,  initializer_range=self.initializer_range ,  pad_token_id=1 ,  new_decoder_architecture=lowerCAmelCase_ ,  )

    def 					_UpperCAmelCase	(				self			:				List[str] ,  lowerCAmelCase_			:				Any ,  lowerCAmelCase_			:				Tuple ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				Optional[int] ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				Tuple ,  lowerCAmelCase_			:				str):







      """simple docstring"""







      lowercase_								= FalconModel(config=lowerCAmelCase_)
      model.to(lowerCAmelCase_)
      model.eval()
      lowercase_								= model(lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_)
      lowercase_								= model(lowerCAmelCase_)
      self.parent.assertEqual(result.last_hidden_state.shape ,  (self.batch_size, self.seq_length, self.hidden_size))

    def 					_UpperCAmelCase	(				self			:				Tuple ,  lowerCAmelCase_			:				Tuple ,  lowerCAmelCase_			:				Tuple ,  lowerCAmelCase_			:				Optional[Any] ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				Dict ,  lowerCAmelCase_			:				Tuple ,  lowerCAmelCase_			:				List[str] ,  lowerCAmelCase_			:				Tuple ,  lowerCAmelCase_			:				Dict ,  ):







      """simple docstring"""







      lowercase_								= True
      lowercase_								= FalconModel(lowerCAmelCase_)
      model.to(lowerCAmelCase_)
      model.eval()
      lowercase_								= model(
          lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_ ,  encoder_hidden_states=lowerCAmelCase_ ,  encoder_attention_mask=lowerCAmelCase_ ,  )
      lowercase_								= model(
          lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_ ,  encoder_hidden_states=lowerCAmelCase_ ,  )
      lowercase_								= model(lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_)
      self.parent.assertEqual(result.last_hidden_state.shape ,  (self.batch_size, self.seq_length, self.hidden_size))

    def 					_UpperCAmelCase	(				self			:				Tuple ,  lowerCAmelCase_			:				List[str] ,  lowerCAmelCase_			:				str ,  lowerCAmelCase_			:				str ,  lowerCAmelCase_			:				Tuple ,  lowerCAmelCase_			:				Any ,  lowerCAmelCase_			:				Any ,  lowerCAmelCase_			:				Dict ,  lowerCAmelCase_			:				str ,  lowerCAmelCase_			:				Optional[Any] ,  ):







      """simple docstring"""







      lowercase_								= FalconForCausalLM(config=lowerCAmelCase_)
      model.to(lowerCAmelCase_)
      model.eval()
      lowercase_								= model(lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_ ,  labels=lowerCAmelCase_)
      self.parent.assertEqual(result.logits.shape ,  (self.batch_size, self.seq_length, self.vocab_size))

    def 					_UpperCAmelCase	(				self			:				List[str] ,  lowerCAmelCase_			:				Tuple ,  lowerCAmelCase_			:				Optional[Any] ,  lowerCAmelCase_			:				str ,  lowerCAmelCase_			:				Tuple ,  lowerCAmelCase_			:				Any ,  lowerCAmelCase_			:				Optional[Any] ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				Optional[Any] ,  lowerCAmelCase_			:				str ,  ):







      """simple docstring"""







      lowercase_								= True
      lowercase_								= True
      lowercase_								= FalconForCausalLM(config=lowerCAmelCase_)
      model.to(lowerCAmelCase_)
      model.eval()

      # first forward pass
      lowercase_								= model(
          lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_ ,  encoder_hidden_states=lowerCAmelCase_ ,  encoder_attention_mask=lowerCAmelCase_ ,  use_cache=lowerCAmelCase_ ,  )
      lowercase_								= outputs.past_key_values

      # create hypothetical multiple next token and extent to next_input_ids
      lowercase_								= ids_tensor((self.batch_size, 3) ,  config.vocab_size)
      lowercase_								= ids_tensor((self.batch_size, 3) ,  vocab_size=2)

      # append to next input_ids and
      lowercase_								= torch.cat([input_ids, next_tokens] ,  dim=-1)
      lowercase_								= torch.cat([input_mask, next_mask] ,  dim=-1)

      lowercase_								= model(
          lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_ ,  encoder_hidden_states=lowerCAmelCase_ ,  encoder_attention_mask=lowerCAmelCase_ ,  output_hidden_states=lowerCAmelCase_ ,  )["""hidden_states"""][0]
      lowercase_								= model(
          lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_ ,  encoder_hidden_states=lowerCAmelCase_ ,  encoder_attention_mask=lowerCAmelCase_ ,  past_key_values=lowerCAmelCase_ ,  output_hidden_states=lowerCAmelCase_ ,  )["""hidden_states"""][0]

      # select random slice
      lowercase_								= ids_tensor((1,) ,  output_from_past.shape[-1]).item()
      lowercase_								= output_from_no_past[:, -3:, random_slice_idx].detach()
      lowercase_								= output_from_past[:, :, random_slice_idx].detach()

      self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1])

      # test that outputs are equal for slice
      self.parent.assertTrue(torch.allclose(lowerCAmelCase_ ,  lowerCAmelCase_ ,  atol=1E-3))




    def 					_UpperCAmelCase	(				self			:				Union[str, Any]):







      """simple docstring"""







      lowercase_								= self.prepare_config_and_inputs()
      (
          (
          lowercase_
      )					,    (
          lowercase_
      )					,    (
          lowercase_
      )					,    (
          lowercase_
      )					,    (
          lowercase_
      )					,    (
          lowercase_
      )					,    (
          lowercase_
      )					,    
      )								= config_and_inputs
      lowercase_								= {"""input_ids""": input_ids, """attention_mask""": input_mask}
      return config, inputs_dict







@require_torch
class    SCREAMING_SNAKE_CASE__ (							__UpperCAmelCase      ,  __UpperCAmelCase      ,  __UpperCAmelCase      ,  unittest.TestCase   ):
    lowercase__      =   (
        (
            FalconModel,
            FalconForCausalLM,
            FalconForSequenceClassification,
            FalconForTokenClassification,
            FalconForQuestionAnswering,
        )
        if is_torch_available()
        else ()
    )
    lowercase__      =   (FalconForCausalLM,) if is_torch_available() else ()
    lowercase__      =   (
        {
            "feature-extraction": FalconModel,
            "text-classification": FalconForSequenceClassification,
            "text-generation": FalconForCausalLM,
            "question-answering": FalconForQuestionAnswering,
            "token-classification": FalconForTokenClassification,
            "zero-shot": FalconForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    lowercase__      =   False
    lowercase__      =   False

    def 					_UpperCAmelCase	(				self			:				str):







      """simple docstring"""







      lowercase_								= FalconModelTester(self)
      lowercase_								= ConfigTester(self ,  config_class=lowerCAmelCase_ ,  hidden_size=3_7)

    def 					_UpperCAmelCase	(				self			:				List[str]):







      """simple docstring"""







      self.config_tester.run_common_tests()

    def 					_UpperCAmelCase	(				self			:				List[Any]):







      """simple docstring"""







      lowercase_								= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*lowerCAmelCase_)

    def 					_UpperCAmelCase	(				self			:				int):







      """simple docstring"""







      lowercase_					,    *lowercase_								= self.model_tester.prepare_config_and_inputs()
      for alibi in [True, False]:
        lowercase_								= alibi
        self.model_tester.create_and_check_model(lowerCAmelCase_ ,  *lowerCAmelCase_)

    def 					_UpperCAmelCase	(				self			:				Any):







      """simple docstring"""







      lowercase_					,    lowercase_								= self.model_tester.prepare_config_and_inputs_for_common()
      lowercase_								= 3
      lowercase_								= input_dict["""input_ids"""]
      lowercase_								= input_ids.ne(1).to(lowerCAmelCase_)
      lowercase_								= ids_tensor([self.model_tester.batch_size] ,  self.model_tester.type_sequence_label_size)
      lowercase_								= FalconForSequenceClassification(lowerCAmelCase_)
      model.to(lowerCAmelCase_)
      model.eval()
      lowercase_								= model(lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_ ,  labels=lowerCAmelCase_)
      self.assertEqual(result.logits.shape ,  (self.model_tester.batch_size, self.model_tester.num_labels))

    def 					_UpperCAmelCase	(				self			:				List[Any]):







      """simple docstring"""







      lowercase_					,    lowercase_								= self.model_tester.prepare_config_and_inputs_for_common()
      lowercase_								= 3
      lowercase_								= """single_label_classification"""
      lowercase_								= input_dict["""input_ids"""]
      lowercase_								= input_ids.ne(1).to(lowerCAmelCase_)
      lowercase_								= ids_tensor([self.model_tester.batch_size] ,  self.model_tester.type_sequence_label_size)
      lowercase_								= FalconForSequenceClassification(lowerCAmelCase_)
      model.to(lowerCAmelCase_)
      model.eval()
      lowercase_								= model(lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_ ,  labels=lowerCAmelCase_)
      self.assertEqual(result.logits.shape ,  (self.model_tester.batch_size, self.model_tester.num_labels))

    def 					_UpperCAmelCase	(				self			:				str):







      """simple docstring"""







      lowercase_					,    lowercase_								= self.model_tester.prepare_config_and_inputs_for_common()
      lowercase_								= input_dict["""input_ids"""]
      lowercase_								= FalconForCausalLM(lowerCAmelCase_)
      model.to(lowerCAmelCase_)
      model.eval()
      lowercase_								= model(lowerCAmelCase_ ,  use_cache=lowerCAmelCase_)
      lowercase_								= input_ids.shape[0]
      lowercase_								= model._convert_to_rw_cache(result.past_key_values)
      lowercase_								= model._convert_cache_to_standard_format(lowerCAmelCase_ ,  lowerCAmelCase_)
      for layer in range(len(lowerCAmelCase_)):
        for tensor_idx in range(2):
          self.assertTrue(rw_cache[layer][tensor_idx].ndim == 3)
          self.assertTrue(result.past_key_values[layer][tensor_idx].ndim == 4)
          self.assertTrue(
              torch.all(result.past_key_values[layer][tensor_idx] == standard_cache[layer][tensor_idx]))

    def 					_UpperCAmelCase	(				self			:				List[str]):







      """simple docstring"""







      lowercase_					,    lowercase_								= self.model_tester.prepare_config_and_inputs_for_common()
      lowercase_								= 3
      lowercase_								= """multi_label_classification"""
      lowercase_								= input_dict["""input_ids"""]
      lowercase_								= input_ids.ne(1).to(lowerCAmelCase_)
      lowercase_								= ids_tensor(
          [self.model_tester.batch_size, config.num_labels] ,  self.model_tester.type_sequence_label_size).to(torch.float)
      lowercase_								= FalconForSequenceClassification(lowerCAmelCase_)
      model.to(lowerCAmelCase_)
      model.eval()
      lowercase_								= model(lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_ ,  labels=lowerCAmelCase_)
      self.assertEqual(result.logits.shape ,  (self.model_tester.batch_size, self.model_tester.num_labels))




    def 					_UpperCAmelCase	(				self			:				List[str]):







      """simple docstring"""







      for model_class in self.all_generative_model_classes:
        lowercase_					,    lowercase_								= self.model_tester.prepare_config_and_inputs_for_common()

        # If it doesn't support cache, pass the test
        if not hasattr(lowerCAmelCase_ ,  """use_cache"""):
          return

        lowercase_								= model_class(lowerCAmelCase_).to(lowerCAmelCase_)
        if "use_cache" not in inputs:
          lowercase_								= True
        lowercase_								= model(**lowerCAmelCase_)

        # If "past_key_values" is not returned, pass the test (e.g. RWKV uses a different cache name and format)
        if "past_key_values" not in outputs:
          return

        lowercase_								= (
            getattr(lowerCAmelCase_ ,  """decoder_layers""" ,  lowerCAmelCase_)
            or getattr(lowerCAmelCase_ ,  """num_decoder_layers""" ,  lowerCAmelCase_)
            or config.num_hidden_layers
        )
        lowercase_								= getattr(lowerCAmelCase_ ,  """num_kv_heads""" ,  config.num_attention_heads)
        lowercase_								= getattr(lowerCAmelCase_ ,  """d_model""" ,  config.hidden_size)
        lowercase_								= embed_dim // num_attention_heads

        lowercase_								= outputs["""past_key_values"""]
        self.assertEqual(len(lowerCAmelCase_) ,  lowerCAmelCase_)

        lowercase_					,    lowercase_								= inputs["""input_ids"""].shape
        for i in range(lowerCAmelCase_):
          if config.new_decoder_architecture:
            lowercase_								= config.num_attention_heads
          elif config.multi_query:
            lowercase_								= 1
          self.assertEqual(len(past_kv[0]) ,  2)  # K V for the decoder = 2
          self.assertEqual(
              past_kv[i][0].shape ,  (batch_size, num_attention_heads, seq_length, per_head_embed_dim))
          self.assertEqual(
              past_kv[i][1].shape ,  (batch_size, num_attention_heads, seq_length, per_head_embed_dim))







@require_torch
class    SCREAMING_SNAKE_CASE__ (							unittest.TestCase   ):

    @slow
    def 					_UpperCAmelCase	(				self			:				Tuple):







      """simple docstring"""







      lowercase_								= AutoTokenizer.from_pretrained("""Rocketknight1/falcon-rw-1b""")
      lowercase_								= FalconForCausalLM.from_pretrained("""Rocketknight1/falcon-rw-1b""")
      model.eval()
      model.to(lowerCAmelCase_)
      lowercase_								= tokenizer("""My favorite food is""" ,  return_tensors="""pt""").to(lowerCAmelCase_)

      lowercase_								= (
          """My favorite food is pizza. I love it so much that I have a pizza party every year for my birthday."""
      )

      lowercase_								= model.generate(**lowerCAmelCase_ ,  do_sample=lowerCAmelCase_ ,  max_new_tokens=1_9)
      lowercase_								= tokenizer.batch_decode(lowerCAmelCase_)[0]

      self.assertEqual(lowerCAmelCase_ ,  lowerCAmelCase_)

    @slow
    def 					_UpperCAmelCase	(				self			:				List[Any]):







      """simple docstring"""







      for repo in ["Rocketknight1/tiny-random-falcon-7b", "Rocketknight1/tiny-random-falcon-40b"]:
        lowercase_								= AutoTokenizer.from_pretrained(lowerCAmelCase_)
        lowercase_								= FalconForCausalLM.from_pretrained(lowerCAmelCase_)
        model.eval()
        model.to(lowerCAmelCase_)
        lowercase_								= tokenizer("""My favorite food is""" ,  return_tensors="""pt""").to(lowerCAmelCase_)

        # We just test that these run without errors - the models are randomly initialized
        # and so the actual text outputs will be garbage
        model.generate(**lowerCAmelCase_ ,  do_sample=lowerCAmelCase_ ,  max_new_tokens=4)
        model.generate(**lowerCAmelCase_ ,  do_sample=lowerCAmelCase_ ,  max_new_tokens=4)
        model.generate(**lowerCAmelCase_ ,  num_beams=2 ,  max_new_tokens=4)




    @slow
    def 					_UpperCAmelCase	(				self			:				Optional[Any]):







      """simple docstring"""







      with torch.no_grad():
        for repo in [
            "Rocketknight1/falcon-rw-1b",
            "Rocketknight1/tiny-random-falcon-7b",
            "Rocketknight1/tiny-random-falcon-40b",
        ]:
          lowercase_								= AutoTokenizer.from_pretrained(lowerCAmelCase_)
          lowercase_								= FalconForCausalLM.from_pretrained(lowerCAmelCase_)
          model.eval()
          model.to(device=lowerCAmelCase_)
          lowercase_								= tokenizer("""My favorite food is""" ,  return_tensors="""pt""").to(lowerCAmelCase_)

          # Test results are the same with and without cache
          lowercase_								= model.generate(**lowerCAmelCase_ ,  do_sample=lowerCAmelCase_ ,  max_new_tokens=2_0 ,  use_cache=lowerCAmelCase_)
          lowercase_								= model.generate(**lowerCAmelCase_ ,  do_sample=lowerCAmelCase_ ,  max_new_tokens=2_0 ,  use_cache=lowerCAmelCase_)
          self.assertTrue((outputs_cache - outputs_no_cache).sum().item() == 0)






"""simple docstring"""
def 			_SCREAMING_SNAKE_CASE						(__lowerCAmelCase=2_81_23							) ->						Optional[Any]:







  '''simple docstring'''


  lowercase_								= [1] * (limit + 1)

  for i in range(2 ,  int(limit**0.5							) + 1							):
    sum_divs[i * i] += i
    for k in range(i + 1 ,  limit // i + 1							):
      sum_divs[k * i] += k + i

  lowercase_								= set()
  lowercase_								= 0

  for n in range(1 ,  limit + 1							):
    if sum_divs[n] > n:
      abundants.add(__lowerCAmelCase							)

    if not any((n - a in abundants) for a in abundants							):
      res += n

  return res


if __name__ == "__main__":
       print(solution())

