




import os
from pickle import UnpicklingError
from typing import Dict, Tuple

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict, unflatten_dict

import transformers

from .utils import logging


_snake_case            = logging.get_logger(__name__)







def        lowerCAmelCase_			(  snake_case_,snake_case_,snake_case_,snake_case_=False  ):
      try:
            import torch  # noqa: F401
      except ImportError:
            logger.error(
                """Loading a PyTorch model in Flax, requires both PyTorch and Flax to be installed. Please see"""
                """ https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation"""
                """ instructions."""  )
            raise

      if not is_sharded:
            _A							:  List[Any]    =		os.path.abspath(snake_case_  )
            logger.info(f'''Loading PyTorch weights from {pt_path}'''  )

            _A							:  Union[str, Any]    =		torch.load(snake_case_,map_location="""cpu"""  )
            logger.info(f'''PyTorch checkpoint contains {sum(t.numel() for t in pt_state_dict.values()  ):,} parameters.'''  )

            _A							:  List[str]    =		convert_pytorch_state_dict_to_flax(snake_case_,snake_case_  )
      else:
            # model is sharded and pytorch_checkpoint_path already contains the list of .pt shard files
            _A							:  Optional[Any]    =		convert_pytorch_sharded_state_dict_to_flax(snake_case_,snake_case_  )
      return flax_state_dict







def        lowerCAmelCase_			(  snake_case_,snake_case_,snake_case_,snake_case_,):

      def is_key_or_prefix_key_in_dict(snake_case_  ) -> bool:
            return len(set(snake_case_  ) & {key, (model_prefix,) + key}  ) > 0

      # layer norm
      _A							:  int    =		pt_tuple_key[:-1] + ("""scale""",)
      if pt_tuple_key[-1] in ["weight", "gamma"] and is_key_or_prefix_key_in_dict(snake_case_  ):
            return renamed_pt_tuple_key, pt_tensor

      # batch norm layer mean
      _A							:  Optional[int]    =		pt_tuple_key[:-1] + ("""mean""",)
      if pt_tuple_key[-1] == "running_mean" and not is_key_or_prefix_key_in_dict(snake_case_  ):
            return renamed_pt_tuple_key, pt_tensor

      # batch norm layer var
      _A							:  Union[str, Any]    =		pt_tuple_key[:-1] + ("""var""",)
      if pt_tuple_key[-1] == "running_var" and not is_key_or_prefix_key_in_dict(snake_case_  ):
            return renamed_pt_tuple_key, pt_tensor

      # embedding
      _A							:  str    =		pt_tuple_key[:-1] + ("""embedding""",)
      if pt_tuple_key[-1] == "weight" and is_key_or_prefix_key_in_dict(snake_case_  ):
            return renamed_pt_tuple_key, pt_tensor

      # conv layer
      _A							:  str    =		pt_tuple_key[:-1] + ("""kernel""",)
      if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4 and not is_key_or_prefix_key_in_dict(snake_case_  ):
            _A							:  Optional[Any]    =		pt_tensor.transpose(2,3,1,0  )
            return renamed_pt_tuple_key, pt_tensor

      # linear layer
      _A							:  Optional[int]    =		pt_tuple_key[:-1] + ("""kernel""",)
      if pt_tuple_key[-1] == "weight" and not is_key_or_prefix_key_in_dict(snake_case_  ):
            _A							:  str    =		pt_tensor.T
            return renamed_pt_tuple_key, pt_tensor

      # old PyTorch layer norm weight
      _A							:  Tuple    =		pt_tuple_key[:-1] + ("""weight""",)
      if pt_tuple_key[-1] == "gamma":
            return renamed_pt_tuple_key, pt_tensor

      # old PyTorch layer norm bias
      _A							:  str    =		pt_tuple_key[:-1] + ("""bias""",)
      if pt_tuple_key[-1] == "beta":
            return renamed_pt_tuple_key, pt_tensor

      # New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
      _A							:  Tuple    =		None
      if pt_tuple_key[-3::2] == ("parametrizations", "original0"):
            _A							:  Optional[int]    =		pt_tuple_key[-2] + """_g"""
      elif pt_tuple_key[-3::2] == ("parametrizations", "original1"):
            _A							:  Optional[Any]    =		pt_tuple_key[-2] + """_v"""
      if name is not None:
            _A							:  List[Any]    =		pt_tuple_key[:-3] + (name,)
            return renamed_pt_tuple_key, pt_tensor

      return pt_tuple_key, pt_tensor







