






'''simple docstring'''




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)

lowerCamelCase_							      =      logging.getLogger()
lowerCamelCase_							      =      torch.cuda.is_available()
lowerCamelCase_							      =      {
    '''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       __lowercase		(    __lowercase     ,						__lowercase ) ->		str:

      '''simple docstring'''


      _A      =			"\n".join(__lowercase )
      Path(__lowercase ).open("w" ).writelines(__lowercase )


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


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

def       __lowercase		(    __lowercase ) ->		int:

      '''simple docstring'''


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

class      _UpperCAmelCase    (  snake_case_		):




   """simple docstring"""

   @classmethod
   def 		lowerCAmelCase				(	cls				:			List[Any]	):





         '''simple docstring'''





         logging.disable(logging.CRITICAL	)  # remove noisy download output from tracebacks
         return cls

   @slow
   @require_torch_gpu
   def 		lowerCAmelCase				(	self				:			Tuple	):





         '''simple docstring'''





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

   def 		lowerCAmelCase				(	self				:			Optional[int]	):





         '''simple docstring'''





         _A      =			{"student_encoder_layers": 2, "student_decoder_layers": 1, "no_teacher": True}
         self._test_distiller_cli(__UpperCAmelCase	)

   def 		lowerCAmelCase				(	self				:			Any	):





         '''simple docstring'''





         _A      =			{
             "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",
         }
         _A      =			self._test_distiller_cli(__UpperCAmelCase   ,   check_contents=__UpperCAmelCase	)

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

         _A      =			self.get_auto_remove_tmp_dir()
         convert_pl_to_hf(ckpts[0]   ,   transformer_ckpts[0].parent   ,   __UpperCAmelCase	)
         assert os.path.exists(os.path.join(__UpperCAmelCase   ,   "pytorch_model.bin"	)	)

   def 		lowerCAmelCase				(	self				:			List[str]	):





         '''simple docstring'''





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

         _A      =			model(__UpperCAmelCase   ,   attention_mask=__UpperCAmelCase   ,   decoder_input_ids=__UpperCAmelCase   ,   use_cache=__UpperCAmelCase	).logits

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

   def 		lowerCAmelCase				(	self				:			Dict	):





         '''simple docstring'''





         _A      =			{
             "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",
         }
         _A      =			self._test_distiller_cli(__UpperCAmelCase   ,   check_contents=__UpperCAmelCase	)
         assert model.model.config.model_type == "mbart"

         _A      =			list(Path(model.output_dir	).glob("*.ckpt"	)	)
         self.assertEqual(1   ,   len(__UpperCAmelCase	)	)
         _A      =			list(Path(model.output_dir	).glob("**/*.bin"	)	)
         _A      =			list(Path(model.output_dir	).glob("best_tfmr/*"	)	)
         assert len(__UpperCAmelCase	) > 2
         self.assertEqual(len(__UpperCAmelCase	)   ,   2	)

   def 		lowerCAmelCase				(	self				:			Union[str, Any]	):





         '''simple docstring'''





         _A      =			{
             "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(__UpperCAmelCase	)

   def 		lowerCAmelCase				(	self				:			List[str]	):





         '''simple docstring'''





         _A      =			{
             "teacher": T5_TINY,
             "student": T5_TINIER,
             "model_name_or_path": T5_TINIER,
             "tokenizer_name": T5_TINIER,
         }
         self._test_distiller_cli(__UpperCAmelCase	)

   def 		lowerCAmelCase				(	self				:			Tuple   ,   __UpperCAmelCase				:			List[str]   ,   __UpperCAmelCase				:			int=True	):





         '''simple docstring'''





         _A      =			{
             "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(__UpperCAmelCase	)
         _A      =			CHEAP_ARGS.copy()
         _A      =			make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()	)
         _A      =			self.get_auto_remove_tmp_dir()

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

         self.assertIn("test_generations.txt"   ,   __UpperCAmelCase	)
         self.assertIn("test_results.txt"   ,   __UpperCAmelCase	)

         _A      =			load_json(model.metrics_save_path	)
         _A      =			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}''']   ,   __UpperCAmelCase	)
         _A      =			int(args_d["max_epochs"] * (1 / args_d["val_check_interval"]) + 1	)
         self.assertEqual(len(metrics["val"]	)   ,   __UpperCAmelCase	)
         self.assertEqual(len(metrics["test"]	)   ,   1	)
         return model





