



import argparse
import logging
import os
import sys
import tempfile
from pathlib import Path

import lightning_base
import pytest
import pytorch_lightning as pl
import torch
from convert_pl_checkpoint_to_hf import convert_pl_to_hf
from distillation import distill_main
from finetune import SummarizationModule, main
from huggingface_hub import list_models
from parameterized import parameterized
from run_eval import generate_summaries_or_translations
from torch import nn

from transformers import AutoConfig, AutoModelForSeqaSeqLM
from transformers.testing_utils import CaptureStderr, CaptureStdout, TestCasePlus, require_torch_gpu, slow
from utils import label_smoothed_nll_loss, lmap, load_json


logging.basicConfig(level=logging.DEBUG)

__magic_name__       =     logging.getLogger()
__magic_name__       =     torch.cuda.is_available()
__magic_name__       =     {
    '''max_tokens_per_batch''': None,
    '''supervise_forward''': True,
    '''normalize_hidden''': True,
    '''label_smoothing''': 0.2,
    '''eval_max_gen_length''': None,
    '''eval_beams''': 1,
    '''val_metric''': '''loss''',
    '''save_top_k''': 1,
    '''adafactor''': True,
    '''early_stopping_patience''': 2,
    '''logger_name''': '''default''',
    '''length_penalty''': 0.5,
    '''cache_dir''': '''''',
    '''task''': '''summarization''',
    '''num_workers''': 2,
    '''alpha_hid''': 0,
    '''freeze_embeds''': True,
    '''enc_only''': False,
    '''tgt_suffix''': '''''',
    '''resume_from_checkpoint''': None,
    '''sortish_sampler''': True,
    '''student_decoder_layers''': 1,
    '''val_check_interval''': 1.0,
    '''output_dir''': '''''',
    '''fp16''': False,  # TODO(SS): set this to CUDA_AVAILABLE if ci installs apex or start using native amp
    '''no_teacher''': False,
    '''fp16_opt_level''': '''O1''',
    '''gpus''': 1 if CUDA_AVAILABLE else 0,
    '''n_tpu_cores''': 0,
    '''max_grad_norm''': 1.0,
    '''do_train''': True,
    '''do_predict''': True,
    '''accumulate_grad_batches''': 1,
    '''server_ip''': '''''',
    '''server_port''': '''''',
    '''seed''': 42,
    '''model_name_or_path''': '''sshleifer/bart-tiny-random''',
    '''config_name''': '''''',
    '''tokenizer_name''': '''facebook/bart-large''',
    '''do_lower_case''': False,
    '''learning_rate''': 0.3,
    '''lr_scheduler''': '''linear''',
    '''weight_decay''': 0.0,
    '''adam_epsilon''': 1E-08,
    '''warmup_steps''': 0,
    '''max_epochs''': 1,
    '''train_batch_size''': 2,
    '''eval_batch_size''': 2,
    '''max_source_length''': 12,
    '''max_target_length''': 12,
    '''val_max_target_length''': 12,
    '''test_max_target_length''': 12,
    '''fast_dev_run''': False,
    '''no_cache''': False,
    '''n_train''': -1,
    '''n_val''': -1,
    '''n_test''': -1,
    '''student_encoder_layers''': 1,
    '''freeze_encoder''': False,
    '''auto_scale_batch_size''': False,
    '''overwrite_output_dir''': False,
    '''student''': None,
}





def 		__snake_case						(    _UpperCAmelCase       ,			_UpperCAmelCase				):
					"""simple docstring"""






					lowercase						   =       '\n'.join(_UpperCAmelCase				)
					Path(_UpperCAmelCase				).open('w'				).writelines(_UpperCAmelCase				)


