


"""simple docstring"""







import unittest

from parameterized import parameterized

from transformers import LlamaConfig, is_torch_available, set_seed
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 LlamaForCausalLM, LlamaForSequenceClassification, LlamaModel, LlamaTokenizer


class 				SCREAMING_SNAKE_CASE__    :
       """simple docstring"""


       def __init__(						self		,					snake_case__		,					snake_case__=13		,					snake_case__=7		,					snake_case__=True		,					snake_case__=True		,					snake_case__=False		,					snake_case__=True		,					snake_case__=99		,					snake_case__=32		,					snake_case__=5		,					snake_case__=4		,					snake_case__=37		,					snake_case__="gelu"		,					snake_case__=0.1		,					snake_case__=0.1		,					snake_case__=512		,					snake_case__=16		,					snake_case__=2		,					snake_case__=0.02		,					snake_case__=3		,					snake_case__=4		,					snake_case__=None		,					):






         """simple docstring"""





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

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					Tuple								=							ids_tensor([self.batch_size, self.seq_length]		,					self.vocab_size						)

         lowerCAmelCase					:					Optional[int]								=							None
         if self.use_input_mask:
           lowerCAmelCase					:					Any								=							random_attention_mask([self.batch_size, self.seq_length]						)

         lowerCAmelCase					:					Any								=							None
         if self.use_token_type_ids:
           lowerCAmelCase					:					str								=							ids_tensor([self.batch_size, self.seq_length]		,					self.type_vocab_size						)

         lowerCAmelCase					:					Dict								=							None
         lowerCAmelCase					:					Tuple								=							None
         lowerCAmelCase					:					int								=							None
         if self.use_labels:
           lowerCAmelCase					:					List[str]								=							ids_tensor([self.batch_size]		,					self.type_sequence_label_size						)
           lowerCAmelCase					:					Any								=							ids_tensor([self.batch_size, self.seq_length]		,					self.num_labels						)
           lowerCAmelCase					:					str								=							ids_tensor([self.batch_size]		,					self.num_choices						)

         lowerCAmelCase					:					str								=							self.get_config()

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

       def 	lowercase__							(						self						):






         """simple docstring"""





         return LlamaConfig(
             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=snake_case__		,					initializer_range=self.initializer_range		,					)

       def 	lowercase__							(						self		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__						):






         """simple docstring"""





         lowerCAmelCase					:					str								=							LlamaModel(config=snake_case__						)
         model.to(snake_case__						)
         model.eval()
         lowerCAmelCase					:					Optional[Any]								=							model(snake_case__		,					attention_mask=snake_case__						)
         lowerCAmelCase					:					Any								=							model(snake_case__						)
         self.parent.assertEqual(result.last_hidden_state.shape		,					(self.batch_size, self.seq_length, self.hidden_size)						)

       def 	lowercase__							(						self		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__		,					):






         """simple docstring"""





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

       def 	lowercase__							(						self		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__		,					):






         """simple docstring"""





         lowerCAmelCase					:					str								=							LlamaForCausalLM(config=snake_case__						)
         model.to(snake_case__						)
         model.eval()
         lowerCAmelCase					:					Tuple								=							model(snake_case__		,					attention_mask=snake_case__		,					labels=snake_case__						)
         self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.seq_length, self.vocab_size)						)

       def 	lowercase__							(						self		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__		,					):






         """simple docstring"""





         lowerCAmelCase					:					Any								=							True
         lowerCAmelCase					:					Optional[int]								=							True
         lowerCAmelCase					:					Optional[Any]								=							LlamaForCausalLM(config=snake_case__						)
         model.to(snake_case__						)
         model.eval()

         # first forward pass
         lowerCAmelCase					:					Any								=							model(
             snake_case__		,					attention_mask=snake_case__		,					encoder_hidden_states=snake_case__		,					encoder_attention_mask=snake_case__		,					use_cache=snake_case__		,					)
         lowerCAmelCase					:					int								=							outputs.past_key_values

         # create hypothetical multiple next token and extent to next_input_ids
         lowerCAmelCase					:					List[Any]								=							ids_tensor((self.batch_size, 3)		,					config.vocab_size						)
         lowerCAmelCase					:					Union[str, Any]								=							ids_tensor((self.batch_size, 3)		,					vocab_size=2						)

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

         lowerCAmelCase					:					str								=							model(
             snake_case__		,					attention_mask=snake_case__		,					encoder_hidden_states=snake_case__		,					encoder_attention_mask=snake_case__		,					output_hidden_states=snake_case__		,					)["hidden_states"][0]
         lowerCAmelCase					:					List[str]								=							model(
             snake_case__		,					attention_mask=snake_case__		,					encoder_hidden_states=snake_case__		,					encoder_attention_mask=snake_case__		,					past_key_values=snake_case__		,					output_hidden_states=snake_case__		,					)["hidden_states"][0]

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







       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					Union[str, Any]								=							self.prepare_config_and_inputs()
         (
             (
             lowerCAmelCase
         )						,  (
             lowerCAmelCase
         )						,  (
             lowerCAmelCase
         )						,  (
             lowerCAmelCase
         )						,  (
             lowerCAmelCase
         )						,  (
             lowerCAmelCase
         )						,  (
             lowerCAmelCase
         )						,  
         )					:					Optional[int]								=							config_and_inputs
         lowerCAmelCase					:					Union[str, Any]								=							{"input_ids": input_ids, "attention_mask": input_mask}
         return config, inputs_dict