class      _UpperCAmelCase    (  snake_case_		):




   """simple docstring"""

   @parameterized.expand(
       [T5_TINY, BART_TINY, MBART_TINY, MARIAN_TINY, FSMT_TINY]   ,   )
   def 		lowerCAmelCase				(	self				:			Optional[int]   ,   __UpperCAmelCase				:			List[Any]	):





         '''simple docstring'''





         _A      =			CHEAP_ARGS.copy()
         _A      =			"translation" if model in [MBART_TINY, MARIAN_TINY, FSMT_TINY] else "summarization"
         _A      =			0.1 if task == "translation" else 0

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

         _A      =			module.model.get_input_embeddings()
         assert not input_embeds.weight.requires_grad
         if model == T5_TINY:
               _A      =			module.model.lm_head
               assert not lm_head.weight.requires_grad
               assert (lm_head.weight == input_embeds.weight).all().item()
         elif model == FSMT_TINY:
               _A      =			module.model.model
               _A      =			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:
               _A      =			module.model.model
               _A      =			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

         _A      =			load_json(module.output_dir / "text_batch.json"	)
         assert isinstance(__UpperCAmelCase   ,   __UpperCAmelCase	)
         assert len(__UpperCAmelCase	) >= 4

   def 		lowerCAmelCase				(	self				:			Union[str, Any]	):





         '''simple docstring'''





         _A      =			CHEAP_ARGS.copy()

         _A      =			"summarization"
         _A      =			make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()	)

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

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

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

   def 		lowerCAmelCase				(	self				:			Optional[Any]	):





         '''simple docstring'''





         _A      =			CHEAP_ARGS.copy()

         _A      =			"summarization"
         _A      =			make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()	)

         _A      =			BART_TINY
         _A      =			self.get_auto_remove_tmp_dir()

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

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

         # --help test
         with pytest.raises(__UpperCAmelCase	) as excinfo:
               with CaptureStdout() as cs:
                     _A      =			parser.parse_args(__UpperCAmelCase	)
               assert False, "--help is expected to sys.exit"
         assert excinfo.type == SystemExit
         _A      =			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
         _A      =			"non_existing_scheduler"
         _A      =			{f'''--lr_scheduler={unsupported_param}'''}
         with pytest.raises(__UpperCAmelCase	) as excinfo:
               with CaptureStderr() as cs:
                     _A      =			parser.parse_args(__UpperCAmelCase	)
               assert False, "invalid argument is expected to sys.exit"
         assert excinfo.type == SystemExit
         _A      =			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
         _A      =			"cosine"
         _A      =			args_d.copy()
         _A      =			supported_param
         _A      =			argparse.Namespace(**__UpperCAmelCase	)
         _A      =			main(__UpperCAmelCase	)
         assert (
             getattr(model.hparams   ,   "lr_scheduler"	) == supported_param
         ), f'''lr_scheduler={supported_param} shouldn\'t fail'''













'''simple docstring'''




def       __lowercase		(    __lowercase ) ->		int:

      '''simple docstring'''


      _A      =			str(__lowercase ).strip()
      if not bin_string:
            raise ValueError("Empty string was passed to the function" )
      _A      =			bin_string[0] == "-"
      if is_negative:
            _A      =			bin_string[1:]
      if not all(char in "01" for char in bin_string ):
            raise ValueError("Non-binary value was passed to the function" )
      _A      =			0
      for char in bin_string:
            _A      =			2 * decimal_number + int(__lowercase )
      return -decimal_number if is_negative else decimal_number


if __name__ == "__main__":
    from doctest import testmod

    testmod()





