






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()
__magic_name__   :  Dict        =							logging.get_logger(__name__)

# here we list all keys to be renamed (original name on the left, our name on the right)
__magic_name__   :  List[Any]        =							[]
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 		lowercase__  (				_UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase) ->			Union[str, Any]:
      """simple docstring"""

      UpperCamelCase					       =  state_dict.pop(_UpperCamelCase)
      UpperCamelCase					       =  val
def 		lowercase__  (				_UpperCamelCase) ->			Dict:
      """simple docstring"""

      UpperCamelCase					       =  OrderedDict()
      for key, value in state_dict.items():
            if "backbone.0.body" in key:
                  UpperCamelCase					       =  key.replace('backbone.0.body'       ,							'backbone.conv_encoder.model')
                  UpperCamelCase					       =  value
            else:
                  UpperCamelCase					       =  value

      return new_state_dict
def 		lowercase__  (				_UpperCamelCase) ->			List[str]:
      """simple docstring"""

      UpperCamelCase					       =  ''

      # 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					       =  state_dict.pop(F'{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_weight')
            UpperCamelCase					       =  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					       =  in_proj_weight[:2_56, :]
            UpperCamelCase					       =  in_proj_bias[:2_56]
            UpperCamelCase					       =  in_proj_weight[2_56:5_12, :]
            UpperCamelCase					       =  in_proj_bias[2_56:5_12]
            UpperCamelCase					       =  in_proj_weight[-2_56:, :]
            UpperCamelCase					       =  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					       =  state_dict.pop(F'{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_weight')
            UpperCamelCase					       =  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					       =  in_proj_weight[:2_56, :]
            UpperCamelCase					       =  in_proj_bias[:2_56]
            UpperCamelCase					       =  in_proj_weight[2_56:5_12, :]
            UpperCamelCase					       =  in_proj_bias[2_56:5_12]
            UpperCamelCase					       =  in_proj_weight[-2_56:, :]
            UpperCamelCase					       =  in_proj_bias[-2_56:]
            # read in weights + bias of input projection layer of cross-attention
            UpperCamelCase					       =  state_dict.pop(
                F'{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_weight')
            UpperCamelCase					       =  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					       =  in_proj_weight_cross_attn[:2_56, :]
            UpperCamelCase					       =  in_proj_bias_cross_attn[:2_56]
            UpperCamelCase					       =  in_proj_weight_cross_attn[2_56:5_12, :]
            UpperCamelCase					       =  in_proj_bias_cross_attn[2_56:5_12]
            UpperCamelCase					       =  in_proj_weight_cross_attn[-2_56:, :]
            UpperCamelCase					       =  in_proj_bias_cross_attn[-2_56:]
def 		lowercase__  (				_UpperCamelCase       ,							_UpperCamelCase) ->			int:
      """simple docstring"""

      UpperCamelCase       ,  UpperCamelCase					       =  image.size
      UpperCamelCase					       =  max(_UpperCamelCase       ,							_UpperCamelCase)
      UpperCamelCase					       =  8_00 if 'detection' in checkpoint_url else 10_00
      UpperCamelCase					       =  target_max_size / current_max_size
      UpperCamelCase					       =  image.resize((int(round(scale * width)), int(round(scale * height))))

      return resized_image
