




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       UpperCAmelCase       (     _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       UpperCAmelCase       (     _snake_case						, _snake_case							):
							if args.student_type == "roberta":
														lowerCAmelCase   =		False
							elif args.student_type == "gpt2":
														lowerCAmelCase   =		False



def       UpperCAmelCase       (     _snake_case						, _snake_case							):
							if args.student_type == "roberta":
														lowerCAmelCase   =		False



def       UpperCAmelCase       (     ):
							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=_snake_case						, required=_snake_case						, help='''The output directory (log, checkpoints, parameters, etc.)'''							)
							parser.add_argument(
							    '''--data_file'''						, type=_snake_case						, required=_snake_case						, help='''The binarized file (tokenized + tokens_to_ids) and grouped by sequence.'''						, )

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

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

							parser.add_argument('''--temperature'''						, default=2.0						, type=_snake_case						, help='''Temperature for the softmax temperature.'''							)
							parser.add_argument(
							    '''--alpha_ce'''						, default=0.5						, type=_snake_case						, help='''Linear weight for the distillation loss. Must be >=0.'''							)
							parser.add_argument(
							    '''--alpha_mlm'''						, default=0.0						, type=_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=_snake_case						, help='''Linear weight for the CLM loss. Must be >=0.'''							)
							parser.add_argument('''--alpha_mse'''						, default=0.0						, type=_snake_case						, help='''Linear weight of the MSE loss. Must be >=0.'''							)
							parser.add_argument(
							    '''--alpha_cos'''						, default=0.0						, type=_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=_snake_case						, help='''Proportion of tokens for which we need to make a prediction.'''						, )
							parser.add_argument('''--word_mask'''						, default=0.8						, type=_snake_case						, help='''Proportion of tokens to mask out.'''							)
							parser.add_argument('''--word_keep'''						, default=0.1						, type=_snake_case						, help='''Proportion of tokens to keep.'''							)
							parser.add_argument('''--word_rand'''						, default=0.1						, type=_snake_case						, help='''Proportion of tokens to randomly replace.'''							)
							parser.add_argument(
							    '''--mlm_smoothing'''						, default=0.7						, type=_snake_case						, help='''Smoothing parameter to emphasize more rare tokens (see XLM, similar to word2vec).'''						, )
							parser.add_argument('''--token_counts'''						, type=_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=_snake_case						, default=3						, help='''Number of pass on the whole dataset.'''							)
							parser.add_argument('''--batch_size'''						, type=_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=_snake_case						, default=50						, help='''Gradient accumulation for larger training batches.'''						, )
							parser.add_argument('''--warmup_prop'''						, default=0.05						, type=_snake_case						, help='''Linear warmup proportion.'''							)
							parser.add_argument('''--weight_decay'''						, default=0.0						, type=_snake_case						, help='''Weight decay if we apply some.'''							)
							parser.add_argument('''--learning_rate'''						, default=5E-4						, type=_snake_case						, help='''The initial learning rate for Adam.'''							)
							parser.add_argument('''--adam_epsilon'''						, default=1E-6						, type=_snake_case						, help='''Epsilon for Adam optimizer.'''							)
							parser.add_argument('''--max_grad_norm'''						, default=5.0						, type=_snake_case						, help='''Max gradient norm.'''							)
							parser.add_argument('''--initializer_range'''						, default=0.02						, type=_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=_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=_snake_case						, default=1						, help='''Number of GPUs in the node.'''							)
							parser.add_argument('''--local_rank'''						, type=_snake_case						, default=-1						, help='''Distributed training - Local rank'''							)
							parser.add_argument('''--seed'''						, type=_snake_case						, default=56						, help='''Random seed'''							)

							parser.add_argument('''--log_interval'''						, type=_snake_case						, default=500						, help='''Tensorboard logging interval.'''							)
							parser.add_argument('''--checkpoint_interval'''						, type=_snake_case						, default=4000						, help='''Checkpoint interval.'''							)
							lowerCAmelCase   =		parser.parse_args()
							sanity_checks(_snake_case							)

							# ARGS #
							init_gpu_params(_snake_case							)
							set_seed(_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(_snake_case							)						, _snake_case						, indent=4							)
														git_log(args.dump_path							)

							lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   =		MODEL_CLASSES[args.student_type]
							lowerCAmelCase   ,     lowerCAmelCase   ,     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(_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(_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(_snake_case							)

														lowerCAmelCase   =		np.maximum(_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(_snake_case							)
							else:
														lowerCAmelCase   =		None

							lowerCAmelCase   =		LmSeqsDataset(params=_snake_case						, data=_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=_snake_case							)
							else:
														lowerCAmelCase   =		student_model_class(_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=_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(_snake_case						, _snake_case							)
							if args.freeze_token_type_embds:
														freeze_token_type_embeddings(_snake_case						, _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=_snake_case						, dataset=_snake_case						, token_probs=_snake_case						, student=_snake_case						, teacher=_snake_case							)
							distiller.train()
							logger.info('''Let\'s go get some drinks.'''							)


if __name__ == "__main__":
		main()









import unittest

from huggingface_hub import hf_hub_download

from transformers import MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING, VideoMAEFeatureExtractor
from transformers.pipelines import VideoClassificationPipeline, pipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_decord,
    require_tf,
    require_torch,
    require_torch_or_tf,
    require_vision,
)

from .test_pipelines_common import ANY


@is_pipeline_test
@require_torch_or_tf
@require_vision
@require_decord
class 							__UpperCamelCase			( unittest.TestCase  ):





				'''simple docstring'''



				__a							: str										=MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING



				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_						):
											lowerCAmelCase   =		hf_hub_download(
											    repo_id='''nateraw/video-demo'''       , filename='''archery.mp4'''       , repo_type='''dataset'''						)
											lowerCAmelCase   =		VideoClassificationPipeline(model=UpperCAmelCase_       , image_processor=UpperCAmelCase_       , top_k=2						)
											lowerCAmelCase   =		[
											    example_video_filepath,
											    '''https://huggingface.co/datasets/nateraw/video-demo/resolve/main/archery.mp4''',
											]
											return video_classifier, examples



				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_						):
											for example in examples:
																		lowerCAmelCase   =		video_classifier(UpperCAmelCase_						)

																		self.assertEqual(
																		    UpperCAmelCase_       , [
																		        {'''score''': ANY(UpperCAmelCase_						), '''label''': ANY(UpperCAmelCase_						)},
																		        {'''score''': ANY(UpperCAmelCase_						), '''label''': ANY(UpperCAmelCase_						)},
																		    ]       , )



				@require_torch
				def       __snake_case		( self						):
											lowerCAmelCase   =		'''hf-internal-testing/tiny-random-VideoMAEForVideoClassification'''
											lowerCAmelCase   =		VideoMAEFeatureExtractor(
											    size={'''shortest_edge''': 10}       , crop_size={'''height''': 10, '''width''': 10}						)
											lowerCAmelCase   =		pipeline(
											    '''video-classification'''       , model=UpperCAmelCase_       , feature_extractor=UpperCAmelCase_       , frame_sampling_rate=4						)

											lowerCAmelCase   =		hf_hub_download(repo_id='''nateraw/video-demo'''       , filename='''archery.mp4'''       , repo_type='''dataset'''						)
											lowerCAmelCase   =		video_classifier(UpperCAmelCase_       , top_k=2						)
											self.assertEqual(
											    nested_simplify(UpperCAmelCase_       , decimals=4						)       , [{'''score''': 0.5199, '''label''': '''LABEL_0'''}, {'''score''': 0.4801, '''label''': '''LABEL_1'''}]       , )

											lowerCAmelCase   =		video_classifier(
											    [
											        video_file_path,
											        video_file_path,
											    ]       , top_k=2       , )
											self.assertEqual(
											    nested_simplify(UpperCAmelCase_       , decimals=4						)       , [
											        [{'''score''': 0.5199, '''label''': '''LABEL_0'''}, {'''score''': 0.4801, '''label''': '''LABEL_1'''}],
											        [{'''score''': 0.5199, '''label''': '''LABEL_0'''}, {'''score''': 0.4801, '''label''': '''LABEL_1'''}],
											    ]       , )






				@require_tf
				def       __snake_case		( self						):
											pass



