

import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTJConfig, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, tooslow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
  import jax
  import jax.numpy as jnp

  from transformers.modeling_flax_pytorch_utils import (
      convert_pytorch_state_dict_to_flax,
      load_flax_weights_in_pytorch_model,
  )
  from transformers.models.gptj.modeling_flax_gptj import FlaxGPTJForCausalLM, FlaxGPTJModel

if is_torch_available():
  import torch





class 					__lowerCAmelCase	:




  def __init__(							self				,  lowerCAmelCase				,  lowerCAmelCase=14				,  lowerCAmelCase=7				,  lowerCAmelCase=True				,  lowerCAmelCase=True				,  lowerCAmelCase=False				,  lowerCAmelCase=True				,  lowerCAmelCase=99				,  lowerCAmelCase=32				,  lowerCAmelCase=4				,  lowerCAmelCase=4				,  lowerCAmelCase=4				,  lowerCAmelCase=37				,  lowerCAmelCase="gelu"				,  lowerCAmelCase=0.1				,  lowerCAmelCase=0.1				,  lowerCAmelCase=512				,  lowerCAmelCase=0.02				,  )    ->					Optional[Any]:
        '''simple docstring'''






        _lowercase      			=parent
        _lowercase      			=batch_size
        _lowercase      			=seq_length
        _lowercase      			=is_training
        _lowercase      			=use_input_mask
        _lowercase      			=use_token_type_ids
        _lowercase      			=use_labels
        _lowercase      			=vocab_size
        _lowercase      			=hidden_size
        _lowercase      			=rotary_dim
        _lowercase      			=num_hidden_layers
        _lowercase      			=num_attention_heads
        _lowercase      			=intermediate_size
        _lowercase      			=hidden_act
        _lowercase      			=hidden_dropout_prob
        _lowercase      			=attention_probs_dropout_prob
        _lowercase      			=max_position_embeddings
        _lowercase      			=initializer_range
        _lowercase      			=None
        _lowercase      			=vocab_size - 1
        _lowercase      			=vocab_size - 1
        _lowercase      			=vocab_size - 1




  def     A__	(							self   )    ->					Any:
        '''simple docstring'''






        _lowercase      			=ids_tensor([self.batch_size, self.seq_length]				,  self.vocab_size   )

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

        _lowercase      			=GPTJConfig(
            vocab_size=self.vocab_size				,  n_embd=self.hidden_size				,  n_layer=self.num_hidden_layers				,  n_head=self.num_attention_heads				,  n_positions=self.max_position_embeddings				,  use_cache=lowerCAmelCase				,  bos_token_id=self.bos_token_id				,  eos_token_id=self.eos_token_id				,  pad_token_id=self.pad_token_id				,  rotary_dim=self.rotary_dim				,  )

        return (config, input_ids, input_mask)




  def     A__	(							self   )    ->					Optional[Any]:
        '''simple docstring'''






        _lowercase      			=self.prepare_config_and_inputs()
        _lowercase			,			_lowercase			,			_lowercase      			=config_and_inputs
        _lowercase      			={'input_ids': input_ids, 'attention_mask': attention_mask}
        return config, inputs_dict




  def     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase				,  lowerCAmelCase				,  lowerCAmelCase   )    ->					Dict:
        '''simple docstring'''






        _lowercase      			=20
        _lowercase      			=model_class_name(lowerCAmelCase   )

        _lowercase      			=model.init_cache(input_ids.shape[0]				,  lowerCAmelCase   )
        _lowercase      			=jnp.ones((input_ids.shape[0], max_decoder_length)				,  dtype='i4'   )

        _lowercase      			=jnp.broadcast_to(
            jnp.arange(input_ids.shape[-1] - 1   )[None, :]				,  (input_ids.shape[0], input_ids.shape[-1] - 1)   )
        _lowercase      			=model(
            input_ids[:, :-1]				,  attention_mask=lowerCAmelCase				,  past_key_values=lowerCAmelCase				,  position_ids=lowerCAmelCase				,  )

        _lowercase      			=jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]				,  dtype='i4'   )
        _lowercase      			=model(
            input_ids[:, -1:]				,  attention_mask=lowerCAmelCase				,  past_key_values=outputs_cache.past_key_values				,  position_ids=lowerCAmelCase				,  )

        _lowercase      			=model(lowerCAmelCase   )

        _lowercase      			=np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])   )   )
        self.parent.assertTrue(diff < 1e-3				,  msg=F'''Max diff is {diff}'''   )




  def     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase				,  lowerCAmelCase				,  lowerCAmelCase   )    ->					List[str]:
        '''simple docstring'''






        _lowercase      			=20
        _lowercase      			=model_class_name(lowerCAmelCase   )

        _lowercase      			=jnp.concatenate(
            [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])   )]				,  axis=-1				,  )

        _lowercase      			=model.init_cache(input_ids.shape[0]				,  lowerCAmelCase   )
        _lowercase      			=jnp.broadcast_to(
            jnp.arange(input_ids.shape[-1] - 1   )[None, :]				,  (input_ids.shape[0], input_ids.shape[-1] - 1)   )

        _lowercase      			=model(
            input_ids[:, :-1]				,  attention_mask=lowerCAmelCase				,  past_key_values=lowerCAmelCase				,  position_ids=lowerCAmelCase				,  )
        _lowercase      			=jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]				,  dtype='i4'   )
        _lowercase      			=model(
            input_ids[:, -1:]				,  past_key_values=outputs_cache.past_key_values				,  attention_mask=lowerCAmelCase				,  position_ids=lowerCAmelCase				,  )

        _lowercase      			=model(lowerCAmelCase				,  attention_mask=lowerCAmelCase   )

        _lowercase      			=np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])   )   )
        self.parent.assertTrue(diff < 1e-3				,  msg=F'''Max diff is {diff}'''   )





