import re
import tempfile
from pathlib import Path

import pytest
import yaml

from datasets.utils.readme import ReadMe


# @pytest.fixture
# def example_yaml_structure():

__magic_name__								=yaml.safe_load(
    '''\
name: ""
allow_empty: false
allow_empty_text: true
subsections:
  - name: "Dataset Card for X" # First-level markdown heading
    allow_empty: false
    allow_empty_text: true
    subsections:
      - name: "Table of Contents"
        allow_empty: false
        allow_empty_text: false
        subsections: null
      - name: "Dataset Description"
        allow_empty: false
        allow_empty_text: false
        subsections:
          - name: "Dataset Summary"
            allow_empty: false
            allow_empty_text: false
            subsections: null
          - name: "Supported Tasks and Leaderboards"
            allow_empty: true
            allow_empty_text: true
            subsections: null
          - name: Languages
            allow_empty: false
            allow_empty_text: true
            subsections: null
'''
)


__magic_name__								={
    '''name''': '''root''',
    '''text''': '''''',
    '''is_empty_text''': True,
    '''subsections''': [
        {
            '''name''': '''Dataset Card for My Dataset''',
            '''text''': '''''',
            '''is_empty_text''': True,
            '''subsections''': [
                {'''name''': '''Table of Contents''', '''text''': '''Some text here.''', '''is_empty_text''': False, '''subsections''': []},
                {
                    '''name''': '''Dataset Description''',
                    '''text''': '''Some text here.''',
                    '''is_empty_text''': False,
                    '''subsections''': [
                        {
                            '''name''': '''Dataset Summary''',
                            '''text''': '''Some text here.''',
                            '''is_empty_text''': False,
                            '''subsections''': [],
                        },
                        {
                            '''name''': '''Supported Tasks and Leaderboards''',
                            '''text''': '''''',
                            '''is_empty_text''': True,
                            '''subsections''': [],
                        },
                        {'''name''': '''Languages''', '''text''': '''Language Text''', '''is_empty_text''': False, '''subsections''': []},
                    ],
                },
            ],
        }
    ],
}


__magic_name__								='''\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
Language Text
'''


__magic_name__								='''\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
#### Extra Ignored Subsection
### Supported Tasks and Leaderboards
### Languages
Language Text
'''

__magic_name__								={
    '''name''': '''root''',
    '''text''': '''''',
    '''is_empty_text''': True,
    '''subsections''': [
        {
            '''name''': '''Dataset Card for My Dataset''',
            '''text''': '''''',
            '''is_empty_text''': True,
            '''subsections''': [
                {'''name''': '''Table of Contents''', '''text''': '''Some text here.''', '''is_empty_text''': False, '''subsections''': []},
                {
                    '''name''': '''Dataset Description''',
                    '''text''': '''Some text here.''',
                    '''is_empty_text''': False,
                    '''subsections''': [
                        {
                            '''name''': '''Dataset Summary''',
                            '''text''': '''Some text here.''',
                            '''is_empty_text''': False,
                            '''subsections''': [
                                {
                                    '''name''': '''Extra Ignored Subsection''',
                                    '''text''': '''''',
                                    '''is_empty_text''': True,
                                    '''subsections''': [],
                                }
                            ],
                        },
                        {
                            '''name''': '''Supported Tasks and Leaderboards''',
                            '''text''': '''''',
                            '''is_empty_text''': True,
                            '''subsections''': [],
                        },
                        {'''name''': '''Languages''', '''text''': '''Language Text''', '''is_empty_text''': False, '''subsections''': []},
                    ],
                },
            ],
        }
    ],
}

__magic_name__								='''\
---
---
# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
Language Text
'''

__magic_name__								=(
    '''The following issues were found for the README at `{path}`:\n-\tEmpty YAML markers are present in the README.'''
)

__magic_name__								='''\
# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
Language Text
'''

__magic_name__								=(
    '''The following issues were found for the README at `{path}`:\n-\tNo YAML markers are present in the README.'''
)

__magic_name__								='''\
---
# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
Language Text
'''

__magic_name__								='''The following issues were found for the README at `{path}`:\n-\tOnly the start of YAML tags present in the README.'''

__magic_name__								='''\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
### Supported Tasks and Leaderboards
### Languages
Language Text
'''
__magic_name__								='''The following issues were found for the README at `{path}`:\n-\tExpected some content in section `Dataset Summary` but it is empty.\n-\tExpected some text in section `Dataset Summary` but it is empty (text in subsections are ignored).'''


