


'''simple docstring'''



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 			UpperCAmelCase__					(	UpperCAmelCase_    :				Tuple							)     ->		Optional[int]:
     __lowerCamelCase   :	Optional[int]								=		filter(lambda UpperCAmelCase_							: p.requires_grad      ,    model.parameters()							)
     __lowerCamelCase   :	int								=		sum([np.prod(p.size()							) for p in model_parameters]							)
     return params


A__		:					Optional[Any]												=  logging.getLogger(__name__)



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				Optional[Any]      ,    UpperCAmelCase_    :				List[Any]							)     ->		Tuple:
     if metric == "rouge2":
          __lowerCamelCase   :	Dict								=		'{val_avg_rouge2:.4f}-{step_count}'
     elif metric == "bleu":
          __lowerCamelCase   :	List[str]								=		'{val_avg_bleu:.4f}-{step_count}'
     elif metric == "em":
          __lowerCamelCase   :	str								=		'{val_avg_em:.4f}-{step_count}'
     elif metric == "loss":
          __lowerCamelCase   :	Dict								=		'{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.'							)

     __lowerCamelCase   :	List[str]								=		ModelCheckpoint(
         dirpath=UpperCAmelCase_      ,    filename=UpperCAmelCase_      ,    monitor=F'val_{metric}'      ,    mode='max'      ,    save_top_k=1      ,    every_n_epochs=1      ,    )

     return checkpoint_callback



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				Optional[int]      ,    UpperCAmelCase_    :				List[str]							)     ->		str:
     return EarlyStopping(
         monitor=F'val_{metric}'      ,    mode='min' if 'loss' in metric else 'max'      ,    patience=UpperCAmelCase_      ,    verbose=UpperCAmelCase_      ,    )





class    UpperCAmelCase_  (pl.Callback  ):





    """simple docstring"""






    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Optional[Any]:
         __lowerCamelCase   :	Any								=		{f'lr_group_{i}': param['lr'] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups	)}
         pl_module.logger.log_metrics(SCREAMING_SNAKE_CASE_	)







    @rank_zero_only
    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_=True	)				->     None:
         logger.info(f'***** {type_path} results at step {trainer.global_step:05d} *****'	)
         __lowerCamelCase   :	List[str]								=		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
         __lowerCamelCase   :	Optional[Any]								=		Path(pl_module.hparams.output_dir	)
         if type_path == "test":
              __lowerCamelCase   :	int								=		od / 'test_results.txt'
              __lowerCamelCase   :	List[str]								=		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.
              __lowerCamelCase   :	Any								=		od / f'{type_path}_results/{trainer.global_step:05d}.txt'
              __lowerCamelCase   :	Any								=		od / f'{type_path}_generations/{trainer.global_step:05d}.txt'
              results_file.parent.mkdir(exist_ok=SCREAMING_SNAKE_CASE_	)
              generations_file.parent.mkdir(exist_ok=SCREAMING_SNAKE_CASE_	)
         with open(SCREAMING_SNAKE_CASE_			,     'a+'	) as writer:
              for key in sorted(SCREAMING_SNAKE_CASE_	):
                   if key in ["log", "progress_bar", "preds"]:
                        continue
                   __lowerCamelCase   :	Any								=		metrics[key]
                   if isinstance(SCREAMING_SNAKE_CASE_			,     torch.Tensor	):
                        __lowerCamelCase   :	List[Any]								=		val.item()
                   __lowerCamelCase   :	Optional[int]								=		f'{key}: {val:.6f}\n'
                   writer.write(SCREAMING_SNAKE_CASE_	)

         if not save_generations:
              return

         if "preds" in metrics:
              __lowerCamelCase   :	str								=		'\n'.join(metrics['preds']	)
              generations_file.open('w+'	).write(SCREAMING_SNAKE_CASE_	)







    @rank_zero_only
    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     List[Any]:
         try:
              __lowerCamelCase   :	Optional[int]								=		pl_module.model.model.num_parameters()
         except AttributeError:
              __lowerCamelCase   :	Tuple								=		pl_module.model.num_parameters()

         __lowerCamelCase   :	Union[str, Any]								=		count_trainable_parameters(SCREAMING_SNAKE_CASE_	)
         # 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       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     List[Any]:
         save_json(pl_module.metrics			,     pl_module.metrics_save_path	)
         return self._write_logs(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     'test'	)







    @rank_zero_only
    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     int:
         save_json(pl_module.metrics			,     pl_module.metrics_save_path	)
         # Uncommenting this will save val generations
         # return self._write_logs(trainer, pl_module, "valid")