@require_flax
class 					__lowerCAmelCase	(   SCREAMING_SNAKE_CASE     ,     SCREAMING_SNAKE_CASE     ,     unittest.TestCase				):
  _a		     =    (FlaxGPTJModel, FlaxGPTJForCausalLM) if is_flax_available() else ()
  _a		     =    (FlaxGPTJForCausalLM,) if is_flax_available() else ()




  def     A__	(							self   )    ->					List[Any]:
        '''simple docstring'''






        _lowercase      			=FlaxGPTJModelTester(self   )




  def     A__	(							self   )    ->					Tuple:
        '''simple docstring'''






        for model_class_name in self.all_model_classes:
              _lowercase			,			_lowercase			,			_lowercase      			=self.model_tester.prepare_config_and_inputs()
              self.model_tester.check_use_cache_forward(lowerCAmelCase				,  lowerCAmelCase				,  lowerCAmelCase				,  lowerCAmelCase   )




  def     A__	(							self   )    ->					List[str]:
        '''simple docstring'''






        for model_class_name in self.all_model_classes:
              _lowercase			,			_lowercase			,			_lowercase      			=self.model_tester.prepare_config_and_inputs()
              self.model_tester.check_use_cache_forward_with_attn_mask(
                  lowerCAmelCase				,  lowerCAmelCase				,  lowerCAmelCase				,  lowerCAmelCase   )




  @tooslow
  def     A__	(							self   )    ->					str:
        '''simple docstring'''






        _lowercase      			=GPTaTokenizer.from_pretrained('gpt2'				,  pad_token='<|endoftext|>'				,  padding_side='left'   )
        _lowercase      			=tokenizer(['Hello this is a long string', 'Hey']				,  return_tensors='np'				,  padding=lowerCAmelCase				,  truncation=lowerCAmelCase   )

        _lowercase      			=FlaxGPTJForCausalLM.from_pretrained('EleutherAI/gpt-j-6B'   )
        _lowercase      			=False
        _lowercase      			=model.config.eos_token_id

        _lowercase      			=jax.jit(model.generate   )

        _lowercase      			=jit_generate(
            inputs['input_ids']				,  attention_mask=inputs['attention_mask']				,  pad_token_id=tokenizer.pad_token_id   ).sequences

        _lowercase      			=tokenizer.batch_decode(lowerCAmelCase				,  skip_special_tokens=lowerCAmelCase   )

        _lowercase      			=[
            'Hello this is a long string of text.\n\nI\'m trying to get the text of the',
            'Hey, I\'m a little late to the party. I\'m going to',
        ]

        self.assertListEqual(lowerCAmelCase				,  lowerCAmelCase   )




  @is_pt_flax_cross_test
  def     A__	(							self   )    ->					List[Any]:
        '''simple docstring'''






        _lowercase			,			_lowercase      			=self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
              with self.subTest(model_class.__name__   ):
                    # prepare inputs
                    _lowercase      			=self._prepare_for_class(lowerCAmelCase				,  lowerCAmelCase   )
                    _lowercase      			={k: torch.tensor(v.tolist()   ) for k, v in prepared_inputs_dict.items()}

                    # load corresponding PyTorch class
                    _lowercase      			=model_class.__name__[4:]  # Skip the "Flax" at the beginning
                    _lowercase      			=getattr(lowerCAmelCase				,  lowerCAmelCase   )

                    _lowercase			,			_lowercase      			=pt_inputs['input_ids'].shape
                    _lowercase      			=np.random.randint(0				,  seq_length - 1				,  size=(batch_size,)   )
                    for batch_idx, start_index in enumerate(lowerCAmelCase   ):
                          _lowercase      			=0
                          _lowercase      			=1
                          _lowercase      			=0
                          _lowercase      			=1
                    _lowercase      			=pt_model_class(lowerCAmelCase   ).eval()
                    _lowercase      			=model_class(lowerCAmelCase				,  dtype=jnp.floataa   )

                    _lowercase      			=convert_pytorch_state_dict_to_flax(pt_model.state_dict()				,  lowerCAmelCase   )
                    _lowercase      			=fx_state

                    with torch.no_grad():
                          _lowercase      			=pt_model(**lowerCAmelCase   ).to_tuple()

                    _lowercase      			=fx_model(**lowerCAmelCase   ).to_tuple()
                    self.assertEqual(len(lowerCAmelCase   )				,  len(lowerCAmelCase   )				,  'Output lengths differ between Flax and PyTorch'   )
                    for fx_output, pt_output in zip(lowerCAmelCase				,  lowerCAmelCase   ):
                          self.assert_almost_equals(fx_output[:, -1]				,  pt_output[:, -1].numpy()				,  4e-2   )

                    with tempfile.TemporaryDirectory() as tmpdirname:
                          pt_model.save_pretrained(lowerCAmelCase   )
                          _lowercase      			=model_class.from_pretrained(lowerCAmelCase				,  from_pt=lowerCAmelCase   )

                    _lowercase      			=fx_model_loaded(**lowerCAmelCase   ).to_tuple()
                    self.assertEqual(
                        len(lowerCAmelCase   )				,  len(lowerCAmelCase   )				,  'Output lengths differ between Flax and PyTorch'   )
                    for fx_output_loaded, pt_output in zip(lowerCAmelCase				,  lowerCAmelCase   ):
                          self.assert_almost_equals(fx_output_loaded[:, -1]				,  pt_output[:, -1].numpy()				,  4e-2   )




  @is_pt_flax_cross_test
  def     A__	(							self   )    ->					Tuple:
        '''simple docstring'''






        _lowercase			,			_lowercase      			=self.model_tester.prepare_config_and_inputs_for_common()
        for model_class in self.all_model_classes:
              with self.subTest(model_class.__name__   ):
                    # prepare inputs
                    _lowercase      			=self._prepare_for_class(lowerCAmelCase				,  lowerCAmelCase   )
                    _lowercase      			={k: torch.tensor(v.tolist()   ) for k, v in prepared_inputs_dict.items()}

                    # load corresponding PyTorch class
                    _lowercase      			=model_class.__name__[4:]  # Skip the "Flax" at the beginning
                    _lowercase      			=getattr(lowerCAmelCase				,  lowerCAmelCase   )

                    _lowercase      			=pt_model_class(lowerCAmelCase   ).eval()
                    _lowercase      			=model_class(lowerCAmelCase				,  dtype=jnp.floataa   )

                    _lowercase      			=load_flax_weights_in_pytorch_model(lowerCAmelCase				,  fx_model.params   )
                    _lowercase			,			_lowercase      			=pt_inputs['input_ids'].shape
                    _lowercase      			=np.random.randint(0				,  seq_length - 1				,  size=(batch_size,)   )
                    for batch_idx, start_index in enumerate(lowerCAmelCase   ):
                          _lowercase      			=0
                          _lowercase      			=1
                          _lowercase      			=0
                          _lowercase      			=1

                    # make sure weights are tied in PyTorch
                    pt_model.tie_weights()

                    with torch.no_grad():
                          _lowercase      			=pt_model(**lowerCAmelCase   ).to_tuple()

                    _lowercase      			=fx_model(**lowerCAmelCase   ).to_tuple()
                    self.assertEqual(len(lowerCAmelCase   )				,  len(lowerCAmelCase   )				,  'Output lengths differ between Flax and PyTorch'   )
                    for fx_output, pt_output in zip(lowerCAmelCase				,  lowerCAmelCase   ):
                          self.assert_almost_equals(fx_output[:, -1]				,  pt_output[:, -1].numpy()				,  4e-2   )

                    with tempfile.TemporaryDirectory() as tmpdirname:
                          fx_model.save_pretrained(lowerCAmelCase   )
                          _lowercase      			=pt_model_class.from_pretrained(lowerCAmelCase				,  from_flax=lowerCAmelCase   )

                    with torch.no_grad():
                          _lowercase      			=pt_model_loaded(**lowerCAmelCase   ).to_tuple()

                    self.assertEqual(
                        len(lowerCAmelCase   )				,  len(lowerCAmelCase   )				,  'Output lengths differ between Flax and PyTorch'   )
                    for fx_output, pt_output in zip(lowerCAmelCase				,  lowerCAmelCase   ):
                          self.assert_almost_equals(fx_output[:, -1]				,  pt_output[:, -1].numpy()				,  4e-2   )




  @tooslow
  def     A__	(							self   )    ->					List[str]:
        '''simple docstring'''






        for model_class_name in self.all_model_classes:
              _lowercase      			=model_class_name.from_pretrained('EleutherAI/gpt-j-6B'   )
              _lowercase      			=model(np.ones((1, 1)   )   )
              self.assertIsNotNone(lowerCAmelCase   )





