




import textwrap

import pyarrow as pa
import pytest

from datasets import Features, Image
from datasets.packaged_modules.text.text import Text

from ..utils import require_pil
@pytest.fixture
def 							__lowerCAmelCase			(UpperCAmelCase_		):
    a__  			=							tmp_path / """text.txt"""
    a__  			=							textwrap.dedent(
        """\
        Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.
        Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.
        Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur.
        Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.

        Second paragraph:
        Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.
        Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.
        Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur.
        Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.
        """		)
    with open(UpperCamelCase__ ,     """w""" ,     encoding="""utf-8"""		) as f:
        f.write(UpperCamelCase__		)
    return str(UpperCamelCase__		)
@pytest.fixture
def 							__lowerCAmelCase			(UpperCAmelCase_ ,     UpperCAmelCase_		):
    a__  			=							tmp_path / """text_with_image.txt"""
    with open(UpperCamelCase__ ,     """w""" ,     encoding="""utf-8"""		) as f:
        f.write(UpperCamelCase__		)
    return str(UpperCamelCase__		)
@pytest.mark.parametrize("""keep_linebreaks""" ,     [True, False]		)
def 							__lowerCAmelCase			(UpperCAmelCase_ ,     UpperCAmelCase_		):


    with open(UpperCamelCase__ ,     encoding="""utf-8"""		) as f:
        a__  			=							f.read().splitlines(keepends=UpperCamelCase__		)
    a__  			=							Text(keep_linebreaks=UpperCamelCase__ ,     encoding="""utf-8"""		)
    a__  			=							text._generate_tables([[text_file]]		)
    a__  			=							pa.concat_tables([table for _, table in generator]		).to_pydict()["""text"""]
    assert generated_content == expected_content
@require_pil
def 							__lowerCAmelCase			(UpperCAmelCase_		):


    with open(UpperCamelCase__ ,     encoding="""utf-8"""		) as f:
        a__  			=							f.read().splitlines()[0]
    a__  			=							Text(encoding="""utf-8""" ,     features=Features({"""image""": Image()}		)		)
    a__  			=							text._generate_tables([[text_file_with_image]]		)
    a__  			=							pa.concat_tables([table for _, table in generator]		)
    assert pa_table.schema.field("""image"""		).type == Image()()
    a__  			=							pa_table.to_pydict()["""image"""]
    assert generated_content == [{"path": image_file, "bytes": None}]







@pytest.mark.parametrize("""sample_by""" ,     ["""line""", """paragraph""", """document"""]		)
def 							__lowerCAmelCase			(UpperCAmelCase_ ,     UpperCAmelCase_		):


    with open(UpperCamelCase__ ,     encoding="""utf-8"""		) as f:
        a__  			=							f.read()
    if sample_by == "line":
        a__  			=							expected_content.splitlines()
    elif sample_by == "paragraph":
        a__  			=							expected_content.split("""\n\n"""		)
    elif sample_by == "document":
        a__  			=							[expected_content]
    a__  			=							Text(sample_by=UpperCamelCase__ ,     encoding="""utf-8""" ,     chunksize=100		)
    a__  			=							text._generate_tables([[text_file]]		)
    a__  			=							pa.concat_tables([table for _, table in generator]		).to_pydict()["""text"""]
    assert generated_content == expected_content












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:
	__A						:		str          =						None

if is_ray_available():
	import ray  # noqa: E402 # isort:skip
	from distributed_ray_retriever import RagRayDistributedRetriever, RayRetriever  # noqa: E402 # isort:skip
else:
	__A						:		int          =						None
	__A						:		Tuple          =						None
	__A						:		Union[str, Any]          =						None





def      UpperCAmelCase  (						UpperCamelCase__   )     ->     Union[str, Any]:







				'''simple docstring'''


				if not (is_datasets_available() and is_faiss_available() and is_psutil_available()):
								__lowerCAmelCase					=  unittest.skip("""test requires Datasets, Faiss, psutil"""   )(UpperCamelCase__   )
				return test_case

