




"""simple docstring"""




import random
import shutil
import tempfile
import unittest

import numpy as np

from transformers import TaTokenizer, TaTokenizerFast
from transformers.testing_utils import require_sentencepiece, require_torch
from transformers.utils.import_utils import is_speech_available, is_torch_available


if is_torch_available():
    pass

if is_speech_available():
    from transformers import EncodecFeatureExtractor, MusicgenProcessor


lowerCamelCase_  						=       random.Random()



def      UpperCAmelCase_   (     __UpperCamelCase,    __UpperCamelCase=1.0,    __UpperCamelCase=None,    __UpperCamelCase=None     ):





    if rng is None:
        SCREAMING_SNAKE_CASE__       =global_rng

    SCREAMING_SNAKE_CASE__       =[]
    for batch_idx in range(shape[0]     ):
        values.append([]     )
        for _ in range(shape[1]     ):
            values[-1].append(rng.random() * scale     )

    return values




@require_torch
@require_sentencepiece
class 				__a (			unittest.TestCase					):


     """simple docstring"""


     def   __A    (     self      :			Tuple						)		->       str:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__       ='facebook/musicgen-small'
         SCREAMING_SNAKE_CASE__       =tempfile.mkdtemp()




     def   __A    (     self      :			str    ,**_UpperCamelCase      :			str						)		->       Dict:


         '''simple docstring'''





         return TaTokenizer.from_pretrained(self.checkpoint    ,**lowercase_						)




     def   __A    (     self      :			str    ,**_UpperCamelCase      :			Optional[Any]						)		->       List[str]:


         '''simple docstring'''





         return EncodecFeatureExtractor.from_pretrained(self.checkpoint    ,**lowercase_						)




     def   __A    (     self      :			Dict						)		->       List[Any]:


         '''simple docstring'''





         shutil.rmtree(self.tmpdirname						)




     def   __A    (     self      :			Tuple						)		->       str:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__       =self.get_tokenizer()
         SCREAMING_SNAKE_CASE__       =self.get_feature_extractor()

         SCREAMING_SNAKE_CASE__       =MusicgenProcessor(tokenizer=lowercase_    ,feature_extractor=lowercase_						)

         processor.save_pretrained(self.tmpdirname						)
         SCREAMING_SNAKE_CASE__       =MusicgenProcessor.from_pretrained(self.tmpdirname						)

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

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




     def   __A    (     self      :			List[str]						)		->       int:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__       =MusicgenProcessor(tokenizer=self.get_tokenizer()    ,feature_extractor=self.get_feature_extractor()						)
         processor.save_pretrained(self.tmpdirname						)

         SCREAMING_SNAKE_CASE__       =self.get_tokenizer(bos_token="""(BOS)"""    ,eos_token="""(EOS)"""						)
         SCREAMING_SNAKE_CASE__       =self.get_feature_extractor(do_normalize=lowercase_    ,padding_value=1.0						)

         SCREAMING_SNAKE_CASE__       =MusicgenProcessor.from_pretrained(
             self.tmpdirname    ,bos_token="""(BOS)"""    ,eos_token="""(EOS)"""    ,do_normalize=lowercase_    ,padding_value=1.0						)

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

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




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


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__       =self.get_feature_extractor()
         SCREAMING_SNAKE_CASE__       =self.get_tokenizer()

         SCREAMING_SNAKE_CASE__       =MusicgenProcessor(tokenizer=lowercase_    ,feature_extractor=lowercase_						)

         SCREAMING_SNAKE_CASE__       =floats_list((3, 1_0_0_0)						)

         SCREAMING_SNAKE_CASE__       =feature_extractor(lowercase_    ,return_tensors="""np"""						)
         SCREAMING_SNAKE_CASE__       =processor(lowercase_    ,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      :			Optional[int]						)		->       List[str]:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__       =self.get_feature_extractor()
         SCREAMING_SNAKE_CASE__       =self.get_tokenizer()

         SCREAMING_SNAKE_CASE__       =MusicgenProcessor(tokenizer=lowercase_    ,feature_extractor=lowercase_						)

         SCREAMING_SNAKE_CASE__       ='This is a test string'

         SCREAMING_SNAKE_CASE__       =processor(text=lowercase_						)

         SCREAMING_SNAKE_CASE__       =tokenizer(lowercase_						)

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




     def   __A    (     self      :			List[str]						)		->       int:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__       =self.get_feature_extractor()
         SCREAMING_SNAKE_CASE__       =self.get_tokenizer()

         SCREAMING_SNAKE_CASE__       =MusicgenProcessor(tokenizer=lowercase_    ,feature_extractor=lowercase_						)

         SCREAMING_SNAKE_CASE__       =[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

         SCREAMING_SNAKE_CASE__       =processor.batch_decode(sequences=lowercase_						)
         SCREAMING_SNAKE_CASE__       =tokenizer.batch_decode(lowercase_						)

         self.assertListEqual(lowercase_    ,lowercase_						)




     def   __A    (     self      :			List[Any]						)		->       Optional[Any]:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__       =self.get_feature_extractor()
         SCREAMING_SNAKE_CASE__       =self.get_tokenizer()

         SCREAMING_SNAKE_CASE__       =MusicgenProcessor(tokenizer=lowercase_    ,feature_extractor=lowercase_						)

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




     def   __A    (     self      :			Dict						)		->       Tuple:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__       =self.get_feature_extractor(padding_side="""left"""						)
         SCREAMING_SNAKE_CASE__       =self.get_tokenizer()

         SCREAMING_SNAKE_CASE__       =MusicgenProcessor(tokenizer=lowercase_    ,feature_extractor=lowercase_						)

         SCREAMING_SNAKE_CASE__       =[floats_list((1, x)						)[0] for x in range(5    ,2_0    ,5						)]
         SCREAMING_SNAKE_CASE__       =processor(lowercase_						).padding_mask

         SCREAMING_SNAKE_CASE__       =np.asarray(floats_list((3, 2_0)						)						)[:, None, :]
         SCREAMING_SNAKE_CASE__       =processor.batch_decode(lowercase_    ,padding_mask=lowercase_						)

         self.assertIsInstance(lowercase_    ,lowercase_						)

         for audio in decoded_audios:
             self.assertIsInstance(lowercase_    ,np.ndarray						)

         self.assertTrue(decoded_audios[0].shape == (1, 1_0)						)
         self.assertTrue(decoded_audios[1].shape == (1, 1_5)						)
         self.assertTrue(decoded_audios[2].shape == (1, 2_0)						)












import logging
from pathlib import Path

import numpy as np
import pytorch_lightning as pl
import torch
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from pytorch_lightning.utilities import rank_zero_only

from utils import save_json



def      UpperCAmelCase_   (     __UpperCamelCase     ):
    SCREAMING_SNAKE_CASE__       =filter(lambda __UpperCamelCase     : p.requires_grad,    model.parameters()     )
    SCREAMING_SNAKE_CASE__       =sum([np.prod(p.size()     ) for p in model_parameters]     )
    return params


lowerCamelCase_  						=       logging.getLogger(__name__)
class 				__a (			pl.Callback					):


     """simple docstring"""


     def   __A    (     self      :			Optional[int]    ,_UpperCamelCase      :			Optional[Any]    ,_UpperCamelCase      :			Any						)		->       int:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__       ={f"""lr_group_{i}""": param["""lr"""] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups						)}
         pl_module.logger.log_metrics(_UpperCamelCase						)




     @rank_zero_only
     def   __A    (     self      :			Optional[int]    ,_UpperCamelCase      :			pl.Trainer    ,_UpperCamelCase      :			pl.LightningModule    ,_UpperCamelCase      :			str    ,_UpperCamelCase      :			Tuple=True						)		->       None:


         '''simple docstring'''





         logger.info(f"""***** {type_path} results at step {trainer.global_step:05d} *****"""						)
         SCREAMING_SNAKE_CASE__       =trainer.callback_metrics
         trainer.logger.log_metrics({k: v for k, v in metrics.items() if k not in ["""log""", """progress_bar""", """preds"""]}						)
         # Log results
         SCREAMING_SNAKE_CASE__       =Path(pl_module.hparams.output_dir						)
         if type_path == "test":
             SCREAMING_SNAKE_CASE__       =od / """test_results.txt"""
             SCREAMING_SNAKE_CASE__       =od / """test_generations.txt"""
         else:
             # this never gets hit. I prefer not to save intermediate generations, and results are in metrics.json
             # If people want this it will be easy enough to add back.
             SCREAMING_SNAKE_CASE__       =od / f"""{type_path}_results/{trainer.global_step:05d}.txt"""
             SCREAMING_SNAKE_CASE__       =od / f"""{type_path}_generations/{trainer.global_step:05d}.txt"""
             results_file.parent.mkdir(exist_ok=_UpperCamelCase						)
             generations_file.parent.mkdir(exist_ok=_UpperCamelCase						)
         with open(_UpperCamelCase    ,"""a+"""						) as writer:
             for key in sorted(_UpperCamelCase						):
                 if key in ["log", "progress_bar", "preds"]:
                     continue
                 SCREAMING_SNAKE_CASE__       =metrics[key]
                 if isinstance(_UpperCamelCase    ,torch.Tensor						):
                     SCREAMING_SNAKE_CASE__       =val.item()
                 SCREAMING_SNAKE_CASE__       =f"""{key}: {val:.6f}\n"""
                 writer.write(_UpperCamelCase						)

         if not save_generations:
             return

         if "preds" in metrics:
             SCREAMING_SNAKE_CASE__       ="""\n""".join(metrics["""preds"""]						)
             generations_file.open("""w+"""						).write(_UpperCamelCase						)




     @rank_zero_only
     def   __A    (     self      :			Tuple    ,_UpperCamelCase      :			int    ,_UpperCamelCase      :			Dict						)		->       Dict:


         '''simple docstring'''





         try:
             SCREAMING_SNAKE_CASE__       =pl_module.model.model.num_parameters()
         except AttributeError:
             SCREAMING_SNAKE_CASE__       =pl_module.model.num_parameters()

         SCREAMING_SNAKE_CASE__       =count_trainable_parameters(_UpperCamelCase						)
         # mp stands for million parameters
         trainer.logger.log_metrics({"""n_params""": npars, """mp""": npars / 1e6, """grad_mp""": n_trainable_pars / 1e6}						)




     @rank_zero_only
     def   __A    (     self      :			Optional[Any]    ,_UpperCamelCase      :			pl.Trainer    ,_UpperCamelCase      :			pl.LightningModule						)		->       Any:


         '''simple docstring'''





         save_json(pl_module.metrics    ,pl_module.metrics_save_path						)
         return self._write_logs(_UpperCamelCase    ,_UpperCamelCase    ,"""test"""						)




     @rank_zero_only
     def   __A    (     self      :			Optional[Any]    ,_UpperCamelCase      :			pl.Trainer    ,_UpperCamelCase      :			List[Any]						)		->       Dict:


         '''simple docstring'''





         save_json(pl_module.metrics    ,pl_module.metrics_save_path						)
         # Uncommenting this will save val generations
         # return self._write_logs(trainer, pl_module, "valid")






