


'''simple docstring'''
import tempfile
import unittest

import numpy as np

from diffusers import (
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionPipeline,
    PNDMScheduler,
)
from diffusers.utils.testing_utils import is_onnx_available, nightly, require_onnxruntime, require_torch_gpu

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
		import onnxruntime as ort
class    lowercase	(							_a	,		unittest.TestCase					):

		"""simple docstring"""

		_a     =				"""hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"""




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_=0							):

							'''simple docstring'''


							UpperCamelCase__   :str												= np.random.RandomState(UpperCamelCase_							)
							UpperCamelCase__   :List[Any]												= {
							    'prompt': 'A painting of a squirrel eating a burger',
							    'generator': generator,
							    'num_inference_steps': 2,
							    'guidance_scale': 7.5,
							    'output_type': 'numpy',
							}
							return inputs




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							UpperCamelCase__   :int												= OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint    ,   provider='''CPUExecutionProvider'''							)
							pipe.set_progress_bar_config(disable=UpperCamelCase_							)

							UpperCamelCase__   :Tuple												= self.get_dummy_inputs()
							UpperCamelCase__   :Optional[Any]												= pipe(**UpperCamelCase_							).images
							UpperCamelCase__   :Any												= image[0, -3:, -3:, -1]

							assert image.shape == (1, 128, 128, 3)
							UpperCamelCase__   :Tuple												= np.array([0.65072, 0.58492, 0.48219, 0.55521, 0.53180, 0.55939, 0.50697, 0.39800, 0.46455]							)

							assert np.abs(image_slice.flatten() - expected_slice							).max() < 1e-2




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							UpperCamelCase__   :List[Any]												= OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint    ,   provider='''CPUExecutionProvider'''							)
							UpperCamelCase__   :List[Any]												= PNDMScheduler.from_config(pipe.scheduler.config    ,   skip_prk_steps=UpperCamelCase_							)
							pipe.set_progress_bar_config(disable=UpperCamelCase_							)

							UpperCamelCase__   :Dict												= self.get_dummy_inputs()
							UpperCamelCase__   :str												= pipe(**UpperCamelCase_							).images
							UpperCamelCase__   :Union[str, Any]												= image[0, -3:, -3:, -1]

							assert image.shape == (1, 128, 128, 3)
							UpperCamelCase__   :List[str]												= np.array([0.65863, 0.59425, 0.49326, 0.56313, 0.53875, 0.56627, 0.51065, 0.39777, 0.46330]							)

							assert np.abs(image_slice.flatten() - expected_slice							).max() < 1e-2




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							UpperCamelCase__   :List[Any]												= OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint    ,   provider='''CPUExecutionProvider'''							)
							UpperCamelCase__   :List[str]												= LMSDiscreteScheduler.from_config(pipe.scheduler.config							)
							pipe.set_progress_bar_config(disable=UpperCamelCase_							)

							UpperCamelCase__   :List[Any]												= self.get_dummy_inputs()
							UpperCamelCase__   :Union[str, Any]												= pipe(**UpperCamelCase_							).images
							UpperCamelCase__   :Tuple												= image[0, -3:, -3:, -1]

							assert image.shape == (1, 128, 128, 3)
							UpperCamelCase__   :str												= np.array([0.53755, 0.60786, 0.47402, 0.49488, 0.51869, 0.49819, 0.47985, 0.38957, 0.44279]							)

							assert np.abs(image_slice.flatten() - expected_slice							).max() < 1e-2




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							UpperCamelCase__   :Union[str, Any]												= OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint    ,   provider='''CPUExecutionProvider'''							)
							UpperCamelCase__   :Optional[int]												= EulerDiscreteScheduler.from_config(pipe.scheduler.config							)
							pipe.set_progress_bar_config(disable=UpperCamelCase_							)

							UpperCamelCase__   :List[Any]												= self.get_dummy_inputs()
							UpperCamelCase__   :int												= pipe(**UpperCamelCase_							).images
							UpperCamelCase__   :Any												= image[0, -3:, -3:, -1]

							assert image.shape == (1, 128, 128, 3)
							UpperCamelCase__   :Any												= np.array([0.53755, 0.60786, 0.47402, 0.49488, 0.51869, 0.49819, 0.47985, 0.38957, 0.44279]							)

							assert np.abs(image_slice.flatten() - expected_slice							).max() < 1e-2




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							UpperCamelCase__   :Dict												= OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint    ,   provider='''CPUExecutionProvider'''							)
							UpperCamelCase__   :Union[str, Any]												= EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config							)
							pipe.set_progress_bar_config(disable=UpperCamelCase_							)

							UpperCamelCase__   :List[Any]												= self.get_dummy_inputs()
							UpperCamelCase__   :List[str]												= pipe(**UpperCamelCase_							).images
							UpperCamelCase__   :Any												= image[0, -3:, -3:, -1]

							assert image.shape == (1, 128, 128, 3)
							UpperCamelCase__   :Any												= np.array([0.53817, 0.60812, 0.47384, 0.49530, 0.51894, 0.49814, 0.47984, 0.38958, 0.44271]							)

							assert np.abs(image_slice.flatten() - expected_slice							).max() < 1e-2




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							UpperCamelCase__   :int												= OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint    ,   provider='''CPUExecutionProvider'''							)
							UpperCamelCase__   :int												= DPMSolverMultistepScheduler.from_config(pipe.scheduler.config							)
							pipe.set_progress_bar_config(disable=UpperCamelCase_							)

							UpperCamelCase__   :List[Any]												= self.get_dummy_inputs()
							UpperCamelCase__   :Optional[Any]												= pipe(**UpperCamelCase_							).images
							UpperCamelCase__   :List[Any]												= image[0, -3:, -3:, -1]

							assert image.shape == (1, 128, 128, 3)
							UpperCamelCase__   :int												= np.array([0.53895, 0.60808, 0.47933, 0.49608, 0.51886, 0.49950, 0.48053, 0.38957, 0.44200]							)

							assert np.abs(image_slice.flatten() - expected_slice							).max() < 1e-2




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							UpperCamelCase__   :List[Any]												= OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint    ,   provider='''CPUExecutionProvider'''							)
							pipe.set_progress_bar_config(disable=UpperCamelCase_							)

							UpperCamelCase__   :int												= self.get_dummy_inputs()
							UpperCamelCase__   :str												= 3 * [inputs['prompt']]

							# forward
							UpperCamelCase__   :Optional[Any]												= pipe(**UpperCamelCase_							)
							UpperCamelCase__   :Optional[Any]												= output.images[0, -3:, -3:, -1]

							UpperCamelCase__   :List[Any]												= self.get_dummy_inputs()
							UpperCamelCase__   :Any												= 3 * [inputs.pop('''prompt'''							)]

							UpperCamelCase__   :Optional[int]												= pipe.tokenizer(
							    UpperCamelCase_    ,   padding='''max_length'''    ,   max_length=pipe.tokenizer.model_max_length    ,   truncation=UpperCamelCase_    ,   return_tensors='''np'''    ,   )
							UpperCamelCase__   :Optional[int]												= text_inputs['input_ids']

							UpperCamelCase__   :str												= pipe.text_encoder(input_ids=text_inputs.astype(np.intaa							)							)[0]

							UpperCamelCase__   :Tuple												= prompt_embeds

							# forward
							UpperCamelCase__   :Union[str, Any]												= pipe(**UpperCamelCase_							)
							UpperCamelCase__   :Tuple												= output.images[0, -3:, -3:, -1]

							assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()							).max() < 1e-4




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							UpperCamelCase__   :List[str]												= OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint    ,   provider='''CPUExecutionProvider'''							)
							pipe.set_progress_bar_config(disable=UpperCamelCase_							)

							UpperCamelCase__   :Union[str, Any]												= self.get_dummy_inputs()
							UpperCamelCase__   :Any												= 3 * ['this is a negative prompt']
							UpperCamelCase__   :int												= negative_prompt
							UpperCamelCase__   :List[str]												= 3 * [inputs['prompt']]

							# forward
							UpperCamelCase__   :str												= pipe(**UpperCamelCase_							)
							UpperCamelCase__   :Tuple												= output.images[0, -3:, -3:, -1]

							UpperCamelCase__   :Optional[Any]												= self.get_dummy_inputs()
							UpperCamelCase__   :int												= 3 * [inputs.pop('''prompt'''							)]

							UpperCamelCase__   :Optional[int]												= []
							for p in [prompt, negative_prompt]:
												UpperCamelCase__   :int												= pipe.tokenizer(
												    UpperCamelCase_    ,   padding='''max_length'''    ,   max_length=pipe.tokenizer.model_max_length    ,   truncation=UpperCamelCase_    ,   return_tensors='''np'''    ,   )
												UpperCamelCase__   :List[Any]												= text_inputs['input_ids']

												embeds.append(pipe.text_encoder(input_ids=text_inputs.astype(np.intaa							)							)[0]							)

							UpperCamelCase__   :Optional[Any]												= embeds

							# forward
							UpperCamelCase__   :List[str]												= pipe(**UpperCamelCase_							)
							UpperCamelCase__   :str												= output.images[0, -3:, -3:, -1]

							assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()							).max() < 1e-4



