

import unittest

from transformers import BigBirdConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
     import jax

     from transformers.models.big_bird.modeling_flax_big_bird import (
         FlaxBigBirdForCausalLM,
         FlaxBigBirdForMaskedLM,
         FlaxBigBirdForMultipleChoice,
         FlaxBigBirdForPreTraining,
         FlaxBigBirdForQuestionAnswering,
         FlaxBigBirdForSequenceClassification,
         FlaxBigBirdForTokenClassification,
         FlaxBigBirdModel,
     )





class  __UpperCAmelCase					(	unittest.TestCase						):



   def __init__(							self							:					List[Any],		__A							:					List[Any],		__A							:					List[Any]=2,		__A							:					Optional[Any]=5_6,		__A							:					List[Any]=True,		__A							:					List[Any]=True,		__A							:					int=True,		__A							:					str=True,		__A							:					str=9_9,		__A							:					Optional[int]=3_2,		__A							:					List[Any]=2,		__A							:					Tuple=2,		__A							:					int=7,		__A							:					List[Any]="gelu_new",		__A							:					List[Any]=0.1,		__A							:					Union[str, Any]=0.1,		__A							:					str=5_1_2,		__A							:					List[str]=1_6,		__A							:					Union[str, Any]=2,		__A							:					Dict=0.0_2,		__A							:					Any=4,		__A							:					List[str]="block_sparse",		__A							:					int=True,		__A							:					str=False,		__A							:					Optional[Any]=2,		__A							:					List[Any]=3,		):
     UpperCAmelCase	:							List[Any]						=			parent
     UpperCAmelCase	:							Union[str, Any]						=			batch_size
     UpperCAmelCase	:							Optional[int]						=			seq_length
     UpperCAmelCase	:							Optional[int]						=			is_training
     UpperCAmelCase	:							Optional[Any]						=			use_attention_mask
     UpperCAmelCase	:							Dict						=			use_token_type_ids
     UpperCAmelCase	:							Any						=			use_labels
     UpperCAmelCase	:							List[str]						=			vocab_size
     UpperCAmelCase	:							List[Any]						=			hidden_size
     UpperCAmelCase	:							Tuple						=			num_hidden_layers
     UpperCAmelCase	:							Any						=			num_attention_heads
     UpperCAmelCase	:							int						=			intermediate_size
     UpperCAmelCase	:							List[Any]						=			hidden_act
     UpperCAmelCase	:							Optional[Any]						=			hidden_dropout_prob
     UpperCAmelCase	:							int						=			attention_probs_dropout_prob
     UpperCAmelCase	:							Optional[Any]						=			max_position_embeddings
     UpperCAmelCase	:							Tuple						=			type_vocab_size
     UpperCAmelCase	:							Tuple						=			type_sequence_label_size
     UpperCAmelCase	:							Any						=			initializer_range
     UpperCAmelCase	:							List[Any]						=			num_choices

     UpperCAmelCase	:							str						=			rescale_embeddings
     UpperCAmelCase	:							Any						=			attention_type
     UpperCAmelCase	:							Optional[int]						=			use_bias
     UpperCAmelCase	:							Optional[Any]						=			block_size
     UpperCAmelCase	:							Any						=			num_random_blocks



   def 	__magic_name__   (							self							:					Any						):
     UpperCAmelCase	:							Optional[int]						=			ids_tensor([self.batch_size, self.seq_length],		self.vocab_size						)

     UpperCAmelCase	:							List[Any]						=			None
     if self.use_attention_mask:
       UpperCAmelCase	:							Optional[Any]						=			random_attention_mask([self.batch_size, self.seq_length]						)

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

     UpperCAmelCase	:							str						=			BigBirdConfig(
         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=__A,		initializer_range=self.initializer_range,		attention_type=self.attention_type,		block_size=self.block_size,		num_random_blocks=self.num_random_blocks,		use_bias=self.use_bias,		rescale_embeddings=self.rescale_embeddings,		)

     return config, input_ids, token_type_ids, attention_mask





   def 	__magic_name__   (							self							:					Tuple						):
     UpperCAmelCase	:							List[Any]						=			self.prepare_config_and_inputs()
     UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							List[str]						=			config_and_inputs
     UpperCAmelCase	:							Optional[int]						=			{
         '''input_ids''': input_ids,
         '''token_type_ids''': token_type_ids,
         '''attention_mask''': attention_mask,
     }
     return config, inputs_dict






