
"""simple docstring"""




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():
						import tensorflow as tf

						from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
						from ..tf_utils import stable_softmax

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

lowerCAmelCase__							=				logging.get_logger(__name__)





@add_end_docstrings(__lowercase   )
class  snake_case     (							__lowercase   ):

							def __init__(self       ,       *SCREAMING_SNAKE_CASE_       ,       **SCREAMING_SNAKE_CASE_    ):

									"""simple docstring"""







									super().__init__(*SCREAMING_SNAKE_CASE_       ,       **SCREAMING_SNAKE_CASE_    )
									requires_backends(self       ,       '''vision'''    )
									self.check_model_type(
									    TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
									    if self.framework == '''tf'''
									    else MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING    )

							def _lowercase					(self       ,       SCREAMING_SNAKE_CASE_=None    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = {}
									if top_k is not None:
											SCREAMING_SNAKE_CASE_						 = top_k
									return {}, {}, postprocess_params

							def __call__(self       ,       SCREAMING_SNAKE_CASE_       ,       **SCREAMING_SNAKE_CASE_    ):

									"""simple docstring"""







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

							def _lowercase					(self       ,       SCREAMING_SNAKE_CASE_    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = load_image(SCREAMING_SNAKE_CASE_    )
									SCREAMING_SNAKE_CASE_						 = self.image_processor(images=SCREAMING_SNAKE_CASE_       ,       return_tensors=self.framework    )
									return model_inputs

							def _lowercase					(self       ,       SCREAMING_SNAKE_CASE_    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = self.model(**SCREAMING_SNAKE_CASE_    )
									return model_outputs







							def _lowercase					(self       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_=5    ):

									"""simple docstring"""







									if top_k > self.model.config.num_labels:
											SCREAMING_SNAKE_CASE_						 = self.model.config.num_labels

									if self.framework == "pt":
											SCREAMING_SNAKE_CASE_						 = model_outputs.logits.softmax(-1    )[0]
											SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_						 = probs.topk(SCREAMING_SNAKE_CASE_    )
									elif self.framework == "tf":
											SCREAMING_SNAKE_CASE_						 = stable_softmax(model_outputs.logits       ,       axis=-1    )[0]
											SCREAMING_SNAKE_CASE_						 = tf.math.top_k(SCREAMING_SNAKE_CASE_       ,       k=SCREAMING_SNAKE_CASE_    )
											SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_						 = topk.values.numpy(), topk.indices.numpy()
									else:
											raise ValueError(f'Unsupported framework: {self.framework}'    )

									SCREAMING_SNAKE_CASE_						 = scores.tolist()
									SCREAMING_SNAKE_CASE_						 = ids.tolist()
									return [{"score": score, "label": self.model.config.idalabel[_id]} for score, _id in zip(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )]

"""simple docstring"""




import argparse

from transformers import BigBirdConfig, BigBirdForPreTraining, BigBirdForQuestionAnswering, load_tf_weights_in_big_bird
from transformers.utils import logging


logging.set_verbosity_info()



def        _lowerCamelCase ( __a,							__a,							__a,							__a    ):
		# Initialise PyTorch model
		SCREAMING_SNAKE_CASE_						 = BigBirdConfig.from_json_file(__a    )
		print(F'Building PyTorch model from configuration: {config}'    )

		if is_trivia_qa:
				SCREAMING_SNAKE_CASE_						 = BigBirdForQuestionAnswering(__a    )
		else:
				SCREAMING_SNAKE_CASE_						 = BigBirdForPreTraining(__a    )

		# Load weights from tf checkpoint
		load_tf_weights_in_big_bird(__a,							__a,							is_trivia_qa=__a    )

		# Save pytorch-model
		print(F'Save PyTorch model to {pytorch_dump_path}'    )
		model.save_pretrained(__a    )


if __name__ == "__main__":
						lowerCAmelCase__							=				argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    '--tf_checkpoint_path', default=None, type=str, required=True, help='Path to the TensorFlow checkpoint path.'
						)
						parser.add_argument(
						    '--big_bird_config_file',
						    default=None,
						    type=str,
						    required=True,
						    help=(
						        'The config json file corresponding to the pre-trained BERT model. \n'
						        'This specifies the model architecture.'
						    ),
						)
						parser.add_argument(
						    '--pytorch_dump_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
						)
						parser.add_argument(
						    '--is_trivia_qa', action='store_true', help='Whether to convert a model with a trivia_qa head.'
						)
						lowerCAmelCase__							=				parser.parse_args()
						convert_tf_checkpoint_to_pytorch(
						    args.tf_checkpoint_path, args.big_bird_config_file, args.pytorch_dump_path, args.is_trivia_qa
						)