





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      :    Optional[Any]      ,					a      :    str      ,					a      :    Optional[int]=3      ,					a      :    Optional[int]=7      ,					a      :    int=True      ,					a      :    List[Any]=True      ,					a      :    Tuple=False      ,					a      :    List[str]=True      ,					a      :    Dict=99      ,					a      :    str=32      ,					a      :    Union[str, Any]=5      ,					a      :    Union[str, Any]=4      ,					a      :    Union[str, Any]=37      ,					a      :    Tuple="gelu"      ,					a      :    Optional[Any]=0.1      ,					a      :    Tuple=0.1      ,					a      :    Tuple=512      ,					a      :    Optional[int]=16      ,					a      :    Any=2      ,					a      :    Optional[Any]=0.02      ,					a      :    str=3      ,					a      :    Optional[Any]=4      ,					a      :    Optional[int]=None      ,					)->       int:





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





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





     """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=a      ,					initializer_range=self.initializer_range      ,					pad_token_id=1      ,					new_decoder_architecture=a      ,					)


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Optional[Any]      ,					a      :    List[str]      ,					a      :    Union[str, Any]      ,					a      :    List[str]      ,					a      :    Tuple      ,					a      :    Dict      ,					a      :    Optional[Any]      ,					a      :    List[Any]	)->       Union[str, Any]:





     """simple docstring"""







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


   def 	SCREAMING_SNAKE_CASE_							(   self      :    List[Any]      ,					a      :    str      ,					a      :    Union[str, Any]      ,					a      :    int      ,					a      :    Any      ,					a      :    Dict      ,					a      :    Optional[Any]      ,					a      :    List[Any]      ,					a      :    str      ,					a      :    Any      ,					)->       Optional[Any]:





     """simple docstring"""







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


   def 	SCREAMING_SNAKE_CASE_							(   self      :    List[Any]      ,					a      :    int      ,					a      :    int      ,					a      :    Dict      ,					a      :    Optional[int]      ,					a      :    Dict      ,					a      :    Dict      ,					a      :    int      ,					a      :    Dict      ,					a      :    Optional[int]      ,					)->       Tuple:





     """simple docstring"""







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


   def 	SCREAMING_SNAKE_CASE_							(   self      :    List[str]      ,					a      :    Any      ,					a      :    Dict      ,					a      :    Union[str, Any]      ,					a      :    Dict      ,					a      :    str      ,					a      :    int      ,					a      :    Dict      ,					a      :    Any      ,					a      :    str      ,					)->       Union[str, Any]:





     """simple docstring"""







     lowercase__           =  True
     lowercase__           =  True
     lowercase__           =  FalconForCausalLM(config=a	)
     model.to(a	)
     model.eval()

     # first forward pass
     lowercase__           =  model(
         a      ,					attention_mask=a      ,					encoder_hidden_states=a      ,					encoder_attention_mask=a      ,					use_cache=a      ,					)
     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(
         a      ,					attention_mask=a      ,					encoder_hidden_states=a      ,					encoder_attention_mask=a      ,					output_hidden_states=a      ,					)['hidden_states'][0]
     lowercase__           =  model(
         a      ,					attention_mask=a      ,					encoder_hidden_states=a      ,					encoder_attention_mask=a      ,					past_key_values=a      ,					output_hidden_states=a      ,					)['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(a      ,					a      ,					atol=1E-3	)	)


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Dict	)->       Optional[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				):
   _UpperCamelCase   :		Tuple			 =	(
       (
           FalconModel,
           FalconForCausalLM,
           FalconForSequenceClassification,
           FalconForTokenClassification,
           FalconForQuestionAnswering,
       )
       if is_torch_available()
       else ()
   )
   _UpperCamelCase   :		Tuple			 =	(FalconForCausalLM,) if is_torch_available() else ()
   _UpperCamelCase   :		Dict			 =	(
       {
           'feature-extraction': FalconModel,
           'text-classification': FalconForSequenceClassification,
           'text-generation': FalconForCausalLM,
           'question-answering': FalconForQuestionAnswering,
           'token-classification': FalconForTokenClassification,
           'zero-shot': FalconForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )
   _UpperCamelCase   :		Union[str, Any]			 =	False
   _UpperCamelCase   :		Optional[Any]			 =	False


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Optional[int]	)->       Any:





     """simple docstring"""







     lowercase__           =  FalconModelTester(self	)
     lowercase__           =  ConfigTester(self      ,					config_class=a      ,					hidden_size=37	)


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Any	)->       int:





     """simple docstring"""







     self.config_tester.run_common_tests()


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Any	)->       Tuple:





     """simple docstring"""







     lowercase__           =  self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*a	)


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Dict	)->       Tuple:





     """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(a      ,					*a	)


   def 	SCREAMING_SNAKE_CASE_							(   self      :    List[str]	)->       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(a	)
     lowercase__           =  ids_tensor([self.model_tester.batch_size]      ,					self.model_tester.type_sequence_label_size	)
     lowercase__           =  FalconForSequenceClassification(a	)
     model.to(a	)
     model.eval()
     lowercase__           =  model(a      ,					attention_mask=a      ,					labels=a	)
     self.assertEqual(result.logits.shape      ,					(self.model_tester.batch_size, self.model_tester.num_labels)	)


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Optional[Any]	)->       Union[str, 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(a	)
     lowercase__           =  ids_tensor([self.model_tester.batch_size]      ,					self.model_tester.type_sequence_label_size	)
     lowercase__           =  FalconForSequenceClassification(a	)
     model.to(a	)
     model.eval()
     lowercase__           =  model(a      ,					attention_mask=a      ,					labels=a	)
     self.assertEqual(result.logits.shape      ,					(self.model_tester.batch_size, self.model_tester.num_labels)	)


   def 	SCREAMING_SNAKE_CASE_							(   self      :    List[str]	)->       str:





     """simple docstring"""







     lowercase__      ,   lowercase__           =  self.model_tester.prepare_config_and_inputs_for_common()
     lowercase__           =  input_dict['input_ids']
     lowercase__           =  FalconForCausalLM(a	)
     model.to(a	)
     model.eval()
     lowercase__           =  model(a      ,					use_cache=a	)
     lowercase__           =  input_ids.shape[0]
     lowercase__           =  model._convert_to_rw_cache(result.past_key_values	)
     lowercase__           =  model._convert_cache_to_standard_format(a      ,					a	)
     for layer in range(len(a	)	):
       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 	SCREAMING_SNAKE_CASE_							(   self      :    Union[str, Any]	)->       int:





     """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(a	)
     lowercase__           =  ids_tensor(
         [self.model_tester.batch_size, config.num_labels]      ,					self.model_tester.type_sequence_label_size	).to(torch.float	)
     lowercase__           =  FalconForSequenceClassification(a	)
     model.to(a	)
     model.eval()
     lowercase__           =  model(a      ,					attention_mask=a      ,					labels=a	)
     self.assertEqual(result.logits.shape      ,					(self.model_tester.batch_size, self.model_tester.num_labels)	)


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Tuple	)->       int:





     """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(a      ,					'use_cache'	):
         return

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

       # 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(a      ,					'decoder_layers'      ,					a	)
           or getattr(a      ,					'num_decoder_layers'      ,					a	)
           or config.num_hidden_layers
       )
       lowercase__           =  getattr(a      ,					'num_kv_heads'      ,					config.num_attention_heads	)
       lowercase__           =  getattr(a      ,					'd_model'      ,					config.hidden_size	)
       lowercase__           =  embed_dim // num_attention_heads

       lowercase__           =  outputs['past_key_values']
       self.assertEqual(len(a	)      ,					a	)

       lowercase__      ,   lowercase__           =  inputs['input_ids'].shape
       for i in range(a	):
         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 	SCREAMING_SNAKE_CASE_							(   self      :    Union[str, Any]	)->       Dict:





     """simple docstring"""







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

     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(**a      ,					do_sample=a      ,					max_new_tokens=19	)
     lowercase__           =  tokenizer.batch_decode(a	)[0]

     self.assertEqual(a      ,					a	)


   @slow
   def 	SCREAMING_SNAKE_CASE_							(   self      :    str	)->       Union[str, Any]:





     """simple docstring"""







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

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


   @slow
   def 	SCREAMING_SNAKE_CASE_							(   self      :    Dict	)->       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(a	)
         lowercase__           =  FalconForCausalLM.from_pretrained(a	)
         model.eval()
         model.to(device=a	)
         lowercase__           =  tokenizer('My favorite food is'      ,					return_tensors='pt'	).to(a	)

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