@nightly
@require_onnxruntime
@require_torch_gpu
class    lowercase	(							unittest.TestCase					):

		"""simple docstring"""

		@property
		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							return (
							    "CUDAExecutionProvider",
							    {
							        "gpu_mem_limit": "15000000000",  # 15GB
							        "arena_extend_strategy": "kSameAsRequested",
							    },
							)




		@property
		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							UpperCamelCase__   :Union[str, Any]												= ort.SessionOptions()
							UpperCamelCase__   :str												= False
							return options




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							UpperCamelCase__   :Dict												= OnnxStableDiffusionPipeline.from_pretrained(
							    '''CompVis/stable-diffusion-v1-4'''    ,   revision='''onnx'''    ,   safety_checker=UpperCamelCase_    ,   feature_extractor=UpperCamelCase_    ,   provider=self.gpu_provider    ,   sess_options=self.gpu_options    ,   )
							sd_pipe.set_progress_bar_config(disable=UpperCamelCase_							)

							UpperCamelCase__   :Tuple												= 'A painting of a squirrel eating a burger'
							np.random.seed(0							)
							UpperCamelCase__   :int												= sd_pipe([prompt]    ,   guidance_scale=6.0    ,   num_inference_steps=10    ,   output_type='''np'''							)
							UpperCamelCase__   :Union[str, Any]												= output.images

							UpperCamelCase__   :Optional[int]												= image[0, -3:, -3:, -1]

							assert image.shape == (1, 512, 512, 3)
							UpperCamelCase__   :Optional[int]												= np.array([0.0452, 0.0390, 0.0087, 0.0350, 0.0617, 0.0364, 0.0544, 0.0523, 0.0720]							)

							assert np.abs(image_slice.flatten() - expected_slice							).max() < 1e-3




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							UpperCamelCase__   :Union[str, Any]												= DDIMScheduler.from_pretrained(
							    '''runwayml/stable-diffusion-v1-5'''    ,   subfolder='''scheduler'''    ,   revision='''onnx'''							)
							UpperCamelCase__   :int												= OnnxStableDiffusionPipeline.from_pretrained(
							    '''runwayml/stable-diffusion-v1-5'''    ,   revision='''onnx'''    ,   scheduler=UpperCamelCase_    ,   safety_checker=UpperCamelCase_    ,   feature_extractor=UpperCamelCase_    ,   provider=self.gpu_provider    ,   sess_options=self.gpu_options    ,   )
							sd_pipe.set_progress_bar_config(disable=UpperCamelCase_							)

							UpperCamelCase__   :Dict												= 'open neural network exchange'
							UpperCamelCase__   :Union[str, Any]												= np.random.RandomState(0							)
							UpperCamelCase__   :int												= sd_pipe([prompt]    ,   guidance_scale=7.5    ,   num_inference_steps=10    ,   generator=UpperCamelCase_    ,   output_type='''np'''							)
							UpperCamelCase__   :Tuple												= output.images
							UpperCamelCase__   :Union[str, Any]												= image[0, -3:, -3:, -1]

							assert image.shape == (1, 512, 512, 3)
							UpperCamelCase__   :List[Any]												= np.array([0.2867, 0.1974, 0.1481, 0.7294, 0.7251, 0.6667, 0.4194, 0.5642, 0.6486]							)

							assert np.abs(image_slice.flatten() - expected_slice							).max() < 1e-3




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							UpperCamelCase__   :int												= LMSDiscreteScheduler.from_pretrained(
							    '''runwayml/stable-diffusion-v1-5'''    ,   subfolder='''scheduler'''    ,   revision='''onnx'''							)
							UpperCamelCase__   :Optional[int]												= OnnxStableDiffusionPipeline.from_pretrained(
							    '''runwayml/stable-diffusion-v1-5'''    ,   revision='''onnx'''    ,   scheduler=UpperCamelCase_    ,   safety_checker=UpperCamelCase_    ,   feature_extractor=UpperCamelCase_    ,   provider=self.gpu_provider    ,   sess_options=self.gpu_options    ,   )
							sd_pipe.set_progress_bar_config(disable=UpperCamelCase_							)

							UpperCamelCase__   :int												= 'open neural network exchange'
							UpperCamelCase__   :Optional[Any]												= np.random.RandomState(0							)
							UpperCamelCase__   :List[str]												= sd_pipe([prompt]    ,   guidance_scale=7.5    ,   num_inference_steps=10    ,   generator=UpperCamelCase_    ,   output_type='''np'''							)
							UpperCamelCase__   :Dict												= output.images
							UpperCamelCase__   :List[str]												= image[0, -3:, -3:, -1]

							assert image.shape == (1, 512, 512, 3)
							UpperCamelCase__   :str												= np.array([0.2306, 0.1959, 0.1593, 0.6549, 0.6394, 0.5408, 0.5065, 0.6010, 0.6161]							)

							assert np.abs(image_slice.flatten() - expected_slice							).max() < 1e-3




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							UpperCamelCase__   :Dict												= 0

							def test_callback_fn(UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_							) -> None:
												UpperCamelCase__   :int												= True
												nonlocal number_of_steps
												number_of_steps += 1
												if step == 0:
																	assert latents.shape == (1, 4, 64, 64)
																	UpperCamelCase__   :Tuple												= latents[0, -3:, -3:, -1]
																	UpperCamelCase__   :List[str]												= np.array(
																	    [-0.6772, -0.3835, -1.2456, 0.1905, -1.0974, 0.6967, -1.9353, 0.0178, 1.0167]							)

																	assert np.abs(latents_slice.flatten() - expected_slice							).max() < 1e-3
												elif step == 5:
																	assert latents.shape == (1, 4, 64, 64)
																	UpperCamelCase__   :int												= latents[0, -3:, -3:, -1]
																	UpperCamelCase__   :Optional[int]												= np.array(
																	    [-0.3351, 0.2241, -0.1837, -0.2325, -0.6577, 0.3393, -0.0241, 0.5899, 1.3875]							)

																	assert np.abs(latents_slice.flatten() - expected_slice							).max() < 1e-3

							UpperCamelCase__   :Any												= False

							UpperCamelCase__   :List[str]												= OnnxStableDiffusionPipeline.from_pretrained(
							    '''runwayml/stable-diffusion-v1-5'''    ,   revision='''onnx'''    ,   safety_checker=UpperCamelCase_    ,   feature_extractor=UpperCamelCase_    ,   provider=self.gpu_provider    ,   sess_options=self.gpu_options    ,   )
							pipe.set_progress_bar_config(disable=UpperCamelCase_							)

							UpperCamelCase__   :str												= 'Andromeda galaxy in a bottle'

							UpperCamelCase__   :Union[str, Any]												= np.random.RandomState(0							)
							pipe(
							    prompt=UpperCamelCase_    ,   num_inference_steps=5    ,   guidance_scale=7.5    ,   generator=UpperCamelCase_    ,   callback=UpperCamelCase_    ,   callback_steps=1    ,   )
							assert test_callback_fn.has_been_called
							assert number_of_steps == 6




		def     lowerCAmelCase__				(		self							):

							'''simple docstring'''


							UpperCamelCase__   :str												= OnnxStableDiffusionPipeline.from_pretrained(
							    '''runwayml/stable-diffusion-v1-5'''    ,   revision='''onnx'''    ,   safety_checker=UpperCamelCase_    ,   feature_extractor=UpperCamelCase_    ,   provider=self.gpu_provider    ,   sess_options=self.gpu_options    ,   )
							assert isinstance(UpperCamelCase_    ,   UpperCamelCase_							)
							assert pipe.safety_checker is None

							UpperCamelCase__   :Union[str, Any]												= pipe('''example prompt'''    ,   num_inference_steps=2							).images[0]
							assert image is not None

							# check that there's no error when saving a pipeline with one of the models being None
							with tempfile.TemporaryDirectory() as tmpdirname:
												pipe.save_pretrained(UpperCamelCase_							)
												UpperCamelCase__   :int												= OnnxStableDiffusionPipeline.from_pretrained(UpperCamelCase_							)

							# sanity check that the pipeline still works
							assert pipe.safety_checker is None
							UpperCamelCase__   :Optional[int]												= pipe('''example prompt'''    ,   num_inference_steps=2							).images[0]
							assert image is not None



