






import os
import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from huggingface_hub.file_download import http_get
from requests.exceptions import HTTPError

from transformers import (
    AlbertTokenizer,
    AutoTokenizer,
    BertTokenizer,
    BertTokenizerFast,
    GPTaTokenizerFast,
    is_tokenizers_available,
)
from transformers.testing_utils import TOKEN, USER, is_staging_test, require_tokenizers
from transformers.tokenization_utils import Trie


sys.path.append(str(Path(__file__).parent.parent / 'utils'))

from test_module.custom_tokenization import CustomTokenizer  # noqa E402


if is_tokenizers_available():
   from test_module.custom_tokenization_fast import CustomTokenizerFast




class 							_lowerCAmelCase(			unittest.TestCase   ):





       """simple docstring"""







       def      _a       (   self				):
         # A mock response for an HTTP head request to emulate server down
         UpperCamelCase_:   Any        =							mock.Mock()
         UpperCamelCase_:   Optional[int]        =							5_0_0
         UpperCamelCase_:   Optional[int]        =							{}
         UpperCamelCase_:   List[str]        =							HTTPError
         UpperCamelCase_:   str        =							{}

         # Download this model to make sure it's in the cache.
         UpperCamelCase_:   int        =							BertTokenizer.from_pretrained('hf-internal-testing/tiny-random-bert'				)

         # Under the mock environment we get a 500 error when trying to reach the tokenizer.
         with mock.patch('requests.Session.request'    ,     return_value=_lowerCamelCase				) as mock_head:
           UpperCamelCase_:   Dict        =							BertTokenizer.from_pretrained('hf-internal-testing/tiny-random-bert'				)
           # This check we did call the fake head request
           mock_head.assert_called()





       @require_tokenizers
       def      _a       (   self				):
         # A mock response for an HTTP head request to emulate server down
         UpperCamelCase_:   Tuple        =							mock.Mock()
         UpperCamelCase_:   Union[str, Any]        =							5_0_0
         UpperCamelCase_:   Optional[int]        =							{}
         UpperCamelCase_:   str        =							HTTPError
         UpperCamelCase_:   str        =							{}

         # Download this model to make sure it's in the cache.
         UpperCamelCase_:   Dict        =							GPTaTokenizerFast.from_pretrained('gpt2'				)

         # Under the mock environment we get a 500 error when trying to reach the tokenizer.
         with mock.patch('requests.Session.request'    ,     return_value=_lowerCamelCase				) as mock_head:
           UpperCamelCase_:   Optional[int]        =							GPTaTokenizerFast.from_pretrained('gpt2'				)
           # This check we did call the fake head request
           mock_head.assert_called()





       def      _a       (   self				):
         # This test is for deprecated behavior and can be removed in v5
         try:
           UpperCamelCase_:   Optional[int]        =							tempfile.mktemp()
           with open(_lowerCamelCase    ,     'wb'				) as f:
             http_get('https://huggingface.co/albert-base-v1/resolve/main/spiece.model'    ,     _lowerCamelCase				)

           UpperCamelCase_:   Tuple        =							AlbertTokenizer.from_pretrained(_lowerCamelCase				)
         finally:
           os.remove(_lowerCamelCase				)

         # Supporting this legacy load introduced a weird bug where the tokenizer would load local files if they are in
         # the current folder and have the right name.
         if os.path.isfile('tokenizer.json'				):
           # We skip the test if the user has a `tokenizer.json` in this folder to avoid deleting it.
           return
         try:
           with open('tokenizer.json'    ,     'wb'				) as f:
             http_get('https://huggingface.co/hf-internal-testing/tiny-random-bert/blob/main/tokenizer.json'    ,     _lowerCamelCase				)
           UpperCamelCase_:   Union[str, Any]        =							AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2'				)
           # The tiny random BERT has a vocab size of 1024, tiny gpt2 as a vocab size of 1000
           self.assertEqual(tokenizer.vocab_size    ,     1_0_0_0				)
           # Tokenizer should depend on the remote checkpoint, not the local tokenizer.json file.

         finally:
           os.remove('tokenizer.json'				)





       def      _a       (   self				):
         # This test is for deprecated behavior and can be removed in v5
         UpperCamelCase_:   str        =							AlbertTokenizer.from_pretrained('https://huggingface.co/albert-base-v1/resolve/main/spiece.model'				)







