





import time
from dataclasses import dataclass
from multiprocessing import Pool
from unittest import TestCase
from unittest.mock import patch

import multiprocess
import numpy as np
import pytest

from datasets.utils.py_utils import (
    NestedDataStructure,
    asdict,
    iflatmap_unordered,
    map_nested,
    temp_seed,
    temporary_assignment,
    zip_dict,
)

from .utils import require_tf, require_torch
def      _A       (		SCREAMING_SNAKE_CASE     :   Optional[int]						):  # picklable for multiprocessing





  """simple docstring"""



  return x.sum()
def      _A       (		SCREAMING_SNAKE_CASE     :   Tuple						):  # picklable for multiprocessing





  """simple docstring"""



  return i + 1







@dataclass
class 				__lowerCAmelCase       :
      _lowercase				:					int
      _lowercase				:					str







class 				__lowerCAmelCase       (    UpperCamelCase__):







      def    _lowercase  (    self							)   ->						Optional[Any]:



        '''simple docstring'''




        a__	:					Optional[int]										={}
        a__	:					List[Any]										=[]
        a__	:					List[str]										=1
        a__	:					int										=[1, 2]
        a__	:					Optional[Any]										={"a": 1, "b": 2}
        a__	:					str										={"a": [1, 2], "b": [3, 4]}
        a__	:					Union[str, Any]										={"a": {"1": 1}, "b": 2}
        a__	:					Optional[Any]										={"a": 1, "b": 2, "c": 3, "d": 4}
        a__	:					int										={}
        a__	:					List[Any]										=[]
        a__	:					Dict										=2
        a__	:					Dict										=[2, 3]
        a__	:					Union[str, Any]										={"a": 2, "b": 3}
        a__	:					List[str]										={"a": [2, 3], "b": [4, 5]}
        a__	:					Tuple										={"a": {"1": 2}, "b": 3}
        a__	:					str										={"a": 2, "b": 3, "c": 4, "d": 5}
        self.assertEqual(map_nested(lowerCAmelCase__		,      lowerCAmelCase__							)		,      lowerCAmelCase__							)
        self.assertEqual(map_nested(lowerCAmelCase__		,      lowerCAmelCase__							)		,      lowerCAmelCase__							)
        self.assertEqual(map_nested(lowerCAmelCase__		,      lowerCAmelCase__							)		,      lowerCAmelCase__							)
        self.assertEqual(map_nested(lowerCAmelCase__		,      lowerCAmelCase__							)		,      lowerCAmelCase__							)
        self.assertEqual(map_nested(lowerCAmelCase__		,      lowerCAmelCase__							)		,      lowerCAmelCase__							)
        self.assertEqual(map_nested(lowerCAmelCase__		,      lowerCAmelCase__							)		,      lowerCAmelCase__							)
        self.assertEqual(map_nested(lowerCAmelCase__		,      lowerCAmelCase__							)		,      lowerCAmelCase__							)
        self.assertEqual(map_nested(lowerCAmelCase__		,      lowerCAmelCase__							)		,      lowerCAmelCase__							)

        a__	:					str										=2
        self.assertEqual(map_nested(lowerCAmelCase__		,      lowerCAmelCase__		,      num_proc=lowerCAmelCase__							)		,      lowerCAmelCase__							)
        self.assertEqual(map_nested(lowerCAmelCase__		,      lowerCAmelCase__		,      num_proc=lowerCAmelCase__							)		,      lowerCAmelCase__							)
        self.assertEqual(map_nested(lowerCAmelCase__		,      lowerCAmelCase__		,      num_proc=lowerCAmelCase__							)		,      lowerCAmelCase__							)
        self.assertEqual(map_nested(lowerCAmelCase__		,      lowerCAmelCase__		,      num_proc=lowerCAmelCase__							)		,      lowerCAmelCase__							)
        self.assertEqual(map_nested(lowerCAmelCase__		,      lowerCAmelCase__		,      num_proc=lowerCAmelCase__							)		,      lowerCAmelCase__							)
        self.assertEqual(map_nested(lowerCAmelCase__		,      lowerCAmelCase__		,      num_proc=lowerCAmelCase__							)		,      lowerCAmelCase__							)
        self.assertEqual(map_nested(lowerCAmelCase__		,      lowerCAmelCase__		,      num_proc=lowerCAmelCase__							)		,      lowerCAmelCase__							)
        self.assertEqual(map_nested(lowerCAmelCase__		,      lowerCAmelCase__		,      num_proc=lowerCAmelCase__							)		,      lowerCAmelCase__							)

        a__	:					List[Any]										={"a": np.eye(2							), "b": np.zeros(3							), "c": np.ones(2							)}
        a__	:					Optional[int]										={"a": 2, "b": 0, "c": 2}
        a__	:					Any										={
            "a": np.eye(2							).astype(lowerCAmelCase__							),
            "b": np.zeros(3							).astype(lowerCAmelCase__							),
            "c": np.ones(2							).astype(lowerCAmelCase__							),
        }
        self.assertEqual(map_nested(lowerCAmelCase__		,      lowerCAmelCase__		,      map_numpy=lowerCAmelCase__							)		,      lowerCAmelCase__							)
        self.assertEqual(
            {k: v.tolist() for k, v in map_nested(lowerCAmelCase__		,      lowerCAmelCase__		,      map_numpy=lowerCAmelCase__							).items()}		,      {k: v.tolist() for k, v in expected_map_nested_sna_int.items()}		,      )
        self.assertEqual(map_nested(lowerCAmelCase__		,      lowerCAmelCase__		,      map_numpy=lowerCAmelCase__		,      num_proc=lowerCAmelCase__							)		,      lowerCAmelCase__							)
        self.assertEqual(
            {k: v.tolist() for k, v in map_nested(lowerCAmelCase__		,      lowerCAmelCase__		,      map_numpy=lowerCAmelCase__		,      num_proc=lowerCAmelCase__							).items()}		,      {k: v.tolist() for k, v in expected_map_nested_sna_int.items()}		,      )
        with self.assertRaises(lowerCAmelCase__							):  # can't pickle a local lambda
          map_nested(lambda lowerCAmelCase__							: x + 1		,      lowerCAmelCase__		,      num_proc=lowerCAmelCase__							)







      def    _lowercase  (    self							)   ->						List[str]:



        '''simple docstring'''




        a__	:					int										={"a": 1, "b": 2}
        a__	:					str										={"a": 3, "b": 4}
        a__	:					List[str]										={"a": 5, "b": 6}
        a__	:					Tuple										=sorted([("a", (1, 3, 5)), ("b", (2, 4, 6))]							)
        self.assertEqual(sorted(zip_dict(lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__							)							)		,      lowerCAmelCase__							)







      def    _lowercase  (    self							)   ->						Union[str, Any]:



        '''simple docstring'''




        class 				__lowerCAmelCase       :
              _lowercase				:					Optional[Any]						=			"""bar"""
        a__	:					Any										=Foo()
        self.assertEqual(foo.my_attr		,      "bar"							)
        with temporary_assignment(lowerCAmelCase__		,      "my_attr"		,      "BAR"							):
          self.assertEqual(foo.my_attr		,      "BAR"							)
        self.assertEqual(foo.my_attr		,      "bar"							)