import collections
from typing import List, Optional, Union

from ...tokenization_utils_base import BatchEncoding
from ...utils import TensorType, add_end_docstrings, add_start_docstrings, logging
from ..bert.tokenization_bert import BertTokenizer


SCREAMING_SNAKE_CASE       :      Any												= logging.get_logger(__name__)

SCREAMING_SNAKE_CASE       :      Union[str, Any]												= {"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"}

SCREAMING_SNAKE_CASE       :      Union[str, Any]												= {
    "vocab_file": {
        "facebook/dpr-ctx_encoder-single-nq-base": (
            "https://huggingface.co/facebook/dpr-ctx_encoder-single-nq-base/resolve/main/vocab.txt"
        ),
        "facebook/dpr-ctx_encoder-multiset-base": (
            "https://huggingface.co/facebook/dpr-ctx_encoder-multiset-base/resolve/main/vocab.txt"
        ),
    },
    "tokenizer_file": {
        "facebook/dpr-ctx_encoder-single-nq-base": (
            "https://huggingface.co/facebook/dpr-ctx_encoder-single-nq-base/resolve/main/tokenizer.json"
        ),
        "facebook/dpr-ctx_encoder-multiset-base": (
            "https://huggingface.co/facebook/dpr-ctx_encoder-multiset-base/resolve/main/tokenizer.json"
        ),
    },
}
SCREAMING_SNAKE_CASE       :      Dict												= {
    "vocab_file": {
        "facebook/dpr-question_encoder-single-nq-base": (
            "https://huggingface.co/facebook/dpr-question_encoder-single-nq-base/resolve/main/vocab.txt"
        ),
        "facebook/dpr-question_encoder-multiset-base": (
            "https://huggingface.co/facebook/dpr-question_encoder-multiset-base/resolve/main/vocab.txt"
        ),
    },
    "tokenizer_file": {
        "facebook/dpr-question_encoder-single-nq-base": (
            "https://huggingface.co/facebook/dpr-question_encoder-single-nq-base/resolve/main/tokenizer.json"
        ),
        "facebook/dpr-question_encoder-multiset-base": (
            "https://huggingface.co/facebook/dpr-question_encoder-multiset-base/resolve/main/tokenizer.json"
        ),
    },
}
SCREAMING_SNAKE_CASE       :      str												= {
    "vocab_file": {
        "facebook/dpr-reader-single-nq-base": (
            "https://huggingface.co/facebook/dpr-reader-single-nq-base/resolve/main/vocab.txt"
        ),
        "facebook/dpr-reader-multiset-base": (
            "https://huggingface.co/facebook/dpr-reader-multiset-base/resolve/main/vocab.txt"
        ),
    },
    "tokenizer_file": {
        "facebook/dpr-reader-single-nq-base": (
            "https://huggingface.co/facebook/dpr-reader-single-nq-base/resolve/main/tokenizer.json"
        ),
        "facebook/dpr-reader-multiset-base": (
            "https://huggingface.co/facebook/dpr-reader-multiset-base/resolve/main/tokenizer.json"
        ),
    },
}

