






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


__UpperCamelCase					:		int	  =						logging.get_logger(__name__)

__UpperCamelCase					:		Dict	  =						{"""vocab_file""": """vocab.txt""", """tokenizer_file""": """tokenizer.json"""}

__UpperCamelCase					:		Optional[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"""
        ),
    },
}
__UpperCamelCase					:		Tuple	  =						{
    """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"""
        ),
    },
}
__UpperCamelCase					:		Dict	  =						{
    """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"""
        ),
    },
}

__UpperCamelCase					:		Any	  =						{
    """facebook/dpr-ctx_encoder-single-nq-base""": 512,
    """facebook/dpr-ctx_encoder-multiset-base""": 512,
}
__UpperCamelCase					:		Union[str, Any]	  =						{
    """facebook/dpr-question_encoder-single-nq-base""": 512,
    """facebook/dpr-question_encoder-multiset-base""": 512,
}
__UpperCamelCase					:		List[Any]	  =						{
    """facebook/dpr-reader-single-nq-base""": 512,
    """facebook/dpr-reader-multiset-base""": 512,
}


__UpperCamelCase					:		List[str]	  =						{
    """facebook/dpr-ctx_encoder-single-nq-base""": {"""do_lower_case""": True},
    """facebook/dpr-ctx_encoder-multiset-base""": {"""do_lower_case""": True},
}
__UpperCamelCase					:		Tuple	  =						{
    """facebook/dpr-question_encoder-single-nq-base""": {"""do_lower_case""": True},
    """facebook/dpr-question_encoder-multiset-base""": {"""do_lower_case""": True},
}
__UpperCamelCase					:		Tuple	  =						{
    """facebook/dpr-reader-single-nq-base""": {"""do_lower_case""": True},
    """facebook/dpr-reader-multiset-base""": {"""do_lower_case""": True},
}


class   __SCREAMING_SNAKE_CASE(	a_	):
							_UpperCAmelCase     =       VOCAB_FILES_NAMES
							_UpperCAmelCase     =       CONTEXT_ENCODER_PRETRAINED_VOCAB_FILES_MAP
							_UpperCAmelCase     =       CONTEXT_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
							_UpperCAmelCase     =       CONTEXT_ENCODER_PRETRAINED_INIT_CONFIGURATION


class   __SCREAMING_SNAKE_CASE(	a_	):
							_UpperCAmelCase     =       VOCAB_FILES_NAMES
							_UpperCAmelCase     =       QUESTION_ENCODER_PRETRAINED_VOCAB_FILES_MAP
							_UpperCAmelCase     =       QUESTION_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
							_UpperCAmelCase     =       QUESTION_ENCODER_PRETRAINED_INIT_CONFIGURATION


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

__UpperCamelCase					:		Tuple	  =						collections.namedtuple("""DPRReaderOutput""", ["""start_logits""", """end_logits""", """relevance_logits"""])


__UpperCamelCase					:		List[str]	  =						R"""
    Return a dictionary with the token ids of the input strings and other information to give to `.decode_best_spans`.
    It converts the strings of a question and different passages (title and text) in a sequence of IDs (integers),
    using the tokenizer and vocabulary. The resulting `input_ids` is a matrix of size `(n_passages, sequence_length)`
    with the format:

    ```
    [CLS] <question token ids> [SEP] <titles ids> [SEP] <texts ids>
    ```

    Args:
        questions (`str` or `List[str]`):
            The questions to be encoded. You can specify one question for many passages. In this case, the question
            will be duplicated like `[questions] * n_passages`. Otherwise you have to specify as many questions as in
            `titles` or `texts`.
        titles (`str` or `List[str]`):
            The passages titles to be encoded. This can be a string or a list of strings if there are several passages.
        texts (`str` or `List[str]`):
            The passages texts to be encoded. This can be a string or a list of strings if there are several passages.
        padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`):
            Activates and controls padding. Accepts the following values:

            - `True` or `\'longest\'`: Pad to the longest sequence in the batch (or no padding if only a single sequence
              if provided).
            - `\'max_length\'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
              acceptable input length for the model if that argument is not provided.
            - `False` or `\'do_not_pad\'` (default): No padding (i.e., can output a batch with sequences of different
              lengths).
        truncation (`bool`, `str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*, defaults to `False`):
            Activates and controls truncation. Accepts the following values:

            - `True` or `\'longest_first\'`: Truncate to a maximum length specified with the argument `max_length` or to
              the maximum acceptable input length for the model if that argument is not provided. This will truncate
              token by token, removing a token from the longest sequence in the pair if a pair of sequences (or a batch
              of pairs) is provided.
            - `\'only_first\'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum
              acceptable input length for the model if that argument is not provided. This will only truncate the first
              sequence of a pair if a pair of sequences (or a batch of pairs) is provided.
            - `\'only_second\'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum
              acceptable input length for the model if that argument is not provided. This will only truncate the
              second sequence of a pair if a pair of sequences (or a batch of pairs) is provided.
            - `False` or `\'do_not_truncate\'` (default): No truncation (i.e., can output batch with sequence lengths
              greater than the model maximum admissible input size).
        max_length (`int`, *optional*):
                Controls the maximum length to use by one of the truncation/padding parameters.

                If left unset or set to `None`, this will use the predefined model maximum length if a maximum length
                is required by one of the truncation/padding parameters. If the model has no specific maximum input
                length (like XLNet) truncation/padding to a maximum length will be deactivated.
        return_tensors (`str` or [`~utils.TensorType`], *optional*):
                If set, will return tensors instead of list of python integers. Acceptable values are:

                - `\'tf\'`: Return TensorFlow `tf.constant` objects.
                - `\'pt\'`: Return PyTorch `torch.Tensor` objects.
                - `\'np\'`: Return Numpy `np.ndarray` objects.
        return_attention_mask (`bool`, *optional*):
            Whether or not to return the attention mask. If not set, will return the attention mask according to the
            specific tokenizer\'s default, defined by the `return_outputs` attribute.

            [What are attention masks?](../glossary#attention-mask)

    Returns:
        `Dict[str, List[List[int]]]`: A dictionary with the following keys:

        - `input_ids`: List of token ids to be fed to a model.
        - `attention_mask`: List of indices specifying which tokens should be attended to by the model.
    """


@add_start_docstrings(a_	)
class   __SCREAMING_SNAKE_CASE:
							def __call__(     self:      Optional[Any]       ,	UpperCamelCase:      Union[str, Any]       ,	UpperCamelCase:      Optional[str] = None       ,	UpperCamelCase:      Optional[str] = None       ,	UpperCamelCase:      Union[bool, str] = False       ,	UpperCamelCase:      Union[bool, str] = False       ,	UpperCamelCase:      Optional[int] = None       ,	UpperCamelCase:      Optional[Union[str, TensorType]] = None       ,	UpperCamelCase:      Optional[bool] = None       ,	**UpperCamelCase:      Any       ,	)		->		Dict:
											if titles is None and texts is None:
															return super().__call__(
															    __UpperCamelCase       ,	padding=__UpperCamelCase       ,	truncation=__UpperCamelCase       ,	max_length=__UpperCamelCase       ,	return_tensors=__UpperCamelCase       ,	return_attention_mask=__UpperCamelCase       ,	**__UpperCamelCase       ,	)
											elif titles is None or texts is None:
															snake_case__									=     titles if texts is None else texts
															return super().__call__(
															    __UpperCamelCase       ,	__UpperCamelCase       ,	padding=__UpperCamelCase       ,	truncation=__UpperCamelCase       ,	max_length=__UpperCamelCase       ,	return_tensors=__UpperCamelCase       ,	return_attention_mask=__UpperCamelCase       ,	**__UpperCamelCase       ,	)
											snake_case__									=     titles if not isinstance(__UpperCamelCase       ,	__UpperCamelCase						) else [titles]
											snake_case__									=     texts if not isinstance(__UpperCamelCase       ,	__UpperCamelCase						) else [texts]
											snake_case__									=     len(__UpperCamelCase						)
											snake_case__									=     questions if not isinstance(__UpperCamelCase       ,	__UpperCamelCase						) else [questions] * n_passages
											if len(__UpperCamelCase						) != len(__UpperCamelCase						):
															raise ValueError(
															    F'''There should be as many titles than texts but got {len(__UpperCamelCase						)} titles and {len(__UpperCamelCase						)} texts.'''						)
											snake_case__									=     super().__call__(__UpperCamelCase       ,	__UpperCamelCase       ,	padding=__UpperCamelCase       ,	truncation=__UpperCamelCase						)['input_ids']
											snake_case__									=     super().__call__(__UpperCamelCase       ,	add_special_tokens=__UpperCamelCase       ,	padding=__UpperCamelCase       ,	truncation=__UpperCamelCase						)['input_ids']
											snake_case__									=     {
											    '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(__UpperCamelCase       ,	__UpperCamelCase						)
											    ]
											}
											if return_attention_mask is not False:
															snake_case__									=     []
															for input_ids in encoded_inputs["input_ids"]:
																			attention_mask.append([int(input_id != self.pad_token_id						) for input_id in input_ids]						)
															snake_case__									=     attention_mask
											return self.pad(__UpperCamelCase       ,	padding=__UpperCamelCase       ,	max_length=__UpperCamelCase       ,	return_tensors=__UpperCamelCase						)
							def lowerCAmelCase_ (     self:      Optional[int]       ,	UpperCamelCase:      BatchEncoding       ,	UpperCamelCase:      DPRReaderOutput       ,	UpperCamelCase:      int = 16       ,	UpperCamelCase:      int = 64       ,	UpperCamelCase:      int = 4       ,	)		->		str:
											snake_case__									=     reader_input['input_ids']
											snake_case__							,						snake_case__							,						snake_case__									=     reader_output[:3]
											snake_case__									=     len(__UpperCamelCase						)
											snake_case__									=     sorted(range(__UpperCamelCase						)       ,	reverse=__UpperCamelCase       ,	key=relevance_logits.__getitem__						)
											snake_case__									=     []
											for doc_id in sorted_docs:
															snake_case__									=     list(input_ids[doc_id]						)
															# assuming question & title information is at the beginning of the sequence
															snake_case__									=     sequence_ids.index(self.sep_token_id       ,	2						) + 1  # second sep id
															if sequence_ids[-1] == self.pad_token_id:
																			snake_case__									=     sequence_ids.index(self.pad_token_id						)
															else:
																			snake_case__									=     len(__UpperCamelCase						)

															snake_case__									=     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=__UpperCamelCase       ,	top_spans=__UpperCamelCase       ,	)
															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=__UpperCamelCase       ,	start_index=__UpperCamelCase       ,	end_index=__UpperCamelCase       ,	text=self.decode(sequence_ids[start_index : end_index + 1]						)       ,	)						)
															if len(__UpperCamelCase						) >= num_spans:
																			break
											return nbest_spans_predictions[:num_spans]


							def lowerCAmelCase_ (     self:      Tuple       ,	UpperCamelCase:      List[int]       ,	UpperCamelCase:      List[int]       ,	UpperCamelCase:      int       ,	UpperCamelCase:      int       ,	)		->		List[Any]:
											snake_case__									=     []
											for start_index, start_score in enumerate(__UpperCamelCase						):
															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)						)
											snake_case__									=     sorted(__UpperCamelCase       ,	key=lambda UpperCamelCase						: x[1]       ,	reverse=__UpperCamelCase						)
											snake_case__									=     []
											for (start_index, end_index), score in scores:
															if start_index > end_index:
																			raise ValueError(F'''Wrong span indices: [{start_index}:{end_index}]'''						)
															snake_case__									=     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(__UpperCamelCase						) == top_spans:
																			break
											return chosen_span_intervals


@add_end_docstrings(a_	)
class   __SCREAMING_SNAKE_CASE(	a_    ,  a_	):
							_UpperCAmelCase     =       VOCAB_FILES_NAMES
							_UpperCAmelCase     =       READER_PRETRAINED_VOCAB_FILES_MAP
							_UpperCAmelCase     =       READER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
							_UpperCAmelCase     =       READER_PRETRAINED_INIT_CONFIGURATION
							_UpperCAmelCase     =       ["input_ids", "attention_mask"]











import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()


class   __SCREAMING_SNAKE_CASE(	a_    ,  a_    ,  unittest.TestCase	):
							_UpperCAmelCase     =       StableDiffusionSAGPipeline
							_UpperCAmelCase     =       TEXT_TO_IMAGE_PARAMS
							_UpperCAmelCase     =       TEXT_TO_IMAGE_BATCH_PARAMS
							_UpperCAmelCase     =       TEXT_TO_IMAGE_IMAGE_PARAMS
							_UpperCAmelCase     =       TEXT_TO_IMAGE_IMAGE_PARAMS
							_UpperCAmelCase     =       False
							def lowerCAmelCase_ (     self:      int						)		->		Any:
											torch.manual_seed(0						)
											snake_case__									=     UNetaDConditionModel(
											    block_out_channels=(32, 64)       ,	layers_per_block=2       ,	sample_size=32       ,	in_channels=4       ,	out_channels=4       ,	down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')       ,	up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')       ,	cross_attention_dim=32       ,	)
											snake_case__									=     DDIMScheduler(
											    beta_start=0.00_085       ,	beta_end=0.012       ,	beta_schedule='scaled_linear'       ,	clip_sample=UpperCamelCase       ,	set_alpha_to_one=UpperCamelCase       ,	)
											torch.manual_seed(0						)
											snake_case__									=     AutoencoderKL(
											    block_out_channels=[32, 64]       ,	in_channels=3       ,	out_channels=3       ,	down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']       ,	up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']       ,	latent_channels=4       ,	)
											torch.manual_seed(0						)
											snake_case__									=     CLIPTextConfig(
											    bos_token_id=0       ,	eos_token_id=2       ,	hidden_size=32       ,	intermediate_size=37       ,	layer_norm_eps=1e-05       ,	num_attention_heads=4       ,	num_hidden_layers=5       ,	pad_token_id=1       ,	vocab_size=10_00       ,	)
											snake_case__									=     CLIPTextModel(UpperCamelCase						)
											snake_case__									=     CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'						)

											snake_case__									=     {
											    'unet': unet,
											    'scheduler': scheduler,
											    'vae': vae,
											    'text_encoder': text_encoder,
											    'tokenizer': tokenizer,
											    'safety_checker': None,
											    'feature_extractor': None,
											}
											return components
							def lowerCAmelCase_ (     self:      Optional[int]       ,	UpperCamelCase:      Dict       ,	UpperCamelCase:      Optional[int]=0						)		->		Union[str, Any]:
											if str(UpperCamelCase						).startswith('mps'						):
															snake_case__									=     torch.manual_seed(UpperCamelCase						)
											else:
															snake_case__									=     torch.Generator(device=UpperCamelCase						).manual_seed(UpperCamelCase						)
											snake_case__									=     {
											    'prompt': '.',
											    'generator': generator,
											    'num_inference_steps': 2,
											    'guidance_scale': 1.0,
											    'sag_scale': 1.0,
											    'output_type': 'numpy',
											}
											return inputs


							def lowerCAmelCase_ (     self:      Union[str, Any]						)		->		Dict:
											super().test_inference_batch_single_identical(expected_max_diff=3e-3						)


@slow
@require_torch_gpu
class   __SCREAMING_SNAKE_CASE(	unittest.TestCase	):
							def lowerCAmelCase_ (     self:      Optional[int]						)		->		Optional[int]:
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()
							def lowerCAmelCase_ (     self:      Dict						)		->		Tuple:
											snake_case__									=     StableDiffusionSAGPipeline.from_pretrained('CompVis/stable-diffusion-v1-4'						)
											snake_case__									=     sag_pipe.to(UpperCamelCase						)
											sag_pipe.set_progress_bar_config(disable=UpperCamelCase						)

											snake_case__									=     '.'
											snake_case__									=     torch.manual_seed(0						)
											snake_case__									=     sag_pipe(
											    [prompt]       ,	generator=UpperCamelCase       ,	guidance_scale=7.5       ,	sag_scale=1.0       ,	num_inference_steps=20       ,	output_type='np'						)

											snake_case__									=     output.images

											snake_case__									=     image[0, -3:, -3:, -1]

											assert image.shape == (1, 5_12, 5_12, 3)
											snake_case__									=     np.array([0.1_568, 0.1_738, 0.1_695, 0.1_693, 0.1_507, 0.1_705, 0.1_547, 0.1_751, 0.1_949]						)

											assert np.abs(image_slice.flatten() - expected_slice						).max() < 5e-2
							def lowerCAmelCase_ (     self:      Union[str, Any]						)		->		List[Any]:
											snake_case__									=     StableDiffusionSAGPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'						)
											snake_case__									=     sag_pipe.to(UpperCamelCase						)
											sag_pipe.set_progress_bar_config(disable=UpperCamelCase						)

											snake_case__									=     '.'
											snake_case__									=     torch.manual_seed(0						)
											snake_case__									=     sag_pipe(
											    [prompt]       ,	generator=UpperCamelCase       ,	guidance_scale=7.5       ,	sag_scale=1.0       ,	num_inference_steps=20       ,	output_type='np'						)

											snake_case__									=     output.images

											snake_case__									=     image[0, -3:, -3:, -1]

											assert image.shape == (1, 5_12, 5_12, 3)
											snake_case__									=     np.array([0.3_459, 0.2_876, 0.2_537, 0.3_002, 0.2_671, 0.2_160, 0.3_026, 0.2_262, 0.2_371]						)

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


							def lowerCAmelCase_ (     self:      Union[str, Any]						)		->		Optional[Any]:
											snake_case__									=     StableDiffusionSAGPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'						)
											snake_case__									=     sag_pipe.to(UpperCamelCase						)
											sag_pipe.set_progress_bar_config(disable=UpperCamelCase						)

											snake_case__									=     '.'
											snake_case__									=     torch.manual_seed(0						)
											snake_case__									=     sag_pipe(
											    [prompt]       ,	width=7_68       ,	height=5_12       ,	generator=UpperCamelCase       ,	guidance_scale=7.5       ,	sag_scale=1.0       ,	num_inference_steps=20       ,	output_type='np'       ,	)

											snake_case__									=     output.images

											assert image.shape == (1, 5_12, 7_68, 3)




