


'''simple docstring'''




import logging
import os

import quant_trainer
import torch
from torch.utils.data import DataLoader

from transformers import Trainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput


UpperCAmelCase_      :		Dict          =		logging.getLogger(__name__)

if is_torch_tpu_available(check_device=False):
 import torch_xla.core.xla_model as xm
 import torch_xla.debug.metrics as met







class 					a			(				snake_case__  ):
 '''simple docstring'''




 def __init__(    self		,      *lowerCamelCase_		,      lowerCamelCase_=None		,      lowerCamelCase_=None		,      lowerCamelCase_=None		,      **lowerCamelCase_      )      ->    Any:
    super().__init__(*a_		,      **a_      )
    _a					:  Optional[Any]						=  eval_examples
    _a					:  Any						=  post_process_function
    _a					:  Dict						=  quant_trainer_args
    _a					:  List[Any]						=  1_2_8  # default number of calibration samples




 def   __UpperCamelCase   (    self		,      lowerCamelCase_=None      )      ->    List[str]:
    if calib_dataset is None and self.calib_dataset is None:
       raise ValueError('Trainer: calibration requires an calib_dataset.'      )
    _a					:  Dict						=  calib_dataset if calib_dataset is not None else self.calib_dataset

    _a					:  Optional[int]						=  self._remove_unused_columns(a_		,      description='Calibration'      )

    return DataLoader(
        a_		,      batch_size=self.args.eval_batch_size		,      collate_fn=self.data_collator		,      drop_last=self.args.dataloader_drop_last		,      num_workers=self.args.dataloader_num_workers		,      pin_memory=self.args.dataloader_pin_memory		,      shuffle=a_		,      )




 def   __UpperCamelCase   (    self		,      lowerCamelCase_=None      )      ->    str:
    _a					:  Optional[Any]						=  self.train_dataset if calib_dataset is None else calib_dataset
    _a					:  Optional[Any]						=  self.get_calib_dataloader(a_      )

    _a					:  Union[str, Any]						=  self.model
    quant_trainer.configure_model(a_		,      self.quant_trainer_args		,      calib=a_      )
    model.eval()
    quant_trainer.enable_calibration(a_      )

    logger.info('***** Running calibration *****'      )
    logger.info(F'''  Num examples = {self.calib_num}'''      )
    logger.info(F'''  Batch size = {calib_dataloader.batch_size}'''      )

    for step, inputs in enumerate(a_      ):
       # Prediction step
       _a					:  Union[str, Any]						=  self.prediction_step(a_		,      a_		,      prediction_loss_only=a_      )
       if (step + 1) * calib_dataloader.batch_size >= self.calib_num:
          break

    quant_trainer.finish_calibration(a_		,      self.quant_trainer_args      )
    _a					:  List[Any]						=  model




 def   __UpperCamelCase   (    self		,      lowerCamelCase_=None		,      lowerCamelCase_=None		,      lowerCamelCase_=None		,      lowerCamelCase_ = "eval"      )      ->    int:
    _a					:  List[str]						=  self.eval_dataset if eval_dataset is None else eval_dataset
    _a					:  Union[str, Any]						=  self.get_eval_dataloader(a_      )
    _a					:  Tuple						=  self.eval_examples if eval_examples is None else eval_examples

    # Temporarily disable metric computation, we will do it in the loop here.
    _a					:  List[str]						=  self.compute_metrics
    _a					:  Dict						=  None
    _a					:  int						=  self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
    try:
       _a					:  Tuple						=  eval_loop(
           a_		,      description='Evaluation'		,      prediction_loss_only=True if compute_metrics is None else None		,      ignore_keys=a_		,      )
    finally:
       _a					:  List[Any]						=  compute_metrics

    if self.post_process_function is not None and self.compute_metrics is not None:
       _a					:  Union[str, Any]						=  self.post_process_function(a_		,      a_		,      output.predictions      )
       _a					:  Union[str, Any]						=  self.compute_metrics(a_      )

       # Prefix all keys with metric_key_prefix + '_'
       for key in list(metrics.keys()      ):
          if not key.startswith(F'''{metric_key_prefix}_'''      ):
             _a					:  Dict						=  metrics.pop(a_      )

       self.log(a_      )
    else:
       _a					:  Tuple						=  {}

    if self.args.tpu_metrics_debug or self.args.debug:
       # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
       xm.master_print(met.metrics_report()      )

    _a					:  int						=  self.callback_handler.on_evaluate(self.args		,      self.state		,      self.control		,      a_      )
    return metrics




 def   __UpperCamelCase   (    self		,      lowerCamelCase_		,      lowerCamelCase_		,      lowerCamelCase_=None		,      lowerCamelCase_ = "test"      )      ->    Optional[int]:
    _a					:  Dict						=  self.get_test_dataloader(a_      )

    # Temporarily disable metric computation, we will do it in the loop here.
    _a					:  List[str]						=  self.compute_metrics
    _a					:  List[str]						=  None
    _a					:  int						=  self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
    try:
       _a					:  Optional[Any]						=  eval_loop(
           a_		,      description='Prediction'		,      prediction_loss_only=True if compute_metrics is None else None		,      ignore_keys=a_		,      )
    finally:
       _a					:  int						=  compute_metrics

    if self.post_process_function is None or self.compute_metrics is None:
       return output

    _a					:  List[str]						=  self.post_process_function(a_		,      a_		,      output.predictions		,      'predict'      )
    _a					:  Dict						=  self.compute_metrics(a_      )

    # Prefix all keys with metric_key_prefix + '_'
    for key in list(metrics.keys()      ):
       if not key.startswith(F'''{metric_key_prefix}_'''      ):
          _a					:  Union[str, Any]						=  metrics.pop(a_      )

    return PredictionOutput(predictions=predictions.predictions		,      label_ids=predictions.label_ids		,      metrics=a_      )




 def   __UpperCamelCase   (    self		,      lowerCamelCase_="./"      )      ->    Optional[Any]:
    _a					:  str						=  self.eval_dataset
    _a					:  Dict						=  self.get_eval_dataloader(a_      )

    _a					:  Any						=  next(iter(a_      )      )

    # saving device - to make it consistent
    _a					:  List[Any]						=  torch.device('cuda' if torch.cuda.is_available() else 'cpu'      )

    # convert to tuple
    _a					:  Tuple						=  tuple(v.to(a_      ) for k, v in batch.items()      )

    logger.info('Converting model to be onnx compatible'      )
    from pytorch_quantization.nn import TensorQuantizer

    _a					:  Optional[Any]						=  True

    _a					:  int						=  self.model.to(a_      )

    model.eval()
    model.float()

    _a					:  Union[str, Any]						=  model.module if hasattr(a_		,      'module'      ) else model
    quant_trainer.configure_model(a_		,      self.quant_trainer_args      )

    _a					:  List[str]						=  os.path.join(a_		,      'model.onnx'      )
    logger.info(F'''exporting model to {output_model_file}'''      )

    _a					:  Any						=  {0: "batch_size", 1: "seq_len"}

    torch.onnx.export(
        a_		,      a_		,      a_		,      export_params=a_		,      opset_version=1_3		,      do_constant_folding=a_		,      input_names=['input_ids', 'attention_mask', 'token_type_ids']		,      output_names=['output_start_logits', 'output_end_logits']		,      dynamic_axes={
            'input_ids': axes,
            'attention_mask': axes,
            'token_type_ids': axes,
            'output_start_logits': axes,
            'output_end_logits': axes,
        }		,      verbose=a_		,      )
    logger.info('onnx export finished'      )






