


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

     '''simple docstring'''





     def __init__(      self							:     int							,	lowerCamelCase__							:     Dict							,	lowerCamelCase__							:     str=3							,	lowerCamelCase__							:     List[str]=7							,	lowerCamelCase__							:     str=True							,	lowerCamelCase__							:     int=True							,	lowerCamelCase__							:     Optional[Any]=False							,	lowerCamelCase__							:     Optional[Any]=True							,	lowerCamelCase__							:     Optional[int]=99							,	lowerCamelCase__							:     Any=32							,	lowerCamelCase__							:     str=5							,	lowerCamelCase__							:     List[Any]=4							,	lowerCamelCase__							:     List[Any]=37							,	lowerCamelCase__							:     List[str]="gelu"							,	lowerCamelCase__							:     List[Any]=0.1							,	lowerCamelCase__							:     List[str]=0.1							,	lowerCamelCase__							:     str=512							,	lowerCamelCase__							:     str=16							,	lowerCamelCase__							:     Dict=2							,	lowerCamelCase__							:     Any=0.02							,	lowerCamelCase__							:     Dict=3							,	lowerCamelCase__							:     str=4							,	lowerCamelCase__							:     str=None							,	)  ->      Optional[Any]:





      '''simple docstring'''






      UpperCamelCase__					:    int          =   parent
      UpperCamelCase__					:    Optional[int]          =   batch_size
      UpperCamelCase__					:    List[Any]          =   seq_length
      UpperCamelCase__					:    str          =   is_training
      UpperCamelCase__					:    List[Any]          =   use_input_mask
      UpperCamelCase__					:    str          =   use_token_type_ids
      UpperCamelCase__					:    Any          =   use_labels
      UpperCamelCase__					:    str          =   vocab_size
      UpperCamelCase__					:    Optional[Any]          =   hidden_size
      UpperCamelCase__					:    Dict          =   num_hidden_layers
      UpperCamelCase__					:    int          =   num_attention_heads
      UpperCamelCase__					:    Union[str, Any]          =   intermediate_size
      UpperCamelCase__					:    Any          =   hidden_act
      UpperCamelCase__					:    Union[str, Any]          =   hidden_dropout_prob
      UpperCamelCase__					:    str          =   attention_probs_dropout_prob
      UpperCamelCase__					:    Optional[int]          =   max_position_embeddings
      UpperCamelCase__					:    Optional[int]          =   type_vocab_size
      UpperCamelCase__					:    List[str]          =   type_sequence_label_size
      UpperCamelCase__					:    List[str]          =   initializer_range
      UpperCamelCase__					:    Any          =   num_labels
      UpperCamelCase__					:    Tuple          =   num_choices
      UpperCamelCase__					:    Dict          =   scope


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





      '''simple docstring'''






      UpperCamelCase__					:    List[Any]          =   ids_tensor([self.batch_size, self.seq_length]							,	self.vocab_size			)

      UpperCamelCase__					:    List[str]          =   None
      if self.use_input_mask:
       UpperCamelCase__					:    Optional[int]          =   random_attention_mask([self.batch_size, self.seq_length]			)

      UpperCamelCase__					:    Any          =   None

      UpperCamelCase__					:    Any          =   None
      UpperCamelCase__					:    Union[str, Any]          =   None
      UpperCamelCase__					:    List[str]          =   None
      if self.use_labels:
       UpperCamelCase__					:    List[Any]          =   ids_tensor([self.batch_size]							,	self.type_sequence_label_size			)
       UpperCamelCase__					:    Any          =   ids_tensor([self.batch_size, self.seq_length]							,	self.num_labels			)
       UpperCamelCase__					:    Optional[int]          =   ids_tensor([self.batch_size]							,	self.num_choices			)

      UpperCamelCase__					:    List[str]          =   self.get_config()

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


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





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


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





      '''simple docstring'''






      UpperCamelCase__					:    Tuple          =   FalconModel(config=__UpperCamelCase			)
      model.to(__UpperCamelCase			)
      model.eval()
      UpperCamelCase__					:    List[Any]          =   model(__UpperCamelCase							,	attention_mask=__UpperCamelCase			)
      UpperCamelCase__					:    Union[str, Any]          =   model(__UpperCamelCase			)
      self.parent.assertEqual(result.last_hidden_state.shape							,	(self.batch_size, self.seq_length, self.hidden_size)			)


     def      UpperCAmelCase__	(      self							:     List[str]							,	lowerCamelCase__							:     Optional[int]							,	lowerCamelCase__							:     Dict							,	lowerCamelCase__							:     List[Any]							,	lowerCamelCase__							:     int							,	lowerCamelCase__							:     int							,	lowerCamelCase__							:     int							,	lowerCamelCase__							:     List[Any]							,	lowerCamelCase__							:     Tuple							,	lowerCamelCase__							:     Optional[Any]							,	)  ->      List[str]:





      '''simple docstring'''






      UpperCamelCase__					:    str          =   True
      UpperCamelCase__					:    Any          =   FalconModel(__UpperCamelCase			)
      model.to(__UpperCamelCase			)
      model.eval()
      UpperCamelCase__					:    List[Any]          =   model(
          __UpperCamelCase							,	attention_mask=__UpperCamelCase							,	encoder_hidden_states=__UpperCamelCase							,	encoder_attention_mask=__UpperCamelCase							,	)
      UpperCamelCase__					:    str          =   model(
          __UpperCamelCase							,	attention_mask=__UpperCamelCase							,	encoder_hidden_states=__UpperCamelCase							,	)
      UpperCamelCase__					:    Dict          =   model(__UpperCamelCase							,	attention_mask=__UpperCamelCase			)
      self.parent.assertEqual(result.last_hidden_state.shape							,	(self.batch_size, self.seq_length, self.hidden_size)			)


     def      UpperCAmelCase__	(      self							:     Any							,	lowerCamelCase__							:     Union[str, Any]							,	lowerCamelCase__							:     Union[str, Any]							,	lowerCamelCase__							:     Tuple							,	lowerCamelCase__							:     Optional[int]							,	lowerCamelCase__							:     Union[str, Any]							,	lowerCamelCase__							:     Tuple							,	lowerCamelCase__							:     Dict							,	lowerCamelCase__							:     Optional[int]							,	lowerCamelCase__							:     int							,	)  ->      List[Any]:





      '''simple docstring'''






      UpperCamelCase__					:    int          =   FalconForCausalLM(config=__UpperCamelCase			)
      model.to(__UpperCamelCase			)
      model.eval()
      UpperCamelCase__					:    List[str]          =   model(__UpperCamelCase							,	attention_mask=__UpperCamelCase							,	labels=__UpperCamelCase			)
      self.parent.assertEqual(result.logits.shape							,	(self.batch_size, self.seq_length, self.vocab_size)			)


     def      UpperCAmelCase__	(      self							:     Optional[Any]							,	lowerCamelCase__							:     Tuple							,	lowerCamelCase__							:     Optional[Any]							,	lowerCamelCase__							:     Any							,	lowerCamelCase__							:     str							,	lowerCamelCase__							:     Optional[Any]							,	lowerCamelCase__							:     int							,	lowerCamelCase__							:     List[str]							,	lowerCamelCase__							:     Union[str, Any]							,	lowerCamelCase__							:     Union[str, Any]							,	)  ->      Dict:





      '''simple docstring'''






      UpperCamelCase__					:    Tuple          =   True
      UpperCamelCase__					:    Any          =   True
      UpperCamelCase__					:    Dict          =   FalconForCausalLM(config=__UpperCamelCase			)
      model.to(__UpperCamelCase			)
      model.eval()

      # first forward pass
      UpperCamelCase__					:    Any          =   model(
          __UpperCamelCase							,	attention_mask=__UpperCamelCase							,	encoder_hidden_states=__UpperCamelCase							,	encoder_attention_mask=__UpperCamelCase							,	use_cache=__UpperCamelCase							,	)
      UpperCamelCase__					:    str          =   outputs.past_key_values

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

      # append to next input_ids and
      UpperCamelCase__					:    List[Any]          =   torch.cat([input_ids, next_tokens]							,	dim=-1			)
      UpperCamelCase__					:    Union[str, Any]          =   torch.cat([input_mask, next_mask]							,	dim=-1			)

      UpperCamelCase__					:    Dict          =   model(
          __UpperCamelCase							,	attention_mask=__UpperCamelCase							,	encoder_hidden_states=__UpperCamelCase							,	encoder_attention_mask=__UpperCamelCase							,	output_hidden_states=__UpperCamelCase							,	)['''hidden_states'''][0]
      UpperCamelCase__					:    Union[str, Any]          =   model(
          __UpperCamelCase							,	attention_mask=__UpperCamelCase							,	encoder_hidden_states=__UpperCamelCase							,	encoder_attention_mask=__UpperCamelCase							,	past_key_values=__UpperCamelCase							,	output_hidden_states=__UpperCamelCase							,	)['''hidden_states'''][0]

      # select random slice
      UpperCamelCase__					:    int          =   ids_tensor((1,)							,	output_from_past.shape[-1]			).item()
      UpperCamelCase__					:    Optional[int]          =   output_from_no_past[:, -3:, random_slice_idx].detach()
      UpperCamelCase__					:    List[Any]          =   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(__UpperCamelCase							,	__UpperCamelCase							,	atol=1E-3			)			)


     def      UpperCAmelCase__	(      self							:     List[Any]			)  ->      List[Any]:





      '''simple docstring'''






      UpperCamelCase__					:    Optional[int]          =   self.prepare_config_and_inputs()
      (
          (
          UpperCamelCase__   
      )       ,				(
          UpperCamelCase__   
      )       ,				(
          UpperCamelCase__   
      )       ,				(
          UpperCamelCase__   
      )       ,				(
          UpperCamelCase__   
      )       ,				(
          UpperCamelCase__   
      )       ,				(
          UpperCamelCase__   
      )       ,				   
      )					:    List[Any]          =   config_and_inputs
      UpperCamelCase__					:    Optional[int]          =   {'''input_ids''': input_ids, '''attention_mask''': input_mask}
      return config, inputs_dict