@require_torch
class 				SCREAMING_SNAKE_CASE__    (		lowercase       ,   lowercase       ,   lowercase       ,   unittest.TestCase    ):
       """simple docstring"""


       a  :    Dict       =(LlamaModel, LlamaForCausalLM, LlamaForSequenceClassification) if is_torch_available() else ()
       a  :    Union[str, Any]       =(LlamaForCausalLM,) if is_torch_available() else ()
       a  :    str       =(
           {
               "feature-extraction": LlamaModel,
               "text-classification": LlamaForSequenceClassification,
               "text-generation": LlamaForCausalLM,
               "zero-shot": LlamaForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )
       a  :    Dict       =False
       a  :    Tuple       =False

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					List[Any]								=							LlamaModelTester(self						)
         lowerCAmelCase					:					str								=							ConfigTester(self		,					config_class=snake_case__		,					hidden_size=37						)

       def 	lowercase__							(						self						):






         """simple docstring"""





         self.config_tester.run_common_tests()

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					Optional[Any]								=							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*snake_case__						)

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					Any								=							self.model_tester.prepare_config_and_inputs()
         for type in ["absolute", "relative_key", "relative_key_query"]:
           lowerCAmelCase					:					Dict								=							type
           self.model_tester.create_and_check_model(*snake_case__						)

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase						,  lowerCAmelCase					:					Any								=							self.model_tester.prepare_config_and_inputs_for_common()
         lowerCAmelCase					:					List[Any]								=							3
         lowerCAmelCase					:					str								=							input_dict["input_ids"]
         lowerCAmelCase					:					Optional[int]								=							input_ids.ne(1						).to(snake_case__						)
         lowerCAmelCase					:					Tuple								=							ids_tensor([self.model_tester.batch_size]		,					self.model_tester.type_sequence_label_size						)
         lowerCAmelCase					:					int								=							LlamaForSequenceClassification(snake_case__						)
         model.to(snake_case__						)
         model.eval()
         lowerCAmelCase					:					List[Any]								=							model(snake_case__		,					attention_mask=snake_case__		,					labels=snake_case__						)
         self.assertEqual(result.logits.shape		,					(self.model_tester.batch_size, self.model_tester.num_labels)						)

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase						,  lowerCAmelCase					:					Optional[Any]								=							self.model_tester.prepare_config_and_inputs_for_common()
         lowerCAmelCase					:					Any								=							3
         lowerCAmelCase					:					Dict								=							"single_label_classification"
         lowerCAmelCase					:					Any								=							input_dict["input_ids"]
         lowerCAmelCase					:					Dict								=							input_ids.ne(1						).to(snake_case__						)
         lowerCAmelCase					:					Union[str, Any]								=							ids_tensor([self.model_tester.batch_size]		,					self.model_tester.type_sequence_label_size						)
         lowerCAmelCase					:					int								=							LlamaForSequenceClassification(snake_case__						)
         model.to(snake_case__						)
         model.eval()
         lowerCAmelCase					:					Any								=							model(snake_case__		,					attention_mask=snake_case__		,					labels=snake_case__						)
         self.assertEqual(result.logits.shape		,					(self.model_tester.batch_size, self.model_tester.num_labels)						)

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase						,  lowerCAmelCase					:					List[str]								=							self.model_tester.prepare_config_and_inputs_for_common()
         lowerCAmelCase					:					List[str]								=							3
         lowerCAmelCase					:					Union[str, Any]								=							"multi_label_classification"
         lowerCAmelCase					:					List[Any]								=							input_dict["input_ids"]
         lowerCAmelCase					:					Optional[int]								=							input_ids.ne(1						).to(snake_case__						)
         lowerCAmelCase					:					List[Any]								=							ids_tensor(
             [self.model_tester.batch_size, config.num_labels]		,					self.model_tester.type_sequence_label_size						).to(torch.float						)
         lowerCAmelCase					:					List[str]								=							LlamaForSequenceClassification(snake_case__						)
         model.to(snake_case__						)
         model.eval()
         lowerCAmelCase					:					Dict								=							model(snake_case__		,					attention_mask=snake_case__		,					labels=snake_case__						)
         self.assertEqual(result.logits.shape		,					(self.model_tester.batch_size, self.model_tester.num_labels)						)

       @unittest.skip("LLaMA buffers include complex numbers, which breaks this test"						)
       def 	lowercase__							(						self						):






         """simple docstring"""





         pass







       @parameterized.expand([("linear",), ("dynamic",)]						)
       def 	lowercase__							(						self		,					snake_case__						):






         """simple docstring"""





         lowerCAmelCase						,  lowerCAmelCase					:					Tuple								=							self.model_tester.prepare_config_and_inputs_for_common()
         lowerCAmelCase					:					Optional[int]								=							ids_tensor([1, 10]		,					config.vocab_size						)
         lowerCAmelCase					:					Tuple								=							ids_tensor([1, int(config.max_position_embeddings * 1.5						)]		,					config.vocab_size						)

         set_seed(42						)  # Fixed seed at init time so the two models get the same random weights
         lowerCAmelCase					:					Optional[int]								=							LlamaModel(snake_case__						)
         original_model.to(snake_case__						)
         original_model.eval()
         lowerCAmelCase					:					Dict								=							original_model(snake_case__						).last_hidden_state
         lowerCAmelCase					:					Optional[int]								=							original_model(snake_case__						).last_hidden_state

         set_seed(42						)  # Fixed seed at init time so the two models get the same random weights
         lowerCAmelCase					:					Tuple								=							{"type": scaling_type, "factor": 10.0}
         lowerCAmelCase					:					int								=							LlamaModel(snake_case__						)
         scaled_model.to(snake_case__						)
         scaled_model.eval()
         lowerCAmelCase					:					Union[str, Any]								=							scaled_model(snake_case__						).last_hidden_state
         lowerCAmelCase					:					Any								=							scaled_model(snake_case__						).last_hidden_state

         # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
         # maximum sequence length, so the outputs for the short input should match.
         if scaling_type == "dynamic":
           self.assertTrue(torch.allclose(snake_case__		,					snake_case__		,					atol=1e-5						)						)
         else:
           self.assertFalse(torch.allclose(snake_case__		,					snake_case__		,					atol=1e-5						)						)

         # The output should be different for long inputs
         self.assertFalse(torch.allclose(snake_case__		,					snake_case__		,					atol=1e-5						)						)


