






import argparse
import json
import os
import pickle
import shutil

import numpy as np
import torch
from distiller import Distiller
from lm_seqs_dataset import LmSeqsDataset

from transformers import (
    BertConfig,
    BertForMaskedLM,
    BertTokenizer,
    DistilBertConfig,
    DistilBertForMaskedLM,
    DistilBertTokenizer,
    GPTaConfig,
    GPTaLMHeadModel,
    GPTaTokenizer,
    RobertaConfig,
    RobertaForMaskedLM,
    RobertaTokenizer,
)
from utils import git_log, init_gpu_params, logger, set_seed


lowerCamelCase			:					Tuple         =	{
    '''distilbert''': (DistilBertConfig, DistilBertForMaskedLM, DistilBertTokenizer),
    '''roberta''': (RobertaConfig, RobertaForMaskedLM, RobertaTokenizer),
    '''bert''': (BertConfig, BertForMaskedLM, BertTokenizer),
    '''gpt2''': (GPTaConfig, GPTaLMHeadModel, GPTaTokenizer),
}





def       __lowerCAmelCase							(						__snake_case       ):
     assert (args.mlm and args.alpha_mlm > 0.0) or (not args.mlm and args.alpha_mlm == 0.0)
     assert (args.alpha_mlm > 0.0 and args.alpha_clm == 0.0) or (args.alpha_mlm == 0.0 and args.alpha_clm > 0.0)
     if args.mlm:
          assert os.path.isfile(args.token_counts       )
          assert (args.student_type in ["roberta", "distilbert"]) and (args.teacher_type in ["roberta", "bert"])
     else:
          assert (args.student_type in ["gpt2"]) and (args.teacher_type in ["gpt2"])

     assert args.teacher_type == args.student_type or (
         args.student_type == "distilbert" and args.teacher_type == "bert"
     )
     assert os.path.isfile(args.student_config       )
     if args.student_pretrained_weights is not None:
          assert os.path.isfile(args.student_pretrained_weights       )

     if args.freeze_token_type_embds:
          assert args.student_type in ["roberta"]

     assert args.alpha_ce >= 0.0
     assert args.alpha_mlm >= 0.0
     assert args.alpha_clm >= 0.0
     assert args.alpha_mse >= 0.0
     assert args.alpha_cos >= 0.0
     assert args.alpha_ce + args.alpha_mlm + args.alpha_clm + args.alpha_mse + args.alpha_cos > 0.0





def       __lowerCAmelCase							(						__snake_case					, __snake_case       ):






     if args.student_type == "roberta":
          __lowerCAmelCase								=					False
     elif args.student_type == "gpt2":
          __lowerCAmelCase								=					False





def       __lowerCAmelCase							(						__snake_case					, __snake_case       ):






     if args.student_type == "roberta":
          __lowerCAmelCase								=					False





