



import unittest

import numpy as np

from transformers import DistilBertConfig, 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.numpy as jnp

 from transformers.models.distilbert.modeling_flax_distilbert import (
     FlaxDistilBertForMaskedLM,
     FlaxDistilBertForMultipleChoice,
     FlaxDistilBertForQuestionAnswering,
     FlaxDistilBertForSequenceClassification,
     FlaxDistilBertForTokenClassification,
     FlaxDistilBertModel,
 )

class    _A	(     unittest.TestCase   ):



       def __init__(		self							:					Tuple					,		__SCREAMING_SNAKE_CASE							:					Optional[int]					,		__SCREAMING_SNAKE_CASE							:					int=13					,		__SCREAMING_SNAKE_CASE							:					List[str]=7					,		__SCREAMING_SNAKE_CASE							:					str=True					,		__SCREAMING_SNAKE_CASE							:					str=True					,		__SCREAMING_SNAKE_CASE							:					Optional[Any]=True					,		__SCREAMING_SNAKE_CASE							:					Dict=True					,		__SCREAMING_SNAKE_CASE							:					Optional[Any]=99					,		__SCREAMING_SNAKE_CASE							:					Dict=32					,		__SCREAMING_SNAKE_CASE							:					str=5					,		__SCREAMING_SNAKE_CASE							:					str=4					,		__SCREAMING_SNAKE_CASE							:					Union[str, Any]=37					,		__SCREAMING_SNAKE_CASE							:					List[str]="gelu"					,		__SCREAMING_SNAKE_CASE							:					List[str]=0.1					,		__SCREAMING_SNAKE_CASE							:					Optional[int]=0.1					,		__SCREAMING_SNAKE_CASE							:					Optional[Any]=512					,		__SCREAMING_SNAKE_CASE							:					Tuple=16					,		__SCREAMING_SNAKE_CASE							:					Dict=2					,		__SCREAMING_SNAKE_CASE							:					str=0.02					,		__SCREAMING_SNAKE_CASE							:					str=4					,		):


           '''simple docstring'''

           __a							=    parent
           __a							=    batch_size
           __a							=    seq_length
           __a							=    is_training
           __a							=    use_attention_mask
           __a							=    use_token_type_ids
           __a							=    use_labels
           __a							=    vocab_size
           __a							=    hidden_size
           __a							=    num_hidden_layers
           __a							=    num_attention_heads
           __a							=    intermediate_size
           __a							=    hidden_act
           __a							=    hidden_dropout_prob
           __a							=    attention_probs_dropout_prob
           __a							=    max_position_embeddings
           __a							=    type_vocab_size
           __a							=    type_sequence_label_size
           __a							=    initializer_range
           __a							=    num_choices



       def 		_lowerCamelCase      (		self							:					List[str]):


           '''simple docstring'''

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

           __a							=    None
           if self.use_attention_mask:
               __a							=    random_attention_mask([self.batch_size, self.seq_length])

           __a							=    DistilBertConfig(
               vocab_size=self.vocab_size					,		dim=self.hidden_size					,		n_layers=self.num_hidden_layers					,		n_heads=self.num_attention_heads					,		hidden_dim=self.intermediate_size					,		hidden_act=self.hidden_act					,		dropout=self.hidden_dropout_prob					,		attention_dropout=self.attention_probs_dropout_prob					,		max_position_embeddings=self.max_position_embeddings					,		initializer_range=self.initializer_range					,		tie_weights_=__SCREAMING_SNAKE_CASE					,		)

           return config, input_ids, attention_mask





       def 		_lowerCamelCase      (		self							:					Union[str, Any]):


           '''simple docstring'''

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