SCREAMING_SNAKE_CASE       :      Dict												= {
    "facebook/dpr-ctx_encoder-single-nq-base": 512,
    "facebook/dpr-ctx_encoder-multiset-base": 512,
}
SCREAMING_SNAKE_CASE       :      Optional[Any]												= {
    "facebook/dpr-question_encoder-single-nq-base": 512,
    "facebook/dpr-question_encoder-multiset-base": 512,
}
SCREAMING_SNAKE_CASE       :      Dict												= {
    "facebook/dpr-reader-single-nq-base": 512,
    "facebook/dpr-reader-multiset-base": 512,
}


SCREAMING_SNAKE_CASE       :      List[Any]												= {
    "facebook/dpr-ctx_encoder-single-nq-base": {"do_lower_case": True},
    "facebook/dpr-ctx_encoder-multiset-base": {"do_lower_case": True},
}
SCREAMING_SNAKE_CASE       :      Dict												= {
    "facebook/dpr-question_encoder-single-nq-base": {"do_lower_case": True},
    "facebook/dpr-question_encoder-multiset-base": {"do_lower_case": True},
}
SCREAMING_SNAKE_CASE       :      Dict												= {
    "facebook/dpr-reader-single-nq-base": {"do_lower_case": True},
    "facebook/dpr-reader-multiset-base": {"do_lower_case": True},
}