def       __lowerCAmelCase							(						):
     __lowerCAmelCase								=					argparse.ArgumentParser(description="Training"       )
     parser.add_argument("--force"					, action="store_true"					, help="Overwrite dump_path if it already exists."       )

     parser.add_argument(
         "--dump_path"					, type=__SCREAMING_SNAKE_CASE					, required=__SCREAMING_SNAKE_CASE					, help="The output directory (log, checkpoints, parameters, etc.)"       )
     parser.add_argument(
         "--data_file"					, type=__SCREAMING_SNAKE_CASE					, required=__SCREAMING_SNAKE_CASE					, help="The binarized file (tokenized + tokens_to_ids) and grouped by sequence."					, )

     parser.add_argument(
         "--student_type"					, type=__SCREAMING_SNAKE_CASE					, choices=["distilbert", "roberta", "gpt2"]					, required=__SCREAMING_SNAKE_CASE					, help="The student type (DistilBERT, RoBERTa)."					, )
     parser.add_argument("--student_config"					, type=__SCREAMING_SNAKE_CASE					, required=__SCREAMING_SNAKE_CASE					, help="Path to the student configuration."       )
     parser.add_argument(
         "--student_pretrained_weights"					, default=__SCREAMING_SNAKE_CASE					, type=__SCREAMING_SNAKE_CASE					, help="Load student initialization checkpoint."       )

     parser.add_argument(
         "--teacher_type"					, choices=["bert", "roberta", "gpt2"]					, required=__SCREAMING_SNAKE_CASE					, help="Teacher type (BERT, RoBERTa)."       )
     parser.add_argument("--teacher_name"					, type=__SCREAMING_SNAKE_CASE					, required=__SCREAMING_SNAKE_CASE					, help="The teacher model."       )

     parser.add_argument("--temperature"					, default=2.0					, type=__SCREAMING_SNAKE_CASE					, help="Temperature for the softmax temperature."       )
     parser.add_argument(
         "--alpha_ce"					, default=0.5					, type=__SCREAMING_SNAKE_CASE					, help="Linear weight for the distillation loss. Must be >=0."       )
     parser.add_argument(
         "--alpha_mlm"					, default=0.0					, type=__SCREAMING_SNAKE_CASE					, help="Linear weight for the MLM loss. Must be >=0. Should be used in conjunction with `mlm` flag."					, )
     parser.add_argument("--alpha_clm"					, default=0.5					, type=__SCREAMING_SNAKE_CASE					, help="Linear weight for the CLM loss. Must be >=0."       )
     parser.add_argument("--alpha_mse"					, default=0.0					, type=__SCREAMING_SNAKE_CASE					, help="Linear weight of the MSE loss. Must be >=0."       )
     parser.add_argument(
         "--alpha_cos"					, default=0.0					, type=__SCREAMING_SNAKE_CASE					, help="Linear weight of the cosine embedding loss. Must be >=0."       )

     parser.add_argument(
         "--mlm"					, action="store_true"					, help="The LM step: MLM or CLM. If `mlm` is True, the MLM is used over CLM."       )
     parser.add_argument(
         "--mlm_mask_prop"					, default=0.15					, type=__SCREAMING_SNAKE_CASE					, help="Proportion of tokens for which we need to make a prediction."					, )
     parser.add_argument("--word_mask"					, default=0.8					, type=__SCREAMING_SNAKE_CASE					, help="Proportion of tokens to mask out."       )
     parser.add_argument("--word_keep"					, default=0.1					, type=__SCREAMING_SNAKE_CASE					, help="Proportion of tokens to keep."       )
     parser.add_argument("--word_rand"					, default=0.1					, type=__SCREAMING_SNAKE_CASE					, help="Proportion of tokens to randomly replace."       )
     parser.add_argument(
         "--mlm_smoothing"					, default=0.7					, type=__SCREAMING_SNAKE_CASE					, help="Smoothing parameter to emphasize more rare tokens (see XLM, similar to word2vec)."					, )
     parser.add_argument("--token_counts"					, type=__SCREAMING_SNAKE_CASE					, help="The token counts in the data_file for MLM."       )

     parser.add_argument(
         "--restrict_ce_to_mask"					, action="store_true"					, help="If true, compute the distillation loss only the [MLM] prediction distribution."					, )
     parser.add_argument(
         "--freeze_pos_embs"					, action="store_true"					, help="Freeze positional embeddings during distillation. For student_type in ['roberta', 'gpt2'] only."					, )
     parser.add_argument(
         "--freeze_token_type_embds"					, action="store_true"					, help="Freeze token type embeddings during distillation if existent. For student_type in ['roberta'] only."					, )

     parser.add_argument("--n_epoch"					, type=__SCREAMING_SNAKE_CASE					, default=3					, help="Number of pass on the whole dataset."       )
     parser.add_argument("--batch_size"					, type=__SCREAMING_SNAKE_CASE					, default=5					, help="Batch size (for each process)."       )
     parser.add_argument(
         "--group_by_size"					, action="store_false"					, help="If true, group sequences that have similar length into the same batch. Default is true."					, )

     parser.add_argument(
         "--gradient_accumulation_steps"					, type=__SCREAMING_SNAKE_CASE					, default=50					, help="Gradient accumulation for larger training batches."					, )
     parser.add_argument("--warmup_prop"					, default=0.05					, type=__SCREAMING_SNAKE_CASE					, help="Linear warmup proportion."       )
     parser.add_argument("--weight_decay"					, default=0.0					, type=__SCREAMING_SNAKE_CASE					, help="Weight decay if we apply some."       )
     parser.add_argument("--learning_rate"					, default=5E-4					, type=__SCREAMING_SNAKE_CASE					, help="The initial learning rate for Adam."       )
     parser.add_argument("--adam_epsilon"					, default=1E-6					, type=__SCREAMING_SNAKE_CASE					, help="Epsilon for Adam optimizer."       )
     parser.add_argument("--max_grad_norm"					, default=5.0					, type=__SCREAMING_SNAKE_CASE					, help="Max gradient norm."       )
     parser.add_argument("--initializer_range"					, default=0.02					, type=__SCREAMING_SNAKE_CASE					, help="Random initialization range."       )

     parser.add_argument(
         "--fp16"					, action="store_true"					, help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit"					, )
     parser.add_argument(
         "--fp16_opt_level"					, type=__SCREAMING_SNAKE_CASE					, default="O1"					, help=(
             "For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
             "See details at https://nvidia.github.io/apex/amp.html"
         )					, )
     parser.add_argument("--n_gpu"					, type=__SCREAMING_SNAKE_CASE					, default=1					, help="Number of GPUs in the node."       )
     parser.add_argument("--local_rank"					, type=__SCREAMING_SNAKE_CASE					, default=-1					, help="Distributed training - Local rank"       )
     parser.add_argument("--seed"					, type=__SCREAMING_SNAKE_CASE					, default=56					, help="Random seed"       )

     parser.add_argument("--log_interval"					, type=__SCREAMING_SNAKE_CASE					, default=500					, help="Tensorboard logging interval."       )
     parser.add_argument("--checkpoint_interval"					, type=__SCREAMING_SNAKE_CASE					, default=4000					, help="Checkpoint interval."       )
     __lowerCAmelCase								=					parser.parse_args()
     sanity_checks(__SCREAMING_SNAKE_CASE       )

     # ARGS #
     init_gpu_params(__SCREAMING_SNAKE_CASE       )
     set_seed(__SCREAMING_SNAKE_CASE       )
     if args.is_master:
          if os.path.exists(args.dump_path       ):
               if not args.force:
                    raise ValueError(
                        F"""Serialization dir {args.dump_path} already exists, but you have not precised wheter to overwrite"""
                        " itUse `--force` if you want to overwrite it"       )
               else:
                    shutil.rmtree(args.dump_path       )

          if not os.path.exists(args.dump_path       ):
               os.makedirs(args.dump_path       )
          logger.info(F"""Experiment will be dumped and logged in {args.dump_path}"""       )

          # SAVE PARAMS #
          logger.info(F"""Param: {args}"""       )
          with open(os.path.join(args.dump_path					, "parameters.json"       )					, "w"       ) as f:
               json.dump(vars(__SCREAMING_SNAKE_CASE       )					, __SCREAMING_SNAKE_CASE					, indent=4       )
          git_log(args.dump_path       )

     __lowerCAmelCase								=					MODEL_CLASSES[args.student_type]
     __lowerCAmelCase								=					MODEL_CLASSES[args.teacher_type]

     # TOKENIZER #
     __lowerCAmelCase								=					teacher_tokenizer_class.from_pretrained(args.teacher_name       )
     __lowerCAmelCase								=					{}
     for tok_name, tok_symbol in tokenizer.special_tokens_map.items():
          __lowerCAmelCase								=					tokenizer.all_special_tokens.index(__SCREAMING_SNAKE_CASE       )
          __lowerCAmelCase								=					tokenizer.all_special_ids[idx]
     logger.info(F"""Special tokens {special_tok_ids}"""       )
     __lowerCAmelCase								=					special_tok_ids
     __lowerCAmelCase								=					tokenizer.max_model_input_sizes[args.teacher_name]

     # DATA LOADER #
     logger.info(F"""Loading data from {args.data_file}"""       )
     with open(args.data_file					, "rb"       ) as fp:
          __lowerCAmelCase								=					pickle.load(__SCREAMING_SNAKE_CASE       )

     if args.mlm:
          logger.info(F"""Loading token counts from {args.token_counts} (already pre-computed)"""       )
          with open(args.token_counts					, "rb"       ) as fp:
               __lowerCAmelCase								=					pickle.load(__SCREAMING_SNAKE_CASE       )

          __lowerCAmelCase								=					np.maximum(__SCREAMING_SNAKE_CASE					, 1       ) ** -args.mlm_smoothing
          for idx in special_tok_ids.values():
               __lowerCAmelCase								=					0.0  # do not predict special tokens
          __lowerCAmelCase								=					torch.from_numpy(__SCREAMING_SNAKE_CASE       )
     else:
          __lowerCAmelCase								=					None

     __lowerCAmelCase								=					LmSeqsDataset(params=__SCREAMING_SNAKE_CASE					, data=__SCREAMING_SNAKE_CASE       )
     logger.info("Data loader created."       )

     # STUDENT #
     logger.info(F"""Loading student config from {args.student_config}"""       )
     __lowerCAmelCase								=					student_config_class.from_pretrained(args.student_config       )
     __lowerCAmelCase								=					True

     if args.student_pretrained_weights is not None:
          logger.info(F"""Loading pretrained weights from {args.student_pretrained_weights}"""       )
          __lowerCAmelCase								=					student_model_class.from_pretrained(args.student_pretrained_weights					, config=__SCREAMING_SNAKE_CASE       )
     else:
          __lowerCAmelCase								=					student_model_class(__SCREAMING_SNAKE_CASE       )

     if args.n_gpu > 0:
          student.to(F"""cuda:{args.local_rank}"""       )
     logger.info("Student loaded."       )

     # TEACHER #
     __lowerCAmelCase								=					teacher_model_class.from_pretrained(args.teacher_name					, output_hidden_states=__SCREAMING_SNAKE_CASE       )
     if args.n_gpu > 0:
          teacher.to(F"""cuda:{args.local_rank}"""       )
     logger.info(F"""Teacher loaded from {args.teacher_name}."""       )

     # FREEZING #
     if args.freeze_pos_embs:
          freeze_pos_embeddings(__SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE       )
     if args.freeze_token_type_embds:
          freeze_token_type_embeddings(__SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE       )

     # SANITY CHECKS #
     assert student.config.vocab_size == teacher.config.vocab_size
     assert student.config.hidden_size == teacher.config.hidden_size
     assert student.config.max_position_embeddings == teacher.config.max_position_embeddings
     if args.mlm:
          assert token_probs.size(0       ) == stu_architecture_config.vocab_size

     # DISTILLER #
     torch.cuda.empty_cache()
     __lowerCAmelCase								=					Distiller(
         params=__SCREAMING_SNAKE_CASE					, dataset=__SCREAMING_SNAKE_CASE					, token_probs=__SCREAMING_SNAKE_CASE					, student=__SCREAMING_SNAKE_CASE					, teacher=__SCREAMING_SNAKE_CASE       )
     distiller.train()
     logger.info("Let's go get some drinks."       )


if __name__ == "__main__":
 main()














import json
import os
import shutil
import tempfile
import unittest

import numpy as np

from transformers import BertTokenizerFast
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES, BertTokenizer
from transformers.testing_utils import require_tokenizers, require_vision
from transformers.utils import IMAGE_PROCESSOR_NAME, is_vision_available


if is_vision_available():
 from PIL import Image

 from transformers import VisionTextDualEncoderProcessor, ViTImageProcessor



@require_tokenizers
@require_vision
class 							_UpperCamelCase			(unittest.TestCase			):


      def  __UpperCAmelCase		(	self	)->	Optional[Any]:
           __lowerCAmelCase								=					tempfile.mkdtemp()

           # fmt: off
           __lowerCAmelCase								=					["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]", "want", "##want", "##ed", "wa", "un", "runn", "##ing", ",", "low", "lowest"]
           # fmt: on
           __lowerCAmelCase								=					os.path.join(self.tmpdirname					,						VOCAB_FILES_NAMES["vocab_file"]	)
           with open(self.vocab_file					,						"w"					,						encoding="utf-8"	) as vocab_writer:
                vocab_writer.write("".join([x + "\n" for x in vocab_tokens]	)	)

           __lowerCAmelCase								=					{
               "do_resize": True,
               "size": {"height": 1_8, "width": 1_8},
               "do_normalize": True,
               "image_mean": [0.5, 0.5, 0.5],
               "image_std": [0.5, 0.5, 0.5],
           }
           __lowerCAmelCase								=					os.path.join(self.tmpdirname					,						__UpperCamelCase	)
           with open(self.image_processor_file					,						"w"					,						encoding="utf-8"	) as fp:
                json.dump(__UpperCamelCase					,						__UpperCamelCase	)


      def  __UpperCAmelCase		(	self					,						**__UpperCamelCase	)->	Optional[Any]:
           return BertTokenizer.from_pretrained(self.tmpdirname					,						**__UpperCamelCase	)


      def  __UpperCAmelCase		(	self					,						**__UpperCamelCase	)->	Any:
           return ViTImageProcessor.from_pretrained(self.tmpdirname					,						**__UpperCamelCase	)


      def  __UpperCAmelCase		(	self	)->	Any:
           shutil.rmtree(self.tmpdirname	)


      def  __UpperCAmelCase		(	self	)->	int:
           __lowerCAmelCase								=					[np.random.randint(2_5_5					,						size=(3, 3_0, 4_0_0)					,						dtype=np.uinta	)]

           __lowerCAmelCase								=					[Image.fromarray(np.moveaxis(__UpperCamelCase					,						0					,						-1	)	) for x in image_inputs]

           return image_inputs


      def  __UpperCAmelCase		(	self	)->	Optional[Any]:
           __lowerCAmelCase								=					self.get_tokenizer()
           __lowerCAmelCase								=					self.get_image_processor()

           __lowerCAmelCase								=					VisionTextDualEncoderProcessor(tokenizer=__UpperCamelCase					,						image_processor=__UpperCamelCase	)

           processor.save_pretrained(self.tmpdirname	)
           __lowerCAmelCase								=					VisionTextDualEncoderProcessor.from_pretrained(self.tmpdirname	)

           self.assertEqual(processor.tokenizer.get_vocab()					,						tokenizer.get_vocab()	)
           self.assertIsInstance(processor.tokenizer					,						(BertTokenizer, BertTokenizerFast)	)

           self.assertEqual(processor.image_processor.to_json_string()					,						image_processor.to_json_string()	)
           self.assertIsInstance(processor.image_processor					,						__UpperCamelCase	)


      def  __UpperCAmelCase		(	self	)->	str:
           __lowerCAmelCase								=					VisionTextDualEncoderProcessor(
               tokenizer=self.get_tokenizer()					,						image_processor=self.get_image_processor()	)
           processor.save_pretrained(self.tmpdirname	)

           __lowerCAmelCase								=					self.get_tokenizer(bos_token="(BOS)"					,						eos_token="(EOS)"	)
           __lowerCAmelCase								=					self.get_image_processor(do_normalize=__UpperCamelCase					,						padding_value=1.0	)

           __lowerCAmelCase								=					VisionTextDualEncoderProcessor.from_pretrained(
               self.tmpdirname					,						bos_token="(BOS)"					,						eos_token="(EOS)"					,						do_normalize=__UpperCamelCase					,						padding_value=1.0	)

           self.assertEqual(processor.tokenizer.get_vocab()					,						tokenizer_add_kwargs.get_vocab()	)
           self.assertIsInstance(processor.tokenizer					,						(BertTokenizer, BertTokenizerFast)	)

           self.assertEqual(processor.image_processor.to_json_string()					,						image_processor_add_kwargs.to_json_string()	)
           self.assertIsInstance(processor.image_processor					,						__UpperCamelCase	)


      def  __UpperCAmelCase		(	self	)->	Union[str, Any]:
           __lowerCAmelCase								=					self.get_image_processor()
           __lowerCAmelCase								=					self.get_tokenizer()

           __lowerCAmelCase								=					VisionTextDualEncoderProcessor(tokenizer=__UpperCamelCase					,						image_processor=__UpperCamelCase	)

           __lowerCAmelCase								=					self.prepare_image_inputs()

           __lowerCAmelCase								=					image_processor(__UpperCamelCase					,						return_tensors="np"	)
           __lowerCAmelCase								=					processor(images=__UpperCamelCase					,						return_tensors="np"	)

           for key in input_feat_extract.keys():
                self.assertAlmostEqual(input_feat_extract[key].sum()					,						input_processor[key].sum()					,						delta=1e-2	)


      def  __UpperCAmelCase		(	self	)->	Dict:
           __lowerCAmelCase								=					self.get_image_processor()
           __lowerCAmelCase								=					self.get_tokenizer()

           __lowerCAmelCase								=					VisionTextDualEncoderProcessor(tokenizer=__UpperCamelCase					,						image_processor=__UpperCamelCase	)

           __lowerCAmelCase								=					"lower newer"

           __lowerCAmelCase								=					processor(text=__UpperCamelCase	)

           __lowerCAmelCase								=					tokenizer(__UpperCamelCase	)

           for key in encoded_tok.keys():
                self.assertListEqual(encoded_tok[key]					,						encoded_processor[key]	)


      def  __UpperCAmelCase		(	self	)->	List[str]:
           __lowerCAmelCase								=					self.get_image_processor()
           __lowerCAmelCase								=					self.get_tokenizer()

           __lowerCAmelCase								=					VisionTextDualEncoderProcessor(tokenizer=__UpperCamelCase					,						image_processor=__UpperCamelCase	)

           __lowerCAmelCase								=					"lower newer"
           __lowerCAmelCase								=					self.prepare_image_inputs()

           __lowerCAmelCase								=					processor(text=__UpperCamelCase					,						images=__UpperCamelCase	)

           self.assertListEqual(list(inputs.keys()	)					,						["input_ids", "token_type_ids", "attention_mask", "pixel_values"]	)

           # test if it raises when no input is passed
           with self.assertRaises(__UpperCamelCase	):
                processor()


      def  __UpperCAmelCase		(	self	)->	Optional[int]:
           __lowerCAmelCase								=					self.get_image_processor()
           __lowerCAmelCase								=					self.get_tokenizer()

           __lowerCAmelCase								=					VisionTextDualEncoderProcessor(tokenizer=__UpperCamelCase					,						image_processor=__UpperCamelCase	)

           __lowerCAmelCase								=					[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

           __lowerCAmelCase								=					processor.batch_decode(__UpperCamelCase	)
           __lowerCAmelCase								=					tokenizer.batch_decode(__UpperCamelCase	)

           self.assertListEqual(__UpperCamelCase					,						__UpperCamelCase	)


      def  __UpperCAmelCase		(	self	)->	Any:
           __lowerCAmelCase								=					self.get_image_processor()
           __lowerCAmelCase								=					self.get_tokenizer()

           __lowerCAmelCase								=					VisionTextDualEncoderProcessor(tokenizer=__UpperCamelCase					,						image_processor=__UpperCamelCase	)

           __lowerCAmelCase								=					"lower newer"
           __lowerCAmelCase								=					self.prepare_image_inputs()

           __lowerCAmelCase								=					processor(text=__UpperCamelCase					,						images=__UpperCamelCase	)

           self.assertListEqual(list(inputs.keys()	)					,						processor.model_input_names	)






