import unittest
from pathlib import Path
from tempfile import NamedTemporaryFile, TemporaryDirectory

from transformers import BertConfig, BertTokenizerFast, FeatureExtractionPipeline
from transformers.convert_graph_to_onnx import (
    convert,
    ensure_valid_input,
    generate_identified_filename,
    infer_shapes,
    quantize,
)
from transformers.testing_utils import require_tf, require_tokenizers, require_torch, slow





class      lowerCAmelCase						:

     '''simple docstring'''





     def 			a__  (					self				,      lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__       ):
            return None





class      lowerCAmelCase						:

     '''simple docstring'''





     def 			a__  (					self				,      lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__       ):
            return None





class      lowerCAmelCase						(					unittest.TestCase			):

     '''simple docstring'''





     _SCREAMING_SNAKE_CASE     :     Optional[Any]		=[
         # (model_name, model_kwargs)
         ('bert-base-cased', {}),
         ('gpt2', {'use_cache': False}),  # We don't support exporting GPT2 past keys anymore
     ]



     @require_tf
     @slow
     def 			a__  (					self       ):
            for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
                   self._test_export(UpperCAmelCase__				,      'tf'				,      12				,      **UpperCAmelCase__       )



     @require_torch
     @slow
     def 			a__  (					self       ):
            for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
                   self._test_export(UpperCAmelCase__				,      'pt'				,      12				,      **UpperCAmelCase__       )



     @require_torch
     @slow
     def 			a__  (					self       ):
            from transformers import BertModel

            _A=	['''[UNK]''', '''[SEP]''', '''[CLS]''', '''[PAD]''', '''[MASK]''', '''some''', '''other''', '''words''']
            with NamedTemporaryFile(mode='w+t'       ) as vocab_file:
                   vocab_file.write('\n'.join(UpperCAmelCase__       )       )
                   vocab_file.flush()
                   _A=	BertTokenizerFast(vocab_file.name       )

            with TemporaryDirectory() as bert_save_dir:
                   _A=	BertModel(BertConfig(vocab_size=len(UpperCAmelCase__       )       )       )
                   model.save_pretrained(UpperCAmelCase__       )
                   self._test_export(UpperCAmelCase__				,      'pt'				,      12				,      UpperCAmelCase__       )



     @require_tf
     @slow
     def 			a__  (					self       ):
            for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
                   _A=	self._test_export(UpperCAmelCase__				,      'tf'				,      12				,      **UpperCAmelCase__       )
                   _A=	quantize(Path(UpperCAmelCase__       )       )

                   # Ensure the actual quantized model is not bigger than the original one
                   if quantized_path.stat().st_size >= Path(UpperCAmelCase__       ).stat().st_size:
                          self.fail('Quantized model is bigger than initial ONNX model'       )



     @require_torch
     @slow
     def 			a__  (					self       ):
            for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
                   _A=	self._test_export(UpperCAmelCase__				,      'pt'				,      12				,      **UpperCAmelCase__       )
                   _A=	quantize(UpperCAmelCase__       )

                   # Ensure the actual quantized model is not bigger than the original one
                   if quantized_path.stat().st_size >= Path(UpperCAmelCase__       ).stat().st_size:
                          self.fail('Quantized model is bigger than initial ONNX model'       )



     def 			a__  (					self				,      lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__=None				,      **lowerCAmelCase__       ):
            try:
                   # Compute path
                   with TemporaryDirectory() as tempdir:
                          _A=	Path(UpperCAmelCase__       ).joinpath('model.onnx'       )

                   # Remove folder if exists
                   if path.parent.exists():
                          path.parent.rmdir()

                   # Export
                   convert(UpperCAmelCase__				,      UpperCAmelCase__				,      UpperCAmelCase__				,      UpperCAmelCase__				,      UpperCAmelCase__				,      **UpperCAmelCase__       )

                   return path
            except Exception as e:
                   self.fail(UpperCAmelCase__       )



     @require_torch
     @require_tokenizers
     @slow
     def 			a__  (					self       ):
            from transformers import BertModel

            _A=	BertModel(BertConfig.from_pretrained('lysandre/tiny-bert-random'       )       )
            _A=	BertTokenizerFast.from_pretrained('lysandre/tiny-bert-random'       )
            self._test_infer_dynamic_axis(UpperCAmelCase__				,      UpperCAmelCase__				,      'pt'       )



     @require_tf
     @require_tokenizers
     @slow
     def 			a__  (					self       ):
            from transformers import TFBertModel

            _A=	TFBertModel(BertConfig.from_pretrained('lysandre/tiny-bert-random'       )       )
            _A=	BertTokenizerFast.from_pretrained('lysandre/tiny-bert-random'       )
            self._test_infer_dynamic_axis(UpperCAmelCase__				,      UpperCAmelCase__				,      'tf'       )



     def 			a__  (					self				,      lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__       ):
            _A=	FeatureExtractionPipeline(UpperCAmelCase__				,      UpperCAmelCase__       )

            _A=	['''input_ids''', '''token_type_ids''', '''attention_mask''', '''output_0''', '''output_1''']
            _A=	infer_shapes(UpperCAmelCase__				,      UpperCAmelCase__       )

            # Assert all variables are present
            self.assertEqual(len(UpperCAmelCase__       )				,      len(UpperCAmelCase__       )       )
            self.assertTrue(all(var_name in shapes for var_name in variable_names       )       )
            self.assertSequenceEqual(variable_names[:3]				,      UpperCAmelCase__       )
            self.assertSequenceEqual(variable_names[3:]				,      UpperCAmelCase__       )

            # Assert inputs are {0: batch, 1: sequence}
            for var_name in ["input_ids", "token_type_ids", "attention_mask"]:
                   self.assertDictEqual(shapes[var_name]				,      {0: 'batch', 1: 'sequence'}       )

            # Assert outputs are {0: batch, 1: sequence} and {0: batch}
            self.assertDictEqual(shapes['output_0']				,      {0: 'batch', 1: 'sequence'}       )
            self.assertDictEqual(shapes['output_1']				,      {0: 'batch'}       )



     def 			a__  (					self       ):
            _A=	['''input_ids''', '''attention_mask''', '''token_type_ids''']
            _A=	{'''input_ids''': [1, 2, 3, 4], '''attention_mask''': [0, 0, 0, 0], '''token_type_ids''': [1, 1, 1, 1]}
            _A=	ensure_valid_input(FuncContiguousArgs()				,      UpperCAmelCase__				,      UpperCAmelCase__       )

            # Should have exactly the same number of args (all are valid)
            self.assertEqual(len(UpperCAmelCase__       )				,      3       )

            # Should have exactly the same input names
            self.assertEqual(set(UpperCAmelCase__       )				,      set(UpperCAmelCase__       )       )

            # Parameter should be reordered according to their respective place in the function:
            # (input_ids, token_type_ids, attention_mask)
            self.assertEqual(UpperCAmelCase__				,      (tokens['input_ids'], tokens['token_type_ids'], tokens['attention_mask'])       )

            # Generated args are interleaved with another args (for instance parameter "past" in GPT2)
            _A=	ensure_valid_input(FuncNonContiguousArgs()				,      UpperCAmelCase__				,      UpperCAmelCase__       )

            # Should have exactly the one arg (all before the one not provided "some_other_args")
            self.assertEqual(len(UpperCAmelCase__       )				,      1       )
            self.assertEqual(len(UpperCAmelCase__       )				,      1       )

            # Should have only "input_ids"
            self.assertEqual(inputs_args[0]				,      tokens['input_ids']       )
            self.assertEqual(ordered_input_names[0]				,      'input_ids'       )





     def 			a__  (					self       ):
            _A=	generate_identified_filename(Path('/home/something/my_fake_model.onnx'       )				,      '-test'       )
            self.assertEqual('/home/something/my_fake_model-test.onnx'				,      generated.as_posix()       )
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				(    *lowerCAmelCase_       ,					**lowerCAmelCase_   )		->						List[Any]:




             '''simple docstring'''






             pass





      UpperCAmelCase_								=							False


