


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    UpperCAmelCase__	( unittest.TestCase		):



    """simple docstring"""






    def _a (      self   )  ->	int:
         __UpperCamelCase   							=mock.Mock()
         __UpperCamelCase   							=500
         __UpperCamelCase   							={}
         __UpperCamelCase   							=HTTPError
         __UpperCamelCase   							={}

         # Download this model to make sure it's in the cache.
         __UpperCamelCase   							=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=snake_case__   ) as mock_head:
              __UpperCamelCase   							=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   )  ->	List[str]:
         __UpperCamelCase   							=mock.Mock()
         __UpperCamelCase   							=500
         __UpperCamelCase   							={}
         __UpperCamelCase   							=HTTPError
         __UpperCamelCase   							={}

         # Download this model to make sure it's in the cache.
         __UpperCamelCase   							=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=snake_case__   ) as mock_head:
              __UpperCamelCase   							=GPTaTokenizerFast.from_pretrained('gpt2'   )
              # This check we did call the fake head request
              mock_head.assert_called()






    def _a (      self   )  ->	str:





         try:
              __UpperCamelCase   							=tempfile.mktemp()
              with open(snake_case__ ,       'wb'   ) as f:
                   http_get('https://huggingface.co/albert-base-v1/resolve/main/spiece.model' ,       snake_case__   )

              __UpperCamelCase   							=AlbertTokenizer.from_pretrained(snake_case__   )
         finally:
              os.remove(snake_case__   )

         # 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' ,       snake_case__   )
              __UpperCamelCase   							=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 ,       1000   )
              # Tokenizer should depend on the remote checkpoint, not the local tokenizer.json file.

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






    def _a (      self   )  ->	Tuple:
         __UpperCamelCase   							=AlbertTokenizer.from_pretrained('https://huggingface.co/albert-base-v1/resolve/main/spiece.model'   )






@is_staging_test
class    UpperCAmelCase__	( unittest.TestCase		):



    """simple docstring"""






    UpperCAmelCase__						:		List[Any]              =  ["""[UNK]""", """[CLS]""", """[SEP]""", """[PAD]""", """[MASK]""", """bla""", """blou"""]






    @classmethod
    def _a (      cls   )  ->	Dict:
         __UpperCamelCase   							=TOKEN
         HfFolder.save_token(snake_case__   )






    @classmethod
    def _a (      cls   )  ->	Optional[int]:





         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   )  ->	Dict:





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

         tokenizer.push_to_hub('test-tokenizer' ,       use_auth_token=self._token   )
         __UpperCamelCase   							=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(snake_case__ ,       repo_id='test-tokenizer' ,       push_to_hub=snake_case__ ,       use_auth_token=self._token   )

         __UpperCamelCase   							=BertTokenizer.from_pretrained(f'{USER}/test-tokenizer'   )
         self.assertDictEqual(new_tokenizer.vocab ,       tokenizer.vocab   )






    def _a (      self   )  ->	List[Any]:





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

         tokenizer.push_to_hub('valid_org/test-tokenizer-org' ,       use_auth_token=self._token   )
         __UpperCamelCase   							=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(
                  snake_case__ ,       repo_id='valid_org/test-tokenizer-org' ,       push_to_hub=snake_case__ ,       use_auth_token=self._token   )

         __UpperCamelCase   							=BertTokenizer.from_pretrained('valid_org/test-tokenizer-org'   )
         self.assertDictEqual(new_tokenizer.vocab ,       tokenizer.vocab   )






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

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

         __UpperCamelCase   							=AutoTokenizer.from_pretrained(f'{USER}/test-dynamic-tokenizer' ,       trust_remote_code=snake_case__   )
         # 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   							=os.path.join(snake_case__ ,       'vocab.txt'   )
              with open(snake_case__ ,       'w' ,       encoding='utf-8'   ) as vocab_writer:
                   vocab_writer.write(''.join([x + '\n' for x in self.vocab_tokens]   )   )

              __UpperCamelCase   							=BertTokenizerFast.from_pretrained(snake_case__   )
              bert_tokenizer.save_pretrained(snake_case__   )
              __UpperCamelCase   							=CustomTokenizerFast.from_pretrained(snake_case__   )

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

         __UpperCamelCase   							=AutoTokenizer.from_pretrained(f'{USER}/test-dynamic-tokenizer' ,       trust_remote_code=snake_case__   )
         # 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   							=AutoTokenizer.from_pretrained(
             f'{USER}/test-dynamic-tokenizer' ,       use_fast=snake_case__ ,       trust_remote_code=snake_case__   )
         # 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    UpperCAmelCase__	( unittest.TestCase		):



    """simple docstring"""






    def _a (      self   )  ->	Optional[int]:
         __UpperCamelCase   							=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   )  ->	Optional[int]:
         __UpperCamelCase   							=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   )  ->	Dict:
         __UpperCamelCase   							=Trie()
         trie.add('A'   )
         self.assertEqual(trie.split('ABC'   ) ,       ['A', 'BC']   )
         self.assertEqual(trie.split('BCA'   ) ,       ['BC', 'A']   )






    def _a (      self   )  ->	Tuple:
         __UpperCamelCase   							=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   )  ->	Union[str, Any]:
         __UpperCamelCase   							=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   )  ->	str:
         __UpperCamelCase   							=Trie()
         trie.add('AB'   )
         trie.add('B'   )
         trie.add('C'   )
         self.assertEqual(trie.split('ABC'   ) ,       ['AB', 'C']   )






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






    def _a (      self   )  ->	Tuple:
         __UpperCamelCase   							=Trie()
         __UpperCamelCase   							=trie.cut_text('ABC' ,       [0, 0, 2, 1, 2, 3]   )
         self.assertEqual(snake_case__ ,       ['AB', 'C']   )









