


'''simple docstring'''
import argparse
from collections import OrderedDict
from pathlib import Path

import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from torchvision.transforms import functional as F

from transformers import DetrImageProcessor, TableTransformerConfig, TableTransformerForObjectDetection
from transformers.utils import logging


logging.set_verbosity_info()
UpperCamelCase__    					=      logging.get_logger(__name__)

# here we list all keys to be renamed (original name on the left, our name on the right)
UpperCamelCase__    					=      []
for i in range(6):
  # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
  rename_keys.append(
      (F"""transformer.encoder.layers.{i}.self_attn.out_proj.weight""", F"""encoder.layers.{i}.self_attn.out_proj.weight""")
  )
  rename_keys.append(
      (F"""transformer.encoder.layers.{i}.self_attn.out_proj.bias""", F"""encoder.layers.{i}.self_attn.out_proj.bias""")
  )
  rename_keys.append((F"""transformer.encoder.layers.{i}.linear1.weight""", F"""encoder.layers.{i}.fc1.weight"""))
  rename_keys.append((F"""transformer.encoder.layers.{i}.linear1.bias""", F"""encoder.layers.{i}.fc1.bias"""))
  rename_keys.append((F"""transformer.encoder.layers.{i}.linear2.weight""", F"""encoder.layers.{i}.fc2.weight"""))
  rename_keys.append((F"""transformer.encoder.layers.{i}.linear2.bias""", F"""encoder.layers.{i}.fc2.bias"""))
  rename_keys.append(
      (F"""transformer.encoder.layers.{i}.norm1.weight""", F"""encoder.layers.{i}.self_attn_layer_norm.weight""")
  )
  rename_keys.append((F"""transformer.encoder.layers.{i}.norm1.bias""", F"""encoder.layers.{i}.self_attn_layer_norm.bias"""))
  rename_keys.append((F"""transformer.encoder.layers.{i}.norm2.weight""", F"""encoder.layers.{i}.final_layer_norm.weight"""))
  rename_keys.append((F"""transformer.encoder.layers.{i}.norm2.bias""", F"""encoder.layers.{i}.final_layer_norm.bias"""))
  # decoder layers: 2 times output projection, 2 feedforward neural networks and 3 layernorms
  rename_keys.append(
      (F"""transformer.decoder.layers.{i}.self_attn.out_proj.weight""", F"""decoder.layers.{i}.self_attn.out_proj.weight""")
  )
  rename_keys.append(
      (F"""transformer.decoder.layers.{i}.self_attn.out_proj.bias""", F"""decoder.layers.{i}.self_attn.out_proj.bias""")
  )
  rename_keys.append(
      (
          F"""transformer.decoder.layers.{i}.multihead_attn.out_proj.weight""",
          F"""decoder.layers.{i}.encoder_attn.out_proj.weight""",
      )
  )
  rename_keys.append(
      (
          F"""transformer.decoder.layers.{i}.multihead_attn.out_proj.bias""",
          F"""decoder.layers.{i}.encoder_attn.out_proj.bias""",
      )
  )
  rename_keys.append((F"""transformer.decoder.layers.{i}.linear1.weight""", F"""decoder.layers.{i}.fc1.weight"""))
  rename_keys.append((F"""transformer.decoder.layers.{i}.linear1.bias""", F"""decoder.layers.{i}.fc1.bias"""))
  rename_keys.append((F"""transformer.decoder.layers.{i}.linear2.weight""", F"""decoder.layers.{i}.fc2.weight"""))
  rename_keys.append((F"""transformer.decoder.layers.{i}.linear2.bias""", F"""decoder.layers.{i}.fc2.bias"""))
  rename_keys.append(
      (F"""transformer.decoder.layers.{i}.norm1.weight""", F"""decoder.layers.{i}.self_attn_layer_norm.weight""")
  )
  rename_keys.append((F"""transformer.decoder.layers.{i}.norm1.bias""", F"""decoder.layers.{i}.self_attn_layer_norm.bias"""))
  rename_keys.append(
      (F"""transformer.decoder.layers.{i}.norm2.weight""", F"""decoder.layers.{i}.encoder_attn_layer_norm.weight""")
  )
  rename_keys.append(
      (F"""transformer.decoder.layers.{i}.norm2.bias""", F"""decoder.layers.{i}.encoder_attn_layer_norm.bias""")
  )
  rename_keys.append((F"""transformer.decoder.layers.{i}.norm3.weight""", F"""decoder.layers.{i}.final_layer_norm.weight"""))
  rename_keys.append((F"""transformer.decoder.layers.{i}.norm3.bias""", F"""decoder.layers.{i}.final_layer_norm.bias"""))

