'''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  lowerCAmelCase__ (      self			:       str  )				->Any:
        UpperCAmelCase_            =  '| <pad> <unk> <s> </s> a b c d e f g h i j k'.split()
        UpperCAmelCase_            =  dict(zip(UpperCAmelCase_  ,							range(len(UpperCAmelCase_  )  )  )  )

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

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

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

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






  def  lowerCAmelCase__ (      self			:       Optional[Any]  ,							**UpperCAmelCase__			:       Optional[int]  )				->Any:
        UpperCAmelCase_            =  self.add_kwargs_tokens_map.copy()
        kwargs.update(UpperCAmelCase_  )
        return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname  ,							**UpperCAmelCase_  )






  def  lowerCAmelCase__ (      self			:       Tuple  ,							**UpperCAmelCase__			:       Tuple  )				->int:
        return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname  ,							**UpperCAmelCase_  )






  def  lowerCAmelCase__ (      self			:       List[Any]  ,							**UpperCAmelCase__			:       List[Any]  )				->List[str]:
        return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name  ,							**UpperCAmelCase_  )






  def  lowerCAmelCase__ (      self			:       Union[str, Any]  )				->Union[str, Any]:
        shutil.rmtree(self.tmpdirname  )






  def  lowerCAmelCase__ (      self			:       Optional[Any]  )				->Dict:
        UpperCAmelCase_            =  self.get_tokenizer()
        UpperCAmelCase_            =  self.get_feature_extractor()
        UpperCAmelCase_            =  self.get_decoder()

        UpperCAmelCase_            =  WavaVecaProcessorWithLM(tokenizer=UpperCAmelCase_  ,							feature_extractor=UpperCAmelCase_  ,							decoder=UpperCAmelCase_  )

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

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

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

        # decoder
        self.assertEqual(processor.decoder._alphabet.labels  ,							decoder._alphabet.labels  )
        self.assertEqual(
            processor.decoder.model_container[decoder._model_key]._unigram_set  ,							decoder.model_container[decoder._model_key]._unigram_set  ,							)
        self.assertIsInstance(processor.decoder  ,							UpperCAmelCase_  )






  def  lowerCAmelCase__ (      self			:       List[str]  )				->List[str]:
        UpperCAmelCase_            =  WavaVecaProcessorWithLM(
            tokenizer=self.get_tokenizer()  ,							feature_extractor=self.get_feature_extractor()  ,							decoder=self.get_decoder()  )
        processor.save_pretrained(self.tmpdirname  )

        # make sure that error is thrown when decoder alphabet doesn't match
        UpperCAmelCase_            =  WavaVecaProcessorWithLM.from_pretrained(
            self.tmpdirname  ,							alpha=5.0  ,							beta=3.0  ,							score_boundary=-7.0  ,							unk_score_offset=3  )

        # decoder
        self.assertEqual(processor.language_model.alpha  ,							5.0  )
        self.assertEqual(processor.language_model.beta  ,							3.0  )
        self.assertEqual(processor.language_model.score_boundary  ,							-7.0  )
        self.assertEqual(processor.language_model.unk_score_offset  ,							3  )






  def  lowerCAmelCase__ (      self			:       Any  )				->List[str]:
        UpperCAmelCase_            =  self.get_tokenizer()
        # add token to trigger raise
        tokenizer.add_tokens(['''xx''']  )
        with self.assertRaisesRegex(UpperCAmelCase_  ,							'''include'''  ):
              WavaVecaProcessorWithLM(
                  tokenizer=UpperCAmelCase_  ,							feature_extractor=self.get_feature_extractor()  ,							decoder=self.get_decoder()  )






  def  lowerCAmelCase__ (      self			:       List[Any]  )				->str:
        UpperCAmelCase_            =  self.get_feature_extractor()
        UpperCAmelCase_            =  self.get_tokenizer()
        UpperCAmelCase_            =  self.get_decoder()

        UpperCAmelCase_            =  WavaVecaProcessorWithLM(tokenizer=UpperCAmelCase_  ,							feature_extractor=UpperCAmelCase_  ,							decoder=UpperCAmelCase_  )

        UpperCAmelCase_            =  floats_list((3, 1000)  )

        UpperCAmelCase_            =  feature_extractor(UpperCAmelCase_  ,							return_tensors='''np'''  )
        UpperCAmelCase_            =  processor(UpperCAmelCase_  ,							return_tensors='''np'''  )

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






  def  lowerCAmelCase__ (      self			:       Dict  )				->Tuple:
        UpperCAmelCase_            =  self.get_feature_extractor()
        UpperCAmelCase_            =  self.get_tokenizer()
        UpperCAmelCase_            =  self.get_decoder()

        UpperCAmelCase_            =  WavaVecaProcessorWithLM(tokenizer=UpperCAmelCase_  ,							feature_extractor=UpperCAmelCase_  ,							decoder=UpperCAmelCase_  )

        UpperCAmelCase_            =  'This is a test string'

        UpperCAmelCase_            =  processor(text=UpperCAmelCase_  )

        UpperCAmelCase_            =  tokenizer(UpperCAmelCase_  )

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






  def  lowerCAmelCase__ (      self			:       Union[str, Any]  ,							UpperCAmelCase__			:       int=(2, 10, 16)  ,							UpperCAmelCase__			:       str=77  )				->str:
        np.random.seed(UpperCAmelCase_  )
        return np.random.rand(*UpperCAmelCase_  )






  def  lowerCAmelCase__ (      self			:       Any  )				->List[Any]:
        UpperCAmelCase_            =  self.get_feature_extractor()
        UpperCAmelCase_            =  self.get_tokenizer()
        UpperCAmelCase_            =  self.get_decoder()

        UpperCAmelCase_            =  WavaVecaProcessorWithLM(tokenizer=UpperCAmelCase_  ,							feature_extractor=UpperCAmelCase_  ,							decoder=UpperCAmelCase_  )

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

        UpperCAmelCase_            =  processor.decode(UpperCAmelCase_  )

        UpperCAmelCase_            =  decoder.decode_beams(UpperCAmelCase_  )[0]

        self.assertEqual(decoded_decoder[0]  ,							decoded_processor.text  )
        self.assertEqual('''</s> <s> </s>'''  ,							decoded_processor.text  )
        self.assertEqual(decoded_decoder[-2]  ,							decoded_processor.logit_score  )
        self.assertEqual(decoded_decoder[-1]  ,							decoded_processor.lm_score  )






  @parameterized.expand([[None], ['''fork'''], ['''spawn''']]  )
  def  lowerCAmelCase__ (      self			:       Tuple  ,							UpperCAmelCase__			:       Tuple  )				->str:
        UpperCAmelCase_            =  self.get_feature_extractor()
        UpperCAmelCase_            =  self.get_tokenizer()
        UpperCAmelCase_            =  self.get_decoder()

        UpperCAmelCase_            =  WavaVecaProcessorWithLM(tokenizer=UpperCAmelCase_  ,							feature_extractor=UpperCAmelCase_  ,							decoder=UpperCAmelCase_  )

        UpperCAmelCase_            =  self._get_dummy_logits()

        # note: pool should be instantiated *after* Wav2Vec2ProcessorWithLM.
        #       otherwise, the LM won't be available to the pool's sub-processes.
        # manual logic used to allow parameterized test for both pool=None and pool=Pool(...)
        if pool_context is None:
              UpperCAmelCase_            =  processor.batch_decode(UpperCAmelCase_  )
        else:
              with get_context(UpperCAmelCase_  ).Pool() as pool:
                    UpperCAmelCase_            =  processor.batch_decode(UpperCAmelCase_  ,							UpperCAmelCase_  )

        UpperCAmelCase_            =  list(UpperCAmelCase_  )

        with get_context('''fork'''  ).Pool() as p:
              UpperCAmelCase_            =  decoder.decode_beams_batch(UpperCAmelCase_  ,							UpperCAmelCase_  )

        UpperCAmelCase_            =  [], [], []
        for beams in decoded_beams:
              texts_decoder.append(beams[0][0]  )
              logit_scores_decoder.append(beams[0][-2]  )
              lm_scores_decoder.append(beams[0][-1]  )

        self.assertListEqual(UpperCAmelCase_  ,							decoded_processor.text  )
        self.assertListEqual(['''<s> <s> </s>''', '''<s> <s> <s>''']  ,							decoded_processor.text  )
        self.assertListEqual(UpperCAmelCase_  ,							decoded_processor.logit_score  )
        self.assertListEqual(UpperCAmelCase_  ,							decoded_processor.lm_score  )






  def  lowerCAmelCase__ (      self			:       str  )				->Optional[Any]:
        UpperCAmelCase_            =  self.get_feature_extractor()
        UpperCAmelCase_            =  self.get_tokenizer()
        UpperCAmelCase_            =  self.get_decoder()

        UpperCAmelCase_            =  WavaVecaProcessorWithLM(tokenizer=UpperCAmelCase_  ,							feature_extractor=UpperCAmelCase_  ,							decoder=UpperCAmelCase_  )

        UpperCAmelCase_            =  self._get_dummy_logits()

        UpperCAmelCase_            =  15
        UpperCAmelCase_            =  -20.0
        UpperCAmelCase_            =  -4.0

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

        UpperCAmelCase_            =  list(UpperCAmelCase_  )

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

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

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

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

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






  def  lowerCAmelCase__ (      self			:       List[Any]  )				->Any:
        UpperCAmelCase_            =  self.get_feature_extractor()
        UpperCAmelCase_            =  self.get_tokenizer()
        UpperCAmelCase_            =  self.get_decoder()

        UpperCAmelCase_            =  WavaVecaProcessorWithLM(tokenizer=UpperCAmelCase_  ,							feature_extractor=UpperCAmelCase_  ,							decoder=UpperCAmelCase_  )

        UpperCAmelCase_            =  self._get_dummy_logits()

        UpperCAmelCase_            =  2.0
        UpperCAmelCase_            =  5.0
        UpperCAmelCase_            =  -20.0
        UpperCAmelCase_            =  True

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

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

        with get_context('''fork'''  ).Pool() as pool:
              UpperCAmelCase_            =  decoder.decode_beams_batch(
                  UpperCAmelCase_  ,							UpperCAmelCase_  ,							)

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

        self.assertListEqual(UpperCAmelCase_  ,							UpperCAmelCase_  )
        self.assertListEqual(['''<s> </s> <s> </s> </s>''', '''</s> </s> <s> </s> </s>''']  ,							UpperCAmelCase_  )
        UpperCAmelCase_            =  processor.decoder.model_container[processor.decoder._model_key]
        self.assertEqual(lm_model.alpha  ,							2.0  )
        self.assertEqual(lm_model.beta  ,							5.0  )
        self.assertEqual(lm_model.unk_score_offset  ,							-20.0  )
        self.assertEqual(lm_model.score_boundary  ,							UpperCAmelCase_  )






  def  lowerCAmelCase__ (      self			:       Dict  )				->Union[str, Any]:
        UpperCAmelCase_            =  WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''  )

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

        UpperCAmelCase_            =  os.listdir(UpperCAmelCase_  )
        UpperCAmelCase_            =  ['alphabet.json', 'language_model']

        downloaded_decoder_files.sort()
        expected_decoder_files.sort()

        # test that only decoder relevant files from
        # https://huggingface.co/hf-internal-testing/processor_with_lm/tree/main
        # are downloaded and none of the rest (e.g. README.md, ...)
        self.assertListEqual(UpperCAmelCase_  ,							UpperCAmelCase_  )






  def  lowerCAmelCase__ (      self			:       List[str]  )				->Any:
        UpperCAmelCase_            =  snapshot_download('''hf-internal-testing/processor_with_lm'''  )

        UpperCAmelCase_            =  WavaVecaProcessorWithLM.from_pretrained(UpperCAmelCase_  )

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

        UpperCAmelCase_            =  os.listdir(UpperCAmelCase_  )
        UpperCAmelCase_            =  os.listdir(UpperCAmelCase_  )

        local_decoder_files.sort()
        expected_decoder_files.sort()

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






  def  lowerCAmelCase__ (      self			:       Tuple  )				->Union[str, Any]:
        UpperCAmelCase_            =  WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''  )
        UpperCAmelCase_            =  AutoProcessor.from_pretrained('''hf-internal-testing/processor_with_lm'''  )

        UpperCAmelCase_            =  floats_list((3, 1000)  )

        UpperCAmelCase_            =  processor_wavaveca(UpperCAmelCase_  ,							return_tensors='''np'''  )
        UpperCAmelCase_            =  processor_auto(UpperCAmelCase_  ,							return_tensors='''np'''  )

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

        UpperCAmelCase_            =  self._get_dummy_logits()

        UpperCAmelCase_            =  processor_wavaveca.batch_decode(UpperCAmelCase_  )
        UpperCAmelCase_            =  processor_auto.batch_decode(UpperCAmelCase_  )

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






  def  lowerCAmelCase__ (      self			:       Dict  )				->List[str]:
        UpperCAmelCase_            =  self.get_feature_extractor()
        UpperCAmelCase_            =  self.get_tokenizer()
        UpperCAmelCase_            =  self.get_decoder()

        UpperCAmelCase_            =  WavaVecaProcessorWithLM(tokenizer=UpperCAmelCase_  ,							feature_extractor=UpperCAmelCase_  ,							decoder=UpperCAmelCase_  )

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






  @staticmethod
  def  lowerCAmelCase__ (      UpperCAmelCase__			:       Optional[Any]  ,							UpperCAmelCase__			:       List[Any]  )				->List[str]:
        UpperCAmelCase_            =  [d[key] for d in offsets]
        return retrieved_list






  def  lowerCAmelCase__ (      self			:       Dict  )				->Optional[int]:
        UpperCAmelCase_            =  WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''  )
        UpperCAmelCase_            =  self._get_dummy_logits()[0]

        UpperCAmelCase_            =  processor.decode(UpperCAmelCase_  ,							output_word_offsets=UpperCAmelCase_  )
        # check Wav2Vec2CTCTokenizerOutput keys for word
        self.assertEqual(len(outputs.keys()  )  ,							4  )
        self.assertTrue('''text''' in outputs  )
        self.assertTrue('''word_offsets''' in outputs  )
        self.assertTrue(isinstance(UpperCAmelCase_  ,							UpperCAmelCase_  )  )

        self.assertEqual(''' '''.join(self.get_from_offsets(outputs['''word_offsets''']  ,							'''word'''  )  )  ,							outputs.text  )
        self.assertListEqual(self.get_from_offsets(outputs['''word_offsets''']  ,							'''word'''  )  ,							['''<s>''', '''<s>''', '''</s>''']  )
        self.assertListEqual(self.get_from_offsets(outputs['''word_offsets''']  ,							'''start_offset'''  )  ,							[0, 2, 4]  )
        self.assertListEqual(self.get_from_offsets(outputs['''word_offsets''']  ,							'''end_offset'''  )  ,							[1, 3, 5]  )






  def  lowerCAmelCase__ (      self			:       List[Any]  )				->int:
        UpperCAmelCase_            =  WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''  )
        UpperCAmelCase_            =  self._get_dummy_logits()

        UpperCAmelCase_            =  processor.batch_decode(UpperCAmelCase_  ,							output_word_offsets=UpperCAmelCase_  )

        # check Wav2Vec2CTCTokenizerOutput keys for word
        self.assertEqual(len(outputs.keys()  )  ,							4  )
        self.assertTrue('''text''' in outputs  )
        self.assertTrue('''word_offsets''' in outputs  )
        self.assertTrue(isinstance(UpperCAmelCase_  ,							UpperCAmelCase_  )  )

        self.assertListEqual(
            [''' '''.join(self.get_from_offsets(UpperCAmelCase_  ,							'''word'''  )  ) for o in outputs['''word_offsets''']]  ,							outputs.text  )
        self.assertListEqual(self.get_from_offsets(outputs['''word_offsets'''][0]  ,							'''word'''  )  ,							['''<s>''', '''<s>''', '''</s>''']  )
        self.assertListEqual(self.get_from_offsets(outputs['''word_offsets'''][0]  ,							'''start_offset'''  )  ,							[0, 2, 4]  )
        self.assertListEqual(self.get_from_offsets(outputs['''word_offsets'''][0]  ,							'''end_offset'''  )  ,							[1, 3, 5]  )






  @slow
  @require_torch
  @require_torchaudio
  def  lowerCAmelCase__ (      self			:       Optional[int]  )				->Union[str, Any]:
        import torch

        UpperCAmelCase_            =  load_dataset('''common_voice'''  ,							'''en'''  ,							split='''train'''  ,							streaming=UpperCAmelCase_  )
        UpperCAmelCase_            =  ds.cast_column('''audio'''  ,							datasets.Audio(sampling_rate=1_6000  )  )
        UpperCAmelCase_            =  iter(UpperCAmelCase_  )
        UpperCAmelCase_            =  next(UpperCAmelCase_  )

        UpperCAmelCase_            =  AutoProcessor.from_pretrained('''patrickvonplaten/wav2vec2-base-100h-with-lm'''  )
        UpperCAmelCase_            =  WavaVecaForCTC.from_pretrained('''patrickvonplaten/wav2vec2-base-100h-with-lm'''  )

        # compare to filename `common_voice_en_100038.mp3` of dataset viewer on https://huggingface.co/datasets/common_voice/viewer/en/train
        UpperCAmelCase_            =  processor(sample['''audio''']['''array''']  ,							return_tensors='''pt'''  ).input_values

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

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

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

        UpperCAmelCase_            =  'WHY DOES MILISANDRA LOOK LIKE SHE WANTS TO CONSUME JOHN SNOW ON THE RIVER AT THE WALL'

        # output words
        self.assertEqual(''' '''.join(self.get_from_offsets(UpperCAmelCase_  ,							'''word'''  )  )  ,							UpperCAmelCase_  )
        self.assertEqual(''' '''.join(self.get_from_offsets(UpperCAmelCase_  ,							'''word'''  )  )  ,							output.text  )

        # output times
        UpperCAmelCase_            =  torch.tensor(self.get_from_offsets(UpperCAmelCase_  ,							'''start_time'''  )  )
        UpperCAmelCase_            =  torch.tensor(self.get_from_offsets(UpperCAmelCase_  ,							'''end_time'''  )  )

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

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






'''simple docstring'''
from __future__ import annotations

import os
import tempfile
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import is_tensorflow_text_available, is_tf_available
from transformers.testing_utils import require_tensorflow_text, require_tf, slow

from ..test_modeling_tf_common import floats_tensor
from .test_framework_agnostic import GenerationIntegrationTestsMixin


if is_tf_available():
      import tensorflow as tf

      from transformers import (
          AutoTokenizer,
          TFAutoModelForCausalLM,
          TFAutoModelForSeqaSeqLM,
          TFAutoModelForSpeechSeqaSeq,
          TFAutoModelForVisionaSeq,
          TFBartForConditionalGeneration,
          TFLogitsProcessorList,
          TFMinLengthLogitsProcessor,
          tf_top_k_top_p_filtering,
      )

if is_tensorflow_text_available():
      import tensorflow_text as text







@require_tf
class lowerCamelCase							(	unittest.TestCase					):




  '''simple docstring'''



  def  lowerCAmelCase__ (      self			:       Dict  )				->List[str]:
        UpperCAmelCase_            =  tf.convert_to_tensor(
            [
                [
                    8.222_0991,  # 3rd highest value; idx. 0
                    -0.562_0044,
                    5.2322_9752,
                    4.038_6393,
                    -6.879_8378,
                    -0.5478_5802,
                    -3.201_2153,
                    2.9277_7176,
                    1.8817_1953,
                    7.3534_1276,  # 5th highest value; idx. 9
                    8.4320_7833,  # 2nd highest value; idx. 10
                    -9.8571_1836,
                    -5.9620_9236,
                    -1.1303_9161,
                    -7.111_5294,
                    -0.836_9633,
                    -5.318_6408,
                    7.0642_7407,
                    0.8136_9344,
                    -0.8202_3817,
                    -5.917_9796,
                    0.5881_3443,
                    -6.9977_8438,
                    4.7155_1189,
                    -0.1877_1637,
                    7.4402_0759,  # 4th highest value; idx. 25
                    9.3845_0987,  # 1st highest value; idx. 26
                    2.1266_2941,
                    -9.3256_2038,
                    2.3565_2522,
                ],  # cummulative prob of 5 highest values <= 0.6
                [
                    0.5842_5518,
                    4.5313_9238,
                    -5.5751_0464,
                    -6.2803_0699,
                    -7.1952_9503,
                    -4.0212_2551,
                    1.3933_7037,
                    -6.0670_7057,
                    1.5948_0517,
                    -9.64_3119,
                    0.0390_7799,
                    0.6723_1762,
                    -8.8820_6726,
                    6.2711_5922,  # 4th highest value; idx. 13
                    2.2852_0723,
                    4.8276_7506,
                    4.3042_1368,
                    8.827_5313,  # 2nd highest value; idx. 17
                    5.4402_9958,  # 5th highest value; idx. 18
                    -4.473_5794,
                    7.3857_9536,  # 3rd highest value; idx. 20
                    -2.9105_1663,
                    2.6194_6077,
                    -2.567_4762,
                    -9.4895_9302,
                    -4.0292_2645,
                    -1.3541_6918,
                    9.6770_2323,  # 1st highest value; idx. 27
                    -5.8947_8553,
                    1.8537_0467,
                ],  # cummulative prob of 5 highest values <= 0.6
            ]  ,							dtype=tf.floataa  ,							)

        UpperCAmelCase_            =  tf.convert_to_tensor(
            [[0, 0], [0, 9], [0, 10], [0, 25], [0, 26], [1, 13], [1, 17], [1, 18], [1, 20], [1, 27]]  ,							dtype=tf.intaa  ,							)  # expected non filtered idx as noted above

        UpperCAmelCase_            =  tf.convert_to_tensor(
            [8.22_2099, 7.353_4126, 8.43_2078, 7.440_2075, 9.3_8451, 6.27_1159, 8.82_7531, 5.440_2995, 7.385_7956, 9.67_7023]  ,							dtype=tf.floataa  ,							)  # expected non filtered values as noted above

        UpperCAmelCase_            =  tf_top_k_top_p_filtering(UpperCAmelCase__  ,							top_k=10  ,							top_p=0.6  ,							min_tokens_to_keep=4  )

        UpperCAmelCase_            =  output[output != -float('''inf'''  )]
        UpperCAmelCase_            =  tf.cast(
            tf.where(tf.not_equal(UpperCAmelCase__  ,							tf.constant(-float('''inf'''  )  ,							dtype=tf.floataa  )  )  )  ,							dtype=tf.intaa  ,							)

        tf.debugging.assert_near(UpperCAmelCase__  ,							UpperCAmelCase__  ,							rtol=1e-12  )
        tf.debugging.assert_equal(UpperCAmelCase__  ,							UpperCAmelCase__  )







@require_tf
class lowerCamelCase							(	unittest.TestCase     ,		lowerCamelCase					):




  '''simple docstring'''



  if is_tf_available():
    lowerCAmelCase__										=	{
        '''AutoModelForCausalLM''': TFAutoModelForCausalLM,
        '''AutoModelForSpeechSeq2Seq''': TFAutoModelForSpeechSeqaSeq,
        '''AutoModelForSeq2SeqLM''': TFAutoModelForSeqaSeqLM,
        '''AutoModelForVision2Seq''': TFAutoModelForVisionaSeq,
        '''LogitsProcessorList''': TFLogitsProcessorList,
        '''MinLengthLogitsProcessor''': TFMinLengthLogitsProcessor,
        '''create_tensor_fn''': tf.convert_to_tensor,
        '''floats_tensor''': floats_tensor,
        '''return_tensors''': '''tf''',
    }






  @slow
  def  lowerCAmelCase__ (      self			:       Dict  )				->Optional[Any]:
        # TF-only test: tf.saved_model export
        UpperCAmelCase_            =  TFAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''  )
        UpperCAmelCase_            =  2
        UpperCAmelCase_            =  2







        class lowerCamelCase							(	tf.Module					):




          '''simple docstring'''



          def __init__(      self			:       List[str]  ,							UpperCAmelCase__			:       List[str]  )				->Dict:
                super(UpperCAmelCase__  ,							self  ).__init__()
                UpperCAmelCase_            =  model






          @tf.function(
              input_signature=(
                  tf.TensorSpec((None, input_length)  ,							tf.intaa  ,							name='''input_ids'''  ),
                  tf.TensorSpec((None, input_length)  ,							tf.intaa  ,							name='''attention_mask'''  ),
              )  ,							jit_compile=UpperCAmelCase__  ,							)
          def  lowerCAmelCase__ (      self			:       Union[str, Any]  ,							UpperCAmelCase__			:       Union[str, Any]  ,							UpperCAmelCase__			:       List[str]  )				->int:
                UpperCAmelCase_            =  self.model.generate(
                    input_ids=UpperCAmelCase__  ,							attention_mask=UpperCAmelCase__  ,							max_new_tokens=UpperCAmelCase__  ,							return_dict_in_generate=UpperCAmelCase__  ,							)
                return {"sequences": outputs["sequences"]}






        UpperCAmelCase_            =  [[2, 0], [102, 103]]
        UpperCAmelCase_            =  [[1, 0], [1, 1]]
        UpperCAmelCase_            =  DummyModel(model=UpperCAmelCase__  )
        with tempfile.TemporaryDirectory() as tmp_dir:
              tf.saved_model.save(UpperCAmelCase__  ,							UpperCAmelCase__  ,							signatures={'''serving_default''': dummy_model.serving}  )
              UpperCAmelCase_            =  tf.saved_model.load(UpperCAmelCase__  ).signatures['''serving_default''']
              for batch_size in range(1  ,							len(UpperCAmelCase__  ) + 1  ):
                    UpperCAmelCase_            =  {
                        '''input_ids''': tf.constant(dummy_input_ids[:batch_size]  ),
                        '''attention_mask''': tf.constant(dummy_attention_masks[:batch_size]  ),
                    }
                    UpperCAmelCase_            =  serving_func(**UpperCAmelCase__  )['''sequences''']
                    UpperCAmelCase_            =  test_model.generate(**UpperCAmelCase__  ,							max_new_tokens=UpperCAmelCase__  )
                    tf.debugging.assert_equal(UpperCAmelCase__  ,							UpperCAmelCase__  )






  @slow
  def  lowerCAmelCase__ (      self			:       Tuple  )				->Optional[int]:
        # TF-only test: tf.saved_model export
        UpperCAmelCase_            =  TFAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''  )
        UpperCAmelCase_            =  1
        UpperCAmelCase_            =  2







        class lowerCamelCase							(	tf.Module					):




          '''simple docstring'''



          def __init__(      self			:       int  ,							UpperCAmelCase__			:       int  )				->List[str]:
                super(UpperCAmelCase__  ,							self  ).__init__()
                UpperCAmelCase_            =  model






          @tf.function(
              input_signature=(
                  tf.TensorSpec((batch_size, None)  ,							tf.intaa  ,							name='''input_ids'''  ),
                  tf.TensorSpec((batch_size, None)  ,							tf.intaa  ,							name='''attention_mask'''  ),
              )  ,							jit_compile=UpperCAmelCase__  ,							)
          def  lowerCAmelCase__ (      self			:       Dict  ,							UpperCAmelCase__			:       List[str]  ,							UpperCAmelCase__			:       Optional[Any]  )				->int:
                UpperCAmelCase_            =  self.model.generate(
                    input_ids=UpperCAmelCase__  ,							attention_mask=UpperCAmelCase__  ,							max_new_tokens=UpperCAmelCase__  ,							return_dict_in_generate=UpperCAmelCase__  ,							)
                return {"sequences": outputs["sequences"]}






        UpperCAmelCase_            =  [[2], [102, 103]]
        UpperCAmelCase_            =  [[1], [1, 1]]
        UpperCAmelCase_            =  DummyModel(model=UpperCAmelCase__  )
        with tempfile.TemporaryDirectory() as tmp_dir:
              tf.saved_model.save(UpperCAmelCase__  ,							UpperCAmelCase__  ,							signatures={'''serving_default''': dummy_model.serving}  )
              UpperCAmelCase_            =  tf.saved_model.load(UpperCAmelCase__  ).signatures['''serving_default''']
              for input_row in range(len(UpperCAmelCase__  )  ):
                    UpperCAmelCase_            =  {
                        '''input_ids''': tf.constant([dummy_input_ids[input_row]]  ),
                        '''attention_mask''': tf.constant([dummy_attention_masks[input_row]]  ),
                    }
                    UpperCAmelCase_            =  serving_func(**UpperCAmelCase__  )['''sequences''']
                    UpperCAmelCase_            =  test_model.generate(**UpperCAmelCase__  ,							max_new_tokens=UpperCAmelCase__  )
                    tf.debugging.assert_equal(UpperCAmelCase__  ,							UpperCAmelCase__  )






  @slow
  @require_tensorflow_text
  def  lowerCAmelCase__ (      self			:       Optional[Any]  )				->int:
        # TF-only test: tf.saved_model export
        with tempfile.TemporaryDirectory() as tmp_dir:
              # file needed to load the TF tokenizer
              hf_hub_download(repo_id='''google/flan-t5-small'''  ,							filename='''spiece.model'''  ,							local_dir=UpperCAmelCase__  )







              class lowerCamelCase							(	tf.keras.layers.Layer					):




                '''simple docstring'''



                def __init__(      self			:       List[str]  )				->Any:
                      super().__init__()
                      UpperCAmelCase_            =  text.SentencepieceTokenizer(
                          model=tf.io.gfile.GFile(os.path.join(UpperCAmelCase__  ,							'''spiece.model'''  )  ,							'''rb'''  ).read()  )
                      UpperCAmelCase_            =  TFAutoModelForSeqaSeqLM.from_pretrained('''hf-internal-testing/tiny-random-t5'''  )






                def  lowerCAmelCase__ (      self			:       Any  ,							UpperCAmelCase__			:       Union[str, Any]  ,							*UpperCAmelCase__			:       Union[str, Any]  ,							**UpperCAmelCase__			:       str  )				->List[str]:
                      UpperCAmelCase_            =  self.tokenizer.tokenize(UpperCAmelCase__  )
                      UpperCAmelCase_      ,		UpperCAmelCase_            =  text.pad_model_inputs(
                          UpperCAmelCase__  ,							max_seq_length=64  ,							pad_value=self.model.config.pad_token_id  )
                      UpperCAmelCase_            =  self.model.generate(input_ids=UpperCAmelCase__  ,							attention_mask=UpperCAmelCase__  )
                      return self.tokenizer.detokenize(UpperCAmelCase__  )






              UpperCAmelCase_            =  CompleteSentenceTransformer()
              UpperCAmelCase_            =  tf.keras.layers.Input(shape=(1,)  ,							dtype=tf.string  ,							name='''inputs'''  )
              UpperCAmelCase_            =  complete_model(UpperCAmelCase__  )
              UpperCAmelCase_            =  tf.keras.Model(UpperCAmelCase__  ,							UpperCAmelCase__  )
              keras_model.save(UpperCAmelCase__  )






  def  lowerCAmelCase__ (      self			:       Optional[int]  )				->Tuple:
        # Has PT equivalent: this test relies on random sampling
        UpperCAmelCase_            =  {
            '''do_sample''': True,
            '''num_beams''': 1,
            '''top_p''': 0.7,
            '''top_k''': 10,
            '''temperature''': 0.7,
        }
        UpperCAmelCase_            =  14

        UpperCAmelCase_            =  AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''  )
        UpperCAmelCase_            =  '''Hello, my dog is cute and'''
        UpperCAmelCase_            =  tokenizer(UpperCAmelCase__  ,							return_tensors='''tf'''  )
        UpperCAmelCase_            =  TFAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''  )

        UpperCAmelCase_            =  638
        # forces the generation to happen on CPU, to avoid GPU-related quirks
        with tf.device(''':/CPU:0'''  ):
              tf.random.set_seed(0  )
              UpperCAmelCase_            =  model.generate(**UpperCAmelCase__  ,							eos_token_id=UpperCAmelCase__  ,							**UpperCAmelCase__  )
        self.assertTrue(expectation == len(generated_tokens[0]  )  )

        UpperCAmelCase_            =  [638, 198]
        with tf.device(''':/CPU:0'''  ):
              tf.random.set_seed(0  )
              UpperCAmelCase_            =  model.generate(**UpperCAmelCase__  ,							eos_token_id=UpperCAmelCase__  ,							**UpperCAmelCase__  )
        self.assertTrue(expectation == len(generated_tokens[0]  )  )






  def  lowerCAmelCase__ (      self			:       Optional[int]  )				->List[str]:
        # Has PT equivalent: ample use of framework-specific code
        UpperCAmelCase_            =  AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-bart'''  )
        UpperCAmelCase_            =  '''Hugging Face is a technology company based in New York and Paris.'''
        UpperCAmelCase_            =  bart_tokenizer(UpperCAmelCase__  ,							return_tensors='''tf'''  ).input_ids
        UpperCAmelCase_            =  TFBartForConditionalGeneration.from_pretrained('''hf-internal-testing/tiny-random-bart'''  )
        UpperCAmelCase_            =  bart_model.generate(UpperCAmelCase__  ).numpy()







        class lowerCamelCase							(	lowerCamelCase					):




          '''simple docstring'''



          def  lowerCAmelCase__ (      self			:       int  ,							UpperCAmelCase__			:       Dict  ,							UpperCAmelCase__			:       List[Any]=None  ,							**UpperCAmelCase__			:       int  )				->List[str]:
                return super().call(UpperCAmelCase__  ,							**UpperCAmelCase__  )






        UpperCAmelCase_            =  FakeBart.from_pretrained('''hf-internal-testing/tiny-random-bart'''  )
        UpperCAmelCase_            =  bart_model.generate(UpperCAmelCase__  ,							foo='''bar'''  ).numpy()
        self.assertTrue(np.array_equal(UpperCAmelCase__  ,							UpperCAmelCase__  )  )







        class lowerCamelCase							(	bart_model.model.encoder.__class__					):




          '''simple docstring'''



          def  lowerCAmelCase__ (      self			:       str  ,							UpperCAmelCase__			:       List[str]  ,							**UpperCAmelCase__			:       Dict  )				->Any:
                return super().call(UpperCAmelCase__  ,							**UpperCAmelCase__  )






        UpperCAmelCase_            =  FakeEncoder(bart_model.config  ,							bart_model.model.shared  )
        UpperCAmelCase_            =  fake_encoder

        # Normal generation still works (the output will be different because the encoder weights are different)
        UpperCAmelCase_            =  bart_model.generate(UpperCAmelCase__  ).numpy()
        with self.assertRaises(UpperCAmelCase__  ):
              # FakeEncoder.call() accepts **kwargs -> no filtering -> value error due to unexpected input "foo"
              bart_model.generate(UpperCAmelCase__  ,							foo='''bar'''  )