def        lowerCAmelCase_			(  snake_case_,snake_case_  ):
      # convert pytorch tensor to numpy
      _A							:  Any    =		{k: v.numpy() for k, v in pt_state_dict.items()}

      _A							:  Optional[Any]    =		flax_model.base_model_prefix

      # use params dict if the model contains batch norm layers
      if "params" in flax_model.params:
            _A							:  List[Any]    =		flax_model.params["""params"""]
      else:
            _A							:  Dict    =		flax_model.params
      _A							:  Optional[int]    =		flatten_dict(snake_case_  )

      # add batch_stats keys,values to dict
      if "batch_stats" in flax_model.params:
            _A							:  str    =		flatten_dict(flax_model.params["""batch_stats"""]  )
            random_flax_state_dict.update(snake_case_  )

      _A							:  Any    =		{}

      _A							:  List[str]    =		(model_prefix not in flax_model_params) and (
          model_prefix in {k.split("""."""  )[0] for k in pt_state_dict.keys()}
      )
      _A							:  Optional[Any]    =		(model_prefix in flax_model_params) and (
          model_prefix not in {k.split("""."""  )[0] for k in pt_state_dict.keys()}
      )

      # Need to change some parameters name to match Flax names
      for pt_key, pt_tensor in pt_state_dict.items():
            _A							:  Dict    =		tuple(pt_key.split("""."""  )  )

            # remove base model prefix if necessary
            _A							:  Union[str, Any]    =		pt_tuple_key[0] == model_prefix
            if load_model_with_head_into_base_model and has_base_model_prefix:
                  _A							:  List[Any]    =		pt_tuple_key[1:]

            # Correctly rename weight parameters
            _A    ,     _A							:  Dict    =		rename_key_and_reshape_tensor(
                snake_case_,snake_case_,snake_case_,snake_case_  )

            # add model prefix if necessary
            _A							:  str    =		(model_prefix,) + flax_key in random_flax_state_dict
            if load_base_model_into_model_with_head and require_base_model_prefix:
                  _A							:  Tuple    =		(model_prefix,) + flax_key

            if flax_key in random_flax_state_dict:
                  if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
                        raise ValueError(
                            f'''PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '''
                            f'''{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.'''  )

        # add batch stats if the model contains batchnorm layers
            if "batch_stats" in flax_model.params:
                  if "mean" in flax_key[-1] or "var" in flax_key[-1]:
                        _A							:  Tuple    =		jnp.asarray(snake_case_  )
                        continue
                  # remove num_batches_tracked key
                  if "num_batches_tracked" in flax_key[-1]:
                        flax_state_dict.pop(snake_case_,snake_case_  )
                        continue

                  # also add unexpected weight so that warning is thrown
                  _A							:  Any    =		jnp.asarray(snake_case_  )

            else:
                  # also add unexpected weight so that warning is thrown
                  _A							:  List[Any]    =		jnp.asarray(snake_case_  )

      return unflatten_dict(snake_case_  )







