




"""simple docstring"""







import logging
import os
from typing import List, Tuple

import numpy as np
import psutil
import torch
import torch.distributed as dist

from transformers import RagRetriever


__A   :  Dict					=      logging.getLogger(__name__)




class       lowerCAmelCase__       (			lowerCAmelCase_    ):


				"""simple docstring"""





				def __init__(     self					:	List[str]    ,   lowercase__					:	Tuple    ,   lowercase__					:	Optional[int]    ,   lowercase__					:	List[Any]    ,   lowercase__					:	Tuple=None						):
						super().__init__(
						    lowercase__    ,   question_encoder_tokenizer=lowercase__    ,   generator_tokenizer=lowercase__    ,   index=lowercase__    ,   init_retrieval=lowercase__    ,   )
						__lowercase				:		Optional[int]												=				None



				def       snake_case    (     self					:	List[Any]    ,   lowercase__					:	int						):
						logger.info("initializing retrieval"						)

						# initializing a separate process group for retrieval as the default
						# nccl backend doesn't support gather/scatter operations while gloo
						# is too slow to replace nccl for the core gpu communication
						if dist.is_initialized():
								logger.info("dist initialized"						)
								# needs to be set manually
								__lowercase				:		Optional[Any]												=				self._infer_socket_ifname()
								# avoid clash with the NCCL port
								__lowercase				:		Optional[Any]												=				str(distributed_port + 1						)
								__lowercase				:		Optional[int]												=				dist.new_group(ranks=lowercase__    ,   backend="gloo"						)

						# initialize retriever only on the main worker
						if not dist.is_initialized() or self._is_main():
								logger.info("dist not initialized / main"						)
								self.index.init_index()

						# all processes wait untill the retriever is initialized by the main process
						if dist.is_initialized():
								torch.distributed.barrier(group=self.process_group						)



				def       snake_case    (     self					:	Union[str, Any]						):
						return dist.get_rank(group=self.process_group						) == 0



				def       snake_case    (     self					:	Union[str, Any]    ,   lowercase__					:	Tuple    ,   lowercase__					:	Dict    ,   lowercase__					:	Any=torch.floataa						):
						__lowercase				:		str												=				torch.empty(lowercase__    ,   dtype=lowercase__						)
						dist.scatter(lowercase__    ,   src=0    ,   scatter_list=lowercase__    ,   group=self.process_group						)
						return target_tensor



				def       snake_case    (     self					:	int						):
						__lowercase				:		Union[str, Any]												=				psutil.net_if_addrs()
						# a hacky way to deal with varying network interface names
						__lowercase				:		Optional[Any]												=				next((addr for addr in addrs if addr.startswith("e"						))    ,   lowercase__						)
						return ifname



				def       snake_case    (     self					:	Union[str, Any]    ,   lowercase__					:	np.ndarray    ,   lowercase__					:	int						):

						# single GPU training
						if not dist.is_initialized():
								__lowercase						,__lowercase				:		Tuple												=				self._main_retrieve(lowercase__    ,   lowercase__						)
								return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(lowercase__						)

						# distributed training
						__lowercase				:		List[str]												=				dist.get_world_size(group=self.process_group						)

						# gather logic
						__lowercase				:		Tuple												=				None
						if self._is_main():
								__lowercase				:		Optional[Any]												=				[torch.empty(question_hidden_states.shape    ,   dtype=torch.floataa						) for _ in range(lowercase__						)]
						dist.gather(torch.tensor(lowercase__						)    ,   dst=0    ,   gather_list=lowercase__    ,   group=self.process_group						)

						# scatter logic
						__lowercase				:		Tuple												=				question_hidden_states.shape[0]
						__lowercase				:		Tuple												=				[]
						__lowercase				:		Union[str, Any]												=				[]
						if self._is_main():
								assert len(lowercase__						) == world_size
								__lowercase						,__lowercase				:		List[str]												=				self._main_retrieve(torch.cat(lowercase__						).numpy()    ,   lowercase__						)
								__lowercase						,__lowercase				:		int												=				torch.tensor(lowercase__						), torch.tensor(lowercase__						)
								__lowercase				:		Optional[int]												=				self._chunk_tensor(lowercase__    ,   lowercase__						)
								__lowercase				:		Union[str, Any]												=				self._chunk_tensor(lowercase__    ,   lowercase__						)
						__lowercase				:		Dict												=				self._scattered(lowercase__    ,   [n_queries, n_docs]    ,   target_type=torch.intaa						)
						__lowercase				:		int												=				self._scattered(lowercase__    ,   [n_queries, n_docs, question_hidden_states.shape[1]]						)

						return retrieved_doc_embeds.numpy(), doc_ids.numpy(), self.index.get_doc_dicts(lowercase__						)






"""simple docstring"""







from scipy.stats import spearmanr

import datasets