from collections.abc import Iterator, MutableMapping
from dataclasses import dataclass
from typing import Generic, TypeVar

lowercase_          =		TypeVar("""KEY""")
lowercase_          =		TypeVar("""VAL""")
@dataclass(frozen=UpperCAmelCase      ,     slots=UpperCAmelCase				)
class     SCREAMING_SNAKE_CASE       (Generic[KEY, VAL]				):
   _UpperCamelCase   :		KEY
   _UpperCamelCase   :		VAL






class     SCREAMING_SNAKE_CASE       (_Item				):


   def __init__(   self      :    Optional[int]	)->       None:





     """simple docstring"""







     super().__init__(a      ,					a	)


   def __bool__(   self      :    str	)->       bool:





     """simple docstring"""







     return False






lowercase_          =		_DeletedItem()
class     SCREAMING_SNAKE_CASE       (MutableMapping[KEY, VAL]				):


   def __init__(   self      :    Tuple      ,					a      :    int = 8      ,					a      :    float = 0.75	)->       None:





     """simple docstring"""







     lowercase__           =  initial_block_size
     lowercase__           =  [None] * initial_block_size
     assert 0.0 < capacity_factor < 1.0
     lowercase__           =  capacity_factor
     lowercase__           =  0


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Any      ,					a      :    KEY	)->       int:





     """simple docstring"""







     return hash(a	) % len(self._buckets	)


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Optional[Any]      ,					a      :    int	)->       int:





     """simple docstring"""







     return (ind + 1) % len(self._buckets	)


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Any      ,					a      :    int      ,					a      :    KEY      ,					a      :    VAL	)->       bool:





     """simple docstring"""







     lowercase__           =  self._buckets[ind]
     if not stored:
       lowercase__           =  _Item(a      ,					a	)
       self._len += 1
       return True
     elif stored.key == key:
       lowercase__           =  _Item(a      ,					a	)
       return True
     else:
       return False


   def 	SCREAMING_SNAKE_CASE_							(   self      :    str	)->       bool:





     """simple docstring"""







     lowercase__           =  len(self._buckets	) * self._capacity_factor
     return len(self	) >= int(a	)


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Tuple	)->       bool:





     """simple docstring"""







     if len(self._buckets	) <= self._initial_block_size:
       return False
     lowercase__           =  len(self._buckets	) * self._capacity_factor / 2
     return len(self	) < limit


   def 	SCREAMING_SNAKE_CASE_							(   self      :    List[str]      ,					a      :    int	)->       None:





     """simple docstring"""







     lowercase__           =  self._buckets
     lowercase__           =  [None] * new_size
     lowercase__           =  0
     for item in old_buckets:
       if item:
         self._add_item(item.key      ,					item.val	)


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Optional[int]	)->       None:





     """simple docstring"""







     self._resize(len(self._buckets	) * 2	)


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Optional[int]	)->       None:





     """simple docstring"""







     self._resize(len(self._buckets	) // 2	)


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Any      ,					a      :    KEY	)->       Iterator[int]:





     """simple docstring"""







     lowercase__           =  self._get_bucket_index(a	)
     for _ in range(len(self._buckets	)	):
       yield ind
       lowercase__           =  self._get_next_ind(a	)


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Tuple      ,					a      :    KEY      ,					a      :    VAL	)->       None:





     """simple docstring"""







     for ind in self._iterate_buckets(a	):
       if self._try_set(a      ,					a      ,					a	):
         break


   def __setitem__(   self      :    List[Any]      ,					a      :    KEY      ,					a      :    VAL	)->       None:





     """simple docstring"""







     if self._is_full():
       self._size_up()

     self._add_item(a      ,					a	)


   def __delitem__(   self      :    str      ,					a      :    KEY	)->       None:





     """simple docstring"""







     for ind in self._iterate_buckets(a	):
       lowercase__           =  self._buckets[ind]
       if item is None:
         raise KeyError(a	)
       if item is _deleted:
         continue
       if item.key == key:
         lowercase__           =  _deleted
         self._len -= 1
         break
     if self._is_sparse():
       self._size_down()


   def __getitem__(   self      :    List[str]      ,					a      :    KEY	)->       VAL:





     """simple docstring"""







     for ind in self._iterate_buckets(a	):
       lowercase__           =  self._buckets[ind]
       if item is None:
         break
       if item is _deleted:
         continue
       if item.key == key:
         return item.val
     raise KeyError(a	)


   def __len__(   self      :    Tuple	)->       int:





     """simple docstring"""







     return self._len


   def __iter__(   self      :    int	)->       Iterator[KEY]:





     """simple docstring"""







     yield from (item.key for item in self._buckets if item)


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





     """simple docstring"""







     lowercase__           =  ' ,'.join(
         f"""{item.key}: {item.val}""" for item in self._buckets if item	)
     return f"""HashMap({val_string})"""


