


import logging
from pathlib import Path

import numpy as np
import pytorch_lightning as pl
import torch
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from pytorch_lightning.utilities import rank_zero_only
from utils_rag import save_json




def        snake_case_					(  snake_case					)  ->    int:
       lowercase__:       Optional[Any]        =       filter(lambda snake_case					: p.requires_grad      ,							model.parameters()					)
       lowercase__:       str        =       sum([np.prod(p.size()					) for p in model_parameters]					)
       return params


__lowerCAmelCase			=							logging.getLogger(__name__)




def        snake_case_					(  snake_case      ,							snake_case					)  ->    List[str]:
       if metric == "rouge2":
              lowercase__:       Tuple        =       '{val_avg_rouge2:.4f}-{step_count}'
       elif metric == "bleu":
              lowercase__:       List[str]        =       '{val_avg_bleu:.4f}-{step_count}'
       elif metric == "em":
              lowercase__:       Union[str, Any]        =       '{val_avg_em:.4f}-{step_count}'
       elif metric == "loss":
              lowercase__:       List[str]        =       '{val_avg_loss:.4f}-{step_count}'
       else:
              raise NotImplementedError(
                  f'seq2seq callbacks only support rouge2 and bleu, got {metric}, You can make your own by adding to this'
                  ' function.'					)

       lowercase__:       List[str]        =       ModelCheckpoint(
           dirpath=snake_case      ,							filename=snake_case      ,							monitor=f'val_{metric}'      ,							mode='max'      ,							save_top_k=1      ,							every_n_epochs=1      ,							)

       return checkpoint_callback




def        snake_case_					(  snake_case      ,							snake_case					)  ->    List[str]:
       return EarlyStopping(
           monitor=f'val_{metric}'      ,							mode='min' if 'loss' in metric else 'max'      ,							patience=snake_case      ,							verbose=snake_case      ,							)





class   __a      (   pl.Callback				):
     def 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__    ,			lowerCAmelCase__			) ->    Any:




            '''simple docstring'''
            lowercase__:       Union[str, Any]        =       {F'lr_group_{i}': param['lr'] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups			)}
            pl_module.logger.log_metrics(lowerCAmelCase__			)
     @rank_zero_only
     def 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__    ,			lowerCAmelCase__    ,			lowerCAmelCase__    ,			lowerCAmelCase__=True			) ->    None:




            '''simple docstring'''
            logger.info(F'***** {type_path} results at step {trainer.global_step:05d} *****'			)
            lowercase__:       Any        =       trainer.callback_metrics
            trainer.logger.log_metrics({k: v for k, v in metrics.items() if k not in ['log', 'progress_bar', 'preds']}			)
            # Log results
            lowercase__:       Any        =       Path(pl_module.hparams.output_dir			)
            if type_path == "test":
                   lowercase__:       Optional[Any]        =       od / 'test_results.txt'
                   lowercase__:       Optional[Any]        =       od / 'test_generations.txt'
            else:
                   # this never gets hit. I prefer not to save intermediate generations, and results are in metrics.json
                   # If people want this it will be easy enough to add back.
                   lowercase__:       str        =       od / F'{type_path}_results/{trainer.global_step:05d}.txt'
                   lowercase__:       Optional[int]        =       od / F'{type_path}_generations/{trainer.global_step:05d}.txt'
                   results_file.parent.mkdir(exist_ok=lowerCAmelCase__			)
                   generations_file.parent.mkdir(exist_ok=lowerCAmelCase__			)
            with open(lowerCAmelCase__    ,			'a+'			) as writer:
                   for key in sorted(lowerCAmelCase__			):
                          if key in ["log", "progress_bar", "preds"]:
                                 continue
                          lowercase__:       Dict        =       metrics[key]
                          if isinstance(lowerCAmelCase__    ,			torch.Tensor			):
                                 lowercase__:       Any        =       val.item()
                          lowercase__:       List[str]        =       F'{key}: {val:.6f}\n'
                          writer.write(lowerCAmelCase__			)

            if not save_generations:
                   return

            if "preds" in metrics:
                   lowercase__:       Optional[Any]        =       '\n'.join(metrics['preds']			)
                   generations_file.open('w+'			).write(lowerCAmelCase__			)
     @rank_zero_only
     def 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__    ,			lowerCAmelCase__			) ->    Tuple:




            '''simple docstring'''
            try:
                   lowercase__:       List[str]        =       pl_module.model.model.num_parameters()
            except AttributeError:
                   lowercase__:       List[Any]        =       pl_module.model.num_parameters()

            lowercase__:       str        =       count_trainable_parameters(lowerCAmelCase__			)
            # mp stands for million parameters
            trainer.logger.log_metrics({'n_params': npars, 'mp': npars / 1E6, 'grad_mp': n_trainable_pars / 1E6}			)
     @rank_zero_only
     def 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__    ,			lowerCAmelCase__			) ->    Dict:




            '''simple docstring'''
            save_json(pl_module.metrics    ,			pl_module.metrics_save_path			)
            return self._write_logs(lowerCAmelCase__    ,			lowerCAmelCase__    ,			'test'			)
     @rank_zero_only
     def 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__    ,			lowerCAmelCase__			) ->    Tuple:




            '''simple docstring'''
            save_json(pl_module.metrics    ,			pl_module.metrics_save_path			)
            # Uncommenting this will save val generations
            # return self._write_logs(trainer, pl_module, "valid")