@pytest.mark.parametrize(
    "iterable_length, num_proc, expected_num_proc"       ,       [
        (1, None, 1),
        (1, 1, 1),
        (2, None, 1),
        (2, 1, 1),
        (2, 2, 1),
        (2, 3, 1),
        (3, 2, 1),
        (16, 16, 16),
        (16, 17, 16),
        (17, 16, 16),
    ]       ,       )
def      _A       (		SCREAMING_SNAKE_CASE     :   Dict       ,       SCREAMING_SNAKE_CASE     :   Dict       ,       SCREAMING_SNAKE_CASE     :   str						):





  """simple docstring"""



  with patch("datasets.utils.py_utils._single_map_nested"						) as mock_single_map_nested, patch(
      "datasets.parallel.parallel.Pool"						) as mock_multiprocessing_pool:
    a__	:					Any										={f'''{i}''': i for i in range(SCREAMING_SNAKE_CASE						)}
    a__	:					Union[str, Any]										=map_nested(lambda SCREAMING_SNAKE_CASE						: x + 10       ,       SCREAMING_SNAKE_CASE       ,       num_proc=SCREAMING_SNAKE_CASE       ,       parallel_min_length=16						)
    if expected_num_proc == 1:
      assert mock_single_map_nested.called
      assert not mock_multiprocessing_pool.called
    else:
      assert not mock_single_map_nested.called
      assert mock_multiprocessing_pool.called
      assert mock_multiprocessing_pool.call_args[0][0] == expected_num_proc







