





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:
							__lowerCAmelCase			:Any	 =     None

if is_ray_available():
							import ray  # noqa: E402 # isort:skip
							from distributed_ray_retriever import RagRayDistributedRetriever, RayRetriever  # noqa: E402 # isort:skip
else:
							__lowerCAmelCase			:List[str]	 =     None
							__lowerCAmelCase			:Optional[int]	 =     None
							__lowerCAmelCase			:Dict	 =     None







def  A					(				UpperCAmelCase					):
					if not (is_datasets_available() and is_faiss_available() and is_psutil_available()):
										_snake_case							:       Dict												= unittest.skip("test requires Datasets, Faiss, psutil"					)(UpperCAmelCase					)
					return test_case






@require_distributed_retrieval
class 							_a(     __A     ):







		def 							lowercase							(     self					)   ->   List[str]:


							'''simple docstring'''

							_snake_case							:       Optional[Any]												= tempfile.mkdtemp()
							_snake_case							:       int												= 8

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

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







		def 							lowercase							(     self					)   ->   DPRQuestionEncoderTokenizer:


							'''simple docstring'''

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







		def 							lowercase							(     self					)   ->   BartTokenizer:


							'''simple docstring'''

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







		def 							lowercase							(     self					)   ->   Union[str, Any]:


							'''simple docstring'''

							shutil.rmtree(self.tmpdirname					)







		def 							lowercase							(     self					)   ->   Dict:


							'''simple docstring'''

							_snake_case							:       List[Any]												= 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			,    __snake_case			,    __snake_case=1_2_3_4_5					)   ->   RagPyTorchDistributedRetriever:


							'''simple docstring'''

							_snake_case							:       List[Any]												= self.get_dummy_dataset()
							_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:
												_snake_case							:       Optional[int]												= dataset
												_snake_case							:       Union[str, Any]												= RagPyTorchDistributedRetriever(
												    __snake_case			,    question_encoder_tokenizer=self.get_dpr_tokenizer()			,    generator_tokenizer=self.get_bart_tokenizer()			,    )
												if init_retrieval:
																	retriever.init_retrieval(__snake_case					)
							return retriever







		def 							lowercase							(     self			,    __snake_case					)   ->   RagRayDistributedRetriever:


							'''simple docstring'''

							ray.init(local_mode=__snake_case					)
							_snake_case							:       Union[str, Any]												= RagConfig(
							    retrieval_vector_size=self.retrieval_vector_size			,    question_encoder=DPRConfig().to_dict()			,    generator=BartConfig().to_dict()			,    )
							_snake_case							:       Optional[Any]												= ray.remote(__snake_case					)
							_snake_case							:       List[Any]												= [remote_cls.remote() for _ in range(1					)]
							with patch("transformers.models.rag.retrieval_rag.load_dataset"					) as mock_load_dataset:
												_snake_case							:       Union[str, Any]												= self.get_dummy_dataset()
												_snake_case							:       Optional[Any]												= RagRayDistributedRetriever(
												    __snake_case			,    question_encoder_tokenizer=self.get_dpr_tokenizer()			,    generator_tokenizer=self.get_bart_tokenizer()			,    retrieval_workers=__snake_case			,    )
												if init_retrieval:
																	retriever.init_retrieval()
							return retriever







		def 							lowercase							(     self			,    __snake_case			,    __snake_case			,    __snake_case=1_2_3_4_5					)   ->   Dict:


							'''simple docstring'''

							_snake_case							:       Dict												= self.get_dummy_dataset()
							_snake_case							:       List[str]												= RagConfig(
							    retrieval_vector_size=self.retrieval_vector_size			,    question_encoder=DPRConfig().to_dict()			,    generator=BartConfig().to_dict()			,    index_name="custom"			,    )
							if from_disk:
												_snake_case							:       Tuple												= os.path.join(self.tmpdirname			,    "dataset"					)
												_snake_case							:       Optional[Any]												= 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
												_snake_case							:       Tuple												= RagPyTorchDistributedRetriever(
												    __snake_case			,    question_encoder_tokenizer=self.get_dpr_tokenizer()			,    generator_tokenizer=self.get_bart_tokenizer()			,    )
							else:
												_snake_case							:       Any												= RagPyTorchDistributedRetriever(
												    __snake_case			,    question_encoder_tokenizer=self.get_dpr_tokenizer()			,    generator_tokenizer=self.get_bart_tokenizer()			,    index=CustomHFIndex(config.retrieval_vector_size			,    __snake_case					)			,    )
							if init_retrieval:
												retriever.init_retrieval(__snake_case					)
							return retriever







		def 							lowercase							(     self			,    __snake_case			,    __snake_case					)   ->   List[Any]:


							'''simple docstring'''

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







		def 							lowercase							(     self			,    __snake_case			,    __snake_case			,    __snake_case					)   ->   None:


							'''simple docstring'''

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


							'''simple docstring'''

							_snake_case							:       List[str]												= 1
							_snake_case							:       List[str]												= 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=__snake_case					)			,    __snake_case			,    __snake_case					)







		def 							lowercase							(     self					)   ->   Any:


							'''simple docstring'''

							_snake_case							:       Any												= 1
							_snake_case							:       List[Any]												= 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=__snake_case			,    from_disk=__snake_case					)			,    __snake_case			,    __snake_case			,    )







		def 							lowercase							(     self					)   ->   str:


							'''simple docstring'''

							_snake_case							:       str												= 1
							_snake_case							:       Union[str, Any]												= 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=__snake_case			,    from_disk=__snake_case					)			,    __snake_case			,    __snake_case			,    )







		@require_ray
		def 							lowercase							(     self					)   ->   Optional[Any]:


							'''simple docstring'''

							_snake_case							:       Optional[int]												= 1
							_snake_case							:       int												= 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=__snake_case					)			,    __snake_case			,    __snake_case					)
							ray.shutdown()







		@require_ray
		def 							lowercase							(     self					)   ->   Dict:


							'''simple docstring'''

							_snake_case							:       List[Any]												= 1
							_snake_case							:       Dict												= np.array(
							    [np.ones(self.retrieval_vector_size					), -np.ones(self.retrieval_vector_size					)]			,    dtype=np.floataa					)
							with self.assertRaises(__snake_case					):
												self.distributed_retriever_check(
												    self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=__snake_case			,    from_disk=__snake_case					)			,    __snake_case			,    __snake_case			,    )
							ray.shutdown()







		@require_ray
		def 							lowercase							(     self					)   ->   Tuple:


							'''simple docstring'''

							_snake_case							:       str												= 1
							_snake_case							:       Tuple												= 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=__snake_case			,    from_disk=__snake_case					)			,    __snake_case			,    __snake_case					)
							ray.shutdown()






