


'''simple docstring'''

import contextlib
import importlib
import io
import unittest

import transformers

# Try to import everything from transformers to ensure every object can be loaded.
from transformers import *  # noqa F406
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, require_tf, require_torch
from transformers.utils import ContextManagers, find_labels, is_flax_available, is_tf_available, is_torch_available


if is_torch_available():
   from transformers import BertForPreTraining, BertForQuestionAnswering, BertForSequenceClassification

if is_tf_available():
   from transformers import TFBertForPreTraining, TFBertForQuestionAnswering, TFBertForSequenceClassification

if is_flax_available():
   from transformers import FlaxBertForPreTraining, FlaxBertForQuestionAnswering, FlaxBertForSequenceClassification


lowercase__				     =		DUMMY_UNKNOWN_IDENTIFIER
# An actual model hosted on huggingface.co

lowercase__				     =		"main"
# Default branch name
lowercase__				     =		"f2c752cfc5c0ab6f4bdec59acea69eefbee381c2"
# One particular commit (not the top of `main`)
lowercase__				     =		"aaaaaaa"
# This commit does not exist, so we should 404.

lowercase__				     =		"d9e9f15bc825e4b2c9249e9578f884bbcb5e3684"
# Sha-1 of config.json on the top of `main`, for checking purposes
lowercase__				     =		"4b243c475af8d0a7754e87d7d096c92e5199ec2fe168a2ee7998e3b8e9bcb1d3"






@contextlib.contextmanager
def __UpperCamelCase   (			)      ->     Any:







     '''simple docstring'''






     print("Welcome!"				)
     yield
     print("Bye!"				)






@contextlib.contextmanager
def __UpperCamelCase   (			)      ->     Dict:







     '''simple docstring'''






     print("Bonjour!"				)
     yield
     print("Au revoir!"				)





class 						__SCREAMING_SNAKE_CASE					(      unittest.TestCase ):






     def  a_					(			self		)       ->    Optional[int]:
          # If the spec is missing, importlib would not be able to import the module dynamically.
          assert transformers.__spec__ is not None
          assert importlib.util.find_spec("transformers"		) is not None







class 						__SCREAMING_SNAKE_CASE					(      unittest.TestCase ):






     @unittest.mock.patch("sys.stdout"     ,							new_callable=io.StringIO		)
     def  a_					(			self     ,							__UpperCamelCase		)       ->    List[Any]:
          with ContextManagers([]		):
               print("Transformers are awesome!"		)
          # The print statement adds a new line at the end of the output
          self.assertEqual(mock_stdout.getvalue()     ,							"Transformers are awesome!\n"		)






     @unittest.mock.patch("sys.stdout"     ,							new_callable=io.StringIO		)
     def  a_					(			self     ,							__UpperCamelCase		)       ->    Dict:
          with ContextManagers([context_en()]		):
               print("Transformers are awesome!"		)
          # The output should be wrapped with an English welcome and goodbye
          self.assertEqual(mock_stdout.getvalue()     ,							"Welcome!\nTransformers are awesome!\nBye!\n"		)






     @unittest.mock.patch("sys.stdout"     ,							new_callable=io.StringIO		)
     def  a_					(			self     ,							__UpperCamelCase		)       ->    Optional[Any]:
          with ContextManagers([context_fr(), context_en()]		):
               print("Transformers are awesome!"		)
          # The output should be wrapped with an English and French welcome and goodbye
          self.assertEqual(mock_stdout.getvalue()     ,							"Bonjour!\nWelcome!\nTransformers are awesome!\nBye!\nAu revoir!\n"		)






     @require_torch
     def  a_					(			self		)       ->    int:
          self.assertEqual(find_labels(__UpperCamelCase		)     ,							["labels"]		)
          self.assertEqual(find_labels(__UpperCamelCase		)     ,							["labels", "next_sentence_label"]		)
          self.assertEqual(find_labels(__UpperCamelCase		)     ,							["start_positions", "end_positions"]		)





          class 						__SCREAMING_SNAKE_CASE					(      lowerCamelCase__ ):
               pass







          self.assertEqual(find_labels(__UpperCamelCase		)     ,							["labels"]		)






     @require_tf
     def  a_					(			self		)       ->    Union[str, Any]:
          self.assertEqual(find_labels(__UpperCamelCase		)     ,							["labels"]		)
          self.assertEqual(find_labels(__UpperCamelCase		)     ,							["labels", "next_sentence_label"]		)
          self.assertEqual(find_labels(__UpperCamelCase		)     ,							["start_positions", "end_positions"]		)





          class 						__SCREAMING_SNAKE_CASE					(      lowerCamelCase__ ):
               pass







          self.assertEqual(find_labels(__UpperCamelCase		)     ,							["labels"]		)






     @require_flax
     def  a_					(			self		)       ->    Any:
          # Flax models don't have labels
          self.assertEqual(find_labels(__UpperCamelCase		)     ,							[]		)
          self.assertEqual(find_labels(__UpperCamelCase		)     ,							[]		)
          self.assertEqual(find_labels(__UpperCamelCase		)     ,							[]		)





          class 						__SCREAMING_SNAKE_CASE					(      lowerCamelCase__ ):
               pass







          self.assertEqual(find_labels(__UpperCamelCase		)     ,							[]		)





