





import argparse
import json
import os

import torch
from torch import nn

from transformers import NllbMoeConfig, NllbMoeModel
from transformers.modeling_utils import dtype_byte_size
from transformers.utils import WEIGHTS_INDEX_NAME, WEIGHTS_NAME







def   __lowerCamelCase   (				__a  :Optional[int]       )						->							List[str]:






		"""simple docstring"""

		A__     =  [
		    """encoder.version""",
		    """decoder.version""",
		    """model.encoder.version""",
		    """model.decoder.version""",
		    """decoder.output_projection.weight""",
		    """_float_tensor""",
		    """encoder.embed_positions._float_tensor""",
		    """decoder.embed_positions._float_tensor""",
		]
		for k in ignore_keys:
				state_dict.pop(__a       ,	__a       )







def   __lowerCamelCase   (				__a  :Optional[Any]       )						->							Optional[int]:






		"""simple docstring"""

		A__ ,     A__     =  emb.weight.shape
		A__     =  nn.Linear(__a       ,	__a       ,	bias=__a       )
		A__     =  emb.weight.data
		return lin_layer







def   __lowerCamelCase   (				__a  :Optional[Any]       ,	__a  :int=None       )						->							Union[str, Any]:






		"""simple docstring"""

		A__     =  {}
		for old_key in state_dict.keys():
				A__     =  old_key
				if "moe_layer.experts." in key:
						if expert_idx is not None:
								A__     =  key.replace("""moe_layer.experts.0"""       ,	F'ffn.experts.expert_{expert_idx}'       )
						else:
								A__     =  key.replace("""moe_layer.experts."""       ,	"""ffn.experts.expert_"""       )
				if "gate" in key:
						A__     =  key.replace(""".moe_layer.gate.wg"""       ,	""".ffn.router.classifier"""       )
				if "fc2" and "experts" not in key:
						A__     =  key.replace(""".fc2."""       ,	""".ffn.fc2."""       )
				if "fc1" and "experts" not in key:
						A__     =  key.replace(""".fc1."""       ,	""".ffn.fc1."""       )
				if ".encoder_attn." in key:
						A__     =  key.replace(""".encoder_attn."""       ,	""".cross_attention."""       )
				if "encoder_attn_layer_norm" in key:
						A__     =  key.replace("""encoder_attn_layer_norm"""       ,	"""cross_attention_layer_norm"""       )
				if "final_layer_norm" in key:
						A__     =  key.replace("""final_layer_norm"""       ,	"""ff_layer_norm"""       )
				A__     =  state_dict[old_key]
		return new_dict







def   __lowerCamelCase   (				__a  :Union[str, Any]       ,	__a  :Dict       ,	__a  :str       ,	__a  :List[str]       ,	__a  :str = WEIGHTS_NAME       )						->							int:






		"""simple docstring"""

		A__     =  []
		A__     =  0
		os.makedirs(__a       ,	exist_ok=__a       )

		for expert in range(__a       ):
				A__     =  switch_checkpoint_path + F'-rank-{expert}.pt'
				if os.path.isfile(__a       ):
						A__     =  torch.load(__a       )["""model"""]
						remove_ignore_keys_(__a       )
						A__     =  rename_fairseq_keys(__a       ,	__a       )
						A__     =  os.path.join(
						    __a       ,	weights_name.replace(""".bin"""       ,	F'-{len(__a       )+1:05d}-of-???.bin'       )       )
						torch.save(__a       ,	__a       )
						sharded_state_dicts.append(expert_state.keys()       )
						total_size += sum([value.numel() for key, value in expert_state.items()]       ) * dtype_byte_size(
						    expert_state[list(__a       )[0]].dtype       )

    # Add the last block
		A__     =  os.path.join(__a       ,	weights_name.replace(""".bin"""       ,	F'-{len(__a       )+1:05d}-of-???.bin'       )       )
		A__     =  torch.load(switch_checkpoint_path + """-shared.pt"""       )["""model"""]
		remove_ignore_keys_(__a       )
		A__     =  rename_fairseq_keys(__a       ,	__a       )
		A__     =  shared_weights["""decoder.embed_tokens.weight"""]
		sharded_state_dicts.append(shared_weights.keys()       )

		# If we only have the shared weights (dummy model/experts saved on the same file)
		if len(__a       ) == 1:
				A__     =  os.path.join(__a       ,	__a       )
				torch.save(__a       ,	__a       )
				return {weights_name: sharded_state_dicts[0]}, None
		else:
				torch.save(__a       ,	__a       )
		# Otherwise, let's build the index
		A__     =  {}
		for idx, shard in enumerate(__a       ):
				A__     =  weights_name.replace(""".bin"""       ,	F'-{idx+1:05d}-of-{len(__a       ):05d}.bin'       )
				A__     =  os.path.join(__a       ,	weights_name.replace(""".bin"""       ,	F'-{idx+1:05d}-of-???.bin'       )       )
				os.rename(__a       ,	os.path.join(__a       ,	__a       )       )
				for key in shard:
						A__     =  shard_file

    # Add the metadata
		A__     =  {"""total_size""": total_size}
		A__     =  {"""metadata""": metadata, """weight_map""": weight_map}

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

		return metadata, index


