



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

from ...utils import logging


__UpperCAmelCase			  =			logging.get_logger(__name__)



class 							lowerCAmelCase_						(    UpperCamelCase__				):
       UpperCAmelCase__							:   Optional[int]  	=    CLIPConfig

       UpperCAmelCase__							:   List[str]  	=    ["CLIPEncoderLayer"]
       def __init__(     self,     SCREAMING_SNAKE_CASE_			)	->						Optional[int]:
          super().__init__(lowercase_			)

          UpperCamelCase       :    str		      =      CLIPVisionModelWithProjection(config.vision_config			)

          UpperCamelCase       :    List[Any]		      =      nn.Linear(config.vision_config.projection_dim,     1			)
          UpperCamelCase       :    Any		      =      nn.Linear(config.vision_config.projection_dim,     1			)






       @torch.no_grad()
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_=0.5,     SCREAMING_SNAKE_CASE_=0.5			)	->						List[str]:
          UpperCamelCase       :    List[str]		      =      self.vision_model(lowercase_			)[0]

          UpperCamelCase       :    Union[str, Any]		      =      self.p_head(lowercase_			)
          UpperCamelCase       :    List[str]		      =      nsfw_detected.flatten()
          UpperCamelCase       :    str		      =      nsfw_detected > p_threshold
          UpperCamelCase       :    List[Any]		      =      nsfw_detected.tolist()

          if any(lowercase_			):
             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.'			)

          for idx, nsfw_detected_ in enumerate(lowercase_			):
             if nsfw_detected_:
                UpperCamelCase       :    Dict		      =      np.zeros(images[idx].shape			)

          UpperCamelCase       :    Tuple		      =      self.w_head(lowercase_			)
          UpperCamelCase       :    Dict		      =      watermark_detected.flatten()
          UpperCamelCase       :    List[Any]		      =      watermark_detected > w_threshold
          UpperCamelCase       :    Union[str, Any]		      =      watermark_detected.tolist()

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

          for idx, watermark_detected_ in enumerate(lowercase_			):
             if watermark_detected_:
                UpperCamelCase       :    Optional[int]		      =      np.zeros(images[idx].shape			)

          return images, nsfw_detected, watermark_detected










'''simple docstring'''

import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


__lowerCAmelCase				:      int						      =5_0000
__lowerCAmelCase				:      Union[str, Any]						      =5000

__lowerCAmelCase      ,	__lowerCAmelCase				:      Tuple						      =os.path.split(__file__)
__lowerCAmelCase				:      int						      =os.path.join(RESULTS_BASEPATH, "results", RESULTS_FILENAME.replace(".py", ".json"))





@get_duration
def 		UpperCamelCase							(				_lowerCamelCase   :   datasets.Dataset		,			_lowerCamelCase   :   int			):
    for i in range(_lowerCamelCase			):
        A__							=     dataset[i]





@get_duration
def 		UpperCamelCase							(				_lowerCamelCase   :   datasets.Dataset		,			_lowerCamelCase   :   Optional[int]		,			_lowerCamelCase   :   Any			):
    for i in range(0		,			len(_lowerCamelCase			)		,			_lowerCamelCase			):
        A__							=     dataset[i : i + batch_size]





@get_duration
def 		UpperCamelCase							(				_lowerCamelCase   :   datasets.Dataset		,			_lowerCamelCase   :   Optional[Any]		,			_lowerCamelCase   :   List[str]			):
    with dataset.formatted_as(type=_lowerCamelCase			):
        for i in range(_lowerCamelCase			):
            A__							=     dataset[i]





@get_duration
def 		UpperCamelCase							(				_lowerCamelCase   :   datasets.Dataset		,			_lowerCamelCase   :   Optional[int]		,			_lowerCamelCase   :   int		,			_lowerCamelCase   :   Dict			):
    with dataset.formatted_as(type=_lowerCamelCase			):
        for i in range(0		,			_lowerCamelCase		,			_lowerCamelCase			):
            A__							=     dataset[i : i + batch_size]





def 		UpperCamelCase							(				):
    A__							=     {"num examples": SPEED_TEST_N_EXAMPLES}
    A__							=     [
        (read, {"length": SMALL_TEST}),
        (read, {"length": SPEED_TEST_N_EXAMPLES}),
        (read_batch, {"length": SPEED_TEST_N_EXAMPLES, "batch_size": 10}),
        (read_batch, {"length": SPEED_TEST_N_EXAMPLES, "batch_size": 1_00}),
        (read_batch, {"length": SPEED_TEST_N_EXAMPLES, "batch_size": 10_00}),
        (read_formatted, {"type": "numpy", "length": SMALL_TEST}),
        (read_formatted, {"type": "pandas", "length": SMALL_TEST}),
        (read_formatted, {"type": "torch", "length": SMALL_TEST}),
        (read_formatted, {"type": "tensorflow", "length": SMALL_TEST}),
        (read_formatted_batch, {"type": "numpy", "length": SMALL_TEST, "batch_size": 10}),
        (read_formatted_batch, {"type": "numpy", "length": SMALL_TEST, "batch_size": 10_00}),
    ]

    A__							=     [
        (read, {"length": SMALL_TEST}),
        (read, {"length": SPEED_TEST_N_EXAMPLES}),
        (read_batch, {"length": SPEED_TEST_N_EXAMPLES, "batch_size": 10}),
        (read_batch, {"length": SPEED_TEST_N_EXAMPLES, "batch_size": 1_00}),
        (read_batch, {"length": SPEED_TEST_N_EXAMPLES, "batch_size": 10_00}),
        (read_formatted, {"type": "numpy", "length": SMALL_TEST}),
        (read_formatted_batch, {"type": "numpy", "length": SMALL_TEST, "batch_size": 10}),
        (read_formatted_batch, {"type": "numpy", "length": SMALL_TEST, "batch_size": 10_00}),
    ]
    with tempfile.TemporaryDirectory() as tmp_dir:
        print("generating dataset"			)
        A__							=     datasets.Features(
            {"list": datasets.Sequence(datasets.Value("float32"			)			), "numbers": datasets.Value("float32"			)}			)
        A__							=     generate_example_dataset(
            os.path.join(_lowerCamelCase		,			"dataset.arrow"			)		,			_lowerCamelCase		,			num_examples=_lowerCamelCase		,			seq_shapes={"list": (1_00,)}		,			)
        print("first set of iterations"			)
        for func, kwargs in functions:
            print(func.__name__		,			str(_lowerCamelCase			)			)
            A__							=     func(_lowerCamelCase		,			**_lowerCamelCase			)

        print("shuffling dataset"			)
        A__							=     dataset.shuffle()
        print("Second set of iterations (after shuffling"			)
        for func, kwargs in functions_shuffled:
            print("shuffled "		,			func.__name__		,			str(_lowerCamelCase			)			)
            A__							=     func(
                _lowerCamelCase		,			**_lowerCamelCase			)

    with open(_lowerCamelCase		,			"wb"			) as f:
        f.write(json.dumps(_lowerCamelCase			).encode("utf-8"			)			)


if __name__ == "__main__":  # useful to run the profiler
  benchmark_iterating()



