




import datasets

from .evaluate import evaluate


UpperCamelCase__         =      """\
@inproceedings{Rajpurkar2016SQuAD10,
  title={SQuAD: 100, 000+ Questions for Machine Comprehension of Text},
  author={Pranav Rajpurkar and Jian Zhang and Konstantin Lopyrev and Percy Liang},
  booktitle={EMNLP},
  year={2016}
}
"""

UpperCamelCase__         =      """
This metric wrap the official scoring script for version 1 of the Stanford Question Answering Dataset (SQuAD).

Stanford Question Answering Dataset (SQuAD) is a reading comprehension dataset, consisting of questions posed by
crowdworkers on a set of Wikipedia articles, where the answer to every question is a segment of text, or span,
from the corresponding reading passage, or the question might be unanswerable.
"""

UpperCamelCase__         =      """
Computes SQuAD scores (F1 and EM).
Args:
    predictions: List of question-answers dictionaries with the following key-values:
        - 'id': id of the question-answer pair as given in the references (see below)
        - 'prediction_text': the text of the answer
    references: List of question-answers dictionaries with the following key-values:
        - 'id': id of the question-answer pair (see above),
        - 'answers': a Dict in the SQuAD dataset format
            {
                'text': list of possible texts for the answer, as a list of strings
                'answer_start': list of start positions for the answer, as a list of ints
            }
            Note that answer_start values are not taken into account to compute the metric.
Returns:
    'exact_match': Exact match (the normalized answer exactly match the gold answer)
    'f1': The F-score of predicted tokens versus the gold answer
Examples:

    >>> predictions = [{'prediction_text': '1976', 'id': '56e10a3be3433e1400422b22'}]
    >>> references = [{'answers': {'answer_start': [97], 'text': ['1976']}, 'id': '56e10a3be3433e1400422b22'}]
    >>> squad_metric = datasets.load_metric(\"squad\")
    >>> results = squad_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {'exact_match': 100.0, 'f1': 100.0}
"""
@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION			,    _KWARGS_DESCRIPTION				)
class       a__   ( datasets.Metric				):






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											return datasets.MetricInfo(
											    description=_DESCRIPTION		,      citation=_CITATION		,      inputs_description=_KWARGS_DESCRIPTION		,      features=datasets.Features(
											        {
											            "predictions": {"id": datasets.Value("string"      ), "prediction_text": datasets.Value("string"      )},
											            "references": {
											                "id": datasets.Value("string"      ),
											                "answers": datasets.features.Sequence(
											                    {
											                        "text": datasets.Value("string"      ),
											                        "answer_start": datasets.Value("int32"      ),
											                    }      ),
											            },
											        }      )		,      codebase_urls=["https://rajpurkar.github.io/SQuAD-explorer/"]		,      reference_urls=["https://rajpurkar.github.io/SQuAD-explorer/"]		,      )






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	{prediction["id"]: prediction["prediction_text"] for prediction in predictions}
											__lowerCAmelCase           =	[
											    {
											        "paragraphs": [
											            {
											                "qas": [
											                    {
											                        "answers": [{"text": answer_text} for answer_text in ref["answers"]["text"]],
											                        "id": ref["id"],
											                    }
											                    for ref in references
											                ]
											            }
											        ]
											    }
											]
											__lowerCAmelCase           =	evaluate(dataset=_A		,      predictions=_A      )
											return score










import logging
import random

import ray

from transformers import RagConfig, RagRetriever, RagTokenizer
from transformers.models.rag.retrieval_rag import CustomHFIndex


UpperCamelCase__         =      logging.getLogger(__name__)
class       a__   :






						def __init__(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	False






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A		,      _A		,      _A      ):




											"""simple docstring"""



											if not self.initialized:
																__lowerCAmelCase           =	RagRetriever(
																    _A		,      question_encoder_tokenizer=_A		,      generator_tokenizer=_A		,      index=_A		,      init_retrieval=_A		,      )
																__lowerCAmelCase           =	True






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											self.retriever.index.init_index()






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase       ,							__lowerCAmelCase           =	self.retriever._main_retrieve(_A		,      _A      )
											return doc_ids, retrieved_doc_embeds






class       a__   ( snake_case__				):






						def __init__(       self		,      _A		,      _A		,      _A		,      _A		,      _A=None      ):




											"""simple docstring"""



											if index is not None and index.is_initialized() and len(_A      ) > 0:
																raise ValueError(
																    "When using Ray for distributed fine-tuning, "
																    "you'll need to provide the paths instead, "
																    "as the dataset and the index are loaded "
																    "separately. More info in examples/rag/use_own_knowledge_dataset.py "      )
											super().__init__(
											    _A		,      question_encoder_tokenizer=_A		,      generator_tokenizer=_A		,      index=_A		,      init_retrieval=_A		,      )
											__lowerCAmelCase           =	retrieval_workers
											if len(self.retrieval_workers      ) > 0:
																ray.get(
																    [
																        worker.create_rag_retriever.remote(_A		,      _A		,      _A		,      _A      )
																        for worker in self.retrieval_workers
																    ]      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											logger.info("initializing retrieval"      )

											if len(self.retrieval_workers      ) > 0:
																ray.get([worker.init_retrieval.remote() for worker in self.retrieval_workers]      )
											else:
																# Non-distributed training. Load index into this same process.
																self.index.init_index()






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A      ):




											"""simple docstring"""



											if len(self.retrieval_workers      ) > 0:
																# Select a random retrieval actor.
																__lowerCAmelCase           =	self.retrieval_workers[random.randint(0		,      len(self.retrieval_workers      ) - 1      )]
																__lowerCAmelCase       ,							__lowerCAmelCase           =	ray.get(random_worker.retrieve.remote(_A		,      _A      )      )
											else:
																__lowerCAmelCase       ,							__lowerCAmelCase           =	self._main_retrieve(_A		,      _A      )
											return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(_A      )






						@classmethod
						def     __SCREAMING_SNAKE_CASE(       cls		,      _A		,      _A=None		,      **_A      ):




											"""simple docstring"""



											return super(_A		,      cls      ).get_tokenizers(_A		,      _A		,      **_A      )






						@classmethod
						def     __SCREAMING_SNAKE_CASE(       cls		,      _A		,      _A		,      _A=None		,      **_A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	kwargs.pop("config"		,      _A      ) or RagConfig.from_pretrained(_A		,      **_A      )
											__lowerCAmelCase           =	RagTokenizer.from_pretrained(_A		,      config=_A      )
											__lowerCAmelCase           =	rag_tokenizer.question_encoder
											__lowerCAmelCase           =	rag_tokenizer.generator
											if indexed_dataset is not None:
																__lowerCAmelCase           =	"custom"
																__lowerCAmelCase           =	CustomHFIndex(config.retrieval_vector_size		,      _A      )
											else:
																__lowerCAmelCase           =	cls._build_index(_A      )
											return cls(
											    _A		,      question_encoder_tokenizer=_A		,      generator_tokenizer=_A		,      retrieval_workers=_A		,      index=_A		,      )