# convolutional projection + query embeddings + layernorm of encoder + layernorm of decoder + class and bounding box heads
rename_keys.extend(
    [
        ('''input_proj.weight''', '''input_projection.weight'''),
        ('''input_proj.bias''', '''input_projection.bias'''),
        ('''query_embed.weight''', '''query_position_embeddings.weight'''),
        ('''transformer.encoder.norm.weight''', '''encoder.layernorm.weight'''),
        ('''transformer.encoder.norm.bias''', '''encoder.layernorm.bias'''),
        ('''transformer.decoder.norm.weight''', '''decoder.layernorm.weight'''),
        ('''transformer.decoder.norm.bias''', '''decoder.layernorm.bias'''),
        ('''class_embed.weight''', '''class_labels_classifier.weight'''),
        ('''class_embed.bias''', '''class_labels_classifier.bias'''),
        ('''bbox_embed.layers.0.weight''', '''bbox_predictor.layers.0.weight'''),
        ('''bbox_embed.layers.0.bias''', '''bbox_predictor.layers.0.bias'''),
        ('''bbox_embed.layers.1.weight''', '''bbox_predictor.layers.1.weight'''),
        ('''bbox_embed.layers.1.bias''', '''bbox_predictor.layers.1.bias'''),
        ('''bbox_embed.layers.2.weight''', '''bbox_predictor.layers.2.weight'''),
        ('''bbox_embed.layers.2.bias''', '''bbox_predictor.layers.2.bias'''),
    ]
)





def 					a__       (					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__						)    ->	Union[str, Any]:
    UpperCAmelCase__			:		Optional[Any]													=				state_dict.pop(lowerCAmelCase__						)
    UpperCAmelCase__			:		Optional[Any]													=				val





def 					a__       (					lowerCAmelCase__						)    ->	Optional[Any]:
    UpperCAmelCase__			:		List[str]													=				OrderedDict()
    for key, value in state_dict.items():
        if "backbone.0.body" in key:
            UpperCAmelCase__			:		Dict													=				key.replace('''backbone.0.body'''  ,					'''backbone.conv_encoder.model'''						)
            UpperCAmelCase__			:		Tuple													=				value
        else:
            UpperCAmelCase__			:		str													=				value

    return new_state_dict





