


import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_bert import BertTokenizer


__lowerCAmelCase			=							logging.get_logger(__name__)

__lowerCAmelCase			=							{'''vocab_file''': '''vocab.txt''', '''tokenizer_file''': '''tokenizer.json'''}

__lowerCAmelCase			=							{
    '''vocab_file''': {
        '''bert-base-uncased''': '''https://huggingface.co/bert-base-uncased/resolve/main/vocab.txt''',
        '''bert-large-uncased''': '''https://huggingface.co/bert-large-uncased/resolve/main/vocab.txt''',
        '''bert-base-cased''': '''https://huggingface.co/bert-base-cased/resolve/main/vocab.txt''',
        '''bert-large-cased''': '''https://huggingface.co/bert-large-cased/resolve/main/vocab.txt''',
        '''bert-base-multilingual-uncased''': (
            '''https://huggingface.co/bert-base-multilingual-uncased/resolve/main/vocab.txt'''
        ),
        '''bert-base-multilingual-cased''': '''https://huggingface.co/bert-base-multilingual-cased/resolve/main/vocab.txt''',
        '''bert-base-chinese''': '''https://huggingface.co/bert-base-chinese/resolve/main/vocab.txt''',
        '''bert-base-german-cased''': '''https://huggingface.co/bert-base-german-cased/resolve/main/vocab.txt''',
        '''bert-large-uncased-whole-word-masking''': (
            '''https://huggingface.co/bert-large-uncased-whole-word-masking/resolve/main/vocab.txt'''
        ),
        '''bert-large-cased-whole-word-masking''': (
            '''https://huggingface.co/bert-large-cased-whole-word-masking/resolve/main/vocab.txt'''
        ),
        '''bert-large-uncased-whole-word-masking-finetuned-squad''': (
            '''https://huggingface.co/bert-large-uncased-whole-word-masking-finetuned-squad/resolve/main/vocab.txt'''
        ),
        '''bert-large-cased-whole-word-masking-finetuned-squad''': (
            '''https://huggingface.co/bert-large-cased-whole-word-masking-finetuned-squad/resolve/main/vocab.txt'''
        ),
        '''bert-base-cased-finetuned-mrpc''': (
            '''https://huggingface.co/bert-base-cased-finetuned-mrpc/resolve/main/vocab.txt'''
        ),
        '''bert-base-german-dbmdz-cased''': '''https://huggingface.co/bert-base-german-dbmdz-cased/resolve/main/vocab.txt''',
        '''bert-base-german-dbmdz-uncased''': (
            '''https://huggingface.co/bert-base-german-dbmdz-uncased/resolve/main/vocab.txt'''
        ),
        '''TurkuNLP/bert-base-finnish-cased-v1''': (
            '''https://huggingface.co/TurkuNLP/bert-base-finnish-cased-v1/resolve/main/vocab.txt'''
        ),
        '''TurkuNLP/bert-base-finnish-uncased-v1''': (
            '''https://huggingface.co/TurkuNLP/bert-base-finnish-uncased-v1/resolve/main/vocab.txt'''
        ),
        '''wietsedv/bert-base-dutch-cased''': (
            '''https://huggingface.co/wietsedv/bert-base-dutch-cased/resolve/main/vocab.txt'''
        ),
    },
    '''tokenizer_file''': {
        '''bert-base-uncased''': '''https://huggingface.co/bert-base-uncased/resolve/main/tokenizer.json''',
        '''bert-large-uncased''': '''https://huggingface.co/bert-large-uncased/resolve/main/tokenizer.json''',
        '''bert-base-cased''': '''https://huggingface.co/bert-base-cased/resolve/main/tokenizer.json''',
        '''bert-large-cased''': '''https://huggingface.co/bert-large-cased/resolve/main/tokenizer.json''',
        '''bert-base-multilingual-uncased''': (
            '''https://huggingface.co/bert-base-multilingual-uncased/resolve/main/tokenizer.json'''
        ),
        '''bert-base-multilingual-cased''': (
            '''https://huggingface.co/bert-base-multilingual-cased/resolve/main/tokenizer.json'''
        ),
        '''bert-base-chinese''': '''https://huggingface.co/bert-base-chinese/resolve/main/tokenizer.json''',
        '''bert-base-german-cased''': '''https://huggingface.co/bert-base-german-cased/resolve/main/tokenizer.json''',
        '''bert-large-uncased-whole-word-masking''': (
            '''https://huggingface.co/bert-large-uncased-whole-word-masking/resolve/main/tokenizer.json'''
        ),
        '''bert-large-cased-whole-word-masking''': (
            '''https://huggingface.co/bert-large-cased-whole-word-masking/resolve/main/tokenizer.json'''
        ),
        '''bert-large-uncased-whole-word-masking-finetuned-squad''': (
            '''https://huggingface.co/bert-large-uncased-whole-word-masking-finetuned-squad/resolve/main/tokenizer.json'''
        ),
        '''bert-large-cased-whole-word-masking-finetuned-squad''': (
            '''https://huggingface.co/bert-large-cased-whole-word-masking-finetuned-squad/resolve/main/tokenizer.json'''
        ),
        '''bert-base-cased-finetuned-mrpc''': (
            '''https://huggingface.co/bert-base-cased-finetuned-mrpc/resolve/main/tokenizer.json'''
        ),
        '''bert-base-german-dbmdz-cased''': (
            '''https://huggingface.co/bert-base-german-dbmdz-cased/resolve/main/tokenizer.json'''
        ),
        '''bert-base-german-dbmdz-uncased''': (
            '''https://huggingface.co/bert-base-german-dbmdz-uncased/resolve/main/tokenizer.json'''
        ),
        '''TurkuNLP/bert-base-finnish-cased-v1''': (
            '''https://huggingface.co/TurkuNLP/bert-base-finnish-cased-v1/resolve/main/tokenizer.json'''
        ),
        '''TurkuNLP/bert-base-finnish-uncased-v1''': (
            '''https://huggingface.co/TurkuNLP/bert-base-finnish-uncased-v1/resolve/main/tokenizer.json'''
        ),
        '''wietsedv/bert-base-dutch-cased''': (
            '''https://huggingface.co/wietsedv/bert-base-dutch-cased/resolve/main/tokenizer.json'''
        ),
    },
}