__magic_name__								='''\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
'''
__magic_name__								='''The following issues were found for the README at `{path}`:\n-\tExpected some content in section `Dataset Card for My Dataset` but it is empty.\n-\tSection `Dataset Card for My Dataset` expected the following subsections: `Table of Contents`, `Dataset Description`. Found \'None\'.'''

__magic_name__								='''\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Languages
Language Text
'''

__magic_name__								='''The following issues were found for the README at `{path}`:\n-\tSection `Dataset Description` is missing subsection: `Supported Tasks and Leaderboards`.'''


__magic_name__								='''\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
'''

__magic_name__								='''The following issues were found for the README at `{path}`:\n-\tExpected some content in section `Languages` but it is empty.'''

__magic_name__								='''\
---
language:
- zh
- en
---

## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
Language Text
'''
__magic_name__								='''The following issues were found for the README at `{path}`:\n-\tThe README has no first-level headings. One heading is expected. Skipping further validation for this README.'''

__magic_name__								='''\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
Language Text
# Dataset Card My Dataset
'''

__magic_name__								='''The following issues were found for the README at `{path}`:\n-\tThe README has several first-level headings: `Dataset Card for My Dataset`, `Dataset Card My Dataset`. Only one heading is expected. Skipping further validation for this README.'''

__magic_name__								='''\
---
language:
- zh
- en
---

# Dataset Card My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
Language Text
'''

__magic_name__								='''The following issues were found for the README at `{path}`:\n-\tNo first-level heading starting with `Dataset Card for` found in README. Skipping further validation for this README.'''

__magic_name__								=''''''

__magic_name__								='''The following issues were found for the README at `{path}`:\n-\tThe README has no first-level headings. One heading is expected. Skipping further validation for this README.\n-\tNo YAML markers are present in the README.'''

__magic_name__								='''\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
Language Text
'''

__magic_name__								='''The following issues were found while parsing the README at `{path}`:\n-\tMultiple sections with the same heading `Dataset Card for My Dataset` have been found. Please keep only one of these sections.'''







@pytest.mark.parametrize(
    '''readme_md, expected_dict''' ,							[
        (README_CORRECT, CORRECT_DICT),
        (README_CORRECT_FOUR_LEVEL, CORRECT_DICT_FOUR_LEVEL),
    ] ,							)
def  __UpperCamelCase    (						A ,							A							):
     assert ReadMe.from_string(A ,							A							).to_dict() == expected_dict







@pytest.mark.parametrize(
    '''readme_md, expected_error''' ,							[
        (README_NO_YAML, EXPECTED_ERROR_README_NO_YAML),
        (README_EMPTY_YAML, EXPECTED_ERROR_README_EMPTY_YAML),
        (README_INCORRECT_YAML, EXPECTED_ERROR_README_INCORRECT_YAML),
        (README_EMPTY, EXPECTED_ERROR_README_EMPTY),
        (README_NONE_SUBSECTION, EXPECTED_ERROR_README_NONE_SUBSECTION),
        (README_MISSING_FIRST_LEVEL, EXPECTED_ERROR_README_MISSING_FIRST_LEVEL),
        (README_MISSING_SUBSECTION, EXPECTED_ERROR_README_MISSING_SUBSECTION),
        (README_MISSING_TEXT, EXPECTED_ERROR_README_MISSING_TEXT),
        (README_WRONG_FIRST_LEVEL, EXPECTED_ERROR_README_WRONG_FIRST_LEVEL),
        (README_MULTIPLE_WRONG_FIRST_LEVEL, EXPECTED_ERROR_README_MULTIPLE_WRONG_FIRST_LEVEL),
        (README_MISSING_CONTENT, EXPECTED_ERROR_README_MISSING_CONTENT),
    ] ,							)
def  __UpperCamelCase    (						A ,							A							):
     with pytest.raises(A ,							match=re.escape(expected_error.format(path='''root'''							)							)							):
          UpperCamelCase__										=      ReadMe.from_string(A ,							A							)
          readme.validate()







@pytest.mark.parametrize(
    '''readme_md, expected_error''' ,							[
        (README_MULTIPLE_SAME_HEADING_1, EXPECTED_ERROR_README_MULTIPLE_SAME_HEADING_1),
    ] ,							)