'''simple docstring'''



import argparse
from typing import Dict

import tensorflow as tf
import torch
from tqdm import tqdm

from transformers import BigBirdPegasusConfig, BigBirdPegasusForConditionalGeneration


A__		:					str												=  [
    # tf -> hf
    ("""/""", """."""),
    ("""layer_""", """layers."""),
    ("""kernel""", """weight"""),
    ("""beta""", """bias"""),
    ("""gamma""", """weight"""),
    ("""pegasus""", """model"""),
]
A__		:					Dict												=  [
    (""".output.dense""", """.fc2"""),
    ("""intermediate.LayerNorm""", """final_layer_norm"""),
    ("""intermediate.dense""", """fc1"""),
]

A__		:					int												=  (
    INIT_COMMON
    + [
        ("""attention.self.LayerNorm""", """self_attn_layer_norm"""),
        ("""attention.output.dense""", """self_attn.out_proj"""),
        ("""attention.self""", """self_attn"""),
        ("""attention.encdec.LayerNorm""", """encoder_attn_layer_norm"""),
        ("""attention.encdec_output.dense""", """encoder_attn.out_proj"""),
        ("""attention.encdec""", """encoder_attn"""),
        ("""key""", """k_proj"""),
        ("""value""", """v_proj"""),
        ("""query""", """q_proj"""),
        ("""decoder.LayerNorm""", """decoder.layernorm_embedding"""),
    ]
    + END_COMMON
)

A__		:					Tuple												=  (
    INIT_COMMON
    + [
        ("""embeddings.word_embeddings""", """shared.weight"""),
        ("""embeddings.position_embeddings""", """embed_positions.weight"""),
        ("""attention.self.LayerNorm""", """self_attn_layer_norm"""),
        ("""attention.output.dense""", """self_attn.output"""),
        ("""attention.self""", """self_attn.self"""),
        ("""encoder.LayerNorm""", """encoder.layernorm_embedding"""),
    ]
    + END_COMMON
)

