
"""simple docstring"""




import math
from typing import Any, Callable, List, Optional, Tuple, Union

import numpy as np
import torch

from ...models import TaFilmDecoder
from ...schedulers import DDPMScheduler
from ...utils import is_onnx_available, logging, randn_tensor


if is_onnx_available():
						from ..onnx_utils import OnnxRuntimeModel

from ..pipeline_utils import AudioPipelineOutput, DiffusionPipeline
from .continous_encoder import SpectrogramContEncoder
from .notes_encoder import SpectrogramNotesEncoder


lowerCAmelCase__				=					logging.get_logger(__name__)  # pylint: disable=invalid-name

lowerCAmelCase__				=					256





class 			snake_case  (				__lowercase ):
							UpperCAmelCase__      =   ['''melgan''']

							def __init__(self							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							):

									"""simple docstring"""







									super().__init__()

									# From MELGAN
									SCREAMING_SNAKE_CASE_    			=			math.log(1e-5     )  # Matches MelGAN training.
									SCREAMING_SNAKE_CASE_    			=			4.0  # Largest value for most examples
									SCREAMING_SNAKE_CASE_    			=			1_28

									self.register_modules(
									    notes_encoder=SCREAMING_SNAKE_CASE_							,							continuous_encoder=SCREAMING_SNAKE_CASE_							,							decoder=SCREAMING_SNAKE_CASE_							,							scheduler=SCREAMING_SNAKE_CASE_							,							melgan=SCREAMING_SNAKE_CASE_							,							)

							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_=(-1.0, 1.0)							,							SCREAMING_SNAKE_CASE_=False     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_    			=			output_range
									if clip:
											SCREAMING_SNAKE_CASE_    			=			torch.clip(SCREAMING_SNAKE_CASE_							,							self.min_value							,							self.max_value     )
									# Scale to [0, 1].
									SCREAMING_SNAKE_CASE_    			=			(features - self.min_value) / (self.max_value - self.min_value)
									# Scale to [min_out, max_out].
									return zero_one * (max_out - min_out) + min_out

							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_=(-1.0, 1.0)							,							SCREAMING_SNAKE_CASE_=False     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_    			=			input_range
									SCREAMING_SNAKE_CASE_    			=			torch.clip(SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_     ) if clip else outputs
									# Scale to [0, 1].
									SCREAMING_SNAKE_CASE_    			=			(outputs - min_out) / (max_out - min_out)
									# Scale to [self.min_value, self.max_value].
									return zero_one * (self.max_value - self.min_value) + self.min_value

							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			input_tokens > 0
									SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_    			=			self.notes_encoder(
									    encoder_input_tokens=SCREAMING_SNAKE_CASE_							,							encoder_inputs_mask=SCREAMING_SNAKE_CASE_     )

									SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_    			=			self.continuous_encoder(
									    encoder_inputs=SCREAMING_SNAKE_CASE_							,							encoder_inputs_mask=SCREAMING_SNAKE_CASE_     )

									return [(tokens_encoded, tokens_mask), (continuous_encoded, continuous_mask)]

							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			noise_time
									if not torch.is_tensor(SCREAMING_SNAKE_CASE_     ):
											SCREAMING_SNAKE_CASE_    			=			torch.tensor([timesteps]							,							dtype=torch.long							,							device=input_tokens.device     )
									elif torch.is_tensor(SCREAMING_SNAKE_CASE_     ) and len(timesteps.shape     ) == 0:
											SCREAMING_SNAKE_CASE_    			=			timesteps[None].to(input_tokens.device     )

									# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
									SCREAMING_SNAKE_CASE_    			=			timesteps * torch.ones(input_tokens.shape[0]							,							dtype=timesteps.dtype							,							device=timesteps.device     )

									SCREAMING_SNAKE_CASE_    			=			self.decoder(
									    encodings_and_masks=SCREAMING_SNAKE_CASE_							,							decoder_input_tokens=SCREAMING_SNAKE_CASE_							,							decoder_noise_time=SCREAMING_SNAKE_CASE_     )
									return logits







							@torch.no_grad()
							def __call__(self							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_ = None							,							SCREAMING_SNAKE_CASE_ = 1_00							,							SCREAMING_SNAKE_CASE_ = True							,							SCREAMING_SNAKE_CASE_ = "numpy"							,							SCREAMING_SNAKE_CASE_ = None							,							SCREAMING_SNAKE_CASE_ = 1							,							):

									"""simple docstring"""







									if (callback_steps is None) or (
									    callback_steps is not None and (not isinstance(SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_     ) or callback_steps <= 0)
									):
											raise ValueError(
											    f'`callback_steps` has to be a positive integer but is {callback_steps} of type'
											    f' {type(SCREAMING_SNAKE_CASE_     )}.'     )

									SCREAMING_SNAKE_CASE_    			=			np.zeros([1, TARGET_FEATURE_LENGTH, self.n_dims]							,							dtype=np.floataa     )
									SCREAMING_SNAKE_CASE_    			=			np.zeros([1, 0, self.n_dims]							,							np.floataa     )
									SCREAMING_SNAKE_CASE_    			=			torch.ones((1, TARGET_FEATURE_LENGTH)							,							dtype=SCREAMING_SNAKE_CASE_							,							device=self.device     )

									for i, encoder_input_tokens in enumerate(SCREAMING_SNAKE_CASE_     ):
											if i == 0:
													SCREAMING_SNAKE_CASE_    			=			torch.from_numpy(pred_mel[:1].copy()     ).to(
													    device=self.device							,							dtype=self.decoder.dtype     )
													# The first chunk has no previous context.
													SCREAMING_SNAKE_CASE_    			=			torch.zeros((1, TARGET_FEATURE_LENGTH)							,							dtype=SCREAMING_SNAKE_CASE_							,							device=self.device     )
											else:
													# The full song pipeline does not feed in a context feature, so the mask
													# will be all 0s after the feature converter. Because we know we're
													# feeding in a full context chunk from the previous prediction, set it
													# to all 1s.
													SCREAMING_SNAKE_CASE_    			=			ones

											SCREAMING_SNAKE_CASE_    			=			self.scale_features(
											    SCREAMING_SNAKE_CASE_							,							output_range=[-1.0, 1.0]							,							clip=SCREAMING_SNAKE_CASE_     )

											SCREAMING_SNAKE_CASE_    			=			self.encode(
											    input_tokens=torch.IntTensor([encoder_input_tokens]     ).to(device=self.device     )							,							continuous_inputs=SCREAMING_SNAKE_CASE_							,							continuous_mask=SCREAMING_SNAKE_CASE_							,							)

											# Sample encoder_continuous_inputs shaped gaussian noise to begin loop
											SCREAMING_SNAKE_CASE_    			=			randn_tensor(
											    shape=encoder_continuous_inputs.shape							,							generator=SCREAMING_SNAKE_CASE_							,							device=self.device							,							dtype=self.decoder.dtype							,							)

											# set step values
											self.scheduler.set_timesteps(SCREAMING_SNAKE_CASE_     )

											# Denoising diffusion loop
											for j, t in enumerate(self.progress_bar(self.scheduler.timesteps     )     ):
													SCREAMING_SNAKE_CASE_    			=			self.decode(
													    encodings_and_masks=SCREAMING_SNAKE_CASE_							,							input_tokens=SCREAMING_SNAKE_CASE_							,							noise_time=t / self.scheduler.config.num_train_timesteps							,							)

													# Compute previous output: x_t -> x_t-1
													SCREAMING_SNAKE_CASE_    			=			self.scheduler.step(SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							generator=SCREAMING_SNAKE_CASE_     ).prev_sample

											SCREAMING_SNAKE_CASE_    			=			self.scale_to_features(SCREAMING_SNAKE_CASE_							,							input_range=[-1.0, 1.0]     )
											SCREAMING_SNAKE_CASE_    			=			mel[:1]
											SCREAMING_SNAKE_CASE_    			=			mel.cpu().float().numpy()

											SCREAMING_SNAKE_CASE_    			=			np.concatenate([full_pred_mel, pred_mel[:1]]							,							axis=1     )

											# call the callback, if provided
											if callback is not None and i % callback_steps == 0:
													callback(SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_     )

											logger.info('''Generated segment'''							,							SCREAMING_SNAKE_CASE_     )

									if output_type == "numpy" and not is_onnx_available():
											raise ValueError(
											    '''Cannot return output in \'np\' format if ONNX is not available. Make sure to have ONNX installed or set \'output_type\' to \'mel\'.'''     )
									elif output_type == "numpy" and self.melgan is None:
											raise ValueError(
											    '''Cannot return output in \'np\' format if melgan component is not defined. Make sure to define `self.melgan` or set \'output_type\' to \'mel\'.'''     )

									if output_type == "numpy":
											SCREAMING_SNAKE_CASE_    			=			self.melgan(input_features=full_pred_mel.astype(np.floataa     )     )
									else:
											SCREAMING_SNAKE_CASE_    			=			full_pred_mel

									if not return_dict:
											return (output,)

									return AudioPipelineOutput(audios=SCREAMING_SNAKE_CASE_     )

