


"""simple docstring"""






import argparse

import torch

from transformers import BertForMaskedLM


if __name__ == "__main__":
	_lowerCamelCase     =  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')
	_lowerCamelCase     =  parser.parse_args()

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

	_lowerCamelCase     =  model.state_dict()
	_lowerCamelCase     =  {}

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

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

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

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

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




"""simple docstring"""






from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_lowerCamelCase     =  {
    'configuration_autoformer': [
        'AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'AutoformerConfig',
    ],
}

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_lowerCamelCase     =  [
	    'AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST',
	    'AutoformerForPrediction',
	    'AutoformerModel',
	    'AutoformerPreTrainedModel',
	]


if TYPE_CHECKING:
	from .configuration_autoformer import (
	    AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
	    AutoformerConfig,
	)

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_autoformer import (
		    AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    AutoformerForPrediction,
		    AutoformerModel,
		    AutoformerPreTrainedModel,
		)

else:
	import sys

	_lowerCamelCase     =  _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)
