

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







 """simple docstring"""







 def  _lowercase			(      self							:      Optional[int]     ):
    __lowercase								=						"| <pad> <unk> <s> </s> a b c d e f g h i j k".split()
    __lowercase								=						dict(zip(UpperCAmelCase__, range(len(UpperCAmelCase__     )     )     )     )

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

    __lowercase								=						tempfile.mkdtemp()
    __lowercase								=						os.path.join(self.tmpdirname, VOCAB_FILES_NAMES["vocab_file"]     )
    __lowercase								=						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
    __lowercase								=						"hf-internal-testing/ngram-beam-search-decoder"


 def  _lowercase			(      self							:      Optional[int], **UpperCAmelCase__							:      List[str]     ):
    __lowercase								=						self.add_kwargs_tokens_map.copy()
    kwargs.update(UpperCAmelCase__     )
    return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname, **UpperCAmelCase__     )


 def  _lowercase			(      self							:      Optional[int], **UpperCAmelCase__							:      int     ):
    return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname, **UpperCAmelCase__     )


 def  _lowercase			(      self							:      List[str], **UpperCAmelCase__							:      Dict     ):
    return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name, **UpperCAmelCase__     )


 def  _lowercase			(      self							:      str     ):
    shutil.rmtree(self.tmpdirname     )


 def  _lowercase			(      self							:      List[str]     ):
    __lowercase								=						self.get_tokenizer()
    __lowercase								=						self.get_feature_extractor()
    __lowercase								=						self.get_decoder()

    __lowercase								=						WavaVecaProcessorWithLM(tokenizer=UpperCAmelCase__, feature_extractor=UpperCAmelCase__, decoder=UpperCAmelCase__     )

    processor.save_pretrained(self.tmpdirname     )
    __lowercase								=						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  _lowercase			(      self							:      List[str]     ):
    __lowercase								=						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								=						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  _lowercase			(      self							:      List[Any]     ):
    __lowercase								=						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  _lowercase			(      self							:      Tuple     ):
    __lowercase								=						self.get_feature_extractor()
    __lowercase								=						self.get_tokenizer()
    __lowercase								=						self.get_decoder()

    __lowercase								=						WavaVecaProcessorWithLM(tokenizer=UpperCAmelCase__, feature_extractor=UpperCAmelCase__, decoder=UpperCAmelCase__     )

    __lowercase								=						floats_list((3, 1_0_0_0)     )

    __lowercase								=						feature_extractor(UpperCAmelCase__, return_tensors="np"     )
    __lowercase								=						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  _lowercase			(      self							:      List[Any]     ):
    __lowercase								=						self.get_feature_extractor()
    __lowercase								=						self.get_tokenizer()
    __lowercase								=						self.get_decoder()

    __lowercase								=						WavaVecaProcessorWithLM(tokenizer=UpperCAmelCase__, feature_extractor=UpperCAmelCase__, decoder=UpperCAmelCase__     )

    __lowercase								=						"This is a test string"

    __lowercase								=						processor(text=UpperCAmelCase__     )

    __lowercase								=						tokenizer(UpperCAmelCase__     )

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


 def  _lowercase			(      self							:      Dict, UpperCAmelCase__							:      Any=(2, 1_0, 1_6), UpperCAmelCase__							:      Optional[int]=7_7     ):
    np.random.seed(UpperCAmelCase__     )
    return np.random.rand(*UpperCAmelCase__     )


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

    __lowercase								=						WavaVecaProcessorWithLM(tokenizer=UpperCAmelCase__, feature_extractor=UpperCAmelCase__, decoder=UpperCAmelCase__     )

    __lowercase								=						self._get_dummy_logits(shape=(1_0, 1_6), seed=1_3     )

    __lowercase								=						processor.decode(UpperCAmelCase__     )

    __lowercase								=						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  _lowercase			(      self							:      int, UpperCAmelCase__							:      List[str]     ):
    __lowercase								=						self.get_feature_extractor()
    __lowercase								=						self.get_tokenizer()
    __lowercase								=						self.get_decoder()

    __lowercase								=						WavaVecaProcessorWithLM(tokenizer=UpperCAmelCase__, feature_extractor=UpperCAmelCase__, decoder=UpperCAmelCase__     )

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

    __lowercase								=						list(UpperCAmelCase__     )

    with get_context("fork"     ).Pool() as p:
       __lowercase								=						decoder.decode_beams_batch(UpperCAmelCase__, UpperCAmelCase__     )

    __lowercase		,__lowercase		,__lowercase								=						[], [], []
    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  _lowercase			(      self							:      Any     ):
    __lowercase								=						self.get_feature_extractor()
    __lowercase								=						self.get_tokenizer()
    __lowercase								=						self.get_decoder()

    __lowercase								=						WavaVecaProcessorWithLM(tokenizer=UpperCAmelCase__, feature_extractor=UpperCAmelCase__, decoder=UpperCAmelCase__     )

    __lowercase								=						self._get_dummy_logits()

    __lowercase								=						1_5
    __lowercase								=						-20.0
    __lowercase								=						-4.0

    __lowercase								=						processor.batch_decode(
        UpperCAmelCase__, beam_width=UpperCAmelCase__, beam_prune_logp=UpperCAmelCase__, token_min_logp=UpperCAmelCase__, )
    __lowercase								=						decoded_processor_out.text

    __lowercase								=						list(UpperCAmelCase__     )

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

    __lowercase								=						[d[0][0] for d in decoded_decoder_out]
    __lowercase								=						[d[0][2] for d in decoded_decoder_out]
    __lowercase								=						[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.9_474], UpperCAmelCase__, atol=1E-3     )     )


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

    __lowercase								=						WavaVecaProcessorWithLM(tokenizer=UpperCAmelCase__, feature_extractor=UpperCAmelCase__, decoder=UpperCAmelCase__     )

    __lowercase								=						self._get_dummy_logits()

    __lowercase								=						2.0
    __lowercase								=						5.0
    __lowercase								=						-20.0
    __lowercase								=						True

    __lowercase								=						processor.batch_decode(
        UpperCAmelCase__, alpha=UpperCAmelCase__, beta=UpperCAmelCase__, unk_score_offset=UpperCAmelCase__, lm_score_boundary=UpperCAmelCase__, )
    __lowercase								=						decoded_processor_out.text

    __lowercase								=						list(UpperCAmelCase__     )
    decoder.reset_params(
        alpha=UpperCAmelCase__, beta=UpperCAmelCase__, unk_score_offset=UpperCAmelCase__, lm_score_boundary=UpperCAmelCase__, )

    with get_context("fork"     ).Pool() as pool:
       __lowercase								=						decoder.decode_beams_batch(
           UpperCAmelCase__, UpperCAmelCase__, )

    __lowercase								=						[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__     )
    __lowercase								=						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  _lowercase			(      self							:      str     ):
    __lowercase								=						WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"     )

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

    __lowercase								=						os.listdir(UpperCAmelCase__     )
    __lowercase								=						["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  _lowercase			(      self							:      Union[str, Any]     ):
    __lowercase								=						snapshot_download("hf-internal-testing/processor_with_lm"     )

    __lowercase								=						WavaVecaProcessorWithLM.from_pretrained(UpperCAmelCase__     )

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

    __lowercase								=						os.listdir(UpperCAmelCase__     )
    __lowercase								=						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  _lowercase			(      self							:      Optional[Any]     ):
    __lowercase								=						WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"     )
    __lowercase								=						AutoProcessor.from_pretrained("hf-internal-testing/processor_with_lm"     )

    __lowercase								=						floats_list((3, 1_0_0_0)     )

    __lowercase								=						processor_wavaveca(UpperCAmelCase__, return_tensors="np"     )
    __lowercase								=						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     )

    __lowercase								=						self._get_dummy_logits()

    __lowercase								=						processor_wavaveca.batch_decode(UpperCAmelCase__     )
    __lowercase								=						processor_auto.batch_decode(UpperCAmelCase__     )

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


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

    __lowercase								=						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  _lowercase			(      UpperCAmelCase__							:      str, UpperCAmelCase__							:      Union[str, Any]     ):
    __lowercase								=						[d[key] for d in offsets]
    return retrieved_list


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

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

    __lowercase								=						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  _lowercase			(      self							:      int     ):
    import torch

    __lowercase								=						load_dataset("common_voice", "en", split="train", streaming=UpperCAmelCase__     )
    __lowercase								=						ds.cast_column("audio", datasets.Audio(sampling_rate=1_6_0_0_0     )     )
    __lowercase								=						iter(UpperCAmelCase__     )
    __lowercase								=						next(UpperCAmelCase__     )

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

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

    __lowercase								=						processor.decode(logits[0], output_word_offsets=UpperCAmelCase__     )

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

    __lowercase								=						"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
    __lowercase								=						torch.tensor(self.get_from_offsets(UpperCAmelCase__, "start_time"     )     )
    __lowercase								=						torch.tensor(self.get_from_offsets(UpperCAmelCase__, "end_time"     )     )

    # fmt: off
    __lowercase								=						torch.tensor([1.4_199, 1.6_599, 2.2_599, 3.0, 3.24, 3.5_999, 3.7_999, 4.0_999, 4.26, 4.94, 5.28, 5.6_599, 5.78, 5.94, 6.32, 6.5_399, 6.6_599]     )
    __lowercase								=						torch.tensor([1.5_399, 1.8_999, 2.9, 3.16, 3.5_399, 3.72, 4.0_199, 4.1_799, 4.76, 5.1_599, 5.5_599, 5.6_999, 5.86, 6.1_999, 6.38, 6.6_199, 6.94]     )
    # fmt: on

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







"""simple docstring"""





from ...utils import is_note_seq_available, is_transformers_available, is_torch_available
from ...utils import OptionalDependencyNotAvailable


try:
   if not (is_transformers_available() and is_torch_available()):
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
   from .notes_encoder import SpectrogramNotesEncoder
   from .continous_encoder import SpectrogramContEncoder
   from .pipeline_spectrogram_diffusion import (
       SpectrogramContEncoder,
       SpectrogramDiffusionPipeline,
       TaFilmDecoder,
   )

try:
   if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   from ...utils.dummy_transformers_and_torch_and_note_seq_objects import *  # noqa F403
else:
   from .midi_utils import MidiProcessor




