



# Lint as: python3
import sys
from collections.abc import Mapping
from typing import TYPE_CHECKING, Dict, Optional

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.logging import get_logger
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter


if TYPE_CHECKING:
    import jax
    import jaxlib

__UpperCAmelCase			  =			get_logger()

__UpperCAmelCase			  =			None



class 							lowerCAmelCase_						(    TensorFormatter[Mapping, "jax.Array", Mapping]				):
       def __init__(     self,     SCREAMING_SNAKE_CASE_=None,     SCREAMING_SNAKE_CASE_=None,     **SCREAMING_SNAKE_CASE_			)	->						List[Any]:
          super().__init__(features=SCREAMING_SNAKE_CASE_			)
          import jax
          from jaxlib.xla_client import Device

          if isinstance(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			):
             raise ValueError(
                 F"""Expected {device} to be a `str` not {type(SCREAMING_SNAKE_CASE_			)}, as `jaxlib.xla_extension.Device` """
                 'is not serializable neither with `pickle` nor with `dill`. Instead you can surround '
                 'the device with `str()` to get its string identifier that will be internally mapped '
                 'to the actual `jaxlib.xla_extension.Device`.'			)
          UpperCamelCase       :    Union[str, Any]		      =      device if isinstance(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			) else str(jax.devices()[0]			)
          # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
          # with `pickle` nor with `dill`, so we need to use a global variable instead
          global DEVICE_MAPPING
          if DEVICE_MAPPING is None:
             UpperCamelCase       :    int		      =      self._map_devices_to_str()
          if self.device not in list(DEVICE_MAPPING.keys()			):
             logger.warning(
                 F"""Device with string identifier {self.device} not listed among the available """
                 F"""devices: {list(DEVICE_MAPPING.keys()			)}, so falling back to the default """
                 F"""device: {str(jax.devices()[0]			)}."""			)
             UpperCamelCase       :    Dict		      =      str(jax.devices()[0]			)
          UpperCamelCase       :    Dict		      =      jnp_array_kwargs
       @staticmethod
       def    snake_case_			(     )	->						Dict[str, "jaxlib.xla_extension.Device"]:
          import jax

          return {str(SCREAMING_SNAKE_CASE_			): device for device in jax.devices()}
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_			)	->						Optional[Any]:
          import jax
          import jax.numpy as jnp

          if isinstance(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			) and column:
             if all(
                 isinstance(SCREAMING_SNAKE_CASE_,     jax.Array			) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column			):
                return jnp.stack(SCREAMING_SNAKE_CASE_,     axis=0			)
          return column
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_			)	->						int:
          import jax
          import jax.numpy as jnp

          if isinstance(SCREAMING_SNAKE_CASE_,     (str, bytes, type(SCREAMING_SNAKE_CASE_			))			):
             return value
          elif isinstance(SCREAMING_SNAKE_CASE_,     (np.character, np.ndarray)			) and np.issubdtype(value.dtype,     np.character			):
             return value.tolist()

          UpperCamelCase       :    Optional[Any]		      =      {}

          if isinstance(SCREAMING_SNAKE_CASE_,     (np.number, np.ndarray)			) and np.issubdtype(value.dtype,     np.integer			):
             # the default int precision depends on the jax config
             # see https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision
             if jax.config.jax_enable_xaa:
                UpperCamelCase       :    List[Any]		      =      {'dtype': jnp.intaa}
             else:
                UpperCamelCase       :    Optional[int]		      =      {'dtype': jnp.intaa}
          elif isinstance(SCREAMING_SNAKE_CASE_,     (np.number, np.ndarray)			) and np.issubdtype(value.dtype,     np.floating			):
             UpperCamelCase       :    Union[str, Any]		      =      {'dtype': jnp.floataa}
          elif config.PIL_AVAILABLE and "PIL" in sys.modules:
             import PIL.Image

             if isinstance(SCREAMING_SNAKE_CASE_,     PIL.Image.Image			):
                UpperCamelCase       :    int		      =      np.asarray(SCREAMING_SNAKE_CASE_			)

        # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
        # with `pickle` nor with `dill`, so we need to use a global variable instead
          global DEVICE_MAPPING
          if DEVICE_MAPPING is None:
             UpperCamelCase       :    Dict		      =      self._map_devices_to_str()

          with jax.default_device(DEVICE_MAPPING[self.device]			):
             # calling jnp.array on a np.ndarray does copy the data
             # see https://github.com/google/jax/issues/4486
             return jnp.array(SCREAMING_SNAKE_CASE_,     **{**default_dtype, **self.jnp_array_kwargs}			)
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_			)	->						Dict:
          import jax

          # support for torch, tf, jax etc.
          if config.TORCH_AVAILABLE and "torch" in sys.modules:
             import torch

             if isinstance(SCREAMING_SNAKE_CASE_,     torch.Tensor			):
                return self._tensorize(data_struct.detach().cpu().numpy()[()]			)
          if hasattr(SCREAMING_SNAKE_CASE_,     '__array__'			) and not isinstance(SCREAMING_SNAKE_CASE_,     jax.Array			):
             UpperCamelCase       :    Any		      =      data_struct.__array__()
          # support for nested types like struct of list of struct
          if isinstance(SCREAMING_SNAKE_CASE_,     np.ndarray			):
             if data_struct.dtype == object:  # jax arrays cannot be instantied from an array of objects
                return self._consolidate([self.recursive_tensorize(SCREAMING_SNAKE_CASE_			) for substruct in data_struct]			)
          elif isinstance(SCREAMING_SNAKE_CASE_,     (list, tuple)			):
             return self._consolidate([self.recursive_tensorize(SCREAMING_SNAKE_CASE_			) for substruct in data_struct]			)
          return self._tensorize(SCREAMING_SNAKE_CASE_			)
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_			)	->						Any:
          return map_nested(self._recursive_tensorize,     SCREAMING_SNAKE_CASE_,     map_list=SCREAMING_SNAKE_CASE_			)
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_			)	->						Mapping:
          UpperCamelCase       :    List[Any]		      =      self.numpy_arrow_extractor().extract_row(SCREAMING_SNAKE_CASE_			)
          UpperCamelCase       :    int		      =      self.python_features_decoder.decode_row(SCREAMING_SNAKE_CASE_			)
          return self.recursive_tensorize(SCREAMING_SNAKE_CASE_			)
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_			)	->						"jax.Array":
          UpperCamelCase       :    str		      =      self.numpy_arrow_extractor().extract_column(SCREAMING_SNAKE_CASE_			)
          UpperCamelCase       :    Union[str, Any]		      =      self.python_features_decoder.decode_column(SCREAMING_SNAKE_CASE_,     pa_table.column_names[0]			)
          UpperCamelCase       :    Tuple		      =      self.recursive_tensorize(SCREAMING_SNAKE_CASE_			)
          UpperCamelCase       :    Union[str, Any]		      =      self._consolidate(SCREAMING_SNAKE_CASE_			)
          return column






       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_			)	->						Mapping:
          UpperCamelCase       :    List[Any]		      =      self.numpy_arrow_extractor().extract_batch(SCREAMING_SNAKE_CASE_			)
          UpperCamelCase       :    Any		      =      self.python_features_decoder.decode_batch(SCREAMING_SNAKE_CASE_			)
          UpperCamelCase       :    Optional[Any]		      =      self.recursive_tensorize(SCREAMING_SNAKE_CASE_			)
          for column_name in batch:
             UpperCamelCase       :    List[Any]		      =      self._consolidate(batch[column_name]			)
          return batch









