
"""simple docstring"""





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


_a								=						{
    """configuration_longformer""": [
        """LONGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """LongformerConfig""",
        """LongformerOnnxConfig""",
    ],
    """tokenization_longformer""": ["""LongformerTokenizer"""],
}

try:
				if not is_tokenizers_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				_a								=						["""LongformerTokenizerFast"""]

try:
				if not is_torch_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				_a								=						[
				    """LONGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST""",
				    """LongformerForMaskedLM""",
				    """LongformerForMultipleChoice""",
				    """LongformerForQuestionAnswering""",
				    """LongformerForSequenceClassification""",
				    """LongformerForTokenClassification""",
				    """LongformerModel""",
				    """LongformerPreTrainedModel""",
				    """LongformerSelfAttention""",
				]

try:
				if not is_tf_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				_a								=						[
				    """TF_LONGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST""",
				    """TFLongformerForMaskedLM""",
				    """TFLongformerForMultipleChoice""",
				    """TFLongformerForQuestionAnswering""",
				    """TFLongformerForSequenceClassification""",
				    """TFLongformerForTokenClassification""",
				    """TFLongformerModel""",
				    """TFLongformerPreTrainedModel""",
				    """TFLongformerSelfAttention""",
				]


if TYPE_CHECKING:
				from .configuration_longformer import (
				    LONGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
				    LongformerConfig,
				    LongformerOnnxConfig,
				)
				from .tokenization_longformer import LongformerTokenizer

				try:
								if not is_tokenizers_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .tokenization_longformer_fast import LongformerTokenizerFast

				try:
								if not is_torch_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_longformer import (
								    LONGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
								    LongformerForMaskedLM,
								    LongformerForMultipleChoice,
								    LongformerForQuestionAnswering,
								    LongformerForSequenceClassification,
								    LongformerForTokenClassification,
								    LongformerModel,
								    LongformerPreTrainedModel,
								    LongformerSelfAttention,
								)

				try:
								if not is_tf_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_tf_longformer import (
								    TF_LONGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
								    TFLongformerForMaskedLM,
								    TFLongformerForMultipleChoice,
								    TFLongformerForQuestionAnswering,
								    TFLongformerForSequenceClassification,
								    TFLongformerForTokenClassification,
								    TFLongformerModel,
								    TFLongformerPreTrainedModel,
								    TFLongformerSelfAttention,
								)

else:
				import sys

				_a								=						_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)








"""simple docstring"""





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)






