
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   :					Tuple     )		->		str:
         snake_case				  =      "| <pad> <unk> <s> </s> a b c d e f g h i j k".split()
         snake_case				  =      dict(zip(_A       ,					range(len(_A     )     )     )     )

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

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

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

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





     def 	UpperCAmelCase(self   :					List[str]       ,					**_A   :					List[str]     )		->		Tuple:
         snake_case				  =      self.add_kwargs_tokens_map.copy()
         kwargs.update(_A     )
         return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname       ,					**_A     )





     def 	UpperCAmelCase(self   :					Optional[Any]       ,					**_A   :					int     )		->		List[Any]:
         return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname       ,					**_A     )





     def 	UpperCAmelCase(self   :					int       ,					**_A   :					str     )		->		Any:
         return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name       ,					**_A     )





     def 	UpperCAmelCase(self   :					Tuple     )		->		Tuple:
         shutil.rmtree(self.tmpdirname     )





     def 	UpperCAmelCase(self   :					List[str]     )		->		Optional[int]:
         snake_case				  =      self.get_tokenizer()
         snake_case				  =      self.get_feature_extractor()
         snake_case				  =      self.get_decoder()

         snake_case				  =      WavaVecaProcessorWithLM(tokenizer=_A       ,					feature_extractor=_A       ,					decoder=_A     )

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

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

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

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





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





     def 	UpperCAmelCase(self   :					Optional[Any]     )		->		Optional[int]:
         snake_case				  =      self.get_feature_extractor()
         snake_case				  =      self.get_tokenizer()
         snake_case				  =      self.get_decoder()

         snake_case				  =      WavaVecaProcessorWithLM(tokenizer=_A       ,					feature_extractor=_A       ,					decoder=_A     )

         snake_case				  =      floats_list((3, 1_0_0_0)     )

         snake_case				  =      feature_extractor(_A       ,					return_tensors="np"     )
         snake_case				  =      processor(_A       ,					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   :					List[Any]     )		->		Optional[int]:
         snake_case				  =      self.get_feature_extractor()
         snake_case				  =      self.get_tokenizer()
         snake_case				  =      self.get_decoder()

         snake_case				  =      WavaVecaProcessorWithLM(tokenizer=_A       ,					feature_extractor=_A       ,					decoder=_A     )

         snake_case				  =      "This is a test string"

         snake_case				  =      processor(text=_A     )

         snake_case				  =      tokenizer(_A     )

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





     def 	UpperCAmelCase(self   :					Optional[Any]       ,					_A   :					Any=(2, 1_0, 1_6)       ,					_A   :					List[Any]=7_7     )		->		Optional[Any]:
         np.random.seed(_A     )
         return np.random.rand(*_A     )





     def 	UpperCAmelCase(self   :					Any     )		->		str:
         snake_case				  =      self.get_feature_extractor()
         snake_case				  =      self.get_tokenizer()
         snake_case				  =      self.get_decoder()

         snake_case				  =      WavaVecaProcessorWithLM(tokenizer=_A       ,					feature_extractor=_A       ,					decoder=_A     )

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

         snake_case				  =      processor.decode(_A     )

         snake_case				  =      decoder.decode_beams(_A     )[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   :					str       ,					_A   :					str     )		->		List[Any]:
         snake_case				  =      self.get_feature_extractor()
         snake_case				  =      self.get_tokenizer()
         snake_case				  =      self.get_decoder()

         snake_case				  =      WavaVecaProcessorWithLM(tokenizer=_A       ,					feature_extractor=_A       ,					decoder=_A     )

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

         snake_case				  =      list(_A     )

         with get_context("fork"     ).Pool() as p:
             snake_case				  =      decoder.decode_beams_batch(_A       ,					_A     )

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





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

         snake_case				  =      WavaVecaProcessorWithLM(tokenizer=_A       ,					feature_extractor=_A       ,					decoder=_A     )

         snake_case				  =      self._get_dummy_logits()

         snake_case				  =      1_5
         snake_case				  =      -20.0
         snake_case				  =      -4.0

         snake_case				  =      processor.batch_decode(
             _A       ,					beam_width=_A       ,					beam_prune_logp=_A       ,					token_min_logp=_A       ,					)
         snake_case				  =      decoded_processor_out.text

         snake_case				  =      list(_A     )

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

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

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

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

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





     def 	UpperCAmelCase(self   :					Optional[int]     )		->		Any:
         snake_case				  =      self.get_feature_extractor()
         snake_case				  =      self.get_tokenizer()
         snake_case				  =      self.get_decoder()

         snake_case				  =      WavaVecaProcessorWithLM(tokenizer=_A       ,					feature_extractor=_A       ,					decoder=_A     )

         snake_case				  =      self._get_dummy_logits()

         snake_case				  =      2.0
         snake_case				  =      5.0
         snake_case				  =      -20.0
         snake_case				  =      True

         snake_case				  =      processor.batch_decode(
             _A       ,					alpha=_A       ,					beta=_A       ,					unk_score_offset=_A       ,					lm_score_boundary=_A       ,					)
         snake_case				  =      decoded_processor_out.text

         snake_case				  =      list(_A     )
         decoder.reset_params(
             alpha=_A       ,					beta=_A       ,					unk_score_offset=_A       ,					lm_score_boundary=_A       ,					)

         with get_context("fork"     ).Pool() as pool:
             snake_case				  =      decoder.decode_beams_batch(
                 _A       ,					_A       ,					)

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

         self.assertListEqual(_A       ,					_A     )
         self.assertListEqual(["<s> </s> <s> </s> </s>", "</s> </s> <s> </s> </s>"]       ,					_A     )
         snake_case				  =      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       ,					_A     )





     def 	UpperCAmelCase(self   :					Optional[Any]     )		->		Any:
         snake_case				  =      WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"     )

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

         snake_case				  =      os.listdir(_A     )
         snake_case				  =      ["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(_A       ,					_A     )





     def 	UpperCAmelCase(self   :					List[Any]     )		->		str:
         snake_case				  =      snapshot_download("hf-internal-testing/processor_with_lm"     )

         snake_case				  =      WavaVecaProcessorWithLM.from_pretrained(_A     )

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

         snake_case				  =      os.listdir(_A     )
         snake_case				  =      os.listdir(_A     )

         local_decoder_files.sort()
         expected_decoder_files.sort()

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





     def 	UpperCAmelCase(self   :					List[str]     )		->		Optional[int]:
         snake_case				  =      WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"     )
         snake_case				  =      AutoProcessor.from_pretrained("hf-internal-testing/processor_with_lm"     )

         snake_case				  =      floats_list((3, 1_0_0_0)     )

         snake_case				  =      processor_wavaveca(_A       ,					return_tensors="np"     )
         snake_case				  =      processor_auto(_A       ,					return_tensors="np"     )

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

         snake_case				  =      self._get_dummy_logits()

         snake_case				  =      processor_wavaveca.batch_decode(_A     )
         snake_case				  =      processor_auto.batch_decode(_A     )

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





     def 	UpperCAmelCase(self   :					Tuple     )		->		List[Any]:
         snake_case				  =      self.get_feature_extractor()
         snake_case				  =      self.get_tokenizer()
         snake_case				  =      self.get_decoder()

         snake_case				  =      WavaVecaProcessorWithLM(tokenizer=_A       ,					feature_extractor=_A       ,					decoder=_A     )

         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   :					List[Any]       ,					_A   :					Union[str, Any]     )		->		Dict:
         snake_case				  =      [d[key] for d in offsets]
         return retrieved_list





     def 	UpperCAmelCase(self   :					Optional[int]     )		->		List[Any]:
         snake_case				  =      WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"     )
         snake_case				  =      self._get_dummy_logits()[0]

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

         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   :					Tuple     )		->		Optional[int]:
         snake_case				  =      WavaVecaProcessorWithLM.from_pretrained("hf-internal-testing/processor_with_lm"     )
         snake_case				  =      self._get_dummy_logits()

         snake_case				  =      processor.batch_decode(_A       ,					output_word_offsets=_A     )

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

         self.assertListEqual(
             [" ".join(self.get_from_offsets(_A       ,					"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   :					Tuple     )		->		Optional[Any]:
         import torch

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

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

         with torch.no_grad():
             snake_case				  =      model(_A     ).logits.cpu().numpy()

         snake_case				  =      processor.decode(logits[0]       ,					output_word_offsets=_A     )

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

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

         # output times
         snake_case				  =      torch.tensor(self.get_from_offsets(_A       ,					"start_time"     )     )
         snake_case				  =      torch.tensor(self.get_from_offsets(_A       ,					"end_time"     )     )

         # fmt: off
         snake_case				  =      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]     )
         snake_case				  =      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(_A       ,					_A       ,					atol=0.01     )     )
         self.assertTrue(torch.allclose(_A       ,					_A       ,					atol=0.01     )     )



def 		lowercase_    (     A__ = 1000	)							->							int:
    """simple docstring"""

    snake_case				  =      -1
    snake_case				  =      0
    for a in range(1				,	n // 3	):
        # Solving the two equations a**2+b**2=c**2 and a+b+c=N eliminating c
        snake_case				  =      (n * n - 2 * a * n) // (2 * n - 2 * a)
        snake_case				  =      n - a - b
        if c * c == (a * a + b * b):
            snake_case				  =      a * b * c
            if candidate >= product:
                snake_case				  =      candidate
    return product


if __name__ == "__main__":
   print(f"{solution() = }")

