





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


   def 					snake_case				( self					):
       A       :       str							=						'''| <pad> <unk> <s> </s> a b c d e f g h i j k'''.split()
       A       :       Optional[int]							=						dict(zip(_UpperCAmelCase    ,   range(len(_UpperCAmelCase					)					)					)					)

       A       :       Tuple							=						{
           '''unk_token''': '''<unk>''',
           '''bos_token''': '''<s>''',
           '''eos_token''': '''</s>''',
       }
       A       :       Optional[int]							=						{
           '''feature_size''': 1,
           '''padding_value''': 0.0,
           '''sampling_rate''': 16_000,
           '''return_attention_mask''': False,
           '''do_normalize''': True,
       }

       A       :       Optional[Any]							=						tempfile.mkdtemp()
       A       :       Union[str, Any]							=						os.path.join(self.tmpdirname    ,   VOCAB_FILES_NAMES['''vocab_file''']					)
       A       :       Optional[Any]							=						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
       A       :       Any							=						'''hf-internal-testing/ngram-beam-search-decoder'''


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


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


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


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


   def 					snake_case				( self					):
       A       :       int							=						self.get_tokenizer()
       A       :       Union[str, Any]							=						self.get_feature_extractor()
       A       :       List[Any]							=						self.get_decoder()

       A       :       str							=						WavaVecaProcessorWithLM(tokenizer=_UpperCAmelCase    ,   feature_extractor=_UpperCAmelCase    ,   decoder=_UpperCAmelCase					)

       processor.save_pretrained(self.tmpdirname					)
       A       :       Any							=						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 					snake_case				( self					):
       A       :       Any							=						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
       A       :       Optional[int]							=						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					):
       A       :       Any							=						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 					snake_case				( self					):
       A       :       Dict							=						self.get_feature_extractor()
       A       :       List[Any]							=						self.get_tokenizer()
       A       :       Optional[int]							=						self.get_decoder()

       A       :       str							=						WavaVecaProcessorWithLM(tokenizer=_UpperCAmelCase    ,   feature_extractor=_UpperCAmelCase    ,   decoder=_UpperCAmelCase					)

       A       :       Tuple							=						floats_list((3, 1_000)					)

       A       :       Optional[int]							=						feature_extractor(_UpperCAmelCase    ,   return_tensors='''np'''					)
       A       :       Dict							=						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 					snake_case				( self					):
       A       :       Union[str, Any]							=						self.get_feature_extractor()
       A       :       int							=						self.get_tokenizer()
       A       :       int							=						self.get_decoder()

       A       :       List[str]							=						WavaVecaProcessorWithLM(tokenizer=_UpperCAmelCase    ,   feature_extractor=_UpperCAmelCase    ,   decoder=_UpperCAmelCase					)

       A       :       int							=						'''This is a test string'''

       A       :       int							=						processor(text=_UpperCAmelCase					)

       A       :       Union[str, Any]							=						tokenizer(_UpperCAmelCase					)

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


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


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

       A       :       Tuple							=						WavaVecaProcessorWithLM(tokenizer=_UpperCAmelCase    ,   feature_extractor=_UpperCAmelCase    ,   decoder=_UpperCAmelCase					)

       A       :       Any							=						self._get_dummy_logits(shape=(10, 16)    ,   seed=13					)

       A       :       Optional[int]							=						processor.decode(_UpperCAmelCase					)

       A       :       int							=						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 					snake_case				( self    ,   _UpperCAmelCase					):
       A       :       Dict							=						self.get_feature_extractor()
       A       :       List[Any]							=						self.get_tokenizer()
       A       :       Tuple							=						self.get_decoder()

       A       :       int							=						WavaVecaProcessorWithLM(tokenizer=_UpperCAmelCase    ,   feature_extractor=_UpperCAmelCase    ,   decoder=_UpperCAmelCase					)

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

       A       :       Optional[int]							=						list(_UpperCAmelCase					)

       with get_context('''fork'''					).Pool() as p:
           A       :       List[Any]							=						decoder.decode_beams_batch(_UpperCAmelCase    ,   _UpperCAmelCase					)

       A,	A,	A       :       int							=						[], [], []
       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 					snake_case				( self					):
       A       :       Tuple							=						self.get_feature_extractor()
       A       :       Optional[Any]							=						self.get_tokenizer()
       A       :       Any							=						self.get_decoder()

       A       :       int							=						WavaVecaProcessorWithLM(tokenizer=_UpperCAmelCase    ,   feature_extractor=_UpperCAmelCase    ,   decoder=_UpperCAmelCase					)

       A       :       List[str]							=						self._get_dummy_logits()

       A       :       Optional[int]							=						15
       A       :       Any							=						-20.0
       A       :       Union[str, Any]							=						-4.0

       A       :       int							=						processor.batch_decode(
           _UpperCAmelCase    ,   beam_width=_UpperCAmelCase    ,   beam_prune_logp=_UpperCAmelCase    ,   token_min_logp=_UpperCAmelCase    ,   )
       A       :       List[str]							=						decoded_processor_out.text

       A       :       Tuple							=						list(_UpperCAmelCase					)

       with get_context('''fork'''					).Pool() as pool:
           A       :       str							=						decoder.decode_beams_batch(
               _UpperCAmelCase    ,   _UpperCAmelCase    ,   beam_width=_UpperCAmelCase    ,   beam_prune_logp=_UpperCAmelCase    ,   token_min_logp=_UpperCAmelCase    ,   )

       A       :       List[Any]							=						[d[0][0] for d in decoded_decoder_out]
       A       :       List[str]							=						[d[0][2] for d in decoded_decoder_out]
       A       :       List[Any]							=						[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 					snake_case				( self					):
       A       :       int							=						self.get_feature_extractor()
       A       :       List[str]							=						self.get_tokenizer()
       A       :       Union[str, Any]							=						self.get_decoder()

       A       :       int							=						WavaVecaProcessorWithLM(tokenizer=_UpperCAmelCase    ,   feature_extractor=_UpperCAmelCase    ,   decoder=_UpperCAmelCase					)

       A       :       List[Any]							=						self._get_dummy_logits()

       A       :       Dict							=						2.0
       A       :       List[Any]							=						5.0
       A       :       List[str]							=						-20.0
       A       :       Any							=						True

       A       :       str							=						processor.batch_decode(
           _UpperCAmelCase    ,   alpha=_UpperCAmelCase    ,   beta=_UpperCAmelCase    ,   unk_score_offset=_UpperCAmelCase    ,   lm_score_boundary=_UpperCAmelCase    ,   )
       A       :       List[Any]							=						decoded_processor_out.text

       A       :       str							=						list(_UpperCAmelCase					)
       decoder.reset_params(
           alpha=_UpperCAmelCase    ,   beta=_UpperCAmelCase    ,   unk_score_offset=_UpperCAmelCase    ,   lm_score_boundary=_UpperCAmelCase    ,   )

       with get_context('''fork'''					).Pool() as pool:
           A       :       List[Any]							=						decoder.decode_beams_batch(
               _UpperCAmelCase    ,   _UpperCAmelCase    ,   )

       A       :       int							=						[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					)
       A       :       List[Any]							=						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 					snake_case				( self					):
       A       :       int							=						WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''					)

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

       A       :       Tuple							=						os.listdir(_UpperCAmelCase					)
       A       :       Any							=						['''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 					snake_case				( self					):
       A       :       str							=						snapshot_download('''hf-internal-testing/processor_with_lm'''					)

       A       :       Dict							=						WavaVecaProcessorWithLM.from_pretrained(_UpperCAmelCase					)

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

       A       :       List[str]							=						os.listdir(_UpperCAmelCase					)
       A       :       Optional[Any]							=						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 					snake_case				( self					):
       A       :       List[Any]							=						WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''					)
       A       :       List[Any]							=						AutoProcessor.from_pretrained('''hf-internal-testing/processor_with_lm'''					)

       A       :       List[str]							=						floats_list((3, 1_000)					)

       A       :       Union[str, Any]							=						processor_wavaveca(_UpperCAmelCase    ,   return_tensors='''np'''					)
       A       :       Any							=						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					)

       A       :       Any							=						self._get_dummy_logits()

       A       :       Optional[int]							=						processor_wavaveca.batch_decode(_UpperCAmelCase					)
       A       :       Union[str, Any]							=						processor_auto.batch_decode(_UpperCAmelCase					)

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


   def 					snake_case				( self					):
       A       :       Optional[int]							=						self.get_feature_extractor()
       A       :       int							=						self.get_tokenizer()
       A       :       Optional[int]							=						self.get_decoder()

       A       :       Optional[int]							=						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 					snake_case				( _UpperCAmelCase    ,   _UpperCAmelCase					):
       A       :       Dict							=						[d[key] for d in offsets]
       return retrieved_list


   def 					snake_case				( self					):
       A       :       List[str]							=						WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''					)
       A       :       Any							=						self._get_dummy_logits()[0]

       A       :       int							=						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 					snake_case				( self					):
       A       :       str							=						WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''					)
       A       :       Any							=						self._get_dummy_logits()

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

       A       :       Any							=						load_dataset('''common_voice'''    ,   '''en'''    ,   split='''train'''    ,   streaming=_UpperCAmelCase					)
       A       :       Dict							=						ds.cast_column('''audio'''    ,   datasets.Audio(sampling_rate=16_000					)					)
       A       :       Union[str, Any]							=						iter(_UpperCAmelCase					)
       A       :       Any							=						next(_UpperCAmelCase					)

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

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

       A       :       List[Any]							=						processor.decode(logits[0]    ,   output_word_offsets=_UpperCAmelCase					)

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

       A       :       Tuple							=						'''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
       A       :       str							=						torch.tensor(self.get_from_offsets(_UpperCAmelCase    ,   '''start_time'''					)					)
       A       :       Tuple							=						torch.tensor(self.get_from_offsets(_UpperCAmelCase    ,   '''end_time'''					)					)

       # fmt: off
       A       :       str							=						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]					)
       A       :       str							=						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 random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
     import os

     import jax.numpy as jnp
     from jax import jit

     from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
     from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

     snake_case_     =	"""0.12"""  # assumed parallelism: 8