@is_staging_test
class 							_lowerCAmelCase(			unittest.TestCase   ):





       """simple docstring"""







       a							:						Dict										=['''[UNK]''', '''[CLS]''', '''[SEP]''', '''[PAD]''', '''[MASK]''', '''bla''', '''blou''']





       @classmethod
       def      _a       (   cls				):
         UpperCamelCase_:   Optional[int]        =							TOKEN
         HfFolder.save_token(_lowerCamelCase				)





       @classmethod
       def      _a       (   cls				):
         try:
           delete_repo(token=cls._token    ,     repo_id='test-tokenizer'				)
         except HTTPError:
           pass

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

         try:
           delete_repo(token=cls._token    ,     repo_id='test-dynamic-tokenizer'				)
         except HTTPError:
           pass





       def      _a       (   self				):
         with tempfile.TemporaryDirectory() as tmp_dir:
           UpperCamelCase_:   Union[str, Any]        =							os.path.join(_lowerCamelCase    ,     'vocab.txt'				)
           with open(_lowerCamelCase    ,     'w'    ,     encoding='utf-8'				) as vocab_writer:
             vocab_writer.write(''.join([x + '\n' for x in self.vocab_tokens]				)				)
           UpperCamelCase_:   Union[str, Any]        =							BertTokenizer(_lowerCamelCase				)

         tokenizer.push_to_hub('test-tokenizer'    ,     use_auth_token=self._token				)
         UpperCamelCase_:   Union[str, Any]        =							BertTokenizer.from_pretrained(f'''{USER}/test-tokenizer'''				)
         self.assertDictEqual(new_tokenizer.vocab    ,     tokenizer.vocab				)

         # Reset repo
         delete_repo(token=self._token    ,     repo_id='test-tokenizer'				)

         # Push to hub via save_pretrained
         with tempfile.TemporaryDirectory() as tmp_dir:
           tokenizer.save_pretrained(_lowerCamelCase    ,     repo_id='test-tokenizer'    ,     push_to_hub=_lowerCamelCase    ,     use_auth_token=self._token				)

         UpperCamelCase_:   Optional[Any]        =							BertTokenizer.from_pretrained(f'''{USER}/test-tokenizer'''				)
         self.assertDictEqual(new_tokenizer.vocab    ,     tokenizer.vocab				)





       def      _a       (   self				):
         with tempfile.TemporaryDirectory() as tmp_dir:
           UpperCamelCase_:   int        =							os.path.join(_lowerCamelCase    ,     'vocab.txt'				)
           with open(_lowerCamelCase    ,     'w'    ,     encoding='utf-8'				) as vocab_writer:
             vocab_writer.write(''.join([x + '\n' for x in self.vocab_tokens]				)				)
           UpperCamelCase_:   Dict        =							BertTokenizer(_lowerCamelCase				)

         tokenizer.push_to_hub('valid_org/test-tokenizer-org'    ,     use_auth_token=self._token				)
         UpperCamelCase_:   Optional[int]        =							BertTokenizer.from_pretrained('valid_org/test-tokenizer-org'				)
         self.assertDictEqual(new_tokenizer.vocab    ,     tokenizer.vocab				)

         # Reset repo
         delete_repo(token=self._token    ,     repo_id='valid_org/test-tokenizer-org'				)

         # Push to hub via save_pretrained
         with tempfile.TemporaryDirectory() as tmp_dir:
           tokenizer.save_pretrained(
               _lowerCamelCase    ,     repo_id='valid_org/test-tokenizer-org'    ,     push_to_hub=_lowerCamelCase    ,     use_auth_token=self._token				)

         UpperCamelCase_:   Optional[int]        =							BertTokenizer.from_pretrained('valid_org/test-tokenizer-org'				)
         self.assertDictEqual(new_tokenizer.vocab    ,     tokenizer.vocab				)





       @require_tokenizers
       def      _a       (   self				):
         CustomTokenizer.register_for_auto_class()
         with tempfile.TemporaryDirectory() as tmp_dir:
           UpperCamelCase_:   Tuple        =							os.path.join(_lowerCamelCase    ,     'vocab.txt'				)
           with open(_lowerCamelCase    ,     'w'    ,     encoding='utf-8'				) as vocab_writer:
             vocab_writer.write(''.join([x + '\n' for x in self.vocab_tokens]				)				)
           UpperCamelCase_:   Optional[int]        =							CustomTokenizer(_lowerCamelCase				)

         # No fast custom tokenizer
         tokenizer.push_to_hub('test-dynamic-tokenizer'    ,     use_auth_token=self._token				)

         UpperCamelCase_:   str        =							AutoTokenizer.from_pretrained(f'''{USER}/test-dynamic-tokenizer'''    ,     trust_remote_code=_lowerCamelCase				)
         # Can't make an isinstance check because the new_model.config is from the CustomTokenizer class of a dynamic module
         self.assertEqual(tokenizer.__class__.__name__    ,     'CustomTokenizer'				)

         # Fast and slow custom tokenizer
         CustomTokenizerFast.register_for_auto_class()
         with tempfile.TemporaryDirectory() as tmp_dir:
           UpperCamelCase_:   Optional[Any]        =							os.path.join(_lowerCamelCase    ,     'vocab.txt'				)
           with open(_lowerCamelCase    ,     'w'    ,     encoding='utf-8'				) as vocab_writer:
             vocab_writer.write(''.join([x + '\n' for x in self.vocab_tokens]				)				)

           UpperCamelCase_:   str        =							BertTokenizerFast.from_pretrained(_lowerCamelCase				)
           bert_tokenizer.save_pretrained(_lowerCamelCase				)
           UpperCamelCase_:   Any        =							CustomTokenizerFast.from_pretrained(_lowerCamelCase				)

         tokenizer.push_to_hub('test-dynamic-tokenizer'    ,     use_auth_token=self._token				)

         UpperCamelCase_:   Dict        =							AutoTokenizer.from_pretrained(f'''{USER}/test-dynamic-tokenizer'''    ,     trust_remote_code=_lowerCamelCase				)
         # Can't make an isinstance check because the new_model.config is from the FakeConfig class of a dynamic module
         self.assertEqual(tokenizer.__class__.__name__    ,     'CustomTokenizerFast'				)
         UpperCamelCase_:   str        =							AutoTokenizer.from_pretrained(
             f'''{USER}/test-dynamic-tokenizer'''    ,     use_fast=_lowerCamelCase    ,     trust_remote_code=_lowerCamelCase				)
         # Can't make an isinstance check because the new_model.config is from the FakeConfig class of a dynamic module
         self.assertEqual(tokenizer.__class__.__name__    ,     'CustomTokenizer'				)







