






import json
import os
import pickle
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers import is_faiss_available
from transformers.models.bart.configuration_bart import BartConfig
from transformers.models.bart.tokenization_bart import BartTokenizer
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
from transformers.models.dpr.configuration_dpr import DPRConfig
from transformers.models.dpr.tokenization_dpr import DPRContextEncoderTokenizer, DPRQuestionEncoderTokenizer
from transformers.models.rag.configuration_rag import RagConfig
from transformers.models.rag.retrieval_rag import CustomHFIndex, RagRetriever
from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
from transformers.testing_utils import require_faiss, require_sentencepiece, require_tokenizers, require_torch


if is_faiss_available():
   import faiss


@require_faiss
class 			SCREAMING_SNAKE_CASE      (      lowerCAmelCase       ):
     '''simple docstring'''


     def 					_A							(					self							:						List[str]   ):
       SCREAMING_SNAKE_CASE							:   Any						    =     tempfile.mkdtemp()
       SCREAMING_SNAKE_CASE							:   Any						    =     8

       # DPR tok
       SCREAMING_SNAKE_CASE							:   str						    =     [
           "[UNK]",
           "[CLS]",
           "[SEP]",
           "[PAD]",
           "[MASK]",
           "want",
           "##want",
           "##ed",
           "wa",
           "un",
           "runn",
           "##ing",
           ",",
           "low",
           "lowest",
       ]
       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     os.path.join(self.tmpdirname       ,			"dpr_tokenizer"   )
       os.makedirs(UpperCAmelCase_       ,			exist_ok=UpperCAmelCase_   )
       SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     os.path.join(UpperCAmelCase_       ,			DPR_VOCAB_FILES_NAMES["vocab_file"]   )
       with open(self.vocab_file       ,			"w"       ,			encoding="utf-8"   ) as vocab_writer:
         vocab_writer.write("".join([x + "\n" for x in vocab_tokens]   )   )

       # BART tok
       SCREAMING_SNAKE_CASE							:   List[str]						    =     [
           "l",
           "o",
           "w",
           "e",
           "r",
           "s",
           "t",
           "i",
           "d",
           "n",
           "\u0120",
           "\u0120l",
           "\u0120n",
           "\u0120lo",
           "\u0120low",
           "er",
           "\u0120lowest",
           "\u0120newer",
           "\u0120wider",
           "<unk>",
       ]
       SCREAMING_SNAKE_CASE							:   Tuple						    =     dict(zip(UpperCAmelCase_       ,			range(len(UpperCAmelCase_   )   )   )   )
       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     ["#version: 0.2", "\u0120 l", "\u0120l o", "\u0120lo w", "e r", ""]
       SCREAMING_SNAKE_CASE							:   Tuple						    =     {"unk_token": "<unk>"}

       SCREAMING_SNAKE_CASE							:   List[str]						    =     os.path.join(self.tmpdirname       ,			"bart_tokenizer"   )
       os.makedirs(UpperCAmelCase_       ,			exist_ok=UpperCAmelCase_   )
       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     os.path.join(UpperCAmelCase_       ,			BART_VOCAB_FILES_NAMES["vocab_file"]   )
       SCREAMING_SNAKE_CASE							:   str						    =     os.path.join(UpperCAmelCase_       ,			BART_VOCAB_FILES_NAMES["merges_file"]   )
       with open(self.vocab_file       ,			"w"       ,			encoding="utf-8"   ) as fp:
         fp.write(json.dumps(UpperCAmelCase_   ) + "\n"   )
       with open(self.merges_file       ,			"w"       ,			encoding="utf-8"   ) as fp:
         fp.write("\n".join(UpperCAmelCase_   )   )


     def 					_A							(					self							:						Union[str, Any]   ):
       return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname       ,			"dpr_tokenizer"   )   )


     def 					_A							(					self							:						Tuple   ):
       return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname       ,			"dpr_tokenizer"   )   )


     def 					_A							(					self							:						Tuple   ):
       return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname       ,			"bart_tokenizer"   )   )


     def 					_A							(					self							:						Union[str, Any]   ):
       shutil.rmtree(self.tmpdirname   )


     def 					_A							(					self							:						int   ):
       SCREAMING_SNAKE_CASE							:   str						    =     Dataset.from_dict(
           {
               "id": ["0", "1"],
               "text": ["foo", "bar"],
               "title": ["Foo", "Bar"],
               "embeddings": [np.ones(self.retrieval_vector_size   ), 2 * np.ones(self.retrieval_vector_size   )],
           }   )
       dataset.add_faiss_index("embeddings"       ,			string_factory="Flat"       ,			metric_type=faiss.METRIC_INNER_PRODUCT   )
       return dataset


     def 					_A							(					self							:						Tuple   ):
       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     self.get_dummy_dataset()
       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     RagConfig(
           retrieval_vector_size=self.retrieval_vector_size       ,			question_encoder=DPRConfig().to_dict()       ,			generator=BartConfig().to_dict()       ,			)
       with patch("transformers.models.rag.retrieval_rag.load_dataset"   ) as mock_load_dataset:
         SCREAMING_SNAKE_CASE							:   List[Any]						    =     dataset
         SCREAMING_SNAKE_CASE							:   int						    =     RagRetriever(
             UpperCAmelCase_       ,			question_encoder_tokenizer=self.get_dpr_tokenizer()       ,			generator_tokenizer=self.get_bart_tokenizer()       ,			)
       return retriever


     def 					_A							(					self							:						List[Any]       ,			UpperCAmelCase_							:						bool   ):
       SCREAMING_SNAKE_CASE							:   List[str]						    =     self.get_dummy_dataset()
       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     RagConfig(
           retrieval_vector_size=self.retrieval_vector_size       ,			question_encoder=DPRConfig().to_dict()       ,			generator=BartConfig().to_dict()       ,			index_name="custom"       ,			)
       if from_disk:
         SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     os.path.join(self.tmpdirname       ,			"dataset"   )
         SCREAMING_SNAKE_CASE							:   Optional[int]						    =     os.path.join(self.tmpdirname       ,			"index.faiss"   )
         dataset.get_index("embeddings"   ).save(os.path.join(self.tmpdirname       ,			"index.faiss"   )   )
         dataset.drop_index("embeddings"   )
         dataset.save_to_disk(os.path.join(self.tmpdirname       ,			"dataset"   )   )
         del dataset
         SCREAMING_SNAKE_CASE							:   Tuple						    =     RagRetriever(
             UpperCAmelCase_       ,			question_encoder_tokenizer=self.get_dpr_tokenizer()       ,			generator_tokenizer=self.get_bart_tokenizer()       ,			)
       else:
         SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     RagRetriever(
             UpperCAmelCase_       ,			question_encoder_tokenizer=self.get_dpr_tokenizer()       ,			generator_tokenizer=self.get_bart_tokenizer()       ,			index=CustomHFIndex(config.retrieval_vector_size       ,			UpperCAmelCase_   )       ,			)
       return retriever


     def 					_A							(					self							:						Dict   ):
       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     Dataset.from_dict(
           {
               "id": ["0", "1"],
               "text": ["foo", "bar"],
               "title": ["Foo", "Bar"],
               "embeddings": [np.ones(self.retrieval_vector_size + 1   ), 2 * np.ones(self.retrieval_vector_size + 1   )],
           }   )
       dataset.add_faiss_index("embeddings"       ,			string_factory="Flat"       ,			metric_type=faiss.METRIC_INNER_PRODUCT   )

       SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     os.path.join(self.tmpdirname       ,			"hf_bert_base.hnswSQ8_correct_phi_128.c_index"   )
       dataset.save_faiss_index("embeddings"       ,			index_file_name + ".index.dpr"   )
       pickle.dump(dataset["id"]       ,			open(index_file_name + ".index_meta.dpr"       ,			"wb"   )   )

       SCREAMING_SNAKE_CASE							:   List[Any]						    =     os.path.join(self.tmpdirname       ,			"psgs_w100.tsv.pkl"   )
       SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     {sample["id"]: [sample["text"], sample["title"]] for sample in dataset}
       pickle.dump(UpperCAmelCase_       ,			open(UpperCAmelCase_       ,			"wb"   )   )

       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     RagConfig(
           retrieval_vector_size=self.retrieval_vector_size       ,			question_encoder=DPRConfig().to_dict()       ,			generator=BartConfig().to_dict()       ,			index_name="legacy"       ,			index_path=self.tmpdirname       ,			)
       SCREAMING_SNAKE_CASE							:   Dict						    =     RagRetriever(
           UpperCAmelCase_       ,			question_encoder_tokenizer=self.get_dpr_tokenizer()       ,			generator_tokenizer=self.get_bart_tokenizer()   )
       return retriever


     def 					_A							(					self							:						Optional[int]   ):
       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     1
       SCREAMING_SNAKE_CASE							:   str						    =     self.get_dummy_canonical_hf_index_retriever()
       SCREAMING_SNAKE_CASE							:   str						    =     np.array(
           [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]       ,			dtype=np.floataa   )
       SCREAMING_SNAKE_CASE    ,	SCREAMING_SNAKE_CASE    ,	SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     retriever.retrieve(UpperCAmelCase_       ,			n_docs=UpperCAmelCase_   )
       self.assertEqual(retrieved_doc_embeds.shape       ,			(2, n_docs, self.retrieval_vector_size)   )
       self.assertEqual(len(UpperCAmelCase_   )       ,			2   )
       self.assertEqual(sorted(doc_dicts[0]   )       ,			["embeddings", "id", "text", "title"]   )
       self.assertEqual(len(doc_dicts[0]["id"]   )       ,			UpperCAmelCase_   )
       self.assertEqual(doc_dicts[0]["id"][0]       ,			"1"   )  # max inner product is reached with second doc
       self.assertEqual(doc_dicts[1]["id"][0]       ,			"0"   )  # max inner product is reached with first doc
       self.assertListEqual(doc_ids.tolist()       ,			[[1], [0]]   )


     def 					_A							(					self							:						Union[str, Any]   ):
       SCREAMING_SNAKE_CASE							:   int						    =     self.get_dummy_canonical_hf_index_retriever()
       with tempfile.TemporaryDirectory() as tmp_dirname:
         with patch("transformers.models.rag.retrieval_rag.load_dataset"   ) as mock_load_dataset:
           SCREAMING_SNAKE_CASE							:   List[str]						    =     self.get_dummy_dataset()
           retriever.save_pretrained(UpperCAmelCase_   )
           SCREAMING_SNAKE_CASE							:   Tuple						    =     RagRetriever.from_pretrained(UpperCAmelCase_   )
           self.assertIsInstance(UpperCAmelCase_       ,			UpperCAmelCase_   )
           SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     np.array(
               [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]       ,			dtype=np.floataa   )
           SCREAMING_SNAKE_CASE							:   List[str]						    =     retriever.retrieve(UpperCAmelCase_       ,			n_docs=1   )
           self.assertTrue(out is not None   )


     def 					_A							(					self							:						Optional[int]   ):
       SCREAMING_SNAKE_CASE							:   Dict						    =     1
       SCREAMING_SNAKE_CASE							:   List[str]						    =     self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_   )
       SCREAMING_SNAKE_CASE							:   str						    =     np.array(
           [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]       ,			dtype=np.floataa   )
       SCREAMING_SNAKE_CASE    ,	SCREAMING_SNAKE_CASE    ,	SCREAMING_SNAKE_CASE							:   str						    =     retriever.retrieve(UpperCAmelCase_       ,			n_docs=UpperCAmelCase_   )
       self.assertEqual(retrieved_doc_embeds.shape       ,			(2, n_docs, self.retrieval_vector_size)   )
       self.assertEqual(len(UpperCAmelCase_   )       ,			2   )
       self.assertEqual(sorted(doc_dicts[0]   )       ,			["embeddings", "id", "text", "title"]   )
       self.assertEqual(len(doc_dicts[0]["id"]   )       ,			UpperCAmelCase_   )
       self.assertEqual(doc_dicts[0]["id"][0]       ,			"1"   )  # max inner product is reached with second doc
       self.assertEqual(doc_dicts[1]["id"][0]       ,			"0"   )  # max inner product is reached with first doc
       self.assertListEqual(doc_ids.tolist()       ,			[[1], [0]]   )


     def 					_A							(					self							:						List[Any]   ):
       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_   )
       with tempfile.TemporaryDirectory() as tmp_dirname:
         retriever.save_pretrained(UpperCAmelCase_   )
         SCREAMING_SNAKE_CASE							:   Dict						    =     RagRetriever.from_pretrained(UpperCAmelCase_   )
         self.assertIsInstance(UpperCAmelCase_       ,			UpperCAmelCase_   )
         SCREAMING_SNAKE_CASE							:   Tuple						    =     np.array(
             [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]       ,			dtype=np.floataa   )
         SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     retriever.retrieve(UpperCAmelCase_       ,			n_docs=1   )
         self.assertTrue(out is not None   )


     def 					_A							(					self							:						Any   ):
       SCREAMING_SNAKE_CASE							:   List[Any]						    =     1
       SCREAMING_SNAKE_CASE							:   Tuple						    =     self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_   )
       SCREAMING_SNAKE_CASE							:   Any						    =     np.array(
           [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]       ,			dtype=np.floataa   )
       SCREAMING_SNAKE_CASE    ,	SCREAMING_SNAKE_CASE    ,	SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     retriever.retrieve(UpperCAmelCase_       ,			n_docs=UpperCAmelCase_   )
       self.assertEqual(retrieved_doc_embeds.shape       ,			(2, n_docs, self.retrieval_vector_size)   )
       self.assertEqual(len(UpperCAmelCase_   )       ,			2   )
       self.assertEqual(sorted(doc_dicts[0]   )       ,			["embeddings", "id", "text", "title"]   )
       self.assertEqual(len(doc_dicts[0]["id"]   )       ,			UpperCAmelCase_   )
       self.assertEqual(doc_dicts[0]["id"][0]       ,			"1"   )  # max inner product is reached with second doc
       self.assertEqual(doc_dicts[1]["id"][0]       ,			"0"   )  # max inner product is reached with first doc
       self.assertListEqual(doc_ids.tolist()       ,			[[1], [0]]   )


     def 					_A							(					self							:						Optional[Any]   ):
       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_   )
       with tempfile.TemporaryDirectory() as tmp_dirname:
         retriever.save_pretrained(UpperCAmelCase_   )
         SCREAMING_SNAKE_CASE							:   Dict						    =     RagRetriever.from_pretrained(UpperCAmelCase_   )
         self.assertIsInstance(UpperCAmelCase_       ,			UpperCAmelCase_   )
         SCREAMING_SNAKE_CASE							:   List[Any]						    =     np.array(
             [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]       ,			dtype=np.floataa   )
         SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     retriever.retrieve(UpperCAmelCase_       ,			n_docs=1   )
         self.assertTrue(out is not None   )


     def 					_A							(					self							:						Optional[int]   ):
       SCREAMING_SNAKE_CASE							:   Tuple						    =     1
       SCREAMING_SNAKE_CASE							:   List[Any]						    =     self.get_dummy_legacy_index_retriever()
       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     np.array(
           [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]       ,			dtype=np.floataa   )
       SCREAMING_SNAKE_CASE    ,	SCREAMING_SNAKE_CASE    ,	SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     retriever.retrieve(UpperCAmelCase_       ,			n_docs=UpperCAmelCase_   )
       self.assertEqual(retrieved_doc_embeds.shape       ,			(2, n_docs, self.retrieval_vector_size)   )
       self.assertEqual(len(UpperCAmelCase_   )       ,			2   )
       self.assertEqual(sorted(doc_dicts[0]   )       ,			["text", "title"]   )
       self.assertEqual(len(doc_dicts[0]["text"]   )       ,			UpperCAmelCase_   )
       self.assertEqual(doc_dicts[0]["text"][0]       ,			"bar"   )  # max inner product is reached with second doc
       self.assertEqual(doc_dicts[1]["text"][0]       ,			"foo"   )  # max inner product is reached with first doc
       self.assertListEqual(doc_ids.tolist()       ,			[[1], [0]]   )


     def 					_A							(					self							:						Dict   ):
       SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     self.get_dummy_legacy_index_retriever()
       with tempfile.TemporaryDirectory() as tmp_dirname:
         retriever.save_pretrained(UpperCAmelCase_   )
         SCREAMING_SNAKE_CASE							:   Tuple						    =     RagRetriever.from_pretrained(UpperCAmelCase_   )
         self.assertIsInstance(UpperCAmelCase_       ,			UpperCAmelCase_   )
         SCREAMING_SNAKE_CASE							:   str						    =     np.array(
             [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]       ,			dtype=np.floataa   )
         SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     retriever.retrieve(UpperCAmelCase_       ,			n_docs=1   )
         self.assertTrue(out is not None   )


     @require_torch
     @require_tokenizers
     @require_sentencepiece
     def 					_A							(					self							:						List[Any]   ):
       import torch

       SCREAMING_SNAKE_CASE							:   List[Any]						    =     1
       SCREAMING_SNAKE_CASE							:   List[Any]						    =     self.get_dummy_canonical_hf_index_retriever()
       SCREAMING_SNAKE_CASE							:   Tuple						    =     [[5, 7], [10, 11]]
       SCREAMING_SNAKE_CASE							:   Tuple						    =     np.array(
           [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]       ,			dtype=np.floataa   )
       SCREAMING_SNAKE_CASE							:   str						    =     retriever(UpperCAmelCase_       ,			UpperCAmelCase_       ,			prefix=retriever.config.generator.prefix       ,			n_docs=UpperCAmelCase_   )
       SCREAMING_SNAKE_CASE    ,	SCREAMING_SNAKE_CASE    ,	SCREAMING_SNAKE_CASE							:   Optional[int]						    =     (
           out["context_input_ids"],
           out["context_attention_mask"],
           out["retrieved_doc_embeds"],
       )
       self.assertEqual(retrieved_doc_embeds.shape       ,			(2, n_docs, self.retrieval_vector_size)   )
       self.assertIsInstance(UpperCAmelCase_       ,			UpperCAmelCase_   )
       self.assertIsInstance(UpperCAmelCase_       ,			UpperCAmelCase_   )
       self.assertIsInstance(UpperCAmelCase_       ,			np.ndarray   )

       SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     retriever(
           UpperCAmelCase_       ,			UpperCAmelCase_       ,			prefix=retriever.config.generator.prefix       ,			n_docs=UpperCAmelCase_       ,			return_tensors="pt"       ,			)
       SCREAMING_SNAKE_CASE    ,	SCREAMING_SNAKE_CASE    ,	SCREAMING_SNAKE_CASE    ,	SCREAMING_SNAKE_CASE							:   List[str]						    =     (  # noqa: F841
           out["context_input_ids"],
           out["context_attention_mask"],
           out["retrieved_doc_embeds"],
           out["doc_ids"],
       )
       self.assertEqual(retrieved_doc_embeds.shape       ,			(2, n_docs, self.retrieval_vector_size)   )
       self.assertIsInstance(UpperCAmelCase_       ,			torch.Tensor   )
       self.assertIsInstance(UpperCAmelCase_       ,			torch.Tensor   )
       self.assertIsInstance(UpperCAmelCase_       ,			torch.Tensor   )






     @require_torch
     @require_tokenizers
     @require_sentencepiece
     def 					_A							(					self							:						Union[str, Any]   ):
       SCREAMING_SNAKE_CASE							:   Dict						    =     self.get_dpr_ctx_encoder_tokenizer()
       SCREAMING_SNAKE_CASE							:   str						    =     1
       SCREAMING_SNAKE_CASE							:   str						    =     self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_   )
       retriever.set_ctx_encoder_tokenizer(UpperCAmelCase_   )

       SCREAMING_SNAKE_CASE							:   List[str]						    =     [[5, 7], [10, 11]]
       SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     np.array(
           [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]       ,			dtype=np.floataa   )
       SCREAMING_SNAKE_CASE							:   str						    =     retriever(UpperCAmelCase_       ,			UpperCAmelCase_       ,			prefix=retriever.config.generator.prefix       ,			n_docs=UpperCAmelCase_   )

       self.assertEqual(
           len(UpperCAmelCase_   )       ,			6   )  # check whether the retriever output consist of 6 attributes including tokenized docs
       self.assertEqual(
           all(k in out for k in ("tokenized_doc_ids", "tokenized_doc_attention_mask")   )       ,			UpperCAmelCase_   )  # check for doc token related keys in dictionary.