__lowerCAmelCase			=							{
    '''bert-base-uncased''': 5_12,
    '''bert-large-uncased''': 5_12,
    '''bert-base-cased''': 5_12,
    '''bert-large-cased''': 5_12,
    '''bert-base-multilingual-uncased''': 5_12,
    '''bert-base-multilingual-cased''': 5_12,
    '''bert-base-chinese''': 5_12,
    '''bert-base-german-cased''': 5_12,
    '''bert-large-uncased-whole-word-masking''': 5_12,
    '''bert-large-cased-whole-word-masking''': 5_12,
    '''bert-large-uncased-whole-word-masking-finetuned-squad''': 5_12,
    '''bert-large-cased-whole-word-masking-finetuned-squad''': 5_12,
    '''bert-base-cased-finetuned-mrpc''': 5_12,
    '''bert-base-german-dbmdz-cased''': 5_12,
    '''bert-base-german-dbmdz-uncased''': 5_12,
    '''TurkuNLP/bert-base-finnish-cased-v1''': 5_12,
    '''TurkuNLP/bert-base-finnish-uncased-v1''': 5_12,
    '''wietsedv/bert-base-dutch-cased''': 5_12,
}

__lowerCAmelCase			=							{
    '''bert-base-uncased''': {'''do_lower_case''': True},
    '''bert-large-uncased''': {'''do_lower_case''': True},
    '''bert-base-cased''': {'''do_lower_case''': False},
    '''bert-large-cased''': {'''do_lower_case''': False},
    '''bert-base-multilingual-uncased''': {'''do_lower_case''': True},
    '''bert-base-multilingual-cased''': {'''do_lower_case''': False},
    '''bert-base-chinese''': {'''do_lower_case''': False},
    '''bert-base-german-cased''': {'''do_lower_case''': False},
    '''bert-large-uncased-whole-word-masking''': {'''do_lower_case''': True},
    '''bert-large-cased-whole-word-masking''': {'''do_lower_case''': False},
    '''bert-large-uncased-whole-word-masking-finetuned-squad''': {'''do_lower_case''': True},
    '''bert-large-cased-whole-word-masking-finetuned-squad''': {'''do_lower_case''': False},
    '''bert-base-cased-finetuned-mrpc''': {'''do_lower_case''': False},
    '''bert-base-german-dbmdz-cased''': {'''do_lower_case''': False},
    '''bert-base-german-dbmdz-uncased''': {'''do_lower_case''': True},
    '''TurkuNLP/bert-base-finnish-cased-v1''': {'''do_lower_case''': False},
    '''TurkuNLP/bert-base-finnish-uncased-v1''': {'''do_lower_case''': True},
    '''wietsedv/bert-base-dutch-cased''': {'''do_lower_case''': False},
}