class 							_lowerCAmelCase(			unittest.TestCase   ):





       """simple docstring"""







       def      _a       (   self				):
         UpperCamelCase_:   str        =							Trie()
         trie.add('Hello 友達'				)
         self.assertEqual(trie.data    ,     {'H': {'e': {'l': {'l': {'o': {' ': {'友': {'達': {'': 1}}}}}}}}}				)
         trie.add('Hello'				)
         trie.data
         self.assertEqual(trie.data    ,     {'H': {'e': {'l': {'l': {'o': {'': 1, ' ': {'友': {'達': {'': 1}}}}}}}}}				)





       def      _a       (   self				):
         UpperCamelCase_:   Optional[int]        =							Trie()
         self.assertEqual(trie.split('[CLS] This is a extra_id_100'				)    ,     ['[CLS] This is a extra_id_100']				)
         trie.add('[CLS]'				)
         trie.add('extra_id_1'				)
         trie.add('extra_id_100'				)
         self.assertEqual(trie.split('[CLS] This is a extra_id_100'				)    ,     ['[CLS]', ' This is a ', 'extra_id_100']				)





       def      _a       (   self				):
         UpperCamelCase_:   int        =							Trie()
         trie.add('A'				)
         self.assertEqual(trie.split('ABC'				)    ,     ['A', 'BC']				)
         self.assertEqual(trie.split('BCA'				)    ,     ['BC', 'A']				)





       def      _a       (   self				):
         UpperCamelCase_:   Union[str, Any]        =							Trie()
         trie.add('TOKEN]'				)
         trie.add('[SPECIAL_TOKEN]'				)
         self.assertEqual(trie.split('This is something [SPECIAL_TOKEN]'				)    ,     ['This is something ', '[SPECIAL_TOKEN]']				)





       def      _a       (   self				):
         UpperCamelCase_:   Tuple        =							Trie()
         trie.add('A'				)
         trie.add('P'				)
         trie.add('[SPECIAL_TOKEN]'				)
         self.assertEqual(trie.split('This is something [SPECIAL_TOKEN]'				)    ,     ['This is something ', '[SPECIAL_TOKEN]']				)





       def      _a       (   self				):
         UpperCamelCase_:   List[Any]        =							Trie()
         trie.add('AB'				)
         trie.add('B'				)
         trie.add('C'				)
         self.assertEqual(trie.split('ABC'				)    ,     ['AB', 'C']				)





       def      _a       (   self				):
         UpperCamelCase_:   Dict        =							Trie()
         trie.add('ABC'				)
         trie.add('B'				)
         trie.add('CD'				)
         self.assertEqual(trie.split('ABCD'				)    ,     ['ABC', 'D']				)





       def      _a       (   self				):
         # Even if the offsets are wrong, we necessarily output correct string
         # parts.
         UpperCamelCase_:   Tuple        =							Trie()
         UpperCamelCase_:   Optional[int]        =							trie.cut_text('ABC'    ,     [0, 0, 2, 1, 2, 3]				)
         self.assertEqual(_lowerCamelCase    ,     ['AB', 'C']				)





