




'''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)

A							:				Union[str, Any]			     =	logging.getLogger()
A							:				str			     =	torch.cuda.is_available()
A							:				Optional[Any]			     =	{
    """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_  (       a__		:     Path   ,a__		:     list      ):



    """simple docstring"""




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


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


A							:				Tuple			     =	logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)
logging.disable(logging.CRITICAL)  # remove noisy download output from tracebacks
def     snake_case_  (       a__		:     str      ):



    """simple docstring"""




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






class 	SCREAMING_SNAKE_CASE(    __A					):

       @classmethod
       def    snake_case__      (   cls					)   -> Optional[int]:

           """simple docstring"""



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

       @slow
       @require_torch_gpu
       def    snake_case__      (   self					)   -> Tuple:

           """simple docstring"""



           __lowercase          =						list_models()
           __lowercase          =						"""sshleifer"""
           __lowercase          =						[x.modelId for x in model_list if x.modelId.startswith(lowerCamelCase__					)]
           __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(lowerCamelCase__					)
               except Exception:
                   failures.append(lowerCamelCase__					)
           assert not failures, F'The following models could not be loaded through AutoConfig: {failures}'

       def    snake_case__      (   self					)   -> Any:

           """simple docstring"""



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

       def    snake_case__      (   self					)   -> Any:

           """simple docstring"""



           __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(lowerCamelCase__ ,		check_contents=lowerCamelCase__					)

           __lowercase          =						list(Path(model.output_dir					).glob("""*.ckpt"""					)					)
           self.assertEqual(1 ,		len(lowerCamelCase__					)					)
           __lowercase          =						list(Path(model.output_dir					).glob("""**/*.bin"""					)					)
           self.assertEqual(len(lowerCamelCase__					) ,		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(lowerCamelCase__ ,		lowerCamelCase__ ,		str(model.output_dir / """best_tfmr"""					)					)
           self.assertTrue(Path(lowerCamelCase__					).exists()					)

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

       def    snake_case__      (   self					)   -> List[Any]:

           """simple docstring"""



           __lowercase          =						AutoModelForSeqaSeqLM.from_pretrained(lowerCamelCase__					)
           __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(
               lowerCamelCase__ ,		attention_mask=lowerCamelCase__ ,		decoder_input_ids=lowerCamelCase__ ,		labels=lowerCamelCase__ ,		use_cache=lowerCamelCase__					).loss

           __lowercase          =						model(lowerCamelCase__ ,		attention_mask=lowerCamelCase__ ,		decoder_input_ids=lowerCamelCase__ ,		use_cache=lowerCamelCase__					).logits

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

       def    snake_case__      (   self					)   -> Tuple:

           """simple docstring"""



           __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(lowerCamelCase__ ,		check_contents=lowerCamelCase__					)
           assert model.model.config.model_type == "mbart"

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

       def    snake_case__      (   self					)   -> Dict:

           """simple docstring"""



           __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(lowerCamelCase__					)

       def    snake_case__      (   self					)   -> List[Any]:

           """simple docstring"""



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




       def    snake_case__      (   self ,		lowerCamelCase__ ,		lowerCamelCase__=True					)   -> int:

           """simple docstring"""



           __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(lowerCamelCase__					)
           __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=lowerCamelCase__ ,		output_dir=lowerCamelCase__ ,		**lowerCamelCase__					)
           __lowercase          =						distill_main(argparse.Namespace(**lowerCamelCase__					)					)
           if not check_contents:
               return model
           __lowercase          =						os.listdir(lowerCamelCase__					)
           __lowercase          =						{os.path.basename(lowerCamelCase__					) for p in contents}
           __lowercase          =						[p for p in contents if p.endswith("""ckpt"""					)]
           assert len(lowerCamelCase__					) > 0

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

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






class 	SCREAMING_SNAKE_CASE(    __A					):

       @parameterized.expand(
           [T5_TINY, BART_TINY, MBART_TINY, MARIAN_TINY, FSMT_TINY] ,		)
       def    snake_case__      (   self ,		lowerCamelCase__					)   -> str:

           """simple docstring"""



           __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=lowerCamelCase__ ,		model_name_or_path=lowerCamelCase__ ,		tokenizer_name=lowerCamelCase__ ,		train_batch_size=2 ,		eval_batch_size=2 ,		output_dir=lowerCamelCase__ ,		do_predict=lowerCamelCase__ ,		task=lowerCamelCase__ ,		src_lang="""en_XX""" ,		tgt_lang="""ro_RO""" ,		freeze_encoder=lowerCamelCase__ ,		freeze_embeds=lowerCamelCase__ ,		)
           assert "n_train" in args_d
           __lowercase          =						argparse.Namespace(**lowerCamelCase__					)
           __lowercase          =						main(lowerCamelCase__					)

           __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(lowerCamelCase__ ,		lowerCamelCase__					)
           assert len(lowerCamelCase__					) >= 4

       def    snake_case__      (   self					)   -> str:

           """simple docstring"""



           __lowercase          =						CHEAP_ARGS.copy()

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

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

           # 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=lowerCamelCase__ ,		output_dir=lowerCamelCase__ ,		)
           __lowercase          =						("""encoder_layerdrop""", """decoder_layerdrop""", """dropout""", """attention_dropout""")
           for p in extra_model_params:
               __lowercase          =						0.5
           __lowercase          =						argparse.Namespace(**lowerCamelCase__					)
           __lowercase          =						main(lowerCamelCase__					)
           for p in extra_model_params:
               assert getattr(model.config ,		lowerCamelCase__					) == 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=lowerCamelCase__ ,		output_dir=lowerCamelCase__ ,		)
           __lowercase          =						"""encoder_layerdrop"""
           __lowercase          =						0.5
           __lowercase          =						argparse.Namespace(**lowerCamelCase__					)
           with pytest.raises(lowerCamelCase__					) as excinfo:
               __lowercase          =						main(lowerCamelCase__					)
           assert str(excinfo.value					) == F'model config doesn\'t have a `{unsupported_param}` attribute'




       def    snake_case__      (   self					)   -> Optional[Any]:

           """simple docstring"""



           __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=lowerCamelCase__ ,		model_name_or_path=lowerCamelCase__ ,		output_dir=lowerCamelCase__ ,		tokenizer_name=lowerCamelCase__ ,		train_batch_size=2 ,		eval_batch_size=2 ,		do_predict=lowerCamelCase__ ,		task=lowerCamelCase__ ,		src_lang="""en_XX""" ,		tgt_lang="""ro_RO""" ,		freeze_encoder=lowerCamelCase__ ,		freeze_embeds=lowerCamelCase__ ,		)

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

           # --help test
           with pytest.raises(lowerCamelCase__					) as excinfo:
               with CaptureStdout() as cs:
                   __lowercase          =						parser.parse_args(lowerCamelCase__					)
               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(lowerCamelCase__					) as excinfo:
               with CaptureStderr() as cs:
                   __lowercase          =						parser.parse_args(lowerCamelCase__					)
               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(**lowerCamelCase__					)
           __lowercase          =						main(lowerCamelCase__					)
           assert (
               getattr(model.hparams ,		"""lr_scheduler"""					) == supported_param
           ), F'lr_scheduler={supported_param} shouldn\'t fail'










'''simple docstring'''
# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from ..models.auto import AutoModelForSeqaSeqLM, AutoTokenizer
from .base import PipelineTool


A							:				Dict			     =	"""Here is a text containing a lot of information: '''{text}'''.

Can you answer this question about the text: '{question}'"""






class 	SCREAMING_SNAKE_CASE(    __A					):
       snake_case_			:					List[Any]   = """google/flan-t5-base"""
       snake_case_			:					int   = (
           """This is a tool that answers questions related to a text. It takes two arguments named `text`, which is the """
           """text where to find the answer, and `question`, which is the question, and returns the answer to the question."""
       )
       snake_case_			:					Optional[int]   = """text_qa"""
       snake_case_			:					List[Any]   = AutoTokenizer
       snake_case_			:					List[str]   = AutoModelForSeqaSeqLM

       snake_case_			:					int   = ["""text""", """text"""]
       snake_case_			:					Dict   = ["""text"""]

       def    snake_case__      (   self ,		lowerCamelCase__ ,		lowerCamelCase__					)   -> Optional[int]:

           """simple docstring"""



           __lowercase          =						QA_PROMPT.format(text=lowerCamelCase__ ,		question=lowerCamelCase__					)
           return self.pre_processor(lowerCamelCase__ ,		return_tensors="""pt"""					)

       def    snake_case__      (   self ,		lowerCamelCase__					)   -> int:

           """simple docstring"""



           __lowercase          =						self.model.generate(**lowerCamelCase__					)

           __lowercase    ,__lowercase          =						inputs["""input_ids"""].shape
           __lowercase          =						output_ids.shape[0]

           return output_ids.reshape(lowerCamelCase__ ,		out_b // in_b ,		*output_ids.shape[1:]					)[0][0]




       def    snake_case__      (   self ,		lowerCamelCase__					)   -> Tuple:

           """simple docstring"""



           return self.pre_processor.decode(lowerCamelCase__ ,		skip_special_tokens=lowerCamelCase__ ,		clean_up_tokenization_spaces=lowerCamelCase__					)




