



"""simple docstring"""







import argparse
import json
import os
from pathlib import Path

import requests
import torch

from transformers import JukeboxConfig, JukeboxModel
from transformers.utils import logging


logging.set_verbosity_info()
lowercase__		:    Tuple				   =     logging.get_logger(__name__)


lowercase__		:    Optional[Any]				   =     """https://openaipublic.azureedge.net/jukebox/models/"""
lowercase__		:    str				   =     {
    """jukebox-1b-lyrics""": [
        """5b/vqvae.pth.tar""",
        """5b/prior_level_0.pth.tar""",
        """5b/prior_level_1.pth.tar""",
        """1b_lyrics/prior_level_2.pth.tar""",
    ],
    """jukebox-5b-lyrics""": [
        """5b/vqvae.pth.tar""",
        """5b/prior_level_0.pth.tar""",
        """5b/prior_level_1.pth.tar""",
        """5b_lyrics/prior_level_2.pth.tar""",
    ],
}

def     UpperCamelCase_   (   lowerCAmelCase__				:    str						)						->       Tuple:

  """simple docstring"""




  if key.endswith('.model.1.bias'						) and len(key.split('.'						)						) > 10:
    lowerCAmelCase_       :    Dict								=			key.replace('.model.1.bias'   ,       '.conv1d_1.bias'						)
  elif key.endswith('.model.1.weight'						) and len(key.split('.'						)						) > 10:
    lowerCAmelCase_       :    List[Any]								=			key.replace('.model.1.weight'   ,       '.conv1d_1.weight'						)
  elif key.endswith('.model.3.bias'						) and len(key.split('.'						)						) > 10:
    lowerCAmelCase_       :    Tuple								=			key.replace('.model.3.bias'   ,       '.conv1d_2.bias'						)
  elif key.endswith('.model.3.weight'						) and len(key.split('.'						)						) > 10:
    lowerCAmelCase_       :    List[str]								=			key.replace('.model.3.weight'   ,       '.conv1d_2.weight'						)

  if "conditioner_blocks.0." in key:
    lowerCAmelCase_       :    Union[str, Any]								=			key.replace('conditioner_blocks.0'   ,       'conditioner_blocks'						)

  if "prime_prior" in key:
    lowerCAmelCase_       :    Dict								=			key.replace('prime_prior'   ,       'encoder'						)

  if ".emb." in key and "total" not in key and "absolute" not in key and "relative" not in key:
    lowerCAmelCase_       :    Optional[int]								=			key.replace('.emb.'   ,       '.'						)

  if key.endswith('k'						):  # replace vqvae.X.k with vqvae.X.codebook
    return key.replace('.k'   ,       '.codebook'						)
  if "y_emb." in key:
    return key.replace('y_emb.'   ,       'metadata_embedding.'						)

  if "x_emb.emb." in key:
    lowerCAmelCase_       :    str								=			key.replace('0.x_emb.emb'   ,       'embed_tokens'						)

  if "prime_state_ln" in key:
    return key.replace('prime_state_ln'   ,       'encoder.final_layer_norm'						)
  if ".ln" in key:
    return key.replace('.ln'   ,       '.layer_norm'						)
  if "_ln" in key:
    return key.replace('_ln'   ,       '_layer_norm'						)

  if "prime_state_proj" in key:
    return key.replace('prime_state_proj'   ,       'encoder.proj_in'						)
  if "prime_x_out" in key:
    return key.replace('prime_x_out'   ,       'encoder.lm_head'						)
  if "prior.x_out" in key:
    return key.replace('x_out'   ,       'fc_proj_out'						)
  if "x_emb" in key:
    return key.replace('x_emb'   ,       'embed_tokens'						)

  return key