class 				__lowerCAmelCase       (    UpperCamelCase__):







      @require_tf
      def    _lowercase  (    self							)   ->						int:



        '''simple docstring'''




        import tensorflow as tf
        from tensorflow.keras import layers

        a__	:					List[str]										=layers.Dense(2							)

        def gen_random_output():
          a__	:					Dict										=tf.random.uniform((1, 3)							)
          return model(lowerCAmelCase__							).numpy()

        with temp_seed(4_2		,      set_tensorflow=lowerCAmelCase__							):
          a__	:					str										=gen_random_output()
        with temp_seed(4_2		,      set_tensorflow=lowerCAmelCase__							):
          a__	:					Optional[Any]										=gen_random_output()
        a__	:					Optional[Any]										=gen_random_output()

        np.testing.assert_equal(lowerCAmelCase__		,      lowerCAmelCase__							)
        self.assertGreater(np.abs(outa - outa							).sum()		,      0							)







      @require_torch
      def    _lowercase  (    self							)   ->						Any:



        '''simple docstring'''




        import torch

        def gen_random_output():
          a__	:					Dict										=torch.nn.Linear(3		,      2							)
          a__	:					str										=torch.rand(1		,      3							)
          return model(lowerCAmelCase__							).detach().numpy()

        with temp_seed(4_2		,      set_pytorch=lowerCAmelCase__							):
          a__	:					Optional[int]										=gen_random_output()
        with temp_seed(4_2		,      set_pytorch=lowerCAmelCase__							):
          a__	:					Optional[Any]										=gen_random_output()
        a__	:					Optional[Any]										=gen_random_output()

        np.testing.assert_equal(lowerCAmelCase__		,      lowerCAmelCase__							)
        self.assertGreater(np.abs(outa - outa							).sum()		,      0							)







      def    _lowercase  (    self							)   ->						Any:



        '''simple docstring'''




        def gen_random_output():
          return np.random.rand(1		,      3							)

        with temp_seed(4_2							):
          a__	:					List[str]										=gen_random_output()
        with temp_seed(4_2							):
          a__	:					Optional[int]										=gen_random_output()
        a__	:					List[str]										=gen_random_output()

        np.testing.assert_equal(lowerCAmelCase__		,      lowerCAmelCase__							)
        self.assertGreater(np.abs(outa - outa							).sum()		,      0							)
@pytest.mark.parametrize("input_data"       ,       [{}]						)
def      _A       (		SCREAMING_SNAKE_CASE     :   Union[str, Any]						):





  """simple docstring"""



  a__	:					int										=NestedDataStructure(SCREAMING_SNAKE_CASE						).data
  assert output_data == input_data
@pytest.mark.parametrize(
    "data, expected_output"       ,       [
        ({}, []),
        ([], []),
        ("foo", ["foo"]),
        (["foo", "bar"], ["foo", "bar"]),
        ([["foo", "bar"]], ["foo", "bar"]),
        ([[["foo"], ["bar"]]], ["foo", "bar"]),
        ([[["foo"], "bar"]], ["foo", "bar"]),
        ({"a": 1, "b": 2}, [1, 2]),
        ({"a": [1, 2], "b": [3, 4]}, [1, 2, 3, 4]),
        ({"a": [[1, 2]], "b": [[3, 4]]}, [1, 2, 3, 4]),
        ({"a": [[1, 2]], "b": [3, 4]}, [1, 2, 3, 4]),
        ({"a": [[[1], [2]]], "b": [[[3], [4]]]}, [1, 2, 3, 4]),
        ({"a": [[[1], [2]]], "b": [[3, 4]]}, [1, 2, 3, 4]),
        ({"a": [[[1], [2]]], "b": [3, 4]}, [1, 2, 3, 4]),
        ({"a": [[[1], [2]]], "b": [3, [4]]}, [1, 2, 3, 4]),
        ({"a": {"1": 1}, "b": 2}, [1, 2]),
        ({"a": {"1": [1]}, "b": 2}, [1, 2]),
        ({"a": {"1": [1]}, "b": [2]}, [1, 2]),
    ]       ,       )
def      _A       (		SCREAMING_SNAKE_CASE     :   int       ,       SCREAMING_SNAKE_CASE     :   List[str]						):





  """simple docstring"""



  a__	:					Any										=NestedDataStructure(SCREAMING_SNAKE_CASE						).flatten()
  assert output == expected_output