def      UpperCAmelCase_   (     __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase=1,    __UpperCamelCase=False     ):
    if metric == "rouge2":
        SCREAMING_SNAKE_CASE__       ="""{val_avg_rouge2:.4f}-{step_count}"""
    elif metric == "bleu":
        SCREAMING_SNAKE_CASE__       ="""{val_avg_bleu:.4f}-{step_count}"""
    elif metric == "loss":
        SCREAMING_SNAKE_CASE__       ="""{val_avg_loss:.4f}-{step_count}"""
    else:
        raise NotImplementedError(
            f"""seq2seq callbacks only support rouge2, bleu and loss, got {metric}, You can make your own by adding to"""
            """ this function."""     )

    SCREAMING_SNAKE_CASE__       =ModelCheckpoint(
        dirpath=__UpperCamelCase,    filename=__UpperCamelCase,    monitor=f"""val_{metric}""",    mode="""min""" if """loss""" in metric else """max""",    save_top_k=__UpperCamelCase,    )
    return checkpoint_callback




def      UpperCAmelCase_   (     __UpperCamelCase,    __UpperCamelCase     ):
    return EarlyStopping(
        monitor=f"""val_{metric}""",    mode="""min""" if """loss""" in metric else """max""",    patience=__UpperCamelCase,    verbose=__UpperCamelCase,    )