def     UpperCamelCase_   (   lowerCAmelCase__				:    List[Any]   ,       lowerCAmelCase__				:    Optional[int]   ,       lowerCAmelCase__				:    int   ,       lowerCAmelCase__				:    Optional[int]						)						->       List[str]:

  """simple docstring"""




  lowerCAmelCase_       :    List[Any]								=			{}
  import re

  lowerCAmelCase_       :    Tuple								=			re.compile(R'encoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).(bias|weight)'						)
  lowerCAmelCase_       :    Optional[int]								=			re.compile(
      R'encoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).model.(\d*).model.(\d*).(bias|weight)'						)
  lowerCAmelCase_       :    Dict								=			re.compile(R'encoders.(\d*).level_blocks.(\d*).model.(\d*).(bias|weight)'						)

  lowerCAmelCase_       :    int								=			re.compile(R'decoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).(bias|weight)'						)
  lowerCAmelCase_       :    Dict								=			re.compile(
      R'decoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).model.(\d*).model.(\d*).(bias|weight)'						)
  lowerCAmelCase_       :    Union[str, Any]								=			re.compile(R'decoders.(\d*).level_blocks.(\d*).model.(\d*).(bias|weight)'						)

  lowerCAmelCase_       :    Tuple								=			re.compile(R'conditioner_blocks.(\d*).cond.model.(\d*).(\d).(bias|weight)'						)
  lowerCAmelCase_       :    Dict								=			re.compile(
      R'conditioner_blocks.(\d*).cond.model.(\d*).(\d).model.(\d*).model.(\d*).(bias|weight)'						)
  lowerCAmelCase_       :    Dict								=			re.compile(R'conditioner_blocks.(\d*).cond.model.(\d*).(bias|weight)'						)

  for original_key, value in state_dict.items():
    # rename vqvae.encoder keys
    if re_encoder_block_conv_in.fullmatch(lowerCAmelCase__						):
      lowerCAmelCase_       :    Dict								=			re_encoder_block_conv_in.match(lowerCAmelCase__						)
      lowerCAmelCase_       :    int								=			regex_match.groups()
      lowerCAmelCase_       :    str								=			int(groups[2]						) * 2 + int(groups[3]						)
      lowerCAmelCase_       :    List[Any]								=			f"encoders.{groups[0]}.level_blocks.{groups[1]}.downsample_block.{block_index}.{groups[-1]}"
      lowerCAmelCase_       :    int								=			re_encoder_block_conv_in.sub(lowerCAmelCase__   ,       lowerCAmelCase__						)

    elif re_encoder_block_resnet.fullmatch(lowerCAmelCase__						):
      lowerCAmelCase_       :    Optional[Any]								=			re_encoder_block_resnet.match(lowerCAmelCase__						)
      lowerCAmelCase_       :    Union[str, Any]								=			regex_match.groups()
      lowerCAmelCase_       :    List[str]								=			int(groups[2]						) * 2 + int(groups[3]						)
      lowerCAmelCase_       :    List[str]								=			{'1': 1, '3': 2}[groups[-2]]
      lowerCAmelCase_       :    Any								=			f"encoders.{groups[0]}.level_blocks.{groups[1]}.downsample_block.{block_index}."
      lowerCAmelCase_       :    List[str]								=			f"resnet_block.{groups[-3]}.conv1d_{conv_index}.{groups[-1]}"
      lowerCAmelCase_       :    Any								=			prefix + resnet_block
      lowerCAmelCase_       :    Optional[int]								=			re_encoder_block_resnet.sub(lowerCAmelCase__   ,       lowerCAmelCase__						)

    elif re_encoder_block_proj_out.fullmatch(lowerCAmelCase__						):
      lowerCAmelCase_       :    Any								=			re_encoder_block_proj_out.match(lowerCAmelCase__						)
      lowerCAmelCase_       :    List[str]								=			regex_match.groups()
      lowerCAmelCase_       :    Any								=			f"encoders.{groups[0]}.level_blocks.{groups[1]}.proj_out.{groups[-1]}"
      lowerCAmelCase_       :    int								=			re_encoder_block_proj_out.sub(lowerCAmelCase__   ,       lowerCAmelCase__						)

    # rename vqvae.decoder keys
    elif re_decoder_block_conv_out.fullmatch(lowerCAmelCase__						):
      lowerCAmelCase_       :    str								=			re_decoder_block_conv_out.match(lowerCAmelCase__						)
      lowerCAmelCase_       :    Optional[Any]								=			regex_match.groups()
      lowerCAmelCase_       :    List[str]								=			int(groups[2]						) * 2 + int(groups[3]						) - 2
      lowerCAmelCase_       :    Tuple								=			f"decoders.{groups[0]}.level_blocks.{groups[1]}.upsample_block.{block_index}.{groups[-1]}"
      lowerCAmelCase_       :    Optional[int]								=			re_decoder_block_conv_out.sub(lowerCAmelCase__   ,       lowerCAmelCase__						)

    elif re_decoder_block_resnet.fullmatch(lowerCAmelCase__						):
      lowerCAmelCase_       :    Tuple								=			re_decoder_block_resnet.match(lowerCAmelCase__						)
      lowerCAmelCase_       :    Optional[int]								=			regex_match.groups()
      lowerCAmelCase_       :    Tuple								=			int(groups[2]						) * 2 + int(groups[3]						) - 2
      lowerCAmelCase_       :    Tuple								=			{'1': 1, '3': 2}[groups[-2]]
      lowerCAmelCase_       :    Tuple								=			f"decoders.{groups[0]}.level_blocks.{groups[1]}.upsample_block.{block_index}."
      lowerCAmelCase_       :    List[str]								=			f"resnet_block.{groups[-3]}.conv1d_{conv_index}.{groups[-1]}"
      lowerCAmelCase_       :    Any								=			prefix + resnet_block
      lowerCAmelCase_       :    Tuple								=			re_decoder_block_resnet.sub(lowerCAmelCase__   ,       lowerCAmelCase__						)

    elif re_decoder_block_proj_in.fullmatch(lowerCAmelCase__						):
      lowerCAmelCase_       :    Optional[int]								=			re_decoder_block_proj_in.match(lowerCAmelCase__						)
      lowerCAmelCase_       :    List[str]								=			regex_match.groups()
      lowerCAmelCase_       :    Optional[Any]								=			f"decoders.{groups[0]}.level_blocks.{groups[1]}.proj_in.{groups[-1]}"
      lowerCAmelCase_       :    Any								=			re_decoder_block_proj_in.sub(lowerCAmelCase__   ,       lowerCAmelCase__						)

    # rename prior cond.model to upsampler.upsample_block and resnet
    elif re_prior_cond_conv_out.fullmatch(lowerCAmelCase__						):
      lowerCAmelCase_       :    List[str]								=			re_prior_cond_conv_out.match(lowerCAmelCase__						)
      lowerCAmelCase_       :    Tuple								=			regex_match.groups()
      lowerCAmelCase_       :    str								=			int(groups[1]						) * 2 + int(groups[2]						) - 2
      lowerCAmelCase_       :    int								=			f"conditioner_blocks.upsampler.upsample_block.{block_index}.{groups[-1]}"
      lowerCAmelCase_       :    List[Any]								=			re_prior_cond_conv_out.sub(lowerCAmelCase__   ,       lowerCAmelCase__						)

    elif re_prior_cond_resnet.fullmatch(lowerCAmelCase__						):
      lowerCAmelCase_       :    Optional[Any]								=			re_prior_cond_resnet.match(lowerCAmelCase__						)
      lowerCAmelCase_       :    Optional[Any]								=			regex_match.groups()
      lowerCAmelCase_       :    str								=			int(groups[1]						) * 2 + int(groups[2]						) - 2
      lowerCAmelCase_       :    int								=			{'1': 1, '3': 2}[groups[-2]]
      lowerCAmelCase_       :    str								=			f"conditioner_blocks.upsampler.upsample_block.{block_index}."
      lowerCAmelCase_       :    Union[str, Any]								=			f"resnet_block.{groups[-3]}.conv1d_{conv_index}.{groups[-1]}"
      lowerCAmelCase_       :    Tuple								=			prefix + resnet_block
      lowerCAmelCase_       :    List[str]								=			re_prior_cond_resnet.sub(lowerCAmelCase__   ,       lowerCAmelCase__						)

    elif re_prior_cond_proj_in.fullmatch(lowerCAmelCase__						):
      lowerCAmelCase_       :    Union[str, Any]								=			re_prior_cond_proj_in.match(lowerCAmelCase__						)
      lowerCAmelCase_       :    List[str]								=			regex_match.groups()
      lowerCAmelCase_       :    str								=			f"conditioner_blocks.upsampler.proj_in.{groups[-1]}"
      lowerCAmelCase_       :    List[Any]								=			re_prior_cond_proj_in.sub(lowerCAmelCase__   ,       lowerCAmelCase__						)

    # keep original key
    else:
      lowerCAmelCase_       :    Optional[int]								=			original_key

    lowerCAmelCase_       :    List[Any]								=			replace_key(lowerCAmelCase__						)

    if f"{key_prefix}.{key}" not in model_state_dict or key is None:
      print(f"failed converting {original_key} to {key}, does not match"						)

    # handle missmatched shape
    elif value.shape != model_state_dict[f"{key_prefix}.{key}"].shape:
      lowerCAmelCase_       :    Tuple								=			model_state_dict[f"{key_prefix}.{key}"]
      print(f"{original_key}-> {key} : \nshape {val.shape} and { value.shape}, do not match"						)
      lowerCAmelCase_       :    Optional[int]								=			original_key

    lowerCAmelCase_       :    Union[str, Any]								=			original_key
    lowerCAmelCase_       :    Optional[int]								=			value

  return new_dict