import argparse
import os
import re

import tensorflow as tf
import torch

from transformers import BertConfig, BertModel
from transformers.utils import logging


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




def        snake_case_					(  snake_case      ,							snake_case      ,							snake_case					)  ->    Any:
       lowercase__:       Dict        =       os.path.abspath(snake_case					)
       logger.info(f'Converting TensorFlow checkpoint from {tf_path}'					)
       # Load weights from TF model
       lowercase__:       Optional[Any]        =       tf.train.list_variables(snake_case					)
       lowercase__:       List[Any]        =       []
       lowercase__:       Tuple        =       []
       lowercase__:       int        =       []
       for full_name, shape in init_vars:
              # logger.info(f"Loading TF weight {name} with shape {shape}")
              lowercase__:       Union[str, Any]        =       full_name.split('/'					)
              if full_name == "_CHECKPOINTABLE_OBJECT_GRAPH" or name[0] in ["global_step", "save_counter"]:
                     logger.info(f'Skipping non-model layer {full_name}'					)
                     continue
              if "optimizer" in full_name:
                     logger.info(f'Skipping optimization layer {full_name}'					)
                     continue
              if name[0] == "model":
                     # ignore initial 'model'
                     lowercase__:       str        =       name[1:]
              # figure out how many levels deep the name is
              lowercase__:       Optional[Any]        =       0
              for _name in name:
                     if _name.startswith('layer_with_weights'					):
                            depth += 1
                     else:
                            break
              layer_depth.append(snake_case					)
              # read data
              lowercase__:       Optional[Any]        =       tf.train.load_variable(snake_case      ,							snake_case					)
              names.append('/'.join(snake_case					)					)
              arrays.append(snake_case					)
       logger.info(f'Read a total of {len(snake_case					):,} layers'					)

       # Sanity check
       if len(set(snake_case					)					) != 1:
              raise ValueError(f'Found layer names with different depths (layer depth {list(set(snake_case					)					)})'					)
       lowercase__:       Any        =       list(set(snake_case					)					)[0]
       if layer_depth != 1:
              raise ValueError(
                  'The model contains more than just the embedding/encoder layers. This script does not handle MLM/NSP'
                  ' heads.'					)

       # convert layers
       logger.info('Converting weights...'					)
       for full_name, array in zip(snake_case      ,							snake_case					):
              lowercase__:       Optional[int]        =       full_name.split('/'					)
              lowercase__:       List[Any]        =       model
              lowercase__:       Any        =       []
              for i, m_name in enumerate(snake_case					):
                     if m_name == ".ATTRIBUTES":
                            # variable names end with .ATTRIBUTES/VARIABLE_VALUE
                            break
                     if m_name.startswith('layer_with_weights'					):
                            lowercase__:       Any        =       int(m_name.split('-'					)[-1]					)
                            if layer_num <= 2:
                                   # embedding layers
                                   # layer_num 0: word_embeddings
                                   # layer_num 1: position_embeddings
                                   # layer_num 2: token_type_embeddings
                                   continue
                            elif layer_num == 3:
                                   # embedding LayerNorm
                                   trace.extend(['embeddings', 'LayerNorm']					)
                                   lowercase__:       str        =       getattr(snake_case      ,							'embeddings'					)
                                   lowercase__:       int        =       getattr(snake_case      ,							'LayerNorm'					)
                            elif layer_num > 3 and layer_num < config.num_hidden_layers + 4:
                                   # encoder layers
                                   trace.extend(['encoder', 'layer', str(layer_num - 4					)]					)
                                   lowercase__:       int        =       getattr(snake_case      ,							'encoder'					)
                                   lowercase__:       List[str]        =       getattr(snake_case      ,							'layer'					)
                                   lowercase__:       Union[str, Any]        =       pointer[layer_num - 4]
                            elif layer_num == config.num_hidden_layers + 4:
                                   # pooler layer
                                   trace.extend(['pooler', 'dense']					)
                                   lowercase__:       Tuple        =       getattr(snake_case      ,							'pooler'					)
                                   lowercase__:       Tuple        =       getattr(snake_case      ,							'dense'					)
                     elif m_name == "embeddings":
                            trace.append('embeddings'					)
                            lowercase__:       Union[str, Any]        =       getattr(snake_case      ,							'embeddings'					)
                            if layer_num == 0:
                                   trace.append('word_embeddings'					)
                                   lowercase__:       Union[str, Any]        =       getattr(snake_case      ,							'word_embeddings'					)
                            elif layer_num == 1:
                                   trace.append('position_embeddings'					)
                                   lowercase__:       Dict        =       getattr(snake_case      ,							'position_embeddings'					)
                            elif layer_num == 2:
                                   trace.append('token_type_embeddings'					)
                                   lowercase__:       Optional[Any]        =       getattr(snake_case      ,							'token_type_embeddings'					)
                            else:
                                   raise ValueError(f'Unknown embedding layer with name {full_name}'					)
                            trace.append('weight'					)
                            lowercase__:       List[str]        =       getattr(snake_case      ,							'weight'					)
                     elif m_name == "_attention_layer":
                            # self-attention layer
                            trace.extend(['attention', 'self']					)
                            lowercase__:       int        =       getattr(snake_case      ,							'attention'					)
                            lowercase__:       Union[str, Any]        =       getattr(snake_case      ,							'self'					)
                     elif m_name == "_attention_layer_norm":
                            # output attention norm
                            trace.extend(['attention', 'output', 'LayerNorm']					)
                            lowercase__:       str        =       getattr(snake_case      ,							'attention'					)
                            lowercase__:       Optional[Any]        =       getattr(snake_case      ,							'output'					)
                            lowercase__:       List[Any]        =       getattr(snake_case      ,							'LayerNorm'					)
                     elif m_name == "_attention_output_dense":
                            # output attention dense
                            trace.extend(['attention', 'output', 'dense']					)
                            lowercase__:       Optional[Any]        =       getattr(snake_case      ,							'attention'					)
                            lowercase__:       Optional[int]        =       getattr(snake_case      ,							'output'					)
                            lowercase__:       Optional[Any]        =       getattr(snake_case      ,							'dense'					)
                     elif m_name == "_output_dense":
                            # output dense
                            trace.extend(['output', 'dense']					)
                            lowercase__:       Union[str, Any]        =       getattr(snake_case      ,							'output'					)
                            lowercase__:       List[Any]        =       getattr(snake_case      ,							'dense'					)
                     elif m_name == "_output_layer_norm":
                            # output dense
                            trace.extend(['output', 'LayerNorm']					)
                            lowercase__:       Any        =       getattr(snake_case      ,							'output'					)
                            lowercase__:       str        =       getattr(snake_case      ,							'LayerNorm'					)
                     elif m_name == "_key_dense":
                            # attention key
                            trace.append('key'					)
                            lowercase__:       Tuple        =       getattr(snake_case      ,							'key'					)
                     elif m_name == "_query_dense":
                            # attention query
                            trace.append('query'					)
                            lowercase__:       List[str]        =       getattr(snake_case      ,							'query'					)
                     elif m_name == "_value_dense":
                            # attention value
                            trace.append('value'					)
                            lowercase__:       Optional[int]        =       getattr(snake_case      ,							'value'					)
                     elif m_name == "_intermediate_dense":
                            # attention intermediate dense
                            trace.extend(['intermediate', 'dense']					)
                            lowercase__:       Any        =       getattr(snake_case      ,							'intermediate'					)
                            lowercase__:       str        =       getattr(snake_case      ,							'dense'					)
                     elif m_name == "_output_layer_norm":
                            # output layer norm
                            trace.append('output'					)
                            lowercase__:       Union[str, Any]        =       getattr(snake_case      ,							'output'					)
                     # weights & biases
                     elif m_name in ["bias", "beta"]:
                            trace.append('bias'					)
                            lowercase__:       str        =       getattr(snake_case      ,							'bias'					)
                     elif m_name in ["kernel", "gamma"]:
                            trace.append('weight'					)
                            lowercase__:       Tuple        =       getattr(snake_case      ,							'weight'					)
                     else:
                            logger.warning(f'Ignored {m_name}'					)
        # for certain layers reshape is necessary
              lowercase__:       Any        =       '.'.join(snake_case					)
              if re.match(R'(\S+)\.attention\.self\.(key|value|query)\.(bias|weight)'      ,							snake_case					) or re.match(
                  R'(\S+)\.attention\.output\.dense\.weight'      ,							snake_case					):
                     lowercase__:       Union[str, Any]        =       array.reshape(pointer.data.shape					)
              if "kernel" in full_name:
                     lowercase__:       str        =       array.transpose()
              if pointer.shape == array.shape:
                     lowercase__:       Union[str, Any]        =       torch.from_numpy(snake_case					)
              else:
                     raise ValueError(
                         f'Shape mismatch in layer {full_name}: Model expects shape {pointer.shape} but layer contains shape:'
                         f' {array.shape}'					)
              logger.info(f'Successfully set variable {full_name} to PyTorch layer {trace}'					)
       return model