def        lowerCAmelCase_			(  snake_case_,snake_case_  ):
      import torch

      # Load the index
      _A							:  List[Any]    =		{}
      for shard_file in shard_filenames:
            # load using msgpack utils
            _A							:  Union[str, Any]    =		torch.load(snake_case_  )
            _A							:  List[str]    =		{k: v.numpy() for k, v in pt_state_dict.items()}

            _A							:  Optional[Any]    =		flax_model.base_model_prefix

            # use params dict if the model contains batch norm layers and then add batch_stats keys,values to dict
            if "batch_stats" in flax_model.params:
                  _A							:  int    =		flax_model.params["""params"""]

                  _A							:  Dict    =		flatten_dict(snake_case_  )
                  random_flax_state_dict.update(flatten_dict(flax_model.params["""batch_stats"""]  )  )
            else:
                  _A							:  int    =		flax_model.params
                  _A							:  Optional[Any]    =		flatten_dict(snake_case_  )

            _A							:  Union[str, Any]    =		(model_prefix not in flax_model_params) and (
                model_prefix in {k.split("""."""  )[0] for k in pt_state_dict.keys()}
            )
            _A							:  List[Any]    =		(model_prefix in flax_model_params) and (
                model_prefix not in {k.split("""."""  )[0] for k in pt_state_dict.keys()}
            )
            # Need to change some parameters name to match Flax names
            for pt_key, pt_tensor in pt_state_dict.items():
                  _A							:  Any    =		tuple(pt_key.split("""."""  )  )

                  # remove base model prefix if necessary
                  _A							:  List[Any]    =		pt_tuple_key[0] == model_prefix
                  if load_model_with_head_into_base_model and has_base_model_prefix:
                        _A							:  List[Any]    =		pt_tuple_key[1:]

                  # Correctly rename weight parameters
                  _A    ,     _A							:  Union[str, Any]    =		rename_key_and_reshape_tensor(
                      snake_case_,snake_case_,snake_case_,snake_case_  )
                  # add model prefix if necessary
                  _A							:  Dict    =		(model_prefix,) + flax_key in random_flax_state_dict
                  if load_base_model_into_model_with_head and require_base_model_prefix:
                        _A							:  Tuple    =		(model_prefix,) + flax_key

                  if flax_key in random_flax_state_dict:
                        if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
                              raise ValueError(
                                  f'''PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '''
                                  f'''{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.'''  )

            # add batch stats if the model contains batchnorm layers
                  if "batch_stats" in flax_model.params:
                        if "mean" in flax_key[-1]:
                              _A							:  Optional[Any]    =		jnp.asarray(snake_case_  )
                              continue
                        if "var" in flax_key[-1]:
                              _A							:  Any    =		jnp.asarray(snake_case_  )
                              continue
                        # remove num_batches_tracked key
                        if "num_batches_tracked" in flax_key[-1]:
                              flax_state_dict.pop(snake_case_,snake_case_  )
                              continue

                        # also add unexpected weight so that warning is thrown
                        _A							:  List[str]    =		jnp.asarray(snake_case_  )

                  else:
                        # also add unexpected weight so that warning is thrown
                        _A							:  Dict    =		jnp.asarray(snake_case_  )
      return unflatten_dict(snake_case_  )







def        lowerCAmelCase_			(  snake_case_,snake_case_  ):
      _A							:  Union[str, Any]    =		os.path.abspath(snake_case_  )
      logger.info(f'''Loading Flax weights from {flax_checkpoint_path}'''  )

      # import correct flax class
      _A							:  Optional[Any]    =		getattr(snake_case_,"""Flax""" + model.__class__.__name__  )

      # load flax weight dict
      with open(snake_case_,"""rb"""  ) as state_f:
            try:
                  _A							:  Dict    =		from_bytes(snake_case_,state_f.read()  )
            except UnpicklingError:
                  raise EnvironmentError(f'''Unable to convert {flax_checkpoint_path} to Flax deserializable object. '''  )

      return load_flax_weights_in_pytorch_model(snake_case_,snake_case_  )







