



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      A_ (		__a       :			Tuple						):

				"""simple docstring"""
				if not (is_datasets_available() and is_faiss_available() and is_psutil_available()):
								a__			 =						unittest.skip("""test requires Datasets, Faiss, psutil"""						)(__a						)
				return test_case




@require_distributed_retrieval
class 	__snake_case (				SCREAMING_SNAKE_CASE):







							'''simple docstring'''






							def _a (							self						):
											a__			 =						tempfile.mkdtemp()
											a__			 =						8

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

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


							def _a (							self						):
											return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname					,					"""dpr_tokenizer"""						)						)


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


							def _a (							self						):
											shutil.rmtree(self.tmpdirname						)


							def _a (							self						):
											a__			 =						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					,					a_					,					a_=12_345						):
											a__			 =						self.get_dummy_dataset()
											a__			 =						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:
															a__			 =						dataset
															a__			 =						RagPyTorchDistributedRetriever(
															    a_					,					question_encoder_tokenizer=self.get_dpr_tokenizer()					,					generator_tokenizer=self.get_bart_tokenizer()					,					)
															if init_retrieval:
																			retriever.init_retrieval(a_						)
											return retriever


							def _a (							self					,					a_						):
											# 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=a_						)
											a__			 =						RagConfig(
											    retrieval_vector_size=self.retrieval_vector_size					,					question_encoder=DPRConfig().to_dict()					,					generator=BartConfig().to_dict()					,					)
											a__			 =						ray.remote(a_						)
											a__			 =						[remote_cls.remote() for _ in range(1						)]
											with patch("""transformers.models.rag.retrieval_rag.load_dataset"""						) as mock_load_dataset:
															a__			 =						self.get_dummy_dataset()
															a__			 =						RagRayDistributedRetriever(
															    a_					,					question_encoder_tokenizer=self.get_dpr_tokenizer()					,					generator_tokenizer=self.get_bart_tokenizer()					,					retrieval_workers=a_					,					)
															if init_retrieval:
																			retriever.init_retrieval()
											return retriever


							def _a (							self					,					a_					,					a_					,					a_=12_345						):
											a__			 =						self.get_dummy_dataset()
											a__			 =						RagConfig(
											    retrieval_vector_size=self.retrieval_vector_size					,					question_encoder=DPRConfig().to_dict()					,					generator=BartConfig().to_dict()					,					index_name="""custom"""					,					)
											if from_disk:
															a__			 =						os.path.join(self.tmpdirname					,					"""dataset"""						)
															a__			 =						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
															a__			 =						RagPyTorchDistributedRetriever(
															    a_					,					question_encoder_tokenizer=self.get_dpr_tokenizer()					,					generator_tokenizer=self.get_bart_tokenizer()					,					)
											else:
															a__			 =						RagPyTorchDistributedRetriever(
															    a_					,					question_encoder_tokenizer=self.get_dpr_tokenizer()					,					generator_tokenizer=self.get_bart_tokenizer()					,					index=CustomHFIndex(config.retrieval_vector_size					,					a_						)					,					)
											if init_retrieval:
															retriever.init_retrieval(a_						)
											return retriever


							def _a (							self					,					a_					,					a_						):
											# 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=a_						)
											a__			 =						self.get_dummy_dataset()
											a__			 =						RagConfig(
											    retrieval_vector_size=self.retrieval_vector_size					,					question_encoder=DPRConfig().to_dict()					,					generator=BartConfig().to_dict()					,					index_name="""custom"""					,					)
											a__			 =						ray.remote(a_						)
											a__			 =						[remote_cls.remote() for _ in range(1						)]
											if from_disk:
															a__			 =						os.path.join(self.tmpdirname					,					"""dataset"""						)
															a__			 =						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
															a__			 =						RagRayDistributedRetriever(
															    a_					,					question_encoder_tokenizer=self.get_dpr_tokenizer()					,					generator_tokenizer=self.get_bart_tokenizer()					,					retrieval_workers=a_					,					index=CustomHFIndex.load_from_disk(
															        vector_size=config.retrieval_vector_size					,					dataset_path=config.passages_path					,					index_path=config.index_path					,					)					,					)
											else:
															a__			 =						RagRayDistributedRetriever(
															    a_					,					question_encoder_tokenizer=self.get_dpr_tokenizer()					,					generator_tokenizer=self.get_bart_tokenizer()					,					retrieval_workers=a_					,					index=CustomHFIndex(config.retrieval_vector_size					,					a_						)					,					)
											if init_retrieval:
															retriever.init_retrieval()
											return retriever


							def _a (							self					,					a_					,					a_					,					a_						):
											a__ ,     a__ ,     a__			 =						retriever.retrieve(a_					,					n_docs=a_						)
											self.assertEqual(retrieved_doc_embeds.shape					,					(2, n_docs, self.retrieval_vector_size)						)
											self.assertEqual(len(a_						)					,					2						)
											self.assertEqual(sorted(doc_dicts[0]						)					,					["""embeddings""", """id""", """text""", """title"""]						)
											self.assertEqual(len(doc_dicts[0]["""id"""]						)					,					a_						)
											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						):
											a__			 =						1
											a__			 =						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=a_						)					,					a_					,					a_						)


							def _a (							self						):
											a__			 =						1
											a__			 =						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=a_					,					from_disk=a_						)					,					a_					,					a_					,					)


							def _a (							self						):
											a__			 =						1
											a__			 =						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=a_					,					from_disk=a_						)					,					a_					,					a_					,					)


							@require_ray
							def _a (							self						):
											a__			 =						1
											a__			 =						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=a_						)					,					a_					,					a_						)
											ray.shutdown()


							@require_ray
							def _a (							self						):
											a__			 =						1
											a__			 =						np.array(
											    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]					,					dtype=np.floataa						)
											with self.assertRaises(a_						):
															self.distributed_retriever_check(
															    self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=a_					,					from_disk=a_						)					,					a_					,					a_					,					)
											ray.shutdown()







							@require_ray
							def _a (							self						):
											a__			 =						1
											a__			 =						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=a_					,					from_disk=a_						)					,					a_					,					a_						)
											ray.shutdown()