if is_torch_available():
     import torch
def 							_lowerCamelCase(  UpperCamelCase__    : int    ,       UpperCamelCase__    : Optional[int]    ,       UpperCamelCase__    : Any=None		)		->  List[Any]:
    if rng is None:
        A       :       int							=						random.Random()

    A       :       int							=						1
    for dim in shape:
        total_dims *= dim

    A       :       Union[str, Any]							=						[]
    for _ in range(UpperCamelCase__		):
        values.append(rng.randint(0    ,       vocab_size - 1		)		)

    A       :       Union[str, Any]							=						np.array(UpperCamelCase__    ,       dtype=jnp.intaa		).reshape(UpperCamelCase__		)

    return output
def 							_lowerCamelCase(  UpperCamelCase__    : Dict    ,       UpperCamelCase__    : Optional[Any]=None		)		->  Dict:
    A       :       Union[str, Any]							=						ids_tensor(UpperCamelCase__    ,       vocab_size=2    ,       rng=UpperCamelCase__		)
    # make sure that at least one token is attended to for each batch
    A       :       Any							=						1
    return attn_mask


@require_flax
class   _lowercase						:
   _UpperCamelCase       		=   None
   _UpperCamelCase       		=   ()


   def 					snake_case				( self					):
       A,	A       :       Optional[int]							=						self.model_tester.prepare_config_and_inputs_for_common()

       # cut to half length & take max batch_size 3
       A       :       Dict							=						2
       A       :       List[str]							=						inputs['''input_ids'''].shape[-1] // 2
       A       :       int							=						inputs['''input_ids'''][:max_batch_size, :sequence_length]

       A       :       Optional[Any]							=						jnp.ones_like(_UpperCAmelCase					)
       A       :       int							=						attention_mask[:max_batch_size, :sequence_length]

       # generate max 5 tokens
       A       :       Optional[int]							=						input_ids.shape[-1] + 5
       if config.eos_token_id is not None and config.pad_token_id is None:
           # hack to allow generate for models such as GPT2 as is done in `generate()`
           A       :       str							=						config.eos_token_id
       return config, input_ids, attention_mask, max_length


   @is_pt_flax_cross_test
   def 					snake_case				( self					):
       A,	A,	A,	A       :       Optional[Any]							=						self._get_input_ids_and_config()
       A       :       Optional[int]							=						False
       A       :       Union[str, Any]							=						max_length
       A       :       str							=						0

       for model_class in self.all_generative_model_classes:
           A       :       int							=						model_class(_UpperCAmelCase					)

           A       :       Dict							=						model_class.__name__[4:]  # Skip the "Flax" at the beginning
           A       :       int							=						getattr(_UpperCAmelCase    ,   _UpperCAmelCase					)
           A       :       Union[str, Any]							=						pt_model_class(_UpperCAmelCase					).eval()
           A       :       Tuple							=						load_flax_weights_in_pytorch_model(_UpperCAmelCase    ,   flax_model.params					)

           A       :       Union[str, Any]							=						flax_model.generate(_UpperCAmelCase					).sequences
           A       :       Optional[int]							=						pt_model.generate(torch.tensor(_UpperCAmelCase    ,   dtype=torch.long					)					)

           if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
               A       :       List[Any]							=						flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

           self.assertListEqual(pt_generation_outputs.numpy().tolist()    ,   flax_generation_outputs.tolist()					)


   def 					snake_case				( self					):
       A,	A,	A,	A       :       str							=						self._get_input_ids_and_config()
       A       :       Optional[int]							=						False
       A       :       Dict							=						max_length

       for model_class in self.all_generative_model_classes:
           A       :       Any							=						model_class(_UpperCAmelCase					)

           A       :       Dict							=						model.generate(_UpperCAmelCase					).sequences
           self.assertEqual(generation_outputs.shape[-1]    ,   _UpperCAmelCase					)

           A       :       str							=						jit(model.generate					)
           A       :       List[Any]							=						jit_generate(_UpperCAmelCase					).sequences

           self.assertListEqual(generation_outputs.tolist()    ,   jit_generation_outputs.tolist()					)


   def 					snake_case				( self					):
       A,	A,	A,	A       :       List[Any]							=						self._get_input_ids_and_config()
       A       :       List[str]							=						True
       A       :       List[Any]							=						max_length

       for model_class in self.all_generative_model_classes:
           A       :       Any							=						model_class(_UpperCAmelCase					)

           A       :       Tuple							=						model.generate(_UpperCAmelCase					).sequences
           self.assertEqual(generation_outputs.shape[-1]    ,   _UpperCAmelCase					)

           A       :       Tuple							=						jit(model.generate					)
           A       :       int							=						jit_generate(_UpperCAmelCase					).sequences

           self.assertListEqual(generation_outputs.tolist()    ,   jit_generation_outputs.tolist()					)


   def 					snake_case				( self					):
       A,	A,	A,	A       :       List[str]							=						self._get_input_ids_and_config()
       A       :       Any							=						False
       A       :       str							=						max_length
       A       :       Optional[int]							=						2

       for model_class in self.all_generative_model_classes:
           A       :       Dict							=						model_class(_UpperCAmelCase					)

           A       :       Any							=						model.generate(_UpperCAmelCase					).sequences
           self.assertEqual(generation_outputs.shape[-1]    ,   _UpperCAmelCase					)

           A       :       List[Any]							=						jit(model.generate					)
           A       :       Optional[Any]							=						jit_generate(_UpperCAmelCase					).sequences

           self.assertListEqual(generation_outputs.tolist()    ,   jit_generation_outputs.tolist()					)


   def 					snake_case				( self					):
       A,	A,	A,	A       :       Optional[int]							=						self._get_input_ids_and_config()
       A       :       Dict							=						False
       A       :       List[Any]							=						max_length
       A       :       Optional[int]							=						2
       A       :       int							=						2

       for model_class in self.all_generative_model_classes:
           A       :       List[str]							=						model_class(_UpperCAmelCase					)

           A       :       List[str]							=						model.generate(_UpperCAmelCase					).sequences
           self.assertEqual(generation_outputs.shape[0]    ,   input_ids.shape[0] * config.num_return_sequences					)


   def 					snake_case				( self					):
       A,	A,	A,	A       :       List[Any]							=						self._get_input_ids_and_config()
       A       :       Any							=						True
       A       :       Any							=						max_length
       A       :       Any							=						0.8
       A       :       List[Any]							=						10
       A       :       List[Any]							=						0.3
       A       :       Union[str, Any]							=						1
       A       :       Any							=						8
       A       :       List[Any]							=						9

       for model_class in self.all_generative_model_classes:
           A       :       List[Any]							=						model_class(_UpperCAmelCase					)

           A       :       Any							=						model.generate(_UpperCAmelCase					).sequences
           self.assertEqual(generation_outputs.shape[-1]    ,   _UpperCAmelCase					)

           A       :       int							=						jit(model.generate					)
           A       :       List[Any]							=						jit_generate(_UpperCAmelCase					).sequences

           self.assertListEqual(generation_outputs.tolist()    ,   jit_generation_outputs.tolist()					)


   def 					snake_case				( self					):
       A,	A,	A,	A       :       Optional[Any]							=						self._get_input_ids_and_config()
       A       :       int							=						max_length
       A       :       Dict							=						1
       A       :       Tuple							=						8
       A       :       Optional[Any]							=						9

       for model_class in self.all_generative_model_classes:
           A       :       Dict							=						model_class(_UpperCAmelCase					)

           A       :       str							=						model.generate(_UpperCAmelCase					).sequences
           self.assertEqual(generation_outputs.shape[-1]    ,   _UpperCAmelCase					)

           A       :       int							=						jit(model.generate					)
           A       :       Optional[int]							=						jit_generate(_UpperCAmelCase					).sequences

           self.assertListEqual(generation_outputs.tolist()    ,   jit_generation_outputs.tolist()					)


   def 					snake_case				( self					):
       A,	A,	A,	A       :       int							=						self._get_input_ids_and_config()
       A       :       List[Any]							=						max_length
       A       :       Optional[Any]							=						2
       A       :       Tuple							=						1
       A       :       List[Any]							=						8
       A       :       Dict							=						9

       for model_class in self.all_generative_model_classes:
           A       :       Union[str, Any]							=						model_class(_UpperCAmelCase					)

           A       :       Any							=						model.generate(_UpperCAmelCase					).sequences
           self.assertEqual(generation_outputs.shape[-1]    ,   _UpperCAmelCase					)

           A       :       Union[str, Any]							=						jit(model.generate					)
           A       :       int							=						jit_generate(_UpperCAmelCase					).sequences

           self.assertListEqual(generation_outputs.tolist()    ,   jit_generation_outputs.tolist()					)


   def 					snake_case				( self					):
       A,	A,	A,	A       :       Dict							=						self._get_input_ids_and_config()

       # pad attention mask on the left
       A       :       Union[str, Any]							=						attention_mask.at[(0, 0)].set(0					)

       A       :       int							=						False
       A       :       List[Any]							=						max_length

       for model_class in self.all_generative_model_classes:
           A       :       List[str]							=						model_class(_UpperCAmelCase					)

           A       :       Dict							=						model.generate(_UpperCAmelCase    ,   attention_mask=_UpperCAmelCase					).sequences
           self.assertEqual(generation_outputs.shape[-1]    ,   _UpperCAmelCase					)

           A       :       Dict							=						jit(model.generate					)
           A       :       Optional[Any]							=						jit_generate(_UpperCAmelCase    ,   attention_mask=_UpperCAmelCase					).sequences

           self.assertListEqual(generation_outputs.tolist()    ,   jit_generation_outputs.tolist()					)


   def 					snake_case				( self					):
       A,	A,	A,	A       :       Optional[Any]							=						self._get_input_ids_and_config()

       # pad attention mask on the left
       A       :       Optional[Any]							=						attention_mask.at[(0, 0)].set(0					)

       A       :       Optional[int]							=						True
       A       :       int							=						max_length

       for model_class in self.all_generative_model_classes:
           A       :       int							=						model_class(_UpperCAmelCase					)

           A       :       Union[str, Any]							=						model.generate(_UpperCAmelCase    ,   attention_mask=_UpperCAmelCase					).sequences
           self.assertEqual(generation_outputs.shape[-1]    ,   _UpperCAmelCase					)

           A       :       Optional[Any]							=						jit(model.generate					)
           A       :       Dict							=						jit_generate(_UpperCAmelCase    ,   attention_mask=_UpperCAmelCase					).sequences

           self.assertListEqual(generation_outputs.tolist()    ,   jit_generation_outputs.tolist()					)


   def 					snake_case				( self					):
       A,	A,	A,	A       :       Optional[Any]							=						self._get_input_ids_and_config()

       # pad attention mask on the left
       A       :       Optional[int]							=						attention_mask.at[(0, 0)].set(0					)

       A       :       Any							=						2
       A       :       Any							=						max_length

       for model_class in self.all_generative_model_classes:
           A       :       Dict							=						model_class(_UpperCAmelCase					)

           A       :       Optional[Any]							=						model.generate(_UpperCAmelCase    ,   attention_mask=_UpperCAmelCase					).sequences
           self.assertEqual(generation_outputs.shape[-1]    ,   _UpperCAmelCase					)

           A       :       str							=						jit(model.generate					)
           A       :       List[str]							=						jit_generate(_UpperCAmelCase    ,   attention_mask=_UpperCAmelCase					).sequences

           self.assertListEqual(generation_outputs.tolist()    ,   jit_generation_outputs.tolist()					)







@require_flax
class   _lowercase						(							unittest.TestCase				):


   def 					snake_case				( self					):
       A       :       List[Any]							=						AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-bert'''					)
       A       :       Dict							=						FlaxAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''					)

       A       :       Union[str, Any]							=						'''Hello world'''
       A       :       Union[str, Any]							=						tokenizer(_UpperCAmelCase    ,   return_tensors='''np'''					).input_ids

       # typos are quickly detected (the correct argument is `do_sample`)
       with self.assertRaisesRegex(_UpperCAmelCase    ,   '''do_samples'''					):
           model.generate(_UpperCAmelCase    ,   do_samples=_UpperCAmelCase					)

       # arbitrary arguments that will not be used anywhere are also not accepted
       with self.assertRaisesRegex(_UpperCAmelCase    ,   '''foo'''					):
           A       :       List[Any]							=						{'''foo''': '''bar'''}
           model.generate(_UpperCAmelCase    ,   **_UpperCAmelCase					)
