





import os

from ...utils import logging
from .utils import DataProcessor, InputExample


__A       : Union[str, Any]						    = logging.get_logger(__name__)







class     lowerCamelCase(	__snake_case  ):
     '''simple docstring'''


     def __init__(						self							,    snake_case_							,    snake_case_=None       ):
         _A    =       language
         _A    =       train_language




     def 					lowerCAmelCase__		(						self							,    snake_case_       ):
         _A    =       self.language if self.train_language is None else self.train_language
         _A    =       self._read_tsv(os.path.join(snake_case_							,    F"XNLI-MT-1.0/multinli/multinli.train.{lg}.tsv"       )       )
         _A    =       []
         for i, line in enumerate(snake_case_       ):
             if i == 0:
                 continue
             _A    =       F"train-{i}"
             _A    =       line[0]
             _A    =       line[1]
             _A    =       'contradiction' if line[2] == 'contradictory' else line[2]
             if not isinstance(snake_case_							,    snake_case_       ):
                 raise ValueError(F"Training input {text_a} is not a string"       )
             if not isinstance(snake_case_							,    snake_case_       ):
                 raise ValueError(F"Training input {text_b} is not a string"       )
             if not isinstance(snake_case_							,    snake_case_       ):
                 raise ValueError(F"Training label {label} is not a string"       )
             examples.append(InputExample(guid=snake_case_							,    text_a=snake_case_							,    text_b=snake_case_							,    label=snake_case_       )       )
         return examples




     def 					lowerCAmelCase__		(						self							,    snake_case_       ):
         _A    =       self._read_tsv(os.path.join(snake_case_							,    'XNLI-1.0/xnli.test.tsv'       )       )
         _A    =       []
         for i, line in enumerate(snake_case_       ):
             if i == 0:
                 continue
             _A    =       line[0]
             if language != self.language:
                 continue
             _A    =       F"test-{i}"
             _A    =       line[6]
             _A    =       line[7]
             _A    =       line[1]
             if not isinstance(snake_case_							,    snake_case_       ):
                 raise ValueError(F"Training input {text_a} is not a string"       )
             if not isinstance(snake_case_							,    snake_case_       ):
                 raise ValueError(F"Training input {text_b} is not a string"       )
             if not isinstance(snake_case_							,    snake_case_       ):
                 raise ValueError(F"Training label {label} is not a string"       )
             examples.append(InputExample(guid=snake_case_							,    text_a=snake_case_							,    text_b=snake_case_							,    label=snake_case_       )       )
         return examples




     def 					lowerCAmelCase__		(						self       ):
         return ["contradiction", "entailment", "neutral"]




__A       : str						    = {
    "xnli": XnliProcessor,
}

__A       : List[Any]						    = {
    "xnli": "classification",
}

__A       : Tuple						    = {
    "xnli": 3,
}












import argparse

import torch
from flax.training.checkpoints import restore_checkpoint

from transformers import FNetConfig, FNetForPreTraining
from transformers.utils import logging


logging.set_verbosity_info()





