




import itertools
import random
import unittest

import numpy as np

from transformers import BatchFeature, SpeechTaFeatureExtractor
from transformers.testing_utils import require_torch
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_torch_available():
  import torch


_snake_case            = random.Random()







def        lowerCAmelCase_			(  snake_case_,snake_case_=1.0,snake_case_=None,snake_case_=None  ):
      if rng is None:
            _A							:  List[Any]    =		global_rng

      _A							:  Union[str, Any]    =		[]
      for batch_idx in range(shape[0]  ):
            values.append([]  )
            for _ in range(shape[1]  ):
                  values[-1].append(rng.random() * scale  )

      return values







@require_torch
class 		lowercase							(     unittest.TestCase	):

     def __init__(						self	,  _a	,  _a=7	,  _a=400	,  _a=2000	,  _a=1	,  _a=0.0	,  _a=1_6000	,  _a=True	,  _a=80	,  _a=16	,  _a=64	,  _a="hann_window"	,  _a=80	,  _a=7600	,  _a=1e-10	,  _a=True	,  )		->							List[str]:
           _A							:  Optional[Any]    =		parent
           _A							:  Union[str, Any]    =		batch_size
           _A							:  int    =		min_seq_length
           _A							:  Tuple    =		max_seq_length
           _A							:  List[Any]    =		(self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
           _A							:  Any    =		feature_size
           _A							:  List[Any]    =		padding_value
           _A							:  Tuple    =		sampling_rate
           _A							:  int    =		do_normalize
           _A							:  str    =		num_mel_bins
           _A							:  List[str]    =		hop_length
           _A							:  Tuple    =		win_length
           _A							:  List[str]    =		win_function
           _A							:  List[str]    =		fmin
           _A							:  Dict    =		fmax
           _A							:  Tuple    =		mel_floor
           _A							:  Tuple    =		return_attention_mask

     def       a__			(						self       )		->							Tuple:
           return {
               "feature_size": self.feature_size,
               "padding_value": self.padding_value,
               "sampling_rate": self.sampling_rate,
               "do_normalize": self.do_normalize,
               "num_mel_bins": self.num_mel_bins,
               "hop_length": self.hop_length,
               "win_length": self.win_length,
               "win_function": self.win_function,
               "fmin": self.fmin,
               "fmax": self.fmax,
               "mel_floor": self.mel_floor,
               "return_attention_mask": self.return_attention_mask,
           }

     def       a__			(						self	,  _a=False	,  _a=False       )		->							Optional[int]:
           def _flatten(_a       ):
                 return list(itertools.chain(*_a       )       )

           if equal_length:
                 _A							:  List[Any]    =		floats_list((self.batch_size, self.max_seq_length)       )
           else:
                 # make sure that inputs increase in size
                 _A							:  List[str]    =		[
                     _flatten(floats_list((x, self.feature_size)       )       )
                     for x in range(self.min_seq_length	,  self.max_seq_length	,  self.seq_length_diff       )
                 ]

           if numpify:
                 _A							:  int    =		[np.asarray(_a       ) for x in speech_inputs]

           return speech_inputs



     def       a__			(						self	,  _a=False	,  _a=False       )		->							str:
           if equal_length:
                 _A							:  int    =		[floats_list((self.max_seq_length, self.num_mel_bins)       ) for _ in range(self.batch_size       )]
           else:
                 # make sure that inputs increase in size
                 _A							:  str    =		[
                     floats_list((x, self.num_mel_bins)       )
                     for x in range(self.min_seq_length	,  self.max_seq_length	,  self.seq_length_diff       )
                 ]

           if numpify:
                 _A							:  List[str]    =		[np.asarray(_a       ) for x in speech_inputs]

           return speech_inputs







@require_torch
class 		lowercase							(     UpperCamelCase__,unittest.TestCase	):
     _a											=		SpeechTaFeatureExtractor

     def       a__			(						self       )		->							Union[str, Any]:
           _A							:  Any    =		SpeechTaFeatureExtractionTester(self       )

     def       a__			(						self	,  _a       )		->							Union[str, Any]:
           self.assertTrue(np.all(np.mean(_a	,  axis=0       ) < 1e-3       )       )
           self.assertTrue(np.all(np.abs(np.var(_a	,  axis=0       ) - 1       ) < 1e-3       )       )

     def       a__			(						self       )		->							int:
           # Tests that all call wrap to encode_plus and batch_encode_plus
           _A							:  List[Any]    =		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()       )
           # create three inputs of length 800, 1000, and 1200
           _A							:  int    =		[floats_list((1, x)       )[0] for x in range(800	,  1400	,  200       )]
           _A							:  int    =		[np.asarray(_a       ) for speech_input in speech_inputs]

           # Test not batched input
           _A							:  List[Any]    =		feat_extract(speech_inputs[0]	,  return_tensors="""np"""       ).input_values
           _A							:  Dict    =		feat_extract(np_speech_inputs[0]	,  return_tensors="""np"""       ).input_values
           self.assertTrue(np.allclose(_a	,  _a	,  atol=1e-3       )       )

           # Test batched
           _A							:  List[str]    =		feat_extract(_a	,  return_tensors="""np"""       ).input_values
           _A							:  List[str]    =		feat_extract(_a	,  return_tensors="""np"""       ).input_values
           for enc_seq_a, enc_seq_a in zip(_a	,  _a       ):
                 self.assertTrue(np.allclose(_a	,  _a	,  atol=1e-3       )       )

     def       a__			(						self       )		->							Dict:
           _A							:  str    =		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()       )
           _A							:  Tuple    =		[floats_list((1, x)       )[0] for x in range(800	,  1400	,  200       )]

           _A							:  Any    =		["""longest""", """max_length""", """do_not_pad"""]
           _A							:  Union[str, Any]    =		[None, 1600, None]
           for max_length, padding in zip(_a	,  _a       ):
                 _A							:  Optional[Any]    =		feat_extract(_a	,  padding=_a	,  max_length=_a	,  return_tensors="""np"""       )
                 _A							:  List[str]    =		processed.input_values

                 self._check_zero_mean_unit_variance(input_values[0][:800]       )
                 self.assertTrue(input_values[0][800:].sum() < 1e-6       )
                 self._check_zero_mean_unit_variance(input_values[1][:1000]       )
                 self.assertTrue(input_values[0][1000:].sum() < 1e-6       )
                 self._check_zero_mean_unit_variance(input_values[2][:1200]       )

     def       a__			(						self       )		->							Union[str, Any]:
           _A							:  List[Any]    =		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()       )
           _A							:  List[str]    =		range(800	,  1400	,  200       )
           _A							:  Any    =		[floats_list((1, x)       )[0] for x in lengths]

           _A							:  List[Any]    =		["""longest""", """max_length""", """do_not_pad"""]
           _A							:  Optional[Any]    =		[None, 1600, None]

           for max_length, padding in zip(_a	,  _a       ):
                 _A							:  Optional[Any]    =		feat_extract(_a	,  max_length=_a	,  padding=_a       )
                 _A							:  int    =		processed.input_values

                 self._check_zero_mean_unit_variance(input_values[0][:800]       )
                 self._check_zero_mean_unit_variance(input_values[1][:1000]       )
                 self._check_zero_mean_unit_variance(input_values[2][:1200]       )

     def       a__			(						self       )		->							Tuple:
           _A							:  Tuple    =		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()       )
           _A							:  str    =		[floats_list((1, x)       )[0] for x in range(800	,  1400	,  200       )]
           _A							:  Any    =		feat_extract(
               _a	,  truncation=_a	,  max_length=1000	,  padding="""max_length"""	,  return_tensors="""np"""       )
           _A							:  Optional[Any]    =		processed.input_values

           self._check_zero_mean_unit_variance(input_values[0, :800]       )
           self._check_zero_mean_unit_variance(input_values[1]       )
           self._check_zero_mean_unit_variance(input_values[2]       )

     def       a__			(						self       )		->							Optional[int]:
           _A							:  Union[str, Any]    =		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()       )
           _A							:  Dict    =		[floats_list((1, x)       )[0] for x in range(800	,  1400	,  200       )]
           _A							:  str    =		feat_extract(
               _a	,  truncation=_a	,  max_length=1000	,  padding="""longest"""	,  return_tensors="""np"""       )
           _A							:  Any    =		processed.input_values

           self._check_zero_mean_unit_variance(input_values[0, :800]       )
           self._check_zero_mean_unit_variance(input_values[1, :1000]       )
           self._check_zero_mean_unit_variance(input_values[2]       )

           # make sure that if max_length < longest -> then pad to max_length
           self.assertTrue(input_values.shape == (3, 1000)       )

           _A							:  List[Any]    =		[floats_list((1, x)       )[0] for x in range(800	,  1400	,  200       )]
           _A							:  Dict    =		feat_extract(
               _a	,  truncation=_a	,  max_length=2000	,  padding="""longest"""	,  return_tensors="""np"""       )
           _A							:  List[str]    =		processed.input_values

           self._check_zero_mean_unit_variance(input_values[0, :800]       )
           self._check_zero_mean_unit_variance(input_values[1, :1000]       )
           self._check_zero_mean_unit_variance(input_values[2]       )

           # make sure that if max_length > longest -> then pad to longest
           self.assertTrue(input_values.shape == (3, 1200)       )

     def       a__			(						self       )		->							int:
           _A							:  Dict    =		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()       )
           _A							:  Any    =		np.random.rand(100       ).astype(np.floataa       )
           _A							:  str    =		np_speech_inputs.tolist()

           for inputs in [py_speech_inputs, np_speech_inputs]:
                 _A							:  List[Any]    =		feature_extractor.pad([{"""input_values""": inputs}]	,  return_tensors="""np"""       )
                 self.assertTrue(np_processed.input_values.dtype == np.floataa       )
                 _A							:  Optional[Any]    =		feature_extractor.pad([{"""input_values""": inputs}]	,  return_tensors="""pt"""       )
                 self.assertTrue(pt_processed.input_values.dtype == torch.floataa       )

     def       a__			(						self       )		->							Any:
           # Tests that all call wrap to encode_plus and batch_encode_plus
           _A							:  Any    =		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()       )
           # create three inputs of length 800, 1000, and 1200
           _A							:  Dict    =		[floats_list((1, x)       )[0] for x in range(800	,  1400	,  200       )]
           _A							:  Optional[Any]    =		[np.asarray(_a       ) for speech_input in speech_inputs]

           # Test feature size
           _A							:  Dict    =		feature_extractor(audio_target=_a	,  padding=_a	,  return_tensors="""np"""       ).input_values
           self.assertTrue(input_values.ndim == 3       )
           self.assertTrue(input_values.shape[-1] == feature_extractor.num_mel_bins       )

           # Test not batched input
           _A							:  Any    =		feature_extractor(speech_inputs[0]	,  return_tensors="""np"""       ).input_values
           _A							:  List[str]    =		feature_extractor(np_speech_inputs[0]	,  return_tensors="""np"""       ).input_values
           self.assertTrue(np.allclose(_a	,  _a	,  atol=1e-3       )       )

           # Test batched
           _A							:  Dict    =		feature_extractor(_a	,  return_tensors="""np"""       ).input_values
           _A							:  Optional[int]    =		feature_extractor(_a	,  return_tensors="""np"""       ).input_values
           for enc_seq_a, enc_seq_a in zip(_a	,  _a       ):
                 self.assertTrue(np.allclose(_a	,  _a	,  atol=1e-3       )       )

           # Test 2-D numpy arrays are batched.
           _A							:  Tuple    =		[floats_list((1, x)       )[0] for x in (800, 800, 800)]
           _A							:  Tuple    =		np.asarray(_a       )
           _A							:  Union[str, Any]    =		feature_extractor(_a	,  return_tensors="""np"""       ).input_values
           _A							:  Optional[int]    =		feature_extractor(_a	,  return_tensors="""np"""       ).input_values
           for enc_seq_a, enc_seq_a in zip(_a	,  _a       ):
                 self.assertTrue(np.allclose(_a	,  _a	,  atol=1e-3       )       )

     def       a__			(						self       )		->							Tuple:
           _A							:  List[Any]    =		self.feat_extract_tester.prepare_inputs_for_target()
           _A							:  Optional[Any]    =		self.feature_extraction_class(**self.feat_extract_dict       )
           _A							:  List[Any]    =		feat_extract.model_input_names[0]

           _A							:  Union[str, Any]    =		BatchFeature({input_name: speech_inputs}       )

           self.assertTrue(all(len(_a       ) == len(_a       ) for x, y in zip(_a	,  processed_features[input_name]       )       )       )

           _A							:  Any    =		self.feat_extract_tester.prepare_inputs_for_target(equal_length=_a       )
           _A							:  Any    =		BatchFeature({input_name: speech_inputs}	,  tensor_type="""np"""       )

           _A							:  Tuple    =		processed_features[input_name]

           if len(batch_features_input.shape       ) < 3:
                 _A							:  str    =		batch_features_input[:, :, None]

           self.assertTrue(
               batch_features_input.shape
               == (self.feat_extract_tester.batch_size, len(speech_inputs[0]       ), self.feat_extract_tester.num_mel_bins)       )

     @require_torch
     def       a__			(						self       )		->							List[Any]:
           _A							:  Any    =		self.feat_extract_tester.prepare_inputs_for_target(equal_length=_a       )
           _A							:  Union[str, Any]    =		self.feature_extraction_class(**self.feat_extract_dict       )
           _A							:  Any    =		feat_extract.model_input_names[0]

           _A							:  Union[str, Any]    =		BatchFeature({input_name: speech_inputs}	,  tensor_type="""pt"""       )

           _A							:  int    =		processed_features[input_name]

           if len(batch_features_input.shape       ) < 3:
                 _A							:  int    =		batch_features_input[:, :, None]

           self.assertTrue(
               batch_features_input.shape
               == (self.feat_extract_tester.batch_size, len(speech_inputs[0]       ), self.feat_extract_tester.num_mel_bins)       )

     @require_torch
     def       a__			(						self       )		->							Any:
           _A							:  Optional[Any]    =		self.feature_extraction_class(**self.feat_extract_dict       )
           _A							:  Any    =		self.feat_extract_tester.prepare_inputs_for_target()
           _A							:  Any    =		feat_extract.model_input_names[0]

           _A							:  Optional[Any]    =		BatchFeature({input_name: speech_inputs}       )

           _A							:  Dict    =		feat_extract.num_mel_bins  # hack!

           _A							:  Union[str, Any]    =		feat_extract.pad(_a	,  padding="""longest"""	,  return_tensors="""np"""       )[input_name]
           _A							:  Union[str, Any]    =		feat_extract.pad(_a	,  padding="""longest"""	,  return_tensors="""pt"""       )[input_name]

           self.assertTrue(abs(input_np.astype(np.floataa       ).sum() - input_pt.numpy().astype(np.floataa       ).sum()       ) < 1e-2       )

     def       a__			(						self       )		->							List[Any]:
           _A							:  List[str]    =		self.feat_extract_dict
           _A							:  Any    =		True
           _A							:  List[Any]    =		self.feature_extraction_class(**_a       )
           _A							:  List[Any]    =		self.feat_extract_tester.prepare_inputs_for_target()
           _A							:  str    =		[len(_a       ) for x in speech_inputs]
           _A							:  Dict    =		feat_extract.model_input_names[0]

           _A							:  str    =		BatchFeature({input_name: speech_inputs}       )

           _A							:  Tuple    =		feat_extract.num_mel_bins  # hack!

           _A							:  Tuple    =		feat_extract.pad(_a	,  padding="""longest"""	,  return_tensors="""np"""       )
           self.assertIn("""attention_mask"""	,  _a       )
           self.assertListEqual(list(processed.attention_mask.shape       )	,  list(processed[input_name].shape[:2]       )       )
           self.assertListEqual(processed.attention_mask.sum(-1       ).tolist()	,  _a       )

     def       a__			(						self       )		->							str:
           _A							:  Union[str, Any]    =		self.feat_extract_dict
           _A							:  int    =		True
           _A							:  Optional[Any]    =		self.feature_extraction_class(**_a       )
           _A							:  Any    =		self.feat_extract_tester.prepare_inputs_for_target()
           _A							:  Dict    =		[len(_a       ) for x in speech_inputs]
           _A							:  Tuple    =		feat_extract.model_input_names[0]

           _A							:  Optional[Any]    =		BatchFeature({input_name: speech_inputs}       )
           _A							:  str    =		min(_a       )

           _A							:  Any    =		feat_extract.num_mel_bins  # hack!

           _A							:  Union[str, Any]    =		feat_extract.pad(
               _a	,  padding="""max_length"""	,  max_length=_a	,  truncation=_a	,  return_tensors="""np"""       )
           self.assertIn("""attention_mask"""	,  _a       )
           self.assertListEqual(
               list(processed_pad.attention_mask.shape       )	,  [processed_pad[input_name].shape[0], max_length]       )
           self.assertListEqual(
               processed_pad.attention_mask[:, :max_length].sum(-1       ).tolist()	,  [max_length for x in speech_inputs]       )

     def       a__			(						self	,  _a       )		->							Tuple:
           from datasets import load_dataset

           _A							:  Optional[Any]    =		load_dataset("""hf-internal-testing/librispeech_asr_dummy"""	,  """clean"""	,  split="""validation"""       )
           # automatic decoding with librispeech
           _A							:  Optional[Any]    =		ds.sort("""id"""       ).select(range(_a       )       )[:num_samples]["""audio"""]

           return [x["array"] for x in speech_samples]

     def       a__			(						self       )		->							Tuple:
           # fmt: off
           _A							:  int    =		torch.tensor(
               [2.38_04e-03, 2.07_52e-03, 1.98_36e-03, 2.10_57e-03, 1.61_74e-03,
                3.05_18e-04, 9.15_53e-05, 3.35_69e-04, 9.76_56e-04, 1.83_11e-03,
                2.01_42e-03, 2.10_57e-03, 1.73_95e-03, 4.57_76e-04, -3.96_73e-04,
                4.57_76e-04, 1.00_71e-03, 9.15_53e-05, 4.88_28e-04, 1.15_97e-03,
                7.32_42e-04, 9.46_04e-04, 1.80_05e-03, 1.83_11e-03, 8.85_01e-04,
                4.27_25e-04, 4.88_28e-04, 7.32_42e-04, 1.09_86e-03, 2.10_57e-03]       )
           # fmt: on

           _A							:  List[str]    =		self._load_datasamples(1       )
           _A							:  Optional[Any]    =		SpeechTaFeatureExtractor()
           _A							:  Dict    =		feature_extractor(_a	,  return_tensors="""pt"""       ).input_values
           self.assertEquals(input_values.shape	,  (1, 9_3680)       )
           self.assertTrue(torch.allclose(input_values[0, :30]	,  _a	,  atol=1e-6       )       )



     def       a__			(						self       )		->							Dict:
           # fmt: off
           _A							:  Tuple    =		torch.tensor(
               [-2.6870, -3.0104, -3.1356, -3.5352, -3.0044, -3.0353, -3.4719, -3.6777,
                -3.1520, -2.9435, -2.6553, -2.8795, -2.9944, -2.5921, -3.0279, -3.0386,
                -3.0864, -3.1291, -3.2353, -2.7444, -2.6831, -2.7287, -3.1761, -3.1571,
                -3.2726, -3.0582, -3.1007, -3.4533, -3.4695, -3.0998]       )
           # fmt: on

           _A							:  Dict    =		self._load_datasamples(1       )
           _A							:  Any    =		SpeechTaFeatureExtractor()
           _A							:  List[str]    =		feature_extractor(audio_target=_a	,  return_tensors="""pt"""       ).input_values
           self.assertEquals(input_values.shape	,  (1, 366, 80)       )
           self.assertTrue(torch.allclose(input_values[0, 0, :30]	,  _a	,  atol=1e-4       )       )







