


'''simple docstring'''






import warnings
from contextlib import contextmanager

from ...processing_utils import ProcessorMixin




class 		a__				(							a__  ):





							'''simple docstring'''






							lowercase__			:   Dict							=      "AutoFeatureExtractor"
							lowercase__			:   int							=      "Speech2Text2Tokenizer"







							def __init__(		self					,    lowerCamelCase_					,    lowerCamelCase_   )       ->							int:
														super().__init__(lowerCamelCase_					,    lowerCamelCase_   )
														lowerCAmelCase__          =	self.feature_extractor
														lowerCAmelCase__          =	False







							def __call__(		self					,    *lowerCamelCase_					,    **lowerCamelCase_   )       ->							Optional[Any]:
														# For backward compatibility
														if self._in_target_context_manager:
																					return self.current_processor(*lowerCamelCase_					,    **lowerCamelCase_   )

														if "raw_speech" in kwargs:
																					warnings.warn('''Using `raw_speech` as a keyword argument is deprecated. Use `audio` instead.'''   )
																					lowerCAmelCase__          =	kwargs.pop('''raw_speech'''   )
														else:
																					lowerCAmelCase__          =	kwargs.pop('''audio'''					,    lowerCamelCase_   )
														lowerCAmelCase__          =	kwargs.pop('''sampling_rate'''					,    lowerCamelCase_   )
														lowerCAmelCase__          =	kwargs.pop('''text'''					,    lowerCamelCase_   )
														if len(lowerCamelCase_   ) > 0:
																					lowerCAmelCase__          =	args[0]
																					lowerCAmelCase__          =	args[1:]

														if audio is None and text is None:
																					raise ValueError('''You need to specify either an `audio` or `text` input to process.'''   )

														if audio is not None:
																					lowerCAmelCase__          =	self.feature_extractor(lowerCamelCase_					,    *lowerCamelCase_					,    sampling_rate=lowerCamelCase_					,    **lowerCamelCase_   )
														if text is not None:
																					lowerCAmelCase__          =	self.tokenizer(lowerCamelCase_					,    **lowerCamelCase_   )

														if text is None:
																					return inputs
														elif audio is None:
																					return encodings
														else:
																					lowerCAmelCase__          =	encodings['''input_ids''']
																					return inputs







							def   __SCREAMING_SNAKE_CASE							(		self					,    *lowerCamelCase_					,    **lowerCamelCase_   )       ->							str:
														return self.tokenizer.batch_decode(*lowerCamelCase_					,    **lowerCamelCase_   )







							def   __SCREAMING_SNAKE_CASE							(		self					,    *lowerCamelCase_					,    **lowerCamelCase_   )       ->							List[Any]:
														return self.tokenizer.decode(*lowerCamelCase_					,    **lowerCamelCase_   )







							@contextmanager
							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Tuple:
														warnings.warn(
														    '''`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your '''
														    '''labels by using the argument `text` of the regular `__call__` method (either in the same call as '''
														    '''your audio inputs, or in a separate call.'''   )
														lowerCAmelCase__          =	True
														lowerCAmelCase__          =	self.tokenizer
														yield
														lowerCAmelCase__          =	self.feature_extractor
														lowerCAmelCase__          =	False



'''simple docstring'''






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

import faiss
import numpy as np
from datasets import Dataset

from transformers import BartConfig, BartTokenizer, DPRConfig, DPRQuestionEncoderTokenizer, RagConfig
from transformers.file_utils import is_datasets_available, is_faiss_available, is_psutil_available, is_torch_available
from transformers.integrations import is_ray_available
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
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_ray


sys.path.append(os.path.join(os.getcwd()))  # noqa: E402 # noqa: E402 # isort:skip

if is_torch_available():
		from distributed_pytorch_retriever import RagPyTorchDistributedRetriever  # noqa: E402 # isort:skip
else:
		__UpperCAmelCase    				=							None

if is_ray_available():
		import ray  # noqa: E402 # isort:skip
		from distributed_ray_retriever import RagRayDistributedRetriever, RayRetriever  # noqa: E402 # isort:skip