def 					a__       (					lowerCAmelCase__						)    ->	Optional[Any]:
    UpperCAmelCase__			:		Tuple													=				''''''

    # first: transformer encoder
    for i in range(6						):
        # read in weights + bias of input projection layer (in PyTorch's MultiHeadAttention, this is a single matrix + bias)
        UpperCAmelCase__			:		Optional[Any]													=				state_dict.pop(F"""{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_weight"""						)
        UpperCAmelCase__			:		Dict													=				state_dict.pop(F"""{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_bias"""						)
        # next, add query, keys and values (in that order) to the state dict
        UpperCAmelCase__			:		Any													=				in_proj_weight[:2_56, :]
        UpperCAmelCase__			:		Optional[int]													=				in_proj_bias[:2_56]
        UpperCAmelCase__			:		List[str]													=				in_proj_weight[2_56:5_12, :]
        UpperCAmelCase__			:		Optional[int]													=				in_proj_bias[2_56:5_12]
        UpperCAmelCase__			:		Tuple													=				in_proj_weight[-2_56:, :]
        UpperCAmelCase__			:		str													=				in_proj_bias[-2_56:]
    # next: transformer decoder (which is a bit more complex because it also includes cross-attention)
    for i in range(6						):
        # read in weights + bias of input projection layer of self-attention
        UpperCAmelCase__			:		List[Any]													=				state_dict.pop(F"""{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_weight"""						)
        UpperCAmelCase__			:		Tuple													=				state_dict.pop(F"""{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_bias"""						)
        # next, add query, keys and values (in that order) to the state dict
        UpperCAmelCase__			:		Tuple													=				in_proj_weight[:2_56, :]
        UpperCAmelCase__			:		Optional[Any]													=				in_proj_bias[:2_56]
        UpperCAmelCase__			:		List[str]													=				in_proj_weight[2_56:5_12, :]
        UpperCAmelCase__			:		str													=				in_proj_bias[2_56:5_12]
        UpperCAmelCase__			:		Optional[Any]													=				in_proj_weight[-2_56:, :]
        UpperCAmelCase__			:		Union[str, Any]													=				in_proj_bias[-2_56:]
        # read in weights + bias of input projection layer of cross-attention
        UpperCAmelCase__			:		Optional[int]													=				state_dict.pop(
            F"""{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_weight"""						)
        UpperCAmelCase__			:		List[Any]													=				state_dict.pop(F"""{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_bias"""						)
        # next, add query, keys and values (in that order) of cross-attention to the state dict
        UpperCAmelCase__			:		Any													=				in_proj_weight_cross_attn[:2_56, :]
        UpperCAmelCase__			:		Union[str, Any]													=				in_proj_bias_cross_attn[:2_56]
        UpperCAmelCase__			:		Union[str, Any]													=				in_proj_weight_cross_attn[2_56:5_12, :]
        UpperCAmelCase__			:		Optional[int]													=				in_proj_bias_cross_attn[2_56:5_12]
        UpperCAmelCase__			:		Any													=				in_proj_weight_cross_attn[-2_56:, :]
        UpperCAmelCase__			:		Union[str, Any]													=				in_proj_bias_cross_attn[-2_56:]





def 					a__       (					lowerCAmelCase__  ,					lowerCAmelCase__						)    ->	int:
    UpperCAmelCase__						,						UpperCAmelCase__			:		Any													=				image.size
    UpperCAmelCase__			:		Any													=				max(lowerCAmelCase__  ,					lowerCAmelCase__						)
    UpperCAmelCase__			:		Optional[Any]													=				8_00 if '''detection''' in checkpoint_url else 10_00
    UpperCAmelCase__			:		Dict													=				target_max_size / current_max_size
    UpperCAmelCase__			:		Dict													=				image.resize((int(round(scale * width						)						), int(round(scale * height						)						))						)

    return resized_image





def 					a__       (					lowerCAmelCase__						)    ->	Dict:
    UpperCAmelCase__			:		List[Any]													=				F.to_tensor(lowerCAmelCase__						)
    UpperCAmelCase__			:		List[Any]													=				F.normalize(lowerCAmelCase__  ,					mean=[0.4_8_5, 0.4_5_6, 0.4_0_6]  ,					std=[0.2_2_9, 0.2_2_4, 0.2_2_5]						)
    return image