from typing import Any, Dict, List, Optional, Tuple, Union

import torch
from torch import nn
from torch.utils.data import DistributedSampler, RandomSampler

from transformers import PreTrainedModel, Trainer, logging
from transformers.integrations import is_fairscale_available
from transformers.models.fsmt.configuration_fsmt import FSMTConfig
from transformers.optimization import (
    Adafactor,
    AdamW,
    get_constant_schedule,
    get_constant_schedule_with_warmup,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.trainer_pt_utils import get_tpu_sampler
from transformers.training_args import ParallelMode
from transformers.utils import is_torch_tpu_available


if is_fairscale_available():
  from fairscale.optim import OSS


_snake_case            = logging.get_logger(__name__)

_snake_case            = {
    "linear": get_linear_schedule_with_warmup,
    "cosine": get_cosine_schedule_with_warmup,
    "cosine_w_restarts": get_cosine_with_hard_restarts_schedule_with_warmup,
    "polynomial": get_polynomial_decay_schedule_with_warmup,
    "constant": get_constant_schedule,
    "constant_w_warmup": get_constant_schedule_with_warmup,
}







class 		lowercase							(     UpperCamelCase__	):

     def __init__(						self	,  _a=None	,  _a=None	,  *_a	,  **_a       )		->							Optional[int]:
           super().__init__(*_a	,  **_a       )

           if config is None:
                 assert isinstance(self.model	,  _a       ), (
                     "If no `config` is passed the model to be trained has to be of type `PreTrainedModel`, but is"
                     F''' {self.model.__class__}'''
                 )
                 _A							:  Optional[Any]    =		self.model.config
           else:
                 _A							:  int    =		config

           _A							:  Optional[Any]    =		data_args
           _A							:  int    =		self.config.tgt_vocab_size if isinstance(self.config	,  _a       ) else self.config.vocab_size

           if self.args.label_smoothing != 0 or (self.data_args is not None and self.data_args.ignore_pad_token_for_loss):
                 assert self.config.pad_token_id is not None, (
                     "Make sure that `config.pad_token_id` is correcly defined when ignoring `pad_token` for loss"
                     " calculation or doing label smoothing."
                 )

           if self.config.pad_token_id is None and self.config.eos_token_id is not None:
                 logger.warning(
                     F'''The `config.pad_token_id` is `None`. Using `config.eos_token_id` = {self.config.eos_token_id} for'''
                     """ padding.."""       )

           if self.args.label_smoothing == 0:
                 _A							:  Optional[Any]    =		torch.nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id       )
           else:
                 # dynamically import label_smoothed_nll_loss
                 from utils import label_smoothed_nll_loss

                 _A							:  Union[str, Any]    =		label_smoothed_nll_loss

     def       a__			(						self	,  _a       )		->							int:
           if self.optimizer is None:
                 _A							:  List[str]    =		["""bias""", """LayerNorm.weight"""]
                 _A							:  str    =		[
                     {
                         """params""": [p for n, p in self.model.named_parameters() if not any(nd in n for nd in no_decay       )],
                         """weight_decay""": self.args.weight_decay,
                     },
                     {
                         """params""": [p for n, p in self.model.named_parameters() if any(nd in n for nd in no_decay       )],
                         """weight_decay""": 0.0,
                     },
                 ]
                 _A							:  Optional[Any]    =		Adafactor if self.args.adafactor else AdamW
                 if self.args.adafactor:
                       _A							:  Dict    =		Adafactor
                       _A							:  int    =		{"""scale_parameter""": False, """relative_step""": False}
                 else:
                       _A							:  int    =		AdamW
                       _A							:  Any    =		{
                           """betas""": (self.args.adam_betaa, self.args.adam_betaa),
                           """eps""": self.args.adam_epsilon,
                       }
                 _A							:  List[str]    =		self.args.learning_rate
                 if self.sharded_ddp:
                       _A							:  List[str]    =		OSS(
                           params=_a	,  optim=_a	,  **_a	,  )
                 else:
                       _A							:  Tuple    =		optimizer_cls(_a	,  **_a       )

           if self.lr_scheduler is None:
                 _A							:  Union[str, Any]    =		self._get_lr_scheduler(_a       )
           else:  # ignoring --lr_scheduler
                 logger.warning("""scheduler is passed to `Seq2SeqTrainer`, `--lr_scheduler` arg is ignored."""       )

     def       a__			(						self	,  _a       )		->							Dict:
           _A							:  List[Any]    =		arg_to_scheduler[self.args.lr_scheduler]
           if self.args.lr_scheduler == "constant":
                 _A							:  Optional[Any]    =		schedule_func(self.optimizer       )
           elif self.args.lr_scheduler == "constant_w_warmup":
                 _A							:  Optional[int]    =		schedule_func(self.optimizer	,  num_warmup_steps=self.args.warmup_steps       )
           else:
                 _A							:  List[Any]    =		schedule_func(
                     self.optimizer	,  num_warmup_steps=self.args.warmup_steps	,  num_training_steps=_a       )
           return scheduler

     def       a__			(						self       )		->							Optional[torch.utils.data.Sampler]:
           if isinstance(self.train_dataset	,  torch.utils.data.IterableDataset       ):
                 return None
           elif is_torch_tpu_available():
                 return get_tpu_sampler(self.train_dataset       )
           else:
                 if self.args.sortish_sampler:
                       self.train_dataset.make_sortish_sampler(
                           self.args.per_device_train_batch_size	,  distributed=(self.args.parallel_mode == ParallelMode.DISTRIBUTED)	,  )

                 return (
                     RandomSampler(self.train_dataset       )
                     if self.args.local_rank == -1
                     else DistributedSampler(self.train_dataset       )
                 )

     def       a__			(						self	,  _a	,  _a	,  _a       )		->							List[str]:
           if self.args.label_smoothing == 0:
                 if self.data_args is not None and self.data_args.ignore_pad_token_for_loss:
                       # force training to ignore pad token
                       _A							:  List[str]    =		model(**_a	,  use_cache=_a       )[0]
                       _A							:  Dict    =		self.loss_fn(logits.view(-1	,  logits.shape[-1]       )	,  labels.view(-1       )       )
                 else:
                       # compute usual loss via models
                       _A    ,     _A							:  str    =		model(**_a	,  labels=_a	,  use_cache=_a       )[:2]
           else:
                 # compute label smoothed loss
                 _A							:  Any    =		model(**_a	,  use_cache=_a       )[0]
                 _A							:  Union[str, Any]    =		torch.nn.functional.log_softmax(_a	,  dim=-1       )
                 _A    ,     _A							:  List[str]    =		self.loss_fn(_a	,  _a	,  self.args.label_smoothing	,  ignore_index=self.config.pad_token_id       )
           return loss, logits

     def       a__			(						self	,  _a	,  _a       )		->							List[Any]:
           _A							:  Optional[int]    =		inputs.pop("""labels"""       )
           _A    ,     _A							:  Dict    =		self._compute_loss(_a	,  _a	,  _a       )
           return loss

     def       a__			(						self	,  _a	,  _a	,  _a	,  _a = None	,  )		->							Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
           _A							:  int    =		self._prepare_inputs(_a       )

           _A							:  Dict    =		{
               """max_length""": self.data_args.val_max_target_length
               if self.data_args is not None
               else self.config.max_length,
               """num_beams""": self.data_args.eval_beams if self.data_args is not None else self.config.num_beams,
           }

           if self.args.predict_with_generate and not self.args.prediction_loss_only:
                 _A							:  List[str]    =		self.model.generate(
                     inputs["""input_ids"""]	,  attention_mask=inputs["""attention_mask"""]	,  **_a	,  )
                 # in case the batch is shorter than max length, the output should be padded
                 if generated_tokens.shape[-1] < gen_kwargs["max_length"]:
                       _A							:  str    =		self._pad_tensors_to_max_len(_a	,  gen_kwargs["""max_length"""]       )

           _A							:  Any    =		inputs.pop("""labels"""       )
           with torch.no_grad():
                 # compute loss on predict data
                 _A    ,     _A							:  Tuple    =		self._compute_loss(_a	,  _a	,  _a       )

           _A							:  Any    =		loss.mean().detach()
           if self.args.prediction_loss_only:
                 return (loss, None, None)

           _A							:  List[Any]    =		generated_tokens if self.args.predict_with_generate else logits

           if labels.shape[-1] < gen_kwargs["max_length"]:
                 _A							:  Any    =		self._pad_tensors_to_max_len(_a	,  gen_kwargs["""max_length"""]       )

           return (loss, logits, labels)



     def       a__			(						self	,  _a	,  _a       )		->							Union[str, Any]:
           # If PAD token is not defined at least EOS token has to be defined
           _A							:  Optional[Any]    =		self.config.pad_token_id if self.config.pad_token_id is not None else self.config.eos_token_id

           if pad_token_id is None:
                 raise ValueError(
                     """Make sure that either `config.pad_token_id` or `config.eos_token_id` is defined if tensor has to be"""
                     F''' padded to `max_length`={max_length}'''       )

           _A							:  Dict    =		pad_token_id * torch.ones(
               (tensor.shape[0], max_length)	,  dtype=tensor.dtype	,  device=tensor.device       )
           _A							:  Dict    =		tensor
           return padded_tensor