@torch.no_grad()
def     UpperCamelCase_   (   lowerCAmelCase__				:    Any=None   ,       lowerCAmelCase__				:    Any=None						)						->       Optional[Any]:

  """simple docstring"""




  for file in MODEL_MAPPING[model_name]:
    if not os.path.isfile(f"{pytorch_dump_folder_path}/{file.split('/'						)[-1]}"						):
      lowerCAmelCase_       :    List[str]								=			requests.get(f"{PREFIX}{file}"   ,       allow_redirects=lowerCAmelCase__						)
      os.makedirs(f"{pytorch_dump_folder_path}/"   ,       exist_ok=lowerCAmelCase__						)
      open(f"{pytorch_dump_folder_path}/{file.split('/'						)[-1]}"   ,       'wb'						).write(r.content						)

  lowerCAmelCase_       :    Optional[int]								=			MODEL_MAPPING[model_name.split('/'						)[-1]]

  lowerCAmelCase_       :    Tuple								=			JukeboxConfig.from_pretrained(lowerCAmelCase__						)
  lowerCAmelCase_       :    Tuple								=			JukeboxModel(lowerCAmelCase__						)

  lowerCAmelCase_       :    int								=			[]
  lowerCAmelCase_       :    Dict								=			{}
  for i, dict_name in enumerate(lowerCAmelCase__						):
    lowerCAmelCase_       :    Optional[int]								=			torch.load(f"{pytorch_dump_folder_path}/{dict_name.split('/'						)[-1]}"						)['model']

    lowerCAmelCase_       :    List[str]								=			{}
    for k in old_dic.keys():
      if k.endswith('.b'						):
        lowerCAmelCase_       :    List[Any]								=			old_dic[k]
      elif k.endswith('.w'						):
        lowerCAmelCase_       :    str								=			old_dic[k]
      elif "level_2" not in dict_name and "cond.model." in k:
        lowerCAmelCase_       :    Dict								=			old_dic[k]
      else:
        lowerCAmelCase_       :    List[str]								=			old_dic[k]

    lowerCAmelCase_       :    Dict								=			'vqvae' if i == 0 else f"priors.{3 - i}"
    lowerCAmelCase_       :    Dict								=			fix_jukebox_keys(lowerCAmelCase__   ,       model.state_dict()   ,       lowerCAmelCase__   ,       lowerCAmelCase__						)
    weight_dict.append(lowerCAmelCase__						)

  lowerCAmelCase_       :    Optional[Any]								=			weight_dict.pop(0						)
  model.vqvae.load_state_dict(lowerCAmelCase__						)
  for i in range(len(lowerCAmelCase__						)						):
    model.priors[i].load_state_dict(weight_dict[2 - i]						)

  Path(lowerCAmelCase__						).mkdir(exist_ok=lowerCAmelCase__						)
  with open(f"{pytorch_dump_folder_path}/mapping.json"   ,       'w'						) as txtfile:
    json.dump(lowerCAmelCase__   ,       lowerCAmelCase__						)

  print(f"Saving model {model_name} to {pytorch_dump_folder_path}"						)
  model.save_pretrained(lowerCAmelCase__						)

  return weight_dict


