



import argparse
import pickle

import numpy as np
import torch
from torch import nn

from transformers import ReformerConfig, ReformerModelWithLMHead
from transformers.utils import logging


logging.set_verbosity_info()





def      __snake_case	(			_UpperCAmelCase   , _UpperCAmelCase   , _UpperCAmelCase=None				):
    # set parameter of one layer
    assert torch_layer.weight.shape == weight.shape, f'{torch_layer} layer.weight does not match'
    __a						       =				nn.Parameter(_UpperCAmelCase				)
    if bias is not None:
        assert torch_layer.bias.shape == bias.shape, f'{torch_layer} layer.bias does not match'
        __a						       =				nn.Parameter(_UpperCAmelCase				)





def      __snake_case	(			_UpperCAmelCase   , _UpperCAmelCase   , _UpperCAmelCase				):
    # set torch weights for 1-to-1 comparison
    __a						       =				np.asarray(weights[0]				)
    __a						       =				np.asarray(weights[1]				)
    __a						       =				np.asarray(weights[2]				)

    set_param(
        torch_layer.self_attention.query_key   , torch.tensor(_UpperCAmelCase				).transpose(1   , 2				).contiguous().view(-1   , _UpperCAmelCase				)   , )
    set_param(
        torch_layer.self_attention.value   , torch.tensor(_UpperCAmelCase				).transpose(1   , 2				).contiguous().view(-1   , _UpperCAmelCase				)   , )
    set_param(
        torch_layer.output.dense   , torch.tensor(_UpperCAmelCase				).view(-1   , _UpperCAmelCase				).contiguous().transpose(0   , 1				)   , )





def      __snake_case	(			_UpperCAmelCase   , _UpperCAmelCase   , _UpperCAmelCase				):
    # set torch weights for 1-to-1 comparison
    __a						       =				np.asarray(weights[0]				)
    __a						       =				np.asarray(weights[1]				)
    __a						       =				np.asarray(weights[2]				)
    __a						       =				np.asarray(weights[3]				)

    set_param(
        torch_layer.self_attention.query   , torch.tensor(_UpperCAmelCase				).transpose(1   , 2				).contiguous().view(-1   , _UpperCAmelCase				)   , )
    set_param(
        torch_layer.self_attention.key   , torch.tensor(_UpperCAmelCase				).transpose(1   , 2				).contiguous().view(-1   , _UpperCAmelCase				)   , )
    set_param(
        torch_layer.self_attention.value   , torch.tensor(_UpperCAmelCase				).transpose(1   , 2				).contiguous().view(-1   , _UpperCAmelCase				)   , )
    set_param(
        torch_layer.output.dense   , torch.tensor(_UpperCAmelCase				).view(-1   , _UpperCAmelCase				).contiguous().transpose(0   , 1				)   , )





def      __snake_case	(			_UpperCAmelCase   , _UpperCAmelCase   , _UpperCAmelCase				):
    # layernorm 1
    __a						       =				weights[0][0][0]
    __a						       =				np.asarray(layer_norm_a[0]				)
    __a						       =				np.asarray(layer_norm_a[1]				)
    set_param(
        torch_block.attention.layer_norm   , torch.tensor(_UpperCAmelCase				)   , torch.tensor(_UpperCAmelCase				)   , )

    # lsh weights + output
    __a						       =				weights[0][1]
    if len(_UpperCAmelCase				) < 4:
        set_layer_weights_in_torch_lsh(_UpperCAmelCase   , torch_block.attention   , _UpperCAmelCase				)
    else:
        set_layer_weights_in_torch_local(_UpperCAmelCase   , torch_block.attention   , _UpperCAmelCase				)

    # intermediate weighs
    __a						       =				weights[2][0][1][2]

    # Chunked Feed Forward
    if len(_UpperCAmelCase				) == 4:
        __a						       =				intermediate_weights[2]

    # layernorm 2
    __a						       =				np.asarray(intermediate_weights[0][0]				)
    __a						       =				np.asarray(intermediate_weights[0][1]				)
    set_param(
        torch_block.feed_forward.layer_norm   , torch.tensor(_UpperCAmelCase				)   , torch.tensor(_UpperCAmelCase				)   , )

    # intermediate dense
    __a						       =				np.asarray(intermediate_weights[1][0]				)
    __a						       =				np.asarray(intermediate_weights[1][1]				)
    set_param(
        torch_block.feed_forward.dense.dense   , torch.tensor(_UpperCAmelCase				).transpose(0   , 1				).contiguous()   , torch.tensor(_UpperCAmelCase				)   , )

    # intermediate out
    __a						       =				np.asarray(intermediate_weights[4][0]				)
    __a						       =				np.asarray(intermediate_weights[4][1]				)
    set_param(
        torch_block.feed_forward.output.dense   , torch.tensor(_UpperCAmelCase				).transpose(0   , 1				).contiguous()   , torch.tensor(_UpperCAmelCase				)   , )





