


'''simple docstring'''

import argparse

import torch

from transformers import BertForMaskedLM


if __name__ == "__main__":
	_UpperCAmelCase				: Optional[int]							    =				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""")
	_UpperCAmelCase				: Union[str, Any]							    =				parser.parse_args()

	if args.model_type == "bert":
		_UpperCAmelCase				: Tuple							    =				BertForMaskedLM.from_pretrained(args.model_name)
		_UpperCAmelCase				: List[Any]							    =				"""bert"""
	else:
		raise ValueError("""args.model_type should be \"bert\".""")

	_UpperCAmelCase				: Any							    =				model.state_dict()
	_UpperCAmelCase				: int							    =				{}

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

	_UpperCAmelCase				: Optional[Any]							    =				0
	for teacher_idx in [0, 2, 4, 7, 9, 1_1]:
		for w in ["weight", "bias"]:
			_UpperCAmelCase				: List[Any]							    =				state_dict[
			    f"""{prefix}.encoder.layer.{teacher_idx}.attention.self.query.{w}"""
			]
			_UpperCAmelCase				: int							    =				state_dict[
			    f"""{prefix}.encoder.layer.{teacher_idx}.attention.self.key.{w}"""
			]
			_UpperCAmelCase				: Tuple							    =				state_dict[
			    f"""{prefix}.encoder.layer.{teacher_idx}.attention.self.value.{w}"""
			]

			_UpperCAmelCase				: List[str]							    =				state_dict[
			    f"""{prefix}.encoder.layer.{teacher_idx}.attention.output.dense.{w}"""
			]
			_UpperCAmelCase				: Tuple							    =				state_dict[
			    f"""{prefix}.encoder.layer.{teacher_idx}.attention.output.LayerNorm.{w}"""
			]

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

	_UpperCAmelCase				: Any							    =				state_dict["""cls.predictions.decoder.weight"""]
	_UpperCAmelCase				: Any							    =				state_dict["""cls.predictions.bias"""]
	if args.vocab_transform:
		for w in ["weight", "bias"]:
			_UpperCAmelCase				: Tuple							    =				state_dict[f"""cls.predictions.transform.dense.{w}"""]
			_UpperCAmelCase				: Union[str, Any]							    =				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 math import asin, atan, cos, radians, sin, sqrt, tan

_UpperCAmelCase				: str							    =				6_37_81_37.0
_UpperCAmelCase				: Tuple							    =				6_35_67_52.31_42_45
_UpperCAmelCase				: Optional[Any]							    =				6_3_7_8_1_3_7

def 							__magic_name__(						lowerCamelCase,				lowerCamelCase,				lowerCamelCase,				lowerCamelCase):
						__lowerCAmelCase							       =				(AXIS_A - AXIS_B) / AXIS_A
						__lowerCAmelCase							       =				atan((1 - flattening) * tan(radians(lowerCamelCase)))
						__lowerCAmelCase							       =				atan((1 - flattening) * tan(radians(lowerCamelCase)))
						__lowerCAmelCase							       =				radians(lowerCamelCase)
						__lowerCAmelCase							       =				radians(lowerCamelCase)
						# Equation
						__lowerCAmelCase							       =				sin((phi_a - phi_a) / 2)
						__lowerCAmelCase							       =				sin((lambda_a - lambda_a) / 2)
						# Square both values
						sin_sq_phi *= sin_sq_phi
						sin_sq_lambda *= sin_sq_lambda
						__lowerCAmelCase							       =				sqrt(sin_sq_phi + (cos(lowerCamelCase) * cos(lowerCamelCase) * sin_sq_lambda))
						return 2 * RADIUS * asin(lowerCamelCase)


if __name__ == "__main__":
	import doctest

	doctest.testmod()