def  __UpperCamelCase    (						A ,							A							):
     with pytest.raises(A ,							match=re.escape(expected_error.format(path='''root'''							)							)							):
          ReadMe.from_string(A ,							A							)







@pytest.mark.parametrize(
    '''readme_md,''' ,							[
        (README_MULTIPLE_SAME_HEADING_1),
    ] ,							)
def  __UpperCamelCase    (						A							):
     ReadMe.from_string(A ,							A ,							suppress_parsing_errors=A							)







@pytest.mark.parametrize(
    '''readme_md, expected_dict''' ,							[
        (README_CORRECT, CORRECT_DICT),
        (README_CORRECT_FOUR_LEVEL, CORRECT_DICT_FOUR_LEVEL),
    ] ,							)
def  __UpperCamelCase    (						A ,							A							):
     with tempfile.TemporaryDirectory() as tmp_dir:
          UpperCamelCase__										=      Path(A							) / '''README.md'''
          with open(A ,							'''w+'''							) as readme_file:
               readme_file.write(A							)
          UpperCamelCase__										=      ReadMe.from_readme(A ,							A							).to_dict()
          assert out["name"] == path
          assert out["text"] == ""
          assert out["is_empty_text"]
          assert out["subsections"] == expected_dict["subsections"]







@pytest.mark.parametrize(
    '''readme_md, expected_error''' ,							[
        (README_NO_YAML, EXPECTED_ERROR_README_NO_YAML),
        (README_EMPTY_YAML, EXPECTED_ERROR_README_EMPTY_YAML),
        (README_INCORRECT_YAML, EXPECTED_ERROR_README_INCORRECT_YAML),
        (README_EMPTY, EXPECTED_ERROR_README_EMPTY),
        (README_NONE_SUBSECTION, EXPECTED_ERROR_README_NONE_SUBSECTION),
        (README_MISSING_FIRST_LEVEL, EXPECTED_ERROR_README_MISSING_FIRST_LEVEL),
        (README_MISSING_SUBSECTION, EXPECTED_ERROR_README_MISSING_SUBSECTION),
        (README_MISSING_TEXT, EXPECTED_ERROR_README_MISSING_TEXT),
        (README_WRONG_FIRST_LEVEL, EXPECTED_ERROR_README_WRONG_FIRST_LEVEL),
        (README_MULTIPLE_WRONG_FIRST_LEVEL, EXPECTED_ERROR_README_MULTIPLE_WRONG_FIRST_LEVEL),
        (README_MISSING_CONTENT, EXPECTED_ERROR_README_MISSING_CONTENT),
    ] ,							)
def  __UpperCamelCase    (						A ,							A							):
     with tempfile.TemporaryDirectory() as tmp_dir:
          UpperCamelCase__										=      Path(A							) / '''README.md'''
          with open(A ,							'''w+'''							) as readme_file:
               readme_file.write(A							)
          UpperCamelCase__										=      expected_error.format(path=A							)
          with pytest.raises(A ,							match=re.escape(A							)							):
               UpperCamelCase__										=      ReadMe.from_readme(A ,							A							)
               readme.validate()







@pytest.mark.parametrize(
    '''readme_md, expected_error''' ,							[
        (README_MULTIPLE_SAME_HEADING_1, EXPECTED_ERROR_README_MULTIPLE_SAME_HEADING_1),
    ] ,							)
def  __UpperCamelCase    (						A ,							A							):
     with tempfile.TemporaryDirectory() as tmp_dir:
          UpperCamelCase__										=      Path(A							) / '''README.md'''
          with open(A ,							'''w+'''							) as readme_file:
               readme_file.write(A							)
          UpperCamelCase__										=      expected_error.format(path=A							)
          with pytest.raises(A ,							match=re.escape(A							)							):
               ReadMe.from_readme(A ,							A							)







@pytest.mark.parametrize(
    '''readme_md,''' ,							[
        (README_MULTIPLE_SAME_HEADING_1),
    ] ,							)
def  __UpperCamelCase    (						A							):
     with tempfile.TemporaryDirectory() as tmp_dir:
          UpperCamelCase__										=      Path(A							) / '''README.md'''
          with open(A ,							'''w+'''							) as readme_file:
               readme_file.write(A							)
          ReadMe.from_readme(A ,							A ,							suppress_parsing_errors=A							)




import argparse
import json
import os