@torch.no_grad()
def 					a__       (					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__						)    ->	Optional[int]:
    logger.info('''Converting model...'''						)

    # load original state dict
    UpperCAmelCase__			:		Tuple													=				torch.hub.load_state_dict_from_url(lowerCAmelCase__  ,					map_location='''cpu'''						)
    # rename keys
    for src, dest in rename_keys:
        rename_key(lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__						)
    UpperCAmelCase__			:		Union[str, Any]													=				rename_backbone_keys(lowerCAmelCase__						)
    # query, key and value matrices need special treatment
    read_in_q_k_v(lowerCAmelCase__						)
    # important: we need to prepend a prefix to each of the base model keys as the head models use different attributes for them
    UpperCAmelCase__			:		Optional[int]													=				'''model.'''
    for key in state_dict.copy().keys():
        if not key.startswith('''class_labels_classifier'''						) and not key.startswith('''bbox_predictor'''						):
            UpperCAmelCase__			:		Optional[int]													=				state_dict.pop(lowerCAmelCase__						)
            UpperCAmelCase__			:		Dict													=				val
    # create HuggingFace model and load state dict
    UpperCAmelCase__			:		List[Any]													=				TableTransformerConfig(
        backbone='''resnet18'''  ,					mask_loss_coefficient=1  ,					dice_loss_coefficient=1  ,					ce_loss_coefficient=1  ,					bbox_loss_coefficient=5  ,					giou_loss_coefficient=2  ,					eos_coefficient=0.4  ,					class_cost=1  ,					bbox_cost=5  ,					giou_cost=2  ,					)

    if "detection" in checkpoint_url:
        UpperCAmelCase__			:		List[Any]													=				15
        UpperCAmelCase__			:		Optional[int]													=				2
        UpperCAmelCase__			:		Optional[int]													=				{0: '''table''', 1: '''table rotated'''}
        UpperCAmelCase__			:		Any													=				idalabel
        UpperCAmelCase__			:		Union[str, Any]													=				{v: k for k, v in idalabel.items()}
    else:
        UpperCAmelCase__			:		int													=				1_25
        UpperCAmelCase__			:		List[str]													=				6
        UpperCAmelCase__			:		List[Any]													=				{
            0: '''table''',
            1: '''table column''',
            2: '''table row''',
            3: '''table column header''',
            4: '''table projected row header''',
            5: '''table spanning cell''',
        }
        UpperCAmelCase__			:		Optional[Any]													=				idalabel
        UpperCAmelCase__			:		str													=				{v: k for k, v in idalabel.items()}

    UpperCAmelCase__			:		List[Any]													=				DetrImageProcessor(
        format='''coco_detection'''  ,					max_size=8_00 if '''detection''' in checkpoint_url else 10_00						)
    UpperCAmelCase__			:		Any													=				TableTransformerForObjectDetection(lowerCAmelCase__						)
    model.load_state_dict(lowerCAmelCase__						)
    model.eval()

    # verify our conversion
    UpperCAmelCase__			:		Optional[Any]													=				'''example_pdf.png''' if '''detection''' in checkpoint_url else '''example_table.png'''
    UpperCAmelCase__			:		Dict													=				hf_hub_download(repo_id='''nielsr/example-pdf'''  ,					repo_type='''dataset'''  ,					filename=lowerCAmelCase__						)
    UpperCAmelCase__			:		Any													=				Image.open(lowerCAmelCase__						).convert('''RGB'''						)
    UpperCAmelCase__			:		Tuple													=				normalize(resize(lowerCAmelCase__  ,					lowerCAmelCase__						)						).unsqueeze(0						)

    UpperCAmelCase__			:		Optional[int]													=				model(lowerCAmelCase__						)

    if "detection" in checkpoint_url:
        UpperCAmelCase__			:		Any													=				(1, 15, 3)
        UpperCAmelCase__			:		Union[str, Any]													=				torch.tensor(
            [[-6.7_8_9_7, -1_6.9_9_8_5, 6.7_9_3_7], [-8.0_1_8_6, -2_2.2_1_9_2, 6.9_6_7_7], [-7.3_1_1_7, -2_1.0_7_0_8, 7.4_0_5_5]]						)
        UpperCAmelCase__			:		Union[str, Any]													=				torch.tensor([[0.4_8_6_7, 0.1_7_6_7, 0.6_7_3_2], [0.6_7_1_8, 0.4_4_7_9, 0.3_8_3_0], [0.4_7_1_6, 0.1_7_6_0, 0.6_3_6_4]]						)

    else:
        UpperCAmelCase__			:		List[str]													=				(1, 1_25, 7)
        UpperCAmelCase__			:		Any													=				torch.tensor(
            [[-1_8.1_4_3_0, -8.3_2_1_4, 4.8_2_7_4], [-1_8.4_6_8_5, -7.1_3_6_1, -4.2_6_6_7], [-2_6.3_6_9_3, -9.3_4_2_9, -4.9_9_6_2]]						)
        UpperCAmelCase__			:		List[str]													=				torch.tensor([[0.4_9_8_3, 0.5_5_9_5, 0.9_4_4_0], [0.4_9_1_6, 0.6_3_1_5, 0.5_9_5_4], [0.6_1_0_8, 0.8_6_3_7, 0.1_1_3_5]]						)

    assert outputs.logits.shape == expected_shape
    assert torch.allclose(outputs.logits[0, :3, :3]  ,					lowerCAmelCase__  ,					atol=1E-4						)
    assert torch.allclose(outputs.pred_boxes[0, :3, :3]  ,					lowerCAmelCase__  ,					atol=1E-4						)
    print('''Looks ok!'''						)

    if pytorch_dump_folder_path is not None:
        # Save model and image processor
        logger.info(F"""Saving PyTorch model and image processor to {pytorch_dump_folder_path}..."""						)
        Path(lowerCAmelCase__						).mkdir(exist_ok=lowerCAmelCase__						)
        model.save_pretrained(lowerCAmelCase__						)
        image_processor.save_pretrained(lowerCAmelCase__						)

    if push_to_hub:
        # Push model to HF hub
        logger.info('''Pushing model to the hub...'''						)
        UpperCAmelCase__			:		Any													=				(
            '''microsoft/table-transformer-detection'''
            if '''detection''' in checkpoint_url
            else '''microsoft/table-transformer-structure-recognition'''
        )
        model.push_to_hub(lowerCAmelCase__						)
        image_processor.push_to_hub(lowerCAmelCase__						)


