





'''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

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





 def 	__lowerCamelCase			(      *__lowerCAmelCase       :							Tuple      ,    **__lowerCAmelCase       :							Dict		)		->  Optional[Any]:
    pass


 _SCREAMING_SNAKE_CASE							     =			False


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





def 	__lowerCamelCase			(      __lowerCAmelCase       :							Namespace		)		->  Optional[int]:
   snake_case 						=   pipeline(
       task=args.task      ,    model=args.model if args.model else None      ,    config=args.config      ,    tokenizer=args.tokenizer      ,    device=args.device      ,    )
   return ServeCommand(__lowerCAmelCase      ,    args.host      ,    args.port      ,    args.workers		)
class 			_lowerCAmelCase		(			A__	):






     """simple docstring"""






     snake_case_						      =   42




class 			_lowerCAmelCase		(			A__	):






     """simple docstring"""






     snake_case_						      =   42
     snake_case_						      =   42




class 			_lowerCAmelCase		(			A__	):






     """simple docstring"""






     snake_case_						      =   42




class 			_lowerCAmelCase		(			A__	):






     """simple docstring"""






     snake_case_						      =   42




class 			_lowerCAmelCase		(			A__	):






     """simple docstring"""






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



     def __init__(   self							:			str   ,       __snake_case							:			Pipeline   ,       __snake_case							:			str   ,       __snake_case							:			int   ,       __snake_case							:			int						)->	List[Any]:
        snake_case 						=   pipeline

        snake_case 						=   host
        snake_case 						=   port
        snake_case 						=   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}'''						)
           snake_case 						=   FastAPI(
               routes=[
                   APIRoute(
                       """/"""   ,       self.model_info   ,       response_model=__snake_case   ,       response_class=__snake_case   ,       methods=["""GET"""]   ,       ),
                   APIRoute(
                       """/tokenize"""   ,       self.tokenize   ,       response_model=__snake_case   ,       response_class=__snake_case   ,       methods=["""POST"""]   ,       ),
                   APIRoute(
                       """/detokenize"""   ,       self.detokenize   ,       response_model=__snake_case   ,       response_class=__snake_case   ,       methods=["""POST"""]   ,       ),
                   APIRoute(
                       """/forward"""   ,       self.forward   ,       response_model=__snake_case   ,       response_class=__snake_case   ,       methods=["""POST"""]   ,       ),
               ]   ,       timeout=6_00   ,       )



     def      lowerCAmelCase			(   self							:			List[Any]						)->	int:
        run(self._app   ,       host=self.host   ,       port=self.port   ,       workers=self.workers						)



     def      lowerCAmelCase			(   self							:			int						)->	Optional[int]:
        return ServeModelInfoResult(infos=vars(self._pipeline.model.config						)						)



     def      lowerCAmelCase			(   self							:			List[str]   ,       __snake_case							:			str = Body(__snake_case   ,       embed=__snake_case						)   ,       __snake_case							:			bool = Body(__snake_case   ,       embed=__snake_case						)						)->	Tuple:
        try:
           snake_case 						=   self._pipeline.tokenizer.tokenize(__snake_case						)

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

        except Exception as e:
           raise HTTPException(status_code=5_00   ,       detail={"""model""": """""", """error""": str(__snake_case						)}						)



     def      lowerCAmelCase			(   self							:			List[str]   ,       __snake_case							:			List[int] = Body(__snake_case   ,       embed=__snake_case						)   ,       __snake_case							:			bool = Body(__snake_case   ,       embed=__snake_case						)   ,       __snake_case							:			bool = Body(__snake_case   ,       embed=__snake_case						)   ,       )->	Union[str, Any]:
        try:
           snake_case 						=   self._pipeline.tokenizer.decode(__snake_case   ,       __snake_case   ,       __snake_case						)
           return ServeDeTokenizeResult(model=""""""   ,       text=__snake_case						)
        except Exception as e:
           raise HTTPException(status_code=5_00   ,       detail={"""model""": """""", """error""": str(__snake_case						)}						)



     async def      lowerCAmelCase			(   self							:			int   ,       __snake_case							:			Tuple=Body(__snake_case   ,       embed=__snake_case						)						)->	str:

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

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








'''simple docstring'''




import unittest

from transformers import is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, require_torch, slow


if is_flax_available():
 import optax
 from flax.training.common_utils import onehot

 from transformers import AutoTokenizer, FlaxMTaForConditionalGeneration
 from transformers.models.ta.modeling_flax_ta import shift_tokens_right
@require_torch
@require_sentencepiece
@require_tokenizers
@require_flax
class 			_lowerCAmelCase		(			unittest.TestCase	):






     """simple docstring"""






     @slow
     def      lowerCAmelCase			(   self							:			Union[str, Any]						)->	Optional[Any]:
        snake_case 						=   FlaxMTaForConditionalGeneration.from_pretrained("""google/mt5-small"""						)
        snake_case 						=   AutoTokenizer.from_pretrained("""google/mt5-small"""						)

        snake_case 						=   tokenizer("""Hello there"""   ,       return_tensors="""np"""						).input_ids
        snake_case 						=   tokenizer("""Hi I am"""   ,       return_tensors="""np"""						).input_ids

        snake_case 						=   shift_tokens_right(__snake_case   ,       model.config.pad_token_id   ,       model.config.decoder_start_token_id						)

        snake_case 						=   model(__snake_case   ,       decoder_input_ids=__snake_case						).logits
        snake_case 						=   optax.softmax_cross_entropy(__snake_case   ,       onehot(__snake_case   ,       logits.shape[-1]						)						).mean()

        snake_case 						=   -(labels.shape[-1] * loss.item())

        snake_case 						=   -84.91_27
        self.assertTrue(abs(mtf_score - EXPECTED_SCORE						) < 1e-4						)

