



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

UpperCAmelCase:   Optional[int]       			=							logging.getLogger()
UpperCAmelCase:   Dict       			=							torch.cuda.is_available()
UpperCAmelCase:   List[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       __SCREAMING_SNAKE_CASE      (   __UpperCAmelCase					,							__UpperCAmelCase	):
		_lowercase							:      Optional[Any]			=					"""\n""".join(__UpperCAmelCase	)
		Path(__UpperCAmelCase	).open("""w"""	).writelines(__UpperCAmelCase	)


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


UpperCAmelCase:   Optional[int]       			=							logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)
logging.disable(logging.CRITICAL)  # remove noisy download output from tracebacks


def       __SCREAMING_SNAKE_CASE      (   __UpperCAmelCase	):
		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    UpperCamelCase     (			snake_case    ):




	"""simple docstring"""






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






	@slow
	@require_torch_gpu
	def 					lowerCamelCase__ (			self					):
			_lowercase							:      Any			=					list_models()
			_lowercase							:      Optional[Any]			=					"""sshleifer"""
			_lowercase							:      Dict			=					[x.modelId for x in model_list if x.modelId.startswith(UpperCAmelCase_					)]
			_lowercase							:      int			=					["""sshleifer/blenderbot-3B""", """sshleifer/blenderbot-90M"""]
			_lowercase							:      Union[str, Any]			=					[]
			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					):
			_lowercase							:      Any			=					{"""student_encoder_layers""": 2, """student_decoder_layers""": 1, """no_teacher""": True}
			self._test_distiller_cli(UpperCAmelCase_					)






	def 					lowerCamelCase__ (			self					):
			_lowercase							:      Optional[Any]			=					{
			    """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							:      Union[str, Any]			=					self._test_distiller_cli(UpperCAmelCase_							,check_contents=UpperCAmelCase_					)

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

			_lowercase							:      Union[str, Any]			=					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					):
			_lowercase							:      int			=					AutoModelForSeqaSeqLM.from_pretrained(UpperCAmelCase_					)
			_lowercase			,   _lowercase							:      str			=					model.dummy_inputs["""input_ids"""], model.dummy_inputs["""attention_mask"""]
			_lowercase							:      Dict			=					torch.tensor([[0, 4, 8, 2], [0, 8, 2, 1]]							,dtype=torch.long							,device=model.device					)
			_lowercase							:      Union[str, Any]			=					target_ids[:, :-1].contiguous()  # Why this line?
			_lowercase							:      str			=					target_ids[:, 1:].clone()  # why clone?
			_lowercase							:      Tuple			=					model(
			    UpperCAmelCase_							,attention_mask=UpperCAmelCase_							,decoder_input_ids=UpperCAmelCase_							,labels=UpperCAmelCase_							,use_cache=UpperCAmelCase_					).loss

			_lowercase							:      Tuple			=					model(UpperCAmelCase_							,attention_mask=UpperCAmelCase_							,decoder_input_ids=UpperCAmelCase_							,use_cache=UpperCAmelCase_					).logits

			_lowercase							:      Union[str, Any]			=					nn.functional.log_softmax(UpperCAmelCase_							,dim=-1					)
			_lowercase			,   _lowercase							:      Optional[Any]			=					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					):
			_lowercase							:      Optional[Any]			=					{
			    """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							:      str			=					self._test_distiller_cli(UpperCAmelCase_							,check_contents=UpperCAmelCase_					)
			assert model.model.config.model_type == "mbart"

			_lowercase							:      Optional[int]			=					list(Path(model.output_dir					).glob("""*.ckpt"""					)					)
			self.assertEqual(1							,len(UpperCAmelCase_					)					)
			_lowercase							:      int			=					list(Path(model.output_dir					).glob("""**/*.bin"""					)					)
			_lowercase							:      Optional[int]			=					list(Path(model.output_dir					).glob("""best_tfmr/*"""					)					)
			assert len(UpperCAmelCase_					) > 2
			self.assertEqual(len(UpperCAmelCase_					)							,2					)






	def 					lowerCamelCase__ (			self					):
			_lowercase							:      Union[str, Any]			=					{
			    """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					):
			_lowercase							:      Optional[Any]			=					{
			    """teacher""": T5_TINY,
			    """student""": T5_TINIER,
			    """model_name_or_path""": T5_TINIER,
			    """tokenizer_name""": T5_TINIER,
			}
			self._test_distiller_cli(UpperCAmelCase_					)






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_							,UpperCAmelCase_=True					):
			_lowercase							:      List[str]			=					{
			    """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_					)
			_lowercase							:      dict			=					CHEAP_ARGS.copy()
			_lowercase							:      List[str]			=					make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()					)
			_lowercase							:      Tuple			=					self.get_auto_remove_tmp_dir()

			args_d.update(data_dir=UpperCAmelCase_							,output_dir=UpperCAmelCase_							,**UpperCAmelCase_					)
			_lowercase							:      Optional[int]			=					distill_main(argparse.Namespace(**UpperCAmelCase_					)					)
			if not check_contents:
					return model
			_lowercase							:      Tuple			=					os.listdir(UpperCAmelCase_					)
			_lowercase							:      str			=					{os.path.basename(UpperCAmelCase_					) for p in contents}
			_lowercase							:      List[str]			=					[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_					)

			_lowercase							:      Any			=					load_json(model.metrics_save_path					)
			_lowercase							:      Union[str, Any]			=					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_					)
			_lowercase							:      Optional[int]			=					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							,UpperCAmelCase_					):
			_lowercase							:      dict			=					CHEAP_ARGS.copy()
			_lowercase							:      Optional[Any]			=					"""translation""" if model in [MBART_TINY, MARIAN_TINY, FSMT_TINY] else """summarization"""
			_lowercase							:      str			=					0.1 if task == """translation""" else 0

			_lowercase							:      Tuple			=					make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()					)
			_lowercase							:      Any			=					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
			_lowercase							:      List[str]			=					argparse.Namespace(**UpperCAmelCase_					)
			_lowercase							:      Optional[int]			=					main(UpperCAmelCase_					)

			_lowercase							:      Union[str, Any]			=					module.model.get_input_embeddings()
			assert not input_embeds.weight.requires_grad
			if model == T5_TINY:
					_lowercase							:      str			=					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							:      Optional[Any]			=					module.model.model
					_lowercase							:      Optional[Any]			=					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							:      List[str]			=					module.model.model
					_lowercase							:      List[str]			=					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							:      Tuple			=					load_json(module.output_dir / """text_batch.json"""					)
			assert isinstance(UpperCAmelCase_							,UpperCAmelCase_					)
			assert len(UpperCAmelCase_					) >= 4






	def 					lowerCamelCase__ (			self					):
			_lowercase							:      dict			=					CHEAP_ARGS.copy()

			_lowercase							:      str			=					"""summarization"""
			_lowercase							:      Dict			=					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
			_lowercase							:      str			=					BART_TINY
			_lowercase							:      Union[str, Any]			=					self.get_auto_remove_tmp_dir()
			_lowercase							:      Optional[Any]			=					args_d.copy()
			args_da.update(
			    model_name_or_path=UpperCAmelCase_							,output_dir=UpperCAmelCase_							,)
			_lowercase							:      List[Any]			=					("""encoder_layerdrop""", """decoder_layerdrop""", """dropout""", """attention_dropout""")
			for p in extra_model_params:
					_lowercase							:      Optional[int]			=					0.5
			_lowercase							:      str			=					argparse.Namespace(**UpperCAmelCase_					)
			_lowercase							:      Any			=					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
			_lowercase							:      Optional[Any]			=					T5_TINY
			_lowercase							:      str			=					self.get_auto_remove_tmp_dir()
			_lowercase							:      List[Any]			=					args_d.copy()
			args_da.update(
			    model_name_or_path=UpperCAmelCase_							,output_dir=UpperCAmelCase_							,)
			_lowercase							:      Optional[Any]			=					"""encoder_layerdrop"""
			_lowercase							:      Any			=					0.5
			_lowercase							:      List[Any]			=					argparse.Namespace(**UpperCAmelCase_					)
			with pytest.raises(UpperCAmelCase_					) as excinfo:
					_lowercase							:      Any			=					main(UpperCAmelCase_					)
			assert str(excinfo.value					) == f"""model config doesn't have a `{unsupported_param}` attribute"""






	def 					lowerCamelCase__ (			self					):
			_lowercase							:      dict			=					CHEAP_ARGS.copy()

			_lowercase							:      Tuple			=					"""summarization"""
			_lowercase							:      List[str]			=					make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()					)

			_lowercase							:      Optional[int]			=					BART_TINY
			_lowercase							:      Tuple			=					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
			_lowercase							:      Dict			=					argparse.ArgumentParser()
			_lowercase							:      List[Any]			=					pl.Trainer.add_argparse_args(UpperCAmelCase_					)
			_lowercase							:      Tuple			=					SummarizationModule.add_model_specific_args(UpperCAmelCase_							,os.getcwd()					)
			_lowercase							:      Union[str, Any]			=					{"""--help""": True}

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








