




"""simple docstring"""







import unittest

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

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      RemBertConfig,
      RemBertForCausalLM,
      RemBertForMaskedLM,
      RemBertForMultipleChoice,
      RemBertForQuestionAnswering,
      RemBertForSequenceClassification,
      RemBertForTokenClassification,
      RemBertModel,
  )
  from transformers.models.rembert.modeling_rembert import REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST






class  SCREAMING_SNAKE_CASE__   :


      def __init__(							self			:     Optional[Any]  ,  lowerCAmelCase			:     Dict  ,  lowerCAmelCase			:     Union[str, Any]=13  ,  lowerCAmelCase			:     Union[str, Any]=7  ,  lowerCAmelCase			:     Union[str, Any]=True  ,  lowerCAmelCase			:     str=True  ,  lowerCAmelCase			:     List[Any]=True  ,  lowerCAmelCase			:     Dict=True  ,  lowerCAmelCase			:     Optional[int]=99  ,  lowerCAmelCase			:     Optional[int]=32  ,  lowerCAmelCase			:     List[Any]=18  ,  lowerCAmelCase			:     List[str]=43  ,  lowerCAmelCase			:     Dict=5  ,  lowerCAmelCase			:     Tuple=4  ,  lowerCAmelCase			:     Union[str, Any]=37  ,  lowerCAmelCase			:     Optional[Any]="gelu"  ,  lowerCAmelCase			:     List[str]=0.1  ,  lowerCAmelCase			:     Optional[Any]=0.1  ,  lowerCAmelCase			:     Any=512  ,  lowerCAmelCase			:     Optional[Any]=16  ,  lowerCAmelCase			:     Any=2  ,  lowerCAmelCase			:     str=0.02  ,  lowerCAmelCase			:     Any=3  ,  lowerCAmelCase			:     Optional[int]=4  ,  lowerCAmelCase			:     Union[str, Any]=None  ,  ):
          lowerCAmelCase		=			parent
          lowerCAmelCase		=			batch_size
          lowerCAmelCase		=			seq_length
          lowerCAmelCase		=			is_training
          lowerCAmelCase		=			use_input_mask
          lowerCAmelCase		=			use_token_type_ids
          lowerCAmelCase		=			use_labels
          lowerCAmelCase		=			vocab_size
          lowerCAmelCase		=			hidden_size
          lowerCAmelCase		=			input_embedding_size
          lowerCAmelCase		=			output_embedding_size
          lowerCAmelCase		=			num_hidden_layers
          lowerCAmelCase		=			num_attention_heads
          lowerCAmelCase		=			intermediate_size
          lowerCAmelCase		=			hidden_act
          lowerCAmelCase		=			hidden_dropout_prob
          lowerCAmelCase		=			attention_probs_dropout_prob
          lowerCAmelCase		=			max_position_embeddings
          lowerCAmelCase		=			type_vocab_size
          lowerCAmelCase		=			type_sequence_label_size
          lowerCAmelCase		=			initializer_range
          lowerCAmelCase		=			num_labels
          lowerCAmelCase		=			num_choices
          lowerCAmelCase		=			scope


      def 					__lowercase	(							self			:     Optional[Any]   ):
          lowerCAmelCase		=			ids_tensor([self.batch_size, self.seq_length]  ,  self.vocab_size   )

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

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

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

          lowerCAmelCase		=			RemBertConfig(
              vocab_size=self.vocab_size  ,  hidden_size=self.hidden_size  ,  input_embedding_size=self.input_embedding_size  ,  output_embedding_size=self.output_embedding_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  ,  )

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


      def 					__lowercase	(							self			:     int   ):
          (
              (
              lowerCAmelCase   
          )	,  (
              lowerCAmelCase   
          )	,  (
              lowerCAmelCase   
          )	,  (
              lowerCAmelCase   
          )	,  (
              lowerCAmelCase   
          )	,  (
              lowerCAmelCase   
          )	,  (
              lowerCAmelCase   
          )	,     
          )		=			self.prepare_config_and_inputs()

          lowerCAmelCase		=			True
          lowerCAmelCase		=			floats_tensor([self.batch_size, self.seq_length, self.hidden_size]   )
          lowerCAmelCase		=			ids_tensor([self.batch_size, self.seq_length]  ,  vocab_size=2   )

          return (
              config,
              input_ids,
              token_type_ids,
              input_mask,
              sequence_labels,
              token_labels,
              choice_labels,
              encoder_hidden_states,
              encoder_attention_mask,
          )


      def 					__lowercase	(							self			:     List[str]  ,  lowerCAmelCase			:     Any  ,  lowerCAmelCase			:     Optional[int]  ,  lowerCAmelCase			:     Any  ,  lowerCAmelCase			:     str  ,  lowerCAmelCase			:     Any  ,  lowerCAmelCase			:     int  ,  lowerCAmelCase			:     int   ):
          lowerCAmelCase		=			RemBertModel(config=lowerCAmelCase   )
          model.to(lowerCAmelCase   )
          model.eval()
          lowerCAmelCase		=			model(lowerCAmelCase  ,  attention_mask=lowerCAmelCase  ,  token_type_ids=lowerCAmelCase   )
          lowerCAmelCase		=			model(lowerCAmelCase  ,  token_type_ids=lowerCAmelCase   )
          lowerCAmelCase		=			model(lowerCAmelCase   )
          self.parent.assertEqual(result.last_hidden_state.shape  ,  (self.batch_size, self.seq_length, self.hidden_size)   )


      def 					__lowercase	(							self			:     Any  ,  lowerCAmelCase			:     int  ,  lowerCAmelCase			:     str  ,  lowerCAmelCase			:     List[str]  ,  lowerCAmelCase			:     List[str]  ,  lowerCAmelCase			:     Optional[int]  ,  lowerCAmelCase			:     Optional[Any]  ,  lowerCAmelCase			:     Tuple  ,  lowerCAmelCase			:     Any  ,  lowerCAmelCase			:     Optional[Any]  ,  ):
          lowerCAmelCase		=			True
          lowerCAmelCase		=			RemBertModel(lowerCAmelCase   )
          model.to(lowerCAmelCase   )
          model.eval()
          lowerCAmelCase		=			model(
              lowerCAmelCase  ,  attention_mask=lowerCAmelCase  ,  token_type_ids=lowerCAmelCase  ,  encoder_hidden_states=lowerCAmelCase  ,  encoder_attention_mask=lowerCAmelCase  ,  )
          lowerCAmelCase		=			model(
              lowerCAmelCase  ,  attention_mask=lowerCAmelCase  ,  token_type_ids=lowerCAmelCase  ,  encoder_hidden_states=lowerCAmelCase  ,  )
          lowerCAmelCase		=			model(lowerCAmelCase  ,  attention_mask=lowerCAmelCase  ,  token_type_ids=lowerCAmelCase   )
          self.parent.assertEqual(result.last_hidden_state.shape  ,  (self.batch_size, self.seq_length, self.hidden_size)   )


      def 					__lowercase	(							self			:     str  ,  lowerCAmelCase			:     int  ,  lowerCAmelCase			:     Tuple  ,  lowerCAmelCase			:     str  ,  lowerCAmelCase			:     str  ,  lowerCAmelCase			:     Any  ,  lowerCAmelCase			:     Optional[Any]  ,  lowerCAmelCase			:     Union[str, Any]  ,  lowerCAmelCase			:     int  ,  lowerCAmelCase			:     Union[str, Any]  ,  ):
          lowerCAmelCase		=			RemBertForCausalLM(config=lowerCAmelCase   )
          model.to(lowerCAmelCase   )
          model.eval()
          lowerCAmelCase		=			model(lowerCAmelCase  ,  attention_mask=lowerCAmelCase  ,  token_type_ids=lowerCAmelCase  ,  labels=lowerCAmelCase   )
          self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.seq_length, self.vocab_size)   )


      def 					__lowercase	(							self			:     Optional[Any]  ,  lowerCAmelCase			:     Union[str, Any]  ,  lowerCAmelCase			:     Dict  ,  lowerCAmelCase			:     Any  ,  lowerCAmelCase			:     Optional[Any]  ,  lowerCAmelCase			:     Union[str, Any]  ,  lowerCAmelCase			:     Any  ,  lowerCAmelCase			:     Any   ):
          lowerCAmelCase		=			RemBertForMaskedLM(config=lowerCAmelCase   )
          model.to(lowerCAmelCase   )
          model.eval()
          lowerCAmelCase		=			model(lowerCAmelCase  ,  attention_mask=lowerCAmelCase  ,  token_type_ids=lowerCAmelCase  ,  labels=lowerCAmelCase   )
          self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.seq_length, self.vocab_size)   )


      def 					__lowercase	(							self			:     Optional[Any]  ,  lowerCAmelCase			:     str  ,  lowerCAmelCase			:     Any  ,  lowerCAmelCase			:     int  ,  lowerCAmelCase			:     List[str]  ,  lowerCAmelCase			:     Optional[int]  ,  lowerCAmelCase			:     List[str]  ,  lowerCAmelCase			:     Tuple  ,  lowerCAmelCase			:     int  ,  lowerCAmelCase			:     Optional[Any]  ,  ):
          lowerCAmelCase		=			True
          lowerCAmelCase		=			True
          lowerCAmelCase		=			RemBertForCausalLM(config=lowerCAmelCase   )
          model.to(lowerCAmelCase   )
          model.eval()

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

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

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

          lowerCAmelCase		=			model(
              lowerCAmelCase  ,  attention_mask=lowerCAmelCase  ,  encoder_hidden_states=lowerCAmelCase  ,  encoder_attention_mask=lowerCAmelCase  ,  output_hidden_states=lowerCAmelCase  ,  )["""hidden_states"""][0]
          lowerCAmelCase		=			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
          lowerCAmelCase		=			ids_tensor((1,)  ,  output_from_past.shape[-1]   ).item()
          lowerCAmelCase		=			output_from_no_past[:, -3:, random_slice_idx].detach()
          lowerCAmelCase		=			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 					__lowercase	(							self			:     Any  ,  lowerCAmelCase			:     Union[str, Any]  ,  lowerCAmelCase			:     Optional[int]  ,  lowerCAmelCase			:     Dict  ,  lowerCAmelCase			:     Union[str, Any]  ,  lowerCAmelCase			:     str  ,  lowerCAmelCase			:     Dict  ,  lowerCAmelCase			:     Optional[Any]   ):
          lowerCAmelCase		=			RemBertForQuestionAnswering(config=lowerCAmelCase   )
          model.to(lowerCAmelCase   )
          model.eval()
          lowerCAmelCase		=			model(
              lowerCAmelCase  ,  attention_mask=lowerCAmelCase  ,  token_type_ids=lowerCAmelCase  ,  start_positions=lowerCAmelCase  ,  end_positions=lowerCAmelCase  ,  )
          self.parent.assertEqual(result.start_logits.shape  ,  (self.batch_size, self.seq_length)   )
          self.parent.assertEqual(result.end_logits.shape  ,  (self.batch_size, self.seq_length)   )


      def 					__lowercase	(							self			:     Any  ,  lowerCAmelCase			:     Union[str, Any]  ,  lowerCAmelCase			:     List[str]  ,  lowerCAmelCase			:     List[str]  ,  lowerCAmelCase			:     Dict  ,  lowerCAmelCase			:     Union[str, Any]  ,  lowerCAmelCase			:     Union[str, Any]  ,  lowerCAmelCase			:     Dict   ):
          lowerCAmelCase		=			self.num_labels
          lowerCAmelCase		=			RemBertForSequenceClassification(lowerCAmelCase   )
          model.to(lowerCAmelCase   )
          model.eval()
          lowerCAmelCase		=			model(lowerCAmelCase  ,  attention_mask=lowerCAmelCase  ,  token_type_ids=lowerCAmelCase  ,  labels=lowerCAmelCase   )
          self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.num_labels)   )


      def 					__lowercase	(							self			:     Any  ,  lowerCAmelCase			:     Union[str, Any]  ,  lowerCAmelCase			:     Tuple  ,  lowerCAmelCase			:     Dict  ,  lowerCAmelCase			:     Union[str, Any]  ,  lowerCAmelCase			:     str  ,  lowerCAmelCase			:     List[Any]  ,  lowerCAmelCase			:     int   ):
          lowerCAmelCase		=			self.num_labels
          lowerCAmelCase		=			RemBertForTokenClassification(config=lowerCAmelCase   )
          model.to(lowerCAmelCase   )
          model.eval()
          lowerCAmelCase		=			model(lowerCAmelCase  ,  attention_mask=lowerCAmelCase  ,  token_type_ids=lowerCAmelCase  ,  labels=lowerCAmelCase   )
          self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.seq_length, self.num_labels)   )


      def 					__lowercase	(							self			:     Any  ,  lowerCAmelCase			:     Optional[Any]  ,  lowerCAmelCase			:     Optional[Any]  ,  lowerCAmelCase			:     int  ,  lowerCAmelCase			:     int  ,  lowerCAmelCase			:     Any  ,  lowerCAmelCase			:     Optional[int]  ,  lowerCAmelCase			:     List[str]   ):
          lowerCAmelCase		=			self.num_choices
          lowerCAmelCase		=			RemBertForMultipleChoice(config=lowerCAmelCase   )
          model.to(lowerCAmelCase   )
          model.eval()
          lowerCAmelCase		=			input_ids.unsqueeze(1   ).expand(-1  ,  self.num_choices  ,  -1   ).contiguous()
          lowerCAmelCase		=			token_type_ids.unsqueeze(1   ).expand(-1  ,  self.num_choices  ,  -1   ).contiguous()
          lowerCAmelCase		=			input_mask.unsqueeze(1   ).expand(-1  ,  self.num_choices  ,  -1   ).contiguous()
          lowerCAmelCase		=			model(
              lowerCAmelCase  ,  attention_mask=lowerCAmelCase  ,  token_type_ids=lowerCAmelCase  ,  labels=lowerCAmelCase  ,  )
          self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.num_choices)   )




      def 					__lowercase	(							self			:     Optional[int]   ):
          lowerCAmelCase		=			self.prepare_config_and_inputs()
          (
              (
              lowerCAmelCase   
          )	,  (
              lowerCAmelCase   
          )	,  (
              lowerCAmelCase   
          )	,  (
              lowerCAmelCase   
          )	,  (
              lowerCAmelCase   
          )	,  (
              lowerCAmelCase   
          )	,  (
              lowerCAmelCase   
          )	,     
          )		=			config_and_inputs
          lowerCAmelCase		=			{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
          return config, inputs_dict






@require_torch
class  SCREAMING_SNAKE_CASE__   (   _a  ,			_a  ,			unittest.TestCase   ):
      _a        =  (
          (
              RemBertModel,
              RemBertForMaskedLM,
              RemBertForCausalLM,
              RemBertForMultipleChoice,
              RemBertForQuestionAnswering,
              RemBertForSequenceClassification,
              RemBertForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      _a        =  (RemBertForCausalLM,) if is_torch_available() else ()
      _a        =  (
          {
              'feature-extraction': RemBertModel,
              'fill-mask': RemBertForMaskedLM,
              'question-answering': RemBertForQuestionAnswering,
              'text-classification': RemBertForSequenceClassification,
              'text-generation': RemBertForCausalLM,
              'token-classification': RemBertForTokenClassification,
              'zero-shot': RemBertForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )


      def 					__lowercase	(							self			:     Any   ):
          lowerCAmelCase		=			RemBertModelTester(self   )
          lowerCAmelCase		=			ConfigTester(self  ,  config_class=lowerCAmelCase  ,  hidden_size=37   )


      def 					__lowercase	(							self			:     str   ):
          self.config_tester.run_common_tests()


      def 					__lowercase	(							self			:     int   ):
          lowerCAmelCase		=			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*lowerCAmelCase   )


      def 					__lowercase	(							self			:     Union[str, Any]   ):
          lowerCAmelCase		=			self.model_tester.prepare_config_and_inputs()
          for type in ["absolute", "relative_key", "relative_key_query"]:
              lowerCAmelCase		=			type
              self.model_tester.create_and_check_model(*lowerCAmelCase   )


      def 					__lowercase	(							self			:     Optional[int]   ):
          lowerCAmelCase		=			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_lm(*lowerCAmelCase   )


      def 					__lowercase	(							self			:     Any   ):
          lowerCAmelCase		=			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_multiple_choice(*lowerCAmelCase   )


      def 					__lowercase	(							self			:     int   ):
          lowerCAmelCase		=			self.model_tester.prepare_config_and_inputs_for_decoder()
          self.model_tester.create_and_check_decoder_model_past_large_inputs(*lowerCAmelCase   )


      def 					__lowercase	(							self			:     Union[str, Any]   ):
          lowerCAmelCase		=			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_question_answering(*lowerCAmelCase   )


      def 					__lowercase	(							self			:     str   ):
          lowerCAmelCase		=			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_sequence_classification(*lowerCAmelCase   )


      def 					__lowercase	(							self			:     str   ):
          lowerCAmelCase		=			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_token_classification(*lowerCAmelCase   )


      def 					__lowercase	(							self			:     Tuple   ):
          lowerCAmelCase		=			self.model_tester.prepare_config_and_inputs_for_decoder()
          self.model_tester.create_and_check_model_as_decoder(*lowerCAmelCase   )


      def 					__lowercase	(							self			:     Union[str, Any]   ):
          (
              (
              lowerCAmelCase   
          )	,  (
              lowerCAmelCase   
          )	,  (
              lowerCAmelCase   
          )	,  (
              lowerCAmelCase   
          )	,  (
              lowerCAmelCase   
          )	,  (
              lowerCAmelCase   
          )	,  (
              lowerCAmelCase   
          )	,  (
              lowerCAmelCase   
          )	,  (
              lowerCAmelCase   
          )	,     
          )		=			self.model_tester.prepare_config_and_inputs_for_decoder()

          lowerCAmelCase		=			None

          self.model_tester.create_and_check_model_as_decoder(
              lowerCAmelCase  ,  lowerCAmelCase  ,  lowerCAmelCase  ,  lowerCAmelCase  ,  lowerCAmelCase  ,  lowerCAmelCase  ,  lowerCAmelCase  ,  lowerCAmelCase  ,  lowerCAmelCase  ,  )




      @slow
      def 					__lowercase	(							self			:     str   ):



          for model_name in REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              lowerCAmelCase		=			RemBertModel.from_pretrained(lowerCAmelCase   )
              self.assertIsNotNone(lowerCAmelCase   )






@require_torch
class  SCREAMING_SNAKE_CASE__   (   unittest.TestCase   ):


      @slow
      def 					__lowercase	(							self			:     Optional[int]   ):
          lowerCAmelCase		=			RemBertModel.from_pretrained("""google/rembert"""   )
          lowerCAmelCase		=			torch.tensor([[312, 5_6498, 313, 2125, 313]]   )
          lowerCAmelCase		=			torch.tensor([[0, 0, 0, 1, 1]]   )
          with torch.no_grad():
              lowerCAmelCase		=			model(lowerCAmelCase  ,  token_type_ids=lowerCAmelCase  ,  output_hidden_states=lowerCAmelCase   )

          lowerCAmelCase		=			1152

          lowerCAmelCase		=			torch.Size((1, 5, hidden_size)   )
          self.assertEqual(output["""last_hidden_state"""].shape  ,  lowerCAmelCase   )

          lowerCAmelCase		=			torch.tensor(
              [
                  [
                      [0.0754, -0.2022, 0.1904],
                      [-0.3354, -0.3692, -0.4791],
                      [-0.2314, -0.6729, -0.0749],
                      [-0.0396, -0.3105, -0.4234],
                      [-0.1571, -0.0525, 0.5353],
                  ]
              ]   )

          # Running on the original tf implementation gives slightly different results here.
          # Not clear why this variations is present
          # TODO: Find reason for discrepancy
          # expected_original_implementation = [[
          #     [0.07630594074726105, -0.20146065950393677, 0.19107051193714142],
          #     [-0.3405614495277405, -0.36971670389175415, -0.4808273911476135],
          #     [-0.22587086260318756, -0.6656315922737122, -0.07844287157058716],
          #     [-0.04145475849509239, -0.3077218234539032, -0.42316967248916626],
          #     [-0.15887849032878876, -0.054529931396245956, 0.5356100797653198]
          # ]]

          self.assertTrue(torch.allclose(output["""last_hidden_state"""][:, :, :3]  ,  lowerCAmelCase  ,  atol=1e-4   )   )











"""simple docstring"""







import colorsys
import io

import cva
import matplotlib as mpl
import matplotlib.colors as mplc
import matplotlib.figure as mplfigure
import numpy as np
import torch
from matplotlib.backends.backend_agg import FigureCanvasAgg

from utils import img_tensorize


a								=							1_0_0_0






class  SCREAMING_SNAKE_CASE__   :


						def __init__(							self			:     Any  ,  lowerCAmelCase			:     List[Any]  ,  lowerCAmelCase			:     Tuple=1.2  ,  lowerCAmelCase			:     str="g"  ,  lowerCAmelCase			:     List[Any]=0.5  ,  lowerCAmelCase			:     List[Any]="-"  ,  lowerCAmelCase			:     Optional[int]="test_out.jpg"  ,  lowerCAmelCase			:     Union[str, Any]=True  ,  lowerCAmelCase			:     Tuple=False  ,  lowerCAmelCase			:     Optional[int]=None  ,  lowerCAmelCase			:     Tuple=None  ,  lowerCAmelCase			:     str=0.7  ,  ):
										if isinstance(lowerCAmelCase  ,  torch.Tensor   ):
														lowerCAmelCase		=			img.numpy().astype("""np.uint8"""   )
										if isinstance(lowerCAmelCase  ,  lowerCAmelCase   ):
														lowerCAmelCase		=			img_tensorize(lowerCAmelCase   )
										assert isinstance(lowerCAmelCase  ,  np.ndarray   )

										lowerCAmelCase	,  lowerCAmelCase		=			img.shape[1], img.shape[0]
										lowerCAmelCase		=			mplfigure.Figure(frameon=lowerCAmelCase   )
										lowerCAmelCase		=			fig.get_dpi()
										lowerCAmelCase		=			(width * scale + 1e-2) / dpi
										lowerCAmelCase		=			(height * scale + 1e-2) / dpi
										fig.set_size_inches(lowerCAmelCase  ,  lowerCAmelCase   )
										lowerCAmelCase		=			fig.add_axes([0.0, 0.0, 1.0, 1.0]   )
										ax.axis("""off"""   )
										ax.set_xlim(0.0  ,  lowerCAmelCase   )
										ax.set_ylim(lowerCAmelCase   )

										lowerCAmelCase		=			saveas
										lowerCAmelCase		=			rgb
										lowerCAmelCase		=			pynb
										lowerCAmelCase		=			img
										lowerCAmelCase		=			edgecolor
										lowerCAmelCase		=			0.5
										lowerCAmelCase		=			linestyle
										lowerCAmelCase		=			int(np.sqrt(min(lowerCAmelCase  ,  lowerCAmelCase   )   ) * scale // 3   )
										lowerCAmelCase		=			width
										lowerCAmelCase		=			height
										lowerCAmelCase		=			scale
										lowerCAmelCase		=			fig
										lowerCAmelCase		=			ax
										lowerCAmelCase		=			pad
										lowerCAmelCase		=			idaobj
										lowerCAmelCase		=			idaattr
										lowerCAmelCase		=			FigureCanvasAgg(lowerCAmelCase   )


						def 					__lowercase	(							self			:     str  ,  lowerCAmelCase			:     Optional[int]  ,  lowerCAmelCase			:     List[str]=None   ):
										if color is None:
														lowerCAmelCase		=			self.edgecolor
										((lowerCAmelCase)	,  (lowerCAmelCase)	,  (lowerCAmelCase)	,  (lowerCAmelCase))		=			box
										lowerCAmelCase		=			xa - xa
										lowerCAmelCase		=			ya - ya
										self.ax.add_patch(
										    mpl.patches.Rectangle(
										        (xa, ya)  ,  lowerCAmelCase  ,  lowerCAmelCase  ,  fill=lowerCAmelCase  ,  edgecolor=lowerCAmelCase  ,  linewidth=self.font_size // 3  ,  alpha=self.alpha  ,  linestyle=self.linestyle  ,  )   )


						def 					__lowercase	(							self			:     int  ,  lowerCAmelCase			:     Any  ,  lowerCAmelCase			:     int=None  ,  lowerCAmelCase			:     Tuple=None  ,  lowerCAmelCase			:     Optional[Any]=None  ,  lowerCAmelCase			:     Dict=None   ):
										if len(boxes.shape   ) > 2:
														lowerCAmelCase		=			boxes[0]
										if len(obj_ids.shape   ) > 1:
														lowerCAmelCase		=			obj_ids[0]
										if len(obj_scores.shape   ) > 1:
														lowerCAmelCase		=			obj_scores[0]
										if len(attr_ids.shape   ) > 1:
														lowerCAmelCase		=			attr_ids[0]
										if len(attr_scores.shape   ) > 1:
														lowerCAmelCase		=			attr_scores[0]
										if isinstance(lowerCAmelCase  ,  torch.Tensor   ):
														lowerCAmelCase		=			boxes.numpy()
										if isinstance(lowerCAmelCase  ,  lowerCAmelCase   ):
														lowerCAmelCase		=			np.array(lowerCAmelCase   )
										assert isinstance(lowerCAmelCase  ,  np.ndarray   )
										lowerCAmelCase		=			np.prod(boxes[:, 2:] - boxes[:, :2]  ,  axis=1   )
										lowerCAmelCase		=			np.argsort(-areas   ).tolist()
										lowerCAmelCase		=			boxes[sorted_idxs] if boxes is not None else None
										lowerCAmelCase		=			obj_ids[sorted_idxs] if obj_ids is not None else None
										lowerCAmelCase		=			obj_scores[sorted_idxs] if obj_scores is not None else None
										lowerCAmelCase		=			attr_ids[sorted_idxs] if attr_ids is not None else None
										lowerCAmelCase		=			attr_scores[sorted_idxs] if attr_scores is not None else None

										lowerCAmelCase		=			[self._random_color(maximum=1   ) for _ in range(len(lowerCAmelCase   )   )]
										lowerCAmelCase		=			[assigned_colors[idx] for idx in sorted_idxs]
										if obj_ids is not None:
														lowerCAmelCase		=			self._create_text_labels_attr(lowerCAmelCase  ,  lowerCAmelCase  ,  lowerCAmelCase  ,  lowerCAmelCase   )
														for i in range(len(lowerCAmelCase   )   ):
																		lowerCAmelCase		=			assigned_colors[i]
																		self.add_box(boxes[i]  ,  lowerCAmelCase   )
																		self.draw_labels(labels[i]  ,  boxes[i]  ,  lowerCAmelCase   )


						def 					__lowercase	(							self			:     int  ,  lowerCAmelCase			:     Optional[Any]  ,  lowerCAmelCase			:     Any  ,  lowerCAmelCase			:     Dict   ):
										lowerCAmelCase	,  lowerCAmelCase	,  lowerCAmelCase	,  lowerCAmelCase		=			box
										lowerCAmelCase		=			(xa, ya)
										lowerCAmelCase		=			(ya - ya) * (xa - xa)
										lowerCAmelCase		=			_SMALL_OBJ * self.scale
										if instance_area < small or ya - ya < 40 * self.scale:
														if ya >= self.height - 5:
																		lowerCAmelCase		=			(xa, ya)
														else:
																		lowerCAmelCase		=			(xa, ya)

										lowerCAmelCase		=			(ya - ya) / np.sqrt(self.height * self.width   )
										lowerCAmelCase		=			self._change_color_brightness(lowerCAmelCase  ,  brightness_factor=0.7   )
										lowerCAmelCase		=			np.clip((height_ratio - 0.02) / 0.08 + 1  ,  1.2  ,  2   )
										font_size *= 0.75 * self.font_size

										self.draw_text(
										    text=lowerCAmelCase  ,  position=lowerCAmelCase  ,  color=lowerCAmelCase  ,  )


						def 					__lowercase	(							self			:     str  ,  lowerCAmelCase			:     List[str]  ,  lowerCAmelCase			:     Optional[int]  ,  lowerCAmelCase			:     Optional[int]="g"  ,  lowerCAmelCase			:     Optional[int]="left"  ,  ):
										lowerCAmelCase		=			0
										lowerCAmelCase		=			self.font_size
										lowerCAmelCase		=			np.maximum(list(mplc.to_rgb(lowerCAmelCase   )   )  ,  0.2   )
										lowerCAmelCase		=			max(0.8  ,  np.max(lowerCAmelCase   )   )
										lowerCAmelCase		=			{
										    """facecolor""": """black""",
										    """alpha""": self.alpha,
										    """pad""": self.pad,
										    """edgecolor""": """none""",
										}
										lowerCAmelCase	,  lowerCAmelCase		=			position
										self.ax.text(
										    lowerCAmelCase  ,  lowerCAmelCase  ,  lowerCAmelCase  ,  size=font_size * self.scale  ,  family="""sans-serif"""  ,  bbox=lowerCAmelCase  ,  verticalalignment="""top"""  ,  horizontalalignment=lowerCAmelCase  ,  color=lowerCAmelCase  ,  zorder=10  ,  rotation=lowerCAmelCase  ,  )


						def 					__lowercase	(							self			:     str  ,  lowerCAmelCase			:     Tuple=None   ):
										if saveas is None:
														lowerCAmelCase		=			self.saveas
										if saveas.lower().endswith(""".jpg"""   ) or saveas.lower().endswith(""".png"""   ):
														cva.imwrite(
														    lowerCAmelCase  ,  self._get_buffer()[:, :, ::-1]  ,  )
										else:
														self.fig.savefig(lowerCAmelCase   )


						def 					__lowercase	(							self			:     str  ,  lowerCAmelCase			:     int  ,  lowerCAmelCase			:     List[str]  ,  lowerCAmelCase			:     Any  ,  lowerCAmelCase			:     Optional[Any]   ):
										lowerCAmelCase		=			[self.idaobj[i] for i in classes]
										lowerCAmelCase		=			[self.idaattr[i] for i in attr_classes]
										lowerCAmelCase		=			[
										    f'''{label} {score:.2f} {attr} {attr_score:.2f}'''
										    for label, score, attr, attr_score in zip(lowerCAmelCase  ,  lowerCAmelCase  ,  lowerCAmelCase  ,  lowerCAmelCase   )
										]
										return labels


						def 					__lowercase	(							self			:     Any  ,  lowerCAmelCase			:     Dict  ,  lowerCAmelCase			:     Union[str, Any]   ):
										lowerCAmelCase		=			[self.idaobj[i] for i in classes]
										if scores is not None:
														if labels is None:
																		lowerCAmelCase		=			["""{:.0f}%""".format(s * 100   ) for s in scores]
														else:
																		lowerCAmelCase		=			["""{} {:.0f}%""".format(lowerCAmelCase  ,  s * 100   ) for li, s in zip(lowerCAmelCase  ,  lowerCAmelCase   )]
										return labels


						def 					__lowercase	(							self			:     Any  ,  lowerCAmelCase			:     str=255   ):
										lowerCAmelCase		=			np.random.randint(0  ,  len(_COLORS   )   )
										lowerCAmelCase		=			_COLORS[idx] * maximum
										if not self.rgb:
														lowerCAmelCase		=			ret[::-1]
										return ret


						def 					__lowercase	(							self			:     List[str]   ):
										if not self.pynb:
														lowerCAmelCase	,  lowerCAmelCase		=			self.canvas.print_to_buffer()
														if (width, height) != (self.width, self.height):
																		lowerCAmelCase		=			cva.resize(self.img  ,  (width, height)   )
														else:
																		lowerCAmelCase		=			self.img
										else:
														lowerCAmelCase		=			io.BytesIO()  # works for cairo backend
														self.canvas.print_rgba(lowerCAmelCase   )
														lowerCAmelCase	,  lowerCAmelCase		=			self.width, self.height
														lowerCAmelCase		=			buf.getvalue()
														lowerCAmelCase		=			self.img

										lowerCAmelCase		=			np.frombuffer(lowerCAmelCase  ,  dtype="""uint8"""   )
										lowerCAmelCase		=			buffer.reshape(lowerCAmelCase  ,  lowerCAmelCase  ,  4   )
										lowerCAmelCase	,  lowerCAmelCase		=			np.split(lowerCAmelCase  ,  [3]  ,  axis=2   )

										try:
														import numexpr as ne  # fuse them with numexpr

														lowerCAmelCase		=			ne.evaluate("""img * (1 - alpha / 255.0) + rgb * (alpha / 255.0)"""   )
										except ImportError:
														lowerCAmelCase		=			alpha.astype("""float32"""   ) / 255.0
														lowerCAmelCase		=			img * (1 - alpha) + rgb * alpha

										return visualized_image.astype("""uint8"""   )




						def 					__lowercase	(							self			:     List[str]  ,  lowerCAmelCase			:     Optional[Any]  ,  lowerCAmelCase			:     Optional[int]   ):
										assert brightness_factor >= -1.0 and brightness_factor <= 1.0
										lowerCAmelCase		=			mplc.to_rgb(lowerCAmelCase   )
										lowerCAmelCase		=			colorsys.rgb_to_hls(*mplc.to_rgb(lowerCAmelCase   )   )
										lowerCAmelCase		=			polygon_color[1] + (brightness_factor * polygon_color[1])
										lowerCAmelCase		=			0.0 if modified_lightness < 0.0 else modified_lightness
										lowerCAmelCase		=			1.0 if modified_lightness > 1.0 else modified_lightness
										lowerCAmelCase		=			colorsys.hls_to_rgb(polygon_color[0]  ,  lowerCAmelCase  ,  polygon_color[2]   )
										return modified_color





a								=							(
    np.array(
        [
            0.000,
            0.447,
            0.741,
            0.850,
            0.325,
            0.098,
            0.929,
            0.694,
            0.125,
            0.494,
            0.184,
            0.556,
            0.466,
            0.674,
            0.188,
            0.301,
            0.745,
            0.933,
            0.635,
            0.078,
            0.184,
            0.300,
            0.300,
            0.300,
            0.600,
            0.600,
            0.600,
            1.000,
            0.000,
            0.000,
            1.000,
            0.500,
            0.000,
            0.749,
            0.749,
            0.000,
            0.000,
            1.000,
            0.000,
            0.000,
            0.000,
            1.000,
            0.667,
            0.000,
            1.000,
            0.333,
            0.333,
            0.000,
            0.333,
            0.667,
            0.000,
            0.333,
            1.000,
            0.000,
            0.667,
            0.333,
            0.000,
            0.667,
            0.667,
            0.000,
            0.667,
            1.000,
            0.000,
            1.000,
            0.333,
            0.000,
            1.000,
            0.667,
            0.000,
            1.000,
            1.000,
            0.000,
            0.000,
            0.333,
            0.500,
            0.000,
            0.667,
            0.500,
            0.000,
            1.000,
            0.500,
            0.333,
            0.000,
            0.500,
            0.333,
            0.333,
            0.500,
            0.333,
            0.667,
            0.500,
            0.333,
            1.000,
            0.500,
            0.667,
            0.000,
            0.500,
            0.667,
            0.333,
            0.500,
            0.667,
            0.667,
            0.500,
            0.667,
            1.000,
            0.500,
            1.000,
            0.000,
            0.500,
            1.000,
            0.333,
            0.500,
            1.000,
            0.667,
            0.500,
            1.000,
            1.000,
            0.500,
            0.000,
            0.333,
            1.000,
            0.000,
            0.667,
            1.000,
            0.000,
            1.000,
            1.000,
            0.333,
            0.000,
            1.000,
            0.333,
            0.333,
            1.000,
            0.333,
            0.667,
            1.000,
            0.333,
            1.000,
            1.000,
            0.667,
            0.000,
            1.000,
            0.667,
            0.333,
            1.000,
            0.667,
            0.667,
            1.000,
            0.667,
            1.000,
            1.000,
            1.000,
            0.000,
            1.000,
            1.000,
            0.333,
            1.000,
            1.000,
            0.667,
            1.000,
            0.333,
            0.000,
            0.000,
            0.500,
            0.000,
            0.000,
            0.667,
            0.000,
            0.000,
            0.833,
            0.000,
            0.000,
            1.000,
            0.000,
            0.000,
            0.000,
            0.167,
            0.000,
            0.000,
            0.333,
            0.000,
            0.000,
            0.500,
            0.000,
            0.000,
            0.667,
            0.000,
            0.000,
            0.833,
            0.000,
            0.000,
            1.000,
            0.000,
            0.000,
            0.000,
            0.167,
            0.000,
            0.000,
            0.333,
            0.000,
            0.000,
            0.500,
            0.000,
            0.000,
            0.667,
            0.000,
            0.000,
            0.833,
            0.000,
            0.000,
            1.000,
            0.000,
            0.000,
            0.000,
            0.143,
            0.143,
            0.143,
            0.857,
            0.857,
            0.857,
            1.000,
            1.000,
            1.000,
        ]
    )
    .astype(np.floataa)
    .reshape(-1, 3)
)