import tensorstore as ts
import torch
from flax import serialization
from flax.traverse_util import flatten_dict, unflatten_dict
from tensorflow.io import gfile

from transformers.modeling_utils import dtype_byte_size
from transformers.models.switch_transformers.convert_switch_transformers_original_flax_checkpoint_to_pytorch import (
    rename_keys,
)
from transformers.utils import WEIGHTS_INDEX_NAME, WEIGHTS_NAME
from transformers.utils.hub import convert_file_size_to_int







def  __UpperCamelCase    (						A ,							A							):
     if flax_key_tuple[-1] == "kernel" and flax_tensor.ndim == 3:
          # expert layer
          UpperCamelCase__										=      flax_key_tuple[:-1] + ('''weight''',)
          UpperCamelCase__										=      torch.permute(A ,							(0, 2, 1)							)
     elif flax_key_tuple[-1] == "kernel" and ".".join(A							):
          # linear layer
          UpperCamelCase__										=      flax_key_tuple[:-1] + ('''weight''',)
          UpperCamelCase__										=      flax_tensor.T
     elif flax_key_tuple[-1] in ["scale", "embedding"]:
          UpperCamelCase__										=      flax_key_tuple[:-1] + ('''weight''',)

     return flax_key_tuple, flax_tensor







def  __UpperCamelCase    (						A ,							A ,							A							):
     if "metadata" in layer:
          UpperCamelCase__										=      layer.split('''metadata'''							)
          UpperCamelCase__										=      ''''''.join(split_layer[0]							)[:-1]
          UpperCamelCase__										=      [tuple(('''metadata''' + split_layer[1]).split('''/'''							)							)]
     elif "kvstore" in layer:
          UpperCamelCase__										=      layer.split('''kvstore'''							)
          UpperCamelCase__										=      ''''''.join(split_layer[0]							)[:-1]
          UpperCamelCase__										=      [tuple(('''kvstore''' + split_layer[1]).split('''/'''							)							)]

     else:
          UpperCamelCase__										=      layer.split('''/'''							)
          UpperCamelCase__										=      '''/'''.join(split_layer[:-1]							)
          UpperCamelCase__										=      (split_layer[-1],)

     if "kvstore/path" in layer:
          UpperCamelCase__										=      f"{switch_checkpoint_path}/{checkpoint_info[layer]}"
     elif "kvstore/driver" in layer:
          UpperCamelCase__										=      '''file'''
     else:
          UpperCamelCase__										=      checkpoint_info[layer]

     return curr_real_layer_name, split_layer, content







def  __UpperCamelCase    (						A ,							A							):
     UpperCamelCase__										=      rename_keys(A							)
     UpperCamelCase__										=      {}
     for k, v in current_block.items():
          UpperCamelCase__										=      v
     UpperCamelCase__										=      new_current_block
     torch.save(A ,							A							)