class _lowerCamelCase(						_a   ):
       lowercase_ :	Any     =			VOCAB_FILES_NAMES
       lowercase_ :	Optional[int]     =			CONTEXT_ENCODER_PRETRAINED_VOCAB_FILES_MAP
       lowercase_ :	str     =			CONTEXT_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
       lowercase_ :	str     =			CONTEXT_ENCODER_PRETRAINED_INIT_CONFIGURATION






class _lowerCamelCase(						_a   ):
       lowercase_ :	Optional[int]     =			VOCAB_FILES_NAMES
       lowercase_ :	Any     =			QUESTION_ENCODER_PRETRAINED_VOCAB_FILES_MAP
       lowercase_ :	str     =			QUESTION_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
       lowercase_ :	Union[str, Any]     =			QUESTION_ENCODER_PRETRAINED_INIT_CONFIGURATION


SCREAMING_SNAKE_CASE       :      Optional[int]												= collections.namedtuple(
    "DPRSpanPrediction", ["span_score", "relevance_score", "doc_id", "start_index", "end_index", "text"]
)

SCREAMING_SNAKE_CASE       :      Any												= collections.namedtuple("DPRReaderOutput", ["start_logits", "end_logits", "relevance_logits"])


SCREAMING_SNAKE_CASE       :      str												= r"\n    Return a dictionary with the token ids of the input strings and other information to give to `.decode_best_spans`.\n    It converts the strings of a question and different passages (title and text) in a sequence of IDs (integers),\n    using the tokenizer and vocabulary. The resulting `input_ids` is a matrix of size `(n_passages, sequence_length)`\n    with the format:\n\n    ```\n    [CLS] <question token ids> [SEP] <titles ids> [SEP] <texts ids>\n    ```\n\n    Args:\n        questions (`str` or `List[str]`):\n            The questions to be encoded. You can specify one question for many passages. In this case, the question\n            will be duplicated like `[questions] * n_passages`. Otherwise you have to specify as many questions as in\n            `titles` or `texts`.\n        titles (`str` or `List[str]`):\n            The passages titles to be encoded. This can be a string or a list of strings if there are several passages.\n        texts (`str` or `List[str]`):\n            The passages texts to be encoded. This can be a string or a list of strings if there are several passages.\n        padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`):\n            Activates and controls padding. Accepts the following values:\n\n            - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single sequence\n              if provided).\n            - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum\n              acceptable input length for the model if that argument is not provided.\n            - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different\n              lengths).\n        truncation (`bool`, `str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*, defaults to `False`):\n            Activates and controls truncation. Accepts the following values:\n\n            - `True` or `'longest_first'`: Truncate to a maximum length specified with the argument `max_length` or to\n              the maximum acceptable input length for the model if that argument is not provided. This will truncate\n              token by token, removing a token from the longest sequence in the pair if a pair of sequences (or a batch\n              of pairs) is provided.\n            - `'only_first'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum\n              acceptable input length for the model if that argument is not provided. This will only truncate the first\n              sequence of a pair if a pair of sequences (or a batch of pairs) is provided.\n            - `'only_second'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum\n              acceptable input length for the model if that argument is not provided. This will only truncate the\n              second sequence of a pair if a pair of sequences (or a batch of pairs) is provided.\n            - `False` or `'do_not_truncate'` (default): No truncation (i.e., can output batch with sequence lengths\n              greater than the model maximum admissible input size).\n        max_length (`int`, *optional*):\n                Controls the maximum length to use by one of the truncation/padding parameters.\n\n                If left unset or set to `None`, this will use the predefined model maximum length if a maximum length\n                is required by one of the truncation/padding parameters. If the model has no specific maximum input\n                length (like XLNet) truncation/padding to a maximum length will be deactivated.\n        return_tensors (`str` or [`~utils.TensorType`], *optional*):\n                If set, will return tensors instead of list of python integers. Acceptable values are:\n\n                - `'tf'`: Return TensorFlow `tf.constant` objects.\n                - `'pt'`: Return PyTorch `torch.Tensor` objects.\n                - `'np'`: Return Numpy `np.ndarray` objects.\n        return_attention_mask (`bool`, *optional*):\n            Whether or not to return the attention mask. If not set, will return the attention mask according to the\n            specific tokenizer's default, defined by the `return_outputs` attribute.\n\n            [What are attention masks?](../glossary#attention-mask)\n\n    Returns:\n        `Dict[str, List[List[int]]]`: A dictionary with the following keys:\n\n        - `input_ids`: List of token ids to be fed to a model.\n        - `attention_mask`: List of indices specifying which tokens should be attended to by the model.\n    "