def 		lowercase__  (				_UpperCamelCase) ->			Tuple:
      """simple docstring"""

      UpperCamelCase					       =  F.to_tensor(_UpperCamelCase)
      UpperCamelCase					       =  F.normalize(_UpperCamelCase       ,							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 		lowercase__  (				_UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase) ->			List[str]:
      """simple docstring"""

      logger.info('Converting model...')

      # load original state dict
      UpperCamelCase					       =  torch.hub.load_state_dict_from_url(_UpperCamelCase       ,							map_location='cpu')
      # rename keys
      for src, dest in rename_keys:
            rename_key(_UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase)
      UpperCamelCase					       =  rename_backbone_keys(_UpperCamelCase)
      # query, key and value matrices need special treatment
      read_in_q_k_v(_UpperCamelCase)
      # important: we need to prepend a prefix to each of the base model keys as the head models use different attributes for them
      UpperCamelCase					       =  'model.'
      for key in state_dict.copy().keys():
            if not key.startswith('class_labels_classifier') and not key.startswith('bbox_predictor'):
                  UpperCamelCase					       =  state_dict.pop(_UpperCamelCase)
                  UpperCamelCase					       =  val
    # create HuggingFace model and load state dict
      UpperCamelCase					       =  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					       =  15
            UpperCamelCase					       =  2
            UpperCamelCase					       =  {0: 'table', 1: 'table rotated'}
            UpperCamelCase					       =  idalabel
            UpperCamelCase					       =  {v: k for k, v in idalabel.items()}
      else:
            UpperCamelCase					       =  1_25
            UpperCamelCase					       =  6
            UpperCamelCase					       =  {
                0: 'table',
                1: 'table column',
                2: 'table row',
                3: 'table column header',
                4: 'table projected row header',
                5: 'table spanning cell',
            }
            UpperCamelCase					       =  idalabel
            UpperCamelCase					       =  {v: k for k, v in idalabel.items()}

      UpperCamelCase					       =  DetrImageProcessor(
          format='coco_detection'       ,							max_size=8_00 if 'detection' in checkpoint_url else 10_00)
      UpperCamelCase					       =  TableTransformerForObjectDetection(_UpperCamelCase)
      model.load_state_dict(_UpperCamelCase)
      model.eval()

      # verify our conversion
      UpperCamelCase					       =  'example_pdf.png' if 'detection' in checkpoint_url else 'example_table.png'
      UpperCamelCase					       =  hf_hub_download(repo_id='nielsr/example-pdf'       ,							repo_type='dataset'       ,							filename=_UpperCamelCase)
      UpperCamelCase					       =  Image.open(_UpperCamelCase).convert('RGB')
      UpperCamelCase					       =  normalize(resize(_UpperCamelCase       ,							_UpperCamelCase)).unsqueeze(0)

      UpperCamelCase					       =  model(_UpperCamelCase)

      if "detection" in checkpoint_url:
            UpperCamelCase					       =  (1, 15, 3)
            UpperCamelCase					       =  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					       =  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					       =  (1, 1_25, 7)
            UpperCamelCase					       =  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					       =  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]       ,							_UpperCamelCase       ,							atol=1e-4)
      assert torch.allclose(outputs.pred_boxes[0, :3, :3]       ,							_UpperCamelCase       ,							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(_UpperCamelCase).mkdir(exist_ok=_UpperCamelCase)
            model.save_pretrained(_UpperCamelCase)
            image_processor.save_pretrained(_UpperCamelCase)

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


if __name__ == "__main__":
       __magic_name__   :  Optional[Any]        =							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.'''
       )
       __magic_name__   :  int        =							parser.parse_args()
       convert_table_transformer_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub)











import tempfile
import unittest
from pathlib import Path
from shutil import copyfile

from transformers import MaMaaaTokenizer, is_torch_available
from transformers.testing_utils import (
    get_tests_dir,
    nested_simplify,
    require_sentencepiece,
    require_tokenizers,
    require_torch,
    slow,
)
from transformers.utils import is_sentencepiece_available


if is_sentencepiece_available():
       from transformers.models.mam_aaa.tokenization_mam_aaa import VOCAB_FILES_NAMES, save_json

from ...test_tokenization_common import TokenizerTesterMixin


if is_sentencepiece_available():
       __magic_name__   :  Any        =							get_tests_dir('''fixtures/test_sentencepiece.model''')


if is_torch_available():
       from transformers.models.mam_aaa.modeling_mam_aaa import shift_tokens_right

__magic_name__   :  Optional[Any]        =							12_8022
__magic_name__   :  Dict        =							12_8028





@require_sentencepiece
class 					A__     (     __snake_case   ,    unittest.TestCase ):

     '''simple docstring'''






     snake_case__					=					MaMaaaTokenizer
     snake_case__					=					False
     snake_case__					=					False
     snake_case__					=					True







     def   _SCREAMING_SNAKE_CASE     ( self      :    List[str] ):






           """simple docstring"""




           super().setUp()

           UpperCamelCase					       =  ['</s>', '<unk>', '▁This', '▁is', '▁a', '▁t', 'est', '\u0120', '<pad>']
           UpperCamelCase					       =  dict(zip(_SCREAMING_SNAKE_CASE   ,						range(len(_SCREAMING_SNAKE_CASE ) ) ) )
           UpperCamelCase					       =  Path(self.tmpdirname )
           save_json(_SCREAMING_SNAKE_CASE   ,						save_dir / VOCAB_FILES_NAMES['vocab_file'] )
           if not (save_dir / VOCAB_FILES_NAMES["spm_file"]).exists():
                 copyfile(_SCREAMING_SNAKE_CASE   ,						save_dir / VOCAB_FILES_NAMES['spm_file'] )

           UpperCamelCase					       =  MaMaaaTokenizer.from_pretrained(self.tmpdirname )
           tokenizer.save_pretrained(self.tmpdirname )







     def   _SCREAMING_SNAKE_CASE     ( self      :    str   ,						**_SCREAMING_SNAKE_CASE      :    List[str] ):






           """simple docstring"""




           return MaMaaaTokenizer.from_pretrained(self.tmpdirname   ,						**_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Optional[int]   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any] ):






           """simple docstring"""




           return (
               "This is a test",
               "This is a test",
           )







     def   _SCREAMING_SNAKE_CASE     ( self      :    List[Any] ):






           """simple docstring"""




           UpperCamelCase					       =  '</s>'
           UpperCamelCase					       =  0

           self.assertEqual(self.get_tokenizer()._convert_token_to_id(_SCREAMING_SNAKE_CASE )   ,						_SCREAMING_SNAKE_CASE )
           self.assertEqual(self.get_tokenizer()._convert_id_to_token(_SCREAMING_SNAKE_CASE )   ,						_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Optional[Any] ):






           """simple docstring"""




           UpperCamelCase					       =  self.get_tokenizer()
           UpperCamelCase					       =  list(tokenizer.get_vocab().keys() )

           self.assertEqual(vocab_keys[0]   ,						'</s>' )
           self.assertEqual(vocab_keys[1]   ,						'<unk>' )
           self.assertEqual(vocab_keys[-1]   ,						'<s>' )
           self.assertEqual(len(_SCREAMING_SNAKE_CASE )   ,						tokenizer.vocab_size + len(tokenizer.get_added_vocab() ) )







     @unittest.skip('Skip this test while all models are still to be uploaded.' )
     def   _SCREAMING_SNAKE_CASE     ( self      :    str ):






           """simple docstring"""




           pass







     def   _SCREAMING_SNAKE_CASE     ( self      :    Any ):






           """simple docstring"""




           UpperCamelCase					       =  self.get_tokenizer()

           UpperCamelCase					       =  tokenizer.tokenize('This is a test' )
           self.assertListEqual(_SCREAMING_SNAKE_CASE   ,						['▁This', '▁is', '▁a', '▁t', 'est'] )

           self.assertListEqual(
               tokenizer.convert_tokens_to_ids(_SCREAMING_SNAKE_CASE )   ,						[2, 3, 4, 5, 6]   ,						)

           UpperCamelCase					       =  tokenizer.convert_ids_to_tokens([2, 3, 4, 5, 6] )
           self.assertListEqual(_SCREAMING_SNAKE_CASE   ,						['▁This', '▁is', '▁a', '▁t', 'est'] )

           UpperCamelCase					       =  tokenizer.convert_tokens_to_string(_SCREAMING_SNAKE_CASE )
           self.assertEqual(_SCREAMING_SNAKE_CASE   ,						'This is a test' )







     @slow
     def   _SCREAMING_SNAKE_CASE     ( self      :    str ):






           """simple docstring"""




           UpperCamelCase					       =  {'input_ids': [[12_8022, 11_0108, 397, 11, 3_8272, 2247, 12_4811, 285, 1_8105, 1586, 207, 7, 3_9534, 4428, 397, 1019, 1_8105, 1586, 207, 7, 4_1337, 1_6786, 241, 7, 2_0214, 17, 12_5690, 1_0398, 7, 4_4378, 5_8069, 6_8342, 7798, 7343, 11, 299, 3_3310, 4, 158, 3_7350, 9_4077, 4569, 299, 3_3310, 90, 4, 5_2840, 290, 4, 3_1270, 112, 299, 682, 4, 5_2840, 3_9953, 1_4079, 193, 5_2519, 9_0894, 1_7894, 12_0697, 11, 4_0445, 551, 17, 1019, 5_2519, 9_0894, 1_7756, 963, 11, 4_0445, 480, 17, 9792, 1120, 5173, 1393, 6240, 1_6786, 241, 12_0996, 28, 1245, 1393, 11_8240, 1_1123, 1019, 9_3612, 2691, 1_0618, 9_8058, 12_0409, 1928, 279, 4, 4_0683, 367, 178, 207, 1019, 103, 10_3121, 506, 6_5296, 5, 2], [12_8022, 2_1217, 367, 117, 12_5450, 128, 719, 7, 7308, 40, 9_3612, 1_2669, 1116, 1_6704, 71, 1_7785, 3699, 1_5592, 35, 144, 9584, 241, 1_1943, 713, 950, 799, 2247, 8_8427, 150, 149, 11_8813, 12_0706, 1019, 10_6906, 8_1518, 28, 1224, 2_2799, 397, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [12_8022, 1658, 12_3311, 5155, 5578, 4722, 279, 1_4947, 2366, 1120, 1197, 14, 1348, 9232, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], 'attention_mask': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}  # noqa: E501
           # fmt: on

           self.tokenizer_integration_test_util(
               expected_encoding=_SCREAMING_SNAKE_CASE   ,						model_name='facebook/m2m100_418M'   ,						revision='c168bae485c864188cf9aa0e4108b0b6934dc91e'   ,						)





@require_torch
@require_sentencepiece
@require_tokenizers
class 					A__     (     unittest.TestCase ):

     '''simple docstring'''






     snake_case__					=					"""facebook/m2m100_418M"""
     snake_case__					=					[
         """In my opinion, there are two levels of response from the French government.""",
         """NSA Affair Emphasizes Complete Lack of Debate on Intelligence""",
     ]
     snake_case__					=					[
         """Selon moi, il y a deux niveaux de réponse de la part du gouvernement français.""",
         """L'affaire NSA souligne l'absence totale de débat sur le renseignement""",
     ]

     # fmt: off
     snake_case__					=					[EN_CODE, 5_93, 19_49, 11_57_81, 4, 7_15_86, 42_34, 6_06_33, 12_62_33, 4_32, 12_38_08, 1_55_92, 11_97, 11_71_32, 12_06_18, 5, 2]







     @classmethod
     def   _SCREAMING_SNAKE_CASE     ( cls      :    Any ):






           """simple docstring"""




           UpperCamelCase					       =  MaMaaaTokenizer.from_pretrained(
               cls.checkpoint_name   ,						src_lang='en'   ,						tgt_lang='fr' )
           UpperCamelCase					       =  1
           return cls







     def   _SCREAMING_SNAKE_CASE     ( self      :    str ):






           """simple docstring"""




           self.assertEqual(self.tokenizer.get_lang_id('ar' )   ,						12_8006 )
           self.assertEqual(self.tokenizer.get_lang_id('en' )   ,						12_8022 )
           self.assertEqual(self.tokenizer.get_lang_id('ro' )   ,						12_8076 )
           self.assertEqual(self.tokenizer.get_lang_id('mr' )   ,						12_8063 )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Tuple ):






           """simple docstring"""




           UpperCamelCase					       =  self.tokenizer.get_vocab()
           self.assertEqual(len(_SCREAMING_SNAKE_CASE )   ,						self.tokenizer.vocab_size )
           self.assertEqual(vocab['<unk>']   ,						3 )
           self.assertIn(self.tokenizer.get_lang_token('en' )   ,						_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Optional[Any] ):






           """simple docstring"""




           UpperCamelCase					       =  'en'
           UpperCamelCase					       =  self.tokenizer.batch_encode_plus(self.src_text ).input_ids[0]
           self.assertListEqual(self.expected_src_tokens   ,						_SCREAMING_SNAKE_CASE )







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






           """simple docstring"""




           self.assertIn(_SCREAMING_SNAKE_CASE   ,						self.tokenizer.all_special_ids )
           # fmt: off
           UpperCamelCase					       =  [FR_CODE, 5364, 82, 8642, 4, 294, 47, 8, 1_4028, 136, 3286, 9706, 6, 9_0797, 6, 14_4012, 162, 8_8128, 3_0061, 5, 2]
           # fmt: on
           UpperCamelCase					       =  self.tokenizer.decode(_SCREAMING_SNAKE_CASE   ,						skip_special_tokens=_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  self.tokenizer.decode(generated_ids[1:]   ,						skip_special_tokens=_SCREAMING_SNAKE_CASE )
           self.assertEqual(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
           self.assertNotIn(self.tokenizer.eos_token   ,						_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Optional[Any] ):






           """simple docstring"""




           UpperCamelCase					       =  tempfile.mkdtemp()
           UpperCamelCase					       =  self.tokenizer.lang_token_to_id
           self.tokenizer.save_pretrained(_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  MaMaaaTokenizer.from_pretrained(_SCREAMING_SNAKE_CASE )
           self.assertDictEqual(new_tok.lang_token_to_id   ,						_SCREAMING_SNAKE_CASE )







     @require_torch
     def   _SCREAMING_SNAKE_CASE     ( self      :    Union[str, Any] ):






           """simple docstring"""




           UpperCamelCase					       =  'en'
           UpperCamelCase					       =  'fr'

           UpperCamelCase					       =  self.tokenizer(self.src_text   ,						text_target=self.tgt_text   ,						padding=_SCREAMING_SNAKE_CASE   ,						return_tensors='pt' )

           UpperCamelCase					       =  shift_tokens_right(
               batch['labels']   ,						self.tokenizer.pad_token_id   ,						self.tokenizer.eos_token_id )

           for k in batch:
                 UpperCamelCase					       =  batch[k].tolist()
           # batch = {k: v.tolist() for k,v in batch.items()}
           # fairseq batch: https://gist.github.com/sshleifer/cba08bc2109361a74ac3760a7e30e4f4
           # batch.decoder_inputs_ids[0][0] ==
           assert batch.input_ids[1][0] == EN_CODE
           assert batch.input_ids[1][-1] == 2
           assert batch.labels[1][0] == FR_CODE
           assert batch.labels[1][-1] == 2
           assert batch.decoder_input_ids[1][:2] == [2, FR_CODE]







     @require_torch
     def   _SCREAMING_SNAKE_CASE     ( self      :    Any ):






           """simple docstring"""




           UpperCamelCase					       =  'mr'
           self.assertListEqual(self.tokenizer.prefix_tokens   ,						[self.tokenizer.get_lang_id('mr' )] )
           self.assertListEqual(self.tokenizer.suffix_tokens   ,						[self.tokenizer.eos_token_id] )

           UpperCamelCase					       =  'zh'
           self.assertListEqual(self.tokenizer.prefix_tokens   ,						[self.tokenizer.get_lang_id('zh' )] )
           self.assertListEqual(self.tokenizer.suffix_tokens   ,						[self.tokenizer.eos_token_id] )







     @require_torch
     def   _SCREAMING_SNAKE_CASE     ( self      :    Tuple ):






           """simple docstring"""




           UpperCamelCase					       =  'mr'
           self.tokenizer._switch_to_target_mode()
           self.assertListEqual(self.tokenizer.prefix_tokens   ,						[self.tokenizer.get_lang_id('mr' )] )
           self.assertListEqual(self.tokenizer.suffix_tokens   ,						[self.tokenizer.eos_token_id] )
           self.tokenizer._switch_to_input_mode()
           self.assertListEqual(self.tokenizer.prefix_tokens   ,						[self.tokenizer.get_lang_id(self.tokenizer.src_lang )] )

           UpperCamelCase					       =  'zh'
           self.tokenizer._switch_to_target_mode()
           self.assertListEqual(self.tokenizer.prefix_tokens   ,						[self.tokenizer.get_lang_id('zh' )] )
           self.assertListEqual(self.tokenizer.suffix_tokens   ,						[self.tokenizer.eos_token_id] )
           self.tokenizer._switch_to_input_mode()
           self.assertListEqual(self.tokenizer.prefix_tokens   ,						[self.tokenizer.get_lang_id(self.tokenizer.src_lang )] )







     @require_torch
     def   _SCREAMING_SNAKE_CASE     ( self      :    Any ):






           """simple docstring"""




           UpperCamelCase					       =  self.tokenizer._build_translation_inputs('A test'   ,						return_tensors='pt'   ,						src_lang='en'   ,						tgt_lang='ar' )

           self.assertEqual(
               nested_simplify(_SCREAMING_SNAKE_CASE )   ,						{
                   # en_XX, A, test, EOS
                   'input_ids': [[12_8022, 58, 4183, 2]],
                   'attention_mask': [[1, 1, 1, 1]],
                   # ar_AR
                   'forced_bos_token_id': 12_8006,
               }   ,						)