A__		:					Tuple												=  [
    """encdec/key/bias""",
    """encdec/query/bias""",
    """encdec/value/bias""",
    """self/key/bias""",
    """self/query/bias""",
    """self/value/bias""",
    """encdec_output/dense/bias""",
    """attention/output/dense/bias""",
]



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				Union[str, Any]      ,    UpperCAmelCase_    :				Any							)     ->		List[str]:
     for tf_name, hf_name in patterns:
          __lowerCamelCase   :	Optional[int]								=		k.replace(UpperCAmelCase_      ,    UpperCAmelCase_							)
     return k



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				dict      ,    UpperCAmelCase_    :				dict							)     ->		BigBirdPegasusForConditionalGeneration:
     __lowerCamelCase   :	int								=		BigBirdPegasusConfig(**UpperCAmelCase_							)
     __lowerCamelCase   :	Any								=		BigBirdPegasusForConditionalGeneration(UpperCAmelCase_							)
     __lowerCamelCase   :	Union[str, Any]								=		torch_model.state_dict()
     __lowerCamelCase   :	Tuple								=		{}

     # separating decoder weights
     __lowerCamelCase   :	Dict								=		{k: tf_weights[k] for k in tf_weights if k.startswith('pegasus/decoder'							)}
     __lowerCamelCase   :	str								=		{k: tf_weights[k] for k in tf_weights if not k.startswith('pegasus/decoder'							)}

     for k, v in tqdm(decoder_weights.items()      ,    'tf -> hf conversion'							):
          __lowerCamelCase   :	Tuple								=		[k.endswith(UpperCAmelCase_							) for ending in KEYS_TO_IGNORE]
          if any(UpperCAmelCase_							):
               continue
          __lowerCamelCase   :	Tuple								=		DECODER_PATTERNS
          __lowerCamelCase   :	Optional[int]								=		rename_state_dict_key(UpperCAmelCase_      ,    UpperCAmelCase_							)
          if new_k not in state_dict:
               raise ValueError(F'could not find new key {new_k} in state dict. (converted from {k})'							)
          if any(True if i in k else False for i in ['dense', 'query', 'key', 'value']							):
               __lowerCamelCase   :	Union[str, Any]								=		v.T
          __lowerCamelCase   :	str								=		torch.from_numpy(UpperCAmelCase_							)
          assert v.shape == state_dict[new_k].shape, F'{new_k}, {k}, {v.shape}, {state_dict[new_k].shape}'

     for k, v in tqdm(remaining_weights.items()      ,    'tf -> hf conversion'							):
          __lowerCamelCase   :	Optional[Any]								=		[k.endswith(UpperCAmelCase_							) for ending in KEYS_TO_IGNORE]
          if any(UpperCAmelCase_							):
               continue
          __lowerCamelCase   :	Dict								=		REMAINING_PATTERNS
          __lowerCamelCase   :	List[str]								=		rename_state_dict_key(UpperCAmelCase_      ,    UpperCAmelCase_							)
          if new_k not in state_dict and k != "pegasus/embeddings/position_embeddings":
               raise ValueError(F'could not find new key {new_k} in state dict. (converted from {k})'							)
          if any(True if i in k else False for i in ['dense', 'query', 'key', 'value']							):
               __lowerCamelCase   :	List[str]								=		v.T
          __lowerCamelCase   :	List[str]								=		torch.from_numpy(UpperCAmelCase_							)
          if k != "pegasus/embeddings/position_embeddings":
               assert v.shape == state_dict[new_k].shape, F'{new_k}, {k}, {v.shape}, {state_dict[new_k].shape}'

     __lowerCamelCase   :	Any								=		mapping['model.embed_positions.weight']
     __lowerCamelCase   :	Union[str, Any]								=		mapping.pop('model.embed_positions.weight'							)
     __lowerCamelCase				,       __lowerCamelCase   :	List[str]								=		torch_model.load_state_dict(UpperCAmelCase_      ,    strict=UpperCAmelCase_							)
     __lowerCamelCase   :	Any								=		[
         k
         for k in missing
         if k
         not in [
             'final_logits_bias',
             'model.encoder.embed_tokens.weight',
             'model.decoder.embed_tokens.weight',
             'lm_head.weight',
         ]
     ]
     assert unexpected_missing == [], F'no matches found for the following torch keys {unexpected_missing}'
     assert extra == [], F'no matches found for the following tf keys {extra}'
     return torch_model



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				Union[str, Any]							)     ->		Dict:
     __lowerCamelCase   :	int								=		tf.train.list_variables(UpperCAmelCase_							)
     __lowerCamelCase   :	List[str]								=		{}
     __lowerCamelCase   :	List[Any]								=		['global_step']
     for name, shape in tqdm(UpperCAmelCase_      ,    desc='converting tf checkpoint to dict'							):
          __lowerCamelCase   :	Any								=		any(pat in name for pat in ignore_name							)
          if skip_key:
               continue
          __lowerCamelCase   :	Dict								=		tf.train.load_variable(UpperCAmelCase_      ,    UpperCAmelCase_							)
          __lowerCamelCase   :	List[Any]								=		array
     return tf_weights



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				str      ,    UpperCAmelCase_    :				str      ,    UpperCAmelCase_    :				dict							)     ->		Union[str, Any]:
     __lowerCamelCase   :	Optional[int]								=		get_tf_weights_as_numpy(UpperCAmelCase_							)
     __lowerCamelCase   :	List[str]								=		convert_bigbird_pegasus(UpperCAmelCase_      ,    UpperCAmelCase_							)
     torch_model.save_pretrained(UpperCAmelCase_							)


if __name__ == "__main__":
      A__		:					Any												=  argparse.ArgumentParser()
      parser.add_argument("""--tf_ckpt_path""", type=str, help="""passed to tf.train.list_variables""")
      parser.add_argument("""--save_dir""", default=None, type=str, help="""Path to the output PyTorch model.""")
      A__		:					List[str]												=  parser.parse_args()
      A__		:					Optional[int]												=  {}
      convert_bigbird_pegasus_ckpt_to_pytorch(args.tf_ckpt_path, args.save_dir, config_update=config_update)






