


'''simple docstring'''






import tempfile
import unittest

import numpy as np
from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import BertConfig, is_flax_available
from transformers.testing_utils import TOKEN, USER, is_staging_test, require_flax


if is_flax_available():
  import os

  from flax.core.frozen_dict import unfreeze
  from flax.traverse_util import flatten_dict

  from transformers import FlaxBertModel

  __A					:  Optional[int]											=      "0.12"  # assumed parallelism: 8






@require_flax
@is_staging_test
class    __snake_case       (   unittest.TestCase):







    """simple docstring"""



    @classmethod
    def __lowercase   (   cls    :  Optional[int]		)   ->  Optional[int]:
      lowerCAmelCase_			:      List[Any]			=   TOKEN
      HfFolder.save_token(lowerCamelCase		)

    @classmethod
    def __lowercase   (   cls    :  Optional[int]		)   ->  Any:
      try:
        delete_repo(token=cls._token    ,						repo_id="""test-model-flax"""		)
      except HTTPError:
        pass

      try:
        delete_repo(token=cls._token    ,						repo_id="""valid_org/test-model-flax-org"""		)
      except HTTPError:
        pass

    def __lowercase   (   self    :  Union[str, Any]		)   ->  Dict:
      lowerCAmelCase_			:      Optional[Any]			=   BertConfig(
          vocab_size=99    ,						hidden_size=32    ,						num_hidden_layers=5    ,						num_attention_heads=4    ,						intermediate_size=37		)
      lowerCAmelCase_			:      Optional[Any]			=   FlaxBertModel(lowerCamelCase		)
      model.push_to_hub("""test-model-flax"""    ,						use_auth_token=self._token		)

      lowerCAmelCase_			:      int			=   FlaxBertModel.from_pretrained(F'{USER}/test-model-flax'		)

      lowerCAmelCase_			:      Optional[Any]			=   flatten_dict(unfreeze(model.params		)		)
      lowerCAmelCase_			:      Union[str, Any]			=   flatten_dict(unfreeze(new_model.params		)		)

      for key in base_params.keys():
        lowerCAmelCase_			:      Optional[Any]			=   (base_params[key] - new_params[key]).sum().item()
        self.assertLessEqual(lowerCamelCase    ,						1E-3    ,						msg=F'{key} not identical'		)

      # Reset repo
      delete_repo(token=self._token    ,						repo_id="""test-model-flax"""		)

      # Push to hub via save_pretrained
      with tempfile.TemporaryDirectory() as tmp_dir:
        model.save_pretrained(lowerCamelCase    ,						repo_id="""test-model-flax"""    ,						push_to_hub=lowerCamelCase    ,						use_auth_token=self._token		)

      lowerCAmelCase_			:      Optional[Any]			=   FlaxBertModel.from_pretrained(F'{USER}/test-model-flax'		)

      lowerCAmelCase_			:      List[Any]			=   flatten_dict(unfreeze(model.params		)		)
      lowerCAmelCase_			:      int			=   flatten_dict(unfreeze(new_model.params		)		)

      for key in base_params.keys():
        lowerCAmelCase_			:      Any			=   (base_params[key] - new_params[key]).sum().item()
        self.assertLessEqual(lowerCamelCase    ,						1E-3    ,						msg=F'{key} not identical'		)





    def __lowercase   (   self    :  Union[str, Any]		)   ->  Dict:
      lowerCAmelCase_			:      Optional[int]			=   BertConfig(
          vocab_size=99    ,						hidden_size=32    ,						num_hidden_layers=5    ,						num_attention_heads=4    ,						intermediate_size=37		)
      lowerCAmelCase_			:      List[str]			=   FlaxBertModel(lowerCamelCase		)
      model.push_to_hub("""valid_org/test-model-flax-org"""    ,						use_auth_token=self._token		)

      lowerCAmelCase_			:      Tuple			=   FlaxBertModel.from_pretrained("""valid_org/test-model-flax-org"""		)

      lowerCAmelCase_			:      Union[str, Any]			=   flatten_dict(unfreeze(model.params		)		)
      lowerCAmelCase_			:      List[str]			=   flatten_dict(unfreeze(new_model.params		)		)

      for key in base_params.keys():
        lowerCAmelCase_			:      Tuple			=   (base_params[key] - new_params[key]).sum().item()
        self.assertLessEqual(lowerCamelCase    ,						1E-3    ,						msg=F'{key} not identical'		)

      # Reset repo
      delete_repo(token=self._token    ,						repo_id="""valid_org/test-model-flax-org"""		)

      # Push to hub via save_pretrained
      with tempfile.TemporaryDirectory() as tmp_dir:
        model.save_pretrained(
            lowerCamelCase    ,						repo_id="""valid_org/test-model-flax-org"""    ,						push_to_hub=lowerCamelCase    ,						use_auth_token=self._token		)

      lowerCAmelCase_			:      List[Any]			=   FlaxBertModel.from_pretrained("""valid_org/test-model-flax-org"""		)

      lowerCAmelCase_			:      Tuple			=   flatten_dict(unfreeze(model.params		)		)
      lowerCAmelCase_			:      Any			=   flatten_dict(unfreeze(new_model.params		)		)

      for key in base_params.keys():
        lowerCAmelCase_			:      Optional[int]			=   (base_params[key] - new_params[key]).sum().item()
        self.assertLessEqual(lowerCamelCase    ,						1E-3    ,						msg=F'{key} not identical'		)