UpperCAmelCase_								=							logging.get_logger('''transformers-cli/serving''')

def 							UpperCamelCase				(    lowerCAmelCase_   )		->						int:




       '''simple docstring'''






       _A=	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			):
     _SCREAMING_SNAKE_CASE     :     dict





class      lowerCAmelCase						(					_a			):
     _SCREAMING_SNAKE_CASE     :     List[str]
     _SCREAMING_SNAKE_CASE     :     Optional[List[int]]





class      lowerCAmelCase						(					_a			):
     _SCREAMING_SNAKE_CASE     :     str





class      lowerCAmelCase						(					_a			):
     _SCREAMING_SNAKE_CASE     :     Any





class      lowerCAmelCase						(					_a			):



     @staticmethod
     def 			a__  (					lowerCAmelCase__       ):
            _A=	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				,      lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__       ):
            _A=	pipeline

            _A=	host
            _A=	port
            _A=	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}"       )
                   _A=	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 			a__  (					self       ):
            run(self._app				,      host=self.host				,      port=self.port				,      workers=self.workers       )



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



     def 			a__  (					self				,      lowerCAmelCase__ = Body(lowerCAmelCase__				,      embed=lowerCAmelCase__       )				,      lowerCAmelCase__ = Body(lowerCAmelCase__				,      embed=lowerCAmelCase__       )       ):
            try:
                   _A=	self._pipeline.tokenizer.tokenize(lowerCAmelCase__       )

                   if return_ids:
                          _A=	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 			a__  (					self				,      lowerCAmelCase__ = Body(lowerCAmelCase__				,      embed=lowerCAmelCase__       )				,      lowerCAmelCase__ = Body(lowerCAmelCase__				,      embed=lowerCAmelCase__       )				,      lowerCAmelCase__ = Body(lowerCAmelCase__				,      embed=lowerCAmelCase__       )				,      ):
            try:
                   _A=	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 			a__  (					self				,      lowerCAmelCase__=Body(lowerCAmelCase__				,      embed=lowerCAmelCase__       )       ):

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

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