"""simple docstring"""







from typing import TYPE_CHECKING

from ..utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_tf_available, is_torch_available


UpperCAmelCase:   int       			=							{
    """configuration_utils""": ["""GenerationConfig"""],
    """streamers""": ["""TextIteratorStreamer""", """TextStreamer"""],
}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		UpperCAmelCase:   List[Any]       			=							[
		    """Constraint""",
		    """ConstraintListState""",
		    """DisjunctiveConstraint""",
		    """PhrasalConstraint""",
		]
		UpperCAmelCase:   Optional[Any]       			=							[
		    """BeamHypotheses""",
		    """BeamScorer""",
		    """BeamSearchScorer""",
		    """ConstrainedBeamSearchScorer""",
		]
		UpperCAmelCase:   Any       			=							[
		    """EpsilonLogitsWarper""",
		    """EtaLogitsWarper""",
		    """ForcedBOSTokenLogitsProcessor""",
		    """ForcedEOSTokenLogitsProcessor""",
		    """HammingDiversityLogitsProcessor""",
		    """InfNanRemoveLogitsProcessor""",
		    """LogitsProcessor""",
		    """LogitsProcessorList""",
		    """LogitsWarper""",
		    """MinLengthLogitsProcessor""",
		    """MinNewTokensLengthLogitsProcessor""",
		    """NoBadWordsLogitsProcessor""",
		    """NoRepeatNGramLogitsProcessor""",
		    """PrefixConstrainedLogitsProcessor""",
		    """RepetitionPenaltyLogitsProcessor""",
		    """SequenceBiasLogitsProcessor""",
		    """EncoderRepetitionPenaltyLogitsProcessor""",
		    """TemperatureLogitsWarper""",
		    """TopKLogitsWarper""",
		    """TopPLogitsWarper""",
		    """TypicalLogitsWarper""",
		    """EncoderNoRepeatNGramLogitsProcessor""",
		    """ExponentialDecayLengthPenalty""",
		    """LogitNormalization""",
		]
		UpperCAmelCase:   str       			=							[
		    """MaxNewTokensCriteria""",
		    """MaxLengthCriteria""",
		    """MaxTimeCriteria""",
		    """StoppingCriteria""",
		    """StoppingCriteriaList""",
		    """validate_stopping_criteria""",
		]
		UpperCAmelCase:   Optional[Any]       			=							[
		    """GenerationMixin""",
		    """top_k_top_p_filtering""",
		    """GreedySearchEncoderDecoderOutput""",
		    """GreedySearchDecoderOnlyOutput""",
		    """SampleEncoderDecoderOutput""",
		    """SampleDecoderOnlyOutput""",
		    """BeamSearchEncoderDecoderOutput""",
		    """BeamSearchDecoderOnlyOutput""",
		    """BeamSampleEncoderDecoderOutput""",
		    """BeamSampleDecoderOnlyOutput""",
		    """ContrastiveSearchEncoderDecoderOutput""",
		    """ContrastiveSearchDecoderOnlyOutput""",
		]

