'''simple docstring'''






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

from ...utils import logging


_UpperCAmelCase						:    List[str]     					=					logging.get_logger(__name__)





def    _SCREAMING_SNAKE_CASE		(			__snake_case       : Dict      ,		__snake_case       : Union[str, Any]							):
     _A								=     nn.functional.normalize(snake_case_							)
     _A								=     nn.functional.normalize(snake_case_							)
     return torch.mm(snake_case_      ,		normalized_text_embeds.t()							)







class lowercase_					(		UpperCamelCase_ ):







   """simple docstring"""

   __lowerCAmelCase    						=      CLIPConfig

   __lowerCAmelCase    						=      ["CLIPEncoderLayer"]
   def __init__(			self			: Optional[Any],   UpperCamelCase__			: Any    )     -> Union[str, Any]:
        super().__init__(__UpperCamelCase    )

        _A								=     CLIPVisionModel(config.vision_config    )
        _A								=     nn.Linear(config.vision_config.hidden_size,   config.projection_dim,   bias=__UpperCamelCase    )

        _A								=     nn.Parameter(torch.ones(17,   config.projection_dim    ),   requires_grad=__UpperCamelCase    )
        _A								=     nn.Parameter(torch.ones(3,   config.projection_dim    ),   requires_grad=__UpperCamelCase    )

        _A								=     nn.Parameter(torch.ones(17    ),   requires_grad=__UpperCamelCase    )
        _A								=     nn.Parameter(torch.ones(3    ),   requires_grad=__UpperCamelCase    )
   @torch.no_grad()
   def   __UpperCAmelCase      (			self			: List[Any],   UpperCamelCase__			: Any,   UpperCamelCase__			: List[Any]    )     -> List[Any]:
        _A								=     self.vision_model(__UpperCamelCase    )[1]  # pooled_output
        _A								=     self.visual_projection(__UpperCamelCase    )

        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
        _A								=     cosine_distance(__UpperCamelCase,   self.special_care_embeds    ).cpu().float().numpy()
        _A								=     cosine_distance(__UpperCamelCase,   self.concept_embeds    ).cpu().float().numpy()

        _A								=     []
        _A								=     image_embeds.shape[0]
        for i in range(__UpperCamelCase    ):
             _A								=     {'''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
             _A								=     0.0

             for concept_idx in range(len(special_cos_dist[0]    )    ):
                  _A								=     special_cos_dist[i][concept_idx]
                  _A								=     self.special_care_embeds_weights[concept_idx].item()
                  _A								=     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]}    )
                       _A								=     0.01

             for concept_idx in range(len(cos_dist[0]    )    ):
                  _A								=     cos_dist[i][concept_idx]
                  _A								=     self.concept_embeds_weights[concept_idx].item()
                  _A								=     round(concept_cos - concept_threshold + adjustment,   3    )
                  if result_img["concept_scores"][concept_idx] > 0:
                       result_img["bad_concepts"].append(__UpperCamelCase    )

             result.append(__UpperCamelCase    )

        _A								=     [len(res['bad_concepts']    ) > 0 for res in result]

        for idx, has_nsfw_concept in enumerate(__UpperCamelCase    ):
             if has_nsfw_concept:
                  if torch.is_tensor(__UpperCamelCase    ) or torch.is_tensor(images[0]    ):
                       _A								=     torch.zeros_like(images[idx]    )  # black image
                  else:
                       _A								=     np.zeros(images[idx].shape    )  # black image

        if any(__UpperCamelCase    ):
             logger.warning(
                 'Potential NSFW content was detected in one or more images. A black image will be returned instead.'
                 ' Try again with a different prompt and/or seed.'    )

        return images, has_nsfw_concepts

   @torch.no_grad()
   def   __UpperCAmelCase      (			self			: Any,   UpperCamelCase__			: List[str],   UpperCamelCase__			: Any    )     -> Tuple:
        _A								=     self.vision_model(__UpperCamelCase    )[1]  # pooled_output
        _A								=     self.visual_projection(__UpperCamelCase    )

        _A								=     cosine_distance(__UpperCamelCase,   self.special_care_embeds    )
        _A								=     cosine_distance(__UpperCamelCase,   self.concept_embeds    )

        # increase this value to create a stronger `nsfw` filter
        # at the cost of increasing the possibility of filtering benign images
        _A								=     0.0

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

        _A								=     (cos_dist - self.concept_embeds_weights) + special_adjustment
        # concept_scores = concept_scores.round(decimals=3)
        _A								=     torch.any(concept_scores > 0,   dim=1    )

        _A								=     0.0  # black image

        return images, has_nsfw_concepts




'''simple docstring'''






import multiprocessing
import os
from typing import BinaryIO, Optional, Union

from .. import Dataset, Features, NamedSplit, config
from ..formatting import query_table
from ..packaged_modules.csv.csv import Csv
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader







class lowercase_					(		_UpperCamelCase ):







   """simple docstring"""

   def __init__(			self			: Any,   UpperCamelCase__			: NestedDataStructureLike[PathLike],   UpperCamelCase__			: Optional[NamedSplit] = None,   UpperCamelCase__			: Optional[Features] = None,   UpperCamelCase__			: str = None,   UpperCamelCase__			: bool = False,   UpperCamelCase__			: bool = False,   UpperCamelCase__			: Optional[int] = None,   **UpperCamelCase__			: Optional[Any],   )     -> Tuple:
        super().__init__(
            UpperCamelCase__,   split=UpperCamelCase__,   features=UpperCamelCase__,   cache_dir=UpperCamelCase__,   keep_in_memory=UpperCamelCase__,   streaming=UpperCamelCase__,   num_proc=UpperCamelCase__,   **UpperCamelCase__,   )
        _A								=     path_or_paths if isinstance(UpperCamelCase__,   UpperCamelCase__    ) else {self.split: path_or_paths}
        _A								=     Csv(
            cache_dir=UpperCamelCase__,   data_files=UpperCamelCase__,   features=UpperCamelCase__,   **UpperCamelCase__,   )

   def   __UpperCAmelCase      (			self			: List[Any]    )     -> int:
        # Build iterable dataset
        if self.streaming:
             _A								=     self.builder.as_streaming_dataset(split=self.split    )
        # Build regular (map-style) dataset
        else:
             _A								=     None
             _A								=     None
             _A								=     None
             _A								=     None

             self.builder.download_and_prepare(
                 download_config=UpperCamelCase__,   download_mode=UpperCamelCase__,   verification_mode=UpperCamelCase__,   base_path=UpperCamelCase__,   num_proc=self.num_proc,   )
             _A								=     self.builder.as_dataset(
                 split=self.split,   verification_mode=UpperCamelCase__,   in_memory=self.keep_in_memory    )
        return dataset







class lowercase_					:







   """simple docstring"""

   def __init__(			self			: Optional[Any],   UpperCamelCase__			: Dataset,   UpperCamelCase__			: Union[PathLike, BinaryIO],   UpperCamelCase__			: Optional[int] = None,   UpperCamelCase__			: Optional[int] = None,   **UpperCamelCase__			: Dict,   )     -> Any:
        if num_proc is not None and num_proc <= 0:
             raise ValueError(f'num_proc {num_proc} must be an integer > 0.'    )

        _A								=     dataset
        _A								=     path_or_buf
        _A								=     batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
        _A								=     num_proc
        _A								=     'utf-8'
        _A								=     to_csv_kwargs
   def   __UpperCAmelCase      (			self			: List[Any]    )     -> int:
        _A								=     self.to_csv_kwargs.pop('path_or_buf',   UpperCamelCase__    )
        _A								=     self.to_csv_kwargs.pop('header',   UpperCamelCase__    )
        _A								=     self.to_csv_kwargs.pop('index',   UpperCamelCase__    )

        if isinstance(self.path_or_buf,   (str, bytes, os.PathLike)    ):
             with open(self.path_or_buf,   'wb+'    ) as buffer:
                  _A								=     self._write(file_obj=UpperCamelCase__,   header=UpperCamelCase__,   index=UpperCamelCase__,   **self.to_csv_kwargs    )
        else:
             _A								=     self._write(file_obj=self.path_or_buf,   header=UpperCamelCase__,   index=UpperCamelCase__,   **self.to_csv_kwargs    )
        return written
   def   __UpperCAmelCase      (			self			: str,   UpperCamelCase__			: Optional[int]    )     -> str:
        _A		,					_A		,					_A		,					_A								=     args

        _A								=     query_table(
            table=self.dataset.data,   key=slice(UpperCamelCase__,   offset + self.batch_size    ),   indices=self.dataset._indices,   )
        _A								=     batch.to_pandas().to_csv(
            path_or_buf=UpperCamelCase__,   header=header if (offset == 0) else False,   index=UpperCamelCase__,   **UpperCamelCase__    )
        return csv_str.encode(self.encoding    )

   def   __UpperCAmelCase      (			self			: Union[str, Any],   UpperCamelCase__			: BinaryIO,   UpperCamelCase__			: List[str],   UpperCamelCase__			: Dict,   **UpperCamelCase__			: Optional[int]    )     -> int:
        _A								=     0

        if self.num_proc is None or self.num_proc == 1:
             for offset in logging.tqdm(
                 range(0,   len(self.dataset    ),   self.batch_size    ),   unit='ba',   disable=not logging.is_progress_bar_enabled(),   desc='Creating CSV from Arrow format',   ):
                  _A								=     self._batch_csv((offset, header, index, to_csv_kwargs)    )
                  written += file_obj.write(UpperCamelCase__    )

        else:
             _A		,					_A								=     len(self.dataset    ), self.batch_size
             with multiprocessing.Pool(self.num_proc    ) as pool:
                  for csv_str in logging.tqdm(
                      pool.imap(
                          self._batch_csv,   [(offset, header, index, to_csv_kwargs) for offset in range(0,   UpperCamelCase__,   UpperCamelCase__    )],   ),   total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size,   unit='ba',   disable=not logging.is_progress_bar_enabled(),   desc='Creating CSV from Arrow format',   ):
                       written += file_obj.write(UpperCamelCase__    )

        return written