__magic_name__       =     [''' Sam ate lunch today.''', '''Sams lunch ingredients.''']
__magic_name__       =     ['''A very interesting story about what I ate for lunch.''', '''Avocado, celery, turkey, coffee''']
__magic_name__       =     '''patrickvonplaten/t5-tiny-random'''
__magic_name__       =     '''sshleifer/t5-tinier-random'''
__magic_name__       =     '''sshleifer/bart-tiny-random'''
__magic_name__       =     '''sshleifer/tiny-mbart'''
__magic_name__       =     '''sshleifer/tiny-marian-en-de'''
__magic_name__       =     '''stas/tiny-wmt19-en-de'''


__magic_name__       =     logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)
logging.disable(logging.CRITICAL)  # remove noisy download output from tracebacks





def 		__snake_case						(    _UpperCAmelCase				):
					"""simple docstring"""






					for split in ["train", "val", "test"]:
										_dump_articles(os.path.join(_UpperCAmelCase       ,			f"""{split}.source"""				)       ,			_UpperCAmelCase				)
										_dump_articles(os.path.join(_UpperCAmelCase       ,			f"""{split}.target"""				)       ,			_UpperCAmelCase				)
					return tmp_dir


class 	a__       (       _snake_case       ):


					"""simple docstring"""



					@classmethod
					def 				__UpperCAmelCase		(		cls     :Tuple    ):
										logging.disable(logging.CRITICAL    )  # remove noisy download output from tracebacks
										return cls

					@slow
					@require_torch_gpu
					def 				__UpperCAmelCase		(		self     :str    ):
										lowercase						   =       list_models()
										lowercase						   =       'sshleifer'
										lowercase						   =       [x.modelId for x in model_list if x.modelId.startswith(lowercase__    )]
										lowercase						   =       ['sshleifer/blenderbot-3B', 'sshleifer/blenderbot-90M']
										lowercase						   =       []
										for m in model_ids:
															if m in allowed_to_be_broken:
																				continue
															try:
																				AutoConfig.from_pretrained(lowercase__    )
															except Exception:
																				failures.append(lowercase__    )
										assert not failures, F"""The following models could not be loaded through AutoConfig: {failures}"""

					def 				__UpperCAmelCase		(		self     :List[str]    ):
										lowercase						   =       {'student_encoder_layers': 2, 'student_decoder_layers': 1, 'no_teacher': True}
										self._test_distiller_cli(lowercase__    )

					def 				__UpperCAmelCase		(		self     :List[Any]    ):
										lowercase						   =       {
										    'student_encoder_layers': 2,
										    'student_decoder_layers': 1,
										    'max_epochs': 4,
										    'val_check_interval': 0.25,
										    'alpha_hid': 2.0,
										    'model_name_or_path': 'IGNORE_THIS_IT_DOESNT_GET_USED',
										}
										lowercase						   =       self._test_distiller_cli(lowercase__      ,  check_contents=lowercase__    )

										lowercase						   =       list(Path(model.output_dir    ).glob('*.ckpt'    )    )
										self.assertEqual(1      ,  len(lowercase__    )    )
										lowercase						   =       list(Path(model.output_dir    ).glob('**/*.bin'    )    )
										self.assertEqual(len(lowercase__    )      ,  2    )
										lowercase						   =       lmap(str.strip      ,  Path(model.hparams.data_dir    ).joinpath('test.source'    ).open().readlines()    )
										lowercase						   =       tempfile.mktemp()  # XXX: not being cleaned up
										generate_summaries_or_translations(lowercase__      ,  lowercase__      ,  str(model.output_dir / 'best_tfmr'    )    )
										self.assertTrue(Path(lowercase__    ).exists()    )

										lowercase						   =       self.get_auto_remove_tmp_dir()
										convert_pl_to_hf(ckpts[0]      ,  transformer_ckpts[0].parent      ,  lowercase__    )
										assert os.path.exists(os.path.join(lowercase__      ,  'pytorch_model.bin'    )    )

					def 				__UpperCAmelCase		(		self     :Tuple    ):
										lowercase						   =       AutoModelForSeqaSeqLM.from_pretrained(lowercase__    )
										lowercase   ,      lowercase						   =       model.dummy_inputs['input_ids'], model.dummy_inputs['attention_mask']
										lowercase						   =       torch.tensor([[0, 4, 8, 2], [0, 8, 2, 1]]      ,  dtype=torch.long      ,  device=model.device    )
										lowercase						   =       target_ids[:, :-1].contiguous()  # Why this line?
										lowercase						   =       target_ids[:, 1:].clone()  # why clone?
										lowercase						   =       model(
										    lowercase__      ,  attention_mask=lowercase__      ,  decoder_input_ids=lowercase__      ,  labels=lowercase__      ,  use_cache=lowercase__    ).loss

										lowercase						   =       model(lowercase__      ,  attention_mask=lowercase__      ,  decoder_input_ids=lowercase__      ,  use_cache=lowercase__    ).logits

										lowercase						   =       nn.functional.log_softmax(lowercase__      ,  dim=-1    )
										lowercase   ,      lowercase						   =       label_smoothed_nll_loss(
										    lowercase__      ,  lowercase__      ,  0.1      ,  ignore_index=model.config.pad_token_id    )
										with self.assertRaises(lowercase__    ):
															# TODO: understand why this breaks
															self.assertEqual(lowercase__      ,  lowercase__    )

					def 				__UpperCAmelCase		(		self     :str    ):
										lowercase						   =       {
										    'student_encoder_layers': 2,
										    'student_decoder_layers': 1,
										    'num_train_epochs': 4,
										    'val_check_interval': 0.25,
										    'alpha_hid': 2.0,
										    'task': 'translation',
										    'model_name_or_path': 'IGNORE_THIS_IT_DOESNT_GET_USED',
										    'tokenizer_name': MBART_TINY,
										    'teacher': MBART_TINY,
										    'src_lang': 'en_XX',
										    'tgt_lang': 'ro_RO',
										}
										lowercase						   =       self._test_distiller_cli(lowercase__      ,  check_contents=lowercase__    )
										assert model.model.config.model_type == "mbart"

										lowercase						   =       list(Path(model.output_dir    ).glob('*.ckpt'    )    )
										self.assertEqual(1      ,  len(lowercase__    )    )
										lowercase						   =       list(Path(model.output_dir    ).glob('**/*.bin'    )    )
										lowercase						   =       list(Path(model.output_dir    ).glob('best_tfmr/*'    )    )
										assert len(lowercase__    ) > 2
										self.assertEqual(len(lowercase__    )      ,  2    )

					def 				__UpperCAmelCase		(		self     :Tuple    ):
										lowercase						   =       {
										    'student_encoder_layers': 1,
										    'student_decoder_layers': 1,
										    'alpha_hid': 2.0,
										    'teacher': T5_TINY,
										    'model_name_or_path': T5_TINY,
										    'tokenizer_name': T5_TINY,
										}
										self._test_distiller_cli(lowercase__    )

					def 				__UpperCAmelCase		(		self     :str    ):
										lowercase						   =       {
										    'teacher': T5_TINY,
										    'student': T5_TINIER,
										    'model_name_or_path': T5_TINIER,
										    'tokenizer_name': T5_TINIER,
										}
										self._test_distiller_cli(lowercase__    )



					def 				__UpperCAmelCase		(		self     :Optional[Any]      ,  lowercase__     :Optional[Any]      ,  lowercase__     :Dict=True    ):
										lowercase						   =       {
										    'label_smoothing': 0.0,
										    'early_stopping_patience': -1,
										    'train_batch_size': 1,
										    'eval_batch_size': 2,
										    'max_epochs': 2,
										    'alpha_mlm': 0.2,
										    'alpha_ce': 0.8,
										    'do_predict': True,
										    'model_name_or_path': 'sshleifer/tinier_bart',
										    'teacher': CHEAP_ARGS['model_name_or_path'],
										    'val_check_interval': 0.5,
										}
										default_updates.update(lowercase__    )
										lowercase						   =       CHEAP_ARGS.copy()
										lowercase						   =       make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()    )
										lowercase						   =       self.get_auto_remove_tmp_dir()

										args_d.update(data_dir=lowercase__      ,  output_dir=lowercase__      ,  **lowercase__    )
										lowercase						   =       distill_main(argparse.Namespace(**lowercase__    )    )
										if not check_contents:
															return model
										lowercase						   =       os.listdir(lowercase__    )
										lowercase						   =       {os.path.basename(lowercase__    ) for p in contents}
										lowercase						   =       [p for p in contents if p.endswith('ckpt'    )]
										assert len(lowercase__    ) > 0

										self.assertIn('test_generations.txt'      ,  lowercase__    )
										self.assertIn('test_results.txt'      ,  lowercase__    )

										lowercase						   =       load_json(model.metrics_save_path    )
										lowercase						   =       metrics['val'][-1]
										self.assertGreaterEqual(last_step_stats['val_avg_gen_time']      ,  0.01    )
										self.assertGreaterEqual(1.0      ,  last_step_stats['val_avg_gen_time']    )
										self.assertIsInstance(last_step_stats[F"""val_avg_{model.val_metric}"""]      ,  lowercase__    )
										lowercase						   =       int(args_d['max_epochs'] * (1 / args_d['val_check_interval']) + 1    )
										self.assertEqual(len(metrics['val']    )      ,  lowercase__    )
										self.assertEqual(len(metrics['test']    )      ,  1    )
										return model