class   __a      (   __UpperCamelCase				):
 __lowercase			:							List[Any]       					=			VOCAB_FILES_NAMES
 __lowercase			:							Any       					=			PRETRAINED_VOCAB_FILES_MAP
 __lowercase			:							int       					=			PRETRAINED_INIT_CONFIGURATION
 __lowercase			:							int       					=			PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 __lowercase			:							Dict       					=			BertTokenizer
 def __init__(       self    ,			lowerCAmelCase__=None    ,			lowerCAmelCase__=None    ,			lowerCAmelCase__=True    ,			lowerCAmelCase__="[UNK]"    ,			lowerCAmelCase__="[SEP]"    ,			lowerCAmelCase__="[PAD]"    ,			lowerCAmelCase__="[CLS]"    ,			lowerCAmelCase__="[MASK]"    ,			lowerCAmelCase__=True    ,			lowerCAmelCase__=None    ,			**lowerCAmelCase__    ,			) ->    Any:




        '''simple docstring'''
        super().__init__(
            lowerCAmelCase__    ,			tokenizer_file=lowerCAmelCase__    ,			do_lower_case=lowerCAmelCase__    ,			unk_token=lowerCAmelCase__    ,			sep_token=lowerCAmelCase__    ,			pad_token=lowerCAmelCase__    ,			cls_token=lowerCAmelCase__    ,			mask_token=lowerCAmelCase__    ,			tokenize_chinese_chars=lowerCAmelCase__    ,			strip_accents=lowerCAmelCase__    ,			**lowerCAmelCase__    ,			)

        lowercase__:       Union[str, Any]        =       json.loads(self.backend_tokenizer.normalizer.__getstate__()			)
        if (
            normalizer_state.get('lowercase'    ,			lowerCAmelCase__			) != do_lower_case
            or normalizer_state.get('strip_accents'    ,			lowerCAmelCase__			) != strip_accents
            or normalizer_state.get('handle_chinese_chars'    ,			lowerCAmelCase__			) != tokenize_chinese_chars
        ):
               lowercase__:       Optional[int]        =       getattr(lowerCAmelCase__    ,			normalizer_state.pop('type'			)			)
               lowercase__:       Dict        =       do_lower_case
               lowercase__:       List[Any]        =       strip_accents
               lowercase__:       Tuple        =       tokenize_chinese_chars
               lowercase__:       Tuple        =       normalizer_class(**lowerCAmelCase__			)

        lowercase__:       Any        =       do_lower_case
 def 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__    ,			lowerCAmelCase__=None			) ->    Union[str, Any]:




        '''simple docstring'''
        lowercase__:       str        =       [self.cls_token_id] + token_ids_a + [self.sep_token_id]

        if token_ids_a:
               output += token_ids_a + [self.sep_token_id]

        return output
 def 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__    ,			lowerCAmelCase__ = None			) ->    List[int]:




        '''simple docstring'''
        lowercase__:       str        =       [self.sep_token_id]
        lowercase__:       List[str]        =       [self.cls_token_id]
        if token_ids_a is None:
               return len(cls + token_ids_a + sep			) * [0]
        return len(cls + token_ids_a + sep			) * [0] + len(token_ids_a + sep			) * [1]
 def 		SCREAMING_SNAKE_CASE__       (       self    ,			lowerCAmelCase__    ,			lowerCAmelCase__ = None			) ->    Tuple[str]:




        '''simple docstring'''
        lowercase__:       Any        =       self._tokenizer.model.save(lowerCAmelCase__    ,			name=lowerCAmelCase__			)
        return tuple(lowerCAmelCase__			)






import argparse
import os