def        snake_case_					(  snake_case      ,							snake_case      ,							snake_case					)  ->    Any:
       # Instantiate model
       logger.info(f'Loading model based on config from {config_path}...'					)
       lowercase__:       int        =       BertConfig.from_json_file(snake_case					)
       lowercase__:       Tuple        =       BertModel(snake_case					)

       # Load weights from checkpoint
       logger.info(f'Loading weights from checkpoint {tf_checkpoint_path}...'					)
       load_tfa_weights_in_bert(snake_case      ,							snake_case      ,							snake_case					)

       # Save pytorch-model
       logger.info(f'Saving PyTorch model to {pytorch_dump_path}...'					)
       torch.save(model.state_dict()      ,							snake_case					)


if __name__ == "__main__":
 __lowerCAmelCase			=							argparse.ArgumentParser()
 parser.add_argument(
     '''--tf_checkpoint_path''', type=str, required=True, help='''Path to the TensorFlow 2.x checkpoint path.'''
 )
 parser.add_argument(
     '''--bert_config_file''',
     type=str,
     required=True,
     help='''The config json file corresponding to the BERT model. This specifies the model architecture.''',
 )
 parser.add_argument(
     '''--pytorch_dump_path''',
     type=str,
     required=True,
     help='''Path to the output PyTorch model (must include filename).''',
 )
 __lowerCAmelCase			=							parser.parse_args()
 convert_tfa_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)


