


"""simple docstring"""


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

from ...utils import logging


A     :   str			    =			logging.get_logger(__name__)
class   _UpperCamelCase  (   lowerCAmelCase__		):
							'''simple docstring'''







							__UpperCAmelCase						:    List[str]				      =CLIPConfig

							__UpperCAmelCase						:    Dict				      =["""CLIPEncoderLayer"""]

							def __init__(     self    ,   __a	):
										super().__init__(__a	)

										__lowerCAmelCase				    =			CLIPVisionModelWithProjection(config.vision_config	)

										__lowerCAmelCase				    =			nn.Linear(config.vision_config.projection_dim    ,   1	)
										__lowerCAmelCase				    =			nn.Linear(config.vision_config.projection_dim    ,   1	)



							@torch.no_grad()
							def 			snake_case				(     self    ,   __a    ,   __a    ,   __a=0.5    ,   __a=0.5	):
										__lowerCAmelCase				    =			self.vision_model(__a	)[0]

										__lowerCAmelCase				    =			self.p_head(__a	)
										__lowerCAmelCase				    =			nsfw_detected.flatten()
										__lowerCAmelCase				    =			nsfw_detected > p_threshold
										__lowerCAmelCase				    =			nsfw_detected.tolist()

										if any(__a	):
													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(__a	):
													if nsfw_detected_:
																__lowerCAmelCase				    =			np.zeros(images[idx].shape	)

										__lowerCAmelCase				    =			self.w_head(__a	)
										__lowerCAmelCase				    =			watermark_detected.flatten()
										__lowerCAmelCase				    =			watermark_detected > w_threshold
										__lowerCAmelCase				    =			watermark_detected.tolist()

										if any(__a	):
													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(__a	):
													if watermark_detected_:
																__lowerCAmelCase				    =			np.zeros(images[idx].shape	)

										return images, nsfw_detected, watermark_detected








"""simple docstring"""


import unittest

from transformers import (
    MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    TextaTextGenerationPipeline,
    pipeline,
)
from transformers.testing_utils import is_pipeline_test, require_tf, require_torch
from transformers.utils import is_torch_available

from .test_pipelines_common import ANY


if is_torch_available():
	import torch
@is_pipeline_test
class   _UpperCamelCase  (   unittest.TestCase		):
							'''simple docstring'''







							__UpperCAmelCase						:    Union[str, Any]				      =MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
							__UpperCAmelCase						:    Union[str, Any]				      =TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING

							def 			snake_case				(     self    ,   __a    ,   __a    ,   __a	):
										__lowerCAmelCase				    =			TextaTextGenerationPipeline(model=__a    ,   tokenizer=__a	)
										return generator, ["Something to write", "Something else"]

							def 			snake_case				(     self    ,   __a    ,   __a	):
										__lowerCAmelCase				    =			generator("Something there"	)
										self.assertEqual(__a    ,   [{"generated_text": ANY(__a	)}]	)
										# These are encoder decoder, they don't just append to incoming string
										self.assertFalse(outputs[0]["generated_text"].startswith("Something there"	)	)

										__lowerCAmelCase				    =			generator(["This is great !", "Something else"]    ,   num_return_sequences=2    ,   do_sample=__a	)
										self.assertEqual(
										    __a    ,   [
										        [{"generated_text": ANY(__a	)}, {"generated_text": ANY(__a	)}],
										        [{"generated_text": ANY(__a	)}, {"generated_text": ANY(__a	)}],
										    ]    ,   )

										__lowerCAmelCase				    =			generator(
										    ["This is great !", "Something else"]    ,   num_return_sequences=2    ,   batch_size=2    ,   do_sample=__a	)
										self.assertEqual(
										    __a    ,   [
										        [{"generated_text": ANY(__a	)}, {"generated_text": ANY(__a	)}],
										        [{"generated_text": ANY(__a	)}, {"generated_text": ANY(__a	)}],
										    ]    ,   )

										with self.assertRaises(__a	):
													generator(4	)

							@require_torch
							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			pipeline("text2text-generation"    ,   model="patrickvonplaten/t5-tiny-random"    ,   framework="pt"	)
										# do_sample=False necessary for reproducibility
										__lowerCAmelCase				    =			generator("Something there"    ,   do_sample=__a	)
										self.assertEqual(__a    ,   [{"generated_text": ""}]	)

										__lowerCAmelCase				    =			3
										__lowerCAmelCase				    =			generator(
										    "Something there"    ,   num_return_sequences=__a    ,   num_beams=__a    ,   )
										__lowerCAmelCase				    =			[
										    {"generated_text": "Beide Beide Beide Beide Beide Beide Beide Beide Beide"},
										    {"generated_text": "Beide Beide Beide Beide Beide Beide Beide Beide"},
										    {"generated_text": ""},
										]
										self.assertEqual(__a    ,   __a	)

										__lowerCAmelCase				    =			generator("This is a test"    ,   do_sample=__a    ,   num_return_sequences=2    ,   return_tensors=__a	)
										self.assertEqual(
										    __a    ,   [
										        {"generated_token_ids": ANY(torch.Tensor	)},
										        {"generated_token_ids": ANY(torch.Tensor	)},
										    ]    ,   )
										__lowerCAmelCase				    =			generator.model.config.eos_token_id
										__lowerCAmelCase				    =			"<pad>"
										__lowerCAmelCase				    =			generator(
										    ["This is a test", "This is a second test"]    ,   do_sample=__a    ,   num_return_sequences=2    ,   batch_size=2    ,   return_tensors=__a    ,   )
										self.assertEqual(
										    __a    ,   [
										        [
										            {"generated_token_ids": ANY(torch.Tensor	)},
										            {"generated_token_ids": ANY(torch.Tensor	)},
										        ],
										        [
										            {"generated_token_ids": ANY(torch.Tensor	)},
										            {"generated_token_ids": ANY(torch.Tensor	)},
										        ],
										    ]    ,   )



							@require_tf
							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			pipeline("text2text-generation"    ,   model="patrickvonplaten/t5-tiny-random"    ,   framework="tf"	)
										# do_sample=False necessary for reproducibility
										__lowerCAmelCase				    =			generator("Something there"    ,   do_sample=__a	)
										self.assertEqual(__a    ,   [{"generated_text": ""}]	)




