"""simple docstring"""


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        a				(		unittest.TestCase   ):


    def __init__(					self	:  Dict   ,				__lowerCAmelCase	:  List[Any]   ,				__lowerCAmelCase	:  Any=2   ,				__lowerCAmelCase	:  Tuple=56   ,				__lowerCAmelCase	:  Any=True   ,				__lowerCAmelCase	:  Dict=True   ,				__lowerCAmelCase	:  Dict=True   ,				__lowerCAmelCase	:  Dict=True   ,				__lowerCAmelCase	:  Union[str, Any]=99   ,				__lowerCAmelCase	:  str=32   ,				__lowerCAmelCase	:  str=2   ,				__lowerCAmelCase	:  List[Any]=2   ,				__lowerCAmelCase	:  Dict=7   ,				__lowerCAmelCase	:  Optional[int]="gelu_new"   ,				__lowerCAmelCase	:  Tuple=0.1   ,				__lowerCAmelCase	:  Union[str, Any]=0.1   ,				__lowerCAmelCase	:  Optional[Any]=512   ,				__lowerCAmelCase	:  int=16   ,				__lowerCAmelCase	:  List[Any]=2   ,				__lowerCAmelCase	:  Union[str, Any]=0.02   ,				__lowerCAmelCase	:  Tuple=4   ,				__lowerCAmelCase	:  Any="block_sparse"   ,				__lowerCAmelCase	:  int=True   ,				__lowerCAmelCase	:  List[Any]=False   ,				__lowerCAmelCase	:  Tuple=2   ,				__lowerCAmelCase	:  int=3   ,				):
      _UpperCAmelCase						 =      parent
      _UpperCAmelCase						 =      batch_size
      _UpperCAmelCase						 =      seq_length
      _UpperCAmelCase						 =      is_training
      _UpperCAmelCase						 =      use_attention_mask
      _UpperCAmelCase						 =      use_token_type_ids
      _UpperCAmelCase						 =      use_labels
      _UpperCAmelCase						 =      vocab_size
      _UpperCAmelCase						 =      hidden_size
      _UpperCAmelCase						 =      num_hidden_layers
      _UpperCAmelCase						 =      num_attention_heads
      _UpperCAmelCase						 =      intermediate_size
      _UpperCAmelCase						 =      hidden_act
      _UpperCAmelCase						 =      hidden_dropout_prob
      _UpperCAmelCase						 =      attention_probs_dropout_prob
      _UpperCAmelCase						 =      max_position_embeddings
      _UpperCAmelCase						 =      type_vocab_size
      _UpperCAmelCase						 =      type_sequence_label_size
      _UpperCAmelCase						 =      initializer_range
      _UpperCAmelCase						 =      num_choices

      _UpperCAmelCase						 =      rescale_embeddings
      _UpperCAmelCase						 =      attention_type
      _UpperCAmelCase						 =      use_bias
      _UpperCAmelCase						 =      block_size
      _UpperCAmelCase						 =      num_random_blocks


    def      lowerCAmelCase_  (					self	:  int					):
      _UpperCAmelCase						 =      ids_tensor([self.batch_size, self.seq_length]   ,				self.vocab_size					)

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

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

      _UpperCAmelCase						 =      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=__lowerCAmelCase   ,				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      lowerCAmelCase_  (					self	:  List[str]					):
      _UpperCAmelCase						 =      self.prepare_config_and_inputs()
      _UpperCAmelCase ,				_UpperCAmelCase ,				_UpperCAmelCase ,				_UpperCAmelCase						 =      config_and_inputs
      _UpperCAmelCase						 =      {
          """input_ids""": input_ids,
          """token_type_ids""": token_type_ids,
          """attention_mask""": attention_mask,
      }
      return config, inputs_dict