@require_torch
class 				SCREAMING_SNAKE_CASE__    (		unittest.TestCase    ):
       """simple docstring"""


       @unittest.skip("Logits are not exactly the same, once we fix the instabalities somehow, will update!"						)
       @slow
       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					List[str]								=							[1, 306, 4_658, 278, 6_593, 310, 2_834, 338]
         lowerCAmelCase					:					Tuple								=							LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf"		,					device_map="auto"						)
         lowerCAmelCase					:					Any								=							model(torch.tensor([input_ids]						)						)
         # Expected mean on dim = -1
         lowerCAmelCase					:					Any								=							torch.tensor([[-6.6550, -4.1227, -4.9859, -3.2406, 0.8262, -3.0033, 1.2964, -3.3699]]						)
         torch.testing.assert_close(out.mean(-1						)		,					snake_case__		,					atol=1e-2		,					rtol=1e-2						)
         # slicing logits[0, 0, 0:30]
         # fmt: off
         lowerCAmelCase					:					int								=							torch.tensor([-12.8281, -7.4453, -0.4639, -8.0625, -7.2500, -8.0000, -6.4883, -7.7695, -7.8438, -7.0312, -6.2188, -7.1328, -1.8496, 1.9961, -8.6250, -6.7227, -12.8281, -6.9492, -7.0742, -7.7852, -7.5820, -7.9062, -6.9375, -7.9805, -8.3438, -8.1562, -8.0469, -7.6250, -7.7422, -7.3398,]						)
         # fmt: on
         torch.testing.assert_close(out[0, 0, :30]		,					snake_case__		,					atol=1e-5		,					rtol=1e-5						)

       @unittest.skip("Logits are not exactly the same, once we fix the instabalities somehow, will update!"						)
       @slow
       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					Any								=							[1, 306, 4_658, 278, 6_593, 310, 2_834, 338]
         lowerCAmelCase					:					Optional[int]								=							LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-13b-hf"		,					device_map="auto"						)
         lowerCAmelCase					:					int								=							model(torch.tensor(snake_case__						)						)
         # Expected mean on dim = -1
         lowerCAmelCase					:					str								=							torch.tensor([[-2.0622, -1.2794, -1.1638, -0.9788, -1.4603, -1.0238, -1.7893, -1.4411]]						)
         torch.testing.assert_close(out.mean(-1						)		,					snake_case__		,					atol=1e-2		,					rtol=1e-2						)
         # slicing logits[0, 0, 0:30]
         # fmt: off
         lowerCAmelCase					:					List[str]								=							torch.tensor([-8.1406, -8.0547, 2.7461, -1.2344, -0.1448, -1.8262, -1.0020, -1.8154, -1.6895, -1.8516, -2.3574, -0.9277, 3.7598, 6.5742, -1.2998, -0.1177, -8.1406, -2.9688, -2.9199, -3.1699, -3.5254, -2.3555, -2.7988, -3.4141, -2.8262, -4.5195, -3.3379, -3.3164, -2.7832, -3.0273]						)
         # fmt: on
         torch.testing.assert_close(out[0, 0, :30]		,					snake_case__		,					atol=1e-5		,					rtol=1e-5						)

       @unittest.skip("Logits are not exactly the same, once we fix the instabalities somehow, will update!"						)
       @slow
       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					str								=							[1, 306, 4_658, 278, 6_593, 310, 2_834, 338]
         lowerCAmelCase					:					Any								=							LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-13b-chat-hf"		,					device_map="auto"						)
         lowerCAmelCase					:					Optional[Any]								=							model(torch.tensor(snake_case__						)						)
         # Expected mean on dim = -1
         lowerCAmelCase					:					List[str]								=							torch.tensor([[-0.8562, -1.8520, -0.7551, -0.4162, -1.5161, -1.2038, -2.4823, -2.3254]]						)
         torch.testing.assert_close(out.mean(-1						)		,					snake_case__		,					atol=1e-2		,					rtol=1e-2						)
         # slicing logits[0, 0, 0:30]
         # fmt: off
         lowerCAmelCase					:					int								=							torch.tensor([-2.2227, 4.8828, 0.9023, -0.4578, -0.7871, -0.1033, -0.6221, -0.5786, -0.7803, -1.0674, -1.2920, -0.1570, 0.8008, 2.0723, -0.9497, 0.2771, -2.2227, -0.7612, -1.4346, -1.2061, -1.6426, -0.3000, -0.7139, -1.1934, -1.8691, -1.6973, -1.5947, -1.2705, -0.3523, -0.5513]						)
         # fmt: on
         torch.testing.assert_close(out.mean(-1						)		,					snake_case__		,					atol=1e-2		,					rtol=1e-2						)

       @unittest.skip(
           "Logits are not exactly the same, once we fix the instabalities somehow, will update! Also it is gonna be a `too_slow` test"						)
       @slow
       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					Dict								=							[1, 306, 4_658, 278, 6_593, 310, 2_834, 338]
         lowerCAmelCase					:					str								=							LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-70b-hf"		,					device_map="auto"						)
         lowerCAmelCase					:					str								=							model(torch.tensor(snake_case__						)						)

         lowerCAmelCase					:					Dict								=							torch.tensor(
             [[-4.2327, -3.3360, -4.6665, -4.7631, -1.8180, -3.4170, -1.4211, -3.1810]]		,					dtype=torch.floataa						)
         torch.testing.assert_close(out.mean(-1						)		,					snake_case__		,					atol=1e-2		,					rtol=1e-2						)
         # fmt: off
         lowerCAmelCase					:					Any								=							torch.tensor([-9.4922, -3.9551, 1.7998, -5.6758, -5.1055, -5.8984, -4.8320, -6.8086, -6.5391, -5.6172, -5.5820, -5.5352, 1.7881, 3.6289, -6.5117, -3.4785, -9.5000, -6.0352, -6.8125, -6.0195, -6.6836, -5.4727, -6.2812, -6.0391, -7.3398, -7.4297, -7.4844, -6.5820, -5.8789, -5.5312]						)
         # fmt: on
         torch.testing.assert_close(out[0, 0, :30]		,					snake_case__		,					atol=1e-5		,					rtol=1e-5						)







       @unittest.skip("Model is curently gated"						)
       @slow
       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					Optional[int]								=							"Simply put, the theory of relativity states that 1) the laws of physics are the same everywhere in the universe and 2) the passage of time and the length of objects can vary depending on the observer\'s frame of reference.\n\nThe first part of the theory, that the laws of physics are the same everywhere, is known as the \"princi"
         lowerCAmelCase					:					Union[str, Any]								=							"Simply put, the theory of relativity states that "
         lowerCAmelCase					:					Any								=							LlamaTokenizer.from_pretrained("meta-llama/Llama-2-13b-chat-hf"						)
         lowerCAmelCase					:					List[str]								=							tokenizer.encode(snake_case__		,					return_tensors="pt"						)
         lowerCAmelCase					:					int								=							LlamaForCausalLM.from_pretrained(
             "meta-llama/Llama-2-13b-chat-hf"		,					device_map="sequential"		,					use_safetensors=snake_case__						)

         # greedy generation outputs
         lowerCAmelCase					:					Optional[Any]								=							model.generate(snake_case__		,					max_new_tokens=64		,					top_p=snake_case__		,					temperature=1		,					do_sample=snake_case__						)
         lowerCAmelCase					:					Union[str, Any]								=							tokenizer.decode(generated_ids[0]		,					skip_special_tokens=snake_case__						)
         self.assertEqual(snake_case__		,					snake_case__						)