def 				UpperCamelCase_			(    A__					:   Optional[int]							,    A__					:   str					):
  '''simple docstring'''





  lowerCAmelCase_			:      List[str]			=   True
  lowerCAmelCase_			:      Optional[int]			=   flatten_dict(modela.params					)
  lowerCAmelCase_			:      str			=   flatten_dict(modela.params					)
  for key in flat_params_a.keys():
    if np.sum(np.abs(flat_params_a[key] - flat_params_a[key]					)					) > 1E-4:
      lowerCAmelCase_			:      Optional[Any]			=   False

  return models_are_equal






@require_flax
class    __snake_case       (   unittest.TestCase):







    """simple docstring"""



    def __lowercase   (   self    :  Optional[Any]		)   ->  Any:
      lowerCAmelCase_			:      Optional[Any]			=   BertConfig.from_pretrained("""hf-internal-testing/tiny-bert-flax-only"""		)
      lowerCAmelCase_			:      Union[str, Any]			=   FlaxBertModel(lowerCamelCase		)

      lowerCAmelCase_			:      int			=   """bert"""
      with tempfile.TemporaryDirectory() as tmp_dir:
        model.save_pretrained(os.path.join(lowerCamelCase    ,						lowerCamelCase		)		)

        with self.assertRaises(lowerCamelCase		):
          lowerCAmelCase_			:      int			=   FlaxBertModel.from_pretrained(lowerCamelCase		)

        lowerCAmelCase_			:      Tuple			=   FlaxBertModel.from_pretrained(lowerCamelCase    ,						subfolder=lowerCamelCase		)

      self.assertTrue(check_models_equal(lowerCamelCase    ,						lowerCamelCase		)		)

    def __lowercase   (   self    :  List[Any]		)   ->  Any:
      lowerCAmelCase_			:      int			=   BertConfig.from_pretrained("""hf-internal-testing/tiny-bert-flax-only"""		)
      lowerCAmelCase_			:      Tuple			=   FlaxBertModel(lowerCamelCase		)

      lowerCAmelCase_			:      Optional[Any]			=   """bert"""
      with tempfile.TemporaryDirectory() as tmp_dir:
        model.save_pretrained(os.path.join(lowerCamelCase    ,						lowerCamelCase		)    ,						max_shard_size="""10KB"""		)

        with self.assertRaises(lowerCamelCase		):
          lowerCAmelCase_			:      Union[str, Any]			=   FlaxBertModel.from_pretrained(lowerCamelCase		)

        lowerCAmelCase_			:      Optional[Any]			=   FlaxBertModel.from_pretrained(lowerCamelCase    ,						subfolder=lowerCamelCase		)

      self.assertTrue(check_models_equal(lowerCamelCase    ,						lowerCamelCase		)		)

    def __lowercase   (   self    :  Union[str, Any]		)   ->  List[Any]:
      lowerCAmelCase_			:      Union[str, Any]			=   """bert"""
      lowerCAmelCase_			:      Any			=   """hf-internal-testing/tiny-random-bert-subfolder"""

      with self.assertRaises(lowerCamelCase		):
        lowerCAmelCase_			:      Optional[Any]			=   FlaxBertModel.from_pretrained(lowerCamelCase		)

      lowerCAmelCase_			:      Optional[int]			=   FlaxBertModel.from_pretrained(lowerCamelCase    ,						subfolder=lowerCamelCase		)

      self.assertIsNotNone(lowerCamelCase		)





    def __lowercase   (   self    :  Dict		)   ->  Optional[int]:
      lowerCAmelCase_			:      Optional[int]			=   """bert"""
      lowerCAmelCase_			:      List[Any]			=   """hf-internal-testing/tiny-random-bert-sharded-subfolder"""
      with self.assertRaises(lowerCamelCase		):
        lowerCAmelCase_			:      str			=   FlaxBertModel.from_pretrained(lowerCamelCase		)

      lowerCAmelCase_			:      Union[str, Any]			=   FlaxBertModel.from_pretrained(lowerCamelCase    ,						subfolder=lowerCamelCase		)

      self.assertIsNotNone(lowerCamelCase		)