try:
		if not is_tf_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		UpperCAmelCase:   Any       			=							[
		    """TFForcedBOSTokenLogitsProcessor""",
		    """TFForcedEOSTokenLogitsProcessor""",
		    """TFLogitsProcessor""",
		    """TFLogitsProcessorList""",
		    """TFLogitsWarper""",
		    """TFMinLengthLogitsProcessor""",
		    """TFNoBadWordsLogitsProcessor""",
		    """TFNoRepeatNGramLogitsProcessor""",
		    """TFRepetitionPenaltyLogitsProcessor""",
		    """TFTemperatureLogitsWarper""",
		    """TFTopKLogitsWarper""",
		    """TFTopPLogitsWarper""",
		    """TFForceTokensLogitsProcessor""",
		    """TFSuppressTokensAtBeginLogitsProcessor""",
		    """TFSuppressTokensLogitsProcessor""",
		]
		UpperCAmelCase:   Tuple       			=							[
		    """TFGenerationMixin""",
		    """tf_top_k_top_p_filtering""",
		    """TFGreedySearchDecoderOnlyOutput""",
		    """TFGreedySearchEncoderDecoderOutput""",
		    """TFSampleEncoderDecoderOutput""",
		    """TFSampleDecoderOnlyOutput""",
		    """TFBeamSearchEncoderDecoderOutput""",
		    """TFBeamSearchDecoderOnlyOutput""",
		    """TFBeamSampleEncoderDecoderOutput""",
		    """TFBeamSampleDecoderOnlyOutput""",
		    """TFContrastiveSearchEncoderDecoderOutput""",
		    """TFContrastiveSearchDecoderOnlyOutput""",
		]

try:
		if not is_flax_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		UpperCAmelCase:   Tuple       			=							[
		    """FlaxForcedBOSTokenLogitsProcessor""",
		    """FlaxForcedEOSTokenLogitsProcessor""",
		    """FlaxLogitsProcessor""",
		    """FlaxLogitsProcessorList""",
		    """FlaxLogitsWarper""",
		    """FlaxMinLengthLogitsProcessor""",
		    """FlaxTemperatureLogitsWarper""",
		    """FlaxTopKLogitsWarper""",
		    """FlaxTopPLogitsWarper""",
		]
		UpperCAmelCase:   int       			=							[
		    """FlaxGenerationMixin""",
		    """FlaxGreedySearchOutput""",
		    """FlaxSampleOutput""",
		    """FlaxBeamSearchOutput""",
		]

