
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					):



      """simple docstring"""







      def 						snake_case_	(					self						:				str ):
          __lowercase						:	Union[str, Any]					=				'''| <pad> <unk> <s> </s> a b c d e f g h i j k'''.split()
          __lowercase						:	Any					=				dict(zip(lowerCamelCase__ ,				range(len(lowerCamelCase__ ) ) ) )

          __lowercase						:	str					=				{
              '''unk_token''': '''<unk>''',
              '''bos_token''': '''<s>''',
              '''eos_token''': '''</s>''',
          }
          __lowercase						:	Tuple					=				{
              '''feature_size''': 1,
              '''padding_value''': 0.0,
              '''sampling_rate''': 1_6000,
              '''return_attention_mask''': False,
              '''do_normalize''': True,
          }

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

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

          # load decoder from hub
          __lowercase						:	List[Any]					=				'''hf-internal-testing/ngram-beam-search-decoder'''




      def 						snake_case_	(					self						:				int ,				**_snake_case						:				Tuple ):
          __lowercase						:	Union[str, Any]					=				self.add_kwargs_tokens_map.copy()
          kwargs.update(lowerCamelCase__ )
          return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname ,				**lowerCamelCase__ )




      def 						snake_case_	(					self						:				Dict ,				**_snake_case						:				int ):
          return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname ,				**lowerCamelCase__ )




      def 						snake_case_	(					self						:				List[str] ,				**_snake_case						:				List[Any] ):
          return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name ,				**lowerCamelCase__ )




      def 						snake_case_	(					self						:				Optional[Any] ):
          shutil.rmtree(self.tmpdirname )




      def 						snake_case_	(					self						:				Any ):
          __lowercase						:	Tuple					=				self.get_tokenizer()
          __lowercase						:	List[Any]					=				self.get_feature_extractor()
          __lowercase						:	Any					=				self.get_decoder()

          __lowercase						:	Optional[Any]					=				WavaVecaProcessorWithLM(tokenizer=lowerCamelCase__ ,				feature_extractor=lowerCamelCase__ ,				decoder=lowerCamelCase__ )

          processor.save_pretrained(self.tmpdirname )
          __lowercase						:	str					=				WavaVecaProcessorWithLM.from_pretrained(self.tmpdirname )

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

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

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




      def 						snake_case_	(					self						:				Any ):
          __lowercase						:	Tuple					=				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
          __lowercase						:	Any					=				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 						snake_case_	(					self						:				Any ):
          __lowercase						:	Optional[int]					=				self.get_tokenizer()
          # add token to trigger raise
          tokenizer.add_tokens(['''xx'''] )
          with self.assertRaisesRegex(lowerCamelCase__ ,				'''include''' ):
              WavaVecaProcessorWithLM(
                  tokenizer=lowerCamelCase__ ,				feature_extractor=self.get_feature_extractor() ,				decoder=self.get_decoder() )




      def 						snake_case_	(					self						:				Optional[Any] ):
          __lowercase						:	Optional[Any]					=				self.get_feature_extractor()
          __lowercase						:	int					=				self.get_tokenizer()
          __lowercase						:	int					=				self.get_decoder()

          __lowercase						:	Optional[Any]					=				WavaVecaProcessorWithLM(tokenizer=lowerCamelCase__ ,				feature_extractor=lowerCamelCase__ ,				decoder=lowerCamelCase__ )

          __lowercase						:	Tuple					=				floats_list((3, 1000) )

          __lowercase						:	Optional[Any]					=				feature_extractor(lowerCamelCase__ ,				return_tensors='''np''' )
          __lowercase						:	str					=				processor(lowerCamelCase__ ,				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 						snake_case_	(					self						:				Any ):
          __lowercase						:	Union[str, Any]					=				self.get_feature_extractor()
          __lowercase						:	Union[str, Any]					=				self.get_tokenizer()
          __lowercase						:	Optional[int]					=				self.get_decoder()

          __lowercase						:	Dict					=				WavaVecaProcessorWithLM(tokenizer=lowerCamelCase__ ,				feature_extractor=lowerCamelCase__ ,				decoder=lowerCamelCase__ )

          __lowercase						:	Optional[int]					=				'''This is a test string'''

          __lowercase						:	Optional[int]					=				processor(text=lowerCamelCase__ )

          __lowercase						:	Optional[int]					=				tokenizer(lowerCamelCase__ )

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




      def 						snake_case_	(					self						:				str ,				_snake_case						:				Tuple=(2, 10, 16) ,				_snake_case						:				List[Any]=77 ):
          np.random.seed(lowerCamelCase__ )
          return np.random.rand(*lowerCamelCase__ )




      def 						snake_case_	(					self						:				Any ):
          __lowercase						:	Optional[Any]					=				self.get_feature_extractor()
          __lowercase						:	int					=				self.get_tokenizer()
          __lowercase						:	Optional[int]					=				self.get_decoder()

          __lowercase						:	Tuple					=				WavaVecaProcessorWithLM(tokenizer=lowerCamelCase__ ,				feature_extractor=lowerCamelCase__ ,				decoder=lowerCamelCase__ )

          __lowercase						:	Optional[Any]					=				self._get_dummy_logits(shape=(10, 16) ,				seed=13 )

          __lowercase						:	Optional[int]					=				processor.decode(lowerCamelCase__ )

          __lowercase						:	Dict					=				decoder.decode_beams(lowerCamelCase__ )[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 						snake_case_	(					self						:				Optional[int] ,				_snake_case						:				Tuple ):
          __lowercase						:	Any					=				self.get_feature_extractor()
          __lowercase						:	Optional[int]					=				self.get_tokenizer()
          __lowercase						:	Union[str, Any]					=				self.get_decoder()

          __lowercase						:	Optional[int]					=				WavaVecaProcessorWithLM(tokenizer=lowerCamelCase__ ,				feature_extractor=lowerCamelCase__ ,				decoder=lowerCamelCase__ )

          __lowercase						:	Any					=				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:
              __lowercase						:	Dict					=				processor.batch_decode(lowerCamelCase__ )
          else:
              with get_context(lowerCamelCase__ ).Pool() as pool:
                  __lowercase						:	Union[str, Any]					=				processor.batch_decode(lowerCamelCase__ ,				lowerCamelCase__ )

          __lowercase						:	Any					=				list(lowerCamelCase__ )

          with get_context('''fork''' ).Pool() as p:
              __lowercase						:	str					=				decoder.decode_beams_batch(lowerCamelCase__ ,				lowerCamelCase__ )

          __lowercase			,					__lowercase			,					__lowercase						:	Dict					=				[], [], []
          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(lowerCamelCase__ ,				decoded_processor.text )
          self.assertListEqual(['''<s> <s> </s>''', '''<s> <s> <s>'''] ,				decoded_processor.text )
          self.assertListEqual(lowerCamelCase__ ,				decoded_processor.logit_score )
          self.assertListEqual(lowerCamelCase__ ,				decoded_processor.lm_score )




      def 						snake_case_	(					self						:				Tuple ):
          __lowercase						:	str					=				self.get_feature_extractor()
          __lowercase						:	int					=				self.get_tokenizer()
          __lowercase						:	List[Any]					=				self.get_decoder()

          __lowercase						:	List[str]					=				WavaVecaProcessorWithLM(tokenizer=lowerCamelCase__ ,				feature_extractor=lowerCamelCase__ ,				decoder=lowerCamelCase__ )

          __lowercase						:	Optional[Any]					=				self._get_dummy_logits()

          __lowercase						:	Any					=				15
          __lowercase						:	Tuple					=				-20.0
          __lowercase						:	Optional[Any]					=				-4.0

          __lowercase						:	List[Any]					=				processor.batch_decode(
              lowerCamelCase__ ,				beam_width=lowerCamelCase__ ,				beam_prune_logp=lowerCamelCase__ ,				token_min_logp=lowerCamelCase__ ,				)
          __lowercase						:	List[str]					=				decoded_processor_out.text

          __lowercase						:	Union[str, Any]					=				list(lowerCamelCase__ )

          with get_context('''fork''' ).Pool() as pool:
              __lowercase						:	int					=				decoder.decode_beams_batch(
                  lowerCamelCase__ ,				lowerCamelCase__ ,				beam_width=lowerCamelCase__ ,				beam_prune_logp=lowerCamelCase__ ,				token_min_logp=lowerCamelCase__ ,				)

          __lowercase						:	Optional[int]					=				[d[0][0] for d in decoded_decoder_out]
          __lowercase						:	Any					=				[d[0][2] for d in decoded_decoder_out]
          __lowercase						:	Optional[int]					=				[d[0][3] for d in decoded_decoder_out]

          self.assertListEqual(lowerCamelCase__ ,				lowerCamelCase__ )
          self.assertListEqual(['''</s> <s> <s>''', '''<s> <s> <s>'''] ,				lowerCamelCase__ )

          self.assertTrue(np.array_equal(lowerCamelCase__ ,				decoded_processor_out.logit_score ) )
          self.assertTrue(np.allclose([-20.0_54, -18.4_47] ,				lowerCamelCase__ ,				atol=1E-3 ) )

          self.assertTrue(np.array_equal(lowerCamelCase__ ,				decoded_processor_out.lm_score ) )
          self.assertTrue(np.allclose([-15.5_54, -13.94_74] ,				lowerCamelCase__ ,				atol=1E-3 ) )




      def 						snake_case_	(					self						:				Union[str, Any] ):
          __lowercase						:	List[Any]					=				self.get_feature_extractor()
          __lowercase						:	Any					=				self.get_tokenizer()
          __lowercase						:	Union[str, Any]					=				self.get_decoder()

          __lowercase						:	Dict					=				WavaVecaProcessorWithLM(tokenizer=lowerCamelCase__ ,				feature_extractor=lowerCamelCase__ ,				decoder=lowerCamelCase__ )

          __lowercase						:	Any					=				self._get_dummy_logits()

          __lowercase						:	Union[str, Any]					=				2.0
          __lowercase						:	Union[str, Any]					=				5.0
          __lowercase						:	int					=				-20.0
          __lowercase						:	Optional[int]					=				True

          __lowercase						:	List[str]					=				processor.batch_decode(
              lowerCamelCase__ ,				alpha=lowerCamelCase__ ,				beta=lowerCamelCase__ ,				unk_score_offset=lowerCamelCase__ ,				lm_score_boundary=lowerCamelCase__ ,				)
          __lowercase						:	Optional[int]					=				decoded_processor_out.text

          __lowercase						:	int					=				list(lowerCamelCase__ )
          decoder.reset_params(
              alpha=lowerCamelCase__ ,				beta=lowerCamelCase__ ,				unk_score_offset=lowerCamelCase__ ,				lm_score_boundary=lowerCamelCase__ ,				)

          with get_context('''fork''' ).Pool() as pool:
              __lowercase						:	Tuple					=				decoder.decode_beams_batch(
                  lowerCamelCase__ ,				lowerCamelCase__ ,				)

          __lowercase						:	str					=				[d[0][0] for d in decoded_decoder_out]

          self.assertListEqual(lowerCamelCase__ ,				lowerCamelCase__ )
          self.assertListEqual(['''<s> </s> <s> </s> </s>''', '''</s> </s> <s> </s> </s>'''] ,				lowerCamelCase__ )
          __lowercase						:	Dict					=				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 ,				lowerCamelCase__ )




      def 						snake_case_	(					self						:				str ):
          __lowercase						:	Tuple					=				WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm''' )

          __lowercase						:	str					=				processor.decoder.model_container[processor.decoder._model_key]
          __lowercase						:	List[Any]					=				Path(language_model._kenlm_model.path.decode('''utf-8''' ) ).parent.parent.absolute()

          __lowercase						:	int					=				os.listdir(lowerCamelCase__ )
          __lowercase						:	List[str]					=				['''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(lowerCamelCase__ ,				lowerCamelCase__ )




      def 						snake_case_	(					self						:				Optional[Any] ):
          __lowercase						:	Any					=				snapshot_download('''hf-internal-testing/processor_with_lm''' )

          __lowercase						:	List[str]					=				WavaVecaProcessorWithLM.from_pretrained(lowerCamelCase__ )

          __lowercase						:	List[str]					=				processor.decoder.model_container[processor.decoder._model_key]
          __lowercase						:	Optional[Any]					=				Path(language_model._kenlm_model.path.decode('''utf-8''' ) ).parent.parent.absolute()

          __lowercase						:	Any					=				os.listdir(lowerCamelCase__ )
          __lowercase						:	List[str]					=				os.listdir(lowerCamelCase__ )

          local_decoder_files.sort()
          expected_decoder_files.sort()

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




      def 						snake_case_	(					self						:				int ):
          __lowercase						:	List[str]					=				WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm''' )
          __lowercase						:	Tuple					=				AutoProcessor.from_pretrained('''hf-internal-testing/processor_with_lm''' )

          __lowercase						:	Union[str, Any]					=				floats_list((3, 1000) )

          __lowercase						:	int					=				processor_wavaveca(lowerCamelCase__ ,				return_tensors='''np''' )
          __lowercase						:	Optional[Any]					=				processor_auto(lowerCamelCase__ ,				return_tensors='''np''' )

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

          __lowercase						:	Union[str, Any]					=				self._get_dummy_logits()

          __lowercase						:	Optional[Any]					=				processor_wavaveca.batch_decode(lowerCamelCase__ )
          __lowercase						:	Optional[Any]					=				processor_auto.batch_decode(lowerCamelCase__ )

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




      def 						snake_case_	(					self						:				List[str] ):
          __lowercase						:	Optional[int]					=				self.get_feature_extractor()
          __lowercase						:	int					=				self.get_tokenizer()
          __lowercase						:	str					=				self.get_decoder()

          __lowercase						:	Optional[Any]					=				WavaVecaProcessorWithLM(tokenizer=lowerCamelCase__ ,				feature_extractor=lowerCamelCase__ ,				decoder=lowerCamelCase__ )

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




      @staticmethod
      def 						snake_case_	(					_snake_case						:				Tuple ,				_snake_case						:				Any ):
          __lowercase						:	Union[str, Any]					=				[d[key] for d in offsets]
          return retrieved_list




      def 						snake_case_	(					self						:				Union[str, Any] ):
          __lowercase						:	Optional[Any]					=				WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm''' )
          __lowercase						:	str					=				self._get_dummy_logits()[0]

          __lowercase						:	int					=				processor.decode(lowerCamelCase__ ,				output_word_offsets=lowerCamelCase__ )
          # 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(lowerCamelCase__ ,				lowerCamelCase__ ) )

          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 						snake_case_	(					self						:				Tuple ):
          __lowercase						:	Union[str, Any]					=				WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm''' )
          __lowercase						:	List[Any]					=				self._get_dummy_logits()

          __lowercase						:	Any					=				processor.batch_decode(lowerCamelCase__ ,				output_word_offsets=lowerCamelCase__ )

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

          self.assertListEqual(
              [''' '''.join(self.get_from_offsets(lowerCamelCase__ ,				'''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 						snake_case_	(					self						:				Tuple ):
          import torch

          __lowercase						:	Dict					=				load_dataset('''common_voice''' ,				'''en''' ,				split='''train''' ,				streaming=lowerCamelCase__ )
          __lowercase						:	int					=				ds.cast_column('''audio''' ,				datasets.Audio(sampling_rate=1_6000 ) )
          __lowercase						:	List[str]					=				iter(lowerCamelCase__ )
          __lowercase						:	int					=				next(lowerCamelCase__ )

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

          with torch.no_grad():
              __lowercase						:	int					=				model(lowerCamelCase__ ).logits.cpu().numpy()

          __lowercase						:	Any					=				processor.decode(logits[0] ,				output_word_offsets=lowerCamelCase__ )

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

          __lowercase						:	List[str]					=				'''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(lowerCamelCase__ ,				'''word''' ) ) ,				lowerCamelCase__ )
          self.assertEqual(''' '''.join(self.get_from_offsets(lowerCamelCase__ ,				'''word''' ) ) ,				output.text )

          # output times
          __lowercase						:	List[str]					=				torch.tensor(self.get_from_offsets(lowerCamelCase__ ,				'''start_time''' ) )
          __lowercase						:	List[str]					=				torch.tensor(self.get_from_offsets(lowerCamelCase__ ,				'''end_time''' ) )

          # fmt: off
          __lowercase						:	Optional[Any]					=				torch.tensor([1.41_99, 1.65_99, 2.25_99, 3.0, 3.24, 3.59_99, 3.79_99, 4.09_99, 4.26, 4.94, 5.28, 5.65_99, 5.78, 5.94, 6.32, 6.53_99, 6.65_99] )
          __lowercase						:	str					=				torch.tensor([1.53_99, 1.89_99, 2.9, 3.16, 3.53_99, 3.72, 4.01_99, 4.17_99, 4.76, 5.15_99, 5.55_99, 5.69_99, 5.86, 6.19_99, 6.38, 6.61_99, 6.94] )
          # fmt: on

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









import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    TextToVideoSDPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import is_xformers_available, load_numpy, skip_mps, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()

@skip_mps
class 		UpperCamelCase__ (      lowerCAmelCase_							,      unittest.TestCase   ):



     '''simple docstring'''

     __snake_case    :	List[str]			=			TextToVideoSDPipeline
     __snake_case    :	int			=			TEXT_TO_IMAGE_PARAMS
     __snake_case    :	Tuple			=			TEXT_TO_IMAGE_BATCH_PARAMS
     # No `output_type`.
     __snake_case    :	Dict			=			frozenset(
         [
             "num_inference_steps",
             "generator",
             "latents",
             "return_dict",
             "callback",
             "callback_steps",
         ]   )






     def 				SCREAMING_SNAKE_CASE__  (							self						:    Optional[int]       )     ->		Optional[int]:




           '''simple docstring'''


           torch.manual_seed(0       )
           SCREAMING_SNAKE_CASE		=						UNetaDConditionModel(
               block_out_channels=(32, 64, 64, 64)						,layers_per_block=2						,sample_size=32						,in_channels=4						,out_channels=4						,down_block_types=("""CrossAttnDownBlock3D""", """CrossAttnDownBlock3D""", """CrossAttnDownBlock3D""", """DownBlock3D""")						,up_block_types=("""UpBlock3D""", """CrossAttnUpBlock3D""", """CrossAttnUpBlock3D""", """CrossAttnUpBlock3D""")						,cross_attention_dim=32						,attention_head_dim=4						,)
           SCREAMING_SNAKE_CASE		=						DDIMScheduler(
               beta_start=0.00085						,beta_end=0.012						,beta_schedule="""scaled_linear"""						,clip_sample=lowerCamelCase__						,set_alpha_to_one=lowerCamelCase__						,)
           torch.manual_seed(0       )
           SCREAMING_SNAKE_CASE		=						AutoencoderKL(
               block_out_channels=[32, 64]						,in_channels=3						,out_channels=3						,down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]						,up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]						,latent_channels=4						,sample_size=128						,)
           torch.manual_seed(0       )
           SCREAMING_SNAKE_CASE		=						CLIPTextConfig(
               bos_token_id=0						,eos_token_id=2						,hidden_size=32						,intermediate_size=37						,layer_norm_eps=1e-0_5						,num_attention_heads=4						,num_hidden_layers=5						,pad_token_id=1						,vocab_size=1000						,hidden_act="""gelu"""						,projection_dim=512						,)
           SCREAMING_SNAKE_CASE		=						CLIPTextModel(lowerCamelCase__       )
           SCREAMING_SNAKE_CASE		=						CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""       )

           SCREAMING_SNAKE_CASE		=						{
               """unet""": unet,
               """scheduler""": scheduler,
               """vae""": vae,
               """text_encoder""": text_encoder,
               """tokenizer""": tokenizer,
           }
           return components






     def 				SCREAMING_SNAKE_CASE__  (							self						:    int						,lowerCamelCase__						:    Optional[Any]						,lowerCamelCase__						:    int=0       )     ->		List[Any]:




           '''simple docstring'''


           if str(lowerCamelCase__       ).startswith("""mps"""       ):
                 SCREAMING_SNAKE_CASE		=						torch.manual_seed(lowerCamelCase__       )
           else:
                 SCREAMING_SNAKE_CASE		=						torch.Generator(device=lowerCamelCase__       ).manual_seed(lowerCamelCase__       )
           SCREAMING_SNAKE_CASE		=						{
               """prompt""": """A painting of a squirrel eating a burger""",
               """generator""": generator,
               """num_inference_steps""": 2,
               """guidance_scale""": 6.0,
               """output_type""": """pt""",
           }
           return inputs






     def 				SCREAMING_SNAKE_CASE__  (							self						:    Optional[Any]       )     ->		List[str]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
           SCREAMING_SNAKE_CASE		=						self.get_dummy_components()
           SCREAMING_SNAKE_CASE		=						TextToVideoSDPipeline(**lowerCamelCase__       )
           SCREAMING_SNAKE_CASE		=						sd_pipe.to(lowerCamelCase__       )
           sd_pipe.set_progress_bar_config(disable=lowerCamelCase__       )

           SCREAMING_SNAKE_CASE		=						self.get_dummy_inputs(lowerCamelCase__       )
           SCREAMING_SNAKE_CASE		=						"""np"""
           SCREAMING_SNAKE_CASE		=						sd_pipe(**lowerCamelCase__       ).frames
           SCREAMING_SNAKE_CASE		=						frames[0][-3:, -3:, -1]

           assert frames[0].shape == (64, 64, 3)
           SCREAMING_SNAKE_CASE		=						np.array([158.0, 160.0, 153.0, 125.0, 100.0, 121.0, 111.0, 93.0, 113.0]       )

           assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-2






     def 				SCREAMING_SNAKE_CASE__  (							self						:    Dict       )     ->		int:




           '''simple docstring'''


           self._test_attention_slicing_forward_pass(test_mean_pixel_difference=lowerCamelCase__						,expected_max_diff=3e-3       )






     @unittest.skipIf(
         torch_device != """cuda""" or not is_xformers_available()						,reason="""XFormers attention is only available with CUDA and `xformers` installed"""						,)
     def 				SCREAMING_SNAKE_CASE__  (							self						:    Tuple       )     ->		Any:




           '''simple docstring'''


           self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=lowerCamelCase__						,expected_max_diff=1e-2       )






     @unittest.skip(reason="""Batching needs to be properly figured out first for this pipeline."""       )
     def 				SCREAMING_SNAKE_CASE__  (							self						:    Optional[Any]       )     ->		List[Any]:




           '''simple docstring'''


           pass






     @unittest.skip(reason="""Batching needs to be properly figured out first for this pipeline."""       )
     def 				SCREAMING_SNAKE_CASE__  (							self						:    Optional[Any]       )     ->		Any:




           '''simple docstring'''


           pass






     @unittest.skip(reason="""`num_images_per_prompt` argument is not supported for this pipeline."""       )
     def 				SCREAMING_SNAKE_CASE__  (							self						:    Optional[int]       )     ->		Tuple:




           '''simple docstring'''


           pass






     def 				SCREAMING_SNAKE_CASE__  (							self						:    List[str]       )     ->		Optional[Any]:




           '''simple docstring'''


           return super().test_progress_bar()

@slow
@skip_mps
class 		UpperCamelCase__ (      unittest.TestCase   ):



     '''simple docstring'''

     def 				SCREAMING_SNAKE_CASE__  (							self						:    Any       )     ->		Optional[int]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						load_numpy(
               """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text_to_video/video.npy"""       )

           SCREAMING_SNAKE_CASE		=						TextToVideoSDPipeline.from_pretrained("""damo-vilab/text-to-video-ms-1.7b"""       )
           SCREAMING_SNAKE_CASE		=						DPMSolverMultistepScheduler.from_config(pipe.scheduler.config       )
           SCREAMING_SNAKE_CASE		=						pipe.to("""cuda"""       )

           SCREAMING_SNAKE_CASE		=						"""Spiderman is surfing"""
           SCREAMING_SNAKE_CASE		=						torch.Generator(device="""cpu"""       ).manual_seed(0       )

           SCREAMING_SNAKE_CASE		=						pipe(lowerCamelCase__						,generator=lowerCamelCase__						,num_inference_steps=25						,output_type="""pt"""       ).frames
           SCREAMING_SNAKE_CASE		=						video_frames.cpu().numpy()

           assert np.abs(expected_video - video       ).mean() < 5e-2






     def 				SCREAMING_SNAKE_CASE__  (							self						:    Any       )     ->		Any:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						load_numpy(
               """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text_to_video/video_2step.npy"""       )

           SCREAMING_SNAKE_CASE		=						TextToVideoSDPipeline.from_pretrained("""damo-vilab/text-to-video-ms-1.7b"""       )
           SCREAMING_SNAKE_CASE		=						pipe.to("""cuda"""       )

           SCREAMING_SNAKE_CASE		=						"""Spiderman is surfing"""
           SCREAMING_SNAKE_CASE		=						torch.Generator(device="""cpu"""       ).manual_seed(0       )

           SCREAMING_SNAKE_CASE		=						pipe(lowerCamelCase__						,generator=lowerCamelCase__						,num_inference_steps=2						,output_type="""pt"""       ).frames
           SCREAMING_SNAKE_CASE		=						video_frames.cpu().numpy()

           assert np.abs(expected_video - video       ).mean() < 5e-2