import json
import os
import subprocess
import sys
import warnings
from argparse import ArgumentParser
from contextlib import AbstractContextManager
from typing import Dict, List, Optional

import requests

from ..utils import logging
from . import BaseTransformersCLICommand


__lowerCAmelCase			:Tuple	 =     logging.get_logger(__name__)  # pylint: disable=invalid-name


__lowerCAmelCase			:Union[str, Any]	 =     'lfs-multipart-upload'






class 							_a(     __A     ):







		@staticmethod
		def 							lowercase							(     __snake_case					)   ->   Tuple:


							'''simple docstring'''

							_snake_case							:       Tuple												= parser.add_parser(
							    "lfs-enable-largefiles"			,    help=(
							        "Deprecated: use `huggingface-cli` instead. Configure your repository to enable upload of files > 5GB."
							    )			,    )
							enable_parser.add_argument("path"			,    type=__snake_case			,    help="Local path to repository you want to configure."					)
							enable_parser.set_defaults(func=lambda __snake_case					: LfsEnableCommand(__snake_case					)					)

							_snake_case							:       Union[str, Any]												= parser.add_parser(
							    __snake_case			,    help=(
							        "Deprecated: use `huggingface-cli` instead. "
							        "Command will get called by git-lfs, do not call it directly."
							    )			,    )
							upload_parser.set_defaults(func=lambda __snake_case					: LfsUploadCommand(__snake_case					)					)






class 							_a:







		def __init__(     self			,    __snake_case					)   ->   Any:


							'''simple docstring'''

							_snake_case							:       List[Any]												= args







		def 							lowercase							(     self					)   ->   Dict:


							'''simple docstring'''

							warnings.warn(
							    "Managing repositories through transformers-cli is deprecated. Please use `huggingface-cli` instead."					)
							_snake_case							:       Dict												= os.path.abspath(self.args.path					)
							if not os.path.isdir(__snake_case					):
												print("This does not look like a valid git repo."					)
												exit(1					)
							subprocess.run(
							    "git config lfs.customtransfer.multipart.path transformers-cli".split()			,    check=__snake_case			,    cwd=__snake_case					)
							subprocess.run(
							    f"""git config lfs.customtransfer.multipart.args {LFS_MULTIPART_UPLOAD_COMMAND}""".split()			,    check=__snake_case			,    cwd=__snake_case			,    )
							print("Local repo set up for largefiles"					)







def  A					(				UpperCAmelCase					):
					_snake_case							:       str												= json.dumps(UpperCAmelCase					) + "\n"
					sys.stdout.write(UpperCAmelCase					)
					sys.stdout.flush()







