





from typing import List, Optional, Union

import numpy as np
import torch
import torchaudio.compliance.kaldi as ta_kaldi

from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
from ...feature_extraction_utils import BatchFeature
from ...utils import PaddingStrategy, TensorType, logging


_lowerCamelCase 					=      logging.get_logger(__name__)




class 	UpperCamelCase_							(  SCREAMING_SNAKE_CASE_     ):
							lowerCamelCase_							     =       ["input_features", "attention_mask"]


							def __init__(				self     :Dict				,					__A     :List[str]=80				,					__A     :int=1_6000				,					__A     :Dict=80				,					__A     :Optional[Any]=0.0				,					__A     :Optional[int]=True				,					__A     :Union[str, Any]=True				,					__A     :Any=True				,					**__A     :Optional[int]				,					) -> Optional[int]:



									"""simple docstring"""



									super().__init__(feature_size=snake_case__				,					sampling_rate=snake_case__				,					padding_value=snake_case__				,					**snake_case__		)
									SCREAMING_SNAKE_CASE__								=   num_mel_bins
									SCREAMING_SNAKE_CASE__								=   do_ceptral_normalize
									SCREAMING_SNAKE_CASE__								=   normalize_means
									SCREAMING_SNAKE_CASE__								=   normalize_vars
									SCREAMING_SNAKE_CASE__								=   True


							def   _snake_case	(				self     :Tuple				,					__A     :int				,					) -> Dict:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   waveform * (2**15)  # Kaldi compliance: 16-bit signed integers
									SCREAMING_SNAKE_CASE__								=   torch.from_numpy(snake_case__		).unsqueeze(0		)
									SCREAMING_SNAKE_CASE__								=   ta_kaldi.fbank(snake_case__				,					num_mel_bins=self.num_mel_bins				,					sample_frequency=self.sampling_rate		)
									return features.numpy()


							@staticmethod
							def   _snake_case	(				__A     :Tuple				,					__A     :Union[str, Any]				,					__A     :List[str] = True				,					__A     :Optional[Any] = True				,					__A     :List[str] = 0.0				,					) -> List[Any]:



									"""simple docstring"""



									if normalize_means:
											SCREAMING_SNAKE_CASE__								=   x[:input_length].mean(axis=0		)
											SCREAMING_SNAKE_CASE__								=   np.subtract(snake_case__				,					snake_case__		)
									if normalize_vars:
											SCREAMING_SNAKE_CASE__								=   x[:input_length].std(axis=0		)
											SCREAMING_SNAKE_CASE__								=   np.divide(snake_case__				,					snake_case__		)

									if input_length < x.shape[0]:
											SCREAMING_SNAKE_CASE__								=   padding_value

									# make sure array is in float32
									SCREAMING_SNAKE_CASE__								=   x.astype(np.floataa		)

									return x


							def   _snake_case	(				self     :Union[str, Any]				,					__A     :List[str]				,					__A     :int = None		) -> List[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   attention_mask.sum(-1		) if attention_mask is not None else [x.shape[0] for x in input_features]
									return [
									    self.utterance_cmvn(snake_case__				,					snake_case__				,					self.normalize_means				,					self.normalize_vars				,					self.padding_value		)
									    for x, n in zip(snake_case__				,					snake_case__		)
									]


							def __call__(				self     :List[Any]				,					__A     :Dict				,					__A     :Optional[Any] = False				,					__A     :Union[str, Any] = None				,					__A     :int = False				,					__A     :List[Any] = None				,					__A     :Tuple = None				,					__A     :Any = None				,					__A     :Optional[Any] = None				,					**__A     :List[str]				,					) -> List[Any]:



									"""simple docstring"""



									if sampling_rate is not None:
											if sampling_rate != self.sampling_rate:
													raise ValueError(
													    f'''The model corresponding to this feature extractor: {self} was trained using a sampling rate of'''
													    f''' {self.sampling_rate}. Please make sure that the provided `raw_speech` input was sampled with'''
													    f''' {self.sampling_rate} and not {sampling_rate}.'''		)
									else:
											logger.warning(
											    """It is strongly recommended to pass the `sampling_rate` argument to this function. """
											    """Failing to do so can result in silent errors that might be hard to debug."""		)

									SCREAMING_SNAKE_CASE__								=   isinstance(snake_case__				,					np.ndarray		) and len(raw_speech.shape		) > 1
									if is_batched_numpy and len(raw_speech.shape		) > 2:
											raise ValueError(f'''Only mono-channel audio is supported for input to {self}'''		)
									SCREAMING_SNAKE_CASE__								=   is_batched_numpy or (
									    isinstance(snake_case__				,					(list, tuple)		) and (isinstance(raw_speech[0]				,					(np.ndarray, tuple, list)		))
									)

									if is_batched:
											SCREAMING_SNAKE_CASE__								=   [np.asarray(snake_case__				,					dtype=np.floataa		) for speech in raw_speech]
									elif not is_batched and not isinstance(snake_case__				,					np.ndarray		):
											SCREAMING_SNAKE_CASE__								=   np.asarray(snake_case__				,					dtype=np.floataa		)
									elif isinstance(snake_case__				,					np.ndarray		) and raw_speech.dtype is np.dtype(np.floataa		):
											SCREAMING_SNAKE_CASE__								=   raw_speech.astype(np.floataa		)

									# always return batch
									if not is_batched:
											SCREAMING_SNAKE_CASE__								=   [raw_speech]

									# extract fbank features
									SCREAMING_SNAKE_CASE__								=   [self._extract_fbank_features(snake_case__		) for waveform in raw_speech]

									# convert into correct format for padding
									SCREAMING_SNAKE_CASE__								=   BatchFeature({"""input_features""": features}		)

									SCREAMING_SNAKE_CASE__								=   self.pad(
									    snake_case__				,					padding=snake_case__				,					max_length=snake_case__				,					truncation=snake_case__				,					pad_to_multiple_of=snake_case__				,					return_attention_mask=snake_case__				,					**snake_case__				,					)

									# make sure list is in array format
									SCREAMING_SNAKE_CASE__								=   padded_inputs.get("""input_features"""		)
									if isinstance(input_features[0]				,					snake_case__		):
											SCREAMING_SNAKE_CASE__								=   [np.asarray(snake_case__				,					dtype=np.floataa		) for feature in input_features]

									SCREAMING_SNAKE_CASE__								=   padded_inputs.get("""attention_mask"""		)
									if attention_mask is not None:
											SCREAMING_SNAKE_CASE__								=   [np.asarray(snake_case__				,					dtype=np.intaa		) for array in attention_mask]

									# Utterance-level cepstral mean and variance normalization
									if self.do_ceptral_normalize:
											SCREAMING_SNAKE_CASE__								=   (
											    np.array(snake_case__				,					dtype=np.intaa		)
											    if self._get_padding_strategies(snake_case__				,					max_length=snake_case__		) is not PaddingStrategy.DO_NOT_PAD
											    else None
											)
											SCREAMING_SNAKE_CASE__								=   self.normalize(
											    padded_inputs["""input_features"""]				,					attention_mask=snake_case__		)

									if return_tensors is not None:
											SCREAMING_SNAKE_CASE__								=   padded_inputs.convert_to_tensors(snake_case__		)

									return padded_inputs



'''simple docstring'''







import copy

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


lowerCAmelCase     :  Tuple											=    r"""
    [`RagConfig`] stores the configuration of a *RagModel*. Configuration objects inherit from [`PretrainedConfig`] and
    can be used to control the model outputs. Read the documentation from [`PretrainedConfig`] for more information.

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



     """simple docstring"""






     __magic_name__           =     "rag"
     __magic_name__           =     True







     def __init__(			self   ,						snake_case__=None   ,						snake_case__=True   ,						snake_case__=None   ,						snake_case__=None   ,						snake_case__=None   ,						snake_case__=None   ,						snake_case__=None   ,						snake_case__=" / "   ,						snake_case__=" // "   ,						snake_case__=5   ,						snake_case__=300   ,						snake_case__=768   ,						snake_case__=8   ,						snake_case__="wiki_dpr"   ,						snake_case__="train"   ,						snake_case__="compressed"   ,						snake_case__=None   ,						snake_case__=None   ,						snake_case__=False   ,						snake_case__=False   ,						snake_case__=0.0   ,						snake_case__=True   ,						snake_case__=False   ,						snake_case__=False   ,						snake_case__=False   ,						snake_case__=True   ,						snake_case__=None   ,						**snake_case__   ,						):

            '''simple docstring'''

            super().__init__(
                bos_token_id=snake_case__   ,						pad_token_id=snake_case__   ,						eos_token_id=snake_case__   ,						decoder_start_token_id=snake_case__   ,						forced_eos_token_id=snake_case__   ,						is_encoder_decoder=snake_case__   ,						prefix=snake_case__   ,						vocab_size=snake_case__   ,						**snake_case__   ,						)
            assert (
                "question_encoder" in kwargs and "generator" in kwargs
            ), "Config has to be initialized with question_encoder and generator config"
            _lowerCAmelCase     :				Optional[Any]       =							kwargs.pop('question_encoder' )
            _lowerCAmelCase     :				List[Any]       =							question_encoder_config.pop('model_type' )
            _lowerCAmelCase     :				List[str]       =							kwargs.pop('generator' )
            _lowerCAmelCase     :				Tuple       =							decoder_config.pop('model_type' )

            from ..auto.configuration_auto import AutoConfig

            _lowerCAmelCase     :				Union[str, Any]       =							AutoConfig.for_model(snake_case__   ,						**snake_case__ )
            _lowerCAmelCase     :				List[str]       =							AutoConfig.for_model(snake_case__   ,						**snake_case__ )

            _lowerCAmelCase     :				Optional[int]       =							reduce_loss
            _lowerCAmelCase     :				Optional[Any]       =							label_smoothing
            _lowerCAmelCase     :				Dict       =							exclude_bos_score
            _lowerCAmelCase     :				Optional[int]       =							do_marginalize

            _lowerCAmelCase     :				List[str]       =							title_sep
            _lowerCAmelCase     :				Optional[Any]       =							doc_sep
            _lowerCAmelCase     :				str       =							n_docs
            _lowerCAmelCase     :				Optional[int]       =							max_combined_length

            _lowerCAmelCase     :				Dict       =							dataset
            _lowerCAmelCase     :				Optional[int]       =							dataset_split
            _lowerCAmelCase     :				str       =							index_name

            _lowerCAmelCase     :				Tuple       =							retrieval_vector_size
            _lowerCAmelCase     :				Any       =							retrieval_batch_size
            _lowerCAmelCase     :				Any       =							passages_path
            _lowerCAmelCase     :				Tuple       =							index_path
            _lowerCAmelCase     :				Any       =							use_dummy_dataset

            _lowerCAmelCase     :				Union[str, Any]       =							output_retrieved

            _lowerCAmelCase     :				List[str]       =							do_deduplication

            _lowerCAmelCase     :				Optional[int]       =							use_cache

            if self.forced_eos_token_id is None:
                   _lowerCAmelCase     :				Tuple       =							getattr(self.generator   ,						'forced_eos_token_id'   ,						snake_case__ )







     @classmethod
     def a						(			cls   ,						snake_case__   ,						snake_case__   ,						**snake_case__ ):

            '''simple docstring'''

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







     def a						(			self ):

            '''simple docstring'''

            _lowerCAmelCase     :				Union[str, Any]       =							copy.deepcopy(self.__dict__ )
            _lowerCAmelCase     :				List[str]       =							self.question_encoder.to_dict()
            _lowerCAmelCase     :				Optional[Any]       =							self.generator.to_dict()
            _lowerCAmelCase     :				Optional[int]       =							self.__class__.model_type
            return output

