

'''simple docstring'''

import argparse
import json
import os
import re
import shutil

import torch

from transformers import BioGptConfig, BioGptForCausalLM
from transformers.models.biogpt.tokenization_biogpt import VOCAB_FILES_NAMES
from transformers.tokenization_utils_base import TOKENIZER_CONFIG_FILE
from transformers.utils import WEIGHTS_NAME, logging


logging.set_verbosity_warning()

A_       						=							2





class 							_snake_case :
   def __init__( self			:      List[str]     ,*,  # begin keyword-only arguments
       SCREAMING_SNAKE_CASE__			:      List[Any]="<s>"     ,SCREAMING_SNAKE_CASE__			:      Union[str, Any]="<pad>"     ,SCREAMING_SNAKE_CASE__			:      Any="</s>"     ,SCREAMING_SNAKE_CASE__			:      Any="<unk>"     ,SCREAMING_SNAKE_CASE__			:      Optional[Any]=None     ,):
     SCREAMING_SNAKE_CASE      ,    SCREAMING_SNAKE_CASE      ,    SCREAMING_SNAKE_CASE      ,    SCREAMING_SNAKE_CASE:Dict						      =					bos, unk, pad, eos
     SCREAMING_SNAKE_CASE:str						      =					[]
     SCREAMING_SNAKE_CASE:Optional[int]						      =					[]
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					{}
     SCREAMING_SNAKE_CASE:str						      =					self.add_symbol(SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					self.add_symbol(SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Optional[int]						      =					self.add_symbol(SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Dict						      =					self.add_symbol(SCREAMING_SNAKE_CASE__							)
     if extra_special_symbols:
       for s in extra_special_symbols:
         self.add_symbol(SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					len(self.symbols							)
   def __eq__( self			:      List[str]     ,SCREAMING_SNAKE_CASE__			:      Any							):
     return self.indices == other.indices
   def __getitem__( self			:      Tuple     ,SCREAMING_SNAKE_CASE__			:      Optional[int]							):
     if idx < len(self.symbols							):
       return self.symbols[idx]
     return self.unk_word
   def __len__( self			:      Optional[int]							):
     return len(self.symbols							)
   def __contains__( self			:      Tuple     ,SCREAMING_SNAKE_CASE__			:      Optional[int]							):
     return sym in self.indices
   @classmethod
   def     __UpperCamelCase						( cls			:      Dict     ,SCREAMING_SNAKE_CASE__			:      List[str]							):
     SCREAMING_SNAKE_CASE:List[Any]						      =					cls()
     d.add_from_file(SCREAMING_SNAKE_CASE__							)
     return d
   def     __UpperCamelCase						( self			:      Any     ,SCREAMING_SNAKE_CASE__			:      Optional[int]     ,SCREAMING_SNAKE_CASE__			:      Dict=1     ,SCREAMING_SNAKE_CASE__			:      str=False							):
     if word in self.indices and not overwrite:
       SCREAMING_SNAKE_CASE:Tuple						      =					self.indices[word]
       SCREAMING_SNAKE_CASE:Optional[int]						      =					self.count[idx] + n
       return idx
     else:
       SCREAMING_SNAKE_CASE:Optional[Any]						      =					len(self.symbols							)
       SCREAMING_SNAKE_CASE:str						      =					idx
       self.symbols.append(SCREAMING_SNAKE_CASE__							)
       self.count.append(SCREAMING_SNAKE_CASE__							)
       return idx
   def     __UpperCamelCase						( self			:      Optional[int]     ,SCREAMING_SNAKE_CASE__			:      List[Any]							):
     return 0







   def     __UpperCamelCase						( self			:      Union[str, Any]     ,SCREAMING_SNAKE_CASE__			:      int							):
     if isinstance(SCREAMING_SNAKE_CASE__     ,SCREAMING_SNAKE_CASE__							):
       try:
         with open(SCREAMING_SNAKE_CASE__     ,"r"     ,encoding="utf-8"							) as fd:
           self.add_from_file(SCREAMING_SNAKE_CASE__							)
       except FileNotFoundError as fnfe:
         raise fnfe
       except UnicodeError:
         raise Exception("Incorrect encoding detected in {}, please rebuild the dataset".format(SCREAMING_SNAKE_CASE__							)							)
       return

     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					f.readlines()
     SCREAMING_SNAKE_CASE:int						      =					self._load_meta(SCREAMING_SNAKE_CASE__							)

     for line in lines[indices_start_line:]:
       try:
         SCREAMING_SNAKE_CASE      ,    SCREAMING_SNAKE_CASE:Optional[Any]						      =					line.rstrip().rsplit(" "     ,1							)
         if field == "#fairseq:overwrite":
           SCREAMING_SNAKE_CASE:Optional[int]						      =					True
           SCREAMING_SNAKE_CASE      ,    SCREAMING_SNAKE_CASE:Optional[Any]						      =					line.rsplit(" "     ,1							)
         else:
           SCREAMING_SNAKE_CASE:Dict						      =					False
         SCREAMING_SNAKE_CASE:str						      =					int(SCREAMING_SNAKE_CASE__							)
         SCREAMING_SNAKE_CASE:str						      =					line
         if word in self and not overwrite:
           raise RuntimeError(
               "Duplicate word found when loading Dictionary: '{}'. "
               "Duplicate words can overwrite earlier ones by adding the "
               "#fairseq:overwrite flag at the end of the corresponding row "
               "in the dictionary file. If using the Camembert model, please "
               "download an updated copy of the model file.".format(SCREAMING_SNAKE_CASE__							)							)
         self.add_symbol(SCREAMING_SNAKE_CASE__     ,n=SCREAMING_SNAKE_CASE__     ,overwrite=SCREAMING_SNAKE_CASE__							)
       except ValueError:
         raise ValueError("Incorrect dictionary format, expected '<token> <cnt> [flags]'"							)





def    A_  ( snake_case ):
  # (1) remove word breaking symbol, (2) add word ending symbol where the word is not broken up,
  # e.g.: d = {'le@@': 5, 'tt@@': 6, 'er': 7} => {'le': 5, 'tt': 6, 'er</w>': 7}
  SCREAMING_SNAKE_CASE:List[Any]						      =					dict((re.sub(r"@@$"      ,       ""      ,       snake_case ), v) if k.endswith("@@" ) else (re.sub(r"$"      ,       "</w>"      ,       snake_case ), v) for k, v in d.items() )
  SCREAMING_SNAKE_CASE:List[Any]						      =					"<s> <pad> </s> <unk>".split()
  # restore the special tokens
  for k in keep_keys:
    del da[F'''{k}</w>''']
    SCREAMING_SNAKE_CASE:List[Any]						      =					d[k]  # restore
  return da





def    A_  ( snake_case      ,       snake_case ):
  # prep
  if not os.path.exists(snake_case ):
    raise ValueError(F'''path {biogpt_checkpoint_path} does not exist!''' )
  os.makedirs(snake_case      ,       exist_ok=snake_case )
  print(F'''Writing results to {pytorch_dump_folder_path}''' )

  # handle various types of models

  SCREAMING_SNAKE_CASE:List[Any]						      =					os.path.join(snake_case      ,       "checkpoint.pt" )
  if not os.path.isfile(snake_case ):
    raise ValueError(F'''path to the file {checkpoint_file} does not exist!''' )
  SCREAMING_SNAKE_CASE:str						      =					torch.load(snake_case      ,       map_location="cpu" )

  SCREAMING_SNAKE_CASE:List[Any]						      =					chkpt["cfg"]["model"]

  # dicts
  SCREAMING_SNAKE_CASE:List[str]						      =					os.path.join(snake_case      ,       "dict.txt" )
  if not os.path.isfile(snake_case ):
    raise ValueError(F'''path to the file {dict_file} does not exist!''' )
  SCREAMING_SNAKE_CASE:Optional[Any]						      =					Dictionary.load(snake_case )
  SCREAMING_SNAKE_CASE:Any						      =					rewrite_dict_keys(src_dict.indices )
  SCREAMING_SNAKE_CASE:Dict						      =					len(snake_case )
  SCREAMING_SNAKE_CASE:Tuple						      =					os.path.join(snake_case      ,       VOCAB_FILES_NAMES["vocab_file"] )
  print(F'''Generating {src_vocab_file} of {src_vocab_size} records''' )
  with open(snake_case      ,       "w"      ,       encoding="utf-8" ) as f:
    f.write(json.dumps(snake_case      ,       ensure_ascii=snake_case      ,       indent=snake_case ) )

  # merges_file (bpecodes)
  SCREAMING_SNAKE_CASE:Optional[Any]						      =					os.path.join(snake_case      ,       "bpecodes" )
  if not os.path.isfile(snake_case ):
    raise ValueError(F'''path to the file {bpecodes_file} does not exist!''' )

  SCREAMING_SNAKE_CASE:List[Any]						      =					os.path.join(snake_case      ,       VOCAB_FILES_NAMES["merges_file"] )
  shutil.copyfile(snake_case      ,       snake_case )

  # model config
  SCREAMING_SNAKE_CASE:Optional[Any]						      =					os.path.join(snake_case      ,       "config.json" )

  SCREAMING_SNAKE_CASE:Tuple						      =					{
      "activation_dropout": args["activation_dropout"],
      "architectures": ["BioGptForCausalLM"],
      "attention_probs_dropout_prob": args["attention_dropout"],
      "bos_token_id": 0,
      "eos_token_id": 2,
      "hidden_act": args["activation_fn"],
      "hidden_dropout_prob": args["dropout"],
      "hidden_size": args["decoder_embed_dim"],
      "initializer_range": 0.02,
      "intermediate_size": args["decoder_ffn_embed_dim"],
      "layer_norm_eps": 1e-12,
      "layerdrop": args["decoder_layerdrop"],
      "max_position_embeddings": args["max_target_positions"],
      "model_type": "biogpt",
      "num_attention_heads": args["decoder_attention_heads"],
      "num_hidden_layers": args["decoder_layers"],
      "pad_token_id": 1,
      "scale_embedding": not args["no_scale_embedding"],
      "tie_word_embeddings": args["share_decoder_input_output_embed"],
      "vocab_size": src_vocab_size,
  }

  # good hparam defaults to start with

  print(F'''Generating {biogpt_model_config_file}''' )
  with open(snake_case      ,       "w"      ,       encoding="utf-8" ) as f:
    f.write(json.dumps(snake_case      ,       ensure_ascii=snake_case      ,       indent=snake_case ) )

  # tokenizer config
  SCREAMING_SNAKE_CASE:Optional[int]						      =					os.path.join(snake_case      ,       snake_case )

  SCREAMING_SNAKE_CASE:Tuple						      =					{
      "bos_token": "<s>",
      "eos_token": "</s>",
      "model_max_length": 1024,
      "pad_token": "<pad>",
      "special_tokens_map_file": None,
      "tokenizer_class": "BioGptTokenizer",
      "unk_token": "<unk>",
  }

  print(F'''Generating {biogpt_tokenizer_config_file}''' )
  with open(snake_case      ,       "w"      ,       encoding="utf-8" ) as f:
    f.write(json.dumps(snake_case      ,       ensure_ascii=snake_case      ,       indent=snake_case ) )

  # model
  SCREAMING_SNAKE_CASE:Dict						      =					chkpt["model"]

  # remove unneeded keys
  SCREAMING_SNAKE_CASE:Optional[int]						      =					[
      "decoder.version",
  ]
  for k in ignore_keys:
    model_state_dict.pop(snake_case      ,       snake_case )

  SCREAMING_SNAKE_CASE:Tuple						      =					list(model_state_dict.keys() )
  for layer_name in layer_names:
    if layer_name.endswith("output_projection.weight" ):
      SCREAMING_SNAKE_CASE:Dict						      =					model_state_dict.pop(snake_case )
    else:
      SCREAMING_SNAKE_CASE:Optional[int]						      =					model_state_dict.pop(snake_case )

  SCREAMING_SNAKE_CASE:Optional[Any]						      =					BioGptConfig.from_pretrained(snake_case )
  SCREAMING_SNAKE_CASE:Tuple						      =					BioGptForCausalLM(snake_case )

  # check that it loads ok
  model_new.load_state_dict(snake_case )

  # save
  SCREAMING_SNAKE_CASE:Tuple						      =					os.path.join(snake_case      ,       snake_case )
  print(F'''Generating {pytorch_weights_dump_path}''' )
  torch.save(snake_case      ,       snake_case )

  print("Conversion is done!" )


if __name__ == "__main__":
    A_       						=							argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        "--biogpt_checkpoint_path",
        default=None,
        type=str,
        required=True,
        help=(
            "Path to the official PyTorch checkpoint file which is expected to reside in the dump dir with dicts,"
            " bpecodes, etc."
        ),
    )
    parser.add_argument(
        "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
    )
    A_       						=							parser.parse_args()
    convert_biogpt_checkpoint_to_pytorch(args.biogpt_checkpoint_path, 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 							_snake_case (       _a						):
   def __init__( self			:      Tuple     ,SCREAMING_SNAKE_CASE__			:      Optional[int]     ,SCREAMING_SNAKE_CASE__			:      Tuple=13     ,SCREAMING_SNAKE_CASE__			:      Optional[int]=7     ,SCREAMING_SNAKE_CASE__			:      List[Any]=True     ,SCREAMING_SNAKE_CASE__			:      List[Any]=True     ,SCREAMING_SNAKE_CASE__			:      str=True     ,SCREAMING_SNAKE_CASE__			:      Union[str, Any]=True     ,SCREAMING_SNAKE_CASE__			:      Any=True     ,SCREAMING_SNAKE_CASE__			:      Optional[Any]=False     ,SCREAMING_SNAKE_CASE__			:      Union[str, Any]=False     ,SCREAMING_SNAKE_CASE__			:      int=False     ,SCREAMING_SNAKE_CASE__			:      int=2     ,SCREAMING_SNAKE_CASE__			:      Any=99     ,SCREAMING_SNAKE_CASE__			:      Optional[int]=0     ,SCREAMING_SNAKE_CASE__			:      Optional[Any]=32     ,SCREAMING_SNAKE_CASE__			:      List[str]=5     ,SCREAMING_SNAKE_CASE__			:      Tuple=4     ,SCREAMING_SNAKE_CASE__			:      Tuple=0.1     ,SCREAMING_SNAKE_CASE__			:      str=0.1     ,SCREAMING_SNAKE_CASE__			:      Dict=512     ,SCREAMING_SNAKE_CASE__			:      Optional[int]=12     ,SCREAMING_SNAKE_CASE__			:      List[Any]=2     ,SCREAMING_SNAKE_CASE__			:      Optional[Any]=0.02     ,SCREAMING_SNAKE_CASE__			:      Union[str, Any]=3     ,SCREAMING_SNAKE_CASE__			:      Dict=4     ,SCREAMING_SNAKE_CASE__			:      Dict="last"     ,SCREAMING_SNAKE_CASE__			:      Optional[Any]=None     ,SCREAMING_SNAKE_CASE__			:      List[Any]=None     ,):
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					parent
     SCREAMING_SNAKE_CASE:Dict						      =					batch_size
     SCREAMING_SNAKE_CASE:int						      =					seq_length
     SCREAMING_SNAKE_CASE:Any						      =					is_training
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					use_input_lengths
     SCREAMING_SNAKE_CASE:List[Any]						      =					use_token_type_ids
     SCREAMING_SNAKE_CASE:Dict						      =					use_labels
     SCREAMING_SNAKE_CASE:Dict						      =					gelu_activation
     SCREAMING_SNAKE_CASE:Dict						      =					sinusoidal_embeddings
     SCREAMING_SNAKE_CASE:List[str]						      =					causal
     SCREAMING_SNAKE_CASE:List[Any]						      =					asm
     SCREAMING_SNAKE_CASE:List[str]						      =					n_langs
     SCREAMING_SNAKE_CASE:Optional[Any]						      =					vocab_size
     SCREAMING_SNAKE_CASE:Optional[int]						      =					n_special
     SCREAMING_SNAKE_CASE:str						      =					hidden_size
     SCREAMING_SNAKE_CASE:Dict						      =					num_hidden_layers
     SCREAMING_SNAKE_CASE:Dict						      =					num_attention_heads
     SCREAMING_SNAKE_CASE:Optional[Any]						      =					hidden_dropout_prob
     SCREAMING_SNAKE_CASE:Tuple						      =					attention_probs_dropout_prob
     SCREAMING_SNAKE_CASE:Tuple						      =					max_position_embeddings
     SCREAMING_SNAKE_CASE:List[str]						      =					type_vocab_size
     SCREAMING_SNAKE_CASE:int						      =					type_sequence_label_size
     SCREAMING_SNAKE_CASE:Any						      =					initializer_range
     SCREAMING_SNAKE_CASE:str						      =					num_labels
     SCREAMING_SNAKE_CASE:Dict						      =					num_choices
     SCREAMING_SNAKE_CASE:Any						      =					summary_type
     SCREAMING_SNAKE_CASE:int						      =					use_proj
     SCREAMING_SNAKE_CASE:int						      =					scope
   def     __UpperCamelCase						( self			:      Optional[int]							):
     SCREAMING_SNAKE_CASE:int						      =					ids_tensor([self.batch_size, self.seq_length]     ,self.vocab_size							)
     SCREAMING_SNAKE_CASE:str						      =					random_attention_mask([self.batch_size, self.seq_length]							)

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

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

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

     SCREAMING_SNAKE_CASE:List[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     __UpperCamelCase						( self			:      Optional[int]							):
     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     __UpperCamelCase						( self			:      Any     ,SCREAMING_SNAKE_CASE__			:      Union[str, Any]     ,SCREAMING_SNAKE_CASE__			:      str     ,SCREAMING_SNAKE_CASE__			:      str     ,SCREAMING_SNAKE_CASE__			:      Dict     ,SCREAMING_SNAKE_CASE__			:      List[str]     ,SCREAMING_SNAKE_CASE__			:      Optional[int]     ,SCREAMING_SNAKE_CASE__			:      Tuple     ,SCREAMING_SNAKE_CASE__			:      Optional[Any]     ,SCREAMING_SNAKE_CASE__			:      List[Any]     ,):
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					FlaubertModel(config=SCREAMING_SNAKE_CASE__							)
     model.to(SCREAMING_SNAKE_CASE__							)
     model.eval()
     SCREAMING_SNAKE_CASE:Any						      =					model(SCREAMING_SNAKE_CASE__     ,lengths=SCREAMING_SNAKE_CASE__     ,langs=SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Any						      =					model(SCREAMING_SNAKE_CASE__     ,langs=SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Optional[int]						      =					model(SCREAMING_SNAKE_CASE__							)
     self.parent.assertEqual(result.last_hidden_state.shape     ,(self.batch_size, self.seq_length, self.hidden_size)							)
   def     __UpperCamelCase						( self			:      Dict     ,SCREAMING_SNAKE_CASE__			:      Optional[int]     ,SCREAMING_SNAKE_CASE__			:      Union[str, Any]     ,SCREAMING_SNAKE_CASE__			:      Union[str, Any]     ,SCREAMING_SNAKE_CASE__			:      Any     ,SCREAMING_SNAKE_CASE__			:      Tuple     ,SCREAMING_SNAKE_CASE__			:      Any     ,SCREAMING_SNAKE_CASE__			:      Union[str, Any]     ,SCREAMING_SNAKE_CASE__			:      Any     ,SCREAMING_SNAKE_CASE__			:      str     ,):
     SCREAMING_SNAKE_CASE:Optional[int]						      =					FlaubertWithLMHeadModel(SCREAMING_SNAKE_CASE__							)
     model.to(SCREAMING_SNAKE_CASE__							)
     model.eval()

     SCREAMING_SNAKE_CASE:Optional[int]						      =					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     __UpperCamelCase						( self			:      Any     ,SCREAMING_SNAKE_CASE__			:      List[str]     ,SCREAMING_SNAKE_CASE__			:      List[str]     ,SCREAMING_SNAKE_CASE__			:      str     ,SCREAMING_SNAKE_CASE__			:      List[Any]     ,SCREAMING_SNAKE_CASE__			:      Optional[Any]     ,SCREAMING_SNAKE_CASE__			:      Union[str, Any]     ,SCREAMING_SNAKE_CASE__			:      Union[str, Any]     ,SCREAMING_SNAKE_CASE__			:      Optional[Any]     ,SCREAMING_SNAKE_CASE__			:      Optional[int]     ,):
     SCREAMING_SNAKE_CASE:Optional[int]						      =					FlaubertForQuestionAnsweringSimple(SCREAMING_SNAKE_CASE__							)
     model.to(SCREAMING_SNAKE_CASE__							)
     model.eval()

     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					model(SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:Dict						      =					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     __UpperCamelCase						( self			:      List[str]     ,SCREAMING_SNAKE_CASE__			:      List[str]     ,SCREAMING_SNAKE_CASE__			:      Optional[Any]     ,SCREAMING_SNAKE_CASE__			:      Optional[int]     ,SCREAMING_SNAKE_CASE__			:      Tuple     ,SCREAMING_SNAKE_CASE__			:      Any     ,SCREAMING_SNAKE_CASE__			:      Union[str, Any]     ,SCREAMING_SNAKE_CASE__			:      Dict     ,SCREAMING_SNAKE_CASE__			:      Any     ,SCREAMING_SNAKE_CASE__			:      Tuple     ,):
     SCREAMING_SNAKE_CASE:str						      =					FlaubertForQuestionAnswering(SCREAMING_SNAKE_CASE__							)
     model.to(SCREAMING_SNAKE_CASE__							)
     model.eval()

     SCREAMING_SNAKE_CASE:List[Any]						      =					model(SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:Optional[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__     ,)

     SCREAMING_SNAKE_CASE:str						      =					model(
         SCREAMING_SNAKE_CASE__     ,start_positions=SCREAMING_SNAKE_CASE__     ,end_positions=SCREAMING_SNAKE_CASE__     ,cls_index=SCREAMING_SNAKE_CASE__     ,is_impossible=SCREAMING_SNAKE_CASE__     ,)

     ((SCREAMING_SNAKE_CASE)      ,    ):List[Any]						      =					result_with_labels.to_tuple()

     SCREAMING_SNAKE_CASE:Optional[int]						      =					model(SCREAMING_SNAKE_CASE__     ,start_positions=SCREAMING_SNAKE_CASE__     ,end_positions=SCREAMING_SNAKE_CASE__							)

     ((SCREAMING_SNAKE_CASE)      ,    ):List[Any]						      =					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     __UpperCamelCase						( self			:      Tuple     ,SCREAMING_SNAKE_CASE__			:      Union[str, Any]     ,SCREAMING_SNAKE_CASE__			:      Optional[Any]     ,SCREAMING_SNAKE_CASE__			:      Optional[Any]     ,SCREAMING_SNAKE_CASE__			:      str     ,SCREAMING_SNAKE_CASE__			:      str     ,SCREAMING_SNAKE_CASE__			:      Any     ,SCREAMING_SNAKE_CASE__			:      Union[str, Any]     ,SCREAMING_SNAKE_CASE__			:      Tuple     ,SCREAMING_SNAKE_CASE__			:      Dict     ,):
     SCREAMING_SNAKE_CASE:List[Any]						      =					FlaubertForSequenceClassification(SCREAMING_SNAKE_CASE__							)
     model.to(SCREAMING_SNAKE_CASE__							)
     model.eval()

     SCREAMING_SNAKE_CASE:Tuple						      =					model(SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Union[str, 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     __UpperCamelCase						( self			:      int     ,SCREAMING_SNAKE_CASE__			:      List[Any]     ,SCREAMING_SNAKE_CASE__			:      Union[str, Any]     ,SCREAMING_SNAKE_CASE__			:      List[Any]     ,SCREAMING_SNAKE_CASE__			:      List[Any]     ,SCREAMING_SNAKE_CASE__			:      List[Any]     ,SCREAMING_SNAKE_CASE__			:      Any     ,SCREAMING_SNAKE_CASE__			:      Union[str, Any]     ,SCREAMING_SNAKE_CASE__			:      List[Any]     ,SCREAMING_SNAKE_CASE__			:      Optional[Any]     ,):
     SCREAMING_SNAKE_CASE:List[str]						      =					self.num_labels
     SCREAMING_SNAKE_CASE:Optional[int]						      =					FlaubertForTokenClassification(SCREAMING_SNAKE_CASE__							)
     model.to(SCREAMING_SNAKE_CASE__							)
     model.eval()

     SCREAMING_SNAKE_CASE:Optional[Any]						      =					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     __UpperCamelCase						( self			:      Any     ,SCREAMING_SNAKE_CASE__			:      str     ,SCREAMING_SNAKE_CASE__			:      Dict     ,SCREAMING_SNAKE_CASE__			:      Optional[Any]     ,SCREAMING_SNAKE_CASE__			:      Optional[int]     ,SCREAMING_SNAKE_CASE__			:      Union[str, Any]     ,SCREAMING_SNAKE_CASE__			:      List[Any]     ,SCREAMING_SNAKE_CASE__			:      str     ,SCREAMING_SNAKE_CASE__			:      Tuple     ,SCREAMING_SNAKE_CASE__			:      List[str]     ,):
     SCREAMING_SNAKE_CASE:List[str]						      =					self.num_choices
     SCREAMING_SNAKE_CASE:Any						      =					FlaubertForMultipleChoice(config=SCREAMING_SNAKE_CASE__							)
     model.to(SCREAMING_SNAKE_CASE__							)
     model.eval()
     SCREAMING_SNAKE_CASE:str						      =					input_ids.unsqueeze(1							).expand(-1     ,self.num_choices     ,-1							).contiguous()
     SCREAMING_SNAKE_CASE:Dict						      =					token_type_ids.unsqueeze(1							).expand(-1     ,self.num_choices     ,-1							).contiguous()
     SCREAMING_SNAKE_CASE:List[Any]						      =					input_mask.unsqueeze(1							).expand(-1     ,self.num_choices     ,-1							).contiguous()
     SCREAMING_SNAKE_CASE:int						      =					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     __UpperCamelCase						( self			:      Any							):
     SCREAMING_SNAKE_CASE:str						      =					self.prepare_config_and_inputs()
     (
         (
         SCREAMING_SNAKE_CASE
     )      ,    (
         SCREAMING_SNAKE_CASE
     )      ,    (
         SCREAMING_SNAKE_CASE
     )      ,    (
         SCREAMING_SNAKE_CASE
     )      ,    (
         SCREAMING_SNAKE_CASE
     )      ,    (
         SCREAMING_SNAKE_CASE
     )      ,    (
         SCREAMING_SNAKE_CASE
     )      ,    (
         SCREAMING_SNAKE_CASE
     )      ,    (
         SCREAMING_SNAKE_CASE
     )      ,    
     ):Tuple						      =					config_and_inputs
     SCREAMING_SNAKE_CASE:Tuple						      =					{
         "input_ids": input_ids,
         "token_type_ids": token_type_ids,
         "lengths": input_lengths,
         "attention_mask": input_mask,
     }
     return config, inputs_dict





@require_torch
class 							_snake_case (       _a	,  _a	,  unittest.TestCase						):
   _A       :     int       =					(
       (
           FlaubertModel,
           FlaubertWithLMHeadModel,
           FlaubertForQuestionAnswering,
           FlaubertForQuestionAnsweringSimple,
           FlaubertForSequenceClassification,
           FlaubertForTokenClassification,
           FlaubertForMultipleChoice,
       )
       if is_torch_available()
       else ()
   )
   _A       :     Optional[Any]       =					(
       {
           '''feature-extraction''': FlaubertModel,
           '''fill-mask''': FlaubertWithLMHeadModel,
           '''question-answering''': FlaubertForQuestionAnsweringSimple,
           '''text-classification''': FlaubertForSequenceClassification,
           '''token-classification''': FlaubertForTokenClassification,
           '''zero-shot''': FlaubertForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )
   def     __UpperCamelCase						( self			:      str     ,SCREAMING_SNAKE_CASE__			:      Optional[int]     ,SCREAMING_SNAKE_CASE__			:      Tuple     ,SCREAMING_SNAKE_CASE__			:      str     ,SCREAMING_SNAKE_CASE__			:      Optional[int]     ,SCREAMING_SNAKE_CASE__			:      Optional[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			:      Union[str, Any]     ,SCREAMING_SNAKE_CASE__			:      Optional[int]     ,SCREAMING_SNAKE_CASE__			:      Union[str, Any]     ,SCREAMING_SNAKE_CASE__			:      int=False							):
     SCREAMING_SNAKE_CASE:Any						      =					super()._prepare_for_class(SCREAMING_SNAKE_CASE__     ,SCREAMING_SNAKE_CASE__     ,return_labels=SCREAMING_SNAKE_CASE__							)

     if return_labels:
       if model_class.__name__ == "FlaubertForQuestionAnswering":
         SCREAMING_SNAKE_CASE:Dict						      =					torch.zeros(
             self.model_tester.batch_size     ,dtype=torch.long     ,device=SCREAMING_SNAKE_CASE__							)
         SCREAMING_SNAKE_CASE:int						      =					torch.zeros(
             self.model_tester.batch_size     ,dtype=torch.long     ,device=SCREAMING_SNAKE_CASE__							)

     return inputs_dict
   def     __UpperCamelCase						( self			:      Optional[int]							):
     SCREAMING_SNAKE_CASE:int						      =					FlaubertModelTester(self							)
     SCREAMING_SNAKE_CASE:Dict						      =					ConfigTester(self     ,config_class=SCREAMING_SNAKE_CASE__     ,emb_dim=37							)
   def     __UpperCamelCase						( self			:      Any							):
     self.config_tester.run_common_tests()
   def     __UpperCamelCase						( self			:      Any							):
     SCREAMING_SNAKE_CASE:Any						      =					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_model(*SCREAMING_SNAKE_CASE__							)
   def     __UpperCamelCase						( self			:      Optional[int]							):
     SCREAMING_SNAKE_CASE:Dict						      =					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_lm_head(*SCREAMING_SNAKE_CASE__							)
   def     __UpperCamelCase						( self			:      str							):
     SCREAMING_SNAKE_CASE:str						      =					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_simple_qa(*SCREAMING_SNAKE_CASE__							)
   def     __UpperCamelCase						( self			:      Any							):
     SCREAMING_SNAKE_CASE:Dict						      =					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_qa(*SCREAMING_SNAKE_CASE__							)
   def     __UpperCamelCase						( self			:      Optional[Any]							):
     SCREAMING_SNAKE_CASE:Dict						      =					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_sequence_classif(*SCREAMING_SNAKE_CASE__							)
   def     __UpperCamelCase						( self			:      List[Any]							):
     SCREAMING_SNAKE_CASE:Optional[Any]						      =					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_token_classif(*SCREAMING_SNAKE_CASE__							)
   def     __UpperCamelCase						( self			:      int							):
     SCREAMING_SNAKE_CASE:List[Any]						      =					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_multiple_choice(*SCREAMING_SNAKE_CASE__							)
   @slow
   def     __UpperCamelCase						( self			:      List[str]							):
     for model_name in FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       SCREAMING_SNAKE_CASE:Tuple						      =					FlaubertModel.from_pretrained(SCREAMING_SNAKE_CASE__							)
       self.assertIsNotNone(SCREAMING_SNAKE_CASE__							)







   @slow
   @require_torch_gpu
   def     __UpperCamelCase						( self			:      str							):
     SCREAMING_SNAKE_CASE      ,    SCREAMING_SNAKE_CASE:int						      =					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

       SCREAMING_SNAKE_CASE:Union[str, Any]						      =					True
       SCREAMING_SNAKE_CASE:Optional[int]						      =					model_class(config=SCREAMING_SNAKE_CASE__							)

       SCREAMING_SNAKE_CASE:Any						      =					self._prepare_for_class(SCREAMING_SNAKE_CASE__     ,SCREAMING_SNAKE_CASE__							)
       SCREAMING_SNAKE_CASE:Any						      =					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"							)							)
         SCREAMING_SNAKE_CASE: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 							_snake_case (       unittest.TestCase						):
   @slow
   def     __UpperCamelCase						( self			:      List[str]							):
     SCREAMING_SNAKE_CASE:Optional[int]						      =					FlaubertModel.from_pretrained("flaubert/flaubert_base_cased"							)
     SCREAMING_SNAKE_CASE:Optional[int]						      =					torch.tensor([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]]							)
     with torch.no_grad():
       SCREAMING_SNAKE_CASE:Union[str, Any]						      =					model(SCREAMING_SNAKE_CASE__							)[0]
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					torch.Size((1, 11, 768)							)
     self.assertEqual(output.shape     ,SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:int						      =					torch.tensor(
         [[[-2.6_251, -1.4_298, -0.0_227], [-2.8_510, -1.6_387, 0.2_258], [-2.8_114, -1.1_832, -0.3_066]]]							)

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