'''simple docstring'''






import re
import warnings
from contextlib import contextmanager

from ...processing_utils import ProcessorMixin






class    __snake_case       (   _SCREAMING_SNAKE_CASE):







    """simple docstring"""



    lowercase    =					['image_processor', 'tokenizer']
    lowercase    =					'AutoImageProcessor'
    lowercase    =					'AutoTokenizer'

    def __init__(   self    :  int    ,						lowerCamelCase    :  List[str]=None    ,						lowerCamelCase    :  Union[str, Any]=None    ,						**lowerCamelCase    :  str		)   ->  Tuple:
      lowerCAmelCase_			:      str			=   None
      if "feature_extractor" in kwargs:
        warnings.warn(
            """The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"""
            """ instead."""    ,						lowerCamelCase    ,						)
        lowerCAmelCase_			:      Tuple			=   kwargs.pop("""feature_extractor"""		)

      lowerCAmelCase_			:      Tuple			=   image_processor if image_processor is not None else feature_extractor
      if image_processor is None:
        raise ValueError("""You need to specify an `image_processor`."""		)
      if tokenizer is None:
        raise ValueError("""You need to specify a `tokenizer`."""		)

      super().__init__(lowerCamelCase    ,						lowerCamelCase		)
      lowerCAmelCase_			:      Optional[int]			=   self.image_processor
      lowerCAmelCase_			:      Any			=   False

    def __call__(   self    :  List[Any]    ,						*lowerCamelCase    :  str    ,						**lowerCamelCase    :  Tuple		)   ->  Union[str, Any]:
      # For backward compatibility
      if self._in_target_context_manager:
        return self.current_processor(*lowerCamelCase    ,						**lowerCamelCase		)

      lowerCAmelCase_			:      Any			=   kwargs.pop("""images"""    ,						lowerCamelCase		)
      lowerCAmelCase_			:      Dict			=   kwargs.pop("""text"""    ,						lowerCamelCase		)
      if len(lowerCamelCase		) > 0:
        lowerCAmelCase_			:      str			=   args[0]
        lowerCAmelCase_			:      Dict			=   args[1:]

      if images is None and text is None:
        raise ValueError("""You need to specify either an `images` or `text` input to process."""		)

      if images is not None:
        lowerCAmelCase_			:      Any			=   self.image_processor(lowerCamelCase    ,						*lowerCamelCase    ,						**lowerCamelCase		)
      if text is not None:
        lowerCAmelCase_			:      str			=   self.tokenizer(lowerCamelCase    ,						**lowerCamelCase		)

      if text is None:
        return inputs
      elif images is None:
        return encodings
      else:
        lowerCAmelCase_			:      Dict			=   encodings["""input_ids"""]
        return inputs

    def __lowercase   (   self    :  str    ,						*lowerCamelCase    :  Union[str, Any]    ,						**lowerCamelCase    :  Union[str, Any]		)   ->  List[str]:
      return self.tokenizer.batch_decode(*lowerCamelCase    ,						**lowerCamelCase		)

    def __lowercase   (   self    :  List[Any]    ,						*lowerCamelCase    :  Any    ,						**lowerCamelCase    :  Union[str, Any]		)   ->  List[str]:
      return self.tokenizer.decode(*lowerCamelCase    ,						**lowerCamelCase		)

    @contextmanager
    def __lowercase   (   self    :  str		)   ->  Union[str, Any]:
      warnings.warn(
          """`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your """
          """labels by using the argument `text` of the regular `__call__` method (either in the same call as """
          """your images inputs, or in a separate call."""		)
      lowerCAmelCase_			:      List[str]			=   True
      lowerCAmelCase_			:      Optional[Any]			=   self.tokenizer
      yield
      lowerCAmelCase_			:      List[Any]			=   self.image_processor
      lowerCAmelCase_			:      List[str]			=   False

    def __lowercase   (   self    :  Dict    ,						lowerCamelCase    :  List[str]    ,						lowerCamelCase    :  str=False    ,						lowerCamelCase    :  List[Any]=None		)   ->  Optional[int]:
      if added_vocab is None:
        lowerCAmelCase_			:      str			=   self.tokenizer.get_added_vocab()

      lowerCAmelCase_			:      Union[str, Any]			=   {}

      while tokens:
        lowerCAmelCase_			:      Dict			=   re.search(R"""<s_(.*?)>"""    ,						lowerCamelCase    ,						re.IGNORECASE		)
        if start_token is None:
          break
        lowerCAmelCase_			:      Tuple			=   start_token.group(1		)
        lowerCAmelCase_			:      Tuple			=   re.search(RF'</s_{key}>'    ,						lowerCamelCase    ,						re.IGNORECASE		)
        lowerCAmelCase_			:      Tuple			=   start_token.group()
        if end_token is None:
          lowerCAmelCase_			:      str			=   tokens.replace(lowerCamelCase    ,						""""""		)
        else:
          lowerCAmelCase_			:      List[str]			=   end_token.group()
          lowerCAmelCase_			:      Dict			=   re.escape(lowerCamelCase		)
          lowerCAmelCase_			:      int			=   re.escape(lowerCamelCase		)
          lowerCAmelCase_			:      Dict			=   re.search(F'{start_token_escaped}(.*?){end_token_escaped}'    ,						lowerCamelCase    ,						re.IGNORECASE		)
          if content is not None:
            lowerCAmelCase_			:      str			=   content.group(1		).strip()
            if r"<s_" in content and r"</s_" in content:  # non-leaf node
              lowerCAmelCase_			:      str			=   self.tokenajson(lowerCamelCase    ,						is_inner_value=lowerCamelCase    ,						added_vocab=lowerCamelCase		)
              if value:
                if len(lowerCamelCase		) == 1:
                  lowerCAmelCase_			:      List[Any]			=   value[0]
                lowerCAmelCase_			:      Optional[Any]			=   value
            else:  # leaf nodes
              lowerCAmelCase_			:      List[str]			=   []
              for leaf in content.split(R"""<sep/>"""		):
                lowerCAmelCase_			:      Union[str, Any]			=   leaf.strip()
                if leaf in added_vocab and leaf[0] == "<" and leaf[-2:] == "/>":
                  lowerCAmelCase_			:      Any			=   leaf[1:-2]  # for categorical special tokens
                output[key].append(lowerCamelCase		)
              if len(output[key]		) == 1:
                lowerCAmelCase_			:      Optional[Any]			=   output[key][0]

          lowerCAmelCase_			:      List[Any]			=   tokens[tokens.find(lowerCamelCase		) + len(lowerCamelCase		) :].strip()
          if tokens[:6] == r"<sep/>":  # non-leaf nodes
            return [output] + self.tokenajson(tokens[6:]    ,						is_inner_value=lowerCamelCase    ,						added_vocab=lowerCamelCase		)

      if len(lowerCamelCase		):
        return [output] if is_inner_value else output
      else:
        return [] if is_inner_value else {"text_sequence": tokens}

    @property
    def __lowercase   (   self    :  Dict		)   ->  int:
      warnings.warn(
          """`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."""    ,						lowerCamelCase    ,						)
      return self.image_processor_class





    @property
    def __lowercase   (   self    :  Dict		)   ->  Optional[Any]:
      warnings.warn(
          """`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."""    ,						lowerCamelCase    ,						)
      return self.image_processor