"""simple docstring"""




import copy

from ...configuration_utils import PretrainedConfig
from ...utils import add_start_docstrings


lowerCAmelCase__				=					r'\n    [`RagConfig`] stores the configuration of a *RagModel*. Configuration objects inherit from [`PretrainedConfig`] and\n    can be used to control the model outputs. Read the documentation from [`PretrainedConfig`] for more information.\n\n    Args:\n        title_sep (`str`, *optional*, defaults to  `" / "`):\n            Separator inserted between the title and the text of the retrieved document when calling [`RagRetriever`].\n        doc_sep (`str`, *optional*, defaults to  `" // "`):\n            Separator inserted between the text of the retrieved document and the original input when calling\n            [`RagRetriever`].\n        n_docs (`int`, *optional*, defaults to 5):\n            Number of documents to retrieve.\n        max_combined_length (`int`, *optional*, defaults to 300):\n            Max length of contextualized input returned by [`~RagRetriever.__call__`].\n        retrieval_vector_size (`int`, *optional*, defaults to 768):\n            Dimensionality of the document embeddings indexed by [`RagRetriever`].\n        retrieval_batch_size (`int`, *optional*, defaults to 8):\n            Retrieval batch size, defined as the number of queries issues concurrently to the faiss index encapsulated\n            [`RagRetriever`].\n        dataset (`str`, *optional*, defaults to `"wiki_dpr"`):\n            A dataset identifier of the indexed dataset in HuggingFace Datasets (list all available datasets and ids\n            using `datasets.list_datasets()`).\n        dataset_split (`str`, *optional*, defaults to `"train"`)\n            Which split of the `dataset` to load.\n        index_name (`str`, *optional*, defaults to `"compressed"`)\n            The index name of the index associated with the `dataset`. One can choose between `"legacy"`, `"exact"` and\n            `"compressed"`.\n        index_path (`str`, *optional*)\n            The path to the serialized faiss index on disk.\n        passages_path (`str`, *optional*):\n            A path to text passages compatible with the faiss index. Required if using\n            [`~models.rag.retrieval_rag.LegacyIndex`]\n        use_dummy_dataset (`bool`, *optional*, defaults to `False`)\n            Whether to load a "dummy" variant of the dataset specified by `dataset`.\n        label_smoothing (`float`, *optional*, defaults to 0.0):\n            Only relevant if `return_loss` is set to `True`. Controls the `epsilon` parameter value for label smoothing\n            in the loss calculation. If set to 0, no label smoothing is performed.\n        do_marginalize (`bool`, *optional*, defaults to `False`):\n            If `True`, the logits are marginalized over all documents by making use of\n            `torch.nn.functional.log_softmax`.\n        reduce_loss (`bool`, *optional*, defaults to `False`):\n            Whether or not to reduce the NLL loss using the `torch.Tensor.sum` operation.\n        do_deduplication (`bool`, *optional*, defaults to `True`):\n            Whether or not to deduplicate the generations from different context documents for a given input. Has to be\n            set to `False` if used while training with distributed backend.\n        exclude_bos_score (`bool`, *optional*, defaults to `False`):\n            Whether or not to disregard the BOS token when computing the loss.\n        output_retrieved(`bool`, *optional*, defaults to `False`):\n            If set to `True`, `retrieved_doc_embeds`, `retrieved_doc_ids`, `context_input_ids` and\n            `context_attention_mask` are returned. See returned tensors for more detail.\n        use_cache (`bool`, *optional*, defaults to `True`):\n            Whether or not the model should return the last key/values attentions (not used by all models).\n        forced_eos_token_id (`int`, *optional*):\n            The id of the token to force as the last generated token when `max_length` is reached. Usually set to\n            `eos_token_id`.\n'





