




from __future__ import annotations

import unittest

from transformers import AlbertConfig, is_tf_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_tf, 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 tensorflow as tf

   from transformers import TF_MODEL_FOR_PRETRAINING_MAPPING
   from transformers.models.albert.modeling_tf_albert import (
       TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
       TFAlbertForMaskedLM,
       TFAlbertForMultipleChoice,
       TFAlbertForPreTraining,
       TFAlbertForQuestionAnswering,
       TFAlbertForSequenceClassification,
       TFAlbertForTokenClassification,
       TFAlbertModel,
   )





class  lowerCamelCase_							:


 '''simple docstring'''







 def __init__(							self						:	Optional[Any]       ,       _lowerCAmelCase						:	List[str]       ,       _lowerCAmelCase						:	str=13       ,       _lowerCAmelCase						:	int=7       ,       _lowerCAmelCase						:	Dict=True       ,       _lowerCAmelCase						:	Tuple=True       ,       _lowerCAmelCase						:	Tuple=True       ,       _lowerCAmelCase						:	Any=True       ,       _lowerCAmelCase						:	Tuple=99       ,       _lowerCAmelCase						:	List[str]=16       ,       _lowerCAmelCase						:	Any=32       ,       _lowerCAmelCase						:	Optional[int]=2       ,       _lowerCAmelCase						:	Dict=4       ,       _lowerCAmelCase						:	str=37       ,       _lowerCAmelCase						:	Optional[Any]="gelu"       ,       _lowerCAmelCase						:	Union[str, Any]=0.1       ,       _lowerCAmelCase						:	Any=0.1       ,       _lowerCAmelCase						:	Optional[Any]=512       ,       _lowerCAmelCase						:	int=16       ,       _lowerCAmelCase						:	Optional[Any]=2       ,       _lowerCAmelCase						:	Union[str, Any]=0.02       ,       _lowerCAmelCase						:	Dict=3       ,       _lowerCAmelCase						:	Optional[int]=4       ,       _lowerCAmelCase						:	Dict=None       ,       ):
       SCREAMING_SNAKE_CASE_					=    parent
       SCREAMING_SNAKE_CASE_					=    13
       SCREAMING_SNAKE_CASE_					=    7
       SCREAMING_SNAKE_CASE_					=    True
       SCREAMING_SNAKE_CASE_					=    True
       SCREAMING_SNAKE_CASE_					=    True
       SCREAMING_SNAKE_CASE_					=    True
       SCREAMING_SNAKE_CASE_					=    99
       SCREAMING_SNAKE_CASE_					=    16
       SCREAMING_SNAKE_CASE_					=    32
       SCREAMING_SNAKE_CASE_					=    2
       SCREAMING_SNAKE_CASE_					=    4
       SCREAMING_SNAKE_CASE_					=    37
       SCREAMING_SNAKE_CASE_					=    'gelu'
       SCREAMING_SNAKE_CASE_					=    0.1
       SCREAMING_SNAKE_CASE_					=    0.1
       SCREAMING_SNAKE_CASE_					=    512
       SCREAMING_SNAKE_CASE_					=    16
       SCREAMING_SNAKE_CASE_					=    2
       SCREAMING_SNAKE_CASE_					=    0.02
       SCREAMING_SNAKE_CASE_					=    3
       SCREAMING_SNAKE_CASE_					=    4
       SCREAMING_SNAKE_CASE_					=    None



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

       SCREAMING_SNAKE_CASE_					=    None
       if self.use_input_mask:
             SCREAMING_SNAKE_CASE_					=    random_attention_mask([self.batch_size, self.seq_length]		)

       SCREAMING_SNAKE_CASE_					=    None
       if self.use_token_type_ids:
             SCREAMING_SNAKE_CASE_					=    ids_tensor([self.batch_size, self.seq_length]       ,       self.type_vocab_size		)

       SCREAMING_SNAKE_CASE_					=    None
       SCREAMING_SNAKE_CASE_					=    None
       SCREAMING_SNAKE_CASE_					=    None
       if self.use_labels:
             SCREAMING_SNAKE_CASE_					=    ids_tensor([self.batch_size]       ,       self.type_sequence_label_size		)
             SCREAMING_SNAKE_CASE_					=    ids_tensor([self.batch_size, self.seq_length]       ,       self.num_labels		)
             SCREAMING_SNAKE_CASE_					=    ids_tensor([self.batch_size]       ,       self.num_choices		)

       SCREAMING_SNAKE_CASE_					=    AlbertConfig(
           vocab_size=self.vocab_size       ,       hidden_size=self.hidden_size       ,       embedding_size=self.embedding_size       ,       num_hidden_layers=self.num_hidden_layers       ,       num_attention_heads=self.num_attention_heads       ,       intermediate_size=self.intermediate_size       ,       hidden_act=self.hidden_act       ,       hidden_dropout_prob=self.hidden_dropout_prob       ,       attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,       max_position_embeddings=self.max_position_embeddings       ,       type_vocab_size=self.type_vocab_size       ,       initializer_range=self.initializer_range       ,       )

       return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels



 def    lowerCAmelCase_						(							self						:	List[Any]       ,       _lowerCAmelCase						:	Union[str, Any]       ,       _lowerCAmelCase						:	int       ,       _lowerCAmelCase						:	Optional[int]       ,       _lowerCAmelCase						:	Tuple       ,       _lowerCAmelCase						:	Union[str, Any]       ,       _lowerCAmelCase						:	Optional[Any]       ,       _lowerCAmelCase						:	Union[str, Any]		):
       SCREAMING_SNAKE_CASE_					=    TFAlbertModel(config=_lowerCAmelCase		)
       # inputs = {'input_ids': input_ids,
       #           'attention_mask': input_mask,
       #           'token_type_ids': token_type_ids}
       # sequence_output, pooled_output = model(**inputs)
       SCREAMING_SNAKE_CASE_					=    {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
       SCREAMING_SNAKE_CASE_					=    model(_lowerCAmelCase		)

       SCREAMING_SNAKE_CASE_					=    [input_ids, input_mask]
       SCREAMING_SNAKE_CASE_					=    model(_lowerCAmelCase		)

       SCREAMING_SNAKE_CASE_					=    model(_lowerCAmelCase		)

       self.parent.assertEqual(result.last_hidden_state.shape       ,       (self.batch_size, self.seq_length, self.hidden_size)		)
       self.parent.assertEqual(result.pooler_output.shape       ,       (self.batch_size, self.hidden_size)		)



 def    lowerCAmelCase_						(							self						:	Optional[int]       ,       _lowerCAmelCase						:	str       ,       _lowerCAmelCase						:	Tuple       ,       _lowerCAmelCase						:	Optional[int]       ,       _lowerCAmelCase						:	Optional[int]       ,       _lowerCAmelCase						:	Any       ,       _lowerCAmelCase						:	List[Any]       ,       _lowerCAmelCase						:	List[Any]		):
       SCREAMING_SNAKE_CASE_					=    self.num_labels
       SCREAMING_SNAKE_CASE_					=    TFAlbertForPreTraining(config=_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_					=    {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
       SCREAMING_SNAKE_CASE_					=    model(_lowerCAmelCase		)
       self.parent.assertEqual(result.prediction_logits.shape       ,       (self.batch_size, self.seq_length, self.vocab_size)		)
       self.parent.assertEqual(result.sop_logits.shape       ,       (self.batch_size, self.num_labels)		)



 def    lowerCAmelCase_						(							self						:	Optional[Any]       ,       _lowerCAmelCase						:	Tuple       ,       _lowerCAmelCase						:	int       ,       _lowerCAmelCase						:	Any       ,       _lowerCAmelCase						:	List[str]       ,       _lowerCAmelCase						:	Optional[Any]       ,       _lowerCAmelCase						:	Optional[Any]       ,       _lowerCAmelCase						:	Dict		):
       SCREAMING_SNAKE_CASE_					=    TFAlbertForMaskedLM(config=_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_					=    {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
       SCREAMING_SNAKE_CASE_					=    model(_lowerCAmelCase		)
       self.parent.assertEqual(result.logits.shape       ,       (self.batch_size, self.seq_length, self.vocab_size)		)



 def    lowerCAmelCase_						(							self						:	List[str]       ,       _lowerCAmelCase						:	Tuple       ,       _lowerCAmelCase						:	Any       ,       _lowerCAmelCase						:	Tuple       ,       _lowerCAmelCase						:	Tuple       ,       _lowerCAmelCase						:	Optional[int]       ,       _lowerCAmelCase						:	Dict       ,       _lowerCAmelCase						:	Dict		):
       SCREAMING_SNAKE_CASE_					=    self.num_labels
       SCREAMING_SNAKE_CASE_					=    TFAlbertForSequenceClassification(config=_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_					=    {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
       SCREAMING_SNAKE_CASE_					=    model(_lowerCAmelCase		)
       self.parent.assertEqual(result.logits.shape       ,       (self.batch_size, self.num_labels)		)



 def    lowerCAmelCase_						(							self						:	List[Any]       ,       _lowerCAmelCase						:	List[str]       ,       _lowerCAmelCase						:	Optional[Any]       ,       _lowerCAmelCase						:	Dict       ,       _lowerCAmelCase						:	int       ,       _lowerCAmelCase						:	List[str]       ,       _lowerCAmelCase						:	List[str]       ,       _lowerCAmelCase						:	str		):
       SCREAMING_SNAKE_CASE_					=    TFAlbertForQuestionAnswering(config=_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_					=    {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
       SCREAMING_SNAKE_CASE_					=    model(_lowerCAmelCase		)
       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    lowerCAmelCase_						(							self						:	List[str]       ,       _lowerCAmelCase						:	List[str]       ,       _lowerCAmelCase						:	Any       ,       _lowerCAmelCase						:	List[str]       ,       _lowerCAmelCase						:	Union[str, Any]       ,       _lowerCAmelCase						:	str       ,       _lowerCAmelCase						:	Tuple       ,       _lowerCAmelCase						:	Dict		):
       SCREAMING_SNAKE_CASE_					=    self.num_choices
       SCREAMING_SNAKE_CASE_					=    TFAlbertForMultipleChoice(config=_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_					=    tf.tile(tf.expand_dims(_lowerCAmelCase       ,       1		)       ,       (1, self.num_choices, 1)		)
       SCREAMING_SNAKE_CASE_					=    tf.tile(tf.expand_dims(_lowerCAmelCase       ,       1		)       ,       (1, self.num_choices, 1)		)
       SCREAMING_SNAKE_CASE_					=    tf.tile(tf.expand_dims(_lowerCAmelCase       ,       1		)       ,       (1, self.num_choices, 1)		)
       SCREAMING_SNAKE_CASE_					=    {
           'input_ids': multiple_choice_inputs_ids,
           'attention_mask': multiple_choice_input_mask,
           'token_type_ids': multiple_choice_token_type_ids,
       }
       SCREAMING_SNAKE_CASE_					=    model(_lowerCAmelCase		)
       self.parent.assertListEqual(list(result['logits'].shape		)       ,       [self.batch_size, self.num_choices]		)



 def    lowerCAmelCase_						(							self						:	Union[str, Any]       ,       _lowerCAmelCase						:	Optional[Any]       ,       _lowerCAmelCase						:	str       ,       _lowerCAmelCase						:	Optional[Any]       ,       _lowerCAmelCase						:	int       ,       _lowerCAmelCase						:	Dict       ,       _lowerCAmelCase						:	Tuple       ,       _lowerCAmelCase						:	List[Any]		):
       SCREAMING_SNAKE_CASE_					=    self.num_labels
       SCREAMING_SNAKE_CASE_					=    TFAlbertForTokenClassification(config=_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_					=    {
           'input_ids': input_ids,
           'attention_mask': input_mask,
           'token_type_ids': token_type_ids,
       }
       SCREAMING_SNAKE_CASE_					=    model(_lowerCAmelCase		)
       self.parent.assertListEqual(list(result['logits'].shape		)       ,       [self.batch_size, self.seq_length, self.num_labels]		)




 def    lowerCAmelCase_						(							self						:	int		):
       SCREAMING_SNAKE_CASE_					=    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_
       )		,				
       )					=    config_and_inputs
       SCREAMING_SNAKE_CASE_					=    {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
       return config, inputs_dict





@require_tf
class  lowerCamelCase_							(   _SCREAMING_SNAKE_CASE ,     _SCREAMING_SNAKE_CASE ,     unittest.TestCase      ):


 '''simple docstring'''







 lowercase_       					=					(
     (
         TFAlbertModel,
         TFAlbertForPreTraining,
         TFAlbertForMaskedLM,
         TFAlbertForSequenceClassification,
         TFAlbertForQuestionAnswering,
         TFAlbertForTokenClassification,
         TFAlbertForMultipleChoice,
     )
     if is_tf_available()
     else ()
 )
 lowercase_       					=					(
     {
         "feature-extraction": TFAlbertModel,
         "fill-mask": TFAlbertForMaskedLM,
         "question-answering": TFAlbertForQuestionAnswering,
         "text-classification": TFAlbertForSequenceClassification,
         "token-classification": TFAlbertForTokenClassification,
         "zero-shot": TFAlbertForSequenceClassification,
     }
     if is_tf_available()
     else {}
 )
 lowercase_       					=					False
 lowercase_       					=					False



 def    lowerCAmelCase_						(							self						:	Optional[int]       ,       _lowerCAmelCase						:	Dict       ,       _lowerCAmelCase						:	str       ,       _lowerCAmelCase						:	Dict=False		):
       SCREAMING_SNAKE_CASE_					=    super()._prepare_for_class(_lowerCAmelCase       ,       _lowerCAmelCase       ,       return_labels=_lowerCAmelCase		)

       if return_labels:
             if model_class in get_values(_lowerCAmelCase		):
                   SCREAMING_SNAKE_CASE_					=    tf.zeros(self.model_tester.batch_size       ,       dtype=tf.intaa		)

       return inputs_dict



 def    lowerCAmelCase_						(							self						:	Optional[int]		):
       SCREAMING_SNAKE_CASE_					=    TFAlbertModelTester(self		)
       SCREAMING_SNAKE_CASE_					=    ConfigTester(self       ,       config_class=_lowerCAmelCase       ,       hidden_size=37		)



 def    lowerCAmelCase_						(							self						:	Optional[int]		):
       self.config_tester.run_common_tests()



 def    lowerCAmelCase_						(							self						:	List[str]		):
       SCREAMING_SNAKE_CASE_					=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_model(*_lowerCAmelCase		)



 def    lowerCAmelCase_						(							self						:	Tuple		):
       SCREAMING_SNAKE_CASE_					=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_pretraining(*_lowerCAmelCase		)



 def    lowerCAmelCase_						(							self						:	Dict		):
       SCREAMING_SNAKE_CASE_					=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_masked_lm(*_lowerCAmelCase		)



 def    lowerCAmelCase_						(							self						:	Tuple		):
       SCREAMING_SNAKE_CASE_					=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_multiple_choice(*_lowerCAmelCase		)



 def    lowerCAmelCase_						(							self						:	List[Any]		):
       SCREAMING_SNAKE_CASE_					=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_sequence_classification(*_lowerCAmelCase		)



 def    lowerCAmelCase_						(							self						:	Optional[Any]		):
       SCREAMING_SNAKE_CASE_					=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_question_answering(*_lowerCAmelCase		)




 @slow
 def    lowerCAmelCase_						(							self						:	Union[str, Any]		):
       for model_name in TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             SCREAMING_SNAKE_CASE_					=    TFAlbertModel.from_pretrained(_lowerCAmelCase		)
             self.assertIsNotNone(_lowerCAmelCase		)





@require_tf
class  lowerCamelCase_							(   unittest.TestCase      ):


 '''simple docstring'''







 @slow
 def    lowerCAmelCase_						(							self						:	List[str]		):
       SCREAMING_SNAKE_CASE_					=    TFAlbertForPreTraining.from_pretrained('albert-base-v2'		)
       SCREAMING_SNAKE_CASE_					=    tf.constant([[0, 1, 2, 3, 4, 5]]		)
       SCREAMING_SNAKE_CASE_					=    model(_lowerCAmelCase		)[0]

       SCREAMING_SNAKE_CASE_					=    [1, 6, 30_000]
       self.assertEqual(output.shape       ,       _lowerCAmelCase		)

       SCREAMING_SNAKE_CASE_					=    tf.constant(
           [
               [
                   [4.59_5668, 0.7446_2754, -1.81_8147],
                   [4.595_4347, 0.745_4184, -1.818_8258],
                   [4.595_4905, 0.744_8235, -1.818_2316],
               ]
           ]		)
       tf.debugging.assert_near(output[:, :3, :3]       ,       _lowerCAmelCase       ,       atol=1E-4		)





import os
from pathlib import Path
from typing import Dict, List

import fire
import torch

from transformers import AutoModelForSeqaSeqLM, AutoTokenizer
from transformers.utils.logging import get_logger


lowerCamelCase__				:     Optional[Any]			       =							get_logger(__name__)

def       UpperCAmelCase_							(   __UpperCAmelCase		:			str    ,    __UpperCAmelCase		:			str							)      ->				List[str]:
      if text.startswith(__UpperCAmelCase							):
            return text[len(__UpperCAmelCase							) :]
      return text  # or whatever

def       UpperCAmelCase_							(   __UpperCAmelCase		:			Union[str, Any]							)      ->				List[Any]:
      return {remove_prefix(__UpperCAmelCase    ,    'model.'							): v for k, v in sd.items()}

def       UpperCAmelCase_							(   __UpperCAmelCase		:			List[Dict[str, torch.Tensor]]							)      ->				Dict:
      SCREAMING_SNAKE_CASE_					=    {}
      for k in state_dicts[0].keys():
            SCREAMING_SNAKE_CASE_					=    [sd[k] for sd in state_dicts]
            SCREAMING_SNAKE_CASE_					=    sum(__UpperCAmelCase							) / len(__UpperCAmelCase							)
            assert isinstance(__UpperCAmelCase    ,    torch.Tensor							)
            SCREAMING_SNAKE_CASE_					=    new_t
      return new_sd

def       UpperCAmelCase_							(   __UpperCAmelCase		:			str    ,    __UpperCAmelCase		:			str    ,    __UpperCAmelCase		:			str							)      ->				None:
      SCREAMING_SNAKE_CASE_					=    AutoModelForSeqaSeqLM.from_pretrained(__UpperCAmelCase							)
      if os.path.isfile(__UpperCAmelCase							):
            SCREAMING_SNAKE_CASE_					=    [pl_ckpt_path]
      else:
            assert os.path.isdir(__UpperCAmelCase							)
            SCREAMING_SNAKE_CASE_					=    list(Path(__UpperCAmelCase							).glob('*.ckpt'							)							)
            assert ckpt_files, f"could not find any ckpt files inside the {pl_ckpt_path} directory"

      if len(__UpperCAmelCase							) > 1:
            logger.info(f"averaging the weights of {ckpt_files}"							)

      SCREAMING_SNAKE_CASE_					=    [sanitize(torch.load(__UpperCAmelCase    ,    map_location='cpu'							)['state_dict']							) for x in ckpt_files]
      SCREAMING_SNAKE_CASE_					=    average_state_dicts(__UpperCAmelCase							)

      SCREAMING_SNAKE_CASE_		,				SCREAMING_SNAKE_CASE_					=    hf_model.load_state_dict(__UpperCAmelCase    ,    strict=__UpperCAmelCase							)
      assert not missing, f"missing keys: {missing}"
      hf_model.save_pretrained(__UpperCAmelCase							)
      try:
            SCREAMING_SNAKE_CASE_					=    AutoTokenizer.from_pretrained(__UpperCAmelCase							)
            tok.save_pretrained(__UpperCAmelCase							)
      except Exception:
            pass
            # dont copy tokenizer if cant


if __name__ == "__main__":
   fire.Fire(convert_pl_to_hf)