if TYPE_CHECKING:
		from .configuration_utils import GenerationConfig
		from .streamers import TextIteratorStreamer, TextStreamer

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .beam_constraints import Constraint, ConstraintListState, DisjunctiveConstraint, PhrasalConstraint
				from .beam_search import BeamHypotheses, BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer
				from .logits_process import (
				    EncoderNoRepeatNGramLogitsProcessor,
				    EncoderRepetitionPenaltyLogitsProcessor,
				    EpsilonLogitsWarper,
				    EtaLogitsWarper,
				    ExponentialDecayLengthPenalty,
				    ForcedBOSTokenLogitsProcessor,
				    ForcedEOSTokenLogitsProcessor,
				    HammingDiversityLogitsProcessor,
				    InfNanRemoveLogitsProcessor,
				    LogitNormalization,
				    LogitsProcessor,
				    LogitsProcessorList,
				    LogitsWarper,
				    MinLengthLogitsProcessor,
				    MinNewTokensLengthLogitsProcessor,
				    NoBadWordsLogitsProcessor,
				    NoRepeatNGramLogitsProcessor,
				    PrefixConstrainedLogitsProcessor,
				    RepetitionPenaltyLogitsProcessor,
				    SequenceBiasLogitsProcessor,
				    TemperatureLogitsWarper,
				    TopKLogitsWarper,
				    TopPLogitsWarper,
				    TypicalLogitsWarper,
				)
				from .stopping_criteria import (
				    MaxLengthCriteria,
				    MaxNewTokensCriteria,
				    MaxTimeCriteria,
				    StoppingCriteria,
				    StoppingCriteriaList,
				    validate_stopping_criteria,
				)
				from .utils import (
				    BeamSampleDecoderOnlyOutput,
				    BeamSampleEncoderDecoderOutput,
				    BeamSearchDecoderOnlyOutput,
				    BeamSearchEncoderDecoderOutput,
				    ContrastiveSearchDecoderOnlyOutput,
				    ContrastiveSearchEncoderDecoderOutput,
				    GenerationMixin,
				    GreedySearchDecoderOnlyOutput,
				    GreedySearchEncoderDecoderOutput,
				    SampleDecoderOnlyOutput,
				    SampleEncoderDecoderOutput,
				    top_k_top_p_filtering,
				)

		try:
				if not is_tf_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .tf_logits_process import (
				    TFForcedBOSTokenLogitsProcessor,
				    TFForcedEOSTokenLogitsProcessor,
				    TFForceTokensLogitsProcessor,
				    TFLogitsProcessor,
				    TFLogitsProcessorList,
				    TFLogitsWarper,
				    TFMinLengthLogitsProcessor,
				    TFNoBadWordsLogitsProcessor,
				    TFNoRepeatNGramLogitsProcessor,
				    TFRepetitionPenaltyLogitsProcessor,
				    TFSuppressTokensAtBeginLogitsProcessor,
				    TFSuppressTokensLogitsProcessor,
				    TFTemperatureLogitsWarper,
				    TFTopKLogitsWarper,
				    TFTopPLogitsWarper,
				)
				from .tf_utils import (
				    TFBeamSampleDecoderOnlyOutput,
				    TFBeamSampleEncoderDecoderOutput,
				    TFBeamSearchDecoderOnlyOutput,
				    TFBeamSearchEncoderDecoderOutput,
				    TFContrastiveSearchDecoderOnlyOutput,
				    TFContrastiveSearchEncoderDecoderOutput,
				    TFGenerationMixin,
				    TFGreedySearchDecoderOnlyOutput,
				    TFGreedySearchEncoderDecoderOutput,
				    TFSampleDecoderOnlyOutput,
				    TFSampleEncoderDecoderOutput,
				    tf_top_k_top_p_filtering,
				)

		try:
				if not is_flax_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .flax_logits_process import (
				    FlaxForcedBOSTokenLogitsProcessor,
				    FlaxForcedEOSTokenLogitsProcessor,
				    FlaxLogitsProcessor,
				    FlaxLogitsProcessorList,
				    FlaxLogitsWarper,
				    FlaxMinLengthLogitsProcessor,
				    FlaxTemperatureLogitsWarper,
				    FlaxTopKLogitsWarper,
				    FlaxTopPLogitsWarper,
				)
				from .flax_utils import FlaxBeamSearchOutput, FlaxGenerationMixin, FlaxGreedySearchOutput, FlaxSampleOutput
else:
		import sys

		UpperCAmelCase:   Optional[Any]       			=							_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)



