
import json
import os
import shutil
import tempfile
import unittest
from multiprocessing import get_context
from pathlib import Path

import datasets
import numpy as np
from datasets import load_dataset
from parameterized import parameterized

from transformers import AutoProcessor
from transformers.models.wavaveca import WavaVecaCTCTokenizer, WavaVecaFeatureExtractor
from transformers.models.wavaveca.tokenization_wavaveca import VOCAB_FILES_NAMES
from transformers.testing_utils import require_pyctcdecode, require_torch, require_torchaudio, slow
from transformers.utils import FEATURE_EXTRACTOR_NAME, is_pyctcdecode_available, is_torch_available

from ..wavaveca.test_feature_extraction_wavaveca import floats_list


if is_pyctcdecode_available():
			from huggingface_hub import snapshot_download
			from pyctcdecode import BeamSearchDecoderCTC

			from transformers.models.wavaveca_with_lm import WavaVecaProcessorWithLM
			from transformers.models.wavaveca_with_lm.processing_wavaveca_with_lm import WavaVecaDecoderWithLMOutput

if is_torch_available():
			from transformers import WavaVecaForCTC




@require_pyctcdecode
class        _SCREAMING_SNAKE_CASE	( unittest.TestCase		):



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Tuple:
					lowerCamelCase_         =     "| <pad> <unk> <s> </s> a b c d e f g h i j k".split()
					lowerCamelCase_         =     dict(zip(lowercase		,    range(len(lowercase					)					)					)					)

					lowerCamelCase_         =     {
					    "unk_token": "<unk>",
					    "bos_token": "<s>",
					    "eos_token": "</s>",
					}
					lowerCamelCase_         =     {
					    "feature_size": 1,
					    "padding_value": 0.0,
					    "sampling_rate": 16000,
					    "return_attention_mask": False,
					    "do_normalize": True,
					}

					lowerCamelCase_         =     tempfile.mkdtemp()
					lowerCamelCase_         =     os.path.join(self.tmpdirname		,    VOCAB_FILES_NAMES["vocab_file"]					)
					lowerCamelCase_         =     os.path.join(self.tmpdirname		,    lowercase					)
					with open(self.vocab_file		,    "w"		,    encoding="utf-8"					) as fp:
									fp.write(json.dumps(lowercase					) + "\n"					)

					with open(self.feature_extraction_file		,    "w"		,    encoding="utf-8"					) as fp:
									fp.write(json.dumps(lowercase					) + "\n"					)

					# load decoder from hub
					lowerCamelCase_         =     "hf-internal-testing/ngram-beam-search-decoder"



	def        SCREAMING_SNAKE_CASE_(       self		,    **lowercase					)							->			Optional[int]:
					lowerCamelCase_         =     self.add_kwargs_tokens_map.copy()
					kwargs.update(lowercase					)
					return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname		,    **lowercase					)



	def        SCREAMING_SNAKE_CASE_(       self		,    **lowercase					)							->			List[Any]:
					return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname		,    **lowercase					)



	def        SCREAMING_SNAKE_CASE_(       self		,    **lowercase					)							->			List[str]:
					return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name		,    **lowercase					)



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Any:
					shutil.rmtree(self.tmpdirname					)



	def        SCREAMING_SNAKE_CASE_(       self					)							->			int:
					lowerCamelCase_         =     self.get_tokenizer()
					lowerCamelCase_         =     self.get_feature_extractor()
					lowerCamelCase_         =     self.get_decoder()

					lowerCamelCase_         =     WavaVecaProcessorWithLM(tokenizer=lowercase		,    feature_extractor=lowercase		,    decoder=lowercase					)

					processor.save_pretrained(self.tmpdirname					)
					lowerCamelCase_         =     WavaVecaProcessorWithLM.from_pretrained(self.tmpdirname					)

					# tokenizer
					self.assertEqual(processor.tokenizer.get_vocab()		,    tokenizer.get_vocab()					)
					self.assertIsInstance(processor.tokenizer		,    lowercase					)

					# feature extractor
					self.assertEqual(processor.feature_extractor.to_json_string()		,    feature_extractor.to_json_string()					)
					self.assertIsInstance(processor.feature_extractor		,    lowercase					)

					# decoder
					self.assertEqual(processor.decoder._alphabet.labels		,    decoder._alphabet.labels					)
					self.assertEqual(
					    processor.decoder.model_container[decoder._model_key]._unigram_set		,    decoder.model_container[decoder._model_key]._unigram_set		,    )
					self.assertIsInstance(processor.decoder		,    lowercase					)



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Optional[Any]:
					lowerCamelCase_         =     WavaVecaProcessorWithLM(
					    tokenizer=self.get_tokenizer()		,    feature_extractor=self.get_feature_extractor()		,    decoder=self.get_decoder()					)
					processor.save_pretrained(self.tmpdirname					)

					# make sure that error is thrown when decoder alphabet doesn't match
					lowerCamelCase_         =     WavaVecaProcessorWithLM.from_pretrained(
					    self.tmpdirname		,    alpha=5.0		,    beta=3.0		,    score_boundary=-7.0		,    unk_score_offset=3					)

					# decoder
					self.assertEqual(processor.language_model.alpha		,    5.0					)
					self.assertEqual(processor.language_model.beta		,    3.0					)
					self.assertEqual(processor.language_model.score_boundary		,    -7.0					)
					self.assertEqual(processor.language_model.unk_score_offset		,    3					)



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Union[str, Any]:
					lowerCamelCase_         =     self.get_tokenizer()
					# add token to trigger raise
					tokenizer.add_tokens(["xx"]					)
					with self.assertRaisesRegex(lowercase		,    "include"					):
									WavaVecaProcessorWithLM(
									    tokenizer=lowercase		,    feature_extractor=self.get_feature_extractor()		,    decoder=self.get_decoder()					)



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Any:
					lowerCamelCase_         =     self.get_feature_extractor()
					lowerCamelCase_         =     self.get_tokenizer()
					lowerCamelCase_         =     self.get_decoder()

					lowerCamelCase_         =     WavaVecaProcessorWithLM(tokenizer=lowercase		,    feature_extractor=lowercase		,    decoder=lowercase					)

					lowerCamelCase_         =     floats_list((3, 1000)					)

					lowerCamelCase_         =     feature_extractor(lowercase		,    return_tensors="np"					)
					lowerCamelCase_         =     processor(lowercase		,    return_tensors="np"					)

					for key in input_feat_extract.keys():
									self.assertAlmostEqual(input_feat_extract[key].sum()		,    input_processor[key].sum()		,    delta=1e-2					)



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Union[str, Any]:
					lowerCamelCase_         =     self.get_feature_extractor()
					lowerCamelCase_         =     self.get_tokenizer()
					lowerCamelCase_         =     self.get_decoder()

					lowerCamelCase_         =     WavaVecaProcessorWithLM(tokenizer=lowercase		,    feature_extractor=lowercase		,    decoder=lowercase					)

					lowerCamelCase_         =     "This is a test string"

					lowerCamelCase_         =     processor(text=lowercase					)

					lowerCamelCase_         =     tokenizer(lowercase					)

					for key in encoded_tok.keys():
									self.assertListEqual(encoded_tok[key]		,    encoded_processor[key]					)



	def        SCREAMING_SNAKE_CASE_(       self		,    lowercase=(2, 10, 16)		,    lowercase=77					)							->			Any:
					np.random.seed(lowercase					)
					return np.random.rand(*lowercase					)



	def        SCREAMING_SNAKE_CASE_(       self					)							->			List[str]:
					lowerCamelCase_         =     self.get_feature_extractor()
					lowerCamelCase_         =     self.get_tokenizer()
					lowerCamelCase_         =     self.get_decoder()

					lowerCamelCase_         =     WavaVecaProcessorWithLM(tokenizer=lowercase		,    feature_extractor=lowercase		,    decoder=lowercase					)

					lowerCamelCase_         =     self._get_dummy_logits(shape=(10, 16)		,    seed=13					)

					lowerCamelCase_         =     processor.decode(lowercase					)

					lowerCamelCase_         =     decoder.decode_beams(lowercase					)[0]

					self.assertEqual(decoded_decoder[0]		,    decoded_processor.text					)
					self.assertEqual("</s> <s> </s>"		,    decoded_processor.text					)
					self.assertEqual(decoded_decoder[-2]		,    decoded_processor.logit_score					)
					self.assertEqual(decoded_decoder[-1]		,    decoded_processor.lm_score					)



	@parameterized.expand([[None], ["fork"], ["spawn"]]					)
	def        SCREAMING_SNAKE_CASE_(       self		,    lowercase					)							->			Union[str, Any]:
					lowerCamelCase_         =     self.get_feature_extractor()
					lowerCamelCase_         =     self.get_tokenizer()
					lowerCamelCase_         =     self.get_decoder()

					lowerCamelCase_         =     WavaVecaProcessorWithLM(tokenizer=lowercase		,    feature_extractor=lowercase		,    decoder=lowercase					)

					lowerCamelCase_         =     self._get_dummy_logits()

					# note: pool should be instantiated *after* Wav2Vec2ProcessorWithLM.
					#       otherwise, the LM won't be available to the pool's sub-processes.
					# manual logic used to allow parameterized test for both pool=None and pool=Pool(...)
					if pool_context is None:
									lowerCamelCase_         =     processor.batch_decode(lowercase					)
					else:
									with get_context(lowercase					).Pool() as pool:
													lowerCamelCase_         =     processor.batch_decode(lowercase		,    lowercase					)

					lowerCamelCase_         =     list(lowercase					)

					with get_context("fork"					).Pool() as p:
									lowerCamelCase_         =     decoder.decode_beams_batch(lowercase		,    lowercase					)

					lowerCamelCase_      ,       lowerCamelCase_      ,       lowerCamelCase_         =     [], [], []
					for beams in decoded_beams:
									texts_decoder.append(beams[0][0]					)
									logit_scores_decoder.append(beams[0][-2]					)
									lm_scores_decoder.append(beams[0][-1]					)

					self.assertListEqual(lowercase		,    decoded_processor.text					)
					self.assertListEqual(["<s> <s> </s>", "<s> <s> <s>"]		,    decoded_processor.text					)
					self.assertListEqual(lowercase		,    decoded_processor.logit_score					)
					self.assertListEqual(lowercase		,    decoded_processor.lm_score					)



	def        SCREAMING_SNAKE_CASE_(       self					)							->			List[Any]:
					lowerCamelCase_         =     self.get_feature_extractor()
					lowerCamelCase_         =     self.get_tokenizer()
					lowerCamelCase_         =     self.get_decoder()

					lowerCamelCase_         =     WavaVecaProcessorWithLM(tokenizer=lowercase		,    feature_extractor=lowercase		,    decoder=lowercase					)

					lowerCamelCase_         =     self._get_dummy_logits()

					lowerCamelCase_         =     15
					lowerCamelCase_         =     -2_0.0
					lowerCamelCase_         =     -4.0

					lowerCamelCase_         =     processor.batch_decode(
					    lowercase		,    beam_width=lowercase		,    beam_prune_logp=lowercase		,    token_min_logp=lowercase		,    )
					lowerCamelCase_         =     decoded_processor_out.text

					lowerCamelCase_         =     list(lowercase					)

					with get_context("fork"					).Pool() as pool:
									lowerCamelCase_         =     decoder.decode_beams_batch(
									    lowercase		,    lowercase		,    beam_width=lowercase		,    beam_prune_logp=lowercase		,    token_min_logp=lowercase		,    )

					lowerCamelCase_         =     [d[0][0] for d in decoded_decoder_out]
					lowerCamelCase_         =     [d[0][2] for d in decoded_decoder_out]
					lowerCamelCase_         =     [d[0][3] for d in decoded_decoder_out]

					self.assertListEqual(lowercase		,    lowercase					)
					self.assertListEqual(["</s> <s> <s>", "<s> <s> <s>"]		,    lowercase					)

					self.assertTrue(np.array_equal(lowercase		,    decoded_processor_out.logit_score					)					)
					self.assertTrue(np.allclose([-2_0.0_5_4, -1_8.4_4_7]		,    lowercase		,    atol=1e-3					)					)

					self.assertTrue(np.array_equal(lowercase		,    decoded_processor_out.lm_score					)					)
					self.assertTrue(np.allclose([-1_5.5_5_4, -1_3.9_4_7_4]		,    lowercase		,    atol=1e-3					)					)



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Optional[Any]:
					lowerCamelCase_         =     self.get_feature_extractor()
					lowerCamelCase_         =     self.get_tokenizer()
					lowerCamelCase_         =     self.get_decoder()

					lowerCamelCase_         =     WavaVecaProcessorWithLM(tokenizer=lowercase		,    feature_extractor=lowercase		,    decoder=lowercase					)

					lowerCamelCase_         =     self._get_dummy_logits()

					lowerCamelCase_         =     2.0
					lowerCamelCase_         =     5.0
					lowerCamelCase_         =     -2_0.0
					lowerCamelCase_         =     True

					lowerCamelCase_         =     processor.batch_decode(
					    lowercase		,    alpha=lowercase		,    beta=lowercase		,    unk_score_offset=lowercase		,    lm_score_boundary=lowercase		,    )
					lowerCamelCase_         =     decoded_processor_out.text

					lowerCamelCase_         =     list(lowercase					)
					decoder.reset_params(
					    alpha=lowercase		,    beta=lowercase		,    unk_score_offset=lowercase		,    lm_score_boundary=lowercase		,    )

					with get_context("fork"					).Pool() as pool:
									lowerCamelCase_         =     decoder.decode_beams_batch(
									    lowercase		,    lowercase		,    )

					lowerCamelCase_         =     [d[0][0] for d in decoded_decoder_out]

					self.assertListEqual(lowercase		,    lowercase					)
					self.assertListEqual(["<s> </s> <s> </s> </s>", "</s> </s> <s> </s> </s>"]		,    lowercase					)
					lowerCamelCase_         =     processor.decoder.model_container[processor.decoder._model_key]
					self.assertEqual(lm_model.alpha		,    2.0					)
					self.assertEqual(lm_model.beta		,    5.0					)
					self.assertEqual(lm_model.unk_score_offset		,    -2_0.0					)
					self.assertEqual(lm_model.score_boundary		,    lowercase					)



	def        SCREAMING_SNAKE_CASE_(       self					)							->			List[Any]:
					lowerCamelCase_         =     WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"					)

					lowerCamelCase_         =     processor.decoder.model_container[processor.decoder._model_key]
					lowerCamelCase_         =     Path(language_model._kenlm_model.path.decode("utf-8"					)					).parent.parent.absolute()

					lowerCamelCase_         =     os.listdir(lowercase					)
					lowerCamelCase_         =     ["alphabet.json", "language_model"]

					downloaded_decoder_files.sort()
					expected_decoder_files.sort()

					# test that only decoder relevant files from
					# https://huggingface.co/hf-internal-testing/processor_with_lm/tree/main
					# are downloaded and none of the rest (e.g. README.md, ...)
					self.assertListEqual(lowercase		,    lowercase					)



	def        SCREAMING_SNAKE_CASE_(       self					)							->			int:
					lowerCamelCase_         =     snapshot_download("hf-internal-testing/processor_with_lm"					)

					lowerCamelCase_         =     WavaVecaProcessorWithLM.from_pretrained(lowercase					)

					lowerCamelCase_         =     processor.decoder.model_container[processor.decoder._model_key]
					lowerCamelCase_         =     Path(language_model._kenlm_model.path.decode("utf-8"					)					).parent.parent.absolute()

					lowerCamelCase_         =     os.listdir(lowercase					)
					lowerCamelCase_         =     os.listdir(lowercase					)

					local_decoder_files.sort()
					expected_decoder_files.sort()

					# test that both decoder form hub and local files in cache are the same
					self.assertListEqual(lowercase		,    lowercase					)



	def        SCREAMING_SNAKE_CASE_(       self					)							->			int:
					lowerCamelCase_         =     WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"					)
					lowerCamelCase_         =     AutoProcessor.from_pretrained("hf-internal-testing/processor_with_lm"					)

					lowerCamelCase_         =     floats_list((3, 1000)					)

					lowerCamelCase_         =     processor_wavaveca(lowercase		,    return_tensors="np"					)
					lowerCamelCase_         =     processor_auto(lowercase		,    return_tensors="np"					)

					for key in input_wavaveca.keys():
									self.assertAlmostEqual(input_wavaveca[key].sum()		,    input_auto[key].sum()		,    delta=1e-2					)

					lowerCamelCase_         =     self._get_dummy_logits()

					lowerCamelCase_         =     processor_wavaveca.batch_decode(lowercase					)
					lowerCamelCase_         =     processor_auto.batch_decode(lowercase					)

					self.assertListEqual(decoded_wavaveca.text		,    decoded_auto.text					)



	def        SCREAMING_SNAKE_CASE_(       self					)							->			List[Any]:
					lowerCamelCase_         =     self.get_feature_extractor()
					lowerCamelCase_         =     self.get_tokenizer()
					lowerCamelCase_         =     self.get_decoder()

					lowerCamelCase_         =     WavaVecaProcessorWithLM(tokenizer=lowercase		,    feature_extractor=lowercase		,    decoder=lowercase					)

					self.assertListEqual(
					    processor.model_input_names		,    feature_extractor.model_input_names		,    msg="`processor` and `feature_extractor` model input names do not match"		,    )



	@staticmethod
	def        SCREAMING_SNAKE_CASE_(       lowercase		,    lowercase					)							->			Optional[int]:
					lowerCamelCase_         =     [d[key] for d in offsets]
					return retrieved_list



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Dict:
					lowerCamelCase_         =     WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"					)
					lowerCamelCase_         =     self._get_dummy_logits()[0]

					lowerCamelCase_         =     processor.decode(lowercase		,    output_word_offsets=lowercase					)
					# check Wav2Vec2CTCTokenizerOutput keys for word
					self.assertEqual(len(outputs.keys()					)		,    4					)
					self.assertTrue("text" in outputs					)
					self.assertTrue("word_offsets" in outputs					)
					self.assertTrue(isinstance(lowercase		,    lowercase					)					)

					self.assertEqual(" ".join(self.get_from_offsets(outputs["word_offsets"]		,    "word"					)					)		,    outputs.text					)
					self.assertListEqual(self.get_from_offsets(outputs["word_offsets"]		,    "word"					)		,    ["<s>", "<s>", "</s>"]					)
					self.assertListEqual(self.get_from_offsets(outputs["word_offsets"]		,    "start_offset"					)		,    [0, 2, 4]					)
					self.assertListEqual(self.get_from_offsets(outputs["word_offsets"]		,    "end_offset"					)		,    [1, 3, 5]					)



	def        SCREAMING_SNAKE_CASE_(       self					)							->			List[str]:
					lowerCamelCase_         =     WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"					)
					lowerCamelCase_         =     self._get_dummy_logits()

					lowerCamelCase_         =     processor.batch_decode(lowercase		,    output_word_offsets=lowercase					)

					# check Wav2Vec2CTCTokenizerOutput keys for word
					self.assertEqual(len(outputs.keys()					)		,    4					)
					self.assertTrue("text" in outputs					)
					self.assertTrue("word_offsets" in outputs					)
					self.assertTrue(isinstance(lowercase		,    lowercase					)					)

					self.assertListEqual(
					    [" ".join(self.get_from_offsets(lowercase		,    "word"					)					) for o in outputs["word_offsets"]]		,    outputs.text					)
					self.assertListEqual(self.get_from_offsets(outputs["word_offsets"][0]		,    "word"					)		,    ["<s>", "<s>", "</s>"]					)
					self.assertListEqual(self.get_from_offsets(outputs["word_offsets"][0]		,    "start_offset"					)		,    [0, 2, 4]					)
					self.assertListEqual(self.get_from_offsets(outputs["word_offsets"][0]		,    "end_offset"					)		,    [1, 3, 5]					)



	@slow
	@require_torch
	@require_torchaudio
	def        SCREAMING_SNAKE_CASE_(       self					)							->			Optional[int]:
					import torch

					lowerCamelCase_         =     load_dataset("common_voice"		,    "en"		,    split="train"		,    streaming=lowercase					)
					lowerCamelCase_         =     ds.cast_column("audio"		,    datasets.Audio(sampling_rate=16000					)					)
					lowerCamelCase_         =     iter(lowercase					)
					lowerCamelCase_         =     next(lowercase					)

					lowerCamelCase_         =     AutoProcessor.from_pretrained("patrickvonplaten/wav2vec2-base-100h-with-lm"					)
					lowerCamelCase_         =     WavaVecaForCTC.from_pretrained("patrickvonplaten/wav2vec2-base-100h-with-lm"					)

					# compare to filename `common_voice_en_100038.mp3` of dataset viewer on https://huggingface.co/datasets/common_voice/viewer/en/train
					lowerCamelCase_         =     processor(sample["audio"]["array"]		,    return_tensors="pt"					).input_values

					with torch.no_grad():
									lowerCamelCase_         =     model(lowercase					).logits.cpu().numpy()

					lowerCamelCase_         =     processor.decode(logits[0]		,    output_word_offsets=lowercase					)

					lowerCamelCase_         =     model.config.inputs_to_logits_ratio / processor.feature_extractor.sampling_rate
					lowerCamelCase_         =     [
					    {
					        "start_time": d["start_offset"] * time_offset,
					        "end_time": d["end_offset"] * time_offset,
					        "word": d["word"],
					    }
					    for d in output["word_offsets"]
					]

					lowerCamelCase_         =     "WHY DOES MILISANDRA LOOK LIKE SHE WANTS TO CONSUME JOHN SNOW ON THE RIVER AT THE WALL"

					# output words
					self.assertEqual(" ".join(self.get_from_offsets(lowercase		,    "word"					)					)		,    lowercase					)
					self.assertEqual(" ".join(self.get_from_offsets(lowercase		,    "word"					)					)		,    output.text					)

					# output times
					lowerCamelCase_         =     torch.tensor(self.get_from_offsets(lowercase		,    "start_time"					)					)
					lowerCamelCase_         =     torch.tensor(self.get_from_offsets(lowercase		,    "end_time"					)					)

					# fmt: off
					lowerCamelCase_         =     torch.tensor([1.4_1_9_9, 1.6_5_9_9, 2.2_5_9_9, 3.0, 3.2_4, 3.5_9_9_9, 3.7_9_9_9, 4.0_9_9_9, 4.2_6, 4.9_4, 5.2_8, 5.6_5_9_9, 5.7_8, 5.9_4, 6.3_2, 6.5_3_9_9, 6.6_5_9_9]					)
					lowerCamelCase_         =     torch.tensor([1.5_3_9_9, 1.8_9_9_9, 2.9, 3.1_6, 3.5_3_9_9, 3.7_2, 4.0_1_9_9, 4.1_7_9_9, 4.7_6, 5.1_5_9_9, 5.5_5_9_9, 5.6_9_9_9, 5.8_6, 6.1_9_9_9, 6.3_8, 6.6_1_9_9, 6.9_4]					)
					# fmt: on

					self.assertTrue(torch.allclose(lowercase		,    lowercase		,    atol=0.0_1					)					)
					self.assertTrue(torch.allclose(lowercase		,    lowercase		,    atol=0.0_1					)					)




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

from ...utils import logging


__A							=logging.get_logger(__name__)




class        _SCREAMING_SNAKE_CASE	( snake_case_		):
	lowerCAmelCase__   =					CLIPConfig

	lowerCAmelCase__   =					['CLIPEncoderLayer']



	def __init__(       self		,    lowercase					)							->			int:
					super().__init__(lowercase					)

					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        SCREAMING_SNAKE_CASE_(       self		,    lowercase		,    lowercase		,    lowercase=0.5		,    lowercase=0.5					)							->			int:
					lowerCamelCase_         =     self.vision_model(lowercase					)[0]

					lowerCamelCase_         =     self.p_head(lowercase					)
					lowerCamelCase_         =     nsfw_detected.flatten()
					lowerCamelCase_         =     nsfw_detected > p_threshold
					lowerCamelCase_         =     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_:
													lowerCamelCase_         =     np.zeros(images[idx].shape					)

					lowerCamelCase_         =     self.w_head(lowercase					)
					lowerCamelCase_         =     watermark_detected.flatten()
					lowerCamelCase_         =     watermark_detected > w_threshold
					lowerCamelCase_         =     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_:
													lowerCamelCase_         =     np.zeros(images[idx].shape					)

					return images, nsfw_detected, watermark_detected