class 	a__       (       _snake_case       ):


					"""simple docstring"""



					@parameterized.expand(
					    [T5_TINY, BART_TINY, MBART_TINY, MARIAN_TINY, FSMT_TINY]      ,  )
					def 				__UpperCAmelCase		(		self     :Union[str, Any]      ,  lowercase__     :Union[str, Any]    ):
										lowercase						   =       CHEAP_ARGS.copy()
										lowercase						   =       'translation' if model in [MBART_TINY, MARIAN_TINY, FSMT_TINY] else 'summarization'
										lowercase						   =       0.1 if task == 'translation' else 0

										lowercase						   =       make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()    )
										lowercase						   =       self.get_auto_remove_tmp_dir()
										args_d.update(
										    data_dir=lowercase__      ,  model_name_or_path=lowercase__      ,  tokenizer_name=lowercase__      ,  train_batch_size=2      ,  eval_batch_size=2      ,  output_dir=lowercase__      ,  do_predict=lowercase__      ,  task=lowercase__      ,  src_lang='en_XX'      ,  tgt_lang='ro_RO'      ,  freeze_encoder=lowercase__      ,  freeze_embeds=lowercase__      ,  )
										assert "n_train" in args_d
										lowercase						   =       argparse.Namespace(**lowercase__    )
										lowercase						   =       main(lowercase__    )

										lowercase						   =       module.model.get_input_embeddings()
										assert not input_embeds.weight.requires_grad
										if model == T5_TINY:
															lowercase						   =       module.model.lm_head
															assert not lm_head.weight.requires_grad
															assert (lm_head.weight == input_embeds.weight).all().item()
										elif model == FSMT_TINY:
															lowercase						   =       module.model.model
															lowercase						   =       fsmt.decoder.embed_positions
															assert not embed_pos.weight.requires_grad
															assert not fsmt.decoder.embed_tokens.weight.requires_grad
															# check that embeds are not the same
															assert fsmt.decoder.embed_tokens != fsmt.encoder.embed_tokens
										else:
															lowercase						   =       module.model.model
															lowercase						   =       bart.decoder.embed_positions
															assert not embed_pos.weight.requires_grad
															assert not bart.shared.weight.requires_grad
															# check that embeds are the same
															assert bart.decoder.embed_tokens == bart.encoder.embed_tokens
															assert bart.decoder.embed_tokens == bart.shared

										lowercase						   =       load_json(module.output_dir / 'text_batch.json'    )
										assert isinstance(lowercase__      ,  lowercase__    )
										assert len(lowercase__    ) >= 4

					def 				__UpperCAmelCase		(		self     :Optional[int]    ):
										lowercase						   =       CHEAP_ARGS.copy()

										lowercase						   =       'summarization'
										lowercase						   =       make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()    )

										args_d.update(
										    data_dir=lowercase__      ,  tokenizer_name=lowercase__      ,  train_batch_size=2      ,  eval_batch_size=2      ,  do_predict=lowercase__      ,  task=lowercase__      ,  src_lang='en_XX'      ,  tgt_lang='ro_RO'      ,  freeze_encoder=lowercase__      ,  freeze_embeds=lowercase__      ,  )

										# test models whose config includes the extra_model_args
										lowercase						   =       BART_TINY
										lowercase						   =       self.get_auto_remove_tmp_dir()
										lowercase						   =       args_d.copy()
										args_da.update(
										    model_name_or_path=lowercase__      ,  output_dir=lowercase__      ,  )
										lowercase						   =       ('encoder_layerdrop', 'decoder_layerdrop', 'dropout', 'attention_dropout')
										for p in extra_model_params:
															lowercase						   =       0.5
										lowercase						   =       argparse.Namespace(**lowercase__    )
										lowercase						   =       main(lowercase__    )
										for p in extra_model_params:
															assert getattr(model.config      ,  lowercase__    ) == 0.5, F"""failed to override the model config for param {p}"""

										# test models whose config doesn't include the extra_model_args
										lowercase						   =       T5_TINY
										lowercase						   =       self.get_auto_remove_tmp_dir()
										lowercase						   =       args_d.copy()
										args_da.update(
										    model_name_or_path=lowercase__      ,  output_dir=lowercase__      ,  )
										lowercase						   =       'encoder_layerdrop'
										lowercase						   =       0.5
										lowercase						   =       argparse.Namespace(**lowercase__    )
										with pytest.raises(lowercase__    ) as excinfo:
															lowercase						   =       main(lowercase__    )
										assert str(excinfo.value    ) == F"""model config doesn't have a `{unsupported_param}` attribute"""



					def 				__UpperCAmelCase		(		self     :Union[str, Any]    ):
										lowercase						   =       CHEAP_ARGS.copy()

										lowercase						   =       'summarization'
										lowercase						   =       make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()    )

										lowercase						   =       BART_TINY
										lowercase						   =       self.get_auto_remove_tmp_dir()

										args_d.update(
										    data_dir=lowercase__      ,  model_name_or_path=lowercase__      ,  output_dir=lowercase__      ,  tokenizer_name=lowercase__      ,  train_batch_size=2      ,  eval_batch_size=2      ,  do_predict=lowercase__      ,  task=lowercase__      ,  src_lang='en_XX'      ,  tgt_lang='ro_RO'      ,  freeze_encoder=lowercase__      ,  freeze_embeds=lowercase__      ,  )

										# emulate finetune.py
										lowercase						   =       argparse.ArgumentParser()
										lowercase						   =       pl.Trainer.add_argparse_args(lowercase__    )
										lowercase						   =       SummarizationModule.add_model_specific_args(lowercase__      ,  os.getcwd()    )
										lowercase						   =       {'--help': True}

										# --help test
										with pytest.raises(lowercase__    ) as excinfo:
															with CaptureStdout() as cs:
																				lowercase						   =       parser.parse_args(lowercase__    )
															assert False, "--help is expected to sys.exit"
										assert excinfo.type == SystemExit
										lowercase						   =       lightning_base.arg_to_scheduler_metavar
										assert expected in cs.out, "--help is expected to list the supported schedulers"

										# --lr_scheduler=non_existing_scheduler test
										lowercase						   =       'non_existing_scheduler'
										lowercase						   =       {F"""--lr_scheduler={unsupported_param}"""}
										with pytest.raises(lowercase__    ) as excinfo:
															with CaptureStderr() as cs:
																				lowercase						   =       parser.parse_args(lowercase__    )
															assert False, "invalid argument is expected to sys.exit"
										assert excinfo.type == SystemExit
										lowercase						   =       F"""invalid choice: '{unsupported_param}'"""
										assert expected in cs.err, F"""should have bailed on invalid choice of scheduler {unsupported_param}"""

										# --lr_scheduler=existing_scheduler test
										lowercase						   =       'cosine'
										lowercase						   =       args_d.copy()
										lowercase						   =       supported_param
										lowercase						   =       argparse.Namespace(**lowercase__    )
										lowercase						   =       main(lowercase__    )
										assert (
										    getattr(model.hparams      ,  'lr_scheduler'    ) == supported_param
										), F"""lr_scheduler={supported_param} shouldn't fail"""