'''simple docstring'''

import re

from flax.core.frozen_dict import freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.experimental import PartitionSpec as P


# Sentinels
lowercase__				     =		object()

# For specifying empty leaf dict `{}`
lowercase__				     =		object()






def __UpperCamelCase   (			__lowerCamelCase		: Tuple      ,  __lowerCamelCase		: Tuple				)      ->     Union[str, Any]:







     '''simple docstring'''






     _a						=		tuple((re.compile(x + "$"				) for x in qs)				)
     for i in range(len(__lowerCamelCase				) - len(__lowerCamelCase				) + 1				):
          _a						=		[x.match(__lowerCamelCase				) for x, y in zip(__lowerCamelCase      ,  ks[i:]				)]
          if matches and all(__lowerCamelCase				):
               return True
     return False






def __UpperCamelCase   (			__lowerCamelCase		: int				)      ->     Union[str, Any]:







     '''simple docstring'''






     def replace(__lowerCamelCase		: Tuple      ,  __lowerCamelCase		: Any				):
          for rule, replacement in rules:
               if _match(__lowerCamelCase      ,  __lowerCamelCase				):
                    return replacement
          return val

     return replace






def __UpperCamelCase   (			)      ->     Tuple:







     '''simple docstring'''






     return [
         # embeddings
         (("transformer", "wpe", "embedding"), P("mp"      ,  __lowerCamelCase				)),
         (("transformer", "wte", "embedding"), P("mp"      ,  __lowerCamelCase				)),
         # atention
         (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(__lowerCamelCase      ,  "mp"				)),
         (("attention", "out_proj", "kernel"), P("mp"      ,  __lowerCamelCase				)),
         (("attention", "out_proj", "bias"), None),
         # mlp
         (("mlp", "c_fc", "kernel"), P(__lowerCamelCase      ,  "mp"				)),
         (("mlp", "c_fc", "bias"), P("mp"				)),
         (("mlp", "c_proj", "kernel"), P("mp"      ,  __lowerCamelCase				)),
         (("mlp", "c_proj", "bias"), None),
         # layer norms
         ((r"ln_\d+", "bias"), None),
         ((r"\d+", r"ln_\d+", "scale"), None),
         (("ln_f", "bias"), None),
         (("ln_f", "scale"), None),
     ]






def __UpperCamelCase   (			__lowerCamelCase		: Dict				)      ->     Tuple:







     '''simple docstring'''






     _a						=		_get_partition_rules()
     _a						=		_replacement_rules(__lowerCamelCase				)
     _a						=		{k: _unmatched for k in flatten_dict(__lowerCamelCase				)}
     _a						=		{k: replace(__lowerCamelCase      ,  __lowerCamelCase				) for k, v in initd.items()}
     assert _unmatched not in result.values(), "Incomplete partition spec."
     return freeze(unflatten_dict(__lowerCamelCase				)				)

