feat: Put jax checkpoint in a dedicated folder

jax_model/hyperparams.json

@@ -0,0 +1,26 @@
+{"alphabet_size": 4107, 
+"pad_token_id": 1, 
+"mask_token_id": 2, 
+"max_positions": 2048,
+"k_for_kmers": 6,
+"embed_scale": 1.0, 
+"emb_layer_norm_before": false, 
+"attention_heads": 16, 
+"key_size": 32, 
+"embed_dim": 512, 
+"ffn_embed_dim": 2048, 
+"num_layers": 12, 
+"positional_embedding": null, 
+"lm_head": "roberta", 
+"add_bias_kv": false, 
+"add_bias_ffn": false, 
+"use_rotary_embedding": true, 
+"ffn_activation_name": "swish", 
+"use_glu_in_ffn": true, 
+"mask_before_attention": false, 
+"token_dropout": false, 
+"masking_ratio": 0.0, 
+"masking_prob": 0.0, 
+"use_gradient_checkpointing": false, 
+"embeddings_layers_to_save": [], 
+"attention_maps_to_save": []}