def        lowerCAmelCase_			(  snake_case_,snake_case_  ):

      try:
            import torch  # noqa: F401
      except ImportError:
            logger.error(
                """Loading a Flax weights in PyTorch, requires both PyTorch and Flax to be installed. Please see"""
                """ https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation"""
                """ instructions."""  )
            raise

      # check if we have bf16 weights
      _A							:  int    =		flatten_dict(jax.tree_util.tree_map(lambda snake_case_  : x.dtype == jnp.bfloataa,snake_case_  )  ).values()
      if any(snake_case_  ):
            # convert all weights to fp32 if the are bf16 since torch.from_numpy can-not handle bf16
            # and bf16 is not fully supported in PT yet.
            logger.warning(
                """Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` """
                """before loading those in PyTorch model."""  )
            _A							:  str    =		jax.tree_util.tree_map(
                lambda snake_case_  : params.astype(np.floataa  ) if params.dtype == jnp.bfloataa else params,snake_case_  )

      _A							:  Dict    =		flatten_dict(snake_case_  )
      _A							:  str    =		pt_model.state_dict()

      _A							:  List[str]    =		(pt_model.base_model_prefix in flax_state) and (
          pt_model.base_model_prefix not in {k.split("""."""  )[0] for k in pt_model_dict.keys()}
      )
      _A							:  List[str]    =		(pt_model.base_model_prefix not in flax_state) and (
          pt_model.base_model_prefix in {k.split("""."""  )[0] for k in pt_model_dict.keys()}
      )

      # keep track of unexpected & missing keys
      _A							:  List[Any]    =		[]
      _A							:  Dict    =		set(pt_model_dict.keys()  )

      for flax_key_tuple, flax_tensor in flax_state_dict.items():
            _A							:  Dict    =		flax_key_tuple[0] == pt_model.base_model_prefix
            _A							:  List[Any]    =		""".""".join((pt_model.base_model_prefix,) + flax_key_tuple  ) in pt_model_dict

            # adapt flax_key to prepare for loading from/to base model only
            if load_model_with_head_into_base_model and has_base_model_prefix:
                  _A							:  List[Any]    =		flax_key_tuple[1:]
            elif load_base_model_into_model_with_head and require_base_model_prefix:
                  _A							:  List[Any]    =		(pt_model.base_model_prefix,) + flax_key_tuple

            # rename flax weights to PyTorch format
            if flax_key_tuple[-1] == "kernel" and flax_tensor.ndim == 4 and ".".join(snake_case_  ) not in pt_model_dict:
                  # conv layer
                  _A							:  str    =		flax_key_tuple[:-1] + ("""weight""",)
                  _A							:  List[str]    =		jnp.transpose(snake_case_,(3, 2, 0, 1)  )
            elif flax_key_tuple[-1] == "kernel" and ".".join(snake_case_  ) not in pt_model_dict:
                  # linear layer
                  _A							:  List[str]    =		flax_key_tuple[:-1] + ("""weight""",)
                  _A							:  Union[str, Any]    =		flax_tensor.T
            elif flax_key_tuple[-1] in ["scale", "embedding"]:
                  _A							:  Optional[Any]    =		flax_key_tuple[:-1] + ("""weight""",)

            # adding batch stats from flax batch norm to pt
            elif "mean" in flax_key_tuple[-1]:
                  _A							:  Optional[Any]    =		flax_key_tuple[:-1] + ("""running_mean""",)
            elif "var" in flax_key_tuple[-1]:
                  _A							:  List[Any]    =		flax_key_tuple[:-1] + ("""running_var""",)

            if "batch_stats" in flax_state:
                  _A							:  Optional[int]    =		""".""".join(flax_key_tuple[1:]  )  # Remove the params/batch_stats header
            else:
                  _A							:  List[Any]    =		""".""".join(snake_case_  )

            # We also need to look at `pt_model_dict` and see if there are keys requiring further transformation.
            _A							:  Optional[Any]    =		{}
            # New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
            for key in pt_model_dict:
                  _A							:  Optional[Any]    =		key.split("""."""  )
                  _A							:  List[str]    =		None
                  if key_components[-3::2] == ["parametrizations", "original0"]:
                        _A							:  str    =		key_components[-2] + """_g"""
                  elif key_components[-3::2] == ["parametrizations", "original1"]:
                        _A							:  Any    =		key_components[-2] + """_v"""
                  if name is not None:
                        _A							:  int    =		key_components[:-3] + [name]
                        _A							:  List[str]    =		""".""".join(snake_case_  )
                        _A							:  int    =		key

            if flax_key in special_pt_names:
                  _A							:  List[str]    =		special_pt_names[flax_key]

            if flax_key in pt_model_dict:
                  if flax_tensor.shape != pt_model_dict[flax_key].shape:
                        raise ValueError(
                            f'''Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected '''
                            f'''to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}.'''  )
                  else:
                        # add weight to pytorch dict
                        _A							:  Union[str, Any]    =		np.asarray(snake_case_  ) if not isinstance(snake_case_,np.ndarray  ) else flax_tensor
                        _A							:  int    =		torch.from_numpy(snake_case_  )
                        # remove from missing keys
                        missing_keys.remove(snake_case_  )
            else:
                  # weight is not expected by PyTorch model
                  unexpected_keys.append(snake_case_  )

      pt_model.load_state_dict(snake_case_  )

      # re-transform missing_keys to list
      _A							:  Union[str, Any]    =		list(snake_case_  )

      if len(snake_case_  ) > 0:
            logger.warning(
                """Some weights of the Flax model were not used when initializing the PyTorch model"""
                f''' {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing'''
                f''' {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture'''
                """ (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This"""
                f''' IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect'''
                """ to be exactly identical (e.g. initializing a BertForSequenceClassification model from a"""
                """ FlaxBertForSequenceClassification model)."""  )
      else:
            logger.warning(f'''All Flax model weights were used when initializing {pt_model.__class__.__name__}.\n'''  )
      if len(snake_case_  ) > 0:
            logger.warning(
                f'''Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly'''
                f''' initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to'''
                """ use it for predictions and inference."""  )
      else:
            logger.warning(
                f'''All the weights of {pt_model.__class__.__name__} were initialized from the Flax model.\n'''
                """If your task is similar to the task the model of the checkpoint was trained on, """
                f'''you can already use {pt_model.__class__.__name__} for predictions without further training.'''  )

      return pt_model