def      __snake_case	(			_UpperCAmelCase   , _UpperCAmelCase   , _UpperCAmelCase				):
    # reformer model
    __a						       =				torch_model.reformer

    # word embeds
    __a						       =				np.asarray(weights[1]				)
    set_param(
        torch_model_reformer.embeddings.word_embeddings   , torch.tensor(_UpperCAmelCase				)   , )

    if isinstance(weights[3]   , _UpperCAmelCase				):
        __a						       =				torch_model_reformer.embeddings.position_embeddings
        for emb_idx in range(len(position_embeddings.weights				)				):
            __a						       =				np.asarray(weights[3][emb_idx][0]				)
            assert (
                position_embeddings.weights[emb_idx].shape == emb_weights.shape
            ), f'{position_embeddings[emb_idx]} emb does not match'
            __a						       =				nn.Parameter(torch.tensor(_UpperCAmelCase				)				)

    __a						       =				weights[5]
    assert len(torch_model_reformer.encoder.layers				) * 4 == len(
        _UpperCAmelCase				), "HF and trax model do not have the same number of layers"
    for layer_idx, layer in enumerate(torch_model_reformer.encoder.layers				):
        __a						       =				trax_layer_weights[4 * layer_idx : 4 * (layer_idx + 1)]
        set_block_weights_in_torch(_UpperCAmelCase   , _UpperCAmelCase   , _UpperCAmelCase				)

    # output layer norm
    __a						       =				np.asarray(weights[7][0]				)
    __a						       =				np.asarray(weights[7][1]				)
    set_param(
        torch_model_reformer.encoder.layer_norm   , torch.tensor(_UpperCAmelCase				)   , torch.tensor(_UpperCAmelCase				)   , )

    # output embeddings
    __a						       =				np.asarray(weights[9][0]				)
    __a						       =				np.asarray(weights[9][1]				)
    set_param(
        torch_model.lm_head.decoder   , torch.tensor(_UpperCAmelCase				).transpose(0   , 1				).contiguous()   , torch.tensor(_UpperCAmelCase				)   , )





def      __snake_case	(			_UpperCAmelCase   , _UpperCAmelCase   , _UpperCAmelCase				):
    # Initialise PyTorch model
    __a						       =				ReformerConfig.from_json_file(_UpperCAmelCase				)
    print(f'Building PyTorch model from configuration: {config}'				)
    __a						       =				ReformerModelWithLMHead(_UpperCAmelCase				)

    with open(_UpperCAmelCase   , '''rb'''				) as f:
        __a						       =				pickle.load(_UpperCAmelCase				)['''weights''']

    set_model_weights_in_torch(_UpperCAmelCase   , _UpperCAmelCase   , config.hidden_size				)

    # Save pytorch-model
    print(f'Save PyTorch model to {pytorch_dump_path}'				)
    torch.save(model.state_dict()   , _UpperCAmelCase				)


