




"""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 				snake_case_			:







   """simple docstring"""




   def __init__(		self       ,					lowerCamelCase_       ,					lowerCamelCase_=1_3       ,					lowerCamelCase_=7       ,					lowerCamelCase_=True       ,					lowerCamelCase_=True       ,					lowerCamelCase_=False       ,					lowerCamelCase_=True       ,					lowerCamelCase_=9_9       ,					lowerCamelCase_=3_2       ,					lowerCamelCase_=5       ,					lowerCamelCase_=4       ,					lowerCamelCase_=3_7       ,					lowerCamelCase_="gelu"       ,					lowerCamelCase_=0.1       ,					lowerCamelCase_=0.1       ,					lowerCamelCase_=5_1_2       ,					lowerCamelCase_=1_6       ,					lowerCamelCase_=2       ,					lowerCamelCase_=0.02       ,					lowerCamelCase_=3       ,					lowerCamelCase_=4       ,					lowerCamelCase_=None       ,					)				->   List[str]:
       UpperCamelCase              =	parent
       UpperCamelCase              =	batch_size
       UpperCamelCase              =	seq_length
       UpperCamelCase              =	is_training
       UpperCamelCase              =	use_input_mask
       UpperCamelCase              =	use_token_type_ids
       UpperCamelCase              =	use_labels
       UpperCamelCase              =	vocab_size
       UpperCamelCase              =	hidden_size
       UpperCamelCase              =	num_hidden_layers
       UpperCamelCase              =	num_attention_heads
       UpperCamelCase              =	intermediate_size
       UpperCamelCase              =	hidden_act
       UpperCamelCase              =	hidden_dropout_prob
       UpperCamelCase              =	attention_probs_dropout_prob
       UpperCamelCase              =	max_position_embeddings
       UpperCamelCase              =	type_vocab_size
       UpperCamelCase              =	type_sequence_label_size
       UpperCamelCase              =	initializer_range
       UpperCamelCase              =	num_labels
       UpperCamelCase              =	num_choices
       UpperCamelCase              =	scope







   def   UpperCAmelCase__			(		self)				->   int:
       UpperCamelCase              =	ids_tensor([self.batch_size, self.seq_length]       ,					self.vocab_size)

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

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

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

       UpperCamelCase              =	self.get_config()

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







   def   UpperCAmelCase__			(		self)				->   Optional[int]:
       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=lowerCamelCase_       ,					initializer_range=self.initializer_range       ,					)







   def   UpperCAmelCase__			(		self       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_)				->   str:
       UpperCamelCase              =	LlamaModel(config=lowerCamelCase_)
       model.to(lowerCamelCase_)
       model.eval()
       UpperCamelCase              =	model(lowerCamelCase_       ,					attention_mask=lowerCamelCase_)
       UpperCamelCase              =	model(lowerCamelCase_)
       self.parent.assertEqual(result.last_hidden_state.shape       ,					(self.batch_size, self.seq_length, self.hidden_size))







   def   UpperCAmelCase__			(		self       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					)				->   int:
       UpperCamelCase              =	True
       UpperCamelCase              =	LlamaModel(lowerCamelCase_)
       model.to(lowerCamelCase_)
       model.eval()
       UpperCamelCase              =	model(
           lowerCamelCase_       ,					attention_mask=lowerCamelCase_       ,					encoder_hidden_states=lowerCamelCase_       ,					encoder_attention_mask=lowerCamelCase_       ,					)
       UpperCamelCase              =	model(
           lowerCamelCase_       ,					attention_mask=lowerCamelCase_       ,					encoder_hidden_states=lowerCamelCase_       ,					)
       UpperCamelCase              =	model(lowerCamelCase_       ,					attention_mask=lowerCamelCase_)
       self.parent.assertEqual(result.last_hidden_state.shape       ,					(self.batch_size, self.seq_length, self.hidden_size))







   def   UpperCAmelCase__			(		self       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					)				->   List[Any]:
       UpperCamelCase              =	LlamaForCausalLM(config=lowerCamelCase_)
       model.to(lowerCamelCase_)
       model.eval()
       UpperCamelCase              =	model(lowerCamelCase_       ,					attention_mask=lowerCamelCase_       ,					labels=lowerCamelCase_)
       self.parent.assertEqual(result.logits.shape       ,					(self.batch_size, self.seq_length, self.vocab_size))







   def   UpperCAmelCase__			(		self       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					)				->   Optional[int]:
       UpperCamelCase              =	True
       UpperCamelCase              =	True
       UpperCamelCase              =	LlamaForCausalLM(config=lowerCamelCase_)
       model.to(lowerCamelCase_)
       model.eval()

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

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

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

       UpperCamelCase              =	model(
           lowerCamelCase_       ,					attention_mask=lowerCamelCase_       ,					encoder_hidden_states=lowerCamelCase_       ,					encoder_attention_mask=lowerCamelCase_       ,					output_hidden_states=lowerCamelCase_       ,					)['''hidden_states'''][0]
       UpperCamelCase              =	model(
           lowerCamelCase_       ,					attention_mask=lowerCamelCase_       ,					encoder_hidden_states=lowerCamelCase_       ,					encoder_attention_mask=lowerCamelCase_       ,					past_key_values=lowerCamelCase_       ,					output_hidden_states=lowerCamelCase_       ,					)['''hidden_states'''][0]

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

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

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







   def   UpperCAmelCase__			(		self)				->   Optional[int]:
       UpperCamelCase              =	self.prepare_config_and_inputs()
       (
           (
           UpperCamelCase			
       )       ,	(
           UpperCamelCase			
       )       ,	(
           UpperCamelCase			
       )       ,	(
           UpperCamelCase			
       )       ,	(
           UpperCamelCase			
       )       ,	(
           UpperCamelCase			
       )       ,	(
           UpperCamelCase			
       )       ,				
       )              =	config_and_inputs
       UpperCamelCase              =	{'''input_ids''': input_ids, '''attention_mask''': input_mask}
       return config, inputs_dict

@require_torch
class 				snake_case_			(					lowerCamelCase_	, lowerCamelCase_	, lowerCamelCase_	, unittest.TestCase			):







   """simple docstring"""




   A_				      =       (LlamaModel, LlamaForCausalLM, LlamaForSequenceClassification) if is_torch_available() else ()
   A_				      =       (LlamaForCausalLM,) if is_torch_available() else ()
   A_				      =       (
       {
           '''feature-extraction''': LlamaModel,
           '''text-classification''': LlamaForSequenceClassification,
           '''text-generation''': LlamaForCausalLM,
           '''zero-shot''': LlamaForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )
   A_				      =       False
   A_				      =       False







   def   UpperCAmelCase__			(		self)				->   Optional[int]:
       UpperCamelCase              =	LlamaModelTester(self)
       UpperCamelCase              =	ConfigTester(self       ,					config_class=lowerCamelCase_       ,					hidden_size=3_7)







   def   UpperCAmelCase__			(		self)				->   Union[str, Any]:
       self.config_tester.run_common_tests()







   def   UpperCAmelCase__			(		self)				->   List[Any]:
       UpperCamelCase              =	self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*lowerCamelCase_)







   def   UpperCAmelCase__			(		self)				->   int:
       UpperCamelCase              =	self.model_tester.prepare_config_and_inputs()
       for type in ["absolute", "relative_key", "relative_key_query"]:
           UpperCamelCase              =	type
           self.model_tester.create_and_check_model(*lowerCamelCase_)







   def   UpperCAmelCase__			(		self)				->   List[Any]:
       UpperCamelCase       ,	UpperCamelCase              =	self.model_tester.prepare_config_and_inputs_for_common()
       UpperCamelCase              =	3
       UpperCamelCase              =	input_dict['''input_ids''']
       UpperCamelCase              =	input_ids.ne(1).to(lowerCamelCase_)
       UpperCamelCase              =	ids_tensor([self.model_tester.batch_size]       ,					self.model_tester.type_sequence_label_size)
       UpperCamelCase              =	LlamaForSequenceClassification(lowerCamelCase_)
       model.to(lowerCamelCase_)
       model.eval()
       UpperCamelCase              =	model(lowerCamelCase_       ,					attention_mask=lowerCamelCase_       ,					labels=lowerCamelCase_)
       self.assertEqual(result.logits.shape       ,					(self.model_tester.batch_size, self.model_tester.num_labels))







   def   UpperCAmelCase__			(		self)				->   Optional[Any]:
       UpperCamelCase       ,	UpperCamelCase              =	self.model_tester.prepare_config_and_inputs_for_common()
       UpperCamelCase              =	3
       UpperCamelCase              =	'''single_label_classification'''
       UpperCamelCase              =	input_dict['''input_ids''']
       UpperCamelCase              =	input_ids.ne(1).to(lowerCamelCase_)
       UpperCamelCase              =	ids_tensor([self.model_tester.batch_size]       ,					self.model_tester.type_sequence_label_size)
       UpperCamelCase              =	LlamaForSequenceClassification(lowerCamelCase_)
       model.to(lowerCamelCase_)
       model.eval()
       UpperCamelCase              =	model(lowerCamelCase_       ,					attention_mask=lowerCamelCase_       ,					labels=lowerCamelCase_)
       self.assertEqual(result.logits.shape       ,					(self.model_tester.batch_size, self.model_tester.num_labels))







   def   UpperCAmelCase__			(		self)				->   Optional[int]:
       UpperCamelCase       ,	UpperCamelCase              =	self.model_tester.prepare_config_and_inputs_for_common()
       UpperCamelCase              =	3
       UpperCamelCase              =	'''multi_label_classification'''
       UpperCamelCase              =	input_dict['''input_ids''']
       UpperCamelCase              =	input_ids.ne(1).to(lowerCamelCase_)
       UpperCamelCase              =	ids_tensor(
           [self.model_tester.batch_size, config.num_labels]       ,					self.model_tester.type_sequence_label_size).to(torch.float)
       UpperCamelCase              =	LlamaForSequenceClassification(lowerCamelCase_)
       model.to(lowerCamelCase_)
       model.eval()
       UpperCamelCase              =	model(lowerCamelCase_       ,					attention_mask=lowerCamelCase_       ,					labels=lowerCamelCase_)
       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   UpperCAmelCase__			(		self)				->   Optional[int]:
       pass







   @parameterized.expand([('''linear''',), ('''dynamic''',)])
   def   UpperCAmelCase__			(		self       ,					lowerCamelCase_)				->   Optional[Any]:
       UpperCamelCase       ,	UpperCamelCase              =	self.model_tester.prepare_config_and_inputs_for_common()
       UpperCamelCase              =	ids_tensor([1, 1_0]       ,					config.vocab_size)
       UpperCamelCase              =	ids_tensor([1, int(config.max_position_embeddings * 1.5)]       ,					config.vocab_size)

       set_seed(4_2)  # Fixed seed at init time so the two models get the same random weights
       UpperCamelCase              =	LlamaModel(lowerCamelCase_)
       original_model.to(lowerCamelCase_)
       original_model.eval()
       UpperCamelCase              =	original_model(lowerCamelCase_).last_hidden_state
       UpperCamelCase              =	original_model(lowerCamelCase_).last_hidden_state

       set_seed(4_2)  # Fixed seed at init time so the two models get the same random weights
       UpperCamelCase              =	{'''type''': scaling_type, '''factor''': 10.0}
       UpperCamelCase              =	LlamaModel(lowerCamelCase_)
       scaled_model.to(lowerCamelCase_)
       scaled_model.eval()
       UpperCamelCase              =	scaled_model(lowerCamelCase_).last_hidden_state
       UpperCamelCase              =	scaled_model(lowerCamelCase_).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(lowerCamelCase_       ,					lowerCamelCase_       ,					atol=1e-5))
       else:
           self.assertFalse(torch.allclose(lowerCamelCase_       ,					lowerCamelCase_       ,					atol=1e-5))

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

@require_torch
class 				snake_case_			(					unittest.TestCase			):







   """simple docstring"""




   @unittest.skip('''Logits are not exactly the same, once we fix the instabalities somehow, will update!''')
   @slow
   def   UpperCAmelCase__			(		self)				->   List[Any]:
       UpperCamelCase              =	[1, 3_0_6, 4_6_5_8, 2_7_8, 6_5_9_3, 3_1_0, 2_8_3_4, 3_3_8]
       UpperCamelCase              =	LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-7b-hf'''       ,					device_map='''auto''')
       UpperCamelCase              =	model(torch.tensor([input_ids]))
       # Expected mean on dim = -1
       UpperCamelCase              =	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)       ,					lowerCamelCase_       ,					atol=1e-2       ,					rtol=1e-2)
       # slicing logits[0, 0, 0:30]
       # fmt: off
       UpperCamelCase              =	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, :3_0]       ,					lowerCamelCase_       ,					atol=1e-5       ,					rtol=1e-5)







   @unittest.skip('''Logits are not exactly the same, once we fix the instabalities somehow, will update!''')
   @slow
   def   UpperCAmelCase__			(		self)				->   int:
       UpperCamelCase              =	[1, 3_0_6, 4_6_5_8, 2_7_8, 6_5_9_3, 3_1_0, 2_8_3_4, 3_3_8]
       UpperCamelCase              =	LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-13b-hf'''       ,					device_map='''auto''')
       UpperCamelCase              =	model(torch.tensor(lowerCamelCase_))
       # Expected mean on dim = -1
       UpperCamelCase              =	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)       ,					lowerCamelCase_       ,					atol=1e-2       ,					rtol=1e-2)
       # slicing logits[0, 0, 0:30]
       # fmt: off
       UpperCamelCase              =	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, :3_0]       ,					lowerCamelCase_       ,					atol=1e-5       ,					rtol=1e-5)







   @unittest.skip('''Logits are not exactly the same, once we fix the instabalities somehow, will update!''')
   @slow
   def   UpperCAmelCase__			(		self)				->   int:
       UpperCamelCase              =	[1, 3_0_6, 4_6_5_8, 2_7_8, 6_5_9_3, 3_1_0, 2_8_3_4, 3_3_8]
       UpperCamelCase              =	LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-13b-chat-hf'''       ,					device_map='''auto''')
       UpperCamelCase              =	model(torch.tensor(lowerCamelCase_))
       # Expected mean on dim = -1
       UpperCamelCase              =	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)       ,					lowerCamelCase_       ,					atol=1e-2       ,					rtol=1e-2)
       # slicing logits[0, 0, 0:30]
       # fmt: off
       UpperCamelCase              =	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)       ,					lowerCamelCase_       ,					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   UpperCAmelCase__			(		self)				->   Union[str, Any]:
       UpperCamelCase              =	[1, 3_0_6, 4_6_5_8, 2_7_8, 6_5_9_3, 3_1_0, 2_8_3_4, 3_3_8]
       UpperCamelCase              =	LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-70b-hf'''       ,					device_map='''auto''')
       UpperCamelCase              =	model(torch.tensor(lowerCamelCase_))

       UpperCamelCase              =	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)       ,					lowerCamelCase_       ,					atol=1e-2       ,					rtol=1e-2)
       # fmt: off
       UpperCamelCase              =	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, :3_0]       ,					lowerCamelCase_       ,					atol=1e-5       ,					rtol=1e-5)







   @unittest.skip('''Model is curently gated''')
   @slow
   def   UpperCAmelCase__			(		self)				->   Optional[Any]:
       UpperCamelCase              =	'''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'''
       UpperCamelCase              =	'''Simply put, the theory of relativity states that '''
       UpperCamelCase              =	LlamaTokenizer.from_pretrained('''meta-llama/Llama-2-13b-chat-hf''')
       UpperCamelCase              =	tokenizer.encode(lowerCamelCase_       ,					return_tensors='''pt''')
       UpperCamelCase              =	LlamaForCausalLM.from_pretrained(
           '''meta-llama/Llama-2-13b-chat-hf'''       ,					device_map='''sequential'''       ,					use_safetensors=lowerCamelCase_)

       # greedy generation outputs
       UpperCamelCase              =	model.generate(lowerCamelCase_       ,					max_new_tokens=6_4       ,					top_p=lowerCamelCase_       ,					temperature=1       ,					do_sample=lowerCamelCase_)
       UpperCamelCase              =	tokenizer.decode(generated_ids[0]       ,					skip_special_tokens=lowerCamelCase_)
       self.assertEqual(lowerCamelCase_       ,					lowerCamelCase_)







import json
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
  import torch

  if is_vision_available():
    from transformers import OneFormerImageProcessor
    from transformers.models.oneformer.image_processing_oneformer import binary_mask_to_rle
    from transformers.models.oneformer.modeling_oneformer import OneFormerForUniversalSegmentationOutput

if is_vision_available():
  from PIL import Image




def       _SCREAMING_SNAKE_CASE (							__snake_case   ,      __snake_case="shi-labs/oneformer_demo" )				->						Any:
       with open(hf_hub_download(__snake_case   ,      __snake_case   ,      repo_type="""dataset""" )   ,      """r""" ) as f:
              _UpperCAmelCase							=    json.load(__snake_case )
       _UpperCAmelCase							=    {}
       _UpperCAmelCase							=    []
       _UpperCAmelCase							=    []
       for key, info in class_info.items():
              _UpperCAmelCase							=    info["""name"""]
              class_names.append(info["""name"""] )
              if info["isthing"]:
                     thing_ids.append(int(__snake_case ) )
       _UpperCAmelCase							=    thing_ids
       _UpperCAmelCase							=    class_names
       return metadata
class 		SCREAMING_SNAKE_CASE__	(  unittest.TestCase     ):

  '''simple docstring'''







  def __init__(      self		:			Optional[Any]       ,				lowerCamelCase		:			Optional[Any]       ,				lowerCamelCase		:			List[Any]=7       ,				lowerCamelCase		:			str=3       ,				lowerCamelCase		:			Union[str, Any]=30       ,				lowerCamelCase		:			Optional[int]=400       ,				lowerCamelCase		:			Union[str, Any]=None       ,				lowerCamelCase		:			int=True       ,				lowerCamelCase		:			List[str]=True       ,				lowerCamelCase		:			Tuple=[0.5, 0.5, 0.5]       ,				lowerCamelCase		:			Tuple=[0.5, 0.5, 0.5]       ,				lowerCamelCase		:			Tuple=10       ,				lowerCamelCase		:			str=False       ,				lowerCamelCase		:			Union[str, Any]=255       ,				lowerCamelCase		:			Tuple="shi-labs/oneformer_demo"       ,				lowerCamelCase		:			Tuple="ade20k_panoptic.json"       ,				lowerCamelCase		:			Optional[Any]=10       ,				)						->	Optional[int]:


         """simple docstring"""
         _UpperCAmelCase							=    parent
         _UpperCAmelCase							=    batch_size
         _UpperCAmelCase							=    num_channels
         _UpperCAmelCase							=    min_resolution
         _UpperCAmelCase							=    max_resolution
         _UpperCAmelCase							=    do_resize
         _UpperCAmelCase							=    {"""shortest_edge""": 32, """longest_edge""": 1333} if size is None else size
         _UpperCAmelCase							=    do_normalize
         _UpperCAmelCase							=    image_mean
         _UpperCAmelCase							=    image_std
         _UpperCAmelCase							=    class_info_file
         _UpperCAmelCase							=    prepare_metadata(lowerCamelCase       ,				lowerCamelCase							)
         _UpperCAmelCase							=    num_text
         _UpperCAmelCase							=    repo_path

         # for the post_process_functions
         _UpperCAmelCase							=    2
         _UpperCAmelCase							=    10
         _UpperCAmelCase							=    10
         _UpperCAmelCase							=    3
         _UpperCAmelCase							=    4
         _UpperCAmelCase							=    num_labels
         _UpperCAmelCase							=    do_reduce_labels
         _UpperCAmelCase							=    ignore_index


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


         """simple docstring"""
         return {
             "do_resize": self.do_resize,
             "size": self.size,
             "do_normalize": self.do_normalize,
             "image_mean": self.image_mean,
             "image_std": self.image_std,
             "num_labels": self.num_labels,
             "do_reduce_labels": self.do_reduce_labels,
             "ignore_index": self.ignore_index,
             "class_info_file": self.class_info_file,
             "metadata": self.metadata,
             "num_text": self.num_text,
         }


  def  lowerCamelCase       (      self		:			Any       ,				lowerCamelCase		:			int       ,				lowerCamelCase		:			Tuple=False							)						->	Any:


         """simple docstring"""
         if not batched:
                _UpperCAmelCase							=    image_inputs[0]
                if isinstance(lowerCamelCase       ,				Image.Image							):
                       _UpperCAmelCase	,			_UpperCAmelCase							=    image.size
                else:
                       _UpperCAmelCase	,			_UpperCAmelCase							=    image.shape[1], image.shape[2]
                if w < h:
                       _UpperCAmelCase							=    int(self.size["""shortest_edge"""] * h / w							)
                       _UpperCAmelCase							=    self.size["""shortest_edge"""]
                elif w > h:
                       _UpperCAmelCase							=    self.size["""shortest_edge"""]
                       _UpperCAmelCase							=    int(self.size["""shortest_edge"""] * w / h							)
                else:
                       _UpperCAmelCase							=    self.size["""shortest_edge"""]
                       _UpperCAmelCase							=    self.size["""shortest_edge"""]

         else:
                _UpperCAmelCase							=    []
                for image in image_inputs:
                       _UpperCAmelCase	,			_UpperCAmelCase							=    self.get_expected_values([image]							)
                       expected_values.append((expected_height, expected_width)							)
                _UpperCAmelCase							=    max(lowerCamelCase       ,				key=lambda lowerCamelCase							: item[0]							)[0]
                _UpperCAmelCase							=    max(lowerCamelCase       ,				key=lambda lowerCamelCase							: item[1]							)[1]

         return expected_height, expected_width





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


         """simple docstring"""
         return OneFormerForUniversalSegmentationOutput(
             # +1 for null class
             class_queries_logits=torch.randn((self.batch_size, self.num_queries, self.num_classes + 1)							)       ,				masks_queries_logits=torch.randn((self.batch_size, self.num_queries, self.height, self.width)							)       ,				)



@require_torch
@require_vision
class 		SCREAMING_SNAKE_CASE__	(  UpperCAmelCase					,      unittest.TestCase     ):

  '''simple docstring'''







  _lowerCamelCase    =  OneFormerImageProcessor if (is_vision_available() and is_torch_available()) else None
  # only for test_image_processing_common.test_image_proc_to_json_string
  _lowerCamelCase    =  image_processing_class


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


         """simple docstring"""
         _UpperCAmelCase							=    OneFormerImageProcessorTester(self							)


  @property
  def  lowerCamelCase       (      self		:			Union[str, Any]							)						->	Dict:


         """simple docstring"""
         return self.image_processing_tester.prepare_image_processor_dict()


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


         """simple docstring"""
         _UpperCAmelCase							=    self.image_processing_class(**self.image_processor_dict							)
         self.assertTrue(hasattr(lowerCamelCase       ,				"""image_mean"""							)							)
         self.assertTrue(hasattr(lowerCamelCase       ,				"""image_std"""							)							)
         self.assertTrue(hasattr(lowerCamelCase       ,				"""do_normalize"""							)							)
         self.assertTrue(hasattr(lowerCamelCase       ,				"""do_resize"""							)							)
         self.assertTrue(hasattr(lowerCamelCase       ,				"""size"""							)							)
         self.assertTrue(hasattr(lowerCamelCase       ,				"""ignore_index"""							)							)
         self.assertTrue(hasattr(lowerCamelCase       ,				"""class_info_file"""							)							)
         self.assertTrue(hasattr(lowerCamelCase       ,				"""num_text"""							)							)
         self.assertTrue(hasattr(lowerCamelCase       ,				"""repo_path"""							)							)
         self.assertTrue(hasattr(lowerCamelCase       ,				"""metadata"""							)							)
         self.assertTrue(hasattr(lowerCamelCase       ,				"""do_reduce_labels"""							)							)


  def  lowerCamelCase       (      self		:			Tuple							)						->	Any:


         """simple docstring"""
         pass


  def  lowerCamelCase       (      self		:			Optional[int]							)						->	Optional[int]:


         """simple docstring"""
         # Initialize image_processor
         _UpperCAmelCase							=    self.image_processing_class(**self.image_processor_dict							)
         # create random PIL images
         _UpperCAmelCase							=    prepare_image_inputs(self.image_processing_tester       ,				equal_resolution=lowerCamelCase							)
         for image in image_inputs:
                self.assertIsInstance(lowerCamelCase       ,				Image.Image							)

         # Test not batched input
         _UpperCAmelCase							=    image_processor(image_inputs[0]       ,				["""semantic"""]       ,				return_tensors="""pt"""							).pixel_values

         _UpperCAmelCase	,			_UpperCAmelCase							=    self.image_processing_tester.get_expected_values(lowerCamelCase							)

         self.assertEqual(
             encoded_images.shape       ,				(1, self.image_processing_tester.num_channels, expected_height, expected_width)       ,				)

         # Test batched
         _UpperCAmelCase	,			_UpperCAmelCase							=    self.image_processing_tester.get_expected_values(lowerCamelCase       ,				batched=lowerCamelCase							)

         _UpperCAmelCase							=    image_processor(
             lowerCamelCase       ,				["""semantic"""] * len(lowerCamelCase							)       ,				return_tensors="""pt"""							).pixel_values
         self.assertEqual(
             encoded_images.shape       ,				(
                 self.image_processing_tester.batch_size,
                 self.image_processing_tester.num_channels,
                 expected_height,
                 expected_width,
             )       ,				)


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


         """simple docstring"""
         # Initialize image_processor
         _UpperCAmelCase							=    self.image_processing_class(**self.image_processor_dict							)
         # create random numpy tensors
         _UpperCAmelCase							=    prepare_image_inputs(self.image_processing_tester       ,				equal_resolution=lowerCamelCase       ,				numpify=lowerCamelCase							)
         for image in image_inputs:
                self.assertIsInstance(lowerCamelCase       ,				np.ndarray							)

         # Test not batched input
         _UpperCAmelCase							=    image_processor(image_inputs[0]       ,				["""semantic"""]       ,				return_tensors="""pt"""							).pixel_values

         _UpperCAmelCase	,			_UpperCAmelCase							=    self.image_processing_tester.get_expected_values(lowerCamelCase							)

         self.assertEqual(
             encoded_images.shape       ,				(1, self.image_processing_tester.num_channels, expected_height, expected_width)       ,				)

         # Test batched
         _UpperCAmelCase	,			_UpperCAmelCase							=    self.image_processing_tester.get_expected_values(lowerCamelCase       ,				batched=lowerCamelCase							)

         _UpperCAmelCase							=    image_processor(
             lowerCamelCase       ,				["""semantic"""] * len(lowerCamelCase							)       ,				return_tensors="""pt"""							).pixel_values
         self.assertEqual(
             encoded_images.shape       ,				(
                 self.image_processing_tester.batch_size,
                 self.image_processing_tester.num_channels,
                 expected_height,
                 expected_width,
             )       ,				)


  def  lowerCamelCase       (      self		:			List[str]							)						->	Tuple:


         """simple docstring"""
         # Initialize image_processor
         _UpperCAmelCase							=    self.image_processing_class(**self.image_processor_dict							)
         # create random PyTorch tensors
         _UpperCAmelCase							=    prepare_image_inputs(self.image_processing_tester       ,				equal_resolution=lowerCamelCase       ,				torchify=lowerCamelCase							)
         for image in image_inputs:
                self.assertIsInstance(lowerCamelCase       ,				torch.Tensor							)

         # Test not batched input
         _UpperCAmelCase							=    image_processor(image_inputs[0]       ,				["""semantic"""]       ,				return_tensors="""pt"""							).pixel_values

         _UpperCAmelCase	,			_UpperCAmelCase							=    self.image_processing_tester.get_expected_values(lowerCamelCase							)

         self.assertEqual(
             encoded_images.shape       ,				(1, self.image_processing_tester.num_channels, expected_height, expected_width)       ,				)

         # Test batched
         _UpperCAmelCase	,			_UpperCAmelCase							=    self.image_processing_tester.get_expected_values(lowerCamelCase       ,				batched=lowerCamelCase							)

         _UpperCAmelCase							=    image_processor(
             lowerCamelCase       ,				["""semantic"""] * len(lowerCamelCase							)       ,				return_tensors="""pt"""							).pixel_values
         self.assertEqual(
             encoded_images.shape       ,				(
                 self.image_processing_tester.batch_size,
                 self.image_processing_tester.num_channels,
                 expected_height,
                 expected_width,
             )       ,				)


  def  lowerCamelCase       (      self		:			Tuple       ,				lowerCamelCase		:			Tuple=False       ,				lowerCamelCase		:			Union[str, Any]=False       ,				lowerCamelCase		:			Union[str, Any]="np"							)						->	str:


         """simple docstring"""
         _UpperCAmelCase							=    self.image_processing_class(**self.image_processor_dict							)
         # prepare image and target
         _UpperCAmelCase							=    self.image_processing_tester.num_labels
         _UpperCAmelCase							=    None
         _UpperCAmelCase							=    None
         _UpperCAmelCase							=    prepare_image_inputs(self.image_processing_tester       ,				equal_resolution=lowerCamelCase							)
         if with_segmentation_maps:
                _UpperCAmelCase							=    num_labels
                if is_instance_map:
                       _UpperCAmelCase							=    list(range(lowerCamelCase							)							) * 2
                       _UpperCAmelCase							=    dict(enumerate(lowerCamelCase							)							)
                _UpperCAmelCase							=    [
                    np.random.randint(0       ,				high * 2       ,				(img.size[1], img.size[0])							).astype(np.uinta							) for img in image_inputs
                ]
                if segmentation_type == "pil":
                       _UpperCAmelCase							=    [Image.fromarray(lowerCamelCase							) for annotation in annotations]

         _UpperCAmelCase							=    image_processor(
             lowerCamelCase       ,				["""semantic"""] * len(lowerCamelCase							)       ,				lowerCamelCase       ,				return_tensors="""pt"""       ,				instance_id_to_semantic_id=lowerCamelCase       ,				pad_and_return_pixel_mask=lowerCamelCase       ,				)

         return inputs


  def  lowerCamelCase       (      self		:			Any							)						->	Dict:


         """simple docstring"""
         pass


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


         """simple docstring"""
         def common(lowerCamelCase		:			List[Any]=False       ,				lowerCamelCase		:			List[Any]=None							):
                _UpperCAmelCase							=    self.comm_get_image_processor_inputs(
                    with_segmentation_maps=lowerCamelCase       ,				is_instance_map=lowerCamelCase       ,				segmentation_type=lowerCamelCase							)

                _UpperCAmelCase							=    inputs["""mask_labels"""]
                _UpperCAmelCase							=    inputs["""class_labels"""]
                _UpperCAmelCase							=    inputs["""pixel_values"""]
                _UpperCAmelCase							=    inputs["""text_inputs"""]

                # check the batch_size
                for mask_label, class_label, text_input in zip(lowerCamelCase       ,				lowerCamelCase       ,				lowerCamelCase							):
                       self.assertEqual(mask_label.shape[0]       ,				class_label.shape[0]							)
                       # this ensure padding has happened
                       self.assertEqual(mask_label.shape[1:]       ,				pixel_values.shape[2:]							)
                       self.assertEqual(len(lowerCamelCase							)       ,				self.image_processing_tester.num_text							)

         common()
         common(is_instance_map=lowerCamelCase							)
         common(is_instance_map=lowerCamelCase       ,				segmentation_type="""pil"""							)
         common(is_instance_map=lowerCamelCase       ,				segmentation_type="""pil"""							)


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


         """simple docstring"""
         _UpperCAmelCase							=    np.zeros((20, 50)							)
         _UpperCAmelCase							=    1
         _UpperCAmelCase							=    1
         _UpperCAmelCase							=    1

         _UpperCAmelCase							=    binary_mask_to_rle(lowerCamelCase							)
         self.assertEqual(len(lowerCamelCase							)       ,				4							)
         self.assertEqual(rle[0]       ,				21							)
         self.assertEqual(rle[1]       ,				45							)


  def  lowerCamelCase       (      self		:			Optional[Any]							)						->	Optional[Any]:


         """simple docstring"""
         _UpperCAmelCase							=    self.image_processing_class(
             num_labels=self.image_processing_tester.num_classes       ,				max_seq_length=77       ,				task_seq_length=77       ,				class_info_file="""ade20k_panoptic.json"""       ,				num_text=self.image_processing_tester.num_text       ,				repo_path="""shi-labs/oneformer_demo"""       ,				)
         _UpperCAmelCase							=    self.image_processing_tester.get_fake_oneformer_outputs()

         _UpperCAmelCase							=    fature_extractor.post_process_semantic_segmentation(lowerCamelCase							)

         self.assertEqual(len(lowerCamelCase							)       ,				self.image_processing_tester.batch_size							)
         self.assertEqual(
             segmentation[0].shape       ,				(
                 self.image_processing_tester.height,
                 self.image_processing_tester.width,
             )       ,				)

         _UpperCAmelCase							=    [(1, 4) for i in range(self.image_processing_tester.batch_size							)]
         _UpperCAmelCase							=    fature_extractor.post_process_semantic_segmentation(lowerCamelCase       ,				target_sizes=lowerCamelCase							)

         self.assertEqual(segmentation[0].shape       ,				target_sizes[0]							)


  def  lowerCamelCase       (      self		:			Dict							)						->	Any:


         """simple docstring"""
         _UpperCAmelCase							=    self.image_processing_class(
             num_labels=self.image_processing_tester.num_classes       ,				max_seq_length=77       ,				task_seq_length=77       ,				class_info_file="""ade20k_panoptic.json"""       ,				num_text=self.image_processing_tester.num_text       ,				repo_path="""shi-labs/oneformer_demo"""       ,				)
         _UpperCAmelCase							=    self.image_processing_tester.get_fake_oneformer_outputs()
         _UpperCAmelCase							=    image_processor.post_process_instance_segmentation(lowerCamelCase       ,				threshold=0							)

         self.assertTrue(len(lowerCamelCase							) == self.image_processing_tester.batch_size							)
         for el in segmentation:
                self.assertTrue("""segmentation""" in el							)
                self.assertTrue("""segments_info""" in el							)
                self.assertEqual(type(el["""segments_info"""]							)       ,				lowerCamelCase							)
                self.assertEqual(
                    el["""segmentation"""].shape       ,				(self.image_processing_tester.height, self.image_processing_tester.width)							)





  def  lowerCamelCase       (      self		:			Dict							)						->	Optional[Any]:


         """simple docstring"""
         _UpperCAmelCase							=    self.image_processing_class(
             num_labels=self.image_processing_tester.num_classes       ,				max_seq_length=77       ,				task_seq_length=77       ,				class_info_file="""ade20k_panoptic.json"""       ,				num_text=self.image_processing_tester.num_text       ,				repo_path="""shi-labs/oneformer_demo"""       ,				)
         _UpperCAmelCase							=    self.image_processing_tester.get_fake_oneformer_outputs()
         _UpperCAmelCase							=    image_processor.post_process_panoptic_segmentation(lowerCamelCase       ,				threshold=0							)

         self.assertTrue(len(lowerCamelCase							) == self.image_processing_tester.batch_size							)
         for el in segmentation:
                self.assertTrue("""segmentation""" in el							)
                self.assertTrue("""segments_info""" in el							)
                self.assertEqual(type(el["""segments_info"""]							)       ,				lowerCamelCase							)
                self.assertEqual(
                    el["""segmentation"""].shape       ,				(self.image_processing_tester.height, self.image_processing_tester.width)							)