import gc
import unittest

from transformers import MODEL_FOR_MASKED_LM_MAPPING, TF_MODEL_FOR_MASKED_LM_MAPPING, FillMaskPipeline, pipeline
from transformers.pipelines import PipelineException
from transformers.testing_utils import (
    is_pipeline_test,
    is_torch_available,
    nested_simplify,
    require_tf,
    require_torch,
    require_torch_gpu,
    slow,
)

from .test_pipelines_common import ANY







@is_pipeline_test
class 		lowercase							(     unittest.TestCase	):
     _a											=		MODEL_FOR_MASKED_LM_MAPPING
     _a											=		TF_MODEL_FOR_MASKED_LM_MAPPING

     def       a__			(						self       )		->							Tuple:
           super().tearDown()
           # clean-up as much as possible GPU memory occupied by PyTorch
           gc.collect()
           if is_torch_available():
                 import torch

                 torch.cuda.empty_cache()

     @require_tf
     def       a__			(						self       )		->							Any:
           _A							:  Optional[Any]    =		pipeline(task="""fill-mask"""	,  model="""sshleifer/tiny-distilroberta-base"""	,  top_k=2	,  framework="""tf"""       )
           _A							:  Optional[int]    =		unmasker("""My name is <mask>"""       )
           self.assertEqual(
               nested_simplify(_a	,  decimals=6       )	,  [
                   {"""sequence""": """My name is grouped""", """score""": 2.1e-05, """token""": 3_8015, """token_str""": """ grouped"""},
                   {"""sequence""": """My name is accuser""", """score""": 2.1e-05, """token""": 2_5506, """token_str""": """ accuser"""},
               ]	,  )

           _A							:  Tuple    =		unmasker("""The largest city in France is <mask>"""       )
           self.assertEqual(
               nested_simplify(_a	,  decimals=6       )	,  [
                   {
                       """sequence""": """The largest city in France is grouped""",
                       """score""": 2.1e-05,
                       """token""": 3_8015,
                       """token_str""": """ grouped""",
                   },
                   {
                       """sequence""": """The largest city in France is accuser""",
                       """score""": 2.1e-05,
                       """token""": 2_5506,
                       """token_str""": """ accuser""",
                   },
               ]	,  )

           _A							:  List[str]    =		unmasker("""My name is <mask>"""	,  targets=[""" Patrick""", """ Clara""", """ Teven"""]	,  top_k=3       )
           self.assertEqual(
               nested_simplify(_a	,  decimals=6       )	,  [
                   {"""sequence""": """My name is Clara""", """score""": 2e-05, """token""": 1_3606, """token_str""": """ Clara"""},
                   {"""sequence""": """My name is Patrick""", """score""": 2e-05, """token""": 3499, """token_str""": """ Patrick"""},
                   {"""sequence""": """My name is Te""", """score""": 1.9e-05, """token""": 2941, """token_str""": """ Te"""},
               ]	,  )

     @require_torch
     def       a__			(						self       )		->							str:
           _A							:  Any    =		pipeline(task="""fill-mask"""	,  model="""sshleifer/tiny-distilroberta-base"""	,  top_k=2	,  framework="""pt"""       )

           _A							:  List[Any]    =		unmasker("""My name is <mask>"""       )
           self.assertEqual(
               nested_simplify(_a	,  decimals=6       )	,  [
                   {"""sequence""": """My name is Maul""", """score""": 2.2e-05, """token""": 3_5676, """token_str""": """ Maul"""},
                   {"""sequence""": """My name isELS""", """score""": 2.2e-05, """token""": 1_6416, """token_str""": """ELS"""},
               ]	,  )

           _A							:  Optional[Any]    =		unmasker("""The largest city in France is <mask>"""       )
           self.assertEqual(
               nested_simplify(_a	,  decimals=6       )	,  [
                   {
                       """sequence""": """The largest city in France is Maul""",
                       """score""": 2.2e-05,
                       """token""": 3_5676,
                       """token_str""": """ Maul""",
                   },
                   {"""sequence""": """The largest city in France isELS""", """score""": 2.2e-05, """token""": 1_6416, """token_str""": """ELS"""},
               ]	,  )

           _A							:  Optional[int]    =		unmasker("""My name is <mask>"""	,  targets=[""" Patrick""", """ Clara""", """ Teven"""]	,  top_k=3       )
           self.assertEqual(
               nested_simplify(_a	,  decimals=6       )	,  [
                   {"""sequence""": """My name is Patrick""", """score""": 2.1e-05, """token""": 3499, """token_str""": """ Patrick"""},
                   {"""sequence""": """My name is Te""", """score""": 2e-05, """token""": 2941, """token_str""": """ Te"""},
                   {"""sequence""": """My name is Clara""", """score""": 2e-05, """token""": 1_3606, """token_str""": """ Clara"""},
               ]	,  )

           _A							:  str    =		unmasker("""My name is <mask> <mask>"""	,  top_k=2       )

           self.assertEqual(
               nested_simplify(_a	,  decimals=6       )	,  [
                   [
                       {
                           """score""": 2.2e-05,
                           """token""": 3_5676,
                           """token_str""": """ Maul""",
                           """sequence""": """<s>My name is Maul<mask></s>""",
                       },
                       {"""score""": 2.2e-05, """token""": 1_6416, """token_str""": """ELS""", """sequence""": """<s>My name isELS<mask></s>"""},
                   ],
                   [
                       {
                           """score""": 2.2e-05,
                           """token""": 3_5676,
                           """token_str""": """ Maul""",
                           """sequence""": """<s>My name is<mask> Maul</s>""",
                       },
                       {"""score""": 2.2e-05, """token""": 1_6416, """token_str""": """ELS""", """sequence""": """<s>My name is<mask>ELS</s>"""},
                   ],
               ]	,  )

     @require_torch_gpu
     def       a__			(						self       )		->							Union[str, Any]:
           _A							:  int    =		pipeline("""fill-mask"""	,  model="""hf-internal-testing/tiny-random-distilbert"""	,  device=0	,  framework="""pt"""       )

           # convert model to fp16
           pipe.model.half()

           _A							:  Optional[Any]    =		pipe("""Paris is the [MASK] of France."""       )
           # We actually don't care about the result, we just want to make sure
           # it works, meaning the float16 tensor got casted back to float32
           # for postprocessing.
           self.assertIsInstance(_a	,  _a       )

     @slow
     @require_torch
     def       a__			(						self       )		->							Optional[int]:
           _A							:  Optional[Any]    =		pipeline(task="""fill-mask"""	,  model="""distilroberta-base"""	,  top_k=2	,  framework="""pt"""       )
           self.run_large_test(_a       )

     @slow
     @require_tf
     def       a__			(						self       )		->							Tuple:
           _A							:  str    =		pipeline(task="""fill-mask"""	,  model="""distilroberta-base"""	,  top_k=2	,  framework="""tf"""       )
           self.run_large_test(_a       )

     def       a__			(						self	,  _a       )		->							Tuple:
           _A							:  Optional[int]    =		unmasker("""My name is <mask>"""       )
           self.assertEqual(
               nested_simplify(_a       )	,  [
                   {"""sequence""": """My name is John""", """score""": 0.008, """token""": 610, """token_str""": """ John"""},
                   {"""sequence""": """My name is Chris""", """score""": 0.007, """token""": 1573, """token_str""": """ Chris"""},
               ]	,  )
           _A							:  int    =		unmasker("""The largest city in France is <mask>"""       )
           self.assertEqual(
               nested_simplify(_a       )	,  [
                   {
                       """sequence""": """The largest city in France is Paris""",
                       """score""": 0.251,
                       """token""": 2201,
                       """token_str""": """ Paris""",
                   },
                   {
                       """sequence""": """The largest city in France is Lyon""",
                       """score""": 0.214,
                       """token""": 1_2790,
                       """token_str""": """ Lyon""",
                   },
               ]	,  )

           _A							:  Optional[Any]    =		unmasker("""My name is <mask>"""	,  targets=[""" Patrick""", """ Clara""", """ Teven"""]	,  top_k=3       )
           self.assertEqual(
               nested_simplify(_a       )	,  [
                   {"""sequence""": """My name is Patrick""", """score""": 0.005, """token""": 3499, """token_str""": """ Patrick"""},
                   {"""sequence""": """My name is Clara""", """score""": 0.000, """token""": 1_3606, """token_str""": """ Clara"""},
                   {"""sequence""": """My name is Te""", """score""": 0.000, """token""": 2941, """token_str""": """ Te"""},
               ]	,  )

     @require_torch
     def       a__			(						self       )		->							Tuple:
           _A							:  List[str]    =		pipeline(task="""fill-mask"""	,  model="""sshleifer/tiny-distilroberta-base"""	,  framework="""pt"""       )
           _A							:  str    =		None
           _A							:  Union[str, Any]    =		None
           self.run_pipeline_test(_a	,  []       )

     @require_tf
     def       a__			(						self       )		->							Union[str, Any]:
           _A							:  Tuple    =		pipeline(task="""fill-mask"""	,  model="""sshleifer/tiny-distilroberta-base"""	,  framework="""tf"""       )
           _A							:  Any    =		None
           _A							:  Dict    =		None
           self.run_pipeline_test(_a	,  []       )

     def       a__			(						self	,  _a	,  _a	,  _a       )		->							Any:
           if tokenizer is None or tokenizer.mask_token_id is None:
                 self.skipTest("""The provided tokenizer has no mask token, (probably reformer or wav2vec2)"""       )

           _A							:  Optional[Any]    =		FillMaskPipeline(model=_a	,  tokenizer=_a       )
           _A							:  Tuple    =		[
               F'''This is another {tokenizer.mask_token} test''',
           ]
           return fill_masker, examples

     def       a__			(						self	,  _a	,  _a       )		->							Dict:
           _A							:  Dict    =		fill_masker.tokenizer
           _A							:  List[str]    =		fill_masker.model

           _A							:  List[str]    =		fill_masker(
               F'''This is a {tokenizer.mask_token}'''	,  )
           self.assertEqual(
               _a	,  [
                   {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                   {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                   {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                   {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                   {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
               ]	,  )

           _A							:  Optional[Any]    =		fill_masker([F'''This is a {tokenizer.mask_token}''']       )
           self.assertEqual(
               _a	,  [
                   {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                   {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                   {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                   {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                   {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
               ]	,  )

           _A							:  List[str]    =		fill_masker([F'''This is a {tokenizer.mask_token}''', F'''Another {tokenizer.mask_token} great test.''']       )
           self.assertEqual(
               _a	,  [
                   [
                       {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                       {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                       {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                       {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                       {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                   ],
                   [
                       {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                       {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                       {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                       {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                       {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                   ],
               ]	,  )

           with self.assertRaises(_a       ):
                 fill_masker([None]       )
           # No mask_token is not supported
           with self.assertRaises(_a       ):
                 fill_masker("""This is"""       )

           self.run_test_top_k(_a	,  _a       )
           self.run_test_targets(_a	,  _a       )
           self.run_test_top_k_targets(_a	,  _a       )
           self.fill_mask_with_duplicate_targets_and_top_k(_a	,  _a       )
           self.fill_mask_with_multiple_masks(_a	,  _a       )

     def       a__			(						self	,  _a	,  _a       )		->							List[str]:
           _A							:  int    =		tokenizer.get_vocab()
           _A							:  str    =		sorted(vocab.keys()       )[:2]
           # Pipeline argument
           _A							:  Tuple    =		FillMaskPipeline(model=_a	,  tokenizer=_a	,  targets=_a       )
           _A							:  Optional[int]    =		fill_masker(F'''This is a {tokenizer.mask_token}'''       )
           self.assertEqual(
               _a	,  [
                   {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                   {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
               ]	,  )
           _A							:  List[str]    =		{vocab[el] for el in targets}
           self.assertEqual({el["""token"""] for el in outputs}	,  _a       )
           _A							:  Union[str, Any]    =		[tokenizer.decode([x]       ) for x in target_ids]
           self.assertEqual({el["""token_str"""] for el in outputs}	,  set(_a       )       )

           # Call argument
           _A							:  str    =		FillMaskPipeline(model=_a	,  tokenizer=_a       )
           _A							:  Optional[int]    =		fill_masker(F'''This is a {tokenizer.mask_token}'''	,  targets=_a       )
           self.assertEqual(
               _a	,  [
                   {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                   {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
               ]	,  )
           _A							:  int    =		{vocab[el] for el in targets}
           self.assertEqual({el["""token"""] for el in outputs}	,  _a       )
           _A							:  Any    =		[tokenizer.decode([x]       ) for x in target_ids]
           self.assertEqual({el["""token_str"""] for el in outputs}	,  set(_a       )       )

           # Score equivalence
           _A							:  int    =		fill_masker(F'''This is a {tokenizer.mask_token}'''	,  targets=_a       )
           _A							:  Optional[int]    =		[top_mask["""token_str"""] for top_mask in outputs]
           _A							:  Union[str, Any]    =		[top_mask["""score"""] for top_mask in outputs]

           # For some BPE tokenizers, `</w>` is removed during decoding, so `token_str` won't be the same as in `targets`.
           if set(_a       ) == set(_a       ):
                 _A							:  Tuple    =		fill_masker(F'''This is a {tokenizer.mask_token}'''	,  targets=_a       )
                 _A							:  Union[str, Any]    =		[top_mask["""score"""] for top_mask in unmasked_targets]
                 self.assertEqual(nested_simplify(_a       )	,  nested_simplify(_a       )       )

           # Raises with invalid
           with self.assertRaises(_a       ):
                 _A							:  str    =		fill_masker(F'''This is a {tokenizer.mask_token}'''	,  targets=[]       )
           # For some tokenizers, `""` is actually in the vocabulary and the expected error won't raised
           if "" not in tokenizer.get_vocab():
                 with self.assertRaises(_a       ):
                       _A							:  Optional[Any]    =		fill_masker(F'''This is a {tokenizer.mask_token}'''	,  targets=[""""""]       )
                 with self.assertRaises(_a       ):
                       _A							:  int    =		fill_masker(F'''This is a {tokenizer.mask_token}'''	,  targets=""""""       )

     def       a__			(						self	,  _a	,  _a       )		->							Optional[Any]:
           _A							:  str    =		FillMaskPipeline(model=_a	,  tokenizer=_a	,  top_k=2       )
           _A							:  str    =		fill_masker(F'''This is a {tokenizer.mask_token}'''       )
           self.assertEqual(
               _a	,  [
                   {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                   {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
               ]	,  )

           _A							:  Union[str, Any]    =		FillMaskPipeline(model=_a	,  tokenizer=_a       )
           _A							:  Union[str, Any]    =		fill_masker(F'''This is a {tokenizer.mask_token}'''	,  top_k=2       )
           self.assertEqual(
               _a	,  [
                   {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                   {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
               ]	,  )
           self.assertEqual(nested_simplify(_a       )	,  nested_simplify(_a       )       )

     def       a__			(						self	,  _a	,  _a       )		->							List[Any]:
           _A							:  Union[str, Any]    =		tokenizer.get_vocab()
           _A							:  int    =		FillMaskPipeline(model=_a	,  tokenizer=_a       )

           # top_k=2, ntargets=3
           _A							:  List[str]    =		sorted(vocab.keys()       )[:3]
           _A							:  Tuple    =		fill_masker(F'''This is a {tokenizer.mask_token}'''	,  top_k=2	,  targets=_a       )

           # If we use the most probably targets, and filter differently, we should still
           # have the same results
           _A							:  Any    =		[el["""token_str"""] for el in sorted(_a	,  key=lambda _a       : x["score"]	,  reverse=_a       )]
           # For some BPE tokenizers, `</w>` is removed during decoding, so `token_str` won't be the same as in `targets`.
           if set(_a       ).issubset(_a       ):
                 _A							:  Any    =		fill_masker(F'''This is a {tokenizer.mask_token}'''	,  top_k=3	,  targets=_a       )
                 # They should yield exactly the same result
                 self.assertEqual(nested_simplify(_a       )	,  nested_simplify(_a       )       )

     def       a__			(						self	,  _a	,  _a       )		->							str:
           _A							:  Optional[int]    =		FillMaskPipeline(model=_a	,  tokenizer=_a       )
           _A							:  List[Any]    =		tokenizer.get_vocab()
           # String duplicates + id duplicates
           _A							:  Optional[Any]    =		sorted(vocab.keys()       )[:3]
           _A							:  Optional[Any]    =		[targets[0], targets[1], targets[0], targets[2], targets[1]]
           _A							:  Union[str, Any]    =		fill_masker(F'''My name is {tokenizer.mask_token}'''	,  targets=_a	,  top_k=10       )

           # The target list contains duplicates, so we can't output more
           # than them
           self.assertEqual(len(_a       )	,  3       )



     def       a__			(						self	,  _a	,  _a       )		->							Tuple:
           _A							:  Any    =		FillMaskPipeline(model=_a	,  tokenizer=_a       )

           _A							:  Optional[Any]    =		fill_masker(
               F'''This is a {tokenizer.mask_token} {tokenizer.mask_token} {tokenizer.mask_token}'''	,  top_k=2       )
           self.assertEqual(
               _a	,  [
                   [
                       {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                       {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                   ],
                   [
                       {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                       {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                   ],
                   [
                       {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                       {"""sequence""": ANY(_a       ), """score""": ANY(_a       ), """token""": ANY(_a       ), """token_str""": ANY(_a       )},
                   ],
               ]	,  )