if __name__ == "__main__":
     lowercase__		:    Any				   =     argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         """--model_name""",
         default="""jukebox-5b-lyrics""",
         type=str,
         help="""Name of the model you'd like to convert.""",
     )
     parser.add_argument(
         """--pytorch_dump_folder_path""",
         default="""jukebox-5b-lyrics-converted""",
         type=str,
         help="""Path to the output PyTorch model directory.""",
     )
     lowercase__		:    Dict				   =     parser.parse_args()
     convert_openai_checkpoint(args.model_name, args.pytorch_dump_folder_path)










"""simple docstring"""







import os
import tempfile
import unittest

from transformers import FlaubertConfig, is_torch_available
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import (
         FlaubertForMultipleChoice,
         FlaubertForQuestionAnswering,
         FlaubertForQuestionAnsweringSimple,
         FlaubertForSequenceClassification,
         FlaubertForTokenClassification,
         FlaubertModel,
         FlaubertWithLMHeadModel,
     )
     from transformers.models.flaubert.modeling_flaubert import FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST







class 				UpperCamelCase__       (     lowercase_  ):




      """simple docstring"""




      def __init__(				self    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				Tuple   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]=1_3   ,      SCREAMING_SNAKE_CASE_    :				Dict=7   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]=True   ,      SCREAMING_SNAKE_CASE_    :				List[str]=True   ,      SCREAMING_SNAKE_CASE_    :				Dict=True   ,      SCREAMING_SNAKE_CASE_    :				Dict=True   ,      SCREAMING_SNAKE_CASE_    :				List[str]=True   ,      SCREAMING_SNAKE_CASE_    :				str=False   ,      SCREAMING_SNAKE_CASE_    :				str=False   ,      SCREAMING_SNAKE_CASE_    :				Any=False   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]=2   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]=9_9   ,      SCREAMING_SNAKE_CASE_    :				Dict=0   ,      SCREAMING_SNAKE_CASE_    :				int=3_2   ,      SCREAMING_SNAKE_CASE_    :				Dict=5   ,      SCREAMING_SNAKE_CASE_    :				Dict=4   ,      SCREAMING_SNAKE_CASE_    :				List[Any]=0.1   ,      SCREAMING_SNAKE_CASE_    :				List[Any]=0.1   ,      SCREAMING_SNAKE_CASE_    :				List[str]=5_1_2   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]=1_2   ,      SCREAMING_SNAKE_CASE_    :				Tuple=2   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]=0.02   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]=3   ,      SCREAMING_SNAKE_CASE_    :				int=4   ,      SCREAMING_SNAKE_CASE_    :				List[Any]="last"   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]=None   ,      SCREAMING_SNAKE_CASE_    :				str=None   ,      ):
        lowerCAmelCase_       :    Tuple								=			parent
        lowerCAmelCase_       :    Tuple								=			batch_size
        lowerCAmelCase_       :    str								=			seq_length
        lowerCAmelCase_       :    List[Any]								=			is_training
        lowerCAmelCase_       :    Optional[int]								=			use_input_lengths
        lowerCAmelCase_       :    Union[str, Any]								=			use_token_type_ids
        lowerCAmelCase_       :    str								=			use_labels
        lowerCAmelCase_       :    str								=			gelu_activation
        lowerCAmelCase_       :    str								=			sinusoidal_embeddings
        lowerCAmelCase_       :    List[Any]								=			causal
        lowerCAmelCase_       :    Union[str, Any]								=			asm
        lowerCAmelCase_       :    Union[str, Any]								=			n_langs
        lowerCAmelCase_       :    Optional[Any]								=			vocab_size
        lowerCAmelCase_       :    Any								=			n_special
        lowerCAmelCase_       :    Union[str, Any]								=			hidden_size
        lowerCAmelCase_       :    str								=			num_hidden_layers
        lowerCAmelCase_       :    Tuple								=			num_attention_heads
        lowerCAmelCase_       :    Optional[Any]								=			hidden_dropout_prob
        lowerCAmelCase_       :    List[Any]								=			attention_probs_dropout_prob
        lowerCAmelCase_       :    Dict								=			max_position_embeddings
        lowerCAmelCase_       :    Union[str, Any]								=			type_vocab_size
        lowerCAmelCase_       :    List[Any]								=			type_sequence_label_size
        lowerCAmelCase_       :    str								=			initializer_range
        lowerCAmelCase_       :    Dict								=			num_labels
        lowerCAmelCase_       :    Union[str, Any]								=			num_choices
        lowerCAmelCase_       :    Union[str, Any]								=			summary_type
        lowerCAmelCase_       :    Optional[Any]								=			use_proj
        lowerCAmelCase_       :    List[Any]								=			scope



      def 					SCREAMING_SNAKE_CASE__			(				self    :				List[Any]	):
        lowerCAmelCase_       :    List[Any]								=			ids_tensor([self.batch_size, self.seq_length]   ,      self.vocab_size	)
        lowerCAmelCase_       :    str								=			random_attention_mask([self.batch_size, self.seq_length]	)

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

        lowerCAmelCase_       :    Any								=			None
        if self.use_token_type_ids:
          lowerCAmelCase_       :    Union[str, Any]								=			ids_tensor([self.batch_size, self.seq_length]   ,      self.n_langs	)

        lowerCAmelCase_       :    Dict								=			None
        lowerCAmelCase_       :    Optional[int]								=			None
        lowerCAmelCase_       :    List[str]								=			None
        if self.use_labels:
          lowerCAmelCase_       :    Tuple								=			ids_tensor([self.batch_size]   ,      self.type_sequence_label_size	)
          lowerCAmelCase_       :    Tuple								=			ids_tensor([self.batch_size, self.seq_length]   ,      self.num_labels	)
          lowerCAmelCase_       :    Any								=			ids_tensor([self.batch_size]   ,      2	).float()
          lowerCAmelCase_       :    List[Any]								=			ids_tensor([self.batch_size]   ,      self.num_choices	)

        lowerCAmelCase_       :    Union[str, Any]								=			self.get_config()

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



      def 					SCREAMING_SNAKE_CASE__			(				self    :				List[Any]	):
        return 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   ,      )



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				Tuple   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				str   ,      SCREAMING_SNAKE_CASE_    :				str   ,      ):
        lowerCAmelCase_       :    Union[str, Any]								=			FlaubertModel(config=SCREAMING_SNAKE_CASE_	)
        model.to(SCREAMING_SNAKE_CASE_	)
        model.eval()
        lowerCAmelCase_       :    Any								=			model(SCREAMING_SNAKE_CASE_   ,      lengths=SCREAMING_SNAKE_CASE_   ,      langs=SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    Any								=			model(SCREAMING_SNAKE_CASE_   ,      langs=SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    Optional[Any]								=			model(SCREAMING_SNAKE_CASE_	)
        self.parent.assertEqual(result.last_hidden_state.shape   ,      (self.batch_size, self.seq_length, self.hidden_size)	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Any   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				str   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				List[str]   ,      SCREAMING_SNAKE_CASE_    :				Tuple   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]   ,      SCREAMING_SNAKE_CASE_    :				List[Any]   ,      ):
        lowerCAmelCase_       :    Any								=			FlaubertWithLMHeadModel(SCREAMING_SNAKE_CASE_	)
        model.to(SCREAMING_SNAKE_CASE_	)
        model.eval()

        lowerCAmelCase_       :    Dict								=			model(SCREAMING_SNAKE_CASE_   ,      token_type_ids=SCREAMING_SNAKE_CASE_   ,      labels=SCREAMING_SNAKE_CASE_	)
        self.parent.assertEqual(result.loss.shape   ,      ()	)
        self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.seq_length, self.vocab_size)	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				int   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				List[str]   ,      SCREAMING_SNAKE_CASE_    :				List[str]   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]   ,      SCREAMING_SNAKE_CASE_    :				Tuple   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]   ,      SCREAMING_SNAKE_CASE_    :				int   ,      ):
        lowerCAmelCase_       :    Tuple								=			FlaubertForQuestionAnsweringSimple(SCREAMING_SNAKE_CASE_	)
        model.to(SCREAMING_SNAKE_CASE_	)
        model.eval()

        lowerCAmelCase_       :    Union[str, Any]								=			model(SCREAMING_SNAKE_CASE_	)

        lowerCAmelCase_       :    Optional[int]								=			model(SCREAMING_SNAKE_CASE_   ,      start_positions=SCREAMING_SNAKE_CASE_   ,      end_positions=SCREAMING_SNAKE_CASE_	)
        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 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				Any   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				Any   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				Any   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				List[str]   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]   ,      ):
        lowerCAmelCase_       :    Optional[int]								=			FlaubertForQuestionAnswering(SCREAMING_SNAKE_CASE_	)
        model.to(SCREAMING_SNAKE_CASE_	)
        model.eval()

        lowerCAmelCase_       :    int								=			model(SCREAMING_SNAKE_CASE_	)

        lowerCAmelCase_       :    Any								=			model(
            SCREAMING_SNAKE_CASE_   ,      start_positions=SCREAMING_SNAKE_CASE_   ,      end_positions=SCREAMING_SNAKE_CASE_   ,      cls_index=SCREAMING_SNAKE_CASE_   ,      is_impossible=SCREAMING_SNAKE_CASE_   ,      p_mask=SCREAMING_SNAKE_CASE_   ,      )

        lowerCAmelCase_       :    Optional[int]								=			model(
            SCREAMING_SNAKE_CASE_   ,      start_positions=SCREAMING_SNAKE_CASE_   ,      end_positions=SCREAMING_SNAKE_CASE_   ,      cls_index=SCREAMING_SNAKE_CASE_   ,      is_impossible=SCREAMING_SNAKE_CASE_   ,      )

        ((lowerCAmelCase_)						,)       :    int								=			result_with_labels.to_tuple()

        lowerCAmelCase_       :    int								=			model(SCREAMING_SNAKE_CASE_   ,      start_positions=SCREAMING_SNAKE_CASE_   ,      end_positions=SCREAMING_SNAKE_CASE_	)

        ((lowerCAmelCase_)						,)       :    Dict								=			result_with_labels.to_tuple()

        self.parent.assertEqual(result_with_labels.loss.shape   ,      ()	)
        self.parent.assertEqual(result.start_top_log_probs.shape   ,      (self.batch_size, model.config.start_n_top)	)
        self.parent.assertEqual(result.start_top_index.shape   ,      (self.batch_size, model.config.start_n_top)	)
        self.parent.assertEqual(
            result.end_top_log_probs.shape   ,      (self.batch_size, model.config.start_n_top * model.config.end_n_top)	)
        self.parent.assertEqual(
            result.end_top_index.shape   ,      (self.batch_size, model.config.start_n_top * model.config.end_n_top)	)
        self.parent.assertEqual(result.cls_logits.shape   ,      (self.batch_size,)	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Union[str, Any]   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]   ,      SCREAMING_SNAKE_CASE_    :				Any   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				List[Any]   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      ):
        lowerCAmelCase_       :    Optional[int]								=			FlaubertForSequenceClassification(SCREAMING_SNAKE_CASE_	)
        model.to(SCREAMING_SNAKE_CASE_	)
        model.eval()

        lowerCAmelCase_       :    Any								=			model(SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    List[Any]								=			model(SCREAMING_SNAKE_CASE_   ,      labels=SCREAMING_SNAKE_CASE_	)

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



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Any   ,      SCREAMING_SNAKE_CASE_    :				List[str]   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				Tuple   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				Any   ,      SCREAMING_SNAKE_CASE_    :				Tuple   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      ):
        lowerCAmelCase_       :    List[Any]								=			self.num_labels
        lowerCAmelCase_       :    Optional[Any]								=			FlaubertForTokenClassification(SCREAMING_SNAKE_CASE_	)
        model.to(SCREAMING_SNAKE_CASE_	)
        model.eval()

        lowerCAmelCase_       :    int								=			model(SCREAMING_SNAKE_CASE_   ,      attention_mask=SCREAMING_SNAKE_CASE_   ,      labels=SCREAMING_SNAKE_CASE_	)
        self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.seq_length, self.num_labels)	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				str   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				List[Any]   ,      SCREAMING_SNAKE_CASE_    :				List[Any]   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]   ,      SCREAMING_SNAKE_CASE_    :				List[Any]   ,      SCREAMING_SNAKE_CASE_    :				List[str]   ,      ):
        lowerCAmelCase_       :    Dict								=			self.num_choices
        lowerCAmelCase_       :    Optional[Any]								=			FlaubertForMultipleChoice(config=SCREAMING_SNAKE_CASE_	)
        model.to(SCREAMING_SNAKE_CASE_	)
        model.eval()
        lowerCAmelCase_       :    Tuple								=			input_ids.unsqueeze(1	).expand(-1   ,      self.num_choices   ,      -1	).contiguous()
        lowerCAmelCase_       :    Union[str, Any]								=			token_type_ids.unsqueeze(1	).expand(-1   ,      self.num_choices   ,      -1	).contiguous()
        lowerCAmelCase_       :    List[Any]								=			input_mask.unsqueeze(1	).expand(-1   ,      self.num_choices   ,      -1	).contiguous()
        lowerCAmelCase_       :    Optional[Any]								=			model(
            SCREAMING_SNAKE_CASE_   ,      attention_mask=SCREAMING_SNAKE_CASE_   ,      token_type_ids=SCREAMING_SNAKE_CASE_   ,      labels=SCREAMING_SNAKE_CASE_   ,      )
        self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.num_choices)	)





      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[int]	):
        lowerCAmelCase_       :    List[Any]								=			self.prepare_config_and_inputs()
        (
            (
            lowerCAmelCase_
        )						,(
            lowerCAmelCase_
        )						,(
            lowerCAmelCase_
        )						,(
            lowerCAmelCase_
        )						,(
            lowerCAmelCase_
        )						,(
            lowerCAmelCase_
        )						,(
            lowerCAmelCase_
        )						,(
            lowerCAmelCase_
        )						,(
            lowerCAmelCase_
        )						,
        )       :    Dict								=			config_and_inputs
        lowerCAmelCase_       :    List[str]								=			{
            'input_ids': input_ids,
            'token_type_ids': token_type_ids,
            'lengths': input_lengths,
            'attention_mask': input_mask,
        }
        return config, inputs_dict