import torch
import torch.nn as nn
from transformers import CLIPConfig, CLIPVisionModel, PreTrainedModel

from ...utils import logging


lowercase_   							=  logging.get_logger(__name__)
def 	a (     A__      :  Optional[int] ,  A__      :  Union[str, Any]	)       -> int:





      """simple docstring"""



      _lowercase      			=nn.functional.normalize(A__	)
      _lowercase      			=nn.functional.normalize(A__	)
      return torch.mm(A__ ,  normalized_text_embeds.t()	)







class 					__lowerCAmelCase	(   SCREAMING_SNAKE_CASE				):
  _a		     =    CLIPConfig

  _a		     =    ["""CLIPEncoderLayer"""]




  def __init__(							self				,  lowerCAmelCase   )    ->					Union[str, Any]:
        '''simple docstring'''






        super().__init__(lowerCAmelCase   )

        _lowercase      			=CLIPVisionModel(config.vision_config   )
        _lowercase      			=nn.Linear(config.vision_config.hidden_size				,  config.projection_dim				,  bias=lowerCAmelCase   )

        _lowercase      			=nn.Parameter(torch.ones(17				,  config.projection_dim   )				,  requires_grad=lowerCAmelCase   )
        _lowercase      			=nn.Parameter(torch.ones(3				,  config.projection_dim   )				,  requires_grad=lowerCAmelCase   )

        _lowercase      			=nn.Parameter(torch.ones(17   )				,  requires_grad=lowerCAmelCase   )
        _lowercase      			=nn.Parameter(torch.ones(3   )				,  requires_grad=lowerCAmelCase   )




  @torch.no_grad()
  def     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase   )    ->					Optional[int]:
        '''simple docstring'''






        _lowercase      			=self.vision_model(lowerCAmelCase   )[1]  # pooled_output
        _lowercase      			=self.visual_projection(lowerCAmelCase   )

        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
        _lowercase      			=cosine_distance(lowerCAmelCase				,  self.special_care_embeds   ).cpu().float().numpy()
        _lowercase      			=cosine_distance(lowerCAmelCase				,  self.concept_embeds   ).cpu().float().numpy()

        _lowercase      			=[]
        _lowercase      			=image_embeds.shape[0]
        for i in range(lowerCAmelCase   ):
              _lowercase      			={'special_scores': {}, 'special_care': [], 'concept_scores': {}, 'bad_concepts': []}

              # increase this value to create a stronger `nfsw` filter
              # at the cost of increasing the possibility of filtering benign images
              _lowercase      			=0.0

              for concept_idx in range(len(special_cos_dist[0]   )   ):
                    _lowercase      			=special_cos_dist[i][concept_idx]
                    _lowercase      			=self.special_care_embeds_weights[concept_idx].item()
                    _lowercase      			=round(concept_cos - concept_threshold + adjustment				,  3   )
                    if result_img["special_scores"][concept_idx] > 0:
                          result_img["special_care"].append({concept_idx, result_img['special_scores'][concept_idx]}   )
                          _lowercase      			=0.01

              for concept_idx in range(len(cos_dist[0]   )   ):
                    _lowercase      			=cos_dist[i][concept_idx]
                    _lowercase      			=self.concept_embeds_weights[concept_idx].item()
                    _lowercase      			=round(concept_cos - concept_threshold + adjustment				,  3   )
                    if result_img["concept_scores"][concept_idx] > 0:
                          result_img["bad_concepts"].append(lowerCAmelCase   )

              result.append(lowerCAmelCase   )

        _lowercase      			=[len(res['bad_concepts']   ) > 0 for res in result]

        return images, has_nsfw_concepts




  @torch.no_grad()
  def     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase   )    ->					Tuple:
        '''simple docstring'''






        _lowercase      			=self.vision_model(lowerCAmelCase   )[1]  # pooled_output
        _lowercase      			=self.visual_projection(lowerCAmelCase   )

        _lowercase      			=cosine_distance(lowerCAmelCase				,  self.special_care_embeds   )
        _lowercase      			=cosine_distance(lowerCAmelCase				,  self.concept_embeds   )

        # increase this value to create a stronger `nsfw` filter
        # at the cost of increasing the possibility of filtering benign images
        _lowercase      			=0.0

        _lowercase      			=special_cos_dist - self.special_care_embeds_weights + adjustment
        # special_scores = special_scores.round(decimals=3)
        _lowercase      			=torch.any(special_scores > 0				,  dim=1   )
        _lowercase      			=special_care * 0.01
        _lowercase      			=special_adjustment.unsqueeze(1   ).expand(-1				,  cos_dist.shape[1]   )

        _lowercase      			=(cos_dist - self.concept_embeds_weights) + special_adjustment
        # concept_scores = concept_scores.round(decimals=3)
        _lowercase      			=torch.any(concept_scores > 0				,  dim=1   )

        return images, has_nsfw_concepts