@add_start_docstrings(__lowercase )
class 			snake_case  (				__lowercase ):
							UpperCAmelCase__      =   '''rag'''
							UpperCAmelCase__      =   True

							def __init__(self							,							SCREAMING_SNAKE_CASE_=None							,							SCREAMING_SNAKE_CASE_=True							,							SCREAMING_SNAKE_CASE_=None							,							SCREAMING_SNAKE_CASE_=None							,							SCREAMING_SNAKE_CASE_=None							,							SCREAMING_SNAKE_CASE_=None							,							SCREAMING_SNAKE_CASE_=None							,							SCREAMING_SNAKE_CASE_=" / "							,							SCREAMING_SNAKE_CASE_=" // "							,							SCREAMING_SNAKE_CASE_=5							,							SCREAMING_SNAKE_CASE_=3_00							,							SCREAMING_SNAKE_CASE_=7_68							,							SCREAMING_SNAKE_CASE_=8							,							SCREAMING_SNAKE_CASE_="wiki_dpr"							,							SCREAMING_SNAKE_CASE_="train"							,							SCREAMING_SNAKE_CASE_="compressed"							,							SCREAMING_SNAKE_CASE_=None							,							SCREAMING_SNAKE_CASE_=None							,							SCREAMING_SNAKE_CASE_=False							,							SCREAMING_SNAKE_CASE_=False							,							SCREAMING_SNAKE_CASE_=0.0							,							SCREAMING_SNAKE_CASE_=True							,							SCREAMING_SNAKE_CASE_=False							,							SCREAMING_SNAKE_CASE_=False							,							SCREAMING_SNAKE_CASE_=False							,							SCREAMING_SNAKE_CASE_=True							,							SCREAMING_SNAKE_CASE_=None							,							**SCREAMING_SNAKE_CASE_							,							):

									"""simple docstring"""







									super().__init__(
									    bos_token_id=SCREAMING_SNAKE_CASE_							,							pad_token_id=SCREAMING_SNAKE_CASE_							,							eos_token_id=SCREAMING_SNAKE_CASE_							,							decoder_start_token_id=SCREAMING_SNAKE_CASE_							,							forced_eos_token_id=SCREAMING_SNAKE_CASE_							,							is_encoder_decoder=SCREAMING_SNAKE_CASE_							,							prefix=SCREAMING_SNAKE_CASE_							,							vocab_size=SCREAMING_SNAKE_CASE_							,							**SCREAMING_SNAKE_CASE_							,							)
									assert (
									    "question_encoder" in kwargs and "generator" in kwargs
									), "Config has to be initialized with question_encoder and generator config"
									SCREAMING_SNAKE_CASE_    			=			kwargs.pop('''question_encoder'''     )
									SCREAMING_SNAKE_CASE_    			=			question_encoder_config.pop('''model_type'''     )
									SCREAMING_SNAKE_CASE_    			=			kwargs.pop('''generator'''     )
									SCREAMING_SNAKE_CASE_    			=			decoder_config.pop('''model_type'''     )

									from ..auto.configuration_auto import AutoConfig

									SCREAMING_SNAKE_CASE_    			=			AutoConfig.for_model(SCREAMING_SNAKE_CASE_							,							**SCREAMING_SNAKE_CASE_     )
									SCREAMING_SNAKE_CASE_    			=			AutoConfig.for_model(SCREAMING_SNAKE_CASE_							,							**SCREAMING_SNAKE_CASE_     )

									SCREAMING_SNAKE_CASE_    			=			reduce_loss
									SCREAMING_SNAKE_CASE_    			=			label_smoothing
									SCREAMING_SNAKE_CASE_    			=			exclude_bos_score
									SCREAMING_SNAKE_CASE_    			=			do_marginalize

									SCREAMING_SNAKE_CASE_    			=			title_sep
									SCREAMING_SNAKE_CASE_    			=			doc_sep
									SCREAMING_SNAKE_CASE_    			=			n_docs
									SCREAMING_SNAKE_CASE_    			=			max_combined_length

									SCREAMING_SNAKE_CASE_    			=			dataset
									SCREAMING_SNAKE_CASE_    			=			dataset_split
									SCREAMING_SNAKE_CASE_    			=			index_name

									SCREAMING_SNAKE_CASE_    			=			retrieval_vector_size
									SCREAMING_SNAKE_CASE_    			=			retrieval_batch_size
									SCREAMING_SNAKE_CASE_    			=			passages_path
									SCREAMING_SNAKE_CASE_    			=			index_path
									SCREAMING_SNAKE_CASE_    			=			use_dummy_dataset

									SCREAMING_SNAKE_CASE_    			=			output_retrieved

									SCREAMING_SNAKE_CASE_    			=			do_deduplication

									SCREAMING_SNAKE_CASE_    			=			use_cache

									if self.forced_eos_token_id is None:
											SCREAMING_SNAKE_CASE_    			=			getattr(self.generator							,							'''forced_eos_token_id'''							,							SCREAMING_SNAKE_CASE_     )

							@classmethod
							def _lowercase       (cls							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							**SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







									return cls(question_encoder=question_encoder_config.to_dict()							,							generator=generator_config.to_dict()							,							**SCREAMING_SNAKE_CASE_     )







							def _lowercase       (self     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			copy.deepcopy(self.__dict__     )
									SCREAMING_SNAKE_CASE_    			=			self.question_encoder.to_dict()
									SCREAMING_SNAKE_CASE_    			=			self.generator.to_dict()
									SCREAMING_SNAKE_CASE_    			=			self.__class__.model_type
									return output