






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


UpperCAmelCase__						:str          =			logging.get_logger(__name__)

UpperCAmelCase__						:Optional[int]          =			{
    """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 	SCREAMING_SNAKE_CASE	(      _UpperCAmelCase  ):




       def __init__(					self      : str	,    A__      : List[str]=None	,    A__      : List[Any]=None	,    *A__      : Tuple	,    **A__      : Optional[int]					):


            """simple docstring"""







            super().__init__(*lowercase__	,    **lowercase__					)

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

            __lowerCamelCase						:       Any				   =				data_args
            __lowerCamelCase						:       Union[str, Any]				   =				self.config.tgt_vocab_size if isinstance(self.config	,    lowercase__					) 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:
                 __lowerCamelCase						:       Optional[int]				   =				torch.nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id					)
            else:
                 # dynamically import label_smoothed_nll_loss
                 from utils import label_smoothed_nll_loss

                 __lowerCamelCase						:       Optional[int]				   =				label_smoothed_nll_loss




       def 					a_					(					self      : List[Any]	,    A__      : List[str]					):


            """simple docstring"""







            if self.optimizer is None:
                 __lowerCamelCase						:       str				   =				["""bias""", """LayerNorm.weight"""]
                 __lowerCamelCase						:       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,
                     },
                 ]
                 __lowerCamelCase						:       Tuple				   =				Adafactor if self.args.adafactor else AdamW
                 if self.args.adafactor:
                      __lowerCamelCase						:       Dict				   =				Adafactor
                      __lowerCamelCase						:       List[Any]				   =				{"""scale_parameter""": False, """relative_step""": False}
                 else:
                      __lowerCamelCase						:       Optional[int]				   =				AdamW
                      __lowerCamelCase						:       Tuple				   =				{
                          """betas""": (self.args.adam_betaa, self.args.adam_betaa),
                          """eps""": self.args.adam_epsilon,
                      }
                 __lowerCamelCase						:       int				   =				self.args.learning_rate
                 if self.sharded_ddp:
                      __lowerCamelCase						:       str				   =				OSS(
                          params=lowercase__	,    optim=lowercase__	,    **lowercase__	,    )
                 else:
                      __lowerCamelCase						:       Any				   =				optimizer_cls(lowercase__	,    **lowercase__					)

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




       def 					a_					(					self      : Union[str, Any]	,    A__      : List[str]					):


            """simple docstring"""







            __lowerCamelCase						:       str				   =				arg_to_scheduler[self.args.lr_scheduler]
            if self.args.lr_scheduler == "constant":
                 __lowerCamelCase						:       Tuple				   =				schedule_func(self.optimizer					)
            elif self.args.lr_scheduler == "constant_w_warmup":
                 __lowerCamelCase						:       Union[str, Any]				   =				schedule_func(self.optimizer	,    num_warmup_steps=self.args.warmup_steps					)
            else:
                 __lowerCamelCase						:       Optional[Any]				   =				schedule_func(
                     self.optimizer	,    num_warmup_steps=self.args.warmup_steps	,    num_training_steps=lowercase__					)
            return scheduler




       def 					a_					(					self      : Optional[int]					):


            """simple docstring"""







            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      : Optional[int]	,    A__      : List[str]	,    A__      : Dict	,    A__      : int					):


            """simple docstring"""







            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
                      __lowerCamelCase						:       int				   =				model(**lowercase__	,    use_cache=lowercase__					)[0]
                      __lowerCamelCase						:       Dict				   =				self.loss_fn(logits.view(-1	,    logits.shape[-1]					)	,    labels.view(-1					)					)
                 else:
                      # compute usual loss via models
                      __lowerCamelCase						:       Union[str, Any]				   =				model(**lowercase__	,    labels=lowercase__	,    use_cache=lowercase__					)[:2]
            else:
                 # compute label smoothed loss
                 __lowerCamelCase						:       int				   =				model(**lowercase__	,    use_cache=lowercase__					)[0]
                 __lowerCamelCase						:       Optional[int]				   =				torch.nn.functional.log_softmax(lowercase__	,    dim=-1					)
                 __lowerCamelCase						:       Any				   =				self.loss_fn(lowercase__	,    lowercase__	,    self.args.label_smoothing	,    ignore_index=self.config.pad_token_id					)
            return loss, logits




       def 					a_					(					self      : int	,    A__      : Optional[Any]	,    A__      : List[Any]					):


            """simple docstring"""







            __lowerCamelCase						:       Tuple				   =				inputs.pop("""labels"""					)
            __lowerCamelCase						:       Union[str, Any]				   =				self._compute_loss(lowercase__	,    lowercase__	,    lowercase__					)
            return loss




       def 					a_					(					self      : List[Any]	,    A__      : str	,    A__      : str	,    A__      : Any	,    A__      : Any = None	,    ):


            """simple docstring"""







            __lowerCamelCase						:       int				   =				self._prepare_inputs(lowercase__					)

            __lowerCamelCase						:       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:
                 __lowerCamelCase						:       int				   =				self.model.generate(
                     inputs["""input_ids"""]	,    attention_mask=inputs["""attention_mask"""]	,    **lowercase__	,    )
                 # in case the batch is shorter than max length, the output should be padded
                 if generated_tokens.shape[-1] < gen_kwargs["max_length"]:
                      __lowerCamelCase						:       Optional[Any]				   =				self._pad_tensors_to_max_len(lowercase__	,    gen_kwargs["""max_length"""]					)

            __lowerCamelCase						:       int				   =				inputs.pop("""labels"""					)
            with torch.no_grad():
                 # compute loss on predict data
                 __lowerCamelCase						:       str				   =				self._compute_loss(lowercase__	,    lowercase__	,    lowercase__					)

            __lowerCamelCase						:       List[Any]				   =				loss.mean().detach()
            if self.args.prediction_loss_only:
                 return (loss, None, None)

            __lowerCamelCase						:       str				   =				generated_tokens if self.args.predict_with_generate else logits

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

            return (loss, logits, labels)




       def 					a_					(					self      : int	,    A__      : str	,    A__      : Dict					):


            """simple docstring"""







            __lowerCamelCase						:       Tuple				   =				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}"					)

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











'''simple docstring'''

import shutil
import tempfile
import unittest

from transformers import ClapFeatureExtractor, ClapProcessor, RobertaTokenizer, RobertaTokenizerFast
from transformers.testing_utils import require_sentencepiece, require_torchaudio

from .test_feature_extraction_clap import floats_list







@require_torchaudio
@require_sentencepiece
class 	SCREAMING_SNAKE_CASE	(      unittest.TestCase  ):




       def 					a_					(					self      : str					):


            """simple docstring"""







            __lowerCamelCase						:       Any				   =				"""laion/clap-htsat-unfused"""
            __lowerCamelCase						:       List[Any]				   =				tempfile.mkdtemp()




       def 					a_					(					self      : Optional[int]	,    **A__      : Union[str, Any]					):


            """simple docstring"""







            return RobertaTokenizer.from_pretrained(self.checkpoint	,    **A__					)




       def 					a_					(					self      : str	,    **A__      : Any					):


            """simple docstring"""







            return ClapFeatureExtractor.from_pretrained(self.checkpoint	,    **A__					)




       def 					a_					(					self      : List[str]					):


            """simple docstring"""







            shutil.rmtree(self.tmpdirname					)




       def 					a_					(					self      : str					):


            """simple docstring"""







            __lowerCamelCase						:       Dict				   =				self.get_tokenizer()
            __lowerCamelCase						:       Union[str, Any]				   =				self.get_feature_extractor()

            __lowerCamelCase						:       Dict				   =				ClapProcessor(tokenizer=A__	,    feature_extractor=A__					)

            processor.save_pretrained(self.tmpdirname					)
            __lowerCamelCase						:       Optional[Any]				   =				ClapProcessor.from_pretrained(self.tmpdirname					)

            self.assertEqual(processor.tokenizer.get_vocab()	,    tokenizer.get_vocab()					)
            self.assertIsInstance(processor.tokenizer	,    A__					)

            self.assertEqual(processor.feature_extractor.to_json_string()	,    feature_extractor.to_json_string()					)
            self.assertIsInstance(processor.feature_extractor	,    A__					)




       def 					a_					(					self      : str					):


            """simple docstring"""







            __lowerCamelCase						:       str				   =				ClapProcessor(tokenizer=self.get_tokenizer()	,    feature_extractor=self.get_feature_extractor()					)
            processor.save_pretrained(self.tmpdirname					)

            __lowerCamelCase						:       Tuple				   =				self.get_tokenizer(bos_token="""(BOS)"""	,    eos_token="""(EOS)"""					)
            __lowerCamelCase						:       Any				   =				self.get_feature_extractor(do_normalize=A__	,    padding_value=1.0					)

            __lowerCamelCase						:       List[Any]				   =				ClapProcessor.from_pretrained(
                self.tmpdirname	,    bos_token="""(BOS)"""	,    eos_token="""(EOS)"""	,    do_normalize=A__	,    padding_value=1.0					)

            self.assertEqual(processor.tokenizer.get_vocab()	,    tokenizer_add_kwargs.get_vocab()					)
            self.assertIsInstance(processor.tokenizer	,    A__					)

            self.assertEqual(processor.feature_extractor.to_json_string()	,    feature_extractor_add_kwargs.to_json_string()					)
            self.assertIsInstance(processor.feature_extractor	,    A__					)




       def 					a_					(					self      : Optional[int]					):


            """simple docstring"""







            __lowerCamelCase						:       Dict				   =				self.get_feature_extractor()
            __lowerCamelCase						:       Tuple				   =				self.get_tokenizer()

            __lowerCamelCase						:       Optional[int]				   =				ClapProcessor(tokenizer=A__	,    feature_extractor=A__					)

            __lowerCamelCase						:       Dict				   =				floats_list((3, 1000)					)

            __lowerCamelCase						:       str				   =				feature_extractor(A__	,    return_tensors="""np"""					)
            __lowerCamelCase						:       Any				   =				processor(audios=A__	,    return_tensors="""np"""					)

            for key in input_feat_extract.keys():
                 self.assertAlmostEqual(input_feat_extract[key].sum()	,    input_processor[key].sum()	,    delta=1e-2					)




       def 					a_					(					self      : Any					):


            """simple docstring"""







            __lowerCamelCase						:       Optional[Any]				   =				self.get_feature_extractor()
            __lowerCamelCase						:       int				   =				self.get_tokenizer()

            __lowerCamelCase						:       str				   =				ClapProcessor(tokenizer=A__	,    feature_extractor=A__					)

            __lowerCamelCase						:       int				   =				"""This is a test string"""

            __lowerCamelCase						:       Any				   =				processor(text=A__					)

            __lowerCamelCase						:       Any				   =				tokenizer(A__					)

            for key in encoded_tok.keys():
                 self.assertListEqual(encoded_tok[key]	,    encoded_processor[key]					)




       def 					a_					(					self      : List[Any]					):


            """simple docstring"""







            __lowerCamelCase						:       Optional[int]				   =				self.get_feature_extractor()
            __lowerCamelCase						:       List[str]				   =				self.get_tokenizer()

            __lowerCamelCase						:       int				   =				ClapProcessor(tokenizer=A__	,    feature_extractor=A__					)

            __lowerCamelCase						:       Union[str, Any]				   =				[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

            __lowerCamelCase						:       int				   =				processor.batch_decode(A__					)
            __lowerCamelCase						:       List[Any]				   =				tokenizer.batch_decode(A__					)

            self.assertListEqual(A__	,    A__					)




       def 					a_					(					self      : str					):


            """simple docstring"""







            __lowerCamelCase						:       Any				   =				self.get_feature_extractor()
            __lowerCamelCase						:       List[Any]				   =				self.get_tokenizer()

            __lowerCamelCase						:       Any				   =				ClapProcessor(tokenizer=A__	,    feature_extractor=A__					)

            self.assertListEqual(
                processor.model_input_names[2:]	,    feature_extractor.model_input_names	,    msg="""`processor` and `feature_extractor` model input names do not match"""	,    )