import jax as jnp
import numpy as onp
import torch
import torch.nn as nn
from music_spectrogram_diffusion import inference
from tax import checkpoints

from diffusers import DDPMScheduler, OnnxRuntimeModel, SpectrogramDiffusionPipeline
from diffusers.pipelines.spectrogram_diffusion import SpectrogramContEncoder, SpectrogramNotesEncoder, TaFilmDecoder


__lowerCAmelCase			=							'''base_with_context'''




def        snake_case_					(  snake_case      ,							snake_case					)  ->    int:
       lowercase__:       Tuple        =       nn.Parameter(torch.FloatTensor(weights['token_embedder']['embedding']					)					)
       lowercase__:       Optional[int]        =       nn.Parameter(
           torch.FloatTensor(weights['Embed_0']['embedding']					)      ,							requires_grad=snake_case					)
       for lyr_num, lyr in enumerate(model.encoders					):
              lowercase__:       List[str]        =       weights[f'layers_{lyr_num}']
              lowercase__:       List[Any]        =       nn.Parameter(
                  torch.FloatTensor(ly_weight['pre_attention_layer_norm']['scale']					)					)

              lowercase__:       Any        =       ly_weight['attention']
              lowercase__:       int        =       nn.Parameter(torch.FloatTensor(attention_weights['query']['kernel'].T					)					)
              lowercase__:       int        =       nn.Parameter(torch.FloatTensor(attention_weights['key']['kernel'].T					)					)
              lowercase__:       List[str]        =       nn.Parameter(torch.FloatTensor(attention_weights['value']['kernel'].T					)					)
              lowercase__:       Dict        =       nn.Parameter(torch.FloatTensor(attention_weights['out']['kernel'].T					)					)

              lowercase__:       List[Any]        =       nn.Parameter(torch.FloatTensor(ly_weight['pre_mlp_layer_norm']['scale']					)					)

              lowercase__:       Optional[Any]        =       nn.Parameter(torch.FloatTensor(ly_weight['mlp']['wi_0']['kernel'].T					)					)
              lowercase__:       Dict        =       nn.Parameter(torch.FloatTensor(ly_weight['mlp']['wi_1']['kernel'].T					)					)
              lowercase__:       List[str]        =       nn.Parameter(torch.FloatTensor(ly_weight['mlp']['wo']['kernel'].T					)					)

       lowercase__:       Any        =       nn.Parameter(torch.FloatTensor(weights['encoder_norm']['scale']					)					)
       return model




def        snake_case_					(  snake_case      ,							snake_case					)  ->    List[str]:
       lowercase__:       str        =       nn.Parameter(torch.FloatTensor(weights['input_proj']['kernel'].T					)					)

       lowercase__:       Dict        =       nn.Parameter(
           torch.FloatTensor(weights['Embed_0']['embedding']					)      ,							requires_grad=snake_case					)

       for lyr_num, lyr in enumerate(model.encoders					):
              lowercase__:       str        =       weights[f'layers_{lyr_num}']
              lowercase__:       Optional[Any]        =       ly_weight['attention']

              lowercase__:       List[Any]        =       nn.Parameter(torch.FloatTensor(attention_weights['query']['kernel'].T					)					)
              lowercase__:       Dict        =       nn.Parameter(torch.FloatTensor(attention_weights['key']['kernel'].T					)					)
              lowercase__:       int        =       nn.Parameter(torch.FloatTensor(attention_weights['value']['kernel'].T					)					)
              lowercase__:       Dict        =       nn.Parameter(torch.FloatTensor(attention_weights['out']['kernel'].T					)					)
              lowercase__:       Union[str, Any]        =       nn.Parameter(
                  torch.FloatTensor(ly_weight['pre_attention_layer_norm']['scale']					)					)

              lowercase__:       List[str]        =       nn.Parameter(torch.FloatTensor(ly_weight['mlp']['wi_0']['kernel'].T					)					)
              lowercase__:       Dict        =       nn.Parameter(torch.FloatTensor(ly_weight['mlp']['wi_1']['kernel'].T					)					)
              lowercase__:       Tuple        =       nn.Parameter(torch.FloatTensor(ly_weight['mlp']['wo']['kernel'].T					)					)
              lowercase__:       Optional[int]        =       nn.Parameter(torch.FloatTensor(ly_weight['pre_mlp_layer_norm']['scale']					)					)

       lowercase__:       List[str]        =       nn.Parameter(torch.FloatTensor(weights['encoder_norm']['scale']					)					)

       return model