@require_flax
class  __UpperCAmelCase					(	lowerCamelCase__      ,  unittest.TestCase						):
   UpperCamelCase       					=					(
       (
           FlaxBigBirdForCausalLM,
           FlaxBigBirdModel,
           FlaxBigBirdForPreTraining,
           FlaxBigBirdForMaskedLM,
           FlaxBigBirdForMultipleChoice,
           FlaxBigBirdForQuestionAnswering,
           FlaxBigBirdForSequenceClassification,
           FlaxBigBirdForTokenClassification,
       )
       if is_flax_available()
       else ()
   )

   UpperCamelCase       					=					False
   UpperCamelCase       					=					False



   def 	__magic_name__   (							self							:					int						):
     UpperCAmelCase	:							List[str]						=			FlaxBigBirdModelTester(self						)



   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def 	__magic_name__   (							self							:					Any						):
     super().test_from_pretrained_save_pretrained()



   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def 	__magic_name__   (							self							:					List[Any]						):
     super().test_from_pretrained_with_no_automatic_init()



   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def 	__magic_name__   (							self							:					Union[str, Any]						):
     super().test_no_automatic_init()



   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def 	__magic_name__   (							self							:					Union[str, Any]						):
     super().test_hidden_states_output()



   @slow
   def 	__magic_name__   (							self							:					List[str]						):
     for model_class_name in self.all_model_classes:
       UpperCAmelCase	:							List[str]						=			model_class_name.from_pretrained('''google/bigbird-roberta-base'''						)
       self.assertIsNotNone(__A						)



   def 	__magic_name__   (							self							:					Any						):
     if self.test_attn_probs:
       super().test_attention_outputs()



   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def 	__magic_name__   (							self							:					int						):
     UpperCAmelCase			,						UpperCAmelCase	:							Dict						=			self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
       with self.subTest(model_class.__name__						):
         UpperCAmelCase	:							Any						=			self._prepare_for_class(__A,		__A						)
         UpperCAmelCase	:							Optional[Any]						=			model_class(__A						)

         @jax.jit
         def model_jitted(__A							:					Optional[Any],		__A							:					List[str]=None,		**__A							:					List[Any]						):
           return model(input_ids=__A,		attention_mask=__A,		**__A						)

         with self.subTest('''JIT Enabled'''						):
           UpperCAmelCase	:							List[Any]						=			model_jitted(**__A						).to_tuple()

         with self.subTest('''JIT Disabled'''						):
           with jax.disable_jit():
             UpperCAmelCase	:							str						=			model_jitted(**__A						).to_tuple()

         self.assertEqual(len(__A						),		len(__A						)						)
         for jitted_output, output in zip(__A,		__A						):
           self.assertEqual(jitted_output.shape,		output.shape						)





   def 	__magic_name__   (							self							:					Union[str, Any],		__A							:					Union[str, Any],		__A							:					Optional[int],		__A							:					str,		__A							:					Any=1E-5,		__A							:					Tuple="outputs",		__A							:					Optional[Any]=None						):
     # `bigbird_block_sparse_attention` in `FlaxBigBird` returns `attention_probs = None`, while in PyTorch version,
     # an effort was done to return `attention_probs` (yet to be verified).
     if name.startswith('''outputs.attentions'''						):
       return
     else:
       super().check_pt_flax_outputs(__A,		__A,		__A,		__A,		__A,		__A						)








from typing import Union

from ..utils import add_end_docstrings, is_torch_available, is_vision_available, logging
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
     from PIL import Image

     from ..image_utils import load_image

if is_torch_available():
     from ..models.auto.modeling_auto import MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING

_lowerCamelCase    :       Dict        =							logging.get_logger(__name__)





@add_end_docstrings(lowerCamelCase__						)
class  __UpperCAmelCase					(	lowerCamelCase__						):



   def __init__(							self							:					Optional[Any],		*__A							:					Tuple,		**__A							:					Tuple						):
     super().__init__(*__A,		**__A						)
     self.check_model_type(__A						)



   def 	__magic_name__   (							self							:					Union[str, Any],		__A							:					int=None,		__A							:					Tuple=None,		__A							:					Any=None,		**__A							:					Optional[int]						):
     UpperCAmelCase			,						UpperCAmelCase	:							List[Any]						=			{}, {}
     if padding is not None:
       UpperCAmelCase	:							Optional[int]						=			padding
     if truncation is not None:
       UpperCAmelCase	:							Optional[int]						=			truncation
     if top_k is not None:
       UpperCAmelCase	:							Tuple						=			top_k
     return preprocess_params, {}, postprocess_params



   def __call__(							self							:					Union[str, Any],		__A							:					Union["Image.Image", str],		__A							:					str = None,		**__A							:					Optional[int]						):
     if isinstance(__A,		(Image.Image, str)						) and isinstance(__A,		__A						):
       UpperCAmelCase	:							int						=			{'''image''': image, '''question''': question}
     else:
       UpperCAmelCase	:							str						=			image
     UpperCAmelCase	:							Union[str, Any]						=			super().__call__(__A,		**__A						)
     return results



   def 	__magic_name__   (							self							:					List[str],		__A							:					Union[str, Any],		__A							:					Tuple=False,		__A							:					List[Any]=False						):
     UpperCAmelCase	:							int						=			load_image(inputs['''image''']						)
     UpperCAmelCase	:							List[str]						=			self.tokenizer(
         inputs['''question'''],		return_tensors=self.framework,		padding=__A,		truncation=__A						)
     UpperCAmelCase	:							Union[str, Any]						=			self.image_processor(images=__A,		return_tensors=self.framework						)
     model_inputs.update(__A						)
     return model_inputs



   def 	__magic_name__   (							self							:					Optional[Any],		__A							:					List[Any]						):
     UpperCAmelCase	:							Optional[int]						=			self.model(**__A						)
     return model_outputs





   def 	__magic_name__   (							self							:					Any,		__A							:					List[str],		__A							:					Union[str, Any]=5						):
     if top_k > self.model.config.num_labels:
       UpperCAmelCase	:							Any						=			self.model.config.num_labels

     if self.framework == "pt":
       UpperCAmelCase	:							Any						=			model_outputs.logits.sigmoid()[0]
       UpperCAmelCase			,						UpperCAmelCase	:							Union[str, Any]						=			probs.topk(__A						)
     else:
       raise ValueError(F'''Unsupported framework: {self.framework}'''						)

     UpperCAmelCase	:							str						=			scores.tolist()
     UpperCAmelCase	:							Tuple						=			ids.tolist()
     return [{"score": score, "answer": self.model.config.idalabel[_id]} for score, _id in zip(__A,		__A						)]