'''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
UpperCamelCase_		:     Optional[int]          =		object()

# For specifying empty leaf dict `{}`
UpperCamelCase_		:     Optional[int]          =		object()





def       __a      (  _UpperCamelCase:       Union[str, Any]   ,		_UpperCamelCase:       Optional[Any]       )		->					Optional[Any]:

     """simple docstring"""




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





def       __a      (  _UpperCamelCase:       Optional[Any]       )		->					Union[str, Any]:

     """simple docstring"""




     def replace(_UpperCamelCase:       Tuple   ,		_UpperCamelCase:       List[str]       ):
          for rule, replacement in rules:
               if _match(_UpperCamelCase   ,		_UpperCamelCase       ):
                    return replacement
          return val

     return replace





def       __a      (  )		->					Any:

     """simple docstring"""




     return [
         # embeddings
         (("transformer", "wpe", "embedding"), P("mp"   ,		_UpperCamelCase       )),
         (("transformer", "wte", "embedding"), P("mp"   ,		_UpperCamelCase       )),
         # atention
         (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(_UpperCamelCase   ,		"mp"       )),
         (("attention", "out_proj", "kernel"), P("mp"   ,		_UpperCamelCase       )),
         (("attention", "out_proj", "bias"), None),
         # mlp
         (("mlp", "c_fc", "kernel"), P(_UpperCamelCase   ,		"mp"       )),
         (("mlp", "c_fc", "bias"), P("mp"       )),
         (("mlp", "c_proj", "kernel"), P("mp"   ,		_UpperCamelCase       )),
         (("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       __a      (  _UpperCamelCase:       Union[str, Any]       )		->					Any:

     """simple docstring"""




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