def 							__lowerCAmelCase(					_SCREAMING_SNAKE_CASE       ,    _SCREAMING_SNAKE_CASE       ,    _SCREAMING_SNAKE_CASE							)       ->   List[str]:




    """simple docstring"""




    _A    =       FNetConfig.from_json_file(_SCREAMING_SNAKE_CASE							)
    print(F"Building PyTorch model from configuration: {config}"							)
    _A    =       FNetForPreTraining(_SCREAMING_SNAKE_CASE							)

    _A    =       restore_checkpoint(_SCREAMING_SNAKE_CASE       ,    _SCREAMING_SNAKE_CASE							)
    _A    =       checkpoint_dict['target']

    # Embeddings
    # Position IDs
    _A    =       fnet_pretraining_model.state_dict()

    _A    =       state_dict['fnet.embeddings.position_ids']
    _A    =       {'fnet.embeddings.position_ids': position_ids}
    # Embedding Layers
    _A    =       torch.tensor(
        pretrained_model_params['encoder']['embedder']['word']['embedding']							)
    _A    =       torch.tensor(
        pretrained_model_params['encoder']['embedder']['position']['embedding'][0]							)
    _A    =       torch.tensor(
        pretrained_model_params['encoder']['embedder']['type']['embedding']							)
    _A    =       torch.tensor(
        pretrained_model_params['encoder']['embedder']['hidden_mapping_in']['kernel']							).T
    _A    =       torch.tensor(
        pretrained_model_params['encoder']['embedder']['hidden_mapping_in']['bias']							)
    _A    =       torch.tensor(
        pretrained_model_params['encoder']['embedder']['layer_norm']['scale']							)
    _A    =       torch.tensor(
        pretrained_model_params['encoder']['embedder']['layer_norm']['bias']							)

    # Encoder Layers
    for layer in range(config.num_hidden_layers							):
        _A    =       torch.tensor(
            pretrained_model_params['encoder'][F"encoder_{layer}"]['mixing_layer_norm']['scale']							)
        _A    =       torch.tensor(
            pretrained_model_params['encoder'][F"encoder_{layer}"]['mixing_layer_norm']['bias']							)

        _A    =       torch.tensor(
            pretrained_model_params['encoder'][F"feed_forward_{layer}"]['intermediate']['kernel']							).T
        _A    =       torch.tensor(
            pretrained_model_params['encoder'][F"feed_forward_{layer}"]['intermediate']['bias']							)

        _A    =       torch.tensor(
            pretrained_model_params['encoder'][F"feed_forward_{layer}"]['output']['kernel']							).T
        _A    =       torch.tensor(
            pretrained_model_params['encoder'][F"feed_forward_{layer}"]['output']['bias']							)

        _A    =       torch.tensor(
            pretrained_model_params['encoder'][F"encoder_{layer}"]['output_layer_norm']['scale']							)
        _A    =       torch.tensor(
            pretrained_model_params['encoder'][F"encoder_{layer}"]['output_layer_norm']['bias']							)

    # Pooler Layers
    _A    =       torch.tensor(pretrained_model_params['encoder']['pooler']['kernel']							).T
    _A    =       torch.tensor(pretrained_model_params['encoder']['pooler']['bias']							)

    # Masked LM Layers
    _A    =       torch.tensor(
        pretrained_model_params['predictions_dense']['kernel']							).T
    _A    =       torch.tensor(
        pretrained_model_params['predictions_dense']['bias']							)
    _A    =       torch.tensor(
        pretrained_model_params['predictions_layer_norm']['scale']							)
    _A    =       torch.tensor(
        pretrained_model_params['predictions_layer_norm']['bias']							)
    _A    =       torch.tensor(
        pretrained_model_params['encoder']['embedder']['word']['embedding']							)
    _A    =       torch.tensor(
        pretrained_model_params['predictions_output']['output_bias']							)
    _A    =       torch.tensor(pretrained_model_params['predictions_output']['output_bias']							)

    # Seq Relationship Layers
    _A    =       torch.tensor(
        pretrained_model_params['classification']['output_kernel']							)
    _A    =       torch.tensor(
        pretrained_model_params['classification']['output_bias']							)

    # Load State Dict
    fnet_pretraining_model.load_state_dict(_SCREAMING_SNAKE_CASE							)

    # Save PreTrained
    print(F"Saving pretrained model to {save_path}"							)
    fnet_pretraining_model.save_pretrained(_SCREAMING_SNAKE_CASE							)


if __name__ == "__main__":
 __A       : str						    = argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     "--flax_checkpoint_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
 )
 parser.add_argument(
     "--fnet_config_file",
     default=None,
     type=str,
     required=True,
     help=(
         "The config json file corresponding to the pre-trained FNet model. \n"
         "This specifies the model architecture."
     ),
 )
 parser.add_argument("--save_path", default=None, type=str, required=True, help="Path to the output model.")
 __A       : Any						    = parser.parse_args()
 convert_flax_checkpoint_to_pytorch(args.flax_checkpoint_path, args.fnet_config_file, args.save_path)