if __name__ == "__main__":
					A       :   Tuple					= argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    '''--nllb_moe_checkpoint_path''',
					    default='''/home/arthur_huggingface_co/fairseq/weights/checkpoints/model_moe_54b/checkpoint_2_300000''',
					    type=str,
					    required=False,
					    help='''Path to a directory containing a folder per layer. Follows the original Google format.''',
					)
					parser.add_argument('''--dtype''', default='''float32''', type=str, required=False, help='''dtype of the saved model''')
					parser.add_argument(
					    '''--pytorch_dump_folder_path''',
					    default='''/home/arthur_huggingface_co/fairseq/weights/checkpoints/hf-converted-moe-54b''',
					    type=str,
					    required=False,
					    help='''Path to the output pytorch model.''',
					)
					A       :   List[str]					= parser.parse_args()
					A,      A       :   List[str]					= shard_on_the_fly(
					    args.nllb_moe_checkpoint_path,
					    args.pytorch_dump_folder_path,
					    1_2_8,
					    args.dtype,
					)

					A       :   int					= NllbMoeConfig.from_pretrained(
					    '''facebook/nllb-200-3.3B''', encoder_sparse_step=4, decoder_sparse_step=4, num_experts=1_2_8
					)
					config.save_pretrained(args.pytorch_dump_folder_path)
					A       :   str					= NllbMoeModel.from_pretrained(args.pytorch_dump_folder_path)
					print('''Done''')
					model.save_pretrained(args.pytorch_dump_folder_path)










import re
import tempfile
from pathlib import Path

import pytest
import yaml

from datasets.utils.readme import ReadMe


# @pytest.fixture
# def example_yaml_structure():

A       :   Dict					= 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
'''
)


A       :   str					= {
    '''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''': []},
                    ],
                },
            ],
        }
    ],
}


A       :   Union[str, Any]					= '''\
---
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
'''


A       :   int					= '''\
---
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
'''

A       :   Optional[int]					= {
    '''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''': []},
                    ],
                },
            ],
        }
    ],
}

A       :   str					= '''\
---
---
# 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
'''

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

A       :   List[Any]					= '''\
# 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
'''

A       :   Union[str, Any]					= (
    '''The following issues were found for the README at `{path}`:\n-\tNo YAML markers are present in the README.'''
)

A       :   int					= '''\
---
# 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
'''

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

A       :   List[Any]					= '''\
---
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
'''
A       :   str					= '''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).'''


A       :   Optional[int]					= '''\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
'''
A       :   Optional[int]					= '''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\'.'''

A       :   Optional[int]					= '''\
---
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
'''

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


A       :   Optional[int]					= '''\
---
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
'''

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

A       :   List[Any]					= '''\
---
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
'''
A       :   Optional[Any]					= '''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.'''

A       :   int					= '''\
---
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
'''

A       :   Dict					= '''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.'''

A       :   str					= '''\
---
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
'''

A       :   int					= '''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.'''

A       :   str					= ''''''

A       :   Optional[Any]					= '''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.'''

A       :   int					= '''\
---
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
'''

A       :   List[Any]					= '''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   __lowerCamelCase   (				__a  :Union[str, Any]       ,	__a  :List[Any]       )						->							List[Any]:






		"""simple docstring"""

		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   __lowerCamelCase   (				__a  :Any       ,	__a  :Union[str, Any]       )						->							Dict:






		"""simple docstring"""

		with pytest.raises(__a       ,	match=re.escape(expected_error.format(path="""root"""       )       )       ):
				A__     =  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   __lowerCamelCase   (				__a  :List[str]       ,	__a  :Tuple       )						->							Dict:






		"""simple docstring"""

		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   __lowerCamelCase   (				__a  :Any       )						->							List[str]:






		"""simple docstring"""

		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   __lowerCamelCase   (				__a  :List[Any]       ,	__a  :Optional[int]       )						->							Optional[int]:






		"""simple docstring"""

		with tempfile.TemporaryDirectory() as tmp_dir:
				A__     =  Path(__a       ) / """README.md"""
				with open(__a       ,	"""w+"""       ) as readme_file:
						readme_file.write(__a       )
				A__     =  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   __lowerCamelCase   (				__a  :List[str]       ,	__a  :Union[str, Any]       )						->							List[str]:






		"""simple docstring"""

		with tempfile.TemporaryDirectory() as tmp_dir:
				A__     =  Path(__a       ) / """README.md"""
				with open(__a       ,	"""w+"""       ) as readme_file:
						readme_file.write(__a       )
				A__     =  expected_error.format(path=__a       )
				with pytest.raises(__a       ,	match=re.escape(__a       )       ):
						A__     =  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   __lowerCamelCase   (				__a  :List[Any]       ,	__a  :Any       )						->							int:






		"""simple docstring"""

		with tempfile.TemporaryDirectory() as tmp_dir:
				A__     =  Path(__a       ) / """README.md"""
				with open(__a       ,	"""w+"""       ) as readme_file:
						readme_file.write(__a       )
				A__     =  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   __lowerCamelCase   (				__a  :Union[str, Any]       )						->							Optional[Any]:






		"""simple docstring"""

		with tempfile.TemporaryDirectory() as tmp_dir:
				A__     =  Path(__a       ) / """README.md"""
				with open(__a       ,	"""w+"""       ) as readme_file:
						readme_file.write(__a       )
				ReadMe.from_readme(__a       ,	__a       ,	suppress_parsing_errors=__a       )