else:
		__UpperCAmelCase    				=							None
		__UpperCAmelCase    				=							None
		__UpperCAmelCase    				=							None


def 					_snake_case   (     A       )						->    Any:
							if not (is_datasets_available() and is_faiss_available() and is_psutil_available()):
														lowerCAmelCase__          =	unittest.skip('''test requires Datasets, Faiss, psutil'''       )(A       )
							return test_case






@require_distributed_retrieval
class 		a__				(							a__  ):





							'''simple docstring'''






							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Union[str, Any]:
														lowerCAmelCase__          =	tempfile.mkdtemp()
														lowerCAmelCase__          =	8

														# DPR tok
														lowerCAmelCase__          =	[
														    '''[UNK]''',
														    '''[CLS]''',
														    '''[SEP]''',
														    '''[PAD]''',
														    '''[MASK]''',
														    '''want''',
														    '''##want''',
														    '''##ed''',
														    '''wa''',
														    '''un''',
														    '''runn''',
														    '''##ing''',
														    ''',''',
														    '''low''',
														    '''lowest''',
														]
														lowerCAmelCase__          =	os.path.join(self.tmpdirname					,    '''dpr_tokenizer'''   )
														os.makedirs(lowerCamelCase_					,    exist_ok=lowerCamelCase_   )
														lowerCAmelCase__          =	os.path.join(lowerCamelCase_					,    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
														lowerCAmelCase__          =	[
														    '''l''',
														    '''o''',
														    '''w''',
														    '''e''',
														    '''r''',
														    '''s''',
														    '''t''',
														    '''i''',
														    '''d''',
														    '''n''',
														    '''\u0120''',
														    '''\u0120l''',
														    '''\u0120n''',
														    '''\u0120lo''',
														    '''\u0120low''',
														    '''er''',
														    '''\u0120lowest''',
														    '''\u0120newer''',
														    '''\u0120wider''',
														    '''<unk>''',
														]
														lowerCAmelCase__          =	dict(zip(lowerCamelCase_					,    range(len(lowerCamelCase_   )   )   )   )
														lowerCAmelCase__          =	['''#version: 0.2''', '''\u0120 l''', '''\u0120l o''', '''\u0120lo w''', '''e r''', '''''']
														lowerCAmelCase__          =	{'''unk_token''': '''<unk>'''}

														lowerCAmelCase__          =	os.path.join(self.tmpdirname					,    '''bart_tokenizer'''   )
														os.makedirs(lowerCamelCase_					,    exist_ok=lowerCamelCase_   )
														lowerCAmelCase__          =	os.path.join(lowerCamelCase_					,    BART_VOCAB_FILES_NAMES['''vocab_file''']   )
														lowerCAmelCase__          =	os.path.join(lowerCamelCase_					,    BART_VOCAB_FILES_NAMES['''merges_file''']   )
														with open(self.vocab_file					,    '''w'''					,    encoding='''utf-8'''   ) as fp:
																					fp.write(json.dumps(lowerCamelCase_   ) + '''\n'''   )
														with open(self.merges_file					,    '''w'''					,    encoding='''utf-8'''   ) as fp:
																					fp.write('''\n'''.join(lowerCamelCase_   )   )







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							DPRQuestionEncoderTokenizer:
														return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname					,    '''dpr_tokenizer'''   )   )







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							BartTokenizer:
														return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname					,    '''bart_tokenizer'''   )   )







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Optional[int]:
														shutil.rmtree(self.tmpdirname   )







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Tuple:
														lowerCAmelCase__          =	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   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_					,    lowerCamelCase_=1_23_45   )       ->							RagPyTorchDistributedRetriever:
														lowerCAmelCase__          =	self.get_dummy_dataset()
														lowerCAmelCase__          =	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:
																					lowerCAmelCase__          =	dataset
																					lowerCAmelCase__          =	RagPyTorchDistributedRetriever(
																					    lowerCamelCase_					,    question_encoder_tokenizer=self.get_dpr_tokenizer()					,    generator_tokenizer=self.get_bart_tokenizer()					,    )
																					if init_retrieval:
																												retriever.init_retrieval(lowerCamelCase_   )
														return retriever







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_   )       ->							RagRayDistributedRetriever:
														# Have to run in local mode because sys.path modifications at top of
														# file are not propogated to remote workers.
														# https://stackoverflow.com/questions/54338013/parallel-import-a-python-file-from-sibling-folder
														ray.init(local_mode=lowerCamelCase_   )
														lowerCAmelCase__          =	RagConfig(
														    retrieval_vector_size=self.retrieval_vector_size					,    question_encoder=DPRConfig().to_dict()					,    generator=BartConfig().to_dict()					,    )
														lowerCAmelCase__          =	ray.remote(lowerCamelCase_   )
														lowerCAmelCase__          =	[remote_cls.remote() for _ in range(1   )]
														with patch('''transformers.models.rag.retrieval_rag.load_dataset'''   ) as mock_load_dataset:
																					lowerCAmelCase__          =	self.get_dummy_dataset()
																					lowerCAmelCase__          =	RagRayDistributedRetriever(
																					    lowerCamelCase_					,    question_encoder_tokenizer=self.get_dpr_tokenizer()					,    generator_tokenizer=self.get_bart_tokenizer()					,    retrieval_workers=lowerCamelCase_					,    )
																					if init_retrieval:
																												retriever.init_retrieval()
														return retriever







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_					,    lowerCamelCase_					,    lowerCamelCase_=1_23_45   )       ->							Union[str, Any]:
														lowerCAmelCase__          =	self.get_dummy_dataset()
														lowerCAmelCase__          =	RagConfig(
														    retrieval_vector_size=self.retrieval_vector_size					,    question_encoder=DPRConfig().to_dict()					,    generator=BartConfig().to_dict()					,    index_name='''custom'''					,    )
														if from_disk:
																					lowerCAmelCase__          =	os.path.join(self.tmpdirname					,    '''dataset'''   )
																					lowerCAmelCase__          =	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
																					lowerCAmelCase__          =	RagPyTorchDistributedRetriever(
																					    lowerCamelCase_					,    question_encoder_tokenizer=self.get_dpr_tokenizer()					,    generator_tokenizer=self.get_bart_tokenizer()					,    )
														else:
																					lowerCAmelCase__          =	RagPyTorchDistributedRetriever(
																					    lowerCamelCase_					,    question_encoder_tokenizer=self.get_dpr_tokenizer()					,    generator_tokenizer=self.get_bart_tokenizer()					,    index=CustomHFIndex(config.retrieval_vector_size					,    lowerCamelCase_   )					,    )
														if init_retrieval:
																					retriever.init_retrieval(lowerCamelCase_   )
														return retriever







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_					,    lowerCamelCase_   )       ->							Union[str, Any]:
														# Have to run in local mode because sys.path modifications at top of
														# file are not propogated to remote workers.
														# https://stackoverflow.com/questions/54338013/parallel-import-a-python-file-from-sibling-folder
														ray.init(local_mode=lowerCamelCase_   )
														lowerCAmelCase__          =	self.get_dummy_dataset()
														lowerCAmelCase__          =	RagConfig(
														    retrieval_vector_size=self.retrieval_vector_size					,    question_encoder=DPRConfig().to_dict()					,    generator=BartConfig().to_dict()					,    index_name='''custom'''					,    )
														lowerCAmelCase__          =	ray.remote(lowerCamelCase_   )
														lowerCAmelCase__          =	[remote_cls.remote() for _ in range(1   )]
														if from_disk:
																					lowerCAmelCase__          =	os.path.join(self.tmpdirname					,    '''dataset'''   )
																					lowerCAmelCase__          =	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
																					lowerCAmelCase__          =	RagRayDistributedRetriever(
																					    lowerCamelCase_					,    question_encoder_tokenizer=self.get_dpr_tokenizer()					,    generator_tokenizer=self.get_bart_tokenizer()					,    retrieval_workers=lowerCamelCase_					,    index=CustomHFIndex.load_from_disk(
																					        vector_size=config.retrieval_vector_size					,    dataset_path=config.passages_path					,    index_path=config.index_path					,    )					,    )
														else:
																					lowerCAmelCase__          =	RagRayDistributedRetriever(
																					    lowerCamelCase_					,    question_encoder_tokenizer=self.get_dpr_tokenizer()					,    generator_tokenizer=self.get_bart_tokenizer()					,    retrieval_workers=lowerCamelCase_					,    index=CustomHFIndex(config.retrieval_vector_size					,    lowerCamelCase_   )					,    )
														if init_retrieval:
																					retriever.init_retrieval()
														return retriever







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_					,    lowerCamelCase_					,    lowerCamelCase_   )       ->							None:
														lowerCAmelCase__   ,					lowerCAmelCase__   ,					lowerCAmelCase__          =	retriever.retrieve(lowerCamelCase_					,    n_docs=lowerCamelCase_   )
														self.assertEqual(retrieved_doc_embeds.shape					,    (2, n_docs, self.retrieval_vector_size)   )
														self.assertEqual(len(lowerCamelCase_   )					,    2   )
														self.assertEqual(sorted(doc_dicts[0]   )					,    ['''embeddings''', '''id''', '''text''', '''title''']   )
														self.assertEqual(len(doc_dicts[0]['''id''']   )					,    lowerCamelCase_   )
														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   __SCREAMING_SNAKE_CASE							(		self   )       ->							Optional[Any]:
														lowerCAmelCase__          =	1
														lowerCAmelCase__          =	np.array(
														    [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]					,    dtype=np.floataa   )

														self.distributed_retriever_check(
														    self.get_dummy_pytorch_distributed_retriever(init_retrieval=lowerCamelCase_   )					,    lowerCamelCase_					,    lowerCamelCase_   )







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Any:
														lowerCAmelCase__          =	1
														lowerCAmelCase__          =	np.array(
														    [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]					,    dtype=np.floataa   )

														self.distributed_retriever_check(
														    self.get_dummy_custom_hf_index_pytorch_retriever(init_retrieval=lowerCamelCase_					,    from_disk=lowerCamelCase_   )					,    lowerCamelCase_					,    lowerCamelCase_					,    )







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Dict:
														lowerCAmelCase__          =	1
														lowerCAmelCase__          =	np.array(
														    [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]					,    dtype=np.floataa   )

														self.distributed_retriever_check(
														    self.get_dummy_custom_hf_index_pytorch_retriever(init_retrieval=lowerCamelCase_					,    from_disk=lowerCamelCase_   )					,    lowerCamelCase_					,    lowerCamelCase_					,    )







							@require_ray
							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							List[Any]:
														lowerCAmelCase__          =	1
														lowerCAmelCase__          =	np.array(
														    [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]					,    dtype=np.floataa   )

														self.distributed_retriever_check(
														    self.get_dummy_ray_distributed_retriever(init_retrieval=lowerCamelCase_   )					,    lowerCamelCase_					,    lowerCamelCase_   )
														ray.shutdown()







							@require_ray
							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							int:
														lowerCAmelCase__          =	1
														lowerCAmelCase__          =	np.array(
														    [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]					,    dtype=np.floataa   )
														with self.assertRaises(lowerCamelCase_   ):
																					self.distributed_retriever_check(
																					    self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=lowerCamelCase_					,    from_disk=lowerCamelCase_   )					,    lowerCamelCase_					,    lowerCamelCase_					,    )
														ray.shutdown()







							@require_ray
							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							List[str]:
														lowerCAmelCase__          =	1
														lowerCAmelCase__          =	np.array(
														    [np.ones(self.retrieval_vector_size   ), -np.ones(self.retrieval_vector_size   )]					,    dtype=np.floataa   )

														self.distributed_retriever_check(
														    self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=lowerCamelCase_					,    from_disk=lowerCamelCase_   )					,    lowerCamelCase_					,    lowerCamelCase_   )
														ray.shutdown()