__A   :  Optional[int]					=      '\nThe Spearman rank-order correlation coefficient is a measure of the\nrelationship between two datasets. Like other correlation coefficients,\nthis one varies between -1 and +1 with 0 implying no correlation.\nPositive correlations imply that as data in dataset x increases, so\ndoes data in dataset y. Negative correlations imply that as x increases,\ny decreases. Correlations of -1 or +1 imply an exact monotonic relationship.\n\nUnlike the Pearson correlation, the Spearman correlation does not\nassume that both datasets are normally distributed.\n\nThe p-value roughly indicates the probability of an uncorrelated system\nproducing datasets that have a Spearman correlation at least as extreme\nas the one computed from these datasets. The p-values are not entirely\nreliable but are probably reasonable for datasets larger than 500 or so.\n'

__A   :  List[Any]					=      '\nArgs:\n    predictions (`List[float]`): Predicted labels, as returned by a model.\n    references (`List[float]`): Ground truth labels.\n    return_pvalue (`bool`): If `True`, returns the p-value. If `False`, returns\n            only the spearmanr score. Defaults to `False`.\nReturns:\n    spearmanr (`float`): Spearman correlation coefficient.\n    p-value (`float`): p-value. **Note**: is only returned if `return_pvalue=True` is input.\nExamples:\n    Example 1:\n        >>> spearmanr_metric = datasets.load_metric("spearmanr")\n        >>> results = spearmanr_metric.compute(references=[1, 2, 3, 4, 5], predictions=[10, 9, 2.5, 6, 4])\n        >>> print(results)\n        {\'spearmanr\': -0.7}\n\n    Example 2:\n        >>> spearmanr_metric = datasets.load_metric("spearmanr")\n        >>> results = spearmanr_metric.compute(references=[1, 2, 3, 4, 5],\n        ...                                     predictions=[10, 9, 2.5, 6, 4],\n        ...                                     return_pvalue=True)\n        >>> print(results[\'spearmanr\'])\n        -0.7\n        >>> print(round(results[\'spearmanr_pvalue\'], 2))\n        0.19\n'

__A   :  Optional[int]					=      R'\\n@book{kokoska2000crc,\n  title={CRC standard probability and statistics tables and formulae},\n  author={Kokoska, Stephen and Zwillinger, Daniel},\n  year={2000},\n  publisher={Crc Press}\n}\n@article{2020SciPy-NMeth,\n  author  = {Virtanen, Pauli and Gommers, Ralf and Oliphant, Travis E. and\n            Haberland, Matt and Reddy, Tyler and Cournapeau, David and\n            Burovski, Evgeni and Peterson, Pearu and Weckesser, Warren and\n            Bright, Jonathan and {van der Walt}, St{\'e}fan J. and\n            Brett, Matthew and Wilson, Joshua and Millman, K. Jarrod and\n            Mayorov, Nikolay and Nelson, Andrew R. J. and Jones, Eric and\n            Kern, Robert and Larson, Eric and Carey, C J and\n            Polat, {\.I}lhan and Feng, Yu and Moore, Eric W. and\n            {VanderPlas}, Jake and Laxalde, Denis and Perktold, Josef and\n            Cimrman, Robert and Henriksen, Ian and Quintero, E. A. and\n            Harris, Charles R. and Archibald, Anne M. and\n            Ribeiro, Ant{\^o}nio H. and Pedregosa, Fabian and\n            {van Mulbregt}, Paul and {SciPy 1.0 Contributors}},\n  title   = {{{SciPy} 1.0: Fundamental Algorithms for Scientific\n            Computing in Python}},\n  journal = {Nature Methods},\n  year    = {2020},\n  volume  = {17},\n  pages   = {261--272},\n  adsurl  = {https://rdcu.be/b08Wh},\n  doi     = {10.1038/s41592-019-0686-2},\n}\n'




@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION  ,			_KWARGS_DESCRIPTION    )
class       lowerCAmelCase__       (			datasets.Metric    ):


				"""simple docstring"""





				def       snake_case    (     self					:	List[str]						):
						return datasets.MetricInfo(
						    description=_DESCRIPTION    ,   citation=_CITATION    ,   inputs_description=_KWARGS_DESCRIPTION    ,   features=datasets.Features(
						        {
						            "predictions": datasets.Value("float"						),
						            "references": datasets.Value("float"						),
						        }						)    ,   reference_urls=["https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.spearmanr.html"]    ,   )



				def       snake_case    (     self					:	Optional[int]    ,   lowercase__					:	str    ,   lowercase__					:	Any    ,   lowercase__					:	int=False						):
						__lowercase				:		Optional[int]												=				spearmanr(lowercase__    ,   lowercase__						)
						if return_pvalue:
								return {"spearmanr": results[0], "spearmanr_pvalue": results[1]}
						else:
								return {"spearmanr": results[0]}