'''simple docstring'''
import argparse
import pathlib

import fairseq
import torch
from fairseq.models.roberta import RobertaModel as FairseqRobertaModel
from fairseq.modules import TransformerSentenceEncoderLayer
from packaging import version

from transformers import XLMRobertaConfig, XLMRobertaXLForMaskedLM, XLMRobertaXLForSequenceClassification
from transformers.models.bert.modeling_bert import (
    BertIntermediate,
    BertLayer,
    BertOutput,
    BertSelfAttention,
    BertSelfOutput,
)
from transformers.models.roberta.modeling_roberta import RobertaAttention
from transformers.utils import logging


if version.parse(fairseq.__version__) < version.parse("""1.0.0a"""):
					raise Exception("""requires fairseq >= 1.0.0a""")

logging.set_verbosity_info()
lowerCAmelCase   : Optional[int]             =				logging.get_logger(__name__)

lowerCAmelCase   : Tuple             =				"""Hello world! cécé herlolip"""




def 	_A  (       A							,A							,A							)  ->				str:
	lowercase						:		Optional[Any]   					=		FairseqRobertaModel.from_pretrained(A							)
	roberta.eval()  # disable dropout
	lowercase						:		int   					=		roberta.model.encoder.sentence_encoder
	lowercase						:		Tuple   					=		XLMRobertaConfig(
	    vocab_size=roberta_sent_encoder.embed_tokens.num_embeddings							,hidden_size=roberta.cfg.model.encoder_embed_dim							,num_hidden_layers=roberta.cfg.model.encoder_layers							,num_attention_heads=roberta.cfg.model.encoder_attention_heads							,intermediate_size=roberta.cfg.model.encoder_ffn_embed_dim							,max_position_embeddings=5_1_4							,type_vocab_size=1							,layer_norm_eps=1e-5							,)
	if classification_head:
		lowercase						:		Tuple   					=		roberta.model.classification_heads["mnli"].out_proj.weight.shape[0]

	print("Our RoBERTa config:"							,A							)

	lowercase						:		Union[str, Any]   					=		XLMRobertaXLForSequenceClassification(A							) if classification_head else XLMRobertaXLForMaskedLM(A							)
	model.eval()

	# Now let's copy all the weights.
	# Embeddings
	lowercase						:		Optional[Any]   					=		roberta_sent_encoder.embed_tokens.weight
	lowercase						:		Dict   					=		roberta_sent_encoder.embed_positions.weight
	lowercase						:		str   					=		torch.zeros_like(
	    model.roberta.embeddings.token_type_embeddings.weight							)  # just zero them out b/c RoBERTa doesn't use them.

	lowercase						:		Dict   					=		roberta_sent_encoder.layer_norm.weight
	lowercase						:		str   					=		roberta_sent_encoder.layer_norm.bias

	for i in range(config.num_hidden_layers							):
		# Encoder: start of layer
		lowercase						:		BertLayer   					=		model.roberta.encoder.layer[i]
		lowercase						:		TransformerSentenceEncoderLayer   					=		roberta_sent_encoder.layers[i]

		lowercase						:		RobertaAttention   					=		layer.attention
		lowercase						:		str   					=		roberta_layer.self_attn_layer_norm.weight
		lowercase						:		List[str]   					=		roberta_layer.self_attn_layer_norm.bias

		# self attention
		lowercase						:		BertSelfAttention   					=		layer.attention.self
		assert (
		    roberta_layer.self_attn.k_proj.weight.data.shape
		    == roberta_layer.self_attn.q_proj.weight.data.shape
		    == roberta_layer.self_attn.v_proj.weight.data.shape
		    == torch.Size((config.hidden_size, config.hidden_size)							)
		)

		lowercase						:		str   					=		roberta_layer.self_attn.q_proj.weight
		lowercase						:		List[str]   					=		roberta_layer.self_attn.q_proj.bias
		lowercase						:		Union[str, Any]   					=		roberta_layer.self_attn.k_proj.weight
		lowercase						:		List[Any]   					=		roberta_layer.self_attn.k_proj.bias
		lowercase						:		Dict   					=		roberta_layer.self_attn.v_proj.weight
		lowercase						:		Tuple   					=		roberta_layer.self_attn.v_proj.bias

		# self-attention output
		lowercase						:		BertSelfOutput   					=		layer.attention.output
		assert self_output.dense.weight.shape == roberta_layer.self_attn.out_proj.weight.shape
		lowercase						:		Tuple   					=		roberta_layer.self_attn.out_proj.weight
		lowercase						:		str   					=		roberta_layer.self_attn.out_proj.bias

		# this one is final layer norm
		lowercase						:		List[Any]   					=		roberta_layer.final_layer_norm.weight
		lowercase						:		Optional[Any]   					=		roberta_layer.final_layer_norm.bias

		# intermediate
		lowercase						:		BertIntermediate   					=		layer.intermediate
		assert intermediate.dense.weight.shape == roberta_layer.fca.weight.shape
		lowercase						:		Optional[int]   					=		roberta_layer.fca.weight
		lowercase						:		Dict   					=		roberta_layer.fca.bias

		# output
		lowercase						:		BertOutput   					=		layer.output
		assert bert_output.dense.weight.shape == roberta_layer.fca.weight.shape
		lowercase						:		Optional[int]   					=		roberta_layer.fca.weight
		lowercase						:		str   					=		roberta_layer.fca.bias
		# end of layer

	if classification_head:
		lowercase						:		int   					=		roberta.model.classification_heads["mnli"].dense.weight
		lowercase						:		List[Any]   					=		roberta.model.classification_heads["mnli"].dense.bias
		lowercase						:		Union[str, Any]   					=		roberta.model.classification_heads["mnli"].out_proj.weight
		lowercase						:		Any   					=		roberta.model.classification_heads["mnli"].out_proj.bias
	else:
		# LM Head
		lowercase						:		str   					=		roberta.model.encoder.lm_head.dense.weight
		lowercase						:		List[Any]   					=		roberta.model.encoder.lm_head.dense.bias
		lowercase						:		str   					=		roberta.model.encoder.lm_head.layer_norm.weight
		lowercase						:		Optional[Any]   					=		roberta.model.encoder.lm_head.layer_norm.bias
		lowercase						:		int   					=		roberta.model.encoder.lm_head.weight
		lowercase						:		Optional[int]   					=		roberta.model.encoder.lm_head.bias

	# Let's check that we get the same results.
	lowercase						:		torch.Tensor   					=		roberta.encode(A							).unsqueeze(0							)  # batch of size 1

	lowercase						:		Tuple   					=		model(A							)[0]
	if classification_head:
		lowercase						:		Union[str, Any]   					=		roberta.model.classification_heads["mnli"](roberta.extract_features(A							)							)
	else:
		lowercase						:		Optional[Any]   					=		roberta.model(A							)[0]
	print(our_output.shape							,their_output.shape							)
	lowercase						:		Optional[Any]   					=		torch.max(torch.abs(our_output - their_output							)							).item()
	print(F'''max_absolute_diff = {max_absolute_diff}'''							)  # ~ 1e-7
	lowercase						:		Dict   					=		torch.allclose(A							,A							,atol=1e-3							)
	print("Do both models output the same tensors?"							,"🔥" if success else "💩"							)
	if not success:
		raise Exception("Something went wRoNg"							)

	pathlib.Path(A							).mkdir(parents=A							,exist_ok=A							)
	print(F'''Saving model to {pytorch_dump_folder_path}'''							)
	model.save_pretrained(A							)


if __name__ == "__main__":
					lowerCAmelCase   : Union[str, Any]             =				argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    """--roberta_checkpoint_path""", default=None, type=str, required=True, help="""Path the official PyTorch dump."""
					)
					parser.add_argument(
					    """--pytorch_dump_folder_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
					)
					parser.add_argument(
					    """--classification_head""", action="""store_true""", help="""Whether to convert a final classification head."""
					)
					lowerCAmelCase   : Optional[Any]             =				parser.parse_args()
					convert_xlm_roberta_xl_checkpoint_to_pytorch(
					    args.roberta_checkpoint_path, args.pytorch_dump_folder_path, args.classification_head
					)


