


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


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




def 	lowerCAmelCase_					(						__UpperCAmelCase:    List[Any]   )   ->			Union[str, Any]:
    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_					(						__UpperCAmelCase:    str	,     __UpperCAmelCase:    Optional[int]   )   ->			List[str]:
    if args.student_type == "roberta":
        UpperCamelCase__							:       str      =   False
    elif args.student_type == "gpt2":
        UpperCamelCase__							:       Union[str, Any]      =   False




def 	lowerCAmelCase_					(						__UpperCAmelCase:    Dict	,     __UpperCAmelCase:    str   )   ->			List[str]:
    if args.student_type == "roberta":
        UpperCamelCase__							:       Dict      =   False




def 	lowerCAmelCase_					(						)   ->			str:
    UpperCamelCase__							:       str      =   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.'''   )
    UpperCamelCase__							:       str      =   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   )

    UpperCamelCase__     ,UpperCamelCase__     ,UpperCamelCase__							:       str      =   MODEL_CLASSES[args.student_type]
    UpperCamelCase__     ,UpperCamelCase__     ,UpperCamelCase__							:       Union[str, Any]      =   MODEL_CLASSES[args.teacher_type]

    # TOKENIZER #
    UpperCamelCase__							:       Optional[int]      =   teacher_tokenizer_class.from_pretrained(args.teacher_name   )
    UpperCamelCase__							:       Union[str, Any]      =   {}
    for tok_name, tok_symbol in tokenizer.special_tokens_map.items():
        UpperCamelCase__							:       List[Any]      =   tokenizer.all_special_tokens.index(SCREAMING_SNAKE_CASE_   )
        UpperCamelCase__							:       Tuple      =   tokenizer.all_special_ids[idx]
    logger.info(f"Special tokens {special_tok_ids}"   )
    UpperCamelCase__							:       str      =   special_tok_ids
    UpperCamelCase__							:       Any      =   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:
        UpperCamelCase__							:       Any      =   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:
            UpperCamelCase__							:       Optional[int]      =   pickle.load(SCREAMING_SNAKE_CASE_   )

        UpperCamelCase__							:       Optional[Any]      =   np.maximum(SCREAMING_SNAKE_CASE_	,     1   ) ** -args.mlm_smoothing
        for idx in special_tok_ids.values():
            UpperCamelCase__							:       Optional[Any]      =   0.0  # do not predict special tokens
        UpperCamelCase__							:       Union[str, Any]      =   torch.from_numpy(SCREAMING_SNAKE_CASE_   )
    else:
        UpperCamelCase__							:       Union[str, Any]      =   None

    UpperCamelCase__							:       Union[str, Any]      =   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}"   )
    UpperCamelCase__							:       List[Any]      =   student_config_class.from_pretrained(args.student_config   )
    UpperCamelCase__							:       int      =   True

    if args.student_pretrained_weights is not None:
        logger.info(f"Loading pretrained weights from {args.student_pretrained_weights}"   )
        UpperCamelCase__							:       Tuple      =   student_model_class.from_pretrained(args.student_pretrained_weights	,     config=SCREAMING_SNAKE_CASE_   )
    else:
        UpperCamelCase__							:       Dict      =   student_model_class(SCREAMING_SNAKE_CASE_   )

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

    # TEACHER #
    UpperCamelCase__							:       Dict      =   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()
    UpperCamelCase__							:       Any      =   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 copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto.configuration_auto import CONFIG_MAPPING


UpperCAmelCase_  							=				logging.get_logger(__name__)






class 					lowercase__ (			__lowerCamelCase       ):

   '''simple docstring'''


   a		:							Any     	=	"upernet"




   def __init__(					self,	__magic_name__=None,	__magic_name__=512,	__magic_name__=0.02,	__magic_name__=[1, 2, 3, 6],	__magic_name__=True,	__magic_name__=0.4,	__magic_name__=384,	__magic_name__=256,	__magic_name__=1,	__magic_name__=False,	__magic_name__=255,	**__magic_name__,	)				->   Dict:






       """simple docstring"""
       super().__init__(**__magic_name__							)

       if backbone_config is None:
           logger.info('''`backbone_config` is `None`. Initializing the config with the default `ResNet` backbone.'''							)
           UpperCamelCase__							:       str      =   CONFIG_MAPPING['''resnet'''](out_features=['''stage1''', '''stage2''', '''stage3''', '''stage4''']							)
       elif isinstance(__magic_name__,	__magic_name__							):
           UpperCamelCase__							:       Tuple      =   backbone_config.get('''model_type'''							)
           UpperCamelCase__							:       List[str]      =   CONFIG_MAPPING[backbone_model_type]
           UpperCamelCase__							:       Optional[int]      =   config_class.from_dict(__magic_name__							)

       UpperCamelCase__							:       Dict      =   backbone_config
       UpperCamelCase__							:       int      =   hidden_size
       UpperCamelCase__							:       int      =   initializer_range
       UpperCamelCase__							:       Dict      =   pool_scales
       UpperCamelCase__							:       Tuple      =   use_auxiliary_head
       UpperCamelCase__							:       Tuple      =   auxiliary_loss_weight
       UpperCamelCase__							:       Optional[int]      =   auxiliary_in_channels
       UpperCamelCase__							:       Optional[Any]      =   auxiliary_channels
       UpperCamelCase__							:       Any      =   auxiliary_num_convs
       UpperCamelCase__							:       Union[str, Any]      =   auxiliary_concat_input
       UpperCamelCase__							:       List[str]      =   loss_ignore_index




   def    UpperCamelCase__						(					self							)				->   Any:






       """simple docstring"""
       UpperCamelCase__							:       List[str]      =   copy.deepcopy(self.__dict__							)
       UpperCamelCase__							:       int      =   self.backbone_config.to_dict()
       UpperCamelCase__							:       Any      =   self.__class__.model_type
       return output