@require_flax
class    _A	(     __UpperCAmelCase				,unittest.TestCase   ):
       UpperCamelCase__  :   Optional[Any]						 =						(
           (
               FlaxDistilBertModel,
               FlaxDistilBertForMaskedLM,
               FlaxDistilBertForMultipleChoice,
               FlaxDistilBertForQuestionAnswering,
               FlaxDistilBertForSequenceClassification,
               FlaxDistilBertForTokenClassification,
               FlaxDistilBertForQuestionAnswering,
           )
           if is_flax_available()
           else ()
       )



       def 		_lowerCamelCase      (		self							:					str):


           '''simple docstring'''

           __a							=    FlaxDistilBertModelTester(self)





       @slow
       def 		_lowerCamelCase      (		self							:					int):


           '''simple docstring'''

           for model_class_name in self.all_model_classes:
               __a							=    model_class_name.from_pretrained('''distilbert-base-uncased''')
               __a							=    model(np.ones((1, 1)))
               self.assertIsNotNone(__SCREAMING_SNAKE_CASE)


@require_flax
class    _A	(     unittest.TestCase   ):



       @slow
       def 		_lowerCamelCase      (		self							:					Optional[int]):


           '''simple docstring'''

           __a							=    FlaxDistilBertModel.from_pretrained('''distilbert-base-uncased''')
           __a							=    np.array([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]])
           __a							=    np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])
           __a							=    model(__SCREAMING_SNAKE_CASE					,		attention_mask=__SCREAMING_SNAKE_CASE)[0]
           __a							=    (1, 11, 768)
           self.assertEqual(output.shape					,		__SCREAMING_SNAKE_CASE)
           __a							=    np.array([[[-0.16_39, 0.32_99, 0.16_48], [-0.17_46, 0.32_89, 0.17_10], [-0.18_84, 0.33_57, 0.18_10]]])

           self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]					,		__SCREAMING_SNAKE_CASE					,		atol=1E-4))










from typing import List, Union