@require_torch
class       __magic_name__	(     __lowerCAmelCase					,				__lowerCAmelCase					,				__lowerCAmelCase					,				unittest.TestCase):

     '''simple docstring'''





     A:		Tuple			=      (
         (
             FalconModel,
             FalconForCausalLM,
             FalconForSequenceClassification,
             FalconForTokenClassification,
             FalconForQuestionAnswering,
         )
         if is_torch_available()
         else ()
     )
     A:		Dict			=      (FalconForCausalLM,) if is_torch_available() else ()
     A:		int			=      (
         {
             "feature-extraction": FalconModel,
             "text-classification": FalconForSequenceClassification,
             "text-generation": FalconForCausalLM,
             "question-answering": FalconForQuestionAnswering,
             "token-classification": FalconForTokenClassification,
             "zero-shot": FalconForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     A:		Union[str, Any]			=      False
     A:		Union[str, Any]			=      False


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





      '''simple docstring'''






      UpperCamelCase__					:    Optional[int]          =   FalconModelTester(self			)
      UpperCamelCase__					:    str          =   ConfigTester(self							,	config_class=__UpperCamelCase							,	hidden_size=37			)


     def      UpperCAmelCase__	(      self							:     List[Any]			)  ->      int:





      '''simple docstring'''






      self.config_tester.run_common_tests()


     def      UpperCAmelCase__	(      self							:     Tuple			)  ->      Tuple:





      '''simple docstring'''






      UpperCamelCase__					:    int          =   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*__UpperCamelCase			)


     def      UpperCAmelCase__	(      self							:     Tuple			)  ->      int:





      '''simple docstring'''






      UpperCamelCase__       ,				*UpperCamelCase__					:    Union[str, Any]          =   self.model_tester.prepare_config_and_inputs()
      for alibi in [True, False]:
       UpperCamelCase__					:    Tuple          =   alibi
       self.model_tester.create_and_check_model(__UpperCamelCase							,	*__UpperCamelCase			)


     def      UpperCAmelCase__	(      self							:     str			)  ->      Tuple:





      '''simple docstring'''






      UpperCamelCase__       ,				UpperCamelCase__					:    List[Any]          =   self.model_tester.prepare_config_and_inputs_for_common()
      UpperCamelCase__					:    Dict          =   3
      UpperCamelCase__					:    Dict          =   input_dict['''input_ids''']
      UpperCamelCase__					:    Union[str, Any]          =   input_ids.ne(1			).to(__UpperCamelCase			)
      UpperCamelCase__					:    Optional[Any]          =   ids_tensor([self.model_tester.batch_size]							,	self.model_tester.type_sequence_label_size			)
      UpperCamelCase__					:    Union[str, Any]          =   FalconForSequenceClassification(__UpperCamelCase			)
      model.to(__UpperCamelCase			)
      model.eval()
      UpperCamelCase__					:    Dict          =   model(__UpperCamelCase							,	attention_mask=__UpperCamelCase							,	labels=__UpperCamelCase			)
      self.assertEqual(result.logits.shape							,	(self.model_tester.batch_size, self.model_tester.num_labels)			)


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





      '''simple docstring'''






      UpperCamelCase__       ,				UpperCamelCase__					:    List[Any]          =   self.model_tester.prepare_config_and_inputs_for_common()
      UpperCamelCase__					:    int          =   3
      UpperCamelCase__					:    Tuple          =   '''single_label_classification'''
      UpperCamelCase__					:    str          =   input_dict['''input_ids''']
      UpperCamelCase__					:    List[str]          =   input_ids.ne(1			).to(__UpperCamelCase			)
      UpperCamelCase__					:    Tuple          =   ids_tensor([self.model_tester.batch_size]							,	self.model_tester.type_sequence_label_size			)
      UpperCamelCase__					:    List[Any]          =   FalconForSequenceClassification(__UpperCamelCase			)
      model.to(__UpperCamelCase			)
      model.eval()
      UpperCamelCase__					:    Union[str, Any]          =   model(__UpperCamelCase							,	attention_mask=__UpperCamelCase							,	labels=__UpperCamelCase			)
      self.assertEqual(result.logits.shape							,	(self.model_tester.batch_size, self.model_tester.num_labels)			)


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





      '''simple docstring'''






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





      '''simple docstring'''






      UpperCamelCase__       ,				UpperCamelCase__					:    Tuple          =   self.model_tester.prepare_config_and_inputs_for_common()
      UpperCamelCase__					:    str          =   3
      UpperCamelCase__					:    Dict          =   '''multi_label_classification'''
      UpperCamelCase__					:    Union[str, Any]          =   input_dict['''input_ids''']
      UpperCamelCase__					:    Optional[Any]          =   input_ids.ne(1			).to(__UpperCamelCase			)
      UpperCamelCase__					:    int          =   ids_tensor(
          [self.model_tester.batch_size, config.num_labels]							,	self.model_tester.type_sequence_label_size			).to(torch.float			)
      UpperCamelCase__					:    Tuple          =   FalconForSequenceClassification(__UpperCamelCase			)
      model.to(__UpperCamelCase			)
      model.eval()
      UpperCamelCase__					:    int          =   model(__UpperCamelCase							,	attention_mask=__UpperCamelCase							,	labels=__UpperCamelCase			)
      self.assertEqual(result.logits.shape							,	(self.model_tester.batch_size, self.model_tester.num_labels)			)


     def      UpperCAmelCase__	(      self							:     Dict			)  ->      Any:





      '''simple docstring'''






      for model_class in self.all_generative_model_classes:
       UpperCamelCase__       ,				UpperCamelCase__					:    Optional[int]          =   self.model_tester.prepare_config_and_inputs_for_common()

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

       UpperCamelCase__					:    Any          =   model_class(__UpperCamelCase			).to(__UpperCamelCase			)
       if "use_cache" not in inputs:
        UpperCamelCase__					:    Any          =   True
       UpperCamelCase__					:    List[Any]          =   model(**__UpperCamelCase			)

       # 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

       UpperCamelCase__					:    int          =   (
           getattr(__UpperCamelCase							,	'''decoder_layers'''							,	__UpperCamelCase			)
           or getattr(__UpperCamelCase							,	'''num_decoder_layers'''							,	__UpperCamelCase			)
           or config.num_hidden_layers
       )
       UpperCamelCase__					:    List[str]          =   getattr(__UpperCamelCase							,	'''num_kv_heads'''							,	config.num_attention_heads			)
       UpperCamelCase__					:    Dict          =   getattr(__UpperCamelCase							,	'''d_model'''							,	config.hidden_size			)
       UpperCamelCase__					:    Union[str, Any]          =   embed_dim // num_attention_heads

       UpperCamelCase__					:    List[str]          =   outputs['''past_key_values''']
       self.assertEqual(len(__UpperCamelCase			)							,	__UpperCamelCase			)

       UpperCamelCase__       ,				UpperCamelCase__					:    List[Any]          =   inputs['''input_ids'''].shape
       for i in range(__UpperCamelCase			):
        if config.new_decoder_architecture:
         UpperCamelCase__					:    List[str]          =   config.num_attention_heads
        elif config.multi_query:
         UpperCamelCase__					:    Tuple          =   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       __magic_name__	(     unittest.TestCase):

     '''simple docstring'''





     @slow
     def      UpperCAmelCase__	(      self							:     Optional[int]			)  ->      Any:





      '''simple docstring'''






      UpperCamelCase__					:    Union[str, Any]          =   AutoTokenizer.from_pretrained('''Rocketknight1/falcon-rw-1b'''			)
      UpperCamelCase__					:    Optional[Any]          =   FalconForCausalLM.from_pretrained('''Rocketknight1/falcon-rw-1b'''			)
      model.eval()
      model.to(__UpperCamelCase			)
      UpperCamelCase__					:    Optional[int]          =   tokenizer('''My favorite food is'''							,	return_tensors='''pt'''			).to(__UpperCamelCase			)

      UpperCamelCase__					:    int          =   (
          '''My favorite food is pizza. I love it so much that I have a pizza party every year for my birthday.'''
      )

      UpperCamelCase__					:    int          =   model.generate(**__UpperCamelCase							,	do_sample=__UpperCamelCase							,	max_new_tokens=19			)
      UpperCamelCase__					:    Dict          =   tokenizer.batch_decode(__UpperCamelCase			)[0]

      self.assertEqual(__UpperCamelCase							,	__UpperCamelCase			)


     @slow
     def      UpperCAmelCase__	(      self							:     Any			)  ->      Any:





      '''simple docstring'''






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

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


     @slow
     def      UpperCAmelCase__	(      self							:     Optional[Any]			)  ->      List[str]:





      '''simple docstring'''






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

        # Test results are the same with and without cache
        UpperCamelCase__					:    int          =   model.generate(**__UpperCamelCase							,	do_sample=__UpperCamelCase							,	max_new_tokens=20							,	use_cache=__UpperCamelCase			)
        UpperCamelCase__					:    List[Any]          =   model.generate(**__UpperCamelCase							,	do_sample=__UpperCamelCase							,	max_new_tokens=20							,	use_cache=__UpperCamelCase			)
        self.assertTrue((outputs_cache - outputs_no_cache).sum().item() == 0			)






import inspect
import unittest

import torch
import torch.nn as nn

from accelerate.hooks import (
    AlignDevicesHook,
    ModelHook,
    SequentialHook,
    add_hook_to_module,
    attach_align_device_hook,
    remove_hook_from_module,
    remove_hook_from_submodules,
)
from accelerate.test_utils import require_multi_gpu

class       __magic_name__	(     nn.Module):


     def __init__(      self							:     Dict			)  ->      str:





      '''simple docstring'''






      super().__init__()
      UpperCamelCase__					:    int          =   nn.Linear(3							,	4			)
      UpperCamelCase__					:    str          =   nn.BatchNormad(4			)
      UpperCamelCase__					:    List[str]          =   nn.Linear(4							,	5			)


     def      UpperCAmelCase__	(      self							:     int							,	lowerCamelCase__							:     Dict			)  ->      Tuple:





      '''simple docstring'''






      return self.lineara(self.batchnorm(self.lineara(lowerCamelCase__			)			)			)



class       __magic_name__	(     __lowerCAmelCase):


     def      UpperCAmelCase__	(      self							:     Union[str, Any]							,	lowerCamelCase__							:     Any							,	*lowerCamelCase__							:     int							,	**lowerCamelCase__							:     List[str]			)  ->      Any:





      '''simple docstring'''






      return (args[0] + 1,) + args[1:], kwargs



class       __magic_name__	(     __lowerCAmelCase):


     def      UpperCAmelCase__	(      self							:     Any							,	lowerCamelCase__							:     Any							,	lowerCamelCase__							:     Any			)  ->      Any:





      '''simple docstring'''






      return output + 1



class       __magic_name__	(     unittest.TestCase):


     def      UpperCAmelCase__	(      self							:     Tuple			)  ->      Tuple:





      '''simple docstring'''






      UpperCamelCase__					:    Tuple          =   ModelForTest()
      UpperCamelCase__					:    Union[str, Any]          =   ModelHook()

      add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__			)
      self.assertEqual(test_model._hf_hook							,	lowerCamelCase__			)
      self.assertTrue(hasattr(lowerCamelCase__							,	'''_old_forward'''			)			)

      # Check adding the hook did not change the name or the signature
      self.assertEqual(test_model.forward.__name__							,	'''forward'''			)
      self.assertListEqual(list(inspect.signature(test_model.forward			).parameters			)							,	['''x''']			)

      remove_hook_from_module(lowerCamelCase__			)
      self.assertFalse(hasattr(lowerCamelCase__							,	'''_hf_hook'''			)			)
      self.assertFalse(hasattr(lowerCamelCase__							,	'''_old_forward'''			)			)


     def      UpperCAmelCase__	(      self							:     Dict			)  ->      str:





      '''simple docstring'''






      UpperCamelCase__					:    Dict          =   ModelForTest()
      UpperCamelCase__					:    List[str]          =   ModelHook()

      add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__			)
      add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__							,	append=lowerCamelCase__			)

      self.assertEqual(isinstance(test_model._hf_hook							,	lowerCamelCase__			)							,	lowerCamelCase__			)
      self.assertEqual(len(test_model._hf_hook.hooks			)							,	2			)
      self.assertTrue(hasattr(lowerCamelCase__							,	'''_old_forward'''			)			)

      # Check adding the hook did not change the name or the signature
      self.assertEqual(test_model.forward.__name__							,	'''forward'''			)
      self.assertListEqual(list(inspect.signature(test_model.forward			).parameters			)							,	['''x''']			)

      remove_hook_from_module(lowerCamelCase__			)
      self.assertFalse(hasattr(lowerCamelCase__							,	'''_hf_hook'''			)			)
      self.assertFalse(hasattr(lowerCamelCase__							,	'''_old_forward'''			)			)


     def      UpperCAmelCase__	(      self							:     List[str]			)  ->      int:





      '''simple docstring'''






      UpperCamelCase__					:    Optional[Any]          =   ModelForTest()
      UpperCamelCase__					:    List[str]          =   torch.randn(2							,	3			)
      UpperCamelCase__					:    int          =   test_model(x + 1			)
      UpperCamelCase__					:    List[Any]          =   test_model(x + 2			)

      UpperCamelCase__					:    Optional[int]          =   PreForwardHook()
      add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__			)
      UpperCamelCase__					:    List[Any]          =   test_model(lowerCamelCase__			)
      self.assertTrue(torch.allclose(lowerCamelCase__							,	lowerCamelCase__							,	atol=1E-5			)			)

      # Attaching a hook to a model when it already has one replaces, does not chain
      UpperCamelCase__					:    Optional[Any]          =   PreForwardHook()
      add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__			)
      UpperCamelCase__					:    List[str]          =   test_model(lowerCamelCase__			)
      self.assertTrue(torch.allclose(lowerCamelCase__							,	lowerCamelCase__							,	atol=1E-5			)			)

      # You need to use the sequential hook to chain two or more hooks
      UpperCamelCase__					:    Tuple          =   SequentialHook(PreForwardHook()							,	PreForwardHook()			)
      add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__			)

      UpperCamelCase__					:    Any          =   test_model(lowerCamelCase__			)
      assert torch.allclose(lowerCamelCase__							,	lowerCamelCase__							,	atol=1E-5			)


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





      '''simple docstring'''






      UpperCamelCase__					:    Dict          =   ModelForTest()
      UpperCamelCase__					:    str          =   torch.randn(2							,	3			)
      UpperCamelCase__					:    Any          =   test_model(lowerCamelCase__			)

      UpperCamelCase__					:    int          =   PostForwardHook()
      add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__			)
      UpperCamelCase__					:    Dict          =   test_model(lowerCamelCase__			)
      self.assertTrue(torch.allclose(lowerCamelCase__							,	output + 1							,	atol=1E-5			)			)

      # Attaching a hook to a model when it already has one replaces, does not chain
      UpperCamelCase__					:    Dict          =   PostForwardHook()
      add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__			)
      UpperCamelCase__					:    int          =   test_model(lowerCamelCase__			)
      self.assertTrue(torch.allclose(lowerCamelCase__							,	output + 1							,	atol=1E-5			)			)

      # You need to use the sequential hook to chain two or more hooks
      UpperCamelCase__					:    List[Any]          =   SequentialHook(PostForwardHook()							,	PostForwardHook()			)
      add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__			)

      UpperCamelCase__					:    Tuple          =   test_model(lowerCamelCase__			)
      assert torch.allclose(lowerCamelCase__							,	output + 2							,	atol=1E-5			)


     def      UpperCAmelCase__	(      self							:     List[Any]			)  ->      Dict:





      '''simple docstring'''






      UpperCamelCase__					:    List[Any]          =   ModelForTest()
      UpperCamelCase__					:    Tuple          =   torch.randn(2							,	3			)
      UpperCamelCase__					:    Tuple          =   test_model(lowerCamelCase__			)

      UpperCamelCase__					:    Union[str, Any]          =   PostForwardHook()
      add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__			)
      UpperCamelCase__					:    Any          =   test_model(lowerCamelCase__			)
      self.assertTrue(torch.allclose(lowerCamelCase__							,	output + 1			)			)
      self.assertTrue(outputa.requires_grad			)

      UpperCamelCase__					:    Optional[int]          =   True
      UpperCamelCase__					:    Optional[int]          =   test_model(lowerCamelCase__			)
      self.assertFalse(outputa.requires_grad			)


     @require_multi_gpu
     def      UpperCAmelCase__	(      self							:     Any			)  ->      Dict:





      '''simple docstring'''






      UpperCamelCase__					:    int          =   ModelForTest()
      # Everything is on CPU
      self.assertEqual(model.lineara.weight.device							,	torch.device('''cpu'''			)			)
      self.assertEqual(model.batchnorm.weight.device							,	torch.device('''cpu'''			)			)
      self.assertEqual(model.lineara.weight.device							,	torch.device('''cpu'''			)			)

      # This will move each submodule on different devices
      add_hook_to_module(model.lineara							,	AlignDevicesHook(execution_device=0			)			)
      add_hook_to_module(model.batchnorm							,	AlignDevicesHook(execution_device=0			)			)
      add_hook_to_module(model.lineara							,	AlignDevicesHook(execution_device=1			)			)

      self.assertEqual(model.lineara.weight.device							,	torch.device(0			)			)
      self.assertEqual(model.batchnorm.weight.device							,	torch.device(0			)			)
      self.assertEqual(model.batchnorm.running_mean.device							,	torch.device(0			)			)
      self.assertEqual(model.lineara.weight.device							,	torch.device(1			)			)

      # We can still make a forward pass. The input does not need to be on any particular device
      UpperCamelCase__					:    Tuple          =   torch.randn(2							,	3			)
      UpperCamelCase__					:    str          =   model(lowerCamelCase__			)
      self.assertEqual(output.device							,	torch.device(1			)			)

      # We can add a general hook to put back output on same device as input.
      add_hook_to_module(lowerCamelCase__							,	AlignDevicesHook(io_same_device=lowerCamelCase__			)			)
      UpperCamelCase__					:    Tuple          =   torch.randn(2							,	3			).to(0			)
      UpperCamelCase__					:    Optional[int]          =   model(lowerCamelCase__			)
      self.assertEqual(output.device							,	torch.device(0			)			)


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





      '''simple docstring'''






      UpperCamelCase__					:    Dict          =   ModelForTest()

      # Everything is on CPU
      self.assertEqual(model.lineara.weight.device							,	torch.device('''cpu'''			)			)
      self.assertEqual(model.batchnorm.weight.device							,	torch.device('''cpu'''			)			)
      self.assertEqual(model.lineara.weight.device							,	torch.device('''cpu'''			)			)

      # This will move each submodule on different devices
      UpperCamelCase__					:    int          =   {'''execution_device''': 0 if torch.cuda.is_available() else '''cpu''', '''offload''': True}

      add_hook_to_module(model.lineara							,	AlignDevicesHook(**lowerCamelCase__			)			)
      add_hook_to_module(model.batchnorm							,	AlignDevicesHook(**lowerCamelCase__			)			)
      add_hook_to_module(model.lineara							,	AlignDevicesHook(**lowerCamelCase__			)			)

      # Parameters have been offloaded, so on the meta device
      self.assertEqual(model.lineara.weight.device							,	torch.device('''meta'''			)			)
      self.assertEqual(model.batchnorm.weight.device							,	torch.device('''meta'''			)			)
      self.assertEqual(model.lineara.weight.device							,	torch.device('''meta'''			)			)
      # Buffers are not included in the offload by default, so are on the execution device
      UpperCamelCase__					:    int          =   torch.device(hook_kwargs['''execution_device''']			)
      self.assertEqual(model.batchnorm.running_mean.device							,	lowerCamelCase__			)

      UpperCamelCase__					:    Dict          =   torch.randn(2							,	3			)
      UpperCamelCase__					:    List[str]          =   model(lowerCamelCase__			)
      self.assertEqual(output.device							,	lowerCamelCase__			)

      # Removing hooks loads back the weights in the model.
      remove_hook_from_module(model.lineara			)
      remove_hook_from_module(model.batchnorm			)
      remove_hook_from_module(model.lineara			)
      self.assertEqual(model.lineara.weight.device							,	torch.device('''cpu'''			)			)
      self.assertEqual(model.batchnorm.weight.device							,	torch.device('''cpu'''			)			)
      self.assertEqual(model.lineara.weight.device							,	torch.device('''cpu'''			)			)

      # Now test with buffers included in the offload
      UpperCamelCase__					:    Optional[Any]          =   {
          '''execution_device''': 0 if torch.cuda.is_available() else '''cpu''',
          '''offload''': True,
          '''offload_buffers''': True,
      }

      add_hook_to_module(model.lineara							,	AlignDevicesHook(**lowerCamelCase__			)			)
      add_hook_to_module(model.batchnorm							,	AlignDevicesHook(**lowerCamelCase__			)			)
      add_hook_to_module(model.lineara							,	AlignDevicesHook(**lowerCamelCase__			)			)

      # Parameters have been offloaded, so on the meta device, buffers included
      self.assertEqual(model.lineara.weight.device							,	torch.device('''meta'''			)			)
      self.assertEqual(model.batchnorm.weight.device							,	torch.device('''meta'''			)			)
      self.assertEqual(model.lineara.weight.device							,	torch.device('''meta'''			)			)
      self.assertEqual(model.batchnorm.running_mean.device							,	torch.device('''meta'''			)			)

      UpperCamelCase__					:    List[Any]          =   torch.randn(2							,	3			)
      UpperCamelCase__					:    str          =   model(lowerCamelCase__			)
      self.assertEqual(output.device							,	lowerCamelCase__			)

      # Removing hooks loads back the weights in the model.
      remove_hook_from_module(model.lineara			)
      remove_hook_from_module(model.batchnorm			)
      remove_hook_from_module(model.lineara			)
      self.assertEqual(model.lineara.weight.device							,	torch.device('''cpu'''			)			)
      self.assertEqual(model.batchnorm.weight.device							,	torch.device('''cpu'''			)			)
      self.assertEqual(model.lineara.weight.device							,	torch.device('''cpu'''			)			)


     def      UpperCAmelCase__	(      self							:     str			)  ->      Optional[int]:





      '''simple docstring'''






      UpperCamelCase__					:    Optional[int]          =   ModelForTest()

      # Everything is on CPU
      self.assertEqual(model.lineara.weight.device							,	torch.device('''cpu'''			)			)
      self.assertEqual(model.batchnorm.weight.device							,	torch.device('''cpu'''			)			)
      self.assertEqual(model.lineara.weight.device							,	torch.device('''cpu'''			)			)

      # This will move each submodule on different devices
      UpperCamelCase__					:    Any          =   0 if torch.cuda.is_available() else '''cpu'''
      attach_align_device_hook(lowerCamelCase__							,	execution_device=lowerCamelCase__							,	offload=lowerCamelCase__			)

      # Parameters have been offloaded, so on the meta device
      self.assertEqual(model.lineara.weight.device							,	torch.device('''meta'''			)			)
      self.assertEqual(model.batchnorm.weight.device							,	torch.device('''meta'''			)			)
      self.assertEqual(model.lineara.weight.device							,	torch.device('''meta'''			)			)
      # Buffers are not included in the offload by default, so are on the execution device
      UpperCamelCase__					:    str          =   torch.device(lowerCamelCase__			)
      self.assertEqual(model.batchnorm.running_mean.device							,	lowerCamelCase__			)

      UpperCamelCase__					:    int          =   torch.randn(2							,	3			)
      UpperCamelCase__					:    Tuple          =   model(lowerCamelCase__			)
      self.assertEqual(output.device							,	lowerCamelCase__			)

      # Removing hooks loads back the weights in the model.
      remove_hook_from_submodules(lowerCamelCase__			)
      self.assertEqual(model.lineara.weight.device							,	torch.device('''cpu'''			)			)
      self.assertEqual(model.batchnorm.weight.device							,	torch.device('''cpu'''			)			)
      self.assertEqual(model.lineara.weight.device							,	torch.device('''cpu'''			)			)

      # Now test with buffers included in the offload
      attach_align_device_hook(lowerCamelCase__							,	execution_device=lowerCamelCase__							,	offload=lowerCamelCase__							,	offload_buffers=lowerCamelCase__			)

      # Parameters have been offloaded, so on the meta device, buffers included
      self.assertEqual(model.lineara.weight.device							,	torch.device('''meta'''			)			)
      self.assertEqual(model.batchnorm.weight.device							,	torch.device('''meta'''			)			)
      self.assertEqual(model.lineara.weight.device							,	torch.device('''meta'''			)			)
      self.assertEqual(model.batchnorm.running_mean.device							,	torch.device('''meta'''			)			)

      UpperCamelCase__					:    int          =   torch.randn(2							,	3			)
      UpperCamelCase__					:    str          =   model(lowerCamelCase__			)
      self.assertEqual(output.device							,	lowerCamelCase__			)

      # Removing hooks loads back the weights in the model.
      remove_hook_from_submodules(lowerCamelCase__			)
      self.assertEqual(model.lineara.weight.device							,	torch.device('''cpu'''			)			)
      self.assertEqual(model.batchnorm.weight.device							,	torch.device('''cpu'''			)			)
      self.assertEqual(model.lineara.weight.device							,	torch.device('''cpu'''			)			)


     def      UpperCAmelCase__	(      self							:     int			)  ->      List[Any]:





      '''simple docstring'''






      UpperCamelCase__					:    List[str]          =   ModelForTest()

      # Everything is on CPU
      self.assertEqual(model.lineara.weight.device							,	torch.device('''cpu'''			)			)
      self.assertEqual(model.batchnorm.weight.device							,	torch.device('''cpu'''			)			)
      self.assertEqual(model.lineara.weight.device							,	torch.device('''cpu'''			)			)

      # This will move each submodule on different devices
      UpperCamelCase__					:    List[str]          =   0 if torch.cuda.is_available() else '''cpu'''
      attach_align_device_hook(
          lowerCamelCase__							,	execution_device=lowerCamelCase__							,	offload=lowerCamelCase__							,	weights_map=model.state_dict()			)

      # Parameters have been offloaded, so on the meta device
      self.assertEqual(model.lineara.weight.device							,	torch.device('''meta'''			)			)
      self.assertEqual(model.batchnorm.weight.device							,	torch.device('''meta'''			)			)
      self.assertEqual(model.lineara.weight.device							,	torch.device('''meta'''			)			)
      # Buffers are not included in the offload by default, so are on the execution device
      UpperCamelCase__					:    List[str]          =   torch.device(lowerCamelCase__			)
      self.assertEqual(model.batchnorm.running_mean.device							,	lowerCamelCase__			)

      UpperCamelCase__					:    List[str]          =   torch.randn(2							,	3			)
      UpperCamelCase__					:    Any          =   model(lowerCamelCase__			)
      self.assertEqual(output.device							,	lowerCamelCase__			)

      # Removing hooks loads back the weights in the model.
      remove_hook_from_submodules(lowerCamelCase__			)
      self.assertEqual(model.lineara.weight.device							,	torch.device('''cpu'''			)			)
      self.assertEqual(model.batchnorm.weight.device							,	torch.device('''cpu'''			)			)
      self.assertEqual(model.lineara.weight.device							,	torch.device('''cpu'''			)			)

      # Now test with buffers included in the offload
      attach_align_device_hook(
          lowerCamelCase__							,	execution_device=lowerCamelCase__							,	offload=lowerCamelCase__							,	weights_map=model.state_dict()							,	offload_buffers=lowerCamelCase__							,	)

      # Parameters have been offloaded, so on the meta device, buffers included
      self.assertEqual(model.lineara.weight.device							,	torch.device('''meta'''			)			)
      self.assertEqual(model.batchnorm.weight.device							,	torch.device('''meta'''			)			)
      self.assertEqual(model.lineara.weight.device							,	torch.device('''meta'''			)			)
      self.assertEqual(model.batchnorm.running_mean.device							,	torch.device('''meta'''			)			)

      UpperCamelCase__					:    Union[str, Any]          =   torch.randn(2							,	3			)
      UpperCamelCase__					:    Optional[int]          =   model(lowerCamelCase__			)
      self.assertEqual(output.device							,	lowerCamelCase__			)

      # Removing hooks loads back the weights in the model.
      remove_hook_from_submodules(lowerCamelCase__			)
      self.assertEqual(model.lineara.weight.device							,	torch.device('''cpu'''			)			)
      self.assertEqual(model.batchnorm.weight.device							,	torch.device('''cpu'''			)			)
      self.assertEqual(model.lineara.weight.device							,	torch.device('''cpu'''			)			)


