


import argparse
from pathlib import Path

import requests
import torch
from PIL import Image

from transformers import (
    RobertaTokenizer,
    TrOCRConfig,
    TrOCRForCausalLM,
    TrOCRProcessor,
    VisionEncoderDecoderModel,
    ViTConfig,
    ViTImageProcessor,
    ViTModel,
)
from transformers.utils import logging


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






def 		lowercase_(  SCREAMING_SNAKE_CASE_      ,						SCREAMING_SNAKE_CASE_	):


  '''simple docstring'''

  lowerCamelCase :					Dict    =				[]
  for i in range(encoder_config.num_hidden_layers	):
    # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
    rename_keys.append(
        (f"""encoder.deit.blocks.{i}.norm1.weight""", f"""encoder.encoder.layer.{i}.layernorm_before.weight""")	)
    rename_keys.append((f"""encoder.deit.blocks.{i}.norm1.bias""", f"""encoder.encoder.layer.{i}.layernorm_before.bias""")	)
    rename_keys.append(
        (f"""encoder.deit.blocks.{i}.attn.proj.weight""", f"""encoder.encoder.layer.{i}.attention.output.dense.weight""")	)
    rename_keys.append(
        (f"""encoder.deit.blocks.{i}.attn.proj.bias""", f"""encoder.encoder.layer.{i}.attention.output.dense.bias""")	)
    rename_keys.append(
        (f"""encoder.deit.blocks.{i}.norm2.weight""", f"""encoder.encoder.layer.{i}.layernorm_after.weight""")	)
    rename_keys.append((f"""encoder.deit.blocks.{i}.norm2.bias""", f"""encoder.encoder.layer.{i}.layernorm_after.bias""")	)
    rename_keys.append(
        (f"""encoder.deit.blocks.{i}.mlp.fc1.weight""", f"""encoder.encoder.layer.{i}.intermediate.dense.weight""")	)
    rename_keys.append(
        (f"""encoder.deit.blocks.{i}.mlp.fc1.bias""", f"""encoder.encoder.layer.{i}.intermediate.dense.bias""")	)
    rename_keys.append(
        (f"""encoder.deit.blocks.{i}.mlp.fc2.weight""", f"""encoder.encoder.layer.{i}.output.dense.weight""")	)
    rename_keys.append((f"""encoder.deit.blocks.{i}.mlp.fc2.bias""", f"""encoder.encoder.layer.{i}.output.dense.bias""")	)

  # cls token, position embeddings and patch embeddings of encoder
  rename_keys.extend(
      [
          ("encoder.deit.cls_token", "encoder.embeddings.cls_token"),
          ("encoder.deit.pos_embed", "encoder.embeddings.position_embeddings"),
          ("encoder.deit.patch_embed.proj.weight", "encoder.embeddings.patch_embeddings.projection.weight"),
          ("encoder.deit.patch_embed.proj.bias", "encoder.embeddings.patch_embeddings.projection.bias"),
          ("encoder.deit.norm.weight", "encoder.layernorm.weight"),
          ("encoder.deit.norm.bias", "encoder.layernorm.bias"),
      ]	)

  return rename_keys






def 		lowercase_(  SCREAMING_SNAKE_CASE_      ,						SCREAMING_SNAKE_CASE_	):


  '''simple docstring'''

  for i in range(encoder_config.num_hidden_layers	):
    # queries, keys and values (only weights, no biases)
    lowerCamelCase :					Dict    =				state_dict.pop(f"""encoder.deit.blocks.{i}.attn.qkv.weight"""	)

    lowerCamelCase :					Tuple    =				in_proj_weight[
        : encoder_config.hidden_size, :
    ]
    lowerCamelCase :					List[str]    =				in_proj_weight[
        encoder_config.hidden_size : encoder_config.hidden_size * 2, :
    ]
    lowerCamelCase :					List[str]    =				in_proj_weight[
        -encoder_config.hidden_size :, :
    ]






def 		lowercase_(  SCREAMING_SNAKE_CASE_      ,						SCREAMING_SNAKE_CASE_      ,						SCREAMING_SNAKE_CASE_	):


  '''simple docstring'''

  lowerCamelCase :					List[Any]    =				dct.pop(SCREAMING_SNAKE_CASE_	)
  lowerCamelCase :					List[str]    =				val






