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






   '''simple docstring'''







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  '''| <pad> <unk> <s> </s> a b c d e f g h i j k'''.split()
      UpperCamelCase       =  dict(zip(lowerCamelCase__ ,   range(len(lowerCamelCase__   )   )   )   )

      UpperCamelCase       =  {
          '''unk_token''': '''<unk>''',
          '''bos_token''': '''<s>''',
          '''eos_token''': '''</s>''',
      }
      UpperCamelCase       =  {
          '''feature_size''': 1,
          '''padding_value''': 0.0,
          '''sampling_rate''': 1_6_0_0_0,
          '''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 ,   lowerCamelCase__   )
      with open(self.vocab_file ,   '''w''' ,   encoding='''utf-8'''   ) as fp:
         fp.write(json.dumps(lowerCamelCase__   ) + '''\n'''   )

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

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







   def       UpperCAmelCase				(			self ,   **lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase       =  self.add_kwargs_tokens_map.copy()
      kwargs.update(lowerCamelCase__   )
      return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname ,   **lowerCamelCase__   )







   def       UpperCAmelCase				(			self ,   **lowerCamelCase__   ):

      '''simple docstring'''







      return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname ,   **lowerCamelCase__   )







   def       UpperCAmelCase				(			self ,   **lowerCamelCase__   ):

      '''simple docstring'''







      return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name ,   **lowerCamelCase__   )







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      shutil.rmtree(self.tmpdirname   )







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self.get_tokenizer()
      UpperCamelCase       =  self.get_feature_extractor()
      UpperCamelCase       =  self.get_decoder()

      UpperCamelCase       =  WavaVecaProcessorWithLM(tokenizer=lowerCamelCase__ ,   feature_extractor=lowerCamelCase__ ,   decoder=lowerCamelCase__   )

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

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

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







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      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       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self.get_tokenizer()
      # add token to trigger raise
      tokenizer.add_tokens(['''xx''']   )
      with self.assertRaisesRegex(lowerCamelCase__ ,   '''include'''   ):
         WavaVecaProcessorWithLM(
             tokenizer=lowerCamelCase__ ,   feature_extractor=self.get_feature_extractor() ,   decoder=self.get_decoder()   )







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self.get_feature_extractor()
      UpperCamelCase       =  self.get_tokenizer()
      UpperCamelCase       =  self.get_decoder()

      UpperCamelCase       =  WavaVecaProcessorWithLM(tokenizer=lowerCamelCase__ ,   feature_extractor=lowerCamelCase__ ,   decoder=lowerCamelCase__   )

      UpperCamelCase       =  floats_list((3, 1_0_0_0)   )

      UpperCamelCase       =  feature_extractor(lowerCamelCase__ ,   return_tensors='''np'''   )
      UpperCamelCase       =  processor(lowerCamelCase__ ,   return_tensors='''np'''   )

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







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self.get_feature_extractor()
      UpperCamelCase       =  self.get_tokenizer()
      UpperCamelCase       =  self.get_decoder()

      UpperCamelCase       =  WavaVecaProcessorWithLM(tokenizer=lowerCamelCase__ ,   feature_extractor=lowerCamelCase__ ,   decoder=lowerCamelCase__   )

      UpperCamelCase       =  '''This is a test string'''

      UpperCamelCase       =  processor(text=lowerCamelCase__   )

      UpperCamelCase       =  tokenizer(lowerCamelCase__   )

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







   def       UpperCAmelCase				(			self ,   lowerCamelCase__=(2, 1_0, 1_6) ,   lowerCamelCase__=7_7   ):

      '''simple docstring'''







      np.random.seed(lowerCamelCase__   )
      return np.random.rand(*lowerCamelCase__   )







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self.get_feature_extractor()
      UpperCamelCase       =  self.get_tokenizer()
      UpperCamelCase       =  self.get_decoder()

      UpperCamelCase       =  WavaVecaProcessorWithLM(tokenizer=lowerCamelCase__ ,   feature_extractor=lowerCamelCase__ ,   decoder=lowerCamelCase__   )

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

      UpperCamelCase       =  processor.decode(lowerCamelCase__   )

      UpperCamelCase       =  decoder.decode_beams(lowerCamelCase__   )[0]

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







   @parameterized.expand([[None], ['''fork'''], ['''spawn''']]   )
   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase       =  self.get_feature_extractor()
      UpperCamelCase       =  self.get_tokenizer()
      UpperCamelCase       =  self.get_decoder()

      UpperCamelCase       =  WavaVecaProcessorWithLM(tokenizer=lowerCamelCase__ ,   feature_extractor=lowerCamelCase__ ,   decoder=lowerCamelCase__   )

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

      UpperCamelCase       =  list(lowerCamelCase__   )

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

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







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self.get_feature_extractor()
      UpperCamelCase       =  self.get_tokenizer()
      UpperCamelCase       =  self.get_decoder()

      UpperCamelCase       =  WavaVecaProcessorWithLM(tokenizer=lowerCamelCase__ ,   feature_extractor=lowerCamelCase__ ,   decoder=lowerCamelCase__   )

      UpperCamelCase       =  self._get_dummy_logits()

      UpperCamelCase       =  1_5
      UpperCamelCase       =  -20.0
      UpperCamelCase       =  -4.0

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

      UpperCamelCase       =  list(lowerCamelCase__   )

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

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

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

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







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self.get_feature_extractor()
      UpperCamelCase       =  self.get_tokenizer()
      UpperCamelCase       =  self.get_decoder()

      UpperCamelCase       =  WavaVecaProcessorWithLM(tokenizer=lowerCamelCase__ ,   feature_extractor=lowerCamelCase__ ,   decoder=lowerCamelCase__   )

      UpperCamelCase       =  self._get_dummy_logits()

      UpperCamelCase       =  2.0
      UpperCamelCase       =  5.0
      UpperCamelCase       =  -20.0
      UpperCamelCase       =  True

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

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

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

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

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







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







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







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  snapshot_download('''hf-internal-testing/processor_with_lm'''   )

      UpperCamelCase       =  WavaVecaProcessorWithLM.from_pretrained(lowerCamelCase__   )

      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(lowerCamelCase__   )
      UpperCamelCase       =  os.listdir(lowerCamelCase__   )

      local_decoder_files.sort()
      expected_decoder_files.sort()

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







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''   )
      UpperCamelCase       =  AutoProcessor.from_pretrained('''hf-internal-testing/processor_with_lm'''   )

      UpperCamelCase       =  floats_list((3, 1_0_0_0)   )

      UpperCamelCase       =  processor_wavaveca(lowerCamelCase__ ,   return_tensors='''np'''   )
      UpperCamelCase       =  processor_auto(lowerCamelCase__ ,   return_tensors='''np'''   )

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

      UpperCamelCase       =  self._get_dummy_logits()

      UpperCamelCase       =  processor_wavaveca.batch_decode(lowerCamelCase__   )
      UpperCamelCase       =  processor_auto.batch_decode(lowerCamelCase__   )

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







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self.get_feature_extractor()
      UpperCamelCase       =  self.get_tokenizer()
      UpperCamelCase       =  self.get_decoder()

      UpperCamelCase       =  WavaVecaProcessorWithLM(tokenizer=lowerCamelCase__ ,   feature_extractor=lowerCamelCase__ ,   decoder=lowerCamelCase__   )

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







   @staticmethod
   def       UpperCAmelCase				(			lowerCamelCase__ ,   lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase       =  [d[key] for d in offsets]
      return retrieved_list







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''   )
      UpperCamelCase       =  self._get_dummy_logits()[0]

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

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







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm'''   )
      UpperCamelCase       =  self._get_dummy_logits()

      UpperCamelCase       =  processor.batch_decode(lowerCamelCase__ ,   output_word_offsets=lowerCamelCase__   )

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

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







   @slow
   @require_torch
   @require_torchaudio
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      import torch

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

      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(lowerCamelCase__   ).logits.cpu().numpy()

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

      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(lowerCamelCase__ ,   '''word'''   )   ) ,   lowerCamelCase__   )
      self.assertEqual(''' '''.join(self.get_from_offsets(lowerCamelCase__ ,   '''word'''   )   ) ,   output.text   )

      # output times
      UpperCamelCase       =  torch.tensor(self.get_from_offsets(lowerCamelCase__ ,   '''start_time'''   )   )
      UpperCamelCase       =  torch.tensor(self.get_from_offsets(lowerCamelCase__ ,   '''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(lowerCamelCase__ ,   lowerCamelCase__ ,   atol=0.01   )   )
      self.assertTrue(torch.allclose(lowerCamelCase__ ,   lowerCamelCase__ ,   atol=0.01   )   )







'''simple docstring'''


from queue import Queue
from typing import TYPE_CHECKING, Optional


if TYPE_CHECKING:
   from ..models.auto import AutoTokenizer


class 						lowercase__	:






   '''simple docstring'''







   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







      raise NotImplementedError()







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      raise NotImplementedError()




class 						lowercase__	(  snake_case_   ):






   '''simple docstring'''







   def __init__(			self ,   lowerCamelCase__ ,   lowerCamelCase__ = False ,   **lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase       =  tokenizer
      UpperCamelCase       =  skip_prompt
      UpperCamelCase       =  decode_kwargs

      # variables used in the streaming process
      UpperCamelCase       =  []
      UpperCamelCase       =  0
      UpperCamelCase       =  True







   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







      if len(value.shape   ) > 1 and value.shape[0] > 1:
         raise ValueError('''TextStreamer only supports batch size 1'''   )
      elif len(value.shape   ) > 1:
         UpperCamelCase       =  value[0]

      if self.skip_prompt and self.next_tokens_are_prompt:
         UpperCamelCase       =  False
         return

      # Add the new token to the cache and decodes the entire thing.
      self.token_cache.extend(value.tolist()   )
      UpperCamelCase       =  self.tokenizer.decode(self.token_cache ,   **self.decode_kwargs   )

      # After the symbol for a new line, we flush the cache.
      if text.endswith('''\n'''   ):
         UpperCamelCase       =  text[self.print_len :]
         UpperCamelCase       =  []
         UpperCamelCase       =  0
      # If the last token is a CJK character, we print the characters.
      elif len(lowerCamelCase__   ) > 0 and self._is_chinese_char(ord(text[-1]   )   ):
         UpperCamelCase       =  text[self.print_len :]
         self.print_len += len(lowerCamelCase__   )
      # Otherwise, prints until the last space char (simple heuristic to avoid printing incomplete words,
      # which may change with the subsequent token -- there are probably smarter ways to do this!)
      else:
         UpperCamelCase       =  text[self.print_len : text.rfind(''' '''   ) + 1]
         self.print_len += len(lowerCamelCase__   )

      self.on_finalized_text(lowerCamelCase__   )







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      if len(self.token_cache   ) > 0:
         UpperCamelCase       =  self.tokenizer.decode(self.token_cache ,   **self.decode_kwargs   )
         UpperCamelCase       =  text[self.print_len :]
         UpperCamelCase       =  []
         UpperCamelCase       =  0
      else:
         UpperCamelCase       =  ''''''

      UpperCamelCase       =  True
      self.on_finalized_text(lowerCamelCase__ ,   stream_end=lowerCamelCase__   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ = False   ):

      '''simple docstring'''







      print(lowerCamelCase__ ,   flush=lowerCamelCase__ ,   end='''''' if not stream_end else None   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







      if (
          (cp >= 0x4_E_0_0 and cp <= 0x9_F_F_F)
          or (cp >= 0x3_4_0_0 and cp <= 0x4_D_B_F)  #
          or (cp >= 0x2_0_0_0_0 and cp <= 0x2_A_6_D_F)  #
          or (cp >= 0x2_A_7_0_0 and cp <= 0x2_B_7_3_F)  #
          or (cp >= 0x2_B_7_4_0 and cp <= 0x2_B_8_1_F)  #
          or (cp >= 0x2_B_8_2_0 and cp <= 0x2_C_E_A_F)  #
          or (cp >= 0xF_9_0_0 and cp <= 0xF_A_F_F)
          or (cp >= 0x2_F_8_0_0 and cp <= 0x2_F_A_1_F)  #
      ):  #
         return True

      return False




class 						lowercase__	(  snake_case_   ):






   '''simple docstring'''







   def __init__(			self ,   lowerCamelCase__ ,   lowerCamelCase__ = False ,   lowerCamelCase__ = None ,   **lowerCamelCase__   ):

      '''simple docstring'''







      super().__init__(lowerCamelCase__ ,   lowerCamelCase__ ,   **lowerCamelCase__   )
      UpperCamelCase       =  Queue()
      UpperCamelCase       =  None
      UpperCamelCase       =  timeout







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ = False   ):

      '''simple docstring'''







      self.text_queue.put(lowerCamelCase__ ,   timeout=self.timeout   )
      if stream_end:
         self.text_queue.put(self.stop_signal ,   timeout=self.timeout   )







   def __iter__(			self   ):

      '''simple docstring'''







      return self







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self.text_queue.get(timeout=self.timeout   )
      if value == self.stop_signal:
         raise StopIteration()
      else:
         return value