def        snake_case_					(  snake_case      ,							snake_case					)  ->    Any:
       lowercase__:       int        =       nn.Parameter(torch.FloatTensor(weights['time_emb_dense0']['kernel'].T					)					)
       lowercase__:       Any        =       nn.Parameter(torch.FloatTensor(weights['time_emb_dense1']['kernel'].T					)					)

       lowercase__:       int        =       nn.Parameter(
           torch.FloatTensor(weights['Embed_0']['embedding']					)      ,							requires_grad=snake_case					)

       lowercase__:       Dict        =       nn.Parameter(
           torch.FloatTensor(weights['continuous_inputs_projection']['kernel'].T					)					)

       for lyr_num, lyr in enumerate(model.decoders					):
              lowercase__:       Optional[Any]        =       weights[f'layers_{lyr_num}']
              lowercase__:       Any        =       nn.Parameter(
                  torch.FloatTensor(ly_weight['pre_self_attention_layer_norm']['scale']					)					)

              lowercase__:       int        =       nn.Parameter(
                  torch.FloatTensor(ly_weight['FiLMLayer_0']['DenseGeneral_0']['kernel'].T					)					)

              lowercase__:       List[str]        =       ly_weight['self_attention']
              lowercase__:       Optional[int]        =       nn.Parameter(torch.FloatTensor(attention_weights['query']['kernel'].T					)					)
              lowercase__:       Any        =       nn.Parameter(torch.FloatTensor(attention_weights['key']['kernel'].T					)					)
              lowercase__:       Tuple        =       nn.Parameter(torch.FloatTensor(attention_weights['value']['kernel'].T					)					)
              lowercase__:       Union[str, Any]        =       nn.Parameter(torch.FloatTensor(attention_weights['out']['kernel'].T					)					)

              lowercase__:       int        =       ly_weight['MultiHeadDotProductAttention_0']
              lowercase__:       List[Any]        =       nn.Parameter(torch.FloatTensor(attention_weights['query']['kernel'].T					)					)
              lowercase__:       Dict        =       nn.Parameter(torch.FloatTensor(attention_weights['key']['kernel'].T					)					)
              lowercase__:       str        =       nn.Parameter(torch.FloatTensor(attention_weights['value']['kernel'].T					)					)
              lowercase__:       Any        =       nn.Parameter(torch.FloatTensor(attention_weights['out']['kernel'].T					)					)
              lowercase__:       int        =       nn.Parameter(
                  torch.FloatTensor(ly_weight['pre_cross_attention_layer_norm']['scale']					)					)

              lowercase__:       List[str]        =       nn.Parameter(torch.FloatTensor(ly_weight['pre_mlp_layer_norm']['scale']					)					)
              lowercase__:       Union[str, Any]        =       nn.Parameter(
                  torch.FloatTensor(ly_weight['FiLMLayer_1']['DenseGeneral_0']['kernel'].T					)					)
              lowercase__:       List[str]        =       nn.Parameter(torch.FloatTensor(ly_weight['mlp']['wi_0']['kernel'].T					)					)
              lowercase__:       int        =       nn.Parameter(torch.FloatTensor(ly_weight['mlp']['wi_1']['kernel'].T					)					)
              lowercase__:       str        =       nn.Parameter(torch.FloatTensor(ly_weight['mlp']['wo']['kernel'].T					)					)

       lowercase__:       Optional[Any]        =       nn.Parameter(torch.FloatTensor(weights['decoder_norm']['scale']					)					)

       lowercase__:       Union[str, Any]        =       nn.Parameter(torch.FloatTensor(weights['spec_out_dense']['kernel'].T					)					)

       return model




