



"""simple docstring"""






import itertools
import json
import os
import unittest

from transformers import AddedToken, RobertaTokenizer, RobertaTokenizerFast
from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES
from transformers.testing_utils import require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin



@require_tokenizers
class      A__  (    __SCREAMING_SNAKE_CASE    ,						unittest.TestCase				):




	'''simple docstring'''





	SCREAMING_SNAKE_CASE        =      RobertaTokenizer
	SCREAMING_SNAKE_CASE        =      RobertaTokenizerFast
	SCREAMING_SNAKE_CASE        =      True
	SCREAMING_SNAKE_CASE        =      {'cls_token': '<s>'}

	def 				_SCREAMING_SNAKE_CASE  (	self:						List[str])							->			Dict:

				"""simple docstring"""

				super().setUp()

				# Adapted from Sennrich et al. 2015 and https://github.com/rsennrich/subword-nmt
				__lowerCAmelCase :    Optional[Any]				=			[
				    "l",
				    "o",
				    "w",
				    "e",
				    "r",
				    "s",
				    "t",
				    "i",
				    "d",
				    "n",
				    "\u0120",
				    "\u0120l",
				    "\u0120n",
				    "\u0120lo",
				    "\u0120low",
				    "er",
				    "\u0120lowest",
				    "\u0120newer",
				    "\u0120wider",
				    "<unk>",
				]
				__lowerCAmelCase :    List[str]				=			dict(zip(_SCREAMING_SNAKE_CASE			,     range(len(_SCREAMING_SNAKE_CASE))))
				__lowerCAmelCase :    Any				=			["#version: 0.2", "\u0120 l", "\u0120l o", "\u0120lo w", "e r", ""]
				__lowerCAmelCase :    Tuple				=			{"unk_token": "<unk>"}

				__lowerCAmelCase :    int				=			os.path.join(self.tmpdirname			,     VOCAB_FILES_NAMES["vocab_file"])
				__lowerCAmelCase :    Union[str, Any]				=			os.path.join(self.tmpdirname			,     VOCAB_FILES_NAMES["merges_file"])
				with open(self.vocab_file			,     "w"			,     encoding="utf-8") as fp:
							fp.write(json.dumps(_SCREAMING_SNAKE_CASE) + "\n")
				with open(self.merges_file			,     "w"			,     encoding="utf-8") as fp:
							fp.write("\n".join(_SCREAMING_SNAKE_CASE))

	def 				_SCREAMING_SNAKE_CASE  (	self:						Any			,     **_SCREAMING_SNAKE_CASE:						Any)							->			Any:

				"""simple docstring"""

				kwargs.update(self.special_tokens_map)
				return self.tokenizer_class.from_pretrained(self.tmpdirname			,     **_SCREAMING_SNAKE_CASE)

	def 				_SCREAMING_SNAKE_CASE  (	self:						Optional[int]			,     **_SCREAMING_SNAKE_CASE:						Any)							->			Dict:

				"""simple docstring"""

				kwargs.update(self.special_tokens_map)
				return RobertaTokenizerFast.from_pretrained(self.tmpdirname			,     **_SCREAMING_SNAKE_CASE)

	def 				_SCREAMING_SNAKE_CASE  (	self:						List[Any]			,     _SCREAMING_SNAKE_CASE:						List[str])							->			Any:

				"""simple docstring"""

				__lowerCAmelCase :    Optional[Any]				=			"lower newer"
				__lowerCAmelCase :    Optional[int]				=			"lower newer"
				return input_text, output_text

	def 				_SCREAMING_SNAKE_CASE  (	self:						Any)							->			Tuple:

				"""simple docstring"""

				__lowerCAmelCase :    Dict				=			self.tokenizer_class(self.vocab_file			,     self.merges_file			,     **self.special_tokens_map)
				__lowerCAmelCase :    Optional[int]				=			"lower newer"
				__lowerCAmelCase :    Optional[int]				=			["l", "o", "w", "er", "\u0120", "n", "e", "w", "er"]
				__lowerCAmelCase :    int				=			tokenizer.tokenize(_SCREAMING_SNAKE_CASE)  # , add_prefix_space=True)
				self.assertListEqual(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    Optional[Any]				=			tokens + [tokenizer.unk_token]
				__lowerCAmelCase :    Tuple				=			[0, 1, 2, 15, 10, 9, 3, 2, 15, 19]
				self.assertListEqual(tokenizer.convert_tokens_to_ids(_SCREAMING_SNAKE_CASE)			,     _SCREAMING_SNAKE_CASE)

	def 				_SCREAMING_SNAKE_CASE  (	self:						Tuple)							->			List[Any]:

				"""simple docstring"""

				__lowerCAmelCase :    List[Any]				=			self.get_tokenizer()

				self.assertListEqual(tokenizer.encode("Hello world!"			,     add_special_tokens=_SCREAMING_SNAKE_CASE)			,     [0, 3_1414, 232, 328, 2])
				self.assertListEqual(
				    tokenizer.encode("Hello world! cécé herlolip 418"			,     add_special_tokens=_SCREAMING_SNAKE_CASE)			,     [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69, 4_6078, 1588, 2]			,     )

	@slow
	def 				_SCREAMING_SNAKE_CASE  (	self:						Dict)							->			Dict:

				"""simple docstring"""

				__lowerCAmelCase :    List[Any]				=			self.tokenizer_class.from_pretrained("roberta-base")

				__lowerCAmelCase :    Union[str, Any]				=			tokenizer.encode("sequence builders"			,     add_special_tokens=_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    Optional[int]				=			tokenizer.encode("multi-sequence build"			,     add_special_tokens=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    Optional[Any]				=			tokenizer.encode(
				    "sequence builders"			,     add_special_tokens=_SCREAMING_SNAKE_CASE			,     add_prefix_space=_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    Tuple				=			tokenizer.encode(
				    "sequence builders"			,     "multi-sequence build"			,     add_special_tokens=_SCREAMING_SNAKE_CASE			,     add_prefix_space=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    Any				=			tokenizer.build_inputs_with_special_tokens(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    Tuple				=			tokenizer.build_inputs_with_special_tokens(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)

				assert encoded_sentence == encoded_text_from_decode
				assert encoded_pair == encoded_pair_from_decode

	def 				_SCREAMING_SNAKE_CASE  (	self:						List[Any])							->			int:

				"""simple docstring"""

				__lowerCAmelCase :    Any				=			self.get_tokenizer()

				__lowerCAmelCase :    Dict				=			"Encode this sequence."
				__lowerCAmelCase :    Dict				=			tokenizer.byte_encoder[" ".encode("utf-8")[0]]

				# Testing encoder arguments
				__lowerCAmelCase :    int				=			tokenizer.encode(_SCREAMING_SNAKE_CASE			,     add_special_tokens=_SCREAMING_SNAKE_CASE			,     add_prefix_space=_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    Dict				=			tokenizer.convert_ids_to_tokens(encoded[0])[0]
				self.assertNotEqual(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    Any				=			tokenizer.encode(_SCREAMING_SNAKE_CASE			,     add_special_tokens=_SCREAMING_SNAKE_CASE			,     add_prefix_space=_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    Union[str, Any]				=			tokenizer.convert_ids_to_tokens(encoded[0])[0]
				self.assertEqual(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)

				tokenizer.add_special_tokens({"bos_token": "<s>"})
				__lowerCAmelCase :    int				=			tokenizer.encode(_SCREAMING_SNAKE_CASE			,     add_special_tokens=_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    Optional[int]				=			tokenizer.convert_ids_to_tokens(encoded[1])[0]
				self.assertNotEqual(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)

				# Testing spaces after special tokens
				__lowerCAmelCase :    Optional[int]				=			"<mask>"
				tokenizer.add_special_tokens(
				    {"mask_token": AddedToken(_SCREAMING_SNAKE_CASE			,     lstrip=_SCREAMING_SNAKE_CASE			,     rstrip=_SCREAMING_SNAKE_CASE)})  # mask token has a left space
				__lowerCAmelCase :    Dict				=			tokenizer.convert_tokens_to_ids(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    Optional[int]				=			"Encode <mask> sequence"
				__lowerCAmelCase :    Any				=			"Encode <mask>sequence"

				__lowerCAmelCase :    Tuple				=			tokenizer.encode(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    int				=			encoded.index(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    int				=			tokenizer.convert_ids_to_tokens(encoded[mask_loc + 1])[0]
				self.assertEqual(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    Optional[int]				=			tokenizer.encode(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    Tuple				=			encoded.index(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    int				=			tokenizer.convert_ids_to_tokens(encoded[mask_loc + 1])[0]
				self.assertNotEqual(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)

	def 				_SCREAMING_SNAKE_CASE  (	self:						Optional[Any])							->			Optional[Any]:

				"""simple docstring"""

				pass

	def 				_SCREAMING_SNAKE_CASE  (	self:						List[Any])							->			Union[str, Any]:

				"""simple docstring"""

				for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
							with self.subTest(F"""{tokenizer.__class__.__name__} ({pretrained_name})"""):
										__lowerCAmelCase :    List[Any]				=			self.rust_tokenizer_class.from_pretrained(_SCREAMING_SNAKE_CASE			,     **_SCREAMING_SNAKE_CASE)
										__lowerCAmelCase :    Tuple				=			self.tokenizer_class.from_pretrained(_SCREAMING_SNAKE_CASE			,     **_SCREAMING_SNAKE_CASE)
										__lowerCAmelCase :    Dict				=			"A, <mask> AllenNLP sentence."
										__lowerCAmelCase :    Tuple				=			tokenizer_r.encode_plus(_SCREAMING_SNAKE_CASE			,     add_special_tokens=_SCREAMING_SNAKE_CASE			,     return_token_type_ids=_SCREAMING_SNAKE_CASE)
										__lowerCAmelCase :    Optional[int]				=			tokenizer_p.encode_plus(_SCREAMING_SNAKE_CASE			,     add_special_tokens=_SCREAMING_SNAKE_CASE			,     return_token_type_ids=_SCREAMING_SNAKE_CASE)

										# token_type_ids should put 0 everywhere
										self.assertEqual(sum(tokens_r["token_type_ids"])			,     sum(tokens_p["token_type_ids"]))

										# attention_mask should put 1 everywhere, so sum over length should be 1
										self.assertEqual(
										    sum(tokens_r["attention_mask"]) / len(tokens_r["attention_mask"])			,     sum(tokens_p["attention_mask"]) / len(tokens_p["attention_mask"])			,     )

										__lowerCAmelCase :    str				=			tokenizer_r.convert_ids_to_tokens(tokens_r["input_ids"])
										__lowerCAmelCase :    Optional[int]				=			tokenizer_p.convert_ids_to_tokens(tokens_p["input_ids"])

										# Rust correctly handles the space before the mask while python doesnt
										self.assertSequenceEqual(tokens_p["input_ids"]			,     [0, 250, 6, 5_0264, 3823, 487, 2_1992, 3645, 4, 2])
										self.assertSequenceEqual(tokens_r["input_ids"]			,     [0, 250, 6, 5_0264, 3823, 487, 2_1992, 3645, 4, 2])

										self.assertSequenceEqual(
										    _SCREAMING_SNAKE_CASE			,     ["<s>", "A", ",", "<mask>", "ĠAllen", "N", "LP", "Ġsentence", ".", "</s>"])
										self.assertSequenceEqual(
										    _SCREAMING_SNAKE_CASE			,     ["<s>", "A", ",", "<mask>", "ĠAllen", "N", "LP", "Ġsentence", ".", "</s>"])

	def 				_SCREAMING_SNAKE_CASE  (	self:						Optional[Any])							->			Union[str, Any]:

				"""simple docstring"""

				for trim_offsets, add_prefix_space in itertools.product([True, False]			,     repeat=2):
							__lowerCAmelCase :    int				=			self.rust_tokenizer_class.from_pretrained(
							    self.tmpdirname			,     use_fast=_SCREAMING_SNAKE_CASE			,     add_prefix_space=_SCREAMING_SNAKE_CASE			,     trim_offsets=_SCREAMING_SNAKE_CASE)

							__lowerCAmelCase :    List[str]				=			json.loads(tokenizer_r.backend_tokenizer.pre_tokenizer.__getstate__())
							__lowerCAmelCase :    Tuple				=			json.loads(tokenizer_r.backend_tokenizer.post_processor.__getstate__())

							self.assertEqual(pre_tokenizer_state["add_prefix_space"]			,     _SCREAMING_SNAKE_CASE)

							self.assertEqual(post_processor_state["add_prefix_space"]			,     _SCREAMING_SNAKE_CASE)
							self.assertEqual(post_processor_state["trim_offsets"]			,     _SCREAMING_SNAKE_CASE)







	def 				_SCREAMING_SNAKE_CASE  (	self:						List[Any])							->			Tuple:

				"""simple docstring"""

				for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
							with self.subTest(F"""{tokenizer.__class__.__name__} ({pretrained_name})"""):
										__lowerCAmelCase :    Union[str, Any]				=			"hello"  # `hello` is a token in the vocabulary of `pretrained_name`
										__lowerCAmelCase :    Union[str, Any]				=			F"""{text_of_1_token} {text_of_1_token}"""

										__lowerCAmelCase :    Any				=			self.rust_tokenizer_class.from_pretrained(
										    _SCREAMING_SNAKE_CASE			,     use_fast=_SCREAMING_SNAKE_CASE			,     add_prefix_space=_SCREAMING_SNAKE_CASE			,     trim_offsets=_SCREAMING_SNAKE_CASE)
										__lowerCAmelCase :    List[str]				=			tokenizer_r(_SCREAMING_SNAKE_CASE			,     return_offsets_mapping=_SCREAMING_SNAKE_CASE			,     add_special_tokens=_SCREAMING_SNAKE_CASE)
										self.assertEqual(encoding.offset_mapping[0]			,     (0, len(_SCREAMING_SNAKE_CASE)))
										self.assertEqual(
										    encoding.offset_mapping[1]			,     (len(_SCREAMING_SNAKE_CASE) + 1, len(_SCREAMING_SNAKE_CASE) + 1 + len(_SCREAMING_SNAKE_CASE))			,     )

										__lowerCAmelCase :    Dict				=			self.rust_tokenizer_class.from_pretrained(
										    _SCREAMING_SNAKE_CASE			,     use_fast=_SCREAMING_SNAKE_CASE			,     add_prefix_space=_SCREAMING_SNAKE_CASE			,     trim_offsets=_SCREAMING_SNAKE_CASE)
										__lowerCAmelCase :    Optional[int]				=			tokenizer_r(_SCREAMING_SNAKE_CASE			,     return_offsets_mapping=_SCREAMING_SNAKE_CASE			,     add_special_tokens=_SCREAMING_SNAKE_CASE)
										self.assertEqual(encoding.offset_mapping[0]			,     (0, len(_SCREAMING_SNAKE_CASE)))
										self.assertEqual(
										    encoding.offset_mapping[1]			,     (len(_SCREAMING_SNAKE_CASE) + 1, len(_SCREAMING_SNAKE_CASE) + 1 + len(_SCREAMING_SNAKE_CASE))			,     )

										__lowerCAmelCase :    Union[str, Any]				=			self.rust_tokenizer_class.from_pretrained(
										    _SCREAMING_SNAKE_CASE			,     use_fast=_SCREAMING_SNAKE_CASE			,     add_prefix_space=_SCREAMING_SNAKE_CASE			,     trim_offsets=_SCREAMING_SNAKE_CASE)
										__lowerCAmelCase :    Any				=			tokenizer_r(_SCREAMING_SNAKE_CASE			,     return_offsets_mapping=_SCREAMING_SNAKE_CASE			,     add_special_tokens=_SCREAMING_SNAKE_CASE)
										self.assertEqual(encoding.offset_mapping[0]			,     (0, len(_SCREAMING_SNAKE_CASE)))
										self.assertEqual(
										    encoding.offset_mapping[1]			,     (len(_SCREAMING_SNAKE_CASE), len(_SCREAMING_SNAKE_CASE) + 1 + len(_SCREAMING_SNAKE_CASE))			,     )

										__lowerCAmelCase :    Optional[int]				=			self.rust_tokenizer_class.from_pretrained(
										    _SCREAMING_SNAKE_CASE			,     use_fast=_SCREAMING_SNAKE_CASE			,     add_prefix_space=_SCREAMING_SNAKE_CASE			,     trim_offsets=_SCREAMING_SNAKE_CASE)
										__lowerCAmelCase :    List[Any]				=			tokenizer_r(_SCREAMING_SNAKE_CASE			,     return_offsets_mapping=_SCREAMING_SNAKE_CASE			,     add_special_tokens=_SCREAMING_SNAKE_CASE)
										self.assertEqual(encoding.offset_mapping[0]			,     (0, len(_SCREAMING_SNAKE_CASE)))
										self.assertEqual(
										    encoding.offset_mapping[1]			,     (len(_SCREAMING_SNAKE_CASE), len(_SCREAMING_SNAKE_CASE) + 1 + len(_SCREAMING_SNAKE_CASE))			,     )

										__lowerCAmelCase :    Dict				=			F""" {text}"""

										# tokenizer_r = self.rust_tokenizer_class.from_pretrained(
										#     pretrained_name, use_fast=True, add_prefix_space=True, trim_offsets=True
										# )
										# encoding = tokenizer_r(text, return_offsets_mapping=True, add_special_tokens=False)
										# self.assertEqual(encoding.offset_mapping[0], (1, 1 + len(text_of_1_token)))
										# self.assertEqual(
										#     encoding.offset_mapping[1],
										#     (1 + len(text_of_1_token) + 1, 1 + len(text_of_1_token) + 1 + len(text_of_1_token)),
										# )

										__lowerCAmelCase :    str				=			self.rust_tokenizer_class.from_pretrained(
										    _SCREAMING_SNAKE_CASE			,     use_fast=_SCREAMING_SNAKE_CASE			,     add_prefix_space=_SCREAMING_SNAKE_CASE			,     trim_offsets=_SCREAMING_SNAKE_CASE)
										__lowerCAmelCase :    Union[str, Any]				=			tokenizer_r(_SCREAMING_SNAKE_CASE			,     return_offsets_mapping=_SCREAMING_SNAKE_CASE			,     add_special_tokens=_SCREAMING_SNAKE_CASE)
										self.assertEqual(encoding.offset_mapping[0]			,     (1, 1 + len(_SCREAMING_SNAKE_CASE)))
										self.assertEqual(
										    encoding.offset_mapping[1]			,     (1 + len(_SCREAMING_SNAKE_CASE) + 1, 1 + len(_SCREAMING_SNAKE_CASE) + 1 + len(_SCREAMING_SNAKE_CASE))			,     )

										__lowerCAmelCase :    Dict				=			self.rust_tokenizer_class.from_pretrained(
										    _SCREAMING_SNAKE_CASE			,     use_fast=_SCREAMING_SNAKE_CASE			,     add_prefix_space=_SCREAMING_SNAKE_CASE			,     trim_offsets=_SCREAMING_SNAKE_CASE)
										__lowerCAmelCase :    Any				=			tokenizer_r(_SCREAMING_SNAKE_CASE			,     return_offsets_mapping=_SCREAMING_SNAKE_CASE			,     add_special_tokens=_SCREAMING_SNAKE_CASE)
										self.assertEqual(encoding.offset_mapping[0]			,     (0, 1 + len(_SCREAMING_SNAKE_CASE)))
										self.assertEqual(
										    encoding.offset_mapping[1]			,     (1 + len(_SCREAMING_SNAKE_CASE), 1 + len(_SCREAMING_SNAKE_CASE) + 1 + len(_SCREAMING_SNAKE_CASE))			,     )

										__lowerCAmelCase :    Tuple				=			self.rust_tokenizer_class.from_pretrained(
										    _SCREAMING_SNAKE_CASE			,     use_fast=_SCREAMING_SNAKE_CASE			,     add_prefix_space=_SCREAMING_SNAKE_CASE			,     trim_offsets=_SCREAMING_SNAKE_CASE)
										__lowerCAmelCase :    Optional[int]				=			tokenizer_r(_SCREAMING_SNAKE_CASE			,     return_offsets_mapping=_SCREAMING_SNAKE_CASE			,     add_special_tokens=_SCREAMING_SNAKE_CASE)
										self.assertEqual(encoding.offset_mapping[0]			,     (0, 1 + len(_SCREAMING_SNAKE_CASE)))
										self.assertEqual(
										    encoding.offset_mapping[1]			,     (1 + len(_SCREAMING_SNAKE_CASE), 1 + len(_SCREAMING_SNAKE_CASE) + 1 + len(_SCREAMING_SNAKE_CASE))			,     )




"""simple docstring"""






import argparse

import torch

from transformers import BertForMaskedLM


if __name__ == "__main__":
	__snake_case				:  Optional[Any] =       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')
	__snake_case				:  List[Any] =       parser.parse_args()

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

	__snake_case				:  List[Any] =       model.state_dict()
	__snake_case				:  Any =       {}

	for w in ["word_embeddings", "position_embeddings"]:
		__snake_case				:  List[str] =       state_dict[F"""{prefix}.embeddings.{w}.weight"""]
	for w in ["weight", "bias"]:
		__snake_case				:  Optional[int] =       state_dict[F"""{prefix}.embeddings.LayerNorm.{w}"""]

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

			__snake_case				:  Any =       state_dict[
			    F"""{prefix}.encoder.layer.{teacher_idx}.attention.output.dense.{w}"""
			]
			__snake_case				:  Union[str, Any] =       state_dict[
			    F"""{prefix}.encoder.layer.{teacher_idx}.attention.output.LayerNorm.{w}"""
			]

			__snake_case				:  Union[str, Any] =       state_dict[
			    F"""{prefix}.encoder.layer.{teacher_idx}.intermediate.dense.{w}"""
			]
			__snake_case				:  Any =       state_dict[
			    F"""{prefix}.encoder.layer.{teacher_idx}.output.dense.{w}"""
			]
			__snake_case				:  Optional[int] =       state_dict[
			    F"""{prefix}.encoder.layer.{teacher_idx}.output.LayerNorm.{w}"""
			]
		std_idx += 1

	__snake_case				:  Union[str, Any] =       state_dict['cls.predictions.decoder.weight']
	__snake_case				:  Union[str, Any] =       state_dict['cls.predictions.bias']
	if args.vocab_transform:
		for w in ["weight", "bias"]:
			__snake_case				:  Union[str, Any] =       state_dict[F"""cls.predictions.transform.dense.{w}"""]
			__snake_case				:  List[str] =       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)