




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 						__lowercase		:




      """simple docstring"""





      def __init__(      self				,      A				,      )		->				List[str]:
         snake_case      :				Any				=       parent
         snake_case      :				Optional[int]				=       1_3
         snake_case      :				Any				=       7
         snake_case      :				List[str]				=       True
         snake_case      :				Union[str, Any]				=       True
         snake_case      :				Any				=       True
         snake_case      :				Any				=       True
         snake_case      :				int				=       True
         snake_case      :				Tuple				=       False
         snake_case      :				List[str]				=       False
         snake_case      :				List[str]				=       False
         snake_case      :				Tuple				=       2
         snake_case      :				List[str]				=       9_9
         snake_case      :				Dict				=       0
         snake_case      :				Dict				=       3_2
         snake_case      :				Optional[int]				=       2
         snake_case      :				Optional[Any]				=       4
         snake_case      :				Any				=       0.1
         snake_case      :				str				=       0.1
         snake_case      :				Union[str, Any]				=       5_1_2
         snake_case      :				List[str]				=       1_6
         snake_case      :				str				=       2
         snake_case      :				Dict				=       0.02
         snake_case      :				Dict				=       3
         snake_case      :				str				=       4
         snake_case      :				int				=       """last"""
         snake_case      :				List[Any]				=       True
         snake_case      :				Optional[int]				=       None
         snake_case      :				Union[str, Any]				=       0

      def 	UpperCAmelCase      (      self	)		->				Any:
         snake_case      :				List[Any]				=       ids_tensor([self.batch_size, self.seq_length]				,      self.vocab_size	)
         snake_case      :				Tuple				=       random_attention_mask([self.batch_size, self.seq_length]				,      dtype=tf.floataa	)

         snake_case      :				Union[str, Any]				=       None
         if self.use_input_lengths:
            snake_case      :				List[str]				=       (
                ids_tensor([self.batch_size]				,      vocab_size=2	) + self.seq_length - 2
            )  # small variation of seq_length

         snake_case      :				str				=       None
         if self.use_token_type_ids:
            snake_case      :				List[str]				=       ids_tensor([self.batch_size, self.seq_length]				,      self.n_langs	)

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

         snake_case      :				List[str]				=       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 	UpperCAmelCase      (      self				,      A				,      A				,      A				,      A				,      A				,      A				,      A				,      A				,      A				,      )		->				Union[str, Any]:
         snake_case      :				Optional[Any]				=       TFFlaubertModel(config=A	)
         snake_case      :				List[str]				=       {"""input_ids""": input_ids, """lengths""": input_lengths, """langs""": token_type_ids}
         snake_case      :				List[str]				=       model(A	)

         snake_case      :				List[str]				=       [input_ids, input_mask]
         snake_case      :				int				=       model(A	)
         self.parent.assertEqual(result.last_hidden_state.shape				,      (self.batch_size, self.seq_length, self.hidden_size)	)

      def 	UpperCAmelCase      (      self				,      A				,      A				,      A				,      A				,      A				,      A				,      A				,      A				,      A				,      )		->				List[Any]:
         snake_case      :				int				=       TFFlaubertWithLMHeadModel(A	)

         snake_case      :				Optional[int]				=       {"""input_ids""": input_ids, """lengths""": input_lengths, """langs""": token_type_ids}
         snake_case      :				Dict				=       model(A	)

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

      def 	UpperCAmelCase      (      self				,      A				,      A				,      A				,      A				,      A				,      A				,      A				,      A				,      A				,      )		->				Any:
         snake_case      :				Union[str, Any]				=       TFFlaubertForQuestionAnsweringSimple(A	)

         snake_case      :				Union[str, Any]				=       {"""input_ids""": input_ids, """lengths""": input_lengths}

         snake_case      :				Optional[int]				=       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 	UpperCAmelCase      (      self				,      A				,      A				,      A				,      A				,      A				,      A				,      A				,      A				,      A				,      )		->				str:
         snake_case      :				int				=       TFFlaubertForSequenceClassification(A	)

         snake_case      :				List[str]				=       {"""input_ids""": input_ids, """lengths""": input_lengths}

         snake_case      :				Any				=       model(A	)

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

      def 	UpperCAmelCase      (      self				,      A				,      A				,      A				,      A				,      A				,      A				,      A				,      A				,      A				,      )		->				Any:
         snake_case      :				Any				=       self.num_labels
         snake_case      :				List[Any]				=       TFFlaubertForTokenClassification(config=A	)
         snake_case      :				List[str]				=       {"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
         snake_case      :				Optional[int]				=       model(A	)
         self.parent.assertEqual(result.logits.shape				,      (self.batch_size, self.seq_length, self.num_labels)	)

      def 	UpperCAmelCase      (      self				,      A				,      A				,      A				,      A				,      A				,      A				,      A				,      A				,      A				,      )		->				Dict:
         snake_case      :				int				=       self.num_choices
         snake_case      :				Optional[Any]				=       TFFlaubertForMultipleChoice(config=A	)
         snake_case      :				Tuple				=       tf.tile(tf.expand_dims(A				,      1	)				,      (1, self.num_choices, 1)	)
         snake_case      :				List[str]				=       tf.tile(tf.expand_dims(A				,      1	)				,      (1, self.num_choices, 1)	)
         snake_case      :				Optional[Any]				=       tf.tile(tf.expand_dims(A				,      1	)				,      (1, self.num_choices, 1)	)
         snake_case      :				Tuple				=       {
             """input_ids""": multiple_choice_inputs_ids,
             """attention_mask""": multiple_choice_input_mask,
             """token_type_ids""": multiple_choice_token_type_ids,
         }
         snake_case      :				Optional[Any]				=       model(A	)
         self.parent.assertEqual(result.logits.shape				,      (self.batch_size, self.num_choices)	)

      def 	UpperCAmelCase      (      self	)		->				Dict:
         snake_case      :				List[Any]				=       self.prepare_config_and_inputs()
         (
             snake_case    
         )      :				Tuple				=       config_and_inputs
         snake_case      :				Dict				=       {
             """input_ids""": input_ids,
             """token_type_ids""": token_type_ids,
             """langs""": token_type_ids,
             """lengths""": input_lengths,
         }
         return config, inputs_dict




@require_tf
class 						__lowercase		(UpperCAmelCase__     ,		UpperCAmelCase__     ,		unittest.TestCase							):




      """simple docstring"""





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

      def 	UpperCAmelCase      (      self				,      A				,      A				,      A				,      A				,      A	)		->				Union[str, Any]:
         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 	UpperCAmelCase      (      self	)		->				Any:
         snake_case      :				Optional[int]				=       TFFlaubertModelTester(self	)
         snake_case      :				List[str]				=       ConfigTester(self				,      config_class=A				,      emb_dim=3_7	)

      def 	UpperCAmelCase      (      self	)		->				Union[str, Any]:
         self.config_tester.run_common_tests()

      def 	UpperCAmelCase      (      self	)		->				Optional[Any]:
         snake_case      :				str				=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_model(*A	)

      def 	UpperCAmelCase      (      self	)		->				List[str]:
         snake_case      :				Tuple				=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_lm_head(*A	)

      def 	UpperCAmelCase      (      self	)		->				Dict:
         snake_case      :				Dict				=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_qa(*A	)

      def 	UpperCAmelCase      (      self	)		->				Any:
         snake_case      :				int				=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_sequence_classif(*A	)

      def 	UpperCAmelCase      (      self	)		->				Dict:
         snake_case      :				Tuple				=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_for_token_classification(*A	)

      def 	UpperCAmelCase      (      self	)		->				Optional[int]:
         snake_case      :				str				=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_for_multiple_choice(*A	)

      @slow
      def 	UpperCAmelCase      (      self	)		->				str:
         for model_name in TF_FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            snake_case      :				Any				=       TFFlaubertModel.from_pretrained(A	)
            self.assertIsNotNone(A	)




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




      """simple docstring"""





      @slow
      def 	UpperCAmelCase      (      self	)		->				List[Any]:
         snake_case      :				List[str]				=       TFFlaubertModel.from_pretrained("""jplu/tf-flaubert-small-cased"""	)

         snake_case      :				List[Any]				=       tf.convert_to_tensor(
             [[0, 1_5_8, 7_3_5, 2_5_9_2, 1_4_2_4, 6_7_2_7, 8_2, 1]]				,      dtype=tf.intaa				,      )  # "J'aime flaubert !"

         snake_case      :				Optional[Any]				=       model(A	)[0]
         snake_case      :				List[str]				=       tf.TensorShape((1, 8, 5_1_2)	)
         self.assertEqual(output.shape				,      A	)
         # compare the actual values for a slice.
         snake_case      :				List[str]				=       tf.convert_to_tensor(
             [
                 [
                     [-1.8_76_87_73, -1.56_65_55, 0.27_07_24_18],
                     [-1.6_92_00_38, -0.5_87_35_05, 1.9_32_95_99],
                     [-2.9_56_39_85, -1.6_99_38_35, 1.7_97_20_52],
                 ]
             ]				,      dtype=tf.floataa				,      )

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









'''simple docstring'''







from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, apply_forward_hook
from .modeling_utils import ModelMixin
from .vae import Decoder, DecoderOutput, Encoder, VectorQuantizer







@dataclass
class A__       (						UpperCAmelCase__      ):
      __UpperCamelCase       :  torch.FloatTensor







class A__       (						UpperCAmelCase__ ,     UpperCAmelCase__      ):






      @register_to_config
      def __init__( self			:Optional[Any]     ,						SCREAMING_SNAKE_CASE			:int = 3     ,						SCREAMING_SNAKE_CASE			:int = 3     ,						SCREAMING_SNAKE_CASE			:Tuple[str] = ("DownEncoderBlock2D",)     ,						SCREAMING_SNAKE_CASE			:Tuple[str] = ("UpDecoderBlock2D",)     ,						SCREAMING_SNAKE_CASE			:Tuple[int] = (6_4,)     ,						SCREAMING_SNAKE_CASE			:int = 1     ,						SCREAMING_SNAKE_CASE			:str = "silu"     ,						SCREAMING_SNAKE_CASE			:int = 3     ,						SCREAMING_SNAKE_CASE			:int = 3_2     ,						SCREAMING_SNAKE_CASE			:int = 2_5_6     ,						SCREAMING_SNAKE_CASE			:int = 3_2     ,						SCREAMING_SNAKE_CASE			:Optional[int] = None     ,						SCREAMING_SNAKE_CASE			:float = 0.18_215     ,						SCREAMING_SNAKE_CASE			:str = "group"     ,						)       ->					Optional[int]:




             '''simple docstring'''
             super().__init__()

             # pass init params to Encoder
             _a    : Union[str, Any]     					=Encoder(
                 in_channels=SCREAMING_SNAKE_CASE     ,						out_channels=SCREAMING_SNAKE_CASE     ,						down_block_types=SCREAMING_SNAKE_CASE     ,						block_out_channels=SCREAMING_SNAKE_CASE     ,						layers_per_block=SCREAMING_SNAKE_CASE     ,						act_fn=SCREAMING_SNAKE_CASE     ,						norm_num_groups=SCREAMING_SNAKE_CASE     ,						double_z=SCREAMING_SNAKE_CASE     ,						)

             _a    : Optional[int]     					=vq_embed_dim if vq_embed_dim is not None else latent_channels

             _a    : Optional[int]     					=nn.Convad(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE     ,						1	)
             _a    : str     					=VectorQuantizer(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE     ,						beta=0.25     ,						remap=SCREAMING_SNAKE_CASE     ,						sane_index_shape=SCREAMING_SNAKE_CASE	)
             _a    : List[str]     					=nn.Convad(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE     ,						1	)

             # pass init params to Decoder
             _a    : List[str]     					=Decoder(
                 in_channels=SCREAMING_SNAKE_CASE     ,						out_channels=SCREAMING_SNAKE_CASE     ,						up_block_types=SCREAMING_SNAKE_CASE     ,						block_out_channels=SCREAMING_SNAKE_CASE     ,						layers_per_block=SCREAMING_SNAKE_CASE     ,						act_fn=SCREAMING_SNAKE_CASE     ,						norm_num_groups=SCREAMING_SNAKE_CASE     ,						norm_type=SCREAMING_SNAKE_CASE     ,						)






      @apply_forward_hook
      def __UpperCAmelCase       ( self			:Optional[Any]     ,						SCREAMING_SNAKE_CASE			:torch.FloatTensor     ,						SCREAMING_SNAKE_CASE			:bool = True	)       ->					VQEncoderOutput:




             '''simple docstring'''
             _a    : Optional[int]     					=self.encoder(SCREAMING_SNAKE_CASE	)
             _a    : int     					=self.quant_conv(SCREAMING_SNAKE_CASE	)

             if not return_dict:
                    return (h,)

             return VQEncoderOutput(latents=SCREAMING_SNAKE_CASE	)






      @apply_forward_hook
      def __UpperCAmelCase       ( self			:List[Any]     ,						SCREAMING_SNAKE_CASE			:torch.FloatTensor     ,						SCREAMING_SNAKE_CASE			:bool = False     ,						SCREAMING_SNAKE_CASE			:bool = True	)       ->					Union[DecoderOutput, torch.FloatTensor]:




             '''simple docstring'''
             # also go through quantization layer
             if not force_not_quantize:
                    _a					,    _a					,    _a    : Tuple     					=self.quantize(SCREAMING_SNAKE_CASE	)
             else:
                    _a    : str     					=h
             _a    : Dict     					=self.post_quant_conv(SCREAMING_SNAKE_CASE	)
             _a    : Union[str, Any]     					=self.decoder(SCREAMING_SNAKE_CASE     ,						quant if self.config.norm_type == """spatial""" else None	)

             if not return_dict:
                    return (dec,)

             return DecoderOutput(sample=SCREAMING_SNAKE_CASE	)






      def __UpperCAmelCase       ( self			:Optional[Any]     ,						SCREAMING_SNAKE_CASE			:torch.FloatTensor     ,						SCREAMING_SNAKE_CASE			:bool = True	)       ->					Union[DecoderOutput, torch.FloatTensor]:




             '''simple docstring'''
             _a    : Tuple     					=sample
             _a    : int     					=self.encode(SCREAMING_SNAKE_CASE	).latents
             _a    : List[Any]     					=self.decode(SCREAMING_SNAKE_CASE	).sample

             if not return_dict:
                    return (dec,)

             return DecoderOutput(sample=SCREAMING_SNAKE_CASE	)