'''simple docstring'''






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


lowercase				=logging.get_logger(__name__)






class    __magic_name__				(						lowerCAmelCase							):
       UpperCAmelCase		 =["input_features", "attention_mask"]






       def __init__(     self    ,				snake_case=8_0    ,				snake_case=1_6_0_0_0    ,				snake_case=8_0    ,				snake_case=0.0    ,				snake_case=True    ,				snake_case=True    ,				snake_case=True    ,				**snake_case    ,				)		->    str:




          '''simple docstring'''
          super().__init__(feature_size=snake_case    ,				sampling_rate=snake_case    ,				padding_value=snake_case    ,				**snake_case)
          _UpperCAmelCase  :  Optional[Any]				       =num_mel_bins
          _UpperCAmelCase  :  Optional[int]				       =do_ceptral_normalize
          _UpperCAmelCase  :  Optional[Any]				       =normalize_means
          _UpperCAmelCase  :  Optional[Any]				       =normalize_vars
          _UpperCAmelCase  :  Tuple				       =True






       def      lowerCAmelCase						(     self    ,				snake_case    ,				)		->    np.ndarray:




          '''simple docstring'''
          _UpperCAmelCase  :  List[Any]				       =waveform * (2**1_5)  # Kaldi compliance: 16-bit signed integers
          _UpperCAmelCase  :  Dict				       =torch.from_numpy(snake_case).unsqueeze(0)
          _UpperCAmelCase  :  Union[str, Any]				       =ta_kaldi.fbank(snake_case    ,				num_mel_bins=self.num_mel_bins    ,				sample_frequency=self.sampling_rate)
          return features.numpy()






       @staticmethod
       def      lowerCAmelCase						(     snake_case    ,				snake_case    ,				snake_case = True    ,				snake_case = True    ,				snake_case = 0.0    ,				)		->    np.ndarray:




          '''simple docstring'''
          # make sure we normalize float32 arrays
          if normalize_means:
             _UpperCAmelCase  :  int				       =x[:input_length].mean(axis=0)
             _UpperCAmelCase  :  List[Any]				       =np.subtract(snake_case    ,				snake_case)
          if normalize_vars:
             _UpperCAmelCase  :  List[Any]				       =x[:input_length].std(axis=0)
             _UpperCAmelCase  :  Optional[int]				       =np.divide(snake_case    ,				snake_case)

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

          # make sure array is in float32
          _UpperCAmelCase  :  str				       =x.astype(np.floataa)

          return x






       def      lowerCAmelCase						(     self    ,				snake_case    ,				snake_case = None)		->    List[np.ndarray]:




          '''simple docstring'''
          _UpperCAmelCase  :  str				       =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    ,				snake_case    ,				snake_case = False    ,				snake_case = None    ,				snake_case = False    ,				snake_case = None    ,				snake_case = None    ,				snake_case = None    ,				snake_case = None    ,				**snake_case    ,				)		->    BatchFeature:




          '''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.')

          _UpperCAmelCase  :  Optional[int]				       =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}")
          _UpperCAmelCase  :  List[Any]				       =is_batched_numpy or (
              isinstance(snake_case    ,				(list, tuple)) and (isinstance(raw_speech[0]    ,				(np.ndarray, tuple, list)))
          )

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

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

          # extract fbank features
          _UpperCAmelCase  :  Optional[Any]				       =[self._extract_fbank_features(snake_case) for waveform in raw_speech]

          # convert into correct format for padding
          _UpperCAmelCase  :  List[str]				       =BatchFeature({'input_features': features})

          _UpperCAmelCase  :  Any				       =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
          _UpperCAmelCase  :  Dict				       =padded_inputs.get('input_features')
          if isinstance(input_features[0]    ,				snake_case):
             _UpperCAmelCase  :  Any				       =[np.asarray(snake_case    ,				dtype=np.floataa) for feature in input_features]

          _UpperCAmelCase  :  int				       =padded_inputs.get('attention_mask')
          if attention_mask is not None:
             _UpperCAmelCase  :  Tuple				       =[np.asarray(snake_case    ,				dtype=np.intaa) for array in attention_mask]

          # Utterance-level cepstral mean and variance normalization
          if self.do_ceptral_normalize:
             _UpperCAmelCase  :  Optional[Any]				       =(
                 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
             )
             _UpperCAmelCase  :  Optional[int]				       =self.normalize(
                 padded_inputs['input_features']    ,				attention_mask=snake_case)

          if return_tensors is not None:
             _UpperCAmelCase  :  List[Any]				       =padded_inputs.convert_to_tensors(snake_case)

          return padded_inputs