import warnings
from typing import Dict

import numpy as np

from ..utils import ExplicitEnum, add_end_docstrings, is_tf_available, is_torch_available
from .base import PIPELINE_INIT_ARGS, GenericTensor, Pipeline


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

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





def 						__A						(      __lowerCamelCase       ) ->			List[str]:
	return 1.0 / (1.0 + np.exp(-_outputs       ))





def 						__A						(      __lowerCamelCase       ) ->			List[str]:
	a									=				np.max(_outputs				,     axis=-1				,     keepdims=__lowerCamelCase       )
	a									=				np.exp(_outputs - maxes       )
	return shifted_exp / shifted_exp.sum(axis=-1				,     keepdims=__lowerCamelCase       )

class 						__lowerCAmelCase      (      __magic_name__					):
							UpperCamelCase__   	=    '''sigmoid'''
							UpperCamelCase__   	=    '''softmax'''
							UpperCamelCase__   	=    '''none'''







@add_end_docstrings(
    __magic_name__ ,			r'''
        return_all_scores (`bool`, *optional*, defaults to `False`):
            Whether to return all prediction scores or just the one of the predicted class.
        function_to_apply (`str`, *optional*, defaults to `"default"`):
            The function to apply to the model outputs in order to retrieve the scores. Accepts four different values:

            - `"default"`: if the model has a single label, will apply the sigmoid function on the output. If the model
              has several labels, will apply the softmax function on the output.
            - `"sigmoid"`: Applies the sigmoid function on the output.
            - `"softmax"`: Applies the softmax function on the output.
            - `"none"`: Does not apply any function on the output.
    ''' ,			)
