



'''simple docstring'''



import copy

from ...configuration_utils import PretrainedConfig
from ...utils import add_start_docstrings


lowerCAmelCase_  :							Optional[Any]             =    R'\n    [`RagConfig`] stores the configuration of a *RagModel*. Configuration objects inherit from [`PretrainedConfig`] and\n    can be used to control the model outputs. Read the documentation from [`PretrainedConfig`] for more information.\n\n    Args:\n        title_sep (`str`, *optional*, defaults to  `" / "`):\n            Separator inserted between the title and the text of the retrieved document when calling [`RagRetriever`].\n        doc_sep (`str`, *optional*, defaults to  `" // "`):\n            Separator inserted between the text of the retrieved document and the original input when calling\n            [`RagRetriever`].\n        n_docs (`int`, *optional*, defaults to 5):\n            Number of documents to retrieve.\n        max_combined_length (`int`, *optional*, defaults to 300):\n            Max length of contextualized input returned by [`~RagRetriever.__call__`].\n        retrieval_vector_size (`int`, *optional*, defaults to 768):\n            Dimensionality of the document embeddings indexed by [`RagRetriever`].\n        retrieval_batch_size (`int`, *optional*, defaults to 8):\n            Retrieval batch size, defined as the number of queries issues concurrently to the faiss index encapsulated\n            [`RagRetriever`].\n        dataset (`str`, *optional*, defaults to `"wiki_dpr"`):\n            A dataset identifier of the indexed dataset in HuggingFace Datasets (list all available datasets and ids\n            using `datasets.list_datasets()`).\n        dataset_split (`str`, *optional*, defaults to `"train"`)\n            Which split of the `dataset` to load.\n        index_name (`str`, *optional*, defaults to `"compressed"`)\n            The index name of the index associated with the `dataset`. One can choose between `"legacy"`, `"exact"` and\n            `"compressed"`.\n        index_path (`str`, *optional*)\n            The path to the serialized faiss index on disk.\n        passages_path (`str`, *optional*):\n            A path to text passages compatible with the faiss index. Required if using\n            [`~models.rag.retrieval_rag.LegacyIndex`]\n        use_dummy_dataset (`bool`, *optional*, defaults to `False`)\n            Whether to load a "dummy" variant of the dataset specified by `dataset`.\n        label_smoothing (`float`, *optional*, defaults to 0.0):\n            Only relevant if `return_loss` is set to `True`. Controls the `epsilon` parameter value for label smoothing\n            in the loss calculation. If set to 0, no label smoothing is performed.\n        do_marginalize (`bool`, *optional*, defaults to `False`):\n            If `True`, the logits are marginalized over all documents by making use of\n            `torch.nn.functional.log_softmax`.\n        reduce_loss (`bool`, *optional*, defaults to `False`):\n            Whether or not to reduce the NLL loss using the `torch.Tensor.sum` operation.\n        do_deduplication (`bool`, *optional*, defaults to `True`):\n            Whether or not to deduplicate the generations from different context documents for a given input. Has to be\n            set to `False` if used while training with distributed backend.\n        exclude_bos_score (`bool`, *optional*, defaults to `False`):\n            Whether or not to disregard the BOS token when computing the loss.\n        output_retrieved(`bool`, *optional*, defaults to `False`):\n            If set to `True`, `retrieved_doc_embeds`, `retrieved_doc_ids`, `context_input_ids` and\n            `context_attention_mask` are returned. See returned tensors for more detail.\n        use_cache (`bool`, *optional*, defaults to `True`):\n            Whether or not the model should return the last key/values attentions (not used by all models).\n        forced_eos_token_id (`int`, *optional*):\n            The id of the token to force as the last generated token when `max_length` is reached. Usually set to\n            `eos_token_id`.\n'
@add_start_docstrings(snake_case_      )
class 						lowerCamelCase_				(       snake_case_      ):
						_lowerCAmelCase  : Optional[Any]						  =					'rag'
						_lowerCAmelCase  : List[str]						  =					True







						def __init__(	self :   Optional[int]					,     lowerCAmelCase__ :   int=None					,     lowerCAmelCase__ :   List[str]=True					,     lowerCAmelCase__ :   Optional[Any]=None					,     lowerCAmelCase__ :   Any=None					,     lowerCAmelCase__ :   Union[str, Any]=None					,     lowerCAmelCase__ :   List[Any]=None					,     lowerCAmelCase__ :   int=None					,     lowerCAmelCase__ :   int=" / "					,     lowerCAmelCase__ :   Any=" // "					,     lowerCAmelCase__ :   int=5					,     lowerCAmelCase__ :   Optional[Any]=3_00					,     lowerCAmelCase__ :   Tuple=7_68					,     lowerCAmelCase__ :   List[str]=8					,     lowerCAmelCase__ :   str="wiki_dpr"					,     lowerCAmelCase__ :   Dict="train"					,     lowerCAmelCase__ :   List[str]="compressed"					,     lowerCAmelCase__ :   Dict=None					,     lowerCAmelCase__ :   Tuple=None					,     lowerCAmelCase__ :   Any=False					,     lowerCAmelCase__ :   Any=False					,     lowerCAmelCase__ :   List[str]=0.0					,     lowerCAmelCase__ :   List[str]=True					,     lowerCAmelCase__ :   Tuple=False					,     lowerCAmelCase__ :   Any=False					,     lowerCAmelCase__ :   Optional[int]=False					,     lowerCAmelCase__ :   Tuple=True					,     lowerCAmelCase__ :   List[Any]=None					,     **lowerCAmelCase__ :   Optional[Any]					,     ):

													"""simple docstring"""
													super().__init__(
													    bos_token_id=lowerCAmelCase__					,     pad_token_id=lowerCAmelCase__					,     eos_token_id=lowerCAmelCase__					,     decoder_start_token_id=lowerCAmelCase__					,     forced_eos_token_id=lowerCAmelCase__					,     is_encoder_decoder=lowerCAmelCase__					,     prefix=lowerCAmelCase__					,     vocab_size=lowerCAmelCase__					,     **lowerCAmelCase__					,     )
													assert (
													    "question_encoder" in kwargs and "generator" in kwargs
													), "Config has to be initialized with question_encoder and generator config"
													SCREAMING_SNAKE_CASE       :   List[Any]         =      kwargs.pop('''question_encoder'''					)
													SCREAMING_SNAKE_CASE       :   Tuple         =      question_encoder_config.pop('''model_type'''					)
													SCREAMING_SNAKE_CASE       :   int         =      kwargs.pop('''generator'''					)
													SCREAMING_SNAKE_CASE       :   Dict         =      decoder_config.pop('''model_type'''					)

													from ..auto.configuration_auto import AutoConfig

													SCREAMING_SNAKE_CASE       :   str         =      AutoConfig.for_model(lowerCAmelCase__					,     **lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      AutoConfig.for_model(lowerCAmelCase__					,     **lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Optional[int]         =      reduce_loss
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      label_smoothing
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      exclude_bos_score
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      do_marginalize

													SCREAMING_SNAKE_CASE       :   Tuple         =      title_sep
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      doc_sep
													SCREAMING_SNAKE_CASE       :   str         =      n_docs
													SCREAMING_SNAKE_CASE       :   Optional[int]         =      max_combined_length

													SCREAMING_SNAKE_CASE       :   Any         =      dataset
													SCREAMING_SNAKE_CASE       :   Any         =      dataset_split
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      index_name

													SCREAMING_SNAKE_CASE       :   Any         =      retrieval_vector_size
													SCREAMING_SNAKE_CASE       :   Optional[int]         =      retrieval_batch_size
													SCREAMING_SNAKE_CASE       :   Optional[int]         =      passages_path
													SCREAMING_SNAKE_CASE       :   int         =      index_path
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      use_dummy_dataset

													SCREAMING_SNAKE_CASE       :   Optional[int]         =      output_retrieved

													SCREAMING_SNAKE_CASE       :   Optional[Any]         =      do_deduplication

													SCREAMING_SNAKE_CASE       :   Dict         =      use_cache

													if self.forced_eos_token_id is None:
																				SCREAMING_SNAKE_CASE       :   str         =      getattr(self.generator					,     '''forced_eos_token_id'''					,     lowerCAmelCase__					)







						@classmethod
						def        __lowercase     (	cls :   Tuple					,     lowerCAmelCase__ :   PretrainedConfig					,     lowerCAmelCase__ :   PretrainedConfig					,     **lowerCAmelCase__ :   Dict					):

													"""simple docstring"""
													return cls(question_encoder=question_encoder_config.to_dict()					,     generator=generator_config.to_dict()					,     **lowerCAmelCase__					)







						def        __lowercase     (	self :   List[str]					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   List[str]         =      copy.deepcopy(self.__dict__					)
													SCREAMING_SNAKE_CASE       :   List[Any]         =      self.question_encoder.to_dict()
													SCREAMING_SNAKE_CASE       :   str         =      self.generator.to_dict()
													SCREAMING_SNAKE_CASE       :   Optional[Any]         =      self.__class__.model_type
													return output









'''simple docstring'''



import torch
import torch.nn as nn
from transformers import CLIPConfig, CLIPVisionModel, PreTrainedModel

from ...utils import logging


lowerCAmelCase_  :							int             =    logging.get_logger(__name__)
def        UpperCAmelCase (    A	:					List[Any]       ,		A	:					Optional[int]					):
							SCREAMING_SNAKE_CASE       :   Any         =      nn.functional.normalize(A					)
							SCREAMING_SNAKE_CASE       :   List[str]         =      nn.functional.normalize(A					)
							return torch.mm(A       ,		normalized_text_embeds.t()					)






class 						lowerCamelCase_				(       snake_case_      ):
						_lowerCAmelCase  : Optional[Any]						  =					CLIPConfig

						_lowerCAmelCase  : List[str]						  =					['CLIPEncoderLayer']







						def __init__(	self :   int					,     lowerCAmelCase__ :   CLIPConfig					):

													"""simple docstring"""
													super().__init__(lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Any         =      CLIPVisionModel(config.vision_config					)
													SCREAMING_SNAKE_CASE       :   List[str]         =      nn.Linear(config.vision_config.hidden_size					,     config.projection_dim					,     bias=lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Any         =      nn.Parameter(torch.ones(17					,     config.projection_dim					)					,     requires_grad=lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   int         =      nn.Parameter(torch.ones(3					,     config.projection_dim					)					,     requires_grad=lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Dict         =      nn.Parameter(torch.ones(17					)					,     requires_grad=lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   List[str]         =      nn.Parameter(torch.ones(3					)					,     requires_grad=lowerCAmelCase__					)







						@torch.no_grad()
						def        __lowercase     (	self :   List[str]					,     lowerCAmelCase__ :   List[str]					,     lowerCAmelCase__ :   Optional[Any]					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   Any         =      self.vision_model(lowerCAmelCase__					)[1]  # pooled_output
													SCREAMING_SNAKE_CASE       :   Any         =      self.visual_projection(lowerCAmelCase__					)

													# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      cosine_distance(lowerCAmelCase__					,     self.special_care_embeds					).cpu().float().numpy()
													SCREAMING_SNAKE_CASE       :   List[str]         =      cosine_distance(lowerCAmelCase__					,     self.concept_embeds					).cpu().float().numpy()

													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      []
													SCREAMING_SNAKE_CASE       :   int         =      image_embeds.shape[0]
													for i in range(lowerCAmelCase__					):
																				SCREAMING_SNAKE_CASE       :   Any         =      {'''special_scores''': {}, '''special_care''': [], '''concept_scores''': {}, '''bad_concepts''': []}

																				# increase this value to create a stronger `nfsw` filter
																				# at the cost of increasing the possibility of filtering benign images
																				SCREAMING_SNAKE_CASE       :   int         =      0.0

																				for concept_idx in range(len(special_cos_dist[0]					)					):
																											SCREAMING_SNAKE_CASE       :   Optional[int]         =      special_cos_dist[i][concept_idx]
																											SCREAMING_SNAKE_CASE       :   int         =      self.special_care_embeds_weights[concept_idx].item()
																											SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      round(concept_cos - concept_threshold + adjustment					,     3					)
																											if result_img["special_scores"][concept_idx] > 0:
																																		result_img["special_care"].append({concept_idx, result_img['''special_scores'''][concept_idx]}					)
																																		SCREAMING_SNAKE_CASE       :   Tuple         =      0.01

																				for concept_idx in range(len(cos_dist[0]					)					):
																											SCREAMING_SNAKE_CASE       :   Optional[Any]         =      cos_dist[i][concept_idx]
																											SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      self.concept_embeds_weights[concept_idx].item()
																											SCREAMING_SNAKE_CASE       :   Dict         =      round(concept_cos - concept_threshold + adjustment					,     3					)
																											if result_img["concept_scores"][concept_idx] > 0:
																																		result_img["bad_concepts"].append(lowerCAmelCase__					)

																				result.append(lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Dict         =      [len(res['''bad_concepts''']					) > 0 for res in result]

													return images, has_nsfw_concepts







						@torch.no_grad()
						def        __lowercase     (	self :   List[str]					,     lowerCAmelCase__ :   torch.FloatTensor					,     lowerCAmelCase__ :   torch.FloatTensor					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      self.vision_model(lowerCAmelCase__					)[1]  # pooled_output
													SCREAMING_SNAKE_CASE       :   List[Any]         =      self.visual_projection(lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   List[str]         =      cosine_distance(lowerCAmelCase__					,     self.special_care_embeds					)
													SCREAMING_SNAKE_CASE       :   List[str]         =      cosine_distance(lowerCAmelCase__					,     self.concept_embeds					)

													# increase this value to create a stronger `nsfw` filter
													# at the cost of increasing the possibility of filtering benign images
													SCREAMING_SNAKE_CASE       :   Optional[int]         =      0.0

													SCREAMING_SNAKE_CASE       :   Any         =      special_cos_dist - self.special_care_embeds_weights + adjustment
													# special_scores = special_scores.round(decimals=3)
													SCREAMING_SNAKE_CASE       :   List[Any]         =      torch.any(special_scores > 0					,     dim=1					)
													SCREAMING_SNAKE_CASE       :   Dict         =      special_care * 0.01
													SCREAMING_SNAKE_CASE       :   List[str]         =      special_adjustment.unsqueeze(1					).expand(-1					,     cos_dist.shape[1]					)

													SCREAMING_SNAKE_CASE       :   Optional[int]         =      (cos_dist - self.concept_embeds_weights) + special_adjustment
													# concept_scores = concept_scores.round(decimals=3)
													SCREAMING_SNAKE_CASE       :   Any         =      torch.any(concept_scores > 0					,     dim=1					)

													return images, has_nsfw_concepts