def 		lowercase_(  SCREAMING_SNAKE_CASE_	):


  '''simple docstring'''

  if "handwritten" in checkpoint_url:
    lowerCamelCase :					List[Any]    =				"https://fki.tic.heia-fr.ch/static/img/a01-122-02-00.jpg"  # industry
    # url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02-12.jpg" # have
    # url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02-10.jpg" # let
    # url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02.jpg"  #
    # url = "https://fki.tic.heia-fr.ch/static/img/a01-122.jpg"
  elif "printed" in checkpoint_url or "stage1" in checkpoint_url:
    lowerCamelCase :					str    =				"https://www.researchgate.net/profile/Dinh-Sang/publication/338099565/figure/fig8/AS:840413229350922@1577381536857/An-receipt-example-in-the-SROIE-2019-dataset_Q640.jpg"
  lowerCamelCase :					Optional[Any]    =				Image.open(requests.get(SCREAMING_SNAKE_CASE_      ,						stream=SCREAMING_SNAKE_CASE_	).raw	).convert("RGB"	)
  return im






@torch.no_grad()
def 		lowercase_(  SCREAMING_SNAKE_CASE_      ,						SCREAMING_SNAKE_CASE_	):


  '''simple docstring'''

  lowerCamelCase :					List[Any]    =				ViTConfig(image_size=384      ,						qkv_bias=SCREAMING_SNAKE_CASE_	)
  lowerCamelCase :					Optional[int]    =				TrOCRConfig()

  # size of the architecture
  if "base" in checkpoint_url:
    lowerCamelCase :					Any    =				768
  elif "large" in checkpoint_url:
    # use ViT-large encoder
    lowerCamelCase :					Dict    =				1024
    lowerCamelCase :					Optional[int]    =				4096
    lowerCamelCase :					Any    =				24
    lowerCamelCase :					Dict    =				16
    lowerCamelCase :					str    =				1024
  else:
    raise ValueError("Should either find 'base' or 'large' in checkpoint URL"	)

  # the large-printed + stage1 checkpoints uses sinusoidal position embeddings, no layernorm afterwards
  if "large-printed" in checkpoint_url or "stage1" in checkpoint_url:
    lowerCamelCase :					List[Any]    =				False
    lowerCamelCase :					List[str]    =				"relu"
    lowerCamelCase :					Any    =				1024
    lowerCamelCase :					Any    =				True
    lowerCamelCase :					int    =				False
    lowerCamelCase :					List[str]    =				False

  # load HuggingFace model
  lowerCamelCase :					List[str]    =				ViTModel(SCREAMING_SNAKE_CASE_      ,						add_pooling_layer=SCREAMING_SNAKE_CASE_	)
  lowerCamelCase :					Union[str, Any]    =				TrOCRForCausalLM(SCREAMING_SNAKE_CASE_	)
  lowerCamelCase :					str    =				VisionEncoderDecoderModel(encoder=SCREAMING_SNAKE_CASE_      ,						decoder=SCREAMING_SNAKE_CASE_	)
  model.eval()

  # load state_dict of original model, rename some keys
  lowerCamelCase :					str    =				torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE_      ,						map_location="cpu"      ,						check_hash=SCREAMING_SNAKE_CASE_	)["model"]

  lowerCamelCase :					Union[str, Any]    =				create_rename_keys(SCREAMING_SNAKE_CASE_      ,						SCREAMING_SNAKE_CASE_	)
  for src, dest in rename_keys:
    rename_key(SCREAMING_SNAKE_CASE_      ,						SCREAMING_SNAKE_CASE_      ,						SCREAMING_SNAKE_CASE_	)
  read_in_q_k_v(SCREAMING_SNAKE_CASE_      ,						SCREAMING_SNAKE_CASE_	)

  # remove parameters we don't need
  del state_dict["encoder.deit.head.weight"]
  del state_dict["encoder.deit.head.bias"]
  del state_dict["decoder.version"]

  # add prefix to decoder keys
  for key, val in state_dict.copy().items():
    lowerCamelCase :					Any    =				state_dict.pop(SCREAMING_SNAKE_CASE_	)
    if key.startswith("decoder"	) and "output_projection" not in key:
      lowerCamelCase :					Union[str, Any]    =				val
    else:
      lowerCamelCase :					Tuple    =				val

    # load state dict
  model.load_state_dict(SCREAMING_SNAKE_CASE_	)

  # Check outputs on an image
  lowerCamelCase :					List[str]    =				ViTImageProcessor(size=encoder_config.image_size	)
  lowerCamelCase :					Tuple    =				RobertaTokenizer.from_pretrained("roberta-large"	)
  lowerCamelCase :					Union[str, Any]    =				TrOCRProcessor(SCREAMING_SNAKE_CASE_      ,						SCREAMING_SNAKE_CASE_	)

  lowerCamelCase :					int    =				processor(images=prepare_img(SCREAMING_SNAKE_CASE_	)      ,						return_tensors="pt"	).pixel_values

  # verify logits
  lowerCamelCase :					Dict    =				torch.tensor([[model.config.decoder.decoder_start_token_id]]	)
  lowerCamelCase :					int    =				model(pixel_values=SCREAMING_SNAKE_CASE_      ,						decoder_input_ids=SCREAMING_SNAKE_CASE_	)
  lowerCamelCase :					Tuple    =				outputs.logits

  lowerCamelCase :					Optional[int]    =				torch.Size([1, 1, 50265]	)
  if "trocr-base-handwritten" in checkpoint_url:
    lowerCamelCase :					Any    =				torch.tensor(
        [-1.4502, -4.6683, -0.5347, -2.9291, 9.1435, -3.0571, 8.9764, 1.7560, 8.7358, -1.5311]	)
  elif "trocr-large-handwritten" in checkpoint_url:
    lowerCamelCase :					str    =				torch.tensor(
        [-2.6437, -1.3129, -2.2596, -5.3455, 6.3539, 1.7604, 5.4991, 1.4702, 5.6113, 2.0170]	)
  elif "trocr-base-printed" in checkpoint_url:
    lowerCamelCase :					List[Any]    =				torch.tensor(
        [-5.6816, -5.8388, 1.1398, -6.9034, 6.8505, -2.4393, 1.2284, -1.0232, -1.9661, -3.9210]	)
  elif "trocr-large-printed" in checkpoint_url:
    lowerCamelCase :					Optional[Any]    =				torch.tensor(
        [-6.0162, -7.0959, 4.4155, -5.1063, 7.0468, -3.1631, 2.6466, -0.3081, -0.8106, -1.7535]	)

  if "stage1" not in checkpoint_url:
    assert logits.shape == expected_shape, "Shape of logits not as expected"
    assert torch.allclose(logits[0, 0, :10]      ,						SCREAMING_SNAKE_CASE_      ,						atol=1E-3	), "First elements of logits not as expected"

  Path(SCREAMING_SNAKE_CASE_	).mkdir(exist_ok=SCREAMING_SNAKE_CASE_	)
  print(f"""Saving model to {pytorch_dump_folder_path}"""	)
  model.save_pretrained(SCREAMING_SNAKE_CASE_	)
  print(f"""Saving processor to {pytorch_dump_folder_path}"""	)
  processor.save_pretrained(SCREAMING_SNAKE_CASE_	)


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

   parser.add_argument(
       '''--checkpoint_url''',
       default='''https://layoutlm.blob.core.windows.net/trocr/model_zoo/fairseq/trocr-base-handwritten.pt''',
       type=str,
       help='''URL to the original PyTorch checkpoint (.pth file).''',
   )
   parser.add_argument(
       '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the folder to output PyTorch model.'''
   )
   _snake_case             =     parser.parse_args()
   convert_tr_ocr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)