import json
import os
import shutil
import tempfile
import unittest

from transformers import is_speech_available, is_torch_available
from transformers.models.speechta import SpeechTaTokenizer
from transformers.testing_utils import get_tests_dir, require_torch
from transformers.utils import FEATURE_EXTRACTOR_NAME


if is_speech_available() and is_torch_available():
	from transformers import SpeechTaFeatureExtractor, SpeechTaProcessor

	from .test_feature_extraction_speechta import floats_list


UpperCAmelCase								=   get_tests_dir("""fixtures/test_sentencepiece_bpe_char.model""")

@require_torch
class 	__snake_case (				unittest.TestCase):







							'''simple docstring'''






							def _a (							self						):
											a__			 =						tempfile.mkdtemp()

											a__			 =						SpeechTaTokenizer(a_						)
											tokenizer.save_pretrained(self.tmpdirname						)

											a__			 =						{
											    """feature_size""": 1,
											    """padding_value""": 0.0,
											    """sampling_rate""": 16_000,
											    """do_normalize""": False,
											    """num_mel_bins""": 80,
											    """hop_length""": 16,
											    """win_length""": 64,
											    """win_function""": """hann_window""",
											    """fmin""": 80,
											    """fmax""": 7_600,
											    """mel_floor""": 1E-1_0,
											    """reduction_factor""": 2,
											    """return_attention_mask""": True,
											}

											a__			 =						os.path.join(self.tmpdirname					,					a_						)
											with open(self.feature_extraction_file					,					"""w"""					,					encoding="""utf-8"""						) as fp:
															fp.write(json.dumps(a_						) + """\n"""						)


							def _a (							self					,					**a_						):
											return SpeechTaTokenizer.from_pretrained(self.tmpdirname					,					**a_						)


							def _a (							self					,					**a_						):
											return SpeechTaFeatureExtractor.from_pretrained(self.tmpdirname					,					**a_						)


							def _a (							self						):
											shutil.rmtree(self.tmpdirname						)


							def _a (							self						):
											a__			 =						self.get_tokenizer()
											a__			 =						self.get_feature_extractor()

											a__			 =						SpeechTaProcessor(tokenizer=a_					,					feature_extractor=a_						)

											processor.save_pretrained(self.tmpdirname						)
											a__			 =						SpeechTaProcessor.from_pretrained(self.tmpdirname						)

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

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


							def _a (							self						):
											a__			 =						SpeechTaProcessor(tokenizer=self.get_tokenizer()					,					feature_extractor=self.get_feature_extractor()						)
											processor.save_pretrained(self.tmpdirname						)

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

											a__			 =						SpeechTaProcessor.from_pretrained(
											    self.tmpdirname					,					bos_token="""(BOS)"""					,					eos_token="""(EOS)"""					,					do_normalize=a_					,					padding_value=1.0						)

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

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


							def _a (							self						):
											a__			 =						self.get_feature_extractor()
											a__			 =						self.get_tokenizer()

											a__			 =						SpeechTaProcessor(tokenizer=a_					,					feature_extractor=a_						)

											a__			 =						floats_list((3, 1_000)						)

											a__			 =						feature_extractor(audio=a_					,					return_tensors="""np"""						)
											a__			 =						processor(audio=a_					,					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						):
											a__			 =						self.get_feature_extractor()
											a__			 =						self.get_tokenizer()

											a__			 =						SpeechTaProcessor(tokenizer=a_					,					feature_extractor=a_						)

											a__			 =						floats_list((3, 1_000)						)

											a__			 =						feature_extractor(audio_target=a_					,					return_tensors="""np"""						)
											a__			 =						processor(audio_target=a_					,					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						):
											a__			 =						self.get_feature_extractor()
											a__			 =						self.get_tokenizer()

											a__			 =						SpeechTaProcessor(tokenizer=a_					,					feature_extractor=a_						)

											a__			 =						"""This is a test string"""

											a__			 =						processor(text=a_						)
											a__			 =						tokenizer(a_						)

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


							def _a (							self						):
											a__			 =						self.get_feature_extractor()
											a__			 =						self.get_tokenizer()

											a__			 =						SpeechTaProcessor(tokenizer=a_					,					feature_extractor=a_						)

											a__			 =						"""This is a test string"""

											a__			 =						processor(text_target=a_						)
											a__			 =						tokenizer(a_						)

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


							def _a (							self						):
											a__			 =						self.get_feature_extractor()
											a__			 =						self.get_tokenizer()

											a__			 =						SpeechTaProcessor(tokenizer=a_					,					feature_extractor=a_						)

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

											a__			 =						processor.batch_decode(a_						)
											a__			 =						tokenizer.batch_decode(a_						)

											self.assertListEqual(a_					,					a_						)







							def _a (							self						):
											a__			 =						self.get_feature_extractor()
											a__			 =						self.get_tokenizer()

											a__			 =						SpeechTaProcessor(tokenizer=a_					,					feature_extractor=a_						)

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