def  A					(				):
					_snake_case							:       str												= json.loads(sys.stdin.readline().strip()					)

					if "terminate" in (msg.get("type"					), msg.get("event"					)):
										# terminate message received
										return None

					if msg.get("event"					) not in ("download", "upload"):
										logger.critical("Received unexpected message"					)
										sys.exit(1					)

					return msg






class 							_a(     __A     ):







		def __init__(     self			,    __snake_case			,    __snake_case			,    __snake_case					)   ->   Optional[int]:


							'''simple docstring'''

							_snake_case							:       int												= filepath
							_snake_case							:       Tuple												= seek_from
							_snake_case							:       int												= read_limit
							_snake_case							:       int												= 0







		def __enter__(     self					)   ->   Any:


							'''simple docstring'''

							_snake_case							:       Optional[Any]												= open(self.filepath			,    "rb"					)
							self.f.seek(self.seek_from					)
							return self







		def __len__(     self					)   ->   Dict:


							'''simple docstring'''

							_snake_case							:       Dict												= os.fstat(self.f.fileno()					).st_size
							return min(self.read_limit			,    total_length - self.seek_from					)







		def 							lowercase							(     self			,    __snake_case=-1					)   ->   Union[str, Any]:


							'''simple docstring'''

							if self.n_seen >= self.read_limit:
												return b""
							_snake_case							:       Dict												= self.read_limit - self.n_seen
							_snake_case							:       Tuple												= self.f.read(remaining_amount if n < 0 else min(__snake_case			,    __snake_case					)					)
							self.n_seen += len(__snake_case					)
							return data







		def __iter__(     self					)   ->   int:


							'''simple docstring'''

							yield self.read(n=4 * 1_0_2_4 * 1_0_2_4					)







		def __exit__(     self			,    *__snake_case					)   ->   List[str]:


							'''simple docstring'''

							self.f.close()






class 							_a:







		def __init__(     self			,    __snake_case					)   ->   List[str]:


							'''simple docstring'''

							_snake_case							:       Any												= args







		def 							lowercase							(     self					)   ->   str:


							'''simple docstring'''

							_snake_case							:       Optional[Any]												= json.loads(sys.stdin.readline().strip()					)
							if not (init_msg.get("event"					) == "init" and init_msg.get("operation"					) == "upload"):
												write_msg({"error": {"code": 3_2, "message": "Wrong lfs init operation"}}					)
												sys.exit(1					)

							# The transfer process should use the information it needs from the
							# initiation structure, and also perform any one-off setup tasks it
							# needs to do. It should then respond on stdout with a simple empty
							# confirmation structure, as follows:
							write_msg({}					)

							# After the initiation exchange, git-lfs will send any number of
							# transfer requests to the stdin of the transfer process, in a serial sequence.
							while True:
												_snake_case							:       Optional[Any]												= read_msg()
												if msg is None:
																	# When all transfers have been processed, git-lfs will send
																	# a terminate event to the stdin of the transfer process.
																	# On receiving this message the transfer process should
																	# clean up and terminate. No response is expected.
																	sys.exit(0					)

												_snake_case							:       Dict												= msg["oid"]
												_snake_case							:       int												= msg["path"]
												_snake_case							:       Any												= msg["action"]["href"]
												_snake_case							:       Union[str, Any]												= msg["action"]["header"]
												_snake_case							:       Any												= int(header.pop("chunk_size"					)					)
												_snake_case							:       List[str]												= list(header.values()					)

												_snake_case							:       Any												= []
												for i, presigned_url in enumerate(__snake_case					):
																	with FileSlice(__snake_case			,    seek_from=i * chunk_size			,    read_limit=__snake_case					) as data:
																						_snake_case							:       List[str]												= requests.put(__snake_case			,    data=__snake_case					)
																						r.raise_for_status()
																						parts.append(
																						    {
																						        "etag": r.headers.get("etag"					),
																						        "partNumber": i + 1,
																						    }					)
																						# In order to support progress reporting while data is uploading / downloading,
																						# the transfer process should post messages to stdout
																						write_msg(
																						    {
																						        "event": "progress",
																						        "oid": oid,
																						        "bytesSoFar": (i + 1) * chunk_size,
																						        "bytesSinceLast": chunk_size,
																						    }					)
																						# Not precise but that's ok.

												_snake_case							:       Optional[Any]												= requests.post(
												    __snake_case			,    json={
												        "oid": oid,
												        "parts": parts,
												    }			,    )
												r.raise_for_status()

												write_msg({"event": "complete", "oid": oid}					)