from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
   import numpy as np
   import tensorflow as tf

   from transformers import (
       TF_FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
       FlaubertConfig,
       TFFlaubertForMultipleChoice,
       TFFlaubertForQuestionAnsweringSimple,
       TFFlaubertForSequenceClassification,
       TFFlaubertForTokenClassification,
       TFFlaubertModel,
       TFFlaubertWithLMHeadModel,
   )
class 	UpperCAmelCase_					:

   '''simple docstring'''




   def __init__(		self						,				__A						,				):






     """simple docstring"""





     lowerCamelCase :					str    =				parent
     lowerCamelCase :					Union[str, Any]    =				13
     lowerCamelCase :					Optional[Any]    =				7
     lowerCamelCase :					List[str]    =				True
     lowerCamelCase :					Optional[int]    =				True
     lowerCamelCase :					Union[str, Any]    =				True
     lowerCamelCase :					List[Any]    =				True
     lowerCamelCase :					Tuple    =				True
     lowerCamelCase :					Any    =				False
     lowerCamelCase :					int    =				False
     lowerCamelCase :					Tuple    =				False
     lowerCamelCase :					Union[str, Any]    =				2
     lowerCamelCase :					Dict    =				99
     lowerCamelCase :					Tuple    =				0
     lowerCamelCase :					Any    =				32
     lowerCamelCase :					List[Any]    =				2
     lowerCamelCase :					Tuple    =				4
     lowerCamelCase :					List[str]    =				0.1
     lowerCamelCase :					int    =				0.1
     lowerCamelCase :					int    =				512
     lowerCamelCase :					List[Any]    =				16
     lowerCamelCase :					Any    =				2
     lowerCamelCase :					Any    =				0.02
     lowerCamelCase :					List[str]    =				3
     lowerCamelCase :					Tuple    =				4
     lowerCamelCase :					int    =				"last"
     lowerCamelCase :					int    =				True
     lowerCamelCase :					Dict    =				None
     lowerCamelCase :					Tuple    =				0






   def _snake_case    (		self ):






     """simple docstring"""





     lowerCamelCase :					str    =				ids_tensor([self.batch_size, self.seq_length]						,				self.vocab_size )
     lowerCamelCase :					Tuple    =				random_attention_mask([self.batch_size, self.seq_length]						,				dtype=tf.floataa )

     lowerCamelCase :					Tuple    =				None
     if self.use_input_lengths:
       lowerCamelCase :					Optional[Any]    =				(
           ids_tensor([self.batch_size]						,				vocab_size=2 ) + self.seq_length - 2
       )  # small variation of seq_length

     lowerCamelCase :					str    =				None
     if self.use_token_type_ids:
       lowerCamelCase :					Optional[int]    =				ids_tensor([self.batch_size, self.seq_length]						,				self.n_langs )

     lowerCamelCase :					Dict    =				None
     lowerCamelCase :					Dict    =				None
     lowerCamelCase :					Tuple    =				None
     if self.use_labels:
       lowerCamelCase :					str    =				ids_tensor([self.batch_size]						,				self.type_sequence_label_size )
       lowerCamelCase :					Union[str, Any]    =				ids_tensor([self.batch_size, self.seq_length]						,				self.num_labels )
       lowerCamelCase :					int    =				ids_tensor([self.batch_size]						,				2						,				dtype=tf.floataa )
       lowerCamelCase :					Union[str, Any]    =				ids_tensor([self.batch_size]						,				self.num_choices )

     lowerCamelCase :					List[Any]    =				FlaubertConfig(
         vocab_size=self.vocab_size						,				n_special=self.n_special						,				emb_dim=self.hidden_size						,				n_layers=self.num_hidden_layers						,				n_heads=self.num_attention_heads						,				dropout=self.hidden_dropout_prob						,				attention_dropout=self.attention_probs_dropout_prob						,				gelu_activation=self.gelu_activation						,				sinusoidal_embeddings=self.sinusoidal_embeddings						,				asm=self.asm						,				causal=self.causal						,				n_langs=self.n_langs						,				max_position_embeddings=self.max_position_embeddings						,				initializer_range=self.initializer_range						,				summary_type=self.summary_type						,				use_proj=self.use_proj						,				bos_token_id=self.bos_token_id						,				)

     return (
         config,
         input_ids,
         token_type_ids,
         input_lengths,
         sequence_labels,
         token_labels,
         is_impossible_labels,
         choice_labels,
         input_mask,
     )






   def _snake_case    (		self						,				__A						,				__A						,				__A						,				__A						,				__A						,				__A						,				__A						,				__A						,				__A						,				):






     """simple docstring"""





     lowerCamelCase :					Optional[Any]    =				TFFlaubertModel(config=__A )
     lowerCamelCase :					Any    =				{"input_ids": input_ids, "lengths": input_lengths, "langs": token_type_ids}
     lowerCamelCase :					Dict    =				model(__A )

     lowerCamelCase :					Any    =				[input_ids, input_mask]
     lowerCamelCase :					Tuple    =				model(__A )
     self.parent.assertEqual(result.last_hidden_state.shape						,				(self.batch_size, self.seq_length, self.hidden_size) )






   def _snake_case    (		self						,				__A						,				__A						,				__A						,				__A						,				__A						,				__A						,				__A						,				__A						,				__A						,				):






     """simple docstring"""





     lowerCamelCase :					int    =				TFFlaubertWithLMHeadModel(__A )

     lowerCamelCase :					List[str]    =				{"input_ids": input_ids, "lengths": input_lengths, "langs": token_type_ids}
     lowerCamelCase :					int    =				model(__A )

     self.parent.assertEqual(result.logits.shape						,				(self.batch_size, self.seq_length, self.vocab_size) )






   def _snake_case    (		self						,				__A						,				__A						,				__A						,				__A						,				__A						,				__A						,				__A						,				__A						,				__A						,				):






     """simple docstring"""





     lowerCamelCase :					Union[str, Any]    =				TFFlaubertForQuestionAnsweringSimple(__A )

     lowerCamelCase :					Optional[int]    =				{"input_ids": input_ids, "lengths": input_lengths}

     lowerCamelCase :					Union[str, Any]    =				model(__A )

     self.parent.assertEqual(result.start_logits.shape						,				(self.batch_size, self.seq_length) )
     self.parent.assertEqual(result.end_logits.shape						,				(self.batch_size, self.seq_length) )






   def _snake_case    (		self						,				__A						,				__A						,				__A						,				__A						,				__A						,				__A						,				__A						,				__A						,				__A						,				):






     """simple docstring"""





     lowerCamelCase :					Optional[int]    =				TFFlaubertForSequenceClassification(__A )

     lowerCamelCase :					str    =				{"input_ids": input_ids, "lengths": input_lengths}

     lowerCamelCase :					Union[str, Any]    =				model(__A )

     self.parent.assertEqual(result.logits.shape						,				(self.batch_size, self.type_sequence_label_size) )






   def _snake_case    (		self						,				__A						,				__A						,				__A						,				__A						,				__A						,				__A						,				__A						,				__A						,				__A						,				):






     """simple docstring"""





     lowerCamelCase :					Tuple    =				self.num_labels
     lowerCamelCase :					Optional[Any]    =				TFFlaubertForTokenClassification(config=__A )
     lowerCamelCase :					int    =				{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
     lowerCamelCase :					Union[str, Any]    =				model(__A )
     self.parent.assertEqual(result.logits.shape						,				(self.batch_size, self.seq_length, self.num_labels) )






   def _snake_case    (		self						,				__A						,				__A						,				__A						,				__A						,				__A						,				__A						,				__A						,				__A						,				__A						,				):






     """simple docstring"""





     lowerCamelCase :					Any    =				self.num_choices
     lowerCamelCase :					Optional[Any]    =				TFFlaubertForMultipleChoice(config=__A )
     lowerCamelCase :					Tuple    =				tf.tile(tf.expand_dims(__A						,				1 )						,				(1, self.num_choices, 1) )
     lowerCamelCase :					int    =				tf.tile(tf.expand_dims(__A						,				1 )						,				(1, self.num_choices, 1) )
     lowerCamelCase :					List[str]    =				tf.tile(tf.expand_dims(__A						,				1 )						,				(1, self.num_choices, 1) )
     lowerCamelCase :					Optional[int]    =				{
         "input_ids": multiple_choice_inputs_ids,
         "attention_mask": multiple_choice_input_mask,
         "token_type_ids": multiple_choice_token_type_ids,
     }
     lowerCamelCase :					Union[str, Any]    =				model(__A )
     self.parent.assertEqual(result.logits.shape						,				(self.batch_size, self.num_choices) )






   def _snake_case    (		self ):






     """simple docstring"""





     lowerCamelCase :					Dict    =				self.prepare_config_and_inputs()
     (
         (
         lowerCamelCase
     )    ,	(
         lowerCamelCase
     )    ,	(
         lowerCamelCase
     )    ,	(
         lowerCamelCase
     )    ,	(
         lowerCamelCase
     )    ,	(
         lowerCamelCase
     )    ,	(
         lowerCamelCase
     )    ,	(
         lowerCamelCase
     )    ,	(
         lowerCamelCase
     )    ,	
     ) :					Optional[Any]    =				config_and_inputs
     lowerCamelCase :					List[Any]    =				{
         "input_ids": input_ids,
         "token_type_ids": token_type_ids,
         "langs": token_type_ids,
         "lengths": input_lengths,
     }
     return config, inputs_dict



@require_tf
class 	UpperCAmelCase_					(					UpperCamelCase		, UpperCamelCase		, unittest.TestCase    ):

   '''simple docstring'''




   __A :      str     		=		(
       (
           TFFlaubertModel,
           TFFlaubertWithLMHeadModel,
           TFFlaubertForSequenceClassification,
           TFFlaubertForQuestionAnsweringSimple,
           TFFlaubertForTokenClassification,
           TFFlaubertForMultipleChoice,
       )
       if is_tf_available()
       else ()
   )
   __A :      Dict     		=		(
       (TFFlaubertWithLMHeadModel,) if is_tf_available() else ()
   )  # TODO (PVP): Check other models whether language generation is also applicable
   __A :      Any     		=		(
       {
           "feature-extraction": TFFlaubertModel,
           "fill-mask": TFFlaubertWithLMHeadModel,
           "question-answering": TFFlaubertForQuestionAnsweringSimple,
           "text-classification": TFFlaubertForSequenceClassification,
           "token-classification": TFFlaubertForTokenClassification,
           "zero-shot": TFFlaubertForSequenceClassification,
       }
       if is_tf_available()
       else {}
   )
   __A :      List[str]     		=		False
   __A :      List[str]     		=		False






   def _snake_case    (		self						,				__A						,				__A						,				__A						,				__A						,				__A ):






     """simple docstring"""





     if (
         pipeline_test_casse_name == "QAPipelineTests"
         and tokenizer_name is not None
         and not tokenizer_name.endswith("Fast" )
     ):
       # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
       # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
       # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
       return True

     return False






   def _snake_case    (		self ):






     """simple docstring"""





     lowerCamelCase :					Tuple    =				TFFlaubertModelTester(self )
     lowerCamelCase :					Optional[int]    =				ConfigTester(self						,				config_class=__A						,				emb_dim=37 )






   def _snake_case    (		self ):






     """simple docstring"""





     self.config_tester.run_common_tests()






   def _snake_case    (		self ):






     """simple docstring"""





     lowerCamelCase :					Union[str, Any]    =				self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_model(*__A )






   def _snake_case    (		self ):






     """simple docstring"""





     lowerCamelCase :					Union[str, Any]    =				self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_lm_head(*__A )






   def _snake_case    (		self ):






     """simple docstring"""





     lowerCamelCase :					int    =				self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_qa(*__A )






   def _snake_case    (		self ):






     """simple docstring"""





     lowerCamelCase :					Tuple    =				self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_sequence_classif(*__A )






   def _snake_case    (		self ):






     """simple docstring"""





     lowerCamelCase :					Dict    =				self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_for_token_classification(*__A )






   def _snake_case    (		self ):






     """simple docstring"""





     lowerCamelCase :					str    =				self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_for_multiple_choice(*__A )






   @slow
   def _snake_case    (		self ):






     """simple docstring"""





     for model_name in TF_FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       lowerCamelCase :					int    =				TFFlaubertModel.from_pretrained(__A )
       self.assertIsNotNone(__A )



@require_tf
@require_sentencepiece
@require_tokenizers
class 	UpperCAmelCase_					(					unittest.TestCase    ):

   '''simple docstring'''




   @slow
   def _snake_case    (		self ):






     """simple docstring"""





     lowerCamelCase :					Optional[int]    =				TFFlaubertModel.from_pretrained("jplu/tf-flaubert-small-cased" )

     lowerCamelCase :					str    =				tf.convert_to_tensor(
         [[0, 158, 735, 2592, 1424, 6727, 82, 1]]						,				dtype=tf.intaa						,				)  # "J'aime flaubert !"

     lowerCamelCase :					Dict    =				model(__A )[0]
     lowerCamelCase :					List[str]    =				tf.TensorShape((1, 8, 512) )
     self.assertEqual(output.shape						,				__A )
     # compare the actual values for a slice.
     lowerCamelCase :					Tuple    =				tf.convert_to_tensor(
         [
             [
                 [-1.8768773, -1.566555, 0.27072418],
                 [-1.6920038, -0.5873505, 1.9329599],
                 [-2.9563985, -1.6993835, 1.7972052],
             ]
         ]						,				dtype=tf.floataa						,				)

     self.assertTrue(np.allclose(output[:, :3, :3].numpy()						,				expected_slice.numpy()						,				atol=1e-4 ) )



