






"""simple docstring"""







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  lowerCamelCase__ (	unittest.TestCase			):






				def     _UpperCamelCase							(			self ):
										UpperCAmelCase          =	'| <pad> <unk> <s> </s> a b c d e f g h i j k'.split()
										UpperCAmelCase          =	dict(zip(_UpperCAmelCase     ,range(len(_UpperCAmelCase ) ) ) )

										UpperCAmelCase          =	{
										    'unk_token': '<unk>',
										    'bos_token': '<s>',
										    'eos_token': '</s>',
										}
										UpperCAmelCase          =	{
										    'feature_size': 1,
										    'padding_value': 0.0,
										    'sampling_rate': 16_000,
										    'return_attention_mask': False,
										    'do_normalize': True,
										}

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

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

										# load decoder from hub
										UpperCAmelCase          =	'hf-internal-testing/ngram-beam-search-decoder'






				def     _UpperCamelCase							(			self     ,**A ):
										UpperCAmelCase          =	self.add_kwargs_tokens_map.copy()
										kwargs.update(_UpperCAmelCase )
										return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname     ,**_UpperCAmelCase )






				def     _UpperCamelCase							(			self     ,**A ):
										return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname     ,**_UpperCAmelCase )






				def     _UpperCamelCase							(			self     ,**A ):
										return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name     ,**_UpperCAmelCase )






				def     _UpperCamelCase							(			self ):
										shutil.rmtree(self.tmpdirname )






				def     _UpperCamelCase							(			self ):
										UpperCAmelCase          =	self.get_tokenizer()
										UpperCAmelCase          =	self.get_feature_extractor()
										UpperCAmelCase          =	self.get_decoder()

										UpperCAmelCase          =	WavaVecaProcessorWithLM(tokenizer=_UpperCAmelCase     ,feature_extractor=_UpperCAmelCase     ,decoder=_UpperCAmelCase )

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

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

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

										# 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     ,_UpperCAmelCase )






				def     _UpperCamelCase							(			self ):
										UpperCAmelCase          =	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
										UpperCAmelCase          =	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     _UpperCamelCase							(			self ):
										UpperCAmelCase          =	self.get_tokenizer()
										# add token to trigger raise
										tokenizer.add_tokens(["""xx"""] )
										with self.assertRaisesRegex(_UpperCAmelCase     ,"""include""" ):
																WavaVecaProcessorWithLM(
																    tokenizer=_UpperCAmelCase     ,feature_extractor=self.get_feature_extractor()     ,decoder=self.get_decoder() )






				def     _UpperCamelCase							(			self ):
										UpperCAmelCase          =	self.get_feature_extractor()
										UpperCAmelCase          =	self.get_tokenizer()
										UpperCAmelCase          =	self.get_decoder()

										UpperCAmelCase          =	WavaVecaProcessorWithLM(tokenizer=_UpperCAmelCase     ,feature_extractor=_UpperCAmelCase     ,decoder=_UpperCAmelCase )

										UpperCAmelCase          =	floats_list((3, 1_000) )

										UpperCAmelCase          =	feature_extractor(_UpperCAmelCase     ,return_tensors="""np""" )
										UpperCAmelCase          =	processor(_UpperCAmelCase     ,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     _UpperCamelCase							(			self ):
										UpperCAmelCase          =	self.get_feature_extractor()
										UpperCAmelCase          =	self.get_tokenizer()
										UpperCAmelCase          =	self.get_decoder()

										UpperCAmelCase          =	WavaVecaProcessorWithLM(tokenizer=_UpperCAmelCase     ,feature_extractor=_UpperCAmelCase     ,decoder=_UpperCAmelCase )

										UpperCAmelCase          =	'This is a test string'

										UpperCAmelCase          =	processor(text=_UpperCAmelCase )

										UpperCAmelCase          =	tokenizer(_UpperCAmelCase )

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






				def     _UpperCamelCase							(			self     ,A=(2, 10, 16)     ,A=77 ):
										np.random.seed(_UpperCAmelCase )
										return np.random.rand(*_UpperCAmelCase )






				def     _UpperCamelCase							(			self ):
										UpperCAmelCase          =	self.get_feature_extractor()
										UpperCAmelCase          =	self.get_tokenizer()
										UpperCAmelCase          =	self.get_decoder()

										UpperCAmelCase          =	WavaVecaProcessorWithLM(tokenizer=_UpperCAmelCase     ,feature_extractor=_UpperCAmelCase     ,decoder=_UpperCAmelCase )

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

										UpperCAmelCase          =	processor.decode(_UpperCAmelCase )

										UpperCAmelCase          =	decoder.decode_beams(_UpperCAmelCase )[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     _UpperCamelCase							(			self     ,A ):
										UpperCAmelCase          =	self.get_feature_extractor()
										UpperCAmelCase          =	self.get_tokenizer()
										UpperCAmelCase          =	self.get_decoder()

										UpperCAmelCase          =	WavaVecaProcessorWithLM(tokenizer=_UpperCAmelCase     ,feature_extractor=_UpperCAmelCase     ,decoder=_UpperCAmelCase )

										UpperCAmelCase          =	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:
																UpperCAmelCase          =	processor.batch_decode(_UpperCAmelCase )
										else:
																with get_context(_UpperCAmelCase ).Pool() as pool:
																						UpperCAmelCase          =	processor.batch_decode(_UpperCAmelCase     ,_UpperCAmelCase )

										UpperCAmelCase          =	list(_UpperCAmelCase )

										with get_context("""fork""" ).Pool() as p:
																UpperCAmelCase          =	decoder.decode_beams_batch(_UpperCAmelCase     ,_UpperCAmelCase )

										UpperCAmelCase          =	[], [], []
										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(_UpperCAmelCase     ,decoded_processor.text )
										self.assertListEqual(["""<s> <s> </s>""", """<s> <s> <s>"""]     ,decoded_processor.text )
										self.assertListEqual(_UpperCAmelCase     ,decoded_processor.logit_score )
										self.assertListEqual(_UpperCAmelCase     ,decoded_processor.lm_score )






				def     _UpperCamelCase							(			self ):
										UpperCAmelCase          =	self.get_feature_extractor()
										UpperCAmelCase          =	self.get_tokenizer()
										UpperCAmelCase          =	self.get_decoder()

										UpperCAmelCase          =	WavaVecaProcessorWithLM(tokenizer=_UpperCAmelCase     ,feature_extractor=_UpperCAmelCase     ,decoder=_UpperCAmelCase )

										UpperCAmelCase          =	self._get_dummy_logits()

										UpperCAmelCase          =	15
										UpperCAmelCase          =	-20.0
										UpperCAmelCase          =	-4.0

										UpperCAmelCase          =	processor.batch_decode(
										    _UpperCAmelCase     ,beam_width=_UpperCAmelCase     ,beam_prune_logp=_UpperCAmelCase     ,token_min_logp=_UpperCAmelCase     ,)
										UpperCAmelCase          =	decoded_processor_out.text

										UpperCAmelCase          =	list(_UpperCAmelCase )

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

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

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

										self.assertTrue(np.array_equal(_UpperCAmelCase     ,decoded_processor_out.logit_score ) )
										self.assertTrue(np.allclose([-20.054, -18.447]     ,_UpperCAmelCase     ,atol=1e-3 ) )

										self.assertTrue(np.array_equal(_UpperCAmelCase     ,decoded_processor_out.lm_score ) )
										self.assertTrue(np.allclose([-15.554, -13.9474]     ,_UpperCAmelCase     ,atol=1e-3 ) )






				def     _UpperCamelCase							(			self ):
										UpperCAmelCase          =	self.get_feature_extractor()
										UpperCAmelCase          =	self.get_tokenizer()
										UpperCAmelCase          =	self.get_decoder()

										UpperCAmelCase          =	WavaVecaProcessorWithLM(tokenizer=_UpperCAmelCase     ,feature_extractor=_UpperCAmelCase     ,decoder=_UpperCAmelCase )

										UpperCAmelCase          =	self._get_dummy_logits()

										UpperCAmelCase          =	2.0
										UpperCAmelCase          =	5.0
										UpperCAmelCase          =	-20.0
										UpperCAmelCase          =	True

										UpperCAmelCase          =	processor.batch_decode(
										    _UpperCAmelCase     ,alpha=_UpperCAmelCase     ,beta=_UpperCAmelCase     ,unk_score_offset=_UpperCAmelCase     ,lm_score_boundary=_UpperCAmelCase     ,)
										UpperCAmelCase          =	decoded_processor_out.text

										UpperCAmelCase          =	list(_UpperCAmelCase )
										decoder.reset_params(
										    alpha=_UpperCAmelCase     ,beta=_UpperCAmelCase     ,unk_score_offset=_UpperCAmelCase     ,lm_score_boundary=_UpperCAmelCase     ,)

										with get_context("""fork""" ).Pool() as pool:
																UpperCAmelCase          =	decoder.decode_beams_batch(
																    _UpperCAmelCase     ,_UpperCAmelCase     ,)

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

										self.assertListEqual(_UpperCAmelCase     ,_UpperCAmelCase )
										self.assertListEqual(["""<s> </s> <s> </s> </s>""", """</s> </s> <s> </s> </s>"""]     ,_UpperCAmelCase )
										UpperCAmelCase          =	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     ,-20.0 )
										self.assertEqual(lm_model.score_boundary     ,_UpperCAmelCase )






				def     _UpperCamelCase							(			self ):
										UpperCAmelCase          =	WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm""" )

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

										UpperCAmelCase          =	os.listdir(_UpperCAmelCase )
										UpperCAmelCase          =	['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(_UpperCAmelCase     ,_UpperCAmelCase )






				def     _UpperCamelCase							(			self ):
										UpperCAmelCase          =	snapshot_download("""hf-internal-testing/processor_with_lm""" )

										UpperCAmelCase          =	WavaVecaProcessorWithLM.from_pretrained(_UpperCAmelCase )

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

										UpperCAmelCase          =	os.listdir(_UpperCAmelCase )
										UpperCAmelCase          =	os.listdir(_UpperCAmelCase )

										local_decoder_files.sort()
										expected_decoder_files.sort()

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






				def     _UpperCamelCase							(			self ):
										UpperCAmelCase          =	WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm""" )
										UpperCAmelCase          =	AutoProcessor.from_pretrained("""hf-internal-testing/processor_with_lm""" )

										UpperCAmelCase          =	floats_list((3, 1_000) )

										UpperCAmelCase          =	processor_wavaveca(_UpperCAmelCase     ,return_tensors="""np""" )
										UpperCAmelCase          =	processor_auto(_UpperCAmelCase     ,return_tensors="""np""" )

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

										UpperCAmelCase          =	self._get_dummy_logits()

										UpperCAmelCase          =	processor_wavaveca.batch_decode(_UpperCAmelCase )
										UpperCAmelCase          =	processor_auto.batch_decode(_UpperCAmelCase )

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






				def     _UpperCamelCase							(			self ):
										UpperCAmelCase          =	self.get_feature_extractor()
										UpperCAmelCase          =	self.get_tokenizer()
										UpperCAmelCase          =	self.get_decoder()

										UpperCAmelCase          =	WavaVecaProcessorWithLM(tokenizer=_UpperCAmelCase     ,feature_extractor=_UpperCAmelCase     ,decoder=_UpperCAmelCase )

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






				@staticmethod
				def     _UpperCamelCase							(			A     ,A ):
										UpperCAmelCase          =	[d[key] for d in offsets]
										return retrieved_list






				def     _UpperCamelCase							(			self ):
										UpperCAmelCase          =	WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm""" )
										UpperCAmelCase          =	self._get_dummy_logits()[0]

										UpperCAmelCase          =	processor.decode(_UpperCAmelCase     ,output_word_offsets=_UpperCAmelCase )
										# 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(_UpperCAmelCase     ,_UpperCAmelCase ) )

										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     _UpperCamelCase							(			self ):
										UpperCAmelCase          =	WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm""" )
										UpperCAmelCase          =	self._get_dummy_logits()

										UpperCAmelCase          =	processor.batch_decode(_UpperCAmelCase     ,output_word_offsets=_UpperCAmelCase )

										# 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(_UpperCAmelCase     ,_UpperCAmelCase ) )

										self.assertListEqual(
										    [""" """.join(self.get_from_offsets(_UpperCAmelCase     ,"""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     _UpperCamelCase							(			self ):
										import torch

										UpperCAmelCase          =	load_dataset("""common_voice"""     ,"""en"""     ,split="""train"""     ,streaming=_UpperCAmelCase )
										UpperCAmelCase          =	ds.cast_column("""audio"""     ,datasets.Audio(sampling_rate=16_000 ) )
										UpperCAmelCase          =	iter(_UpperCAmelCase )
										UpperCAmelCase          =	next(_UpperCAmelCase )

										UpperCAmelCase          =	AutoProcessor.from_pretrained("""patrickvonplaten/wav2vec2-base-100h-with-lm""" )
										UpperCAmelCase          =	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
										UpperCAmelCase          =	processor(sample["""audio"""]["""array"""]     ,return_tensors="""pt""" ).input_values

										with torch.no_grad():
																UpperCAmelCase          =	model(_UpperCAmelCase ).logits.cpu().numpy()

										UpperCAmelCase          =	processor.decode(logits[0]     ,output_word_offsets=_UpperCAmelCase )

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

										UpperCAmelCase          =	'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(_UpperCAmelCase     ,"""word""" ) )     ,_UpperCAmelCase )
										self.assertEqual(""" """.join(self.get_from_offsets(_UpperCAmelCase     ,"""word""" ) )     ,output.text )

										# output times
										UpperCAmelCase          =	torch.tensor(self.get_from_offsets(_UpperCAmelCase     ,"""start_time""" ) )
										UpperCAmelCase          =	torch.tensor(self.get_from_offsets(_UpperCAmelCase     ,"""end_time""" ) )

										# fmt: off
										UpperCAmelCase          =	torch.tensor([1.4199, 1.6599, 2.2599, 3.0, 3.24, 3.5999, 3.7999, 4.0999, 4.26, 4.94, 5.28, 5.6599, 5.78, 5.94, 6.32, 6.5399, 6.6599] )
										UpperCAmelCase          =	torch.tensor([1.5399, 1.8999, 2.9, 3.16, 3.5399, 3.72, 4.0199, 4.1799, 4.76, 5.1599, 5.5599, 5.6999, 5.86, 6.1999, 6.38, 6.6199, 6.94] )
										# fmt: on

										self.assertTrue(torch.allclose(_UpperCAmelCase     ,_UpperCAmelCase     ,atol=0.01 ) )
										self.assertTrue(torch.allclose(_UpperCAmelCase     ,_UpperCAmelCase     ,atol=0.01 ) )











"""simple docstring"""







import hashlib
import unittest
from typing import Dict

import numpy as np

from transformers import (
    MODEL_FOR_MASK_GENERATION_MAPPING,
    TF_MODEL_FOR_MASK_GENERATION_MAPPING,
    is_vision_available,
    pipeline,
)
from transformers.pipelines import MaskGenerationPipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_tf,
    require_torch,
    require_vision,
    slow,
)


if is_vision_available():
  from PIL import Image
else:







  class  lowerCamelCase__ :






      @staticmethod
      def     _UpperCamelCase							(			*A     ,**A ):
            pass

def _a			(		_snake_case							):







      """simple docstring"""






      UpperCAmelCase          =	hashlib.mda(image.tobytes()							)
      return m.hexdigest()[:10]

def _a			(		_snake_case							):







      """simple docstring"""






      UpperCAmelCase          =	np.array(_snake_case							)
      UpperCAmelCase          =	npimg.shape
      return {"hash": hashimage(_snake_case							), "shape": shape}







@is_pipeline_test
@require_vision
@require_torch
class  lowerCamelCase__ (	unittest.TestCase			):
    SCREAMING_SNAKE_CASE		=      dict(
        (list(MODEL_FOR_MASK_GENERATION_MAPPING.items()			) if MODEL_FOR_MASK_GENERATION_MAPPING else [])			)
    SCREAMING_SNAKE_CASE		=      dict(
        (list(TF_MODEL_FOR_MASK_GENERATION_MAPPING.items()			) if TF_MODEL_FOR_MASK_GENERATION_MAPPING else [])			)






    def     _UpperCamelCase							(			self     ,A     ,A     ,A ):
          UpperCAmelCase          =	MaskGenerationPipeline(model=A     ,image_processor=A )
          return image_segmenter, [
              "./tests/fixtures/tests_samples/COCO/000000039769.png",
              "./tests/fixtures/tests_samples/COCO/000000039769.png",
          ]






    def     _UpperCamelCase							(			self     ,A     ,A ):
          pass






    @require_tf
    @unittest.skip("""Image segmentation not implemented in TF""" )
    def     _UpperCamelCase							(			self ):
          pass






    @slow
    @require_torch
    def     _UpperCamelCase							(			self ):
          UpperCAmelCase          =	pipeline("""mask-generation"""     ,model="""facebook/sam-vit-huge""" )

          UpperCAmelCase          =	image_segmenter("""http://images.cocodataset.org/val2017/000000039769.jpg"""     ,points_per_batch=256 )

          # Shortening by hashing
          UpperCAmelCase          =	[]
          for i, o in enumerate(outputs["""masks"""] ):
                new_outupt += [{"mask": mask_to_test_readable(A ), "scores": outputs["scores"][i]}]

          # fmt: off
          self.assertEqual(
              nested_simplify(A     ,decimals=4 )     ,[
                  {"""mask""": {"""hash""": """115ad19f5f""", """shape""": (480, 640)}, """scores""": 1.0444},
                  {"""mask""": {"""hash""": """6affa964c6""", """shape""": (480, 640)}, """scores""": 1.021},
                  {"""mask""": {"""hash""": """dfe28a0388""", """shape""": (480, 640)}, """scores""": 1.0167},
                  {"""mask""": {"""hash""": """c0a5f4a318""", """shape""": (480, 640)}, """scores""": 1.0132},
                  {"""mask""": {"""hash""": """fe8065c197""", """shape""": (480, 640)}, """scores""": 1.0053},
                  {"""mask""": {"""hash""": """e2d0b7a0b7""", """shape""": (480, 640)}, """scores""": 0.9967},
                  {"""mask""": {"""hash""": """453c7844bd""", """shape""": (480, 640)}, """scores""": 0.993},
                  {"""mask""": {"""hash""": """3d44f2926d""", """shape""": (480, 640)}, """scores""": 0.9909},
                  {"""mask""": {"""hash""": """64033ddc3f""", """shape""": (480, 640)}, """scores""": 0.9879},
                  {"""mask""": {"""hash""": """801064ff79""", """shape""": (480, 640)}, """scores""": 0.9834},
                  {"""mask""": {"""hash""": """6172f276ef""", """shape""": (480, 640)}, """scores""": 0.9716},
                  {"""mask""": {"""hash""": """b49e60e084""", """shape""": (480, 640)}, """scores""": 0.9612},
                  {"""mask""": {"""hash""": """a811e775fd""", """shape""": (480, 640)}, """scores""": 0.9599},
                  {"""mask""": {"""hash""": """a6a8ebcf4b""", """shape""": (480, 640)}, """scores""": 0.9552},
                  {"""mask""": {"""hash""": """9d8257e080""", """shape""": (480, 640)}, """scores""": 0.9532},
                  {"""mask""": {"""hash""": """32de6454a8""", """shape""": (480, 640)}, """scores""": 0.9516},
                  {"""mask""": {"""hash""": """af3d4af2c8""", """shape""": (480, 640)}, """scores""": 0.9499},
                  {"""mask""": {"""hash""": """3c6db475fb""", """shape""": (480, 640)}, """scores""": 0.9483},
                  {"""mask""": {"""hash""": """c290813fb9""", """shape""": (480, 640)}, """scores""": 0.9464},
                  {"""mask""": {"""hash""": """b6f0b8f606""", """shape""": (480, 640)}, """scores""": 0.943},
                  {"""mask""": {"""hash""": """92ce16bfdf""", """shape""": (480, 640)}, """scores""": 0.943},
                  {"""mask""": {"""hash""": """c749b25868""", """shape""": (480, 640)}, """scores""": 0.9408},
                  {"""mask""": {"""hash""": """efb6cab859""", """shape""": (480, 640)}, """scores""": 0.9335},
                  {"""mask""": {"""hash""": """1ff2eafb30""", """shape""": (480, 640)}, """scores""": 0.9326},
                  {"""mask""": {"""hash""": """788b798e24""", """shape""": (480, 640)}, """scores""": 0.9262},
                  {"""mask""": {"""hash""": """abea804f0e""", """shape""": (480, 640)}, """scores""": 0.8999},
                  {"""mask""": {"""hash""": """7b9e8ddb73""", """shape""": (480, 640)}, """scores""": 0.8986},
                  {"""mask""": {"""hash""": """cd24047c8a""", """shape""": (480, 640)}, """scores""": 0.8984},
                  {"""mask""": {"""hash""": """6943e6bcbd""", """shape""": (480, 640)}, """scores""": 0.8873},
                  {"""mask""": {"""hash""": """b5f47c9191""", """shape""": (480, 640)}, """scores""": 0.8871}
              ]     ,)
          # fmt: on






    @require_torch
    @slow
    def     _UpperCamelCase							(			self ):
          UpperCAmelCase          =	"""facebook/sam-vit-huge"""
          UpperCAmelCase          =	pipeline("""mask-generation"""     ,model=A )

          UpperCAmelCase          =	image_segmenter(
              """http://images.cocodataset.org/val2017/000000039769.jpg"""     ,pred_iou_thresh=1     ,points_per_batch=256 )

          # Shortening by hashing
          UpperCAmelCase          =	[]
          for i, o in enumerate(outputs["""masks"""] ):
                new_outupt += [{"mask": mask_to_test_readable(A ), "scores": outputs["scores"][i]}]

          self.assertEqual(
              nested_simplify(A     ,decimals=4 )     ,[
                  {"""mask""": {"""hash""": """115ad19f5f""", """shape""": (480, 640)}, """scores""": 1.0444},
                  {"""mask""": {"""hash""": """6affa964c6""", """shape""": (480, 640)}, """scores""": 1.0210},
                  {"""mask""": {"""hash""": """dfe28a0388""", """shape""": (480, 640)}, """scores""": 1.0167},
                  {"""mask""": {"""hash""": """c0a5f4a318""", """shape""": (480, 640)}, """scores""": 1.0132},
                  {"""mask""": {"""hash""": """fe8065c197""", """shape""": (480, 640)}, """scores""": 1.0053},
              ]     ,)



