






"""simple docstring"""






import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow


if is_torch_available():
     import torch

     from transformers import XLMRobertaModel




@require_sentencepiece
@require_tokenizers
@require_torch
class 							lowerCAmelCase__       (    unittest.TestCase				):






      '''simple docstring'''

      @slow
      def _SCREAMING_SNAKE_CASE				(		self  :						str):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	XLMRobertaModel.from_pretrained('''xlm-roberta-base''')
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	torch.tensor([[0, 581, 10269, 83, 99942, 136, 60742, 23, 70, 80583, 18276, 2]])
          # The dog is cute and lives in the garden house

          SCREAMING_SNAKE_CASE_   :						Dict							=	torch.Size((1, 12, 768))  # batch_size, sequence_length, embedding_vector_dim
          SCREAMING_SNAKE_CASE_   :						int							=	torch.tensor(
              [[-0.01_01, 0.12_18, -0.08_03, 0.08_01, 0.13_27, 0.07_76, -0.12_15, 0.23_83, 0.33_38, 0.31_06, 0.03_00, 0.02_52]])
          #  xlmr = torch.hub.load('pytorch/fairseq', 'xlmr.base')
          #  xlmr.eval()
          #  expected_output_values_last_dim = xlmr.extract_features(input_ids[0])[:, :, -1]
          with torch.no_grad():
              SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	model(lowercase_)['''last_hidden_state'''].detach()
          self.assertEqual(output.shape		,			lowercase_)
          # compare the actual values for a slice of last dim
          self.assertTrue(torch.allclose(output[:, :, -1]		,			lowercase_		,			atol=1e-3))



      @slow
      def _SCREAMING_SNAKE_CASE				(		self  :						Any):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						List[str]							=	XLMRobertaModel.from_pretrained('''xlm-roberta-large''')
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	torch.tensor([[0, 581, 10269, 83, 99942, 136, 60742, 23, 70, 80583, 18276, 2]])
          # The dog is cute and lives in the garden house

          SCREAMING_SNAKE_CASE_   :						Tuple							=	torch.Size((1, 12, 1024))  # batch_size, sequence_length, embedding_vector_dim
          SCREAMING_SNAKE_CASE_   :						str							=	torch.tensor(
              [[-0.06_99, -0.03_18, 0.07_05, -0.12_41, 0.09_99, -0.05_20, 0.10_04, -0.18_38, -0.47_04, 0.14_37, 0.08_21, 0.01_26]])
          #  xlmr = torch.hub.load('pytorch/fairseq', 'xlmr.large')
          #  xlmr.eval()
          #  expected_output_values_last_dim = xlmr.extract_features(input_ids[0])[:, :, -1]
          with torch.no_grad():
              SCREAMING_SNAKE_CASE_   :						Any							=	model(lowercase_)['''last_hidden_state'''].detach()
          self.assertEqual(output.shape		,			lowercase_)
          # compare the actual values for a slice of last dim
          self.assertTrue(torch.allclose(output[:, :, -1]		,			lowercase_		,			atol=1e-3))










"""simple docstring"""






import argparse
import os

import torch

from transformers import (
    XLNetConfig,
    XLNetForQuestionAnswering,
    XLNetForSequenceClassification,
    XLNetLMHeadModel,
    load_tf_weights_in_xlnet,
)
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


UpperCAmelCase_     :				int         =				{
    """cola""": 2,
    """mnli""": 3,
    """mrpc""": 2,
    """sst-2""": 2,
    """sts-b""": 1,
    """qqp""": 2,
    """qnli""": 2,
    """rte""": 2,
    """wnli""": 2,
}


logging.set_verbosity_info()



def       _A    (__a					,       __a					,       __a					,       __a=None							)				-> int:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_   :						str							=	XLNetConfig.from_json_file(__a							)

    SCREAMING_SNAKE_CASE_   :						List[Any]							=	finetuning_task.lower() if finetuning_task is not None else ''''''
    if finetuning_task in GLUE_TASKS_NUM_LABELS:
        print(f'Building PyTorch XLNetForSequenceClassification model from configuration: {config}'							)
        SCREAMING_SNAKE_CASE_   :						List[Any]							=	finetuning_task
        SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	GLUE_TASKS_NUM_LABELS[finetuning_task]
        SCREAMING_SNAKE_CASE_   :						Dict							=	XLNetForSequenceClassification(__a							)
    elif "squad" in finetuning_task:
        SCREAMING_SNAKE_CASE_   :						str							=	finetuning_task
        SCREAMING_SNAKE_CASE_   :						Dict							=	XLNetForQuestionAnswering(__a							)
    else:
        SCREAMING_SNAKE_CASE_   :						List[str]							=	XLNetLMHeadModel(__a							)

    # Load weights from tf checkpoint
    load_tf_weights_in_xlnet(__a					,       __a					,       __a							)

    # Save pytorch-model
    SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	os.path.join(__a					,       __a							)
    SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	os.path.join(__a					,       __a							)
    print(f'Save PyTorch model to {os.path.abspath(__a							)}'							)
    torch.save(model.state_dict()					,       __a							)
    print(f'Save configuration file to {os.path.abspath(__a							)}'							)
    with open(__a					,       '''w'''					,       encoding='''utf-8'''							) as f:
        f.write(config.to_json_string()							)


if __name__ == "__main__":
  UpperCAmelCase_     :				str         =				argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      """--tf_checkpoint_path""", default=None, type=str, required=True, help="""Path to the TensorFlow checkpoint path."""
  )
  parser.add_argument(
      """--xlnet_config_file""",
      default=None,
      type=str,
      required=True,
      help=(
          """The config json file corresponding to the pre-trained XLNet model. \n"""
          """This specifies the model architecture."""
      ),
  )
  parser.add_argument(
      """--pytorch_dump_folder_path""",
      default=None,
      type=str,
      required=True,
      help="""Path to the folder to store the PyTorch model or dataset/vocab.""",
  )
  parser.add_argument(
      """--finetuning_task""",
      default=None,
      type=str,
      help="""Name of a task on which the XLNet TensorFlow model was fine-tuned""",
  )
  UpperCAmelCase_     :				int         =				parser.parse_args()
  print(args)

  convert_xlnet_checkpoint_to_pytorch(
      args.tf_checkpoint_path, args.xlnet_config_file, args.pytorch_dump_folder_path, args.finetuning_task
  )