@require_torch
class 				UpperCamelCase__       (     lowercase_, lowercase_, unittest.TestCase  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE      =					(
          (
              FlaubertModel,
              FlaubertWithLMHeadModel,
              FlaubertForQuestionAnswering,
              FlaubertForQuestionAnsweringSimple,
              FlaubertForSequenceClassification,
              FlaubertForTokenClassification,
              FlaubertForMultipleChoice,
          )
          if is_torch_available()
          else ()
      )
      _SCREAMING_SNAKE_CASE      =					(
          {
              """feature-extraction""": FlaubertModel,
              """fill-mask""": FlaubertWithLMHeadModel,
              """question-answering""": FlaubertForQuestionAnsweringSimple,
              """text-classification""": FlaubertForSequenceClassification,
              """token-classification""": FlaubertForTokenClassification,
              """zero-shot""": FlaubertForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )



      def 					SCREAMING_SNAKE_CASE__			(				self    :				int   ,      SCREAMING_SNAKE_CASE_    :				Tuple   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]   ,      SCREAMING_SNAKE_CASE_    :				str   ,      SCREAMING_SNAKE_CASE_    :				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 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]=False	):
        lowerCAmelCase_       :    Any								=			super()._prepare_for_class(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_   ,      return_labels=SCREAMING_SNAKE_CASE_	)

        if return_labels:
          if model_class.__name__ == "FlaubertForQuestionAnswering":
            lowerCAmelCase_       :    Tuple								=			torch.zeros(
                self.model_tester.batch_size   ,      dtype=torch.long   ,      device=SCREAMING_SNAKE_CASE_	)
            lowerCAmelCase_       :    int								=			torch.zeros(
                self.model_tester.batch_size   ,      dtype=torch.long   ,      device=SCREAMING_SNAKE_CASE_	)

        return inputs_dict



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Any	):
        lowerCAmelCase_       :    Tuple								=			FlaubertModelTester(self	)
        lowerCAmelCase_       :    str								=			ConfigTester(self   ,      config_class=SCREAMING_SNAKE_CASE_   ,      emb_dim=3_7	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Any	):
        self.config_tester.run_common_tests()



      def 					SCREAMING_SNAKE_CASE__			(				self    :				List[str]	):
        lowerCAmelCase_       :    List[Any]								=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_flaubert_model(*SCREAMING_SNAKE_CASE_	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[int]	):
        lowerCAmelCase_       :    str								=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_flaubert_lm_head(*SCREAMING_SNAKE_CASE_	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				List[str]	):
        lowerCAmelCase_       :    List[Any]								=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_flaubert_simple_qa(*SCREAMING_SNAKE_CASE_	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[int]	):
        lowerCAmelCase_       :    List[str]								=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_flaubert_qa(*SCREAMING_SNAKE_CASE_	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Union[str, Any]	):
        lowerCAmelCase_       :    Optional[Any]								=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_flaubert_sequence_classif(*SCREAMING_SNAKE_CASE_	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[Any]	):
        lowerCAmelCase_       :    Dict								=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_flaubert_token_classif(*SCREAMING_SNAKE_CASE_	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[int]	):
        lowerCAmelCase_       :    List[Any]								=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_flaubert_multiple_choice(*SCREAMING_SNAKE_CASE_	)



      @slow
      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[Any]	):
        for model_name in FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          lowerCAmelCase_       :    Optional[Any]								=			FlaubertModel.from_pretrained(SCREAMING_SNAKE_CASE_	)
          self.assertIsNotNone(SCREAMING_SNAKE_CASE_	)





      @slow
      @require_torch_gpu
      def 					SCREAMING_SNAKE_CASE__			(				self    :				Tuple	):
        lowerCAmelCase_						,lowerCAmelCase_       :    Any								=			self.model_tester.prepare_config_and_inputs_for_common()
        for model_class in self.all_model_classes:
          # FlauBertForMultipleChoice behaves incorrectly in JIT environments.
          if model_class == FlaubertForMultipleChoice:
            return

          lowerCAmelCase_       :    int								=			True
          lowerCAmelCase_       :    Union[str, Any]								=			model_class(config=SCREAMING_SNAKE_CASE_	)

          lowerCAmelCase_       :    Any								=			self._prepare_for_class(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	)
          lowerCAmelCase_       :    int								=			torch.jit.trace(
              SCREAMING_SNAKE_CASE_   ,      (inputs_dict['input_ids'].to('cpu'	), inputs_dict['attention_mask'].to('cpu'	))	)

          with tempfile.TemporaryDirectory() as tmp:
            torch.jit.save(SCREAMING_SNAKE_CASE_   ,      os.path.join(SCREAMING_SNAKE_CASE_   ,      'traced_model.pt'	)	)
            lowerCAmelCase_       :    Optional[Any]								=			torch.jit.load(os.path.join(SCREAMING_SNAKE_CASE_   ,      'traced_model.pt'	)   ,      map_location=SCREAMING_SNAKE_CASE_	)
            loaded(inputs_dict['input_ids'].to(SCREAMING_SNAKE_CASE_	)   ,      inputs_dict['attention_mask'].to(SCREAMING_SNAKE_CASE_	)	)







@require_torch
class 				UpperCamelCase__       (     unittest.TestCase  ):




      """simple docstring"""




      @slow
      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[Any]	):
        lowerCAmelCase_       :    int								=			FlaubertModel.from_pretrained('flaubert/flaubert_base_cased'	)
        lowerCAmelCase_       :    List[Any]								=			torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]	)
        with torch.no_grad():
          lowerCAmelCase_       :    Optional[Any]								=			model(SCREAMING_SNAKE_CASE_	)[0]
        lowerCAmelCase_       :    Optional[Any]								=			torch.Size((1, 1_1, 7_6_8)	)
        self.assertEqual(output.shape   ,      SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    Optional[int]								=			torch.tensor(
            [[[-2.62_51, -1.42_98, -0.02_27], [-2.85_10, -1.63_87, 0.22_58], [-2.81_14, -1.18_32, -0.30_66]]]	)

        self.assertTrue(torch.allclose(output[:, :3, :3]   ,      SCREAMING_SNAKE_CASE_   ,      atol=1E-4	)	)