class 						__lowerCAmelCase      (      __magic_name__					):
							UpperCamelCase__   	=    False
							UpperCamelCase__   	=    ClassificationFunction.NONE
							def __init__(		self		:List[str]						, **__magic_name__		:List[Any]							):







								'''simple docstring'''







								super().__init__(**__magic_name__							)

								self.check_model_type(
								    TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING
								    if self.framework == """tf"""
								    else MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING							)
							def  lowerCamelCase__				(		self		:Any						, __magic_name__		:int=None						, __magic_name__		:Any=None						, __magic_name__		:Union[str, Any]=""						, **__magic_name__		:Tuple							):







								'''simple docstring'''







								a									=				tokenizer_kwargs

								a									=				{}
								if hasattr(self.model.config						, """return_all_scores"""							) and return_all_scores is None:
									a									=				self.model.config.return_all_scores

								if isinstance(__magic_name__						, __magic_name__							) or top_k is None:
									a									=				top_k
									a									=				False
								elif return_all_scores is not None:
									warnings.warn(
									    """`return_all_scores` is now deprecated,  if want a similar functionality use `top_k=None` instead of"""
									    """ `return_all_scores=True` or `top_k=1` instead of `return_all_scores=False`."""						, __magic_name__						, )
									if return_all_scores:
										a									=				None
									else:
										a									=				1

								if isinstance(__magic_name__						, __magic_name__							):
									a									=				ClassificationFunction[function_to_apply.upper()]

								if function_to_apply is not None:
									a									=				function_to_apply
								return preprocess_params, {}, postprocess_params
							def __call__(		self		:Dict						, *__magic_name__		:Optional[int]						, **__magic_name__		:Optional[Any]							):







								'''simple docstring'''







								a									=				super().__call__(*__magic_name__						, **__magic_name__							)
								# TODO try and retrieve it in a nicer way from _sanitize_parameters.
								a									=				"""top_k""" not in kwargs
								if isinstance(args[0]						, __magic_name__							) and _legacy:
									# This pipeline is odd, and return a list when single item is run
									return [result]
								else:
									return result
							def  lowerCamelCase__				(		self		:Union[str, Any]						, __magic_name__		:Optional[Any]						, **__magic_name__		:Optional[Any]							):







								'''simple docstring'''







								a									=				self.framework
								if isinstance(__magic_name__						, __magic_name__							):
									return self.tokenizer(**__magic_name__						, return_tensors=__magic_name__						, **__magic_name__							)
								elif isinstance(__magic_name__						, __magic_name__							) and len(__magic_name__							) == 1 and isinstance(inputs[0]						, __magic_name__							) and len(inputs[0]							) == 2:
									# It used to be valid to use a list of list of list for text pairs, keeping this path for BC
									return self.tokenizer(
									    text=inputs[0][0]						, text_pair=inputs[0][1]						, return_tensors=__magic_name__						, **__magic_name__							)
								elif isinstance(__magic_name__						, __magic_name__							):
									# This is likely an invalid usage of the pipeline attempting to pass text pairs.
									raise ValueError(
									    """The pipeline received invalid inputs, if you are trying to send text pairs, you can try to send a"""
									    """ dictionary `{\"text\": \"My text\", \"text_pair\": \"My pair\"}` in order to send a text pair."""							)
								return self.tokenizer(__magic_name__						, return_tensors=__magic_name__						, **__magic_name__							)
							def  lowerCamelCase__				(		self		:List[str]						, __magic_name__		:Tuple							):







								'''simple docstring'''







								return self.model(**__magic_name__							)





							def  lowerCamelCase__				(		self		:Dict						, __magic_name__		:Union[str, Any]						, __magic_name__		:int=None						, __magic_name__		:Union[str, Any]=1						, __magic_name__		:Tuple=True							):







								'''simple docstring'''







								if function_to_apply is None:
									if self.model.config.problem_type == "multi_label_classification" or self.model.config.num_labels == 1:
										a									=				ClassificationFunction.SIGMOID
									elif self.model.config.problem_type == "single_label_classification" or self.model.config.num_labels > 1:
										a									=				ClassificationFunction.SOFTMAX
									elif hasattr(self.model.config						, """function_to_apply"""							) and function_to_apply is None:
										a									=				self.model.config.function_to_apply
									else:
										a									=				ClassificationFunction.NONE

								a									=				model_outputs["""logits"""][0]
								a									=				outputs.numpy()

								if function_to_apply == ClassificationFunction.SIGMOID:
									a									=				sigmoid(__magic_name__							)
								elif function_to_apply == ClassificationFunction.SOFTMAX:
									a									=				softmax(__magic_name__							)
								elif function_to_apply == ClassificationFunction.NONE:
									a									=				outputs
								else:
									raise ValueError(F'Unrecognized `function_to_apply` argument: {function_to_apply}'							)

								if top_k == 1 and _legacy:
									return {"label": self.model.config.idalabel[scores.argmax().item()], "score": scores.max().item()}

								a									=				[
								    {"""label""": self.model.config.idalabel[i], """score""": score.item()} for i, score in enumerate(__magic_name__							)
								]
								if not _legacy:
									dict_scores.sort(key=lambda __magic_name__							: x["score"]						, reverse=__magic_name__							)
									if top_k is not None:
										a									=				dict_scores[:top_k]
								return dict_scores