if __name__ == "__main__":
 __snake_case						:str      		=     argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     '''--trax_model_pkl_path''', default=None, type=str, required=True, help='''Path to the TensorFlow checkpoint path.'''
 )
 parser.add_argument(
     '''--config_file''',
     default=None,
     type=str,
     required=True,
     help=(
         '''The config json file corresponding to the pre-trained Reformer model. \n'''
         '''This specifies the model architecture.'''
     ),
 )
 parser.add_argument(
     '''--pytorch_dump_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
 )
 __snake_case						:Optional[Any]      		=     parser.parse_args()
 convert_trax_checkpoint_to_pytorch(args.trax_model_pkl_path, args.config_file, args.pytorch_dump_path)










import os
from pathlib import Path
from unittest.mock import patch

import pytest
import zstandard as zstd

from datasets.download.download_config import DownloadConfig
from datasets.utils.file_utils import (
    OfflineModeIsEnabled,
    cached_path,
    fsspec_get,
    fsspec_head,
    ftp_get,
    ftp_head,
    get_from_cache,
    http_get,
    http_head,
)


__snake_case						:List[str]      		=     '''\
    Text data.
    Second line of data.'''

__snake_case						:Optional[Any]      		=     '''file'''





@pytest.fixture(scope='''session'''				)
def      __snake_case	(			_UpperCAmelCase				):
    __a						       =				tmp_path_factory.mktemp('''data'''				) / (FILE_PATH + '''.zstd''')
    __a						       =				bytes(_UpperCAmelCase   , '''utf-8'''				)
    with zstd.open(_UpperCAmelCase   , '''wb'''				) as f:
        f.write(_UpperCAmelCase				)
    return path





@pytest.fixture
def      __snake_case	(			_UpperCAmelCase				):
    with open(os.path.join(tmpfs.local_root_dir   , _UpperCAmelCase				)   , '''w'''				) as f:
        f.write(_UpperCAmelCase				)
    return FILE_PATH





@pytest.mark.parametrize('''compression_format'''   , ['''gzip''', '''xz''', '''zstd''']				)
def      __snake_case	(			_UpperCAmelCase   , _UpperCAmelCase   , _UpperCAmelCase   , _UpperCAmelCase   , _UpperCAmelCase   , _UpperCAmelCase				):
    __a						       =				{'''gzip''': gz_file, '''xz''': xz_file, '''zstd''': zstd_path}
    __a						       =				input_paths[compression_format]
    __a						       =				tmp_path / '''cache'''
    __a						       =				DownloadConfig(cache_dir=_UpperCAmelCase   , extract_compressed_file=_UpperCAmelCase				)
    __a						       =				cached_path(_UpperCAmelCase   , download_config=_UpperCAmelCase				)
    with open(_UpperCAmelCase				) as f:
        __a						       =				f.read()
    with open(_UpperCAmelCase				) as f:
        __a						       =				f.read()
    assert extracted_file_content == expected_file_content





@pytest.mark.parametrize('''default_extracted'''   , [True, False]				)
@pytest.mark.parametrize('''default_cache_dir'''   , [True, False]				)
def      __snake_case	(			_UpperCAmelCase   , _UpperCAmelCase   , _UpperCAmelCase   , _UpperCAmelCase   , _UpperCAmelCase				):
    __a						       =				'''custom_cache'''
    __a						       =				'''custom_extracted_dir'''
    __a						       =				tmp_path / '''custom_extracted_path'''
    if default_extracted:
        __a						       =				('''downloads''' if default_cache_dir else custom_cache_dir, '''extracted''')
    else:
        monkeypatch.setattr('''datasets.config.EXTRACTED_DATASETS_DIR'''   , _UpperCAmelCase				)
        monkeypatch.setattr('''datasets.config.EXTRACTED_DATASETS_PATH'''   , str(_UpperCAmelCase				)				)
        __a						       =				custom_extracted_path.parts[-2:] if default_cache_dir else (custom_cache_dir, custom_extracted_dir)

    __a						       =				xz_file
    __a						       =				(
        DownloadConfig(extract_compressed_file=_UpperCAmelCase				)
        if default_cache_dir
        else DownloadConfig(cache_dir=tmp_path / custom_cache_dir   , extract_compressed_file=_UpperCAmelCase				)
    )
    __a						       =				cached_path(_UpperCAmelCase   , download_config=_UpperCAmelCase				)
    assert Path(_UpperCAmelCase				).parent.parts[-2:] == expected





def      __snake_case	(			_UpperCAmelCase				):
    # absolute path
    __a						       =				str(Path(_UpperCAmelCase				).resolve()				)
    assert cached_path(_UpperCAmelCase				) == text_file
    # relative path
    __a						       =				str(Path(_UpperCAmelCase				).resolve().relative_to(Path(os.getcwd()				)				)				)
    assert cached_path(_UpperCAmelCase				) == text_file





def      __snake_case	(			_UpperCAmelCase				):
    # absolute path
    __a						       =				str(tmp_path.resolve() / '''__missing_file__.txt'''				)
    with pytest.raises(_UpperCAmelCase				):
        cached_path(_UpperCAmelCase				)
    # relative path
    __a						       =				'''./__missing_file__.txt'''
    with pytest.raises(_UpperCAmelCase				):
        cached_path(_UpperCAmelCase				)





def      __snake_case	(			_UpperCAmelCase				):
    __a						       =				get_from_cache(f'tmp://{tmpfs_file}'				)
    with open(_UpperCAmelCase				) as f:
        __a						       =				f.read()
    assert output_file_content == FILE_CONTENT





@patch('''datasets.config.HF_DATASETS_OFFLINE'''   , _UpperCAmelCase				)
def      __snake_case	(			):
    with pytest.raises(_UpperCAmelCase				):
        cached_path('''https://huggingface.co'''				)





@patch('''datasets.config.HF_DATASETS_OFFLINE'''   , _UpperCAmelCase				)
def      __snake_case	(			_UpperCAmelCase				):
    __a						       =				tmp_path_factory.mktemp('''data'''				) / '''file.html'''
    with pytest.raises(_UpperCAmelCase				):
        http_get('''https://huggingface.co'''   , temp_file=_UpperCAmelCase				)
    with pytest.raises(_UpperCAmelCase				):
        http_head('''https://huggingface.co'''				)





@patch('''datasets.config.HF_DATASETS_OFFLINE'''   , _UpperCAmelCase				)
def      __snake_case	(			_UpperCAmelCase				):
    __a						       =				tmp_path_factory.mktemp('''data'''				) / '''file.html'''
    with pytest.raises(_UpperCAmelCase				):
        ftp_get('''ftp://huggingface.co'''   , temp_file=_UpperCAmelCase				)
    with pytest.raises(_UpperCAmelCase				):
        ftp_head('''ftp://huggingface.co'''				)





@patch('''datasets.config.HF_DATASETS_OFFLINE'''   , _UpperCAmelCase				)
def      __snake_case	(			_UpperCAmelCase				):
    __a						       =				tmp_path_factory.mktemp('''data'''				) / '''file.html'''
    with pytest.raises(_UpperCAmelCase				):
        fsspec_get('''s3://huggingface.co'''   , temp_file=_UpperCAmelCase				)
    with pytest.raises(_UpperCAmelCase				):
        fsspec_head('''s3://huggingface.co'''				)