def  __UpperCamelCase    (						A ,							A ,							A ,							A ,							A = WEIGHTS_NAME							):
     UpperCamelCase__										=      convert_file_size_to_int(A							)
     UpperCamelCase__										=      []
     UpperCamelCase__										=      {}
     UpperCamelCase__										=      0
     UpperCamelCase__										=      0

     os.makedirs(A ,							exist_ok=A							)
     with gfile.GFile(switch_checkpoint_path + '''/checkpoint''' ,							'''rb'''							) as fp:
          UpperCamelCase__										=      serialization.msgpack_restore(fp.read()							)['''optimizer''']['''target''']
          UpperCamelCase__										=      flatten_dict(A ,							sep='''/'''							)

     UpperCamelCase__										=      {}
     for layer in checkpoint_info.keys():
          UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__										=      get_key_and_tensorstore_dict(
              A ,							A ,							A							)
          if curr_real_layer_name in all_layers:
               UpperCamelCase__										=      content
          else:
               UpperCamelCase__										=      {split_layer[-1]: content}

     for key in all_layers.keys():
          # open tensorstore file
          UpperCamelCase__										=      ts.open(unflatten_dict(all_layers[key]							)							).result().read().result()
          UpperCamelCase__										=      torch.tensor(A							)
          UpperCamelCase__										=      raw_weights.numel() * dtype_byte_size(raw_weights.dtype							)

          # use the renaming pattern from the small conversion scripts
          UpperCamelCase__						,       UpperCamelCase__										=      rename_base_flax_keys(tuple(key.split('''/'''							)							) ,							A							)
          UpperCamelCase__										=      '''/'''.join(A							)

          # If this weight is going to tip up over the maximal size, we split.
          if current_block_size + weight_size > max_shard_size:
               UpperCamelCase__										=      os.path.join(
                   A ,							weights_name.replace('''.bin''' ,							f"-{len(A							)+1:05d}-of-???.bin"							)							)
               rename_and_save_block(A ,							A							)
               sharded_state_dicts.append(current_block.keys()							)
               del current_block
               UpperCamelCase__										=      {}
               UpperCamelCase__										=      0

          UpperCamelCase__										=      raw_weights.to(getattr(A ,							A							)							)
          current_block_size += weight_size
          total_size += weight_size

     # Add the last block
     UpperCamelCase__										=      os.path.join(A ,							weights_name.replace('''.bin''' ,							f"-{len(A							)+1:05d}-of-???.bin"							)							)
     rename_and_save_block(A ,							A							)
     sharded_state_dicts.append(current_block.keys()							)

     # If we only have one shard, we return it
     if len(A							) == 1:
          return {weights_name: sharded_state_dicts[0]}, None

     # Otherwise, let's build the index
     UpperCamelCase__										=      {}
     UpperCamelCase__										=      {}
     for idx, shard in enumerate(A							):
          UpperCamelCase__										=      weights_name.replace(
              '''.bin''' ,							f"-{idx+1:05d}-of-{len(A							):05d}.bin"							)  # len(sharded_state_dicts):05d}
          UpperCamelCase__										=      os.path.join(A ,							weights_name.replace('''.bin''' ,							f"-{idx+1:05d}-of-???.bin"							)							)
          os.rename(A ,							os.path.join(A ,							A							)							)
          UpperCamelCase__										=      shard
          for key in shard:
               UpperCamelCase__										=      shard_file

    # Add the metadata
     UpperCamelCase__										=      {'''total_size''': total_size}
     UpperCamelCase__										=      {'''metadata''': metadata, '''weight_map''': weight_map}

     with open(os.path.join(A ,							A							) ,							'''w''' ,							encoding='''utf-8'''							) as f:
          UpperCamelCase__										=      json.dumps(A ,							indent=2 ,							sort_keys=A							) + '''\n'''
          f.write(A							)

     return metadata, index


if __name__ == "__main__":
     __magic_name__								=argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         '''--switch_t5x_checkpoint_path''',
         default='''/mnt/disks/disk_switch/original_checkpoints/switch-xxl-128/checkpoint_634600''',
         type=str,
         required=False,
         help='''Path to a directory containing a folder per layer. Follows the original Google format.''',
     )
     parser.add_argument('''--max_shard_size''', default='''10GB''', required=False, help='''Max shard size''')
     parser.add_argument('''--dtype''', default='''bfloat16''', type=str, required=False, help='''dtype of the saved model''')
     parser.add_argument(
         '''--pytorch_dump_folder_path''',
         default='''/mnt/disks/disk_switch/original_checkpoints/switch-xxl-128-converted''',
         type=str,
         required=False,
         help='''Path to the output pytorch model.''',
     )
     __magic_name__								=parser.parse_args()
     shard_on_the_fly(
         args.switch_tax_checkpoint_path,
         args.pytorch_dump_folder_path,
         args.max_shard_size,
         args.dtype,
     )







def  __UpperCamelCase    (						):
     from transformers import SwitchTransformersConfig, SwitchTransformersForConditionalGeneration, TaTokenizer

     UpperCamelCase__										=      SwitchTransformersConfig.from_pretrained('''google/switch-base-8'''							)
     config.save_pretrained('''/home/arthur_huggingface_co/transformers/switch_converted'''							)
     UpperCamelCase__										=      SwitchTransformersForConditionalGeneration.from_pretrained(
         '''/home/arthur_huggingface_co/transformers/switch_converted''' ,							device_map='''auto'''							)

     UpperCamelCase__										=      TaTokenizer.from_pretrained('''t5-small'''							)
     UpperCamelCase__										=      '''A <extra_id_0> walks into a bar a orders a <extra_id_1> with <extra_id_2> pinch of <extra_id_3>.'''

     UpperCamelCase__										=      tokenizer(A ,							return_tensors='''pt'''							).input_ids
     UpperCamelCase__										=      model.generate(A ,							decoder_start_token_id=0							)
     print(tokenizer.decode(out[0]							)							)