if __name__ == "__main__":
  UpperCamelCase__    					=      argparse.ArgumentParser()

  parser.add_argument(
      '''--checkpoint_url''',
      default='''https://pubtables1m.blob.core.windows.net/model/pubtables1m_detection_detr_r18.pth''',
      type=str,
      choices=[
          '''https://pubtables1m.blob.core.windows.net/model/pubtables1m_detection_detr_r18.pth''',
          '''https://pubtables1m.blob.core.windows.net/model/pubtables1m_structure_detr_r18.pth''',
      ],
      help='''URL of the Table Transformer checkpoint you\'d like to convert.''',
  )
  parser.add_argument(
      '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the folder to output PyTorch model.'''
  )
  parser.add_argument(
      '''--push_to_hub''', action='''store_true''', help='''Whether or not to push the converted model to the 🤗 hub.'''
  )
  UpperCamelCase__    					=      parser.parse_args()
  convert_table_transformer_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub)




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






      def lowercase_							(					self				:    Any ):


          '''simple docstring'''


          UpperCAmelCase__			:		List[str]													=				'''| <pad> <unk> <s> </s> a b c d e f g h i j k'''.split()
          UpperCAmelCase__			:		str													=				dict(zip(_A						,						range(len(_A ) ) ) )

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

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

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

          # load decoder from hub
          UpperCAmelCase__			:		Tuple													=				'''hf-internal-testing/ngram-beam-search-decoder'''






      def lowercase_							(					self				:    Union[str, Any]						,						**_A				:    Tuple ):


          '''simple docstring'''


          UpperCAmelCase__			:		List[str]													=				self.add_kwargs_tokens_map.copy()
          kwargs.update(_A )
          return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname						,						**_A )






      def lowercase_							(					self				:    Dict						,						**_A				:    str ):


          '''simple docstring'''


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






      def lowercase_							(					self				:    Tuple						,						**_A				:    int ):


          '''simple docstring'''


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






      def lowercase_							(					self				:    Optional[Any] ):


          '''simple docstring'''


          shutil.rmtree(self.tmpdirname )






      def lowercase_							(					self				:    Any ):


          '''simple docstring'''


          UpperCAmelCase__			:		Optional[int]													=				self.get_tokenizer()
          UpperCAmelCase__			:		Any													=				self.get_feature_extractor()
          UpperCAmelCase__			:		Dict													=				self.get_decoder()

          UpperCAmelCase__			:		Dict													=				WavaVecaProcessorWithLM(tokenizer=_A						,						feature_extractor=_A						,						decoder=_A )

          processor.save_pretrained(self.tmpdirname )
          UpperCAmelCase__			:		Optional[Any]													=				WavaVecaProcessorWithLM.from_pretrained(self.tmpdirname )

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

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

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






      def lowercase_							(					self				:    Tuple ):


          '''simple docstring'''


          UpperCAmelCase__			:		Optional[int]													=				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__			:		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 lowercase_							(					self				:    str ):


          '''simple docstring'''


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






      def lowercase_							(					self				:    Any ):


          '''simple docstring'''


          UpperCAmelCase__			:		List[str]													=				self.get_feature_extractor()
          UpperCAmelCase__			:		List[Any]													=				self.get_tokenizer()
          UpperCAmelCase__			:		int													=				self.get_decoder()

          UpperCAmelCase__			:		Optional[int]													=				WavaVecaProcessorWithLM(tokenizer=_A						,						feature_extractor=_A						,						decoder=_A )

          UpperCAmelCase__			:		str													=				floats_list((3, 1_000) )

          UpperCAmelCase__			:		str													=				feature_extractor(_A						,						return_tensors='''np''' )
          UpperCAmelCase__			:		Optional[Any]													=				processor(_A						,						return_tensors='''np''' )

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






      def lowercase_							(					self				:    Optional[Any] ):


          '''simple docstring'''


          UpperCAmelCase__			:		Union[str, Any]													=				self.get_feature_extractor()
          UpperCAmelCase__			:		Any													=				self.get_tokenizer()
          UpperCAmelCase__			:		Any													=				self.get_decoder()

          UpperCAmelCase__			:		str													=				WavaVecaProcessorWithLM(tokenizer=_A						,						feature_extractor=_A						,						decoder=_A )

          UpperCAmelCase__			:		Union[str, Any]													=				'''This is a test string'''

          UpperCAmelCase__			:		Union[str, Any]													=				processor(text=_A )

          UpperCAmelCase__			:		Dict													=				tokenizer(_A )

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






      def lowercase_							(					self				:    str						,						_A				:    int=(2, 10, 16)						,						_A				:    Optional[int]=77 ):


          '''simple docstring'''


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






      def lowercase_							(					self				:    Optional[Any] ):


          '''simple docstring'''


          UpperCAmelCase__			:		Dict													=				self.get_feature_extractor()
          UpperCAmelCase__			:		Dict													=				self.get_tokenizer()
          UpperCAmelCase__			:		Tuple													=				self.get_decoder()

          UpperCAmelCase__			:		Optional[int]													=				WavaVecaProcessorWithLM(tokenizer=_A						,						feature_extractor=_A						,						decoder=_A )

          UpperCAmelCase__			:		Tuple													=				self._get_dummy_logits(shape=(10, 16)						,						seed=13 )

          UpperCAmelCase__			:		Union[str, Any]													=				processor.decode(_A )

          UpperCAmelCase__			:		Union[str, Any]													=				decoder.decode_beams(_A )[0]

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






      @parameterized.expand([[None], ['''fork'''], ['''spawn''']] )
      def lowercase_							(					self				:    Dict						,						_A				:    str ):


          '''simple docstring'''


          UpperCAmelCase__			:		Tuple													=				self.get_feature_extractor()
          UpperCAmelCase__			:		Union[str, Any]													=				self.get_tokenizer()
          UpperCAmelCase__			:		Union[str, Any]													=				self.get_decoder()

          UpperCAmelCase__			:		Optional[int]													=				WavaVecaProcessorWithLM(tokenizer=_A						,						feature_extractor=_A						,						decoder=_A )

          UpperCAmelCase__			:		List[str]													=				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__			:		Union[str, Any]													=				processor.batch_decode(_A )
          else:
              with get_context(_A ).Pool() as pool:
                  UpperCAmelCase__			:		List[Any]													=				processor.batch_decode(_A						,						_A )

          UpperCAmelCase__			:		List[str]													=				list(_A )

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

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

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






      def lowercase_							(					self				:    Tuple ):


          '''simple docstring'''


          UpperCAmelCase__			:		Any													=				self.get_feature_extractor()
          UpperCAmelCase__			:		str													=				self.get_tokenizer()
          UpperCAmelCase__			:		Optional[int]													=				self.get_decoder()

          UpperCAmelCase__			:		Optional[int]													=				WavaVecaProcessorWithLM(tokenizer=_A						,						feature_extractor=_A						,						decoder=_A )

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

          UpperCAmelCase__			:		Optional[Any]													=				15
          UpperCAmelCase__			:		Union[str, Any]													=				-2_0.0
          UpperCAmelCase__			:		List[str]													=				-4.0

          UpperCAmelCase__			:		str													=				processor.batch_decode(
              _A						,						beam_width=_A						,						beam_prune_logp=_A						,						token_min_logp=_A						,						)
          UpperCAmelCase__			:		Union[str, Any]													=				decoded_processor_out.text

          UpperCAmelCase__			:		Tuple													=				list(_A )

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

          UpperCAmelCase__			:		Optional[Any]													=				[d[0][0] for d in decoded_decoder_out]
          UpperCAmelCase__			:		Any													=				[d[0][2] for d in decoded_decoder_out]
          UpperCAmelCase__			:		Tuple													=				[d[0][3] for d in decoded_decoder_out]

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

          self.assertTrue(np.array_equal(_A						,						decoded_processor_out.logit_score ) )
          self.assertTrue(np.allclose([-2_0.0_5_4, -1_8.4_4_7]						,						_A						,						atol=1e-3 ) )

          self.assertTrue(np.array_equal(_A						,						decoded_processor_out.lm_score ) )
          self.assertTrue(np.allclose([-1_5.5_5_4, -1_3.9_4_7_4]						,						_A						,						atol=1e-3 ) )






      def lowercase_							(					self				:    Optional[Any] ):


          '''simple docstring'''


          UpperCAmelCase__			:		int													=				self.get_feature_extractor()
          UpperCAmelCase__			:		str													=				self.get_tokenizer()
          UpperCAmelCase__			:		List[str]													=				self.get_decoder()

          UpperCAmelCase__			:		Optional[int]													=				WavaVecaProcessorWithLM(tokenizer=_A						,						feature_extractor=_A						,						decoder=_A )

          UpperCAmelCase__			:		Optional[Any]													=				self._get_dummy_logits()

          UpperCAmelCase__			:		List[Any]													=				2.0
          UpperCAmelCase__			:		Union[str, Any]													=				5.0
          UpperCAmelCase__			:		Any													=				-2_0.0
          UpperCAmelCase__			:		List[Any]													=				True

          UpperCAmelCase__			:		Any													=				processor.batch_decode(
              _A						,						alpha=_A						,						beta=_A						,						unk_score_offset=_A						,						lm_score_boundary=_A						,						)
          UpperCAmelCase__			:		List[Any]													=				decoded_processor_out.text

          UpperCAmelCase__			:		List[Any]													=				list(_A )
          decoder.reset_params(
              alpha=_A						,						beta=_A						,						unk_score_offset=_A						,						lm_score_boundary=_A						,						)

          with get_context('''fork''' ).Pool() as pool:
              UpperCAmelCase__			:		Any													=				decoder.decode_beams_batch(
                  _A						,						_A						,						)

          UpperCAmelCase__			:		Union[str, Any]													=				[d[0][0] for d in decoded_decoder_out]

          self.assertListEqual(_A						,						_A )
          self.assertListEqual(['''<s> </s> <s> </s> </s>''', '''</s> </s> <s> </s> </s>''']						,						_A )
          UpperCAmelCase__			:		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						,						-2_0.0 )
          self.assertEqual(lm_model.score_boundary						,						_A )






      def lowercase_							(					self				:    Any ):


          '''simple docstring'''


          UpperCAmelCase__			:		Optional[Any]													=				WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm''' )

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

          UpperCAmelCase__			:		Dict													=				os.listdir(_A )
          UpperCAmelCase__			:		str													=				['''alphabet.json''', '''language_model''']

          downloaded_decoder_files.sort()
          expected_decoder_files.sort()

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






      def lowercase_							(					self				:    List[Any] ):


          '''simple docstring'''


          UpperCAmelCase__			:		str													=				snapshot_download('''hf-internal-testing/processor_with_lm''' )

          UpperCAmelCase__			:		Tuple													=				WavaVecaProcessorWithLM.from_pretrained(_A )

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

          UpperCAmelCase__			:		Optional[Any]													=				os.listdir(_A )
          UpperCAmelCase__			:		int													=				os.listdir(_A )

          local_decoder_files.sort()
          expected_decoder_files.sort()

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






      def lowercase_							(					self				:    Dict ):


          '''simple docstring'''


          UpperCAmelCase__			:		Optional[int]													=				WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm''' )
          UpperCAmelCase__			:		Any													=				AutoProcessor.from_pretrained('''hf-internal-testing/processor_with_lm''' )

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

          UpperCAmelCase__			:		Tuple													=				processor_wavaveca(_A						,						return_tensors='''np''' )
          UpperCAmelCase__			:		Union[str, Any]													=				processor_auto(_A						,						return_tensors='''np''' )

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

          UpperCAmelCase__			:		str													=				self._get_dummy_logits()

          UpperCAmelCase__			:		List[Any]													=				processor_wavaveca.batch_decode(_A )
          UpperCAmelCase__			:		List[Any]													=				processor_auto.batch_decode(_A )

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






      def lowercase_							(					self				:    Union[str, Any] ):


          '''simple docstring'''


          UpperCAmelCase__			:		Dict													=				self.get_feature_extractor()
          UpperCAmelCase__			:		Any													=				self.get_tokenizer()
          UpperCAmelCase__			:		List[str]													=				self.get_decoder()

          UpperCAmelCase__			:		Dict													=				WavaVecaProcessorWithLM(tokenizer=_A						,						feature_extractor=_A						,						decoder=_A )

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






      @staticmethod
      def lowercase_							(					_A				:    Tuple						,						_A				:    Tuple ):


          '''simple docstring'''


          UpperCAmelCase__			:		Optional[int]													=				[d[key] for d in offsets]
          return retrieved_list






      def lowercase_							(					self				:    Optional[Any] ):


          '''simple docstring'''


          UpperCAmelCase__			:		Dict													=				WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm''' )
          UpperCAmelCase__			:		Tuple													=				self._get_dummy_logits()[0]

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

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






      def lowercase_							(					self				:    Optional[Any] ):


          '''simple docstring'''


          UpperCAmelCase__			:		List[str]													=				WavaVecaProcessorWithLM.from_pretrained('''hf-internal-testing/processor_with_lm''' )
          UpperCAmelCase__			:		List[str]													=				self._get_dummy_logits()

          UpperCAmelCase__			:		Dict													=				processor.batch_decode(_A						,						output_word_offsets=_A )

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

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






      @slow
      @require_torch
      @require_torchaudio
      def lowercase_							(					self				:    str ):


          '''simple docstring'''


          import torch

          UpperCAmelCase__			:		Dict													=				load_dataset('''common_voice'''						,						'''en'''						,						split='''train'''						,						streaming=_A )
          UpperCAmelCase__			:		Optional[Any]													=				ds.cast_column('''audio'''						,						datasets.Audio(sampling_rate=16_000 ) )
          UpperCAmelCase__			:		Any													=				iter(_A )
          UpperCAmelCase__			:		Dict													=				next(_A )

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

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

          UpperCAmelCase__			:		List[str]													=				processor.decode(logits[0]						,						output_word_offsets=_A )

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

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

          # output times
          UpperCAmelCase__			:		List[Any]													=				torch.tensor(self.get_from_offsets(_A						,						'''start_time''' ) )
          UpperCAmelCase__			:		Dict													=				torch.tensor(self.get_from_offsets(_A						,						'''end_time''' ) )

          # fmt: off
          UpperCAmelCase__			:		Any													=				torch.tensor([1.4_1_9_9, 1.6_5_9_9, 2.2_5_9_9, 3.0, 3.2_4, 3.5_9_9_9, 3.7_9_9_9, 4.0_9_9_9, 4.2_6, 4.9_4, 5.2_8, 5.6_5_9_9, 5.7_8, 5.9_4, 6.3_2, 6.5_3_9_9, 6.6_5_9_9] )
          UpperCAmelCase__			:		Any													=				torch.tensor([1.5_3_9_9, 1.8_9_9_9, 2.9, 3.1_6, 3.5_3_9_9, 3.7_2, 4.0_1_9_9, 4.1_7_9_9, 4.7_6, 5.1_5_9_9, 5.5_5_9_9, 5.6_9_9_9, 5.8_6, 6.1_9_9_9, 6.3_8, 6.6_1_9_9, 6.9_4] )
          # fmt: on

          self.assertTrue(torch.allclose(_A						,						_A						,						atol=0.0_1 ) )
          self.assertTrue(torch.allclose(_A						,						_A						,						atol=0.0_1 ) )