from ..utils import (
    add_end_docstrings,
    is_tf_available,
    is_torch_available,
    is_vision_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
 from PIL import Image

 from ..image_utils import load_image

if is_tf_available():
 from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_VISION_2_SEQ_MAPPING

if is_torch_available():
 import torch

 from ..models.auto.modeling_auto import MODEL_FOR_VISION_2_SEQ_MAPPING

__snake_case						:Dict						=					logging.get_logger(__name__)

@add_end_docstrings(__UpperCAmelCase   )
class    _A	(     __UpperCAmelCase   ):



       def __init__(		self							:					int					,		*__SCREAMING_SNAKE_CASE							:					Optional[int]					,		**__SCREAMING_SNAKE_CASE							:					str):


           '''simple docstring'''

           super().__init__(*__SCREAMING_SNAKE_CASE					,		**__SCREAMING_SNAKE_CASE)
           requires_backends(self					,		'''vision''')
           self.check_model_type(
               TF_MODEL_FOR_VISION_2_SEQ_MAPPING if self.framework == '''tf''' else MODEL_FOR_VISION_2_SEQ_MAPPING)



       def 		_lowerCamelCase      (		self							:					int					,		__SCREAMING_SNAKE_CASE							:					int=None					,		__SCREAMING_SNAKE_CASE							:					int=None					,		__SCREAMING_SNAKE_CASE							:					Tuple=None):


           '''simple docstring'''

           __a							=    {}
           __a							=    {}

           if prompt is not None:
               __a							=    prompt

           if generate_kwargs is not None:
               __a							=    generate_kwargs
           if max_new_tokens is not None:
               if "generate_kwargs" not in forward_kwargs:
                   __a							=    {}
               if "max_new_tokens" in forward_kwargs["generate_kwargs"]:
                   raise ValueError(
                       '''\'max_new_tokens\' is defined twice, once in \'generate_kwargs\' and once as a direct parameter,'''
                       ''' please use only one''')
               __a							=    max_new_tokens
           return preprocess_params, forward_kwargs, {}



       def __call__(		self							:					Any					,		__SCREAMING_SNAKE_CASE							:					Union[str, List[str], "Image.Image", List["Image.Image"]]					,		**__SCREAMING_SNAKE_CASE							:					Any):


           '''simple docstring'''

           return super().__call__(__SCREAMING_SNAKE_CASE					,		**__SCREAMING_SNAKE_CASE)



       def 		_lowerCamelCase      (		self							:					Optional[int]					,		__SCREAMING_SNAKE_CASE							:					Dict					,		__SCREAMING_SNAKE_CASE							:					Optional[Any]=None):


           '''simple docstring'''

           __a							=    load_image(__SCREAMING_SNAKE_CASE)

           if prompt is not None:
               if not isinstance(__SCREAMING_SNAKE_CASE					,		__SCREAMING_SNAKE_CASE):
                   raise ValueError(
                       F'Received an invalid text input, got - {type(__SCREAMING_SNAKE_CASE)} - but expected a single string. '
                       '''Note also that one single text can be provided for conditional image to text generation.''')

               __a							=    self.model.config.model_type

               if model_type == "git":
                   __a							=    self.image_processor(images=__SCREAMING_SNAKE_CASE					,		return_tensors=self.framework)
                   __a							=    self.tokenizer(text=__SCREAMING_SNAKE_CASE					,		add_special_tokens=__SCREAMING_SNAKE_CASE).input_ids
                   __a							=    [self.tokenizer.cls_token_id] + input_ids
                   __a							=    torch.tensor(__SCREAMING_SNAKE_CASE).unsqueeze(0)
                   model_inputs.update({'''input_ids''': input_ids})

               elif model_type == "pix2struct":
                   __a							=    self.image_processor(images=__SCREAMING_SNAKE_CASE					,		header_text=__SCREAMING_SNAKE_CASE					,		return_tensors=self.framework)

               elif model_type != "vision-encoder-decoder":
                   # vision-encoder-decoder does not support conditional generation
                   __a							=    self.image_processor(images=__SCREAMING_SNAKE_CASE					,		return_tensors=self.framework)
                   __a							=    self.tokenizer(__SCREAMING_SNAKE_CASE					,		return_tensors=self.framework)
                   model_inputs.update(__SCREAMING_SNAKE_CASE)

               else:
                   raise ValueError(F'Model type {model_type} does not support conditional text generation')

           else:
               __a							=    self.image_processor(images=__SCREAMING_SNAKE_CASE					,		return_tensors=self.framework)

           if self.model.config.model_type == "git" and prompt is None:
               __a							=    None

           return model_inputs



       def 		_lowerCamelCase      (		self							:					List[Any]					,		__SCREAMING_SNAKE_CASE							:					str					,		__SCREAMING_SNAKE_CASE							:					Tuple=None):


           '''simple docstring'''

           if (
               "input_ids" in model_inputs
               and isinstance(model_inputs['''input_ids''']					,		__SCREAMING_SNAKE_CASE)
               and all(x is None for x in model_inputs['''input_ids'''])
           ):
               __a							=    None

           if generate_kwargs is None:
               __a							=    {}
           # FIXME: We need to pop here due to a difference in how `generation.py` and `generation.tf_utils.py`
           #  parse inputs. In the Tensorflow version, `generate` raises an error if we don't use `input_ids` whereas
           #  the PyTorch version matches it with `self.model.main_input_name` or `self.model.encoder.main_input_name`
           #  in the `_prepare_model_inputs` method.
           __a							=    model_inputs.pop(self.model.main_input_name)
           __a							=    self.model.generate(__SCREAMING_SNAKE_CASE					,		**__SCREAMING_SNAKE_CASE					,		**__SCREAMING_SNAKE_CASE)
           return model_outputs





       def 		_lowerCamelCase      (		self							:					str					,		__SCREAMING_SNAKE_CASE							:					Union[str, Any]):


           '''simple docstring'''

           __a							=    []
           for output_ids in model_outputs:
               __a							=    {
                   '''generated_text''': self.tokenizer.decode(
                       __SCREAMING_SNAKE_CASE					,		skip_special_tokens=__SCREAMING_SNAKE_CASE					,		)
               }
               records.append(__SCREAMING_SNAKE_CASE)
           return records