from typing import List, Optional, Tuple, Union

from . import logging
from .import_utils import is_torch_available, is_torch_version


if is_torch_available():
		import torch

__magic_name__       =     logging.get_logger(__name__)  # pylint: disable=invalid-name

try:
		from torch._dynamo import allow_in_graph as maybe_allow_in_graph
except (ImportError, ModuleNotFoundError):





		def 		__snake_case						(    cls				):
							"""simple docstring"""






							return cls





def 		__snake_case						(    _UpperCAmelCase       ,			_UpperCAmelCase = None       ,			_UpperCAmelCase = None       ,			_UpperCAmelCase = None       ,			_UpperCAmelCase = None       ,			):
					"""simple docstring"""






					lowercase						   =       device
					lowercase						   =       shape[0]

					lowercase						   =       layout or torch.strided
					lowercase						   =       device or torch.device('cpu'				)

					if generator is not None:
										lowercase						   =       generator.device.type if not isinstance(_UpperCAmelCase       ,			_UpperCAmelCase				) else generator[0].device.type
										if gen_device_type != device.type and gen_device_type == "cpu":
															lowercase						   =       'cpu'
															if device != "mps":
																				logger.info(
																				    f"""The passed generator was created on 'cpu' even though a tensor on {device} was expected."""
																				    f""" Tensors will be created on 'cpu' and then moved to {device}. Note that one can probably"""
																				    f""" slighly speed up this function by passing a generator that was created on the {device} device."""				)
										elif gen_device_type != device.type and gen_device_type == "cuda":
															raise ValueError(f"""Cannot generate a {device} tensor from a generator of type {gen_device_type}."""				)

					if isinstance(_UpperCAmelCase       ,			_UpperCAmelCase				):
										lowercase						   =       (1,) + shape[1:]
										lowercase						   =       [
										    torch.randn(_UpperCAmelCase       ,			generator=generator[i]       ,			device=_UpperCAmelCase       ,			dtype=_UpperCAmelCase       ,			layout=_UpperCAmelCase				)
										    for i in range(_UpperCAmelCase				)
										]
										lowercase						   =       torch.cat(_UpperCAmelCase       ,			dim=0				).to(_UpperCAmelCase				)
					else:
										lowercase						   =       torch.randn(_UpperCAmelCase       ,			generator=_UpperCAmelCase       ,			device=_UpperCAmelCase       ,			dtype=_UpperCAmelCase       ,			layout=_UpperCAmelCase				).to(_UpperCAmelCase				)

					return latents





def 		__snake_case						(    _UpperCAmelCase				):
					"""simple docstring"""






					if is_torch_version('<'       ,			'2.0.0'				) or not hasattr(_UpperCAmelCase       ,			'_dynamo'				):
										return False
					return isinstance(_UpperCAmelCase       ,			torch._dynamo.eval_frame.OptimizedModule				)






