



"""simple docstring"""






import unittest
from queue import Empty
from threading import Thread

from transformers import AutoTokenizer, TextIteratorStreamer, TextStreamer, is_torch_available
from transformers.testing_utils import CaptureStdout, require_torch, torch_device

from ..test_modeling_common import ids_tensor


if is_torch_available():
	import torch

	from transformers import AutoModelForCausalLM



@require_torch
class      A__  (    unittest.TestCase				):




	'''simple docstring'''





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

				"""simple docstring"""

				__lowerCAmelCase :    Dict				=			AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
				__lowerCAmelCase :    Optional[Any]				=			AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    Optional[Any]				=			-1

				__lowerCAmelCase :    Any				=			ids_tensor((1, 5)			,     vocab_size=model.config.vocab_size).to(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    int				=			model.generate(_SCREAMING_SNAKE_CASE			,     max_new_tokens=10			,     do_sample=_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    List[Any]				=			tokenizer.decode(greedy_ids[0])

				with CaptureStdout() as cs:
							__lowerCAmelCase :    List[str]				=			TextStreamer(_SCREAMING_SNAKE_CASE)
							model.generate(_SCREAMING_SNAKE_CASE			,     max_new_tokens=10			,     do_sample=_SCREAMING_SNAKE_CASE			,     streamer=_SCREAMING_SNAKE_CASE)
				# The greedy text should be printed to stdout, except for the final "\n" in the streamer
				__lowerCAmelCase :    Union[str, Any]				=			cs.out[:-1]

				self.assertEqual(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)

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

				"""simple docstring"""

				__lowerCAmelCase :    List[Any]				=			AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
				__lowerCAmelCase :    Any				=			AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    int				=			-1

				__lowerCAmelCase :    str				=			ids_tensor((1, 5)			,     vocab_size=model.config.vocab_size).to(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    int				=			model.generate(_SCREAMING_SNAKE_CASE			,     max_new_tokens=10			,     do_sample=_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    List[Any]				=			tokenizer.decode(greedy_ids[0])

				__lowerCAmelCase :    Optional[int]				=			TextIteratorStreamer(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    Optional[Any]				=			{"input_ids": input_ids, "max_new_tokens": 10, "do_sample": False, "streamer": streamer}
				__lowerCAmelCase :    Union[str, Any]				=			Thread(target=model.generate			,     kwargs=_SCREAMING_SNAKE_CASE)
				thread.start()
				__lowerCAmelCase :    List[Any]				=			""
				for new_text in streamer:
							streamer_text += new_text

				self.assertEqual(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)

	def 				_SCREAMING_SNAKE_CASE  (	self:						Any)							->			Any:

				"""simple docstring"""

				__lowerCAmelCase :    Union[str, Any]				=			AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
				__lowerCAmelCase :    List[str]				=			AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    Dict				=			-1

				__lowerCAmelCase :    List[Any]				=			ids_tensor((1, 5)			,     vocab_size=model.config.vocab_size).to(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    List[str]				=			model.generate(_SCREAMING_SNAKE_CASE			,     max_new_tokens=10			,     do_sample=_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    Union[str, Any]				=			greedy_ids[:, input_ids.shape[1] :]
				__lowerCAmelCase :    Optional[int]				=			tokenizer.decode(new_greedy_ids[0])

				with CaptureStdout() as cs:
							__lowerCAmelCase :    Optional[int]				=			TextStreamer(_SCREAMING_SNAKE_CASE			,     skip_prompt=_SCREAMING_SNAKE_CASE)
							model.generate(_SCREAMING_SNAKE_CASE			,     max_new_tokens=10			,     do_sample=_SCREAMING_SNAKE_CASE			,     streamer=_SCREAMING_SNAKE_CASE)
				# The greedy text should be printed to stdout, except for the final "\n" in the streamer
				__lowerCAmelCase :    Dict				=			cs.out[:-1]

				self.assertEqual(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)

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

				"""simple docstring"""

				__lowerCAmelCase :    str				=			AutoTokenizer.from_pretrained("distilgpt2")
				__lowerCAmelCase :    str				=			AutoModelForCausalLM.from_pretrained("distilgpt2").to(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    int				=			-1

				__lowerCAmelCase :    int				=			torch.ones((1, 5)			,     device=_SCREAMING_SNAKE_CASE).long() * model.config.bos_token_id
				with CaptureStdout() as cs:
							__lowerCAmelCase :    Dict				=			TextStreamer(_SCREAMING_SNAKE_CASE			,     skip_special_tokens=_SCREAMING_SNAKE_CASE)
							model.generate(_SCREAMING_SNAKE_CASE			,     max_new_tokens=1			,     do_sample=_SCREAMING_SNAKE_CASE			,     streamer=_SCREAMING_SNAKE_CASE)

				# The prompt contains a special token, so the streamer should not print it. As such, the output text, when
				# re-tokenized, must only contain one token
				__lowerCAmelCase :    str				=			cs.out[:-1]  # Remove the final "\n"
				__lowerCAmelCase :    Optional[int]				=			tokenizer(_SCREAMING_SNAKE_CASE			,     return_tensors="pt")
				self.assertEqual(streamer_text_tokenized.input_ids.shape			,     (1, 1))







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

				"""simple docstring"""

				__lowerCAmelCase :    Any				=			AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
				__lowerCAmelCase :    int				=			AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    Optional[Any]				=			-1

				__lowerCAmelCase :    List[str]				=			ids_tensor((1, 5)			,     vocab_size=model.config.vocab_size).to(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    int				=			TextIteratorStreamer(_SCREAMING_SNAKE_CASE			,     timeout=0.001)
				__lowerCAmelCase :    Optional[Any]				=			{"input_ids": input_ids, "max_new_tokens": 10, "do_sample": False, "streamer": streamer}
				__lowerCAmelCase :    Optional[Any]				=			Thread(target=model.generate			,     kwargs=_SCREAMING_SNAKE_CASE)
				thread.start()

				# The streamer will timeout after 0.001 seconds, so an exception will be raised
				with self.assertRaises(_SCREAMING_SNAKE_CASE):
							__lowerCAmelCase :    Optional[int]				=			""
							for new_text in streamer:
										streamer_text += new_text




"""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)