def        snake_case_					(  snake_case					)  ->    Any:
       lowercase__:       int        =       checkpoints.load_tax_checkpoint(args.checkpoint_path					)
       lowercase__:       Tuple        =       jnp.tree_util.tree_map(onp.array      ,							snake_case					)

       lowercase__:       List[str]        =       [
           'from __gin__ import dynamic_registration',
           'from music_spectrogram_diffusion.models.diffusion import diffusion_utils',
           'diffusion_utils.ClassifierFreeGuidanceConfig.eval_condition_weight = 2.0',
           'diffusion_utils.DiffusionConfig.classifier_free_guidance = @diffusion_utils.ClassifierFreeGuidanceConfig()',
       ]

       lowercase__:       List[Any]        =       os.path.join(args.checkpoint_path      ,							'..'      ,							'config.gin'					)
       lowercase__:       Optional[Any]        =       inference.parse_training_gin_file(snake_case      ,							snake_case					)
       lowercase__:       str        =       inference.InferenceModel(args.checkpoint_path      ,							snake_case					)

       lowercase__:       Dict        =       DDPMScheduler(beta_schedule='squaredcos_cap_v2'      ,							variance_type='fixed_large'					)

       lowercase__:       List[Any]        =       SpectrogramNotesEncoder(
           max_length=synth_model.sequence_length['inputs']      ,							vocab_size=synth_model.model.module.config.vocab_size      ,							d_model=synth_model.model.module.config.emb_dim      ,							dropout_rate=synth_model.model.module.config.dropout_rate      ,							num_layers=synth_model.model.module.config.num_encoder_layers      ,							num_heads=synth_model.model.module.config.num_heads      ,							d_kv=synth_model.model.module.config.head_dim      ,							d_ff=synth_model.model.module.config.mlp_dim      ,							feed_forward_proj='gated-gelu'      ,							)

       lowercase__:       Dict        =       SpectrogramContEncoder(
           input_dims=synth_model.audio_codec.n_dims      ,							targets_context_length=synth_model.sequence_length['targets_context']      ,							d_model=synth_model.model.module.config.emb_dim      ,							dropout_rate=synth_model.model.module.config.dropout_rate      ,							num_layers=synth_model.model.module.config.num_encoder_layers      ,							num_heads=synth_model.model.module.config.num_heads      ,							d_kv=synth_model.model.module.config.head_dim      ,							d_ff=synth_model.model.module.config.mlp_dim      ,							feed_forward_proj='gated-gelu'      ,							)

       lowercase__:       Optional[Any]        =       TaFilmDecoder(
           input_dims=synth_model.audio_codec.n_dims      ,							targets_length=synth_model.sequence_length['targets_context']      ,							max_decoder_noise_time=synth_model.model.module.config.max_decoder_noise_time      ,							d_model=synth_model.model.module.config.emb_dim      ,							num_layers=synth_model.model.module.config.num_decoder_layers      ,							num_heads=synth_model.model.module.config.num_heads      ,							d_kv=synth_model.model.module.config.head_dim      ,							d_ff=synth_model.model.module.config.mlp_dim      ,							dropout_rate=synth_model.model.module.config.dropout_rate      ,							)

       lowercase__:       Dict        =       load_notes_encoder(ta_checkpoint['target']['token_encoder']      ,							snake_case					)
       lowercase__:       int        =       load_continuous_encoder(ta_checkpoint['target']['continuous_encoder']      ,							snake_case					)
       lowercase__:       Optional[int]        =       load_decoder(ta_checkpoint['target']['decoder']      ,							snake_case					)

       lowercase__:       int        =       OnnxRuntimeModel.from_pretrained('kashif/soundstream_mel_decoder'					)

       lowercase__:       List[Any]        =       SpectrogramDiffusionPipeline(
           notes_encoder=snake_case      ,							continuous_encoder=snake_case      ,							decoder=snake_case      ,							scheduler=snake_case      ,							melgan=snake_case      ,							)
       if args.save:
              pipe.save_pretrained(args.output_path					)


if __name__ == "__main__":
 __lowerCAmelCase			=							argparse.ArgumentParser()

 parser.add_argument('''--output_path''', default=None, type=str, required=True, help='''Path to the converted model.''')
 parser.add_argument(
     '''--save''', default=True, type=bool, required=False, help='''Whether to save the converted model or not.'''
 )
 parser.add_argument(
     '''--checkpoint_path''',
     default=F'''{MODEL}/checkpoint_500000''',
     type=str,
     required=False,
     help='''Path to the original jax model checkpoint.''',
 )
 __lowerCAmelCase			=							parser.parse_args()

 main(args)