"""simple docstring"""







import argparse

from diffusers.pipelines.stable_diffusion.convert_from_ckpt import download_controlnet_from_original_ckpt


if __name__ == "__main__":
     lowerCAmelCase__				    =       argparse.ArgumentParser()

     parser.add_argument(
         '''--checkpoint_path''', default=None, type=str, required=True, help='''Path to the checkpoint to convert.'''
     )
     parser.add_argument(
         '''--original_config_file''',
         type=str,
         required=True,
         help='''The YAML config file corresponding to the original architecture.''',
     )
     parser.add_argument(
         '''--num_in_channels''',
         default=None,
         type=int,
         help='''The number of input channels. If `None` number of input channels will be automatically inferred.''',
     )
     parser.add_argument(
         '''--image_size''',
         default=512,
         type=int,
         help=(
             '''The image size that the model was trained on. Use 512 for Stable Diffusion v1.X and Stable Siffusion v2'''
             ''' Base. Use 768 for Stable Diffusion v2.'''
         ),
     )
     parser.add_argument(
         '''--extract_ema''',
         action='''store_true''',
         help=(
             '''Only relevant for checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights'''
             ''' or not. Defaults to `False`. Add `--extract_ema` to extract the EMA weights. EMA weights usually yield'''
             ''' higher quality images for inference. Non-EMA weights are usually better to continue fine-tuning.'''
         ),
     )
     parser.add_argument(
         '''--upcast_attention''',
         action='''store_true''',
         help=(
             '''Whether the attention computation should always be upcasted. This is necessary when running stable'''
             ''' diffusion 2.1.'''
         ),
     )
     parser.add_argument(
         '''--from_safetensors''',
         action='''store_true''',
         help='''If `--checkpoint_path` is in `safetensors` format, load checkpoint with safetensors instead of PyTorch.''',
     )
     parser.add_argument(
         '''--to_safetensors''',
         action='''store_true''',
         help='''Whether to store pipeline in safetensors format or not.''',
     )
     parser.add_argument('''--dump_path''', default=None, type=str, required=True, help='''Path to the output model.''')
     parser.add_argument('''--device''', type=str, help='''Device to use (e.g. cpu, cuda:0, cuda:1, etc.)''')
     def  a__      (    SCREAMING_SNAKE_CASE   :   Tuple    ):


       '''simple docstring'''







       if string == "True":
         return True
       elif string == "False":
         return False
       else:
         raise ValueError(f"""could not parse string as bool {string}"""    )

     parser.add_argument(
         '''--use_linear_projection''', help='''Override for use linear projection''', required=False, type=parse_bool
     )

     parser.add_argument('''--cross_attention_dim''', help='''Override for cross attention_dim''', required=False, type=int)

     lowerCAmelCase__				    =       parser.parse_args()

     lowerCAmelCase__				    =       download_controlnet_from_original_ckpt(
         checkpoint_path=args.checkpoint_path,
         original_config_file=args.original_config_file,
         image_size=args.image_size,
         extract_ema=args.extract_ema,
         num_in_channels=args.num_in_channels,
         upcast_attention=args.upcast_attention,
         from_safetensors=args.from_safetensors,
         device=args.device,
         use_linear_projection=args.use_linear_projection,
         cross_attention_dim=args.cross_attention_dim,
     )





     controlnet.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors)