@add_start_docstrings(_a   )
class _lowerCamelCase:
       def __call__(					self,      lowerCamelCase,      lowerCamelCase = None,      lowerCamelCase = None,      lowerCamelCase = False,      lowerCamelCase = False,      lowerCamelCase = None,      lowerCamelCase = None,      lowerCamelCase = None,      **lowerCamelCase,      )       ->   BatchEncoding:




           """simple docstring"""

           if titles is None and texts is None:
               return super().__call__(
                   lowerCamelCase,      padding=lowerCamelCase,      truncation=lowerCamelCase,      max_length=lowerCamelCase,      return_tensors=lowerCamelCase,      return_attention_mask=lowerCamelCase,      **lowerCamelCase,      )
           elif titles is None or texts is None:
               _lowercase       :							Dict      		=		titles if texts is None else texts
               return super().__call__(
                   lowerCamelCase,      lowerCamelCase,      padding=lowerCamelCase,      truncation=lowerCamelCase,      max_length=lowerCamelCase,      return_tensors=lowerCamelCase,      return_attention_mask=lowerCamelCase,      **lowerCamelCase,      )
           _lowercase       :							Union[str, Any]      		=		titles if not isinstance(lowerCamelCase,      lowerCamelCase) else [titles]
           _lowercase       :							Tuple      		=		texts if not isinstance(lowerCamelCase,      lowerCamelCase) else [texts]
           _lowercase       :							Optional[Any]      		=		len(lowerCamelCase)
           _lowercase       :							Any      		=		questions if not isinstance(lowerCamelCase,      lowerCamelCase) else [questions] * n_passages
           if len(lowerCamelCase) != len(lowerCamelCase):
               raise ValueError(
                   F'''There should be as many titles than texts but got {len(lowerCamelCase)} titles and {len(lowerCamelCase)} texts.''')
           _lowercase       :							Any      		=		super().__call__(lowerCamelCase,      lowerCamelCase,      padding=lowerCamelCase,      truncation=lowerCamelCase)['input_ids']
           _lowercase       :							Tuple      		=		super().__call__(lowerCamelCase,      add_special_tokens=lowerCamelCase,      padding=lowerCamelCase,      truncation=lowerCamelCase)['input_ids']
           _lowercase       :							int      		=		{
               'input_ids': [
                   (encoded_question_and_title + encoded_text)[:max_length]
                   if max_length is not None and truncation
                   else encoded_question_and_title + encoded_text
                   for encoded_question_and_title, encoded_text in zip(lowerCamelCase,      lowerCamelCase)
               ]
           }
           if return_attention_mask is not False:
               _lowercase       :							Optional[Any]      		=		[]
               for input_ids in encoded_inputs["input_ids"]:
                   attention_mask.append([int(input_id != self.pad_token_id) for input_id in input_ids])
               _lowercase       :							Union[str, Any]      		=		attention_mask
           return self.pad(lowerCamelCase,      padding=lowerCamelCase,      max_length=lowerCamelCase,      return_tensors=lowerCamelCase)
       def 			UpperCamelCase     (					self,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase = 16,      lowerCamelCase = 64,      lowerCamelCase = 4,      )       ->   List[DPRSpanPrediction]:




           """simple docstring"""

           _lowercase       :							Union[str, Any]      		=		reader_input['input_ids']
           _lowercase		,				_lowercase		,				_lowercase       :							Tuple      		=		reader_output[:3]
           _lowercase       :							Tuple      		=		len(lowerCamelCase)
           _lowercase       :							str      		=		sorted(range(lowerCamelCase),      reverse=lowerCamelCase,      key=relevance_logits.__getitem__)
           _lowercase       :							List[DPRReaderOutput]      		=		[]
           for doc_id in sorted_docs:
               _lowercase       :							str      		=		list(input_ids[doc_id])
               # assuming question & title information is at the beginning of the sequence
               _lowercase       :							Any      		=		sequence_ids.index(self.sep_token_id,      2) + 1  # second sep id
               if sequence_ids[-1] == self.pad_token_id:
                   _lowercase       :							List[Any]      		=		sequence_ids.index(self.pad_token_id)
               else:
                   _lowercase       :							List[str]      		=		len(lowerCamelCase)

               _lowercase       :							Tuple      		=		self._get_best_spans(
                   start_logits=start_logits[doc_id][passage_offset:sequence_len],      end_logits=end_logits[doc_id][passage_offset:sequence_len],      max_answer_length=lowerCamelCase,      top_spans=lowerCamelCase,      )
               for start_index, end_index in best_spans:
                   start_index += passage_offset
                   end_index += passage_offset
                   nbest_spans_predictions.append(
                       DPRSpanPrediction(
                           span_score=start_logits[doc_id][start_index] + end_logits[doc_id][end_index],      relevance_score=relevance_logits[doc_id],      doc_id=lowerCamelCase,      start_index=lowerCamelCase,      end_index=lowerCamelCase,      text=self.decode(sequence_ids[start_index : end_index + 1]),      ))
               if len(lowerCamelCase) >= num_spans:
                   break
           return nbest_spans_predictions[:num_spans]




       def 			UpperCamelCase     (					self,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      )       ->   List[DPRSpanPrediction]:




           """simple docstring"""

           _lowercase       :							str      		=		[]
           for start_index, start_score in enumerate(lowerCamelCase):
               for answer_length, end_score in enumerate(end_logits[start_index : start_index + max_answer_length]):
                   scores.append(((start_index, start_index + answer_length), start_score + end_score))
           _lowercase       :							Dict      		=		sorted(lowerCamelCase,      key=lambda lowerCamelCase: x[1],      reverse=lowerCamelCase)
           _lowercase       :							List[str]      		=		[]
           for (start_index, end_index), score in scores:
               if start_index > end_index:
                   raise ValueError(F'''Wrong span indices: [{start_index}:{end_index}]''')
               _lowercase       :							Dict      		=		end_index - start_index + 1
               if length > max_answer_length:
                   raise ValueError(F'''Span is too long: {length} > {max_answer_length}''')
               if any(
                   start_index <= prev_start_index <= prev_end_index <= end_index
                   or prev_start_index <= start_index <= end_index <= prev_end_index
                   for (prev_start_index, prev_end_index) in chosen_span_intervals):
                   continue
               chosen_span_intervals.append((start_index, end_index))

               if len(lowerCamelCase) == top_spans:
                   break
           return chosen_span_intervals






@add_end_docstrings(_a   )
class _lowerCamelCase(						_a, _a   ):
       lowercase_ :	Union[str, Any]     =			VOCAB_FILES_NAMES
       lowercase_ :	Any     =			READER_PRETRAINED_VOCAB_FILES_MAP
       lowercase_ :	Dict     =			READER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
       lowercase_ :	Optional[int]     =			READER_PRETRAINED_INIT_CONFIGURATION
       lowercase_ :	str     =			["""input_ids""", """attention_mask"""]


