





def  __UpperCamelCase    (						snake_case      )    ->		tuple[int, int]:


   '''simple docstring'''







   if len(snake_case      ) == 1:
      return a[0], a[0]
   else:
      # split A into half.
      __A    =  a[: len(snake_case      ) // 2]
      __A    =  a[len(snake_case      ) // 2 :]

      # 2 sub problems, 1/2 of original size.
      __A   ,		__A    =  max_difference(snake_case      )
      __A   ,		__A    =  max_difference(snake_case      )

      # get min of first and max of second
      # linear time
      __A    =  min(snake_case      )
      __A    =  max(snake_case      )

      # 3 cases, either (small1, big1),
      # (min_first, max_second), (small2, big2)
      # constant comparisons
      if biga - smalla > max_second - min_first and biga - smalla > biga - smalla:
         return smalla, biga
      elif biga - smalla > max_second - min_first:
         return smalla, biga
      else:
         return min_first, max_second


if __name__ == "__main__":
    import doctest

    doctest.testmod()








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)

_UpperCamelCase       :				Tuple           =			logging.getLogger()
_UpperCamelCase       :				Optional[int]           =			torch.cuda.is_available()
_UpperCamelCase       :				List[str]           =			{
    """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""": 4_2,
    """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-0_8,
    """warmup_steps""": 0,
    """max_epochs""": 1,
    """train_batch_size""": 2,
    """eval_batch_size""": 2,
    """max_source_length""": 1_2,
    """max_target_length""": 1_2,
    """val_max_target_length""": 1_2,
    """test_max_target_length""": 1_2,
    """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  __UpperCamelCase    (						snake_case							,			snake_case      )    ->		List[Any]:


   '''simple docstring'''







   __A    =  '''\n'''.join(snake_case      )
   Path(snake_case      ).open('''w'''      ).writelines(snake_case      )


_UpperCamelCase       :				List[Any]           =			[""" Sam ate lunch today.""", """Sams lunch ingredients."""]
_UpperCamelCase       :				int           =			["""A very interesting story about what I ate for lunch.""", """Avocado, celery, turkey, coffee"""]
_UpperCamelCase       :				List[str]           =			"""patrickvonplaten/t5-tiny-random"""
_UpperCamelCase       :				int           =			"""sshleifer/t5-tinier-random"""
_UpperCamelCase       :				Dict           =			"""sshleifer/bart-tiny-random"""
_UpperCamelCase       :				int           =			"""sshleifer/tiny-mbart"""
_UpperCamelCase       :				Optional[int]           =			"""sshleifer/tiny-marian-en-de"""
_UpperCamelCase       :				List[Any]           =			"""stas/tiny-wmt19-en-de"""


_UpperCamelCase       :				Optional[Any]           =			logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)
logging.disable(logging.CRITICAL)  # remove noisy download output from tracebacks
def  __UpperCamelCase    (						snake_case      )    ->		Optional[int]:


   '''simple docstring'''







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







class 		_lowerCAmelCase(		_a):

  """simple docstring"""







  @classmethod
  def 						SCREAMING_SNAKE_CASE__			(   cls      )->   List[str]:
     logging.disable(logging.CRITICAL      )  # remove noisy download output from tracebacks
     return cls

  @slow
  @require_torch_gpu
  def 						SCREAMING_SNAKE_CASE__			(   self      )->   Optional[int]:
     __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 						SCREAMING_SNAKE_CASE__			(   self      )->   List[str]:
     __A    =  {'''student_encoder_layers''': 2, '''student_decoder_layers''': 1, '''no_teacher''': True}
     self._test_distiller_cli(UpperCAmelCase      )

  def 						SCREAMING_SNAKE_CASE__			(   self      )->   int:
     __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 						SCREAMING_SNAKE_CASE__			(   self      )->   List[str]:
     __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 						SCREAMING_SNAKE_CASE__			(   self      )->   Dict:
     __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 						SCREAMING_SNAKE_CASE__			(   self      )->   str:
     __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 						SCREAMING_SNAKE_CASE__			(   self      )->   Tuple:
     __A    =  {
         '''teacher''': T5_TINY,
         '''student''': T5_TINIER,
         '''model_name_or_path''': T5_TINIER,
         '''tokenizer_name''': T5_TINIER,
     }
     self._test_distiller_cli(UpperCAmelCase      )





  def 						SCREAMING_SNAKE_CASE__			(   self    ,							UpperCAmelCase    ,							UpperCAmelCase=True      )->   Tuple:
     __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 		_lowerCAmelCase(		_a):

  """simple docstring"""







  @parameterized.expand(
      [T5_TINY, BART_TINY, MBART_TINY, MARIAN_TINY, FSMT_TINY]    ,							)
  def 						SCREAMING_SNAKE_CASE__			(   self    ,							UpperCAmelCase      )->   List[str]:
     __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 						SCREAMING_SNAKE_CASE__			(   self      )->   str:
     __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 						SCREAMING_SNAKE_CASE__			(   self      )->   int:
     __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"