def      _A       (		):





  """simple docstring"""



  a__	:					List[str]										=A(x=1       ,       y="foobar"						)
  a__	:					Optional[Any]										={"x": 1, "y": "foobar"}
  assert asdict(SCREAMING_SNAKE_CASE						) == expected_output

  a__	:					str										={"a": {"b": A(x=10       ,       y="foo"						)}, "c": [A(x=20       ,       y="bar"						)]}
  a__	:					Union[str, Any]										={"a": {"b": {"x": 10, "y": "foo"}}, "c": [{"x": 20, "y": "bar"}]}
  assert asdict(SCREAMING_SNAKE_CASE						) == expected_output

  with pytest.raises(SCREAMING_SNAKE_CASE						):
    asdict([1, A(x=10       ,       y="foo"						)]						)
def      _A       (		SCREAMING_SNAKE_CASE     :   str						):





  """simple docstring"""



  return text.split()
def      _A       (		SCREAMING_SNAKE_CASE     :   List[Any]						):





  """simple docstring"""



  yield (time.time(), content)
  time.sleep(2						)
  yield (time.time(), content)





def      _A       (		):





  """simple docstring"""



  with Pool(2						) as pool:
    a__	:					Union[str, Any]										=list(iflatmap_unordered(SCREAMING_SNAKE_CASE       ,       _split_text       ,       kwargs_iterable=[{"text": "hello there"}] * 10						)						)
    assert out.count("hello"						) == 10
    assert out.count("there"						) == 10
    assert len(SCREAMING_SNAKE_CASE						) == 20

  # check multiprocess from pathos (uses dill for pickling)
  with multiprocess.Pool(2						) as pool:
    a__	:					Dict										=list(iflatmap_unordered(SCREAMING_SNAKE_CASE       ,       _split_text       ,       kwargs_iterable=[{"text": "hello there"}] * 10						)						)
    assert out.count("hello"						) == 10
    assert out.count("there"						) == 10
    assert len(SCREAMING_SNAKE_CASE						) == 20

  # check that we get items as fast as possible
  with Pool(2						) as pool:
    a__	:					Dict										=[]
    for yield_time, content in iflatmap_unordered(
        SCREAMING_SNAKE_CASE       ,       _aseconds_generator_of_aitems_with_timing       ,       kwargs_iterable=[{"content": "a"}, {"content": "b"}]						):
      assert yield_time < time.time() + 0.1, "we should each item directly after it was yielded"
      out.append(SCREAMING_SNAKE_CASE						)
    assert out.count("a"						) == 2
    assert out.count("b"						) == 2
    assert len(SCREAMING_SNAKE_CASE						) == 4













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       (		SCREAMING_SNAKE_CASE     :   List[Any]       ,       SCREAMING_SNAKE_CASE     :   Dict       ,       SCREAMING_SNAKE_CASE     :   Dict       ,       SCREAMING_SNAKE_CASE     :   Union[str, Any]=None						):





  """simple docstring"""



  a__	:					Optional[int]										=XLNetConfig.from_json_file(SCREAMING_SNAKE_CASE						)

  a__	:					Dict										=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}'''						)
    a__	:					List[str]										=finetuning_task
    a__	:					Tuple										=GLUE_TASKS_NUM_LABELS[finetuning_task]
    a__	:					List[Any]										=XLNetForSequenceClassification(SCREAMING_SNAKE_CASE						)
  elif "squad" in finetuning_task:
    a__	:					Optional[int]										=finetuning_task
    a__	:					Dict										=XLNetForQuestionAnswering(SCREAMING_SNAKE_CASE						)
  else:
    a__	:					List[Any]										=XLNetLMHeadModel(SCREAMING_SNAKE_CASE						)

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

  # Save pytorch-model
  a__	:					Dict										=os.path.join(SCREAMING_SNAKE_CASE       ,       SCREAMING_SNAKE_CASE						)
  a__	:					Dict										=os.path.join(SCREAMING_SNAKE_CASE       ,       SCREAMING_SNAKE_CASE						)
  print(f'''Save PyTorch model to {os.path.abspath(SCREAMING_SNAKE_CASE						)}'''						)
  torch.save(model.state_dict()       ,       SCREAMING_SNAKE_CASE						)
  print(f'''Save configuration file to {os.path.abspath(SCREAMING_SNAKE_CASE						)}'''						)
  with open(SCREAMING_SNAKE_CASE       ,       "w"       ,       encoding="utf-8"						) as f:
    f.write(config.to_json_string()						)


if __name__ == "__main__":
     UpperCAmelCase	:	List[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
     )






