




'''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 		snake_case_   (UpperCamelCase	:  str				):





							'''simple docstring'''


							_a     					=    filter(lambda UpperCamelCase				: p.requires_grad			,						model.parameters()				)
							_a     					=    sum([np.prod(p.size()				) for p in model_parameters]				)
							return params


_snake_case				:				Any       							=		logging.getLogger(__name__)
def 		snake_case_   (UpperCamelCase	:  Dict			,						UpperCamelCase	:  List[str]				):





							'''simple docstring'''


							if metric == "rouge2":
														_a     					=    '''{val_avg_rouge2:.4f}-{step_count}'''
							elif metric == "bleu":
														_a     					=    '''{val_avg_bleu:.4f}-{step_count}'''
							elif metric == "em":
														_a     					=    '''{val_avg_em:.4f}-{step_count}'''
							elif metric == "loss":
														_a     					=    '''{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.'''				)

							_a     					=    ModelCheckpoint(
							    dirpath=UpperCamelCase			,						filename=UpperCamelCase			,						monitor=f'val_{metric}'			,						mode='''max'''			,						save_top_k=1			,						every_n_epochs=1			,						)

							return checkpoint_callback
def 		snake_case_   (UpperCamelCase	:  int			,						UpperCamelCase	:  str				):





							'''simple docstring'''


							return EarlyStopping(
							    monitor=f'val_{metric}'			,						mode='''min''' if '''loss''' in metric else '''max'''			,						patience=UpperCamelCase			,						verbose=UpperCamelCase			,						)

class       A   (		pl.Callback   ):







							def        __lowerCAmelCase      (  self	:  Optional[Any]							,      lowerCAmelCase_	:  List[str]							,      lowerCAmelCase_	:  int )      ->							Any:


														"""simple docstring"""


														_a     					=    {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        __lowerCAmelCase      (  self	:  Union[str, Any]							,      lowerCAmelCase_	:  pl.Trainer							,      lowerCAmelCase_	:  pl.LightningModule							,      lowerCAmelCase_	:  str							,      lowerCAmelCase_	:  str=True )      ->							None:


														"""simple docstring"""


														logger.info(F'***** {type_path} results at step {trainer.global_step:05d} *****' )
														_a     					=    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
														_a     					=    Path(pl_module.hparams.output_dir )
														if type_path == "test":
																					_a     					=    od / '''test_results.txt'''
																					_a     					=    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.
																					_a     					=    od / F'{type_path}_results/{trainer.global_step:05d}.txt'
																					_a     					=    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
																												_a     					=    metrics[key]
																												if isinstance(lowerCAmelCase_							,      torch.Tensor ):
																																			_a     					=    val.item()
																												_a     					=    F'{key}: {val:.6f}\n'
																												writer.write(lowerCAmelCase_ )

														if not save_generations:
																					return

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







							@rank_zero_only
							def        __lowerCAmelCase      (  self	:  Optional[Any]							,      lowerCAmelCase_	:  Dict							,      lowerCAmelCase_	:  Tuple )      ->							Optional[Any]:


														"""simple docstring"""


														try:
																					_a     					=    pl_module.model.model.num_parameters()
														except AttributeError:
																					_a     					=    pl_module.model.num_parameters()

														_a     					=    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        __lowerCAmelCase      (  self	:  Union[str, Any]							,      lowerCAmelCase_	:  pl.Trainer							,      lowerCAmelCase_	:  pl.LightningModule )      ->							List[Any]:


														"""simple docstring"""


														save_json(pl_module.metrics							,      pl_module.metrics_save_path )
														return self._write_logs(lowerCAmelCase_							,      lowerCAmelCase_							,      '''test''' )







							@rank_zero_only
							def        __lowerCAmelCase      (  self	:  Optional[Any]							,      lowerCAmelCase_	:  pl.Trainer							,      lowerCAmelCase_	:  int )      ->							Optional[int]:


														"""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")







'''simple docstring'''


from argparse import ArgumentParser, Namespace

from ..utils import logging
from . import BaseTransformersCLICommand
def 		snake_case_   (UpperCamelCase	:  Namespace				):





							'''simple docstring'''


							return ConvertCommand(
							    args.model_type			,						args.tf_checkpoint			,						args.pytorch_dump_output			,						args.config			,						args.finetuning_task_name				)


_snake_case				:				List[Any]       							=		'\ntransformers can only be used from the commandline to convert TensorFlow models in PyTorch, In that case, it requires\nTensorFlow to be installed. Please see https://www.tensorflow.org/install/ for installation instructions.\n'

class       A   (		_a   ):







							@staticmethod
							def        __lowerCAmelCase      (  lowerCAmelCase_	:  ArgumentParser )      ->							Any:


														"""simple docstring"""


														_a     					=    parser.add_parser(
														    '''convert'''							,      help='''CLI tool to run convert model from original author checkpoints to Transformers PyTorch checkpoints.'''							,      )
														train_parser.add_argument('''--model_type'''							,      type=lowerCAmelCase_							,      required=lowerCAmelCase_							,      help='''Model\'s type.''' )
														train_parser.add_argument(
														    '''--tf_checkpoint'''							,      type=lowerCAmelCase_							,      required=lowerCAmelCase_							,      help='''TensorFlow checkpoint path or folder.''' )
														train_parser.add_argument(
														    '''--pytorch_dump_output'''							,      type=lowerCAmelCase_							,      required=lowerCAmelCase_							,      help='''Path to the PyTorch saved model output.''' )
														train_parser.add_argument('''--config'''							,      type=lowerCAmelCase_							,      default=''''''							,      help='''Configuration file path or folder.''' )
														train_parser.add_argument(
														    '''--finetuning_task_name'''							,      type=lowerCAmelCase_							,      default=lowerCAmelCase_							,      help='''Optional fine-tuning task name if the TF model was a finetuned model.'''							,      )
														train_parser.set_defaults(func=lowerCAmelCase_ )







							def __init__(  self	:  int							,      lowerCAmelCase_	:  str							,      lowerCAmelCase_	:  str							,      lowerCAmelCase_	:  str							,      lowerCAmelCase_	:  str							,      lowerCAmelCase_	:  str							,      *lowerCAmelCase_	:  Optional[Any]							,      )      ->							List[str]:


														"""simple docstring"""


														_a     					=    logging.get_logger('''transformers-cli/converting''' )

														self._logger.info(F'Loading model {model_type}' )
														_a     					=    model_type
														_a     					=    tf_checkpoint
														_a     					=    pytorch_dump_output
														_a     					=    config
														_a     					=    finetuning_task_name







							def        __lowerCAmelCase      (  self	:  Dict )      ->							Dict:


														"""simple docstring"""


														if self._model_type == "albert":
																					try:
																												from ..models.albert.convert_albert_original_tf_checkpoint_to_pytorch import (
																												    convert_tf_checkpoint_to_pytorch,
																												)
																					except ImportError:
																												raise ImportError(lowerCAmelCase_ )

																					convert_tf_checkpoint_to_pytorch(self._tf_checkpoint							,      self._config							,      self._pytorch_dump_output )
														elif self._model_type == "bert":
																					try:
																												from ..models.bert.convert_bert_original_tf_checkpoint_to_pytorch import (
																												    convert_tf_checkpoint_to_pytorch,
																												)
																					except ImportError:
																												raise ImportError(lowerCAmelCase_ )

																					convert_tf_checkpoint_to_pytorch(self._tf_checkpoint							,      self._config							,      self._pytorch_dump_output )
														elif self._model_type == "funnel":
																					try:
																												from ..models.funnel.convert_funnel_original_tf_checkpoint_to_pytorch import (
																												    convert_tf_checkpoint_to_pytorch,
																												)
																					except ImportError:
																												raise ImportError(lowerCAmelCase_ )

																					convert_tf_checkpoint_to_pytorch(self._tf_checkpoint							,      self._config							,      self._pytorch_dump_output )
														elif self._model_type == "t5":
																					try:
																												from ..models.ta.convert_ta_original_tf_checkpoint_to_pytorch import convert_tf_checkpoint_to_pytorch
																					except ImportError:
																												raise ImportError(lowerCAmelCase_ )

																					convert_tf_checkpoint_to_pytorch(self._tf_checkpoint							,      self._config							,      self._pytorch_dump_output )
														elif self._model_type == "gpt":
																					from ..models.openai.convert_openai_original_tf_checkpoint_to_pytorch import (
																					    convert_openai_checkpoint_to_pytorch,
																					)

																					convert_openai_checkpoint_to_pytorch(self._tf_checkpoint							,      self._config							,      self._pytorch_dump_output )
														elif self._model_type == "transfo_xl":
																					try:
																												from ..models.transfo_xl.convert_transfo_xl_original_tf_checkpoint_to_pytorch import (
																												    convert_transfo_xl_checkpoint_to_pytorch,
																												)
																					except ImportError:
																												raise ImportError(lowerCAmelCase_ )

																					if "ckpt" in self._tf_checkpoint.lower():
																												_a     					=    self._tf_checkpoint
																												_a     					=    ''''''
																					else:
																												_a     					=    self._tf_checkpoint
																												_a     					=    ''''''
																					convert_transfo_xl_checkpoint_to_pytorch(
																					    lowerCAmelCase_							,      self._config							,      self._pytorch_dump_output							,      lowerCAmelCase_ )
														elif self._model_type == "gpt2":
																					try:
																												from ..models.gpta.convert_gpta_original_tf_checkpoint_to_pytorch import (
																												    convert_gpta_checkpoint_to_pytorch,
																												)
																					except ImportError:
																												raise ImportError(lowerCAmelCase_ )

																					convert_gpta_checkpoint_to_pytorch(self._tf_checkpoint							,      self._config							,      self._pytorch_dump_output )
														elif self._model_type == "xlnet":
																					try:
																												from ..models.xlnet.convert_xlnet_original_tf_checkpoint_to_pytorch import (
																												    convert_xlnet_checkpoint_to_pytorch,
																												)
																					except ImportError:
																												raise ImportError(lowerCAmelCase_ )

																					convert_xlnet_checkpoint_to_pytorch(
																					    self._tf_checkpoint							,      self._config							,      self._pytorch_dump_output							,      self._finetuning_task_name )
														elif self._model_type == "xlm":
																					from ..models.xlm.convert_xlm_original_pytorch_checkpoint_to_pytorch import (
																					    convert_xlm_checkpoint_to_pytorch,
																					)

																					convert_xlm_checkpoint_to_pytorch(self._tf_checkpoint							,      self._pytorch_dump_output )
														elif self._model_type == "lxmert":
																					from ..models.lxmert.convert_lxmert_original_tf_checkpoint_to_pytorch import (
																					    convert_lxmert_checkpoint_to_pytorch,
																					)

																					convert_lxmert_checkpoint_to_pytorch(self._tf_checkpoint							,      self._pytorch_dump_output )
														elif self._model_type == "rembert":
																					from ..models.rembert.convert_rembert_tf_checkpoint_to_pytorch import (
																					    convert_rembert_tf_checkpoint_to_pytorch,
																					)

																					convert_rembert_tf_checkpoint_to_pytorch(self._tf_checkpoint							,      self._config							,      self._pytorch_dump_output )
														else:
																					raise ValueError(
																					    '''--model_type should be selected in the list [bert, gpt, gpt2, t5, transfo_xl, xlnet, xlm, lxmert]''' )