@require_distributed_retrieval
class 		lowercase_       (	lowerCAmelCase__				):




					def     _lowercase						(   self:		List[Any]):


									'''simple docstring'''


									__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(_lowercase,     exist_ok=_lowercase)
									__lowerCAmelCase					=  os.path.join(_lowercase,     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(_lowercase,     range(len(_lowercase))))
									__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(_lowercase,     exist_ok=_lowercase)
									__lowerCAmelCase					=  os.path.join(_lowercase,     BART_VOCAB_FILES_NAMES["""vocab_file"""])
									__lowerCAmelCase					=  os.path.join(_lowercase,     BART_VOCAB_FILES_NAMES["""merges_file"""])
									with open(self.vocab_file,     """w""",     encoding="""utf-8""") as fp:
													fp.write(json.dumps(_lowercase) + """\n""")
									with open(self.merges_file,     """w""",     encoding="""utf-8""") as fp:
													fp.write("""\n""".join(_lowercase))




					def     _lowercase						(   self:		Optional[Any]):


									'''simple docstring'''


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




					def     _lowercase						(   self:		int):


									'''simple docstring'''


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




					def     _lowercase						(   self:		Optional[int]):


									'''simple docstring'''


									shutil.rmtree(self.tmpdirname)




					def     _lowercase						(   self:		Tuple):


									'''simple docstring'''


									__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     _lowercase						(   self:		Union[str, Any],     _lowercase:		bool,     _lowercase:		Any=12345):


									'''simple docstring'''


									__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(
													    _lowercase,     question_encoder_tokenizer=self.get_dpr_tokenizer(),     generator_tokenizer=self.get_bart_tokenizer(),     )
													if init_retrieval:
																	retriever.init_retrieval(_lowercase)
									return retriever




					def     _lowercase						(   self:		Union[str, Any],     _lowercase:		bool):


									'''simple docstring'''


									ray.init(local_mode=_lowercase)
									__lowerCAmelCase					=  RagConfig(
									    retrieval_vector_size=self.retrieval_vector_size,     question_encoder=DPRConfig().to_dict(),     generator=BartConfig().to_dict(),     )
									__lowerCAmelCase					=  ray.remote(_lowercase)
									__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(
													    _lowercase,     question_encoder_tokenizer=self.get_dpr_tokenizer(),     generator_tokenizer=self.get_bart_tokenizer(),     retrieval_workers=_lowercase,     )
													if init_retrieval:
																	retriever.init_retrieval()
									return retriever




					def     _lowercase						(   self:		Optional[Any],     _lowercase:		bool,     _lowercase:		bool,     _lowercase:		int=12345):


									'''simple docstring'''


									__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(
													    _lowercase,     question_encoder_tokenizer=self.get_dpr_tokenizer(),     generator_tokenizer=self.get_bart_tokenizer(),     )
									else:
													__lowerCAmelCase					=  RagPyTorchDistributedRetriever(
													    _lowercase,     question_encoder_tokenizer=self.get_dpr_tokenizer(),     generator_tokenizer=self.get_bart_tokenizer(),     index=CustomHFIndex(config.retrieval_vector_size,     _lowercase),     )
									if init_retrieval:
													retriever.init_retrieval(_lowercase)
									return retriever




					def     _lowercase						(   self:		List[Any],     _lowercase:		bool,     _lowercase:		bool):


									'''simple docstring'''


									ray.init(local_mode=_lowercase)
									__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(_lowercase)
									__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(
													    _lowercase,     question_encoder_tokenizer=self.get_dpr_tokenizer(),     generator_tokenizer=self.get_bart_tokenizer(),     retrieval_workers=_lowercase,     index=CustomHFIndex.load_from_disk(
													        vector_size=config.retrieval_vector_size,     dataset_path=config.passages_path,     index_path=config.index_path,     ),     )
									else:
													__lowerCAmelCase					=  RagRayDistributedRetriever(
													    _lowercase,     question_encoder_tokenizer=self.get_dpr_tokenizer(),     generator_tokenizer=self.get_bart_tokenizer(),     retrieval_workers=_lowercase,     index=CustomHFIndex(config.retrieval_vector_size,     _lowercase),     )
									if init_retrieval:
													retriever.init_retrieval()
									return retriever




					def     _lowercase						(   self:		int,     _lowercase:		RagRetriever,     _lowercase:		np.array,     _lowercase:		int):


									'''simple docstring'''


									__lowerCAmelCase	,				__lowerCAmelCase	,				__lowerCAmelCase					=  retriever.retrieve(_lowercase,     n_docs=_lowercase)
									self.assertEqual(retrieved_doc_embeds.shape,     (2, n_docs, self.retrieval_vector_size))
									self.assertEqual(len(_lowercase),     2)
									self.assertEqual(sorted(doc_dicts[0]),     ["""embeddings""", """id""", """text""", """title"""])
									self.assertEqual(len(doc_dicts[0]["""id"""]),     _lowercase)
									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     _lowercase						(   self:		Optional[Any]):


									'''simple docstring'''


									__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=_lowercase),     _lowercase,     _lowercase)




					def     _lowercase						(   self:		int):


									'''simple docstring'''


									__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=_lowercase,     from_disk=_lowercase),     _lowercase,     _lowercase,     )




					def     _lowercase						(   self:		Dict):


									'''simple docstring'''


									__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=_lowercase,     from_disk=_lowercase),     _lowercase,     _lowercase,     )




					@require_ray
					def     _lowercase						(   self:		List[str]):


									'''simple docstring'''


									__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=_lowercase),     _lowercase,     _lowercase)
									ray.shutdown()




					@require_ray
					def     _lowercase						(   self:		Union[str, Any]):


									'''simple docstring'''


									__lowerCAmelCase					=  1
									__lowerCAmelCase					=  np.array(
									    [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)],     dtype=np.floataa)
									with self.assertRaises(_lowercase):
													self.distributed_retriever_check(
													    self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=_lowercase,     from_disk=_lowercase),     _lowercase,     _lowercase,     )
									ray.shutdown()







					@require_ray
					def     _lowercase						(   self:		Tuple):


									'''simple docstring'''


									__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=_lowercase,     from_disk=_lowercase),     _lowercase,     _lowercase)
									ray.shutdown()