@require_flax
class        a				(		lowerCAmelCase_	,							unittest.TestCase   ):
    _snake_case     :  List[Any]     							=							(
        (
            FlaxBigBirdForCausalLM,
            FlaxBigBirdModel,
            FlaxBigBirdForPreTraining,
            FlaxBigBirdForMaskedLM,
            FlaxBigBirdForMultipleChoice,
            FlaxBigBirdForQuestionAnswering,
            FlaxBigBirdForSequenceClassification,
            FlaxBigBirdForTokenClassification,
        )
        if is_flax_available()
        else ()
    )

    _snake_case     :  Optional[Any]     							=							False
    _snake_case     :  Tuple     							=							False


    def      lowerCAmelCase_  (					self	:  Optional[Any]					):
      _UpperCAmelCase						 =      FlaxBigBirdModelTester(self					)


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


    @slow
    # copied from `test_modeling_flax_common` because it takes much longer than other models
    def      lowerCAmelCase_  (					self	:  Union[str, 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      lowerCAmelCase_  (					self	:  List[Any]					):
      super().test_no_automatic_init()


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


    @slow
    def      lowerCAmelCase_  (					self	:  Union[str, Any]					):
      for model_class_name in self.all_model_classes:
        _UpperCAmelCase						 =      model_class_name.from_pretrained("""google/bigbird-roberta-base"""					)
        self.assertIsNotNone(__lowerCAmelCase					)


    def      lowerCAmelCase_  (					self	:  Union[str, 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      lowerCAmelCase_  (					self	:  Optional[int]					):
      _UpperCAmelCase ,				_UpperCAmelCase						 =      self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
        with self.subTest(model_class.__name__					):
          _UpperCAmelCase						 =      self._prepare_for_class(__lowerCAmelCase   ,				__lowerCAmelCase					)
          _UpperCAmelCase						 =      model_class(__lowerCAmelCase					)

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

          with self.subTest("""JIT Enabled"""					):
            _UpperCAmelCase						 =      model_jitted(**__lowerCAmelCase					).to_tuple()

          with self.subTest("""JIT Disabled"""					):
            with jax.disable_jit():
              _UpperCAmelCase						 =      model_jitted(**__lowerCAmelCase					).to_tuple()

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


    def      lowerCAmelCase_  (					self	:  List[Any]   ,				__lowerCAmelCase	:  Optional[Any]   ,				__lowerCAmelCase	:  Tuple   ,				__lowerCAmelCase	:  Tuple   ,				__lowerCAmelCase	:  Any=1e-5   ,				__lowerCAmelCase	:  List[str]="outputs"   ,				__lowerCAmelCase	:  Union[str, 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(__lowerCAmelCase   ,				__lowerCAmelCase   ,				__lowerCAmelCase   ,				__lowerCAmelCase   ,				__lowerCAmelCase   ,				__lowerCAmelCase					)




"""simple docstring"""


from argparse import ArgumentParser, Namespace
from typing import Any, List, Optional

from ..pipelines import Pipeline, get_supported_tasks, pipeline
from ..utils import logging
from . import BaseTransformersCLICommand


try:
    from fastapi import Body, FastAPI, HTTPException
    from fastapi.routing import APIRoute
    from pydantic import BaseModel
    from starlette.responses import JSONResponse
    from uvicorn import run

    UpperCAmelCase__  		=						True
except (ImportError, AttributeError):
    UpperCAmelCase__  		=						object






    def     __UpperCAmelCase					(			*lowercase   ,**lowercase		):







      """simple docstring"""
      pass

    UpperCAmelCase__  		=						False


UpperCAmelCase__  		=						logging.get_logger("""transformers-cli/serving""")






def     __UpperCAmelCase					(			lowercase		):







  """simple docstring"""
  _UpperCAmelCase						 =      pipeline(
      task=args.task   ,model=args.model if args.model else None   ,config=args.config   ,tokenizer=args.tokenizer   ,device=args.device   ,)
  return ServeCommand(lowercase   ,args.host   ,args.port   ,args.workers		)


class        a				(		lowerCAmelCase_   ):
    _snake_case     :  dict



class        a				(		lowerCAmelCase_   ):
    _snake_case     :  List[str]
    _snake_case     :  Optional[List[int]]



class        a				(		lowerCAmelCase_   ):
    _snake_case     :  str



class        a				(		lowerCAmelCase_   ):
    _snake_case     :  Any



class        a				(		lowerCAmelCase_   ):


    @staticmethod
    def      lowerCAmelCase_  (					__lowerCAmelCase	:  ArgumentParser					):
      _UpperCAmelCase						 =      parser.add_parser(
          """serve"""   ,				help="""CLI tool to run inference requests through REST and GraphQL endpoints."""					)
      serve_parser.add_argument(
          """--task"""   ,				type=__lowerCAmelCase   ,				choices=get_supported_tasks()   ,				help="""The task to run the pipeline on"""   ,				)
      serve_parser.add_argument("""--host"""   ,				type=__lowerCAmelCase   ,				default="""localhost"""   ,				help="""Interface the server will listen on."""					)
      serve_parser.add_argument("""--port"""   ,				type=__lowerCAmelCase   ,				default=8888   ,				help="""Port the serving will listen to."""					)
      serve_parser.add_argument("""--workers"""   ,				type=__lowerCAmelCase   ,				default=1   ,				help="""Number of http workers"""					)
      serve_parser.add_argument("""--model"""   ,				type=__lowerCAmelCase   ,				help="""Model's name or path to stored model."""					)
      serve_parser.add_argument("""--config"""   ,				type=__lowerCAmelCase   ,				help="""Model's config name or path to stored model."""					)
      serve_parser.add_argument("""--tokenizer"""   ,				type=__lowerCAmelCase   ,				help="""Tokenizer name to use."""					)
      serve_parser.add_argument(
          """--device"""   ,				type=__lowerCAmelCase   ,				default=-1   ,				help="""Indicate the device to run onto, -1 indicates CPU, >= 0 indicates GPU (default: -1)"""   ,				)
      serve_parser.set_defaults(func=__lowerCAmelCase					)


    def __init__(					self	:  Tuple   ,				__lowerCAmelCase	:  Pipeline   ,				__lowerCAmelCase	:  str   ,				__lowerCAmelCase	:  int   ,				__lowerCAmelCase	:  int					):
      _UpperCAmelCase						 =      pipeline

      _UpperCAmelCase						 =      host
      _UpperCAmelCase						 =      port
      _UpperCAmelCase						 =      workers

      if not _serve_dependencies_installed:
        raise RuntimeError(
            """Using serve command requires FastAPI and uvicorn. """
            """Please install transformers with [serving]: pip install \"transformers[serving]\"."""
            """Or install FastAPI and uvicorn separately."""					)
      else:
        logger.info(f'''Serving model over {host}:{port}'''					)
        _UpperCAmelCase						 =      FastAPI(
            routes=[
                APIRoute(
                    """/"""   ,				self.model_info   ,				response_model=__lowerCAmelCase   ,				response_class=__lowerCAmelCase   ,				methods=["""GET"""]   ,				),
                APIRoute(
                    """/tokenize"""   ,				self.tokenize   ,				response_model=__lowerCAmelCase   ,				response_class=__lowerCAmelCase   ,				methods=["""POST"""]   ,				),
                APIRoute(
                    """/detokenize"""   ,				self.detokenize   ,				response_model=__lowerCAmelCase   ,				response_class=__lowerCAmelCase   ,				methods=["""POST"""]   ,				),
                APIRoute(
                    """/forward"""   ,				self.forward   ,				response_model=__lowerCAmelCase   ,				response_class=__lowerCAmelCase   ,				methods=["""POST"""]   ,				),
            ]   ,				timeout=600   ,				)


    def      lowerCAmelCase_  (					self	:  Optional[int]					):
      run(self._app   ,				host=self.host   ,				port=self.port   ,				workers=self.workers					)


    def      lowerCAmelCase_  (					self	:  str					):
      return ServeModelInfoResult(infos=vars(self._pipeline.model.config					)					)


    def      lowerCAmelCase_  (					self	:  Tuple   ,				__lowerCAmelCase	:  str = Body(__lowerCAmelCase   ,				embed=__lowerCAmelCase					)   ,				__lowerCAmelCase	:  bool = Body(__lowerCAmelCase   ,				embed=__lowerCAmelCase					)					):
      try:
        _UpperCAmelCase						 =      self._pipeline.tokenizer.tokenize(__lowerCAmelCase					)

        if return_ids:
          _UpperCAmelCase						 =      self._pipeline.tokenizer.convert_tokens_to_ids(__lowerCAmelCase					)
          return ServeTokenizeResult(tokens=__lowerCAmelCase   ,				tokens_ids=__lowerCAmelCase					)
        else:
          return ServeTokenizeResult(tokens=__lowerCAmelCase					)

      except Exception as e:
        raise HTTPException(status_code=500   ,				detail={"""model""": """""", """error""": str(__lowerCAmelCase					)}					)


    def      lowerCAmelCase_  (					self	:  int   ,				__lowerCAmelCase	:  List[int] = Body(__lowerCAmelCase   ,				embed=__lowerCAmelCase					)   ,				__lowerCAmelCase	:  bool = Body(__lowerCAmelCase   ,				embed=__lowerCAmelCase					)   ,				__lowerCAmelCase	:  bool = Body(__lowerCAmelCase   ,				embed=__lowerCAmelCase					)   ,				):
      try:
        _UpperCAmelCase						 =      self._pipeline.tokenizer.decode(__lowerCAmelCase   ,				__lowerCAmelCase   ,				__lowerCAmelCase					)
        return ServeDeTokenizeResult(model=""""""   ,				text=__lowerCAmelCase					)
      except Exception as e:
        raise HTTPException(status_code=500   ,				detail={"""model""": """""", """error""": str(__lowerCAmelCase					)}					)


    async def      lowerCAmelCase_  (					self	:  str   ,				__lowerCAmelCase	:  List[Any]=Body(__lowerCAmelCase   ,				embed=__lowerCAmelCase					)					):

      # Check we don't have empty string
      if len(__lowerCAmelCase					) == 0:
        return ServeForwardResult(output=[]   ,				attention=[]					)

      try:
        # Forward through the model
        _UpperCAmelCase						 =      self._pipeline(__lowerCAmelCase					)
        return ServeForwardResult(output=__lowerCAmelCase					)
      except Exception as e:
        raise HTTPException(500   ,				{"""error""": str(__lowerCAmelCase					)}					)