import argparse

import torch

from transformers import BertForMaskedLM


if __name__ == "__main__":
   _A			=						argparse.ArgumentParser(
       description=(
           'Extraction some layers of the full BertForMaskedLM or RObertaForMaskedLM for Transfer Learned'
           ' Distillation'
       )
   )
   parser.add_argument('--model_type', default='bert', choices=['bert'])
   parser.add_argument('--model_name', default='bert-base-uncased', type=str)
   parser.add_argument('--dump_checkpoint', default='serialization_dir/tf_bert-base-uncased_0247911.pth', type=str)
   parser.add_argument('--vocab_transform', action='store_true')
   _A			=						parser.parse_args()

   if args.model_type == "bert":
      _A			=						BertForMaskedLM.from_pretrained(args.model_name)
      _A			=						'bert'
   else:
      raise ValueError('args.model_type should be "bert".')

   _A			=						model.state_dict()
   _A			=						{}

   for w in ["word_embeddings", "position_embeddings"]:
      _A			=						state_dict[f"""{prefix}.embeddings.{w}.weight"""]
   for w in ["weight", "bias"]:
      _A			=						state_dict[f"""{prefix}.embeddings.LayerNorm.{w}"""]

   _A			=						0
   for teacher_idx in [0, 2, 4, 7, 9, 11]:
      for w in ["weight", "bias"]:
         _A			=						state_dict[
             f"""{prefix}.encoder.layer.{teacher_idx}.attention.self.query.{w}"""
         ]
         _A			=						state_dict[
             f"""{prefix}.encoder.layer.{teacher_idx}.attention.self.key.{w}"""
         ]
         _A			=						state_dict[
             f"""{prefix}.encoder.layer.{teacher_idx}.attention.self.value.{w}"""
         ]

         _A			=						state_dict[
             f"""{prefix}.encoder.layer.{teacher_idx}.attention.output.dense.{w}"""
         ]
         _A			=						state_dict[
             f"""{prefix}.encoder.layer.{teacher_idx}.attention.output.LayerNorm.{w}"""
         ]

         _A			=						state_dict[
             f"""{prefix}.encoder.layer.{teacher_idx}.intermediate.dense.{w}"""
         ]
         _A			=						state_dict[
             f"""{prefix}.encoder.layer.{teacher_idx}.output.dense.{w}"""
         ]
         _A			=						state_dict[
             f"""{prefix}.encoder.layer.{teacher_idx}.output.LayerNorm.{w}"""
         ]
      std_idx += 1

   _A			=						state_dict['cls.predictions.decoder.weight']
   _A			=						state_dict['cls.predictions.bias']
   if args.vocab_transform:
      for w in ["weight", "bias"]:
         _A			=						state_dict[f"""cls.predictions.transform.dense.{w}"""]
         _A			=						state_dict[f"""cls.predictions.transform.LayerNorm.{w}"""]

   print(f"""N layers selected for distillation: {std_idx}""")
   print(f"""Number of params transferred for distillation: {len(compressed_sd.keys())}""")

   print(f"""Save transferred checkpoint to {args.dump_checkpoint}.""")
   torch.save(compressed_sd, args.dump_checkpoint)





