





"""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_					:						Tuple      = logging.getLogger()
lowerCamelCase_					:						Optional[Any]      = torch.cuda.is_available()
lowerCamelCase_					:						Union[str, 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 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  A_  :	Optional[int]    						=  '\n'.join(_UpperCAmelCase )
  Path(_UpperCAmelCase ).open('w' ).writelines(_UpperCAmelCase )


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


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





def 						UpperCAmelCase__		(					_UpperCAmelCase ):






  """simple docstring"""





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

class 			_UpperCAmelCase  (       UpperCAmelCase__		):



      '''simple docstring'''

      @classmethod
      def  lowerCamelCase_      (					cls  ):

        """simple docstring"""

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

      @slow
      @require_torch_gpu
      def  lowerCamelCase_      (					self  ):

        """simple docstring"""

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

      def  lowerCamelCase_      (					self  ):

        """simple docstring"""

        A_  :	List[Any]    						=  {'student_encoder_layers': 2, 'student_decoder_layers': 1, 'no_teacher': True}
        self._test_distiller_cli(snake_case_  )

      def  lowerCamelCase_      (					self  ):

        """simple docstring"""

        A_  :	List[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',
        }
        A_  :	Any    						=  self._test_distiller_cli(snake_case_ ,				check_contents=snake_case_  )

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

        A_  :	Dict    						=  self.get_auto_remove_tmp_dir()
        convert_pl_to_hf(ckpts[0] ,				transformer_ckpts[0].parent ,				snake_case_  )
        assert os.path.exists(os.path.join(snake_case_ ,				'pytorch_model.bin'  )  )

      def  lowerCamelCase_      (					self  ):

        """simple docstring"""

        A_  :	Union[str, Any]    						=  AutoModelForSeqaSeqLM.from_pretrained(snake_case_  )
        A_  :	Optional[int]    						=  model.dummy_inputs['input_ids'], model.dummy_inputs['attention_mask']
        A_  :	Union[str, Any]    						=  torch.tensor([[0, 4, 8, 2], [0, 8, 2, 1]] ,				dtype=torch.long ,				device=model.device  )
        A_  :	str    						=  target_ids[:, :-1].contiguous()  # Why this line?
        A_  :	Any    						=  target_ids[:, 1:].clone()  # why clone?
        A_  :	Dict    						=  model(
            snake_case_ ,				attention_mask=snake_case_ ,				decoder_input_ids=snake_case_ ,				labels=snake_case_ ,				use_cache=snake_case_  ).loss

        A_  :	str    						=  model(snake_case_ ,				attention_mask=snake_case_ ,				decoder_input_ids=snake_case_ ,				use_cache=snake_case_  ).logits

        A_  :	Any    						=  nn.functional.log_softmax(snake_case_ ,				dim=-1  )
        A_  :	List[Any]    						=  label_smoothed_nll_loss(
            snake_case_ ,				snake_case_ ,				0.1 ,				ignore_index=model.config.pad_token_id  )
        with self.assertRaises(snake_case_  ):
          # TODO: understand why this breaks
          self.assertEqual(snake_case_ ,				snake_case_  )

      def  lowerCamelCase_      (					self  ):

        """simple docstring"""

        A_  :	List[str]    						=  {
            '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_  :	Optional[Any]    						=  self._test_distiller_cli(snake_case_ ,				check_contents=snake_case_  )
        assert model.model.config.model_type == "mbart"

        A_  :	str    						=  list(Path(model.output_dir  ).glob('*.ckpt'  )  )
        self.assertEqual(1 ,				len(snake_case_  )  )
        A_  :	Optional[Any]    						=  list(Path(model.output_dir  ).glob('**/*.bin'  )  )
        A_  :	int    						=  list(Path(model.output_dir  ).glob('best_tfmr/*'  )  )
        assert len(snake_case_  ) > 2
        self.assertEqual(len(snake_case_  ) ,				2  )

      def  lowerCamelCase_      (					self  ):

        """simple docstring"""

        A_  :	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(snake_case_  )

      def  lowerCamelCase_      (					self  ):

        """simple docstring"""

        A_  :	List[str]    						=  {
            'teacher': T5_TINY,
            'student': T5_TINIER,
            'model_name_or_path': T5_TINIER,
            'tokenizer_name': T5_TINIER,
        }
        self._test_distiller_cli(snake_case_  )

      def  lowerCamelCase_      (					self ,				snake_case_ ,				snake_case_=True  ):

        """simple docstring"""

        A_  :	Dict    						=  {
            '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(snake_case_  )
        A_  :	dict    						=  CHEAP_ARGS.copy()
        A_  :	Any    						=  make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()  )
        A_  :	Union[str, Any]    						=  self.get_auto_remove_tmp_dir()

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

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

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


class 			_UpperCAmelCase  (       UpperCAmelCase__		):



      '''simple docstring'''

      @parameterized.expand(
          [T5_TINY, BART_TINY, MBART_TINY, MARIAN_TINY, FSMT_TINY] ,				)
      def  lowerCamelCase_      (					self ,				snake_case_  ):

        """simple docstring"""

        A_  :	dict    						=  CHEAP_ARGS.copy()
        A_  :	Optional[Any]    						=  'translation' if model in [MBART_TINY, MARIAN_TINY, FSMT_TINY] else 'summarization'
        A_  :	Any    						=  0.1 if task == 'translation' else 0

        A_  :	Union[str, Any]    						=  make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()  )
        A_  :	Dict    						=  self.get_auto_remove_tmp_dir()
        args_d.update(
            data_dir=snake_case_ ,				model_name_or_path=snake_case_ ,				tokenizer_name=snake_case_ ,				train_batch_size=2 ,				eval_batch_size=2 ,				output_dir=snake_case_ ,				do_predict=snake_case_ ,				task=snake_case_ ,				src_lang='en_XX' ,				tgt_lang='ro_RO' ,				freeze_encoder=snake_case_ ,				freeze_embeds=snake_case_ ,				)
        assert "n_train" in args_d
        A_  :	Union[str, Any]    						=  argparse.Namespace(**snake_case_  )
        A_  :	str    						=  main(snake_case_  )

        A_  :	Any    						=  module.model.get_input_embeddings()
        assert not input_embeds.weight.requires_grad
        if model == T5_TINY:
          A_  :	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:
          A_  :	str    						=  module.model.model
          A_  :	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:
          A_  :	Dict    						=  module.model.model
          A_  :	Dict    						=  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_  :	Optional[int]    						=  load_json(module.output_dir / 'text_batch.json'  )
        assert isinstance(snake_case_ ,				snake_case_  )
        assert len(snake_case_  ) >= 4

      def  lowerCamelCase_      (					self  ):

        """simple docstring"""

        A_  :	dict    						=  CHEAP_ARGS.copy()

        A_  :	str    						=  'summarization'
        A_  :	Any    						=  make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()  )

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

        # test models whose config includes the extra_model_args
        A_  :	int    						=  BART_TINY
        A_  :	Any    						=  self.get_auto_remove_tmp_dir()
        A_  :	Dict    						=  args_d.copy()
        args_da.update(
            model_name_or_path=snake_case_ ,				output_dir=snake_case_ ,				)
        A_  :	Optional[Any]    						=  ('encoder_layerdrop', 'decoder_layerdrop', 'dropout', 'attention_dropout')
        for p in extra_model_params:
          A_  :	Tuple    						=  0.5
        A_  :	str    						=  argparse.Namespace(**snake_case_  )
        A_  :	Union[str, Any]    						=  main(snake_case_  )
        for p in extra_model_params:
          assert getattr(model.config ,				snake_case_  ) == 0.5, F"""failed to override the model config for param {p}"""

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

      def  lowerCamelCase_      (					self  ):

        """simple docstring"""

        A_  :	dict    						=  CHEAP_ARGS.copy()

        A_  :	List[Any]    						=  'summarization'
        A_  :	List[Any]    						=  make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()  )

        A_  :	Any    						=  BART_TINY
        A_  :	Tuple    						=  self.get_auto_remove_tmp_dir()

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

        # emulate finetune.py
        A_  :	Optional[int]    						=  argparse.ArgumentParser()
        A_  :	List[str]    						=  pl.Trainer.add_argparse_args(snake_case_  )
        A_  :	Tuple    						=  SummarizationModule.add_model_specific_args(snake_case_ ,				os.getcwd()  )
        A_  :	Tuple    						=  {'--help': True}

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


lowerCamelCase_					:						Any      = {
    'configuration_utils': ['GenerationConfig'],
    'streamers': ['TextIteratorStreamer', 'TextStreamer'],
}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   lowerCamelCase_					:						Optional[Any]      = [
       'Constraint',
       'ConstraintListState',
       'DisjunctiveConstraint',
       'PhrasalConstraint',
   ]
   lowerCamelCase_					:						List[str]      = [
       'BeamHypotheses',
       'BeamScorer',
       'BeamSearchScorer',
       'ConstrainedBeamSearchScorer',
   ]
   lowerCamelCase_					:						Dict      = [
       'EpsilonLogitsWarper',
       'EtaLogitsWarper',
       'ForcedBOSTokenLogitsProcessor',
       'ForcedEOSTokenLogitsProcessor',
       'HammingDiversityLogitsProcessor',
       'InfNanRemoveLogitsProcessor',
       'LogitsProcessor',
       'LogitsProcessorList',
       'LogitsWarper',
       'MinLengthLogitsProcessor',
       'MinNewTokensLengthLogitsProcessor',
       'NoBadWordsLogitsProcessor',
       'NoRepeatNGramLogitsProcessor',
       'PrefixConstrainedLogitsProcessor',
       'RepetitionPenaltyLogitsProcessor',
       'SequenceBiasLogitsProcessor',
       'EncoderRepetitionPenaltyLogitsProcessor',
       'TemperatureLogitsWarper',
       'TopKLogitsWarper',
       'TopPLogitsWarper',
       'TypicalLogitsWarper',
       'EncoderNoRepeatNGramLogitsProcessor',
       'ExponentialDecayLengthPenalty',
       'LogitNormalization',
   ]
   lowerCamelCase_					:						str      = [
       'MaxNewTokensCriteria',
       'MaxLengthCriteria',
       'MaxTimeCriteria',
       'StoppingCriteria',
       'StoppingCriteriaList',
       'validate_stopping_criteria',
   ]
   lowerCamelCase_					:						List[str]      = [
       '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:
   lowerCamelCase_					:						Union[str, Any]      = [
       'TFForcedBOSTokenLogitsProcessor',
       'TFForcedEOSTokenLogitsProcessor',
       'TFLogitsProcessor',
       'TFLogitsProcessorList',
       'TFLogitsWarper',
       'TFMinLengthLogitsProcessor',
       'TFNoBadWordsLogitsProcessor',
       'TFNoRepeatNGramLogitsProcessor',
       'TFRepetitionPenaltyLogitsProcessor',
       'TFTemperatureLogitsWarper',
       'TFTopKLogitsWarper',
       'TFTopPLogitsWarper',
       'TFForceTokensLogitsProcessor',
       'TFSuppressTokensAtBeginLogitsProcessor',
       'TFSuppressTokensLogitsProcessor',
   ]
   lowerCamelCase_					:						Union[str, Any]      = [
       '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:
   lowerCamelCase_					:						int      = [
       'FlaxForcedBOSTokenLogitsProcessor',
       'FlaxForcedEOSTokenLogitsProcessor',
       'FlaxLogitsProcessor',
       'FlaxLogitsProcessorList',
       'FlaxLogitsWarper',
       'FlaxMinLengthLogitsProcessor',
       'FlaxTemperatureLogitsWarper',
       'FlaxTopKLogitsWarper',
       'FlaxTopPLogitsWarper',
   ]
   lowerCamelCase_					:						Tuple      = [
       '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

   lowerCamelCase_					:						Tuple      = _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)