diff --git a/.gitignore b/.gitignore
index 0cff484033362ac4616c0085b4eb726f0154f53d..8fec32c62fb90125c9c61805b28a6539546cde27 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,9 @@
 **/__pycache__
 .ipynb_checkpoints
+scripts
+demo2.ipynb
+*.py
+*.ipynb
+instructany2pix
+assets
 # ckpts
\ No newline at end of file
diff --git a/assets/demo/an antique shop.jpg b/assets/demo/an antique shop.jpg
deleted file mode 100644
index 4ac3133dbd043e916b0a79a100f77824e72c43ca..0000000000000000000000000000000000000000
Binary files a/assets/demo/an antique shop.jpg and /dev/null differ
diff --git a/assets/demo/clock ticking.wav b/assets/demo/clock ticking.wav
deleted file mode 100644
index 07fe53f4f6929e47ec562a8d351e6f1d4fcd3fb3..0000000000000000000000000000000000000000
Binary files a/assets/demo/clock ticking.wav and /dev/null differ
diff --git a/ckpts/imagebind_huge.pth b/ckpts/imagebind_huge.pth
deleted file mode 120000
index b5b417db27cf1245f117359df739c6787b868ab2..0000000000000000000000000000000000000000
--- a/ckpts/imagebind_huge.pth
+++ /dev/null
@@ -1 +0,0 @@
-/home/jacklishufan/AudioLDM2/.checkpoints/imagebind_huge.pth
\ No newline at end of file
diff --git a/ckpts/imagebind_huge.pth b/ckpts/imagebind_huge.pth
new file mode 100644
index 0000000000000000000000000000000000000000..21709a82dbb35fca042d46877b89074bc1885401
--- /dev/null
+++ b/ckpts/imagebind_huge.pth
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d6f6c22bedcc90708448d5d2fbb7b2db9c73f505dc89bd0b2e09b23af1b62157
+size 4803584173
diff --git a/ckpts/llm b/ckpts/llm
deleted file mode 120000
index 0ca57910c38b981db010dd82c60d96f28361ecae..0000000000000000000000000000000000000000
--- a/ckpts/llm
+++ /dev/null
@@ -1 +0,0 @@
-/home/jacklishufan/LLaVA/outputs/llava-v1.5-7b-pretrain-any2any2align/
\ No newline at end of file
diff --git a/ckpts/llm/added_tokens.json b/ckpts/llm/added_tokens.json
new file mode 100644
index 0000000000000000000000000000000000000000..33d7d87bf4ccf7c3f652eb834528c9376657c947
--- /dev/null
+++ b/ckpts/llm/added_tokens.json
@@ -0,0 +1,11 @@
+{
+  "<audio>": 32005,
+  "<audio_gen>": 32003,
+  "<audio_gen_start>": 32004,
+  "<base>": 32007,
+  "<base_null>": 32008,
+  "<im_gen>": 32001,
+  "<im_gen_start>": 32000,
+  "<mask_gen>": 32002,
+  "<video>": 32006
+}
diff --git a/ckpts/llm/config.json b/ckpts/llm/config.json
new file mode 100644
index 0000000000000000000000000000000000000000..624f338f7264ad9c2e41903140ebc0f08d30ef4d
--- /dev/null
+++ b/ckpts/llm/config.json
@@ -0,0 +1,42 @@
+{
+  "_name_or_path": "/home/jacklishufan/LLaVA/outputs/any2anyalign",
+  "architectures": [
+    "LlavaLlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "freeze_mm_mlp_adapter": false,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "image_aspect_ratio": "square",
+  "image_grid_pinpoints": null,
+  "initializer_range": 0.02,
+  "intermediate_size": 11008,
+  "max_position_embeddings": 4096,
+  "mm_hidden_size": 1024,
+  "mm_projector_type": "mlp2x_gelu",
+  "mm_use_im_patch_token": false,
+  "mm_use_im_start_end": false,
+  "mm_vision_select_feature": "projection",
+  "mm_vision_select_layer": -2,
+  "mm_vision_tower": "laion/CLIP-ViT-H-14-laion2B-s32B-b79K",
+  "model_type": "llava",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 32,
+  "num_key_value_heads": 32,
+  "pad_token_id": 0,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": null,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.34.1",
+  "tune_mm_mlp_adapter": false,
+  "use_cache": true,
+  "use_mm_proj": true,
+  "vae_audio": "vq-npz",
+  "vae_image": "vq-npz",
+  "vocab_size": 32009
+}
diff --git a/ckpts/llm/generation_config.json b/ckpts/llm/generation_config.json
new file mode 100644
index 0000000000000000000000000000000000000000..c26d216dc64711b76750a0d69db303a23484a49b
--- /dev/null
+++ b/ckpts/llm/generation_config.json
@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 0,
+  "transformers_version": "4.34.1"
+}
diff --git a/ckpts/llm/pytorch_model-00001-of-00002.bin b/ckpts/llm/pytorch_model-00001-of-00002.bin
new file mode 100644
index 0000000000000000000000000000000000000000..73ac2cf545973af7297d60fb3afaac43b035fc20
--- /dev/null
+++ b/ckpts/llm/pytorch_model-00001-of-00002.bin
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e2135464f292dfe61635e4365706e03f0b116b539e1fd3edc5da952026aafe50
+size 9976693850
diff --git a/ckpts/llm/pytorch_model-00002-of-00002.bin b/ckpts/llm/pytorch_model-00002-of-00002.bin
new file mode 100644
index 0000000000000000000000000000000000000000..475861928ea2f25de9cec6014e39e98cf5394987
--- /dev/null
+++ b/ckpts/llm/pytorch_model-00002-of-00002.bin
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4eb4805933b421b55a1b4bf4096db8eeb5a35ca996305b32e6d4fe3f785bc0e9
+size 4974527229
diff --git a/ckpts/llm/pytorch_model.bin.index.json b/ckpts/llm/pytorch_model.bin.index.json
new file mode 100644
index 0000000000000000000000000000000000000000..5214aa703f2699ff8983dfb64a35fae60664f502
--- /dev/null
+++ b/ckpts/llm/pytorch_model.bin.index.json
@@ -0,0 +1,838 @@
+{
+  "metadata": {
+    "total_size": 14950917120
+  },
+  "weight_map": {
+    "lm_head.weight": "pytorch_model-00002-of-00002.bin",
+    "model.embed_tokens.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.21.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.21.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.21.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.21.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.21.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.21.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.21.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.21.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.21.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.22.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.22.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.22.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.22.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.22.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.22.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.22.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.22.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.22.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.23.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.23.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.23.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.23.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.23.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.23.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.23.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.23.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.23.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.24.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.mlp.down_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.mlp.gate_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.mlp.up_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.self_attn.o_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.mlp.down_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.mlp.gate_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.mlp.up_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.self_attn.o_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.mlp.down_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.mlp.gate_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.mlp.up_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.self_attn.o_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.mlp.down_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.mlp.gate_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.mlp.up_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.self_attn.o_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.28.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.28.mlp.down_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.28.mlp.gate_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.28.mlp.up_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.28.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.28.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.28.self_attn.o_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.28.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.28.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.29.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.29.mlp.down_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.29.mlp.gate_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.29.mlp.up_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.29.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.29.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.29.self_attn.o_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.29.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.29.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.3.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.30.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.30.mlp.down_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.30.mlp.gate_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.30.mlp.up_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.30.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.30.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.30.self_attn.o_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.30.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.30.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.31.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.31.mlp.down_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.31.mlp.gate_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.31.mlp.up_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.31.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.31.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.31.self_attn.o_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.31.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.31.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.4.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.mm_projector.0.bias": "pytorch_model-00002-of-00002.bin",
+    "model.mm_projector.0.weight": "pytorch_model-00002-of-00002.bin",
+    "model.mm_projector.2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.mm_projector.2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.norm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vae_predictor_audio.0.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vae_predictor_audio.0.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vae_predictor_audio.2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vae_predictor_audio.2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vae_predictor_image.0.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vae_predictor_image.0.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vae_predictor_image.2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vae_predictor_image.2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vae_projector_audio.0.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vae_projector_audio.0.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vae_projector_audio.2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vae_projector_audio.2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vae_projector_image.0.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vae_projector_image.0.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vae_projector_image.2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vae_projector_image.2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.embeddings.class_embedding": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.embeddings.patch_embedding.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.embeddings.position_embedding.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.0.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.0.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.0.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.0.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.0.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.0.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.0.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.0.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.0.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.0.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.0.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.0.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.0.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.0.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.0.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.0.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.1.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.1.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.1.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.1.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.1.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.1.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.1.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.1.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.1.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.1.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.1.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.1.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.1.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.1.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.1.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.1.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.10.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.10.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.10.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.10.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.10.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.10.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.10.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.10.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.10.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.10.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.10.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.10.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.10.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.10.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.10.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.10.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.11.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.11.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.11.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.11.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.11.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.11.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.11.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.11.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.11.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.11.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.11.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.11.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.11.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.11.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.11.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.11.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.12.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.12.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.12.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.12.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.12.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.12.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.12.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.12.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.12.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.12.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.12.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.12.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.12.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.12.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.12.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.12.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.13.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.13.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.13.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.13.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.13.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.13.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.13.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.13.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.13.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.13.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.13.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.13.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.13.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.13.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.13.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.13.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.14.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.14.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.14.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.14.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.14.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.14.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.14.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.14.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.14.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.14.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.14.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.14.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.14.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.14.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.14.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.14.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.15.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.15.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.15.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.15.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.15.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.15.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.15.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.15.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.15.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.15.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.15.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.15.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.15.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.15.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.15.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.15.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.16.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.16.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.16.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.16.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.16.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.16.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.16.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.16.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.16.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.16.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.16.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.16.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.16.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.16.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.16.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.16.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.17.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.17.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.17.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.17.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.17.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.17.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.17.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.17.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.17.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.17.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.17.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.17.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.17.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.17.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.17.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.17.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.18.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.18.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.18.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.18.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.18.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.18.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.18.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.18.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.18.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.18.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.18.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.18.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.18.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.18.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.18.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.18.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.19.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.19.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.19.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.19.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.19.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.19.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.19.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.19.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.19.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.19.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.19.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.19.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.19.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.19.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.19.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.19.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.2.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.2.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.2.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.2.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.2.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.2.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.2.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.2.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.2.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.2.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.2.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.2.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.2.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.2.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.2.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.2.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.20.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.20.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.20.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.20.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.20.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.20.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.20.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.20.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.20.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.20.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.20.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.20.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.20.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.20.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.20.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.20.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.21.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.21.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.21.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.21.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.21.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.21.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.21.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.21.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.21.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.21.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.21.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.21.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.21.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.21.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.21.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.21.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.22.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.22.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.22.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.22.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.22.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.22.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.22.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.22.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.22.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.22.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.22.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.22.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.22.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.22.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.22.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.22.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.23.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.23.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.23.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.23.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.23.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.23.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.23.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.23.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.23.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.23.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.23.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.23.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.23.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.23.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.23.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.23.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.24.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.24.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.24.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.24.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.24.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.24.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.24.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.24.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.24.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.24.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.24.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.24.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.24.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.24.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.24.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.24.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.25.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.25.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.25.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.25.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.25.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.25.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.25.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.25.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.25.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.25.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.25.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.25.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.25.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.25.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.25.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.25.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.26.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.26.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.26.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.26.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.26.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.26.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.26.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.26.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.26.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.26.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.26.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.26.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.26.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.26.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.26.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.26.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.27.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.27.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.27.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.27.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.27.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.27.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.27.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.27.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.27.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.27.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.27.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.27.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.27.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.27.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.27.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.27.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.28.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.28.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.28.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.28.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.28.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.28.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.28.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.28.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.28.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.28.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.28.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.28.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.28.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.28.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.28.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.28.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.29.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.29.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.29.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.29.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.29.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.29.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.29.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.29.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.29.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.29.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.29.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.29.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.29.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.29.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.29.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.29.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.3.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.3.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.3.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.3.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.3.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.3.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.3.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.3.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.3.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.3.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.3.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.3.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.3.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.3.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.3.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.3.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.30.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.30.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.30.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.30.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.30.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.30.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.30.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.30.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.30.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.30.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.30.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.30.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.30.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.30.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.30.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.30.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.31.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.31.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.31.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.31.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.31.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.31.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.31.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.31.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.31.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.31.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.31.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.31.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.31.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.31.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.31.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.31.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.4.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.4.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.4.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.4.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.4.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.4.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.4.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.4.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.4.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.4.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.4.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.4.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.4.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.4.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.4.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.4.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.5.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.5.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.5.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.5.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.5.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.5.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.5.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.5.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.5.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.5.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.5.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.5.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.5.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.5.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.5.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.5.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.6.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.6.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.6.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.6.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.6.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.6.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.6.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.6.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.6.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.6.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.6.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.6.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.6.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.6.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.6.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.6.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.7.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.7.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.7.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.7.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.7.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.7.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.7.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.7.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.7.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.7.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.7.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.7.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.7.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.7.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.7.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.7.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.8.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.8.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.8.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.8.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.8.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.8.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.8.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.8.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.8.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.8.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.8.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.8.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.8.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.8.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.8.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.8.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.9.layer_norm1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.9.layer_norm1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.9.layer_norm2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.9.layer_norm2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.9.mlp.fc1.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.9.mlp.fc1.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.9.mlp.fc2.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.9.mlp.fc2.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.9.self_attn.k_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.9.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.9.self_attn.out_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.9.self_attn.out_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.9.self_attn.q_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.9.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.9.self_attn.v_proj.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.encoder.layers.9.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.post_layernorm.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.post_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.pre_layrnorm.bias": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.vision_model.pre_layrnorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.vision_tower.vision_tower.visual_projection.weight": "pytorch_model-00002-of-00002.bin"
+  }
+}
diff --git a/ckpts/llm/special_tokens_map.json b/ckpts/llm/special_tokens_map.json
new file mode 100644
index 0000000000000000000000000000000000000000..c4358cb4804ad12eca61102d370bcd556e4ab923
--- /dev/null
+++ b/ckpts/llm/special_tokens_map.json
@@ -0,0 +1,89 @@
+{
+  "additional_special_tokens": [
+    {
+      "content": "<im_gen_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<im_gen>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<mask_gen>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<audio_gen>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<audio_gen_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<audio>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<video>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<base>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<base_null>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    }
+  ],
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "<unk>",
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}
diff --git a/ckpts/llm/tokenizer.model b/ckpts/llm/tokenizer.model
new file mode 100644
index 0000000000000000000000000000000000000000..6c00c742ce03c627d6cd5b795984876fa49fa899
--- /dev/null
+++ b/ckpts/llm/tokenizer.model
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9e556afd44213b6bd1be2b850ebbbd98f5481437a8021afaf58ee7fb1818d347
+size 499723
diff --git a/ckpts/llm/tokenizer/added_tokens.json b/ckpts/llm/tokenizer/added_tokens.json
new file mode 100644
index 0000000000000000000000000000000000000000..33d7d87bf4ccf7c3f652eb834528c9376657c947
--- /dev/null
+++ b/ckpts/llm/tokenizer/added_tokens.json
@@ -0,0 +1,11 @@
+{
+  "<audio>": 32005,
+  "<audio_gen>": 32003,
+  "<audio_gen_start>": 32004,
+  "<base>": 32007,
+  "<base_null>": 32008,
+  "<im_gen>": 32001,
+  "<im_gen_start>": 32000,
+  "<mask_gen>": 32002,
+  "<video>": 32006
+}
diff --git a/ckpts/llm/tokenizer/special_tokens_map.json b/ckpts/llm/tokenizer/special_tokens_map.json
new file mode 100644
index 0000000000000000000000000000000000000000..c4358cb4804ad12eca61102d370bcd556e4ab923
--- /dev/null
+++ b/ckpts/llm/tokenizer/special_tokens_map.json
@@ -0,0 +1,89 @@
+{
+  "additional_special_tokens": [
+    {
+      "content": "<im_gen_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<im_gen>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<mask_gen>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<audio_gen>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<audio_gen_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<audio>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<video>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<base>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<base_null>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    }
+  ],
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "<unk>",
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}
diff --git a/ckpts/llm/tokenizer/tokenizer.model b/ckpts/llm/tokenizer/tokenizer.model
new file mode 100644
index 0000000000000000000000000000000000000000..6c00c742ce03c627d6cd5b795984876fa49fa899
--- /dev/null
+++ b/ckpts/llm/tokenizer/tokenizer.model
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9e556afd44213b6bd1be2b850ebbbd98f5481437a8021afaf58ee7fb1818d347
+size 499723
diff --git a/ckpts/llm/tokenizer/tokenizer_config.json b/ckpts/llm/tokenizer/tokenizer_config.json
new file mode 100644
index 0000000000000000000000000000000000000000..bb8e174539a4f62eeda7fb096379474db3f4dce4
--- /dev/null
+++ b/ckpts/llm/tokenizer/tokenizer_config.json
@@ -0,0 +1,125 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32000": {
+      "content": "<im_gen_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32001": {
+      "content": "<im_gen>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32002": {
+      "content": "<mask_gen>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32003": {
+      "content": "<audio_gen>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32004": {
+      "content": "<audio_gen_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32005": {
+      "content": "<audio>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32006": {
+      "content": "<video>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32007": {
+      "content": "<base>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32008": {
+      "content": "<base_null>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<im_gen_start>",
+    "<im_gen>",
+    "<mask_gen>",
+    "<audio_gen>",
+    "<audio_gen_start>",
+    "<audio>",
+    "<video>",
+    "<base>",
+    "<base_null>"
+  ],
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "legacy": false,
+  "model_max_length": 2048,
+  "pad_token": "<unk>",
+  "padding_side": "right",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": true
+}
diff --git a/ckpts/llm/tokenizer_config.json b/ckpts/llm/tokenizer_config.json
new file mode 100644
index 0000000000000000000000000000000000000000..bb8e174539a4f62eeda7fb096379474db3f4dce4
--- /dev/null
+++ b/ckpts/llm/tokenizer_config.json
@@ -0,0 +1,125 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32000": {
+      "content": "<im_gen_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32001": {
+      "content": "<im_gen>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32002": {
+      "content": "<mask_gen>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32003": {
+      "content": "<audio_gen>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32004": {
+      "content": "<audio_gen_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32005": {
+      "content": "<audio>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32006": {
+      "content": "<video>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32007": {
+      "content": "<base>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32008": {
+      "content": "<base_null>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<im_gen_start>",
+    "<im_gen>",
+    "<mask_gen>",
+    "<audio_gen>",
+    "<audio_gen_start>",
+    "<audio>",
+    "<video>",
+    "<base>",
+    "<base_null>"
+  ],
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "legacy": false,
+  "model_max_length": 2048,
+  "pad_token": "<unk>",
+  "padding_side": "right",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": true
+}
diff --git a/ckpts/llm/trainer_state.json b/ckpts/llm/trainer_state.json
new file mode 100644
index 0000000000000000000000000000000000000000..d2bb658a88e2628d05fcc114a89d4ce74b13708a
--- /dev/null
+++ b/ckpts/llm/trainer_state.json
@@ -0,0 +1,23182 @@
+{
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 1.0,
+  "eval_steps": 500,
+  "global_step": 3859,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.0,
+      "learning_rate": 1.7241379310344828e-07,
+      "loss": 3.6719,
+      "step": 1
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 3.4482758620689656e-07,
+      "loss": 3.5625,
+      "step": 2
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 5.172413793103449e-07,
+      "loss": 3.6426,
+      "step": 3
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 6.896551724137931e-07,
+      "loss": 3.6504,
+      "step": 4
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 8.620689655172415e-07,
+      "loss": 3.5977,
+      "step": 5
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 1.0344827586206898e-06,
+      "loss": 3.5508,
+      "step": 6
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 1.2068965517241381e-06,
+      "loss": 3.5293,
+      "step": 7
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 1.3793103448275862e-06,
+      "loss": 3.5215,
+      "step": 8
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 1.5517241379310346e-06,
+      "loss": 3.4219,
+      "step": 9
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 1.724137931034483e-06,
+      "loss": 3.168,
+      "step": 10
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 1.896551724137931e-06,
+      "loss": 3.2773,
+      "step": 11
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 2.0689655172413796e-06,
+      "loss": 2.8066,
+      "step": 12
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 2.241379310344828e-06,
+      "loss": 2.6738,
+      "step": 13
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 2.4137931034482762e-06,
+      "loss": 2.5586,
+      "step": 14
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 2.5862068965517246e-06,
+      "loss": 2.3398,
+      "step": 15
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 2.7586206896551725e-06,
+      "loss": 1.8877,
+      "step": 16
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 2.931034482758621e-06,
+      "loss": 1.8281,
+      "step": 17
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 3.103448275862069e-06,
+      "loss": 1.7666,
+      "step": 18
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 3.2758620689655175e-06,
+      "loss": 1.6055,
+      "step": 19
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 3.448275862068966e-06,
+      "loss": 1.5791,
+      "step": 20
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 3.620689655172414e-06,
+      "loss": 1.2471,
+      "step": 21
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 3.793103448275862e-06,
+      "loss": 1.2168,
+      "step": 22
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 3.96551724137931e-06,
+      "loss": 1.1494,
+      "step": 23
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 4.137931034482759e-06,
+      "loss": 1.0649,
+      "step": 24
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 4.310344827586207e-06,
+      "loss": 1.0879,
+      "step": 25
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 4.482758620689656e-06,
+      "loss": 1.0229,
+      "step": 26
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 4.655172413793104e-06,
+      "loss": 0.9854,
+      "step": 27
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 4.8275862068965525e-06,
+      "loss": 0.9907,
+      "step": 28
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 5e-06,
+      "loss": 1.0132,
+      "step": 29
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 5.172413793103449e-06,
+      "loss": 0.9487,
+      "step": 30
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 5.344827586206896e-06,
+      "loss": 0.9448,
+      "step": 31
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 5.517241379310345e-06,
+      "loss": 0.9409,
+      "step": 32
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 5.689655172413794e-06,
+      "loss": 0.8623,
+      "step": 33
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 5.862068965517242e-06,
+      "loss": 0.8936,
+      "step": 34
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 6.03448275862069e-06,
+      "loss": 0.9243,
+      "step": 35
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 6.206896551724138e-06,
+      "loss": 0.9146,
+      "step": 36
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 6.379310344827587e-06,
+      "loss": 0.8799,
+      "step": 37
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 6.551724137931035e-06,
+      "loss": 0.8101,
+      "step": 38
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 6.724137931034484e-06,
+      "loss": 0.8496,
+      "step": 39
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 6.896551724137932e-06,
+      "loss": 0.8496,
+      "step": 40
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 7.0689655172413796e-06,
+      "loss": 0.8193,
+      "step": 41
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 7.241379310344828e-06,
+      "loss": 0.8018,
+      "step": 42
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 7.413793103448277e-06,
+      "loss": 0.8716,
+      "step": 43
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 7.586206896551724e-06,
+      "loss": 0.7988,
+      "step": 44
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 7.758620689655173e-06,
+      "loss": 0.8701,
+      "step": 45
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 7.93103448275862e-06,
+      "loss": 0.8213,
+      "step": 46
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 8.103448275862069e-06,
+      "loss": 0.8076,
+      "step": 47
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 8.275862068965518e-06,
+      "loss": 0.8062,
+      "step": 48
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 8.448275862068966e-06,
+      "loss": 0.8286,
+      "step": 49
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 8.620689655172414e-06,
+      "loss": 0.7734,
+      "step": 50
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 8.793103448275862e-06,
+      "loss": 0.8203,
+      "step": 51
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 8.965517241379312e-06,
+      "loss": 0.77,
+      "step": 52
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 9.13793103448276e-06,
+      "loss": 0.8145,
+      "step": 53
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 9.310344827586207e-06,
+      "loss": 0.7617,
+      "step": 54
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 9.482758620689655e-06,
+      "loss": 0.7979,
+      "step": 55
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 9.655172413793105e-06,
+      "loss": 0.7676,
+      "step": 56
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 9.827586206896553e-06,
+      "loss": 0.7837,
+      "step": 57
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1e-05,
+      "loss": 0.8037,
+      "step": 58
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.0172413793103449e-05,
+      "loss": 0.7817,
+      "step": 59
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.0344827586206898e-05,
+      "loss": 0.7432,
+      "step": 60
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.0517241379310346e-05,
+      "loss": 0.7686,
+      "step": 61
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.0689655172413792e-05,
+      "loss": 0.7383,
+      "step": 62
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.0862068965517242e-05,
+      "loss": 0.7637,
+      "step": 63
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.103448275862069e-05,
+      "loss": 0.7476,
+      "step": 64
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.1206896551724138e-05,
+      "loss": 0.7256,
+      "step": 65
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.1379310344827587e-05,
+      "loss": 0.7495,
+      "step": 66
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.1551724137931035e-05,
+      "loss": 0.728,
+      "step": 67
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.1724137931034483e-05,
+      "loss": 0.7515,
+      "step": 68
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.1896551724137933e-05,
+      "loss": 0.7749,
+      "step": 69
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.206896551724138e-05,
+      "loss": 0.6997,
+      "step": 70
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.2241379310344827e-05,
+      "loss": 0.7412,
+      "step": 71
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.2413793103448277e-05,
+      "loss": 0.7153,
+      "step": 72
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.2586206896551725e-05,
+      "loss": 0.7549,
+      "step": 73
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.2758620689655174e-05,
+      "loss": 0.7407,
+      "step": 74
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.2931034482758622e-05,
+      "loss": 0.751,
+      "step": 75
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.310344827586207e-05,
+      "loss": 0.7158,
+      "step": 76
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.327586206896552e-05,
+      "loss": 0.7607,
+      "step": 77
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.3448275862068967e-05,
+      "loss": 0.6968,
+      "step": 78
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.3620689655172414e-05,
+      "loss": 0.7114,
+      "step": 79
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.3793103448275863e-05,
+      "loss": 0.7109,
+      "step": 80
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.3965517241379311e-05,
+      "loss": 0.7368,
+      "step": 81
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.4137931034482759e-05,
+      "loss": 0.7046,
+      "step": 82
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.4310344827586209e-05,
+      "loss": 0.7041,
+      "step": 83
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.4482758620689657e-05,
+      "loss": 0.7212,
+      "step": 84
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.4655172413793105e-05,
+      "loss": 0.6968,
+      "step": 85
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.4827586206896554e-05,
+      "loss": 0.707,
+      "step": 86
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.5000000000000002e-05,
+      "loss": 0.731,
+      "step": 87
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.5172413793103448e-05,
+      "loss": 0.7056,
+      "step": 88
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.5344827586206898e-05,
+      "loss": 0.7046,
+      "step": 89
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.5517241379310346e-05,
+      "loss": 0.7104,
+      "step": 90
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.5689655172413794e-05,
+      "loss": 0.7334,
+      "step": 91
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.586206896551724e-05,
+      "loss": 0.6533,
+      "step": 92
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.603448275862069e-05,
+      "loss": 0.7456,
+      "step": 93
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.6206896551724137e-05,
+      "loss": 0.6899,
+      "step": 94
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.637931034482759e-05,
+      "loss": 0.7217,
+      "step": 95
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 1.6551724137931037e-05,
+      "loss": 0.6895,
+      "step": 96
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.6724137931034485e-05,
+      "loss": 0.7407,
+      "step": 97
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.6896551724137932e-05,
+      "loss": 0.6841,
+      "step": 98
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.706896551724138e-05,
+      "loss": 0.7275,
+      "step": 99
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.7241379310344828e-05,
+      "loss": 0.7114,
+      "step": 100
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.7413793103448276e-05,
+      "loss": 0.7065,
+      "step": 101
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.7586206896551724e-05,
+      "loss": 0.6797,
+      "step": 102
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.7758620689655175e-05,
+      "loss": 0.7227,
+      "step": 103
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.7931034482758623e-05,
+      "loss": 0.6948,
+      "step": 104
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.810344827586207e-05,
+      "loss": 0.6743,
+      "step": 105
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.827586206896552e-05,
+      "loss": 0.689,
+      "step": 106
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.8448275862068967e-05,
+      "loss": 0.73,
+      "step": 107
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.8620689655172415e-05,
+      "loss": 0.73,
+      "step": 108
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.8793103448275863e-05,
+      "loss": 0.6963,
+      "step": 109
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.896551724137931e-05,
+      "loss": 0.6807,
+      "step": 110
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.913793103448276e-05,
+      "loss": 0.7041,
+      "step": 111
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.931034482758621e-05,
+      "loss": 0.6611,
+      "step": 112
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.9482758620689658e-05,
+      "loss": 0.6948,
+      "step": 113
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.9655172413793106e-05,
+      "loss": 0.7085,
+      "step": 114
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.9827586206896554e-05,
+      "loss": 0.6646,
+      "step": 115
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 2e-05,
+      "loss": 0.6431,
+      "step": 116
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.999999647766977e-05,
+      "loss": 0.6821,
+      "step": 117
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.9999985910681565e-05,
+      "loss": 0.6743,
+      "step": 118
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.9999968299042825e-05,
+      "loss": 0.6772,
+      "step": 119
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.9999943642765957e-05,
+      "loss": 0.6538,
+      "step": 120
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.9999911941868336e-05,
+      "loss": 0.6562,
+      "step": 121
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.9999873196372284e-05,
+      "loss": 0.6528,
+      "step": 122
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.9999827406305105e-05,
+      "loss": 0.6855,
+      "step": 123
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.9999774571699057e-05,
+      "loss": 0.6313,
+      "step": 124
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.9999714692591352e-05,
+      "loss": 0.6582,
+      "step": 125
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.999964776902418e-05,
+      "loss": 0.6875,
+      "step": 126
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.999957380104468e-05,
+      "loss": 0.686,
+      "step": 127
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.9999492788704966e-05,
+      "loss": 0.6851,
+      "step": 128
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.9999404732062107e-05,
+      "loss": 0.6777,
+      "step": 129
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.9999309631178135e-05,
+      "loss": 0.6704,
+      "step": 130
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.9999207486120048e-05,
+      "loss": 0.6816,
+      "step": 131
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.9999098296959798e-05,
+      "loss": 0.7021,
+      "step": 132
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.9998982063774307e-05,
+      "loss": 0.7251,
+      "step": 133
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.9998858786645462e-05,
+      "loss": 0.6465,
+      "step": 134
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.9998728465660105e-05,
+      "loss": 0.6748,
+      "step": 135
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.9998591100910038e-05,
+      "loss": 0.6875,
+      "step": 136
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.9998446692492037e-05,
+      "loss": 0.6904,
+      "step": 137
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.999829524050783e-05,
+      "loss": 0.6748,
+      "step": 138
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.9998136745064105e-05,
+      "loss": 0.6689,
+      "step": 139
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.9997971206272524e-05,
+      "loss": 0.6787,
+      "step": 140
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.99977986242497e-05,
+      "loss": 0.6797,
+      "step": 141
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.999761899911721e-05,
+      "loss": 0.7271,
+      "step": 142
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.9997432331001597e-05,
+      "loss": 0.6802,
+      "step": 143
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.999723862003436e-05,
+      "loss": 0.6392,
+      "step": 144
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.9997037866351966e-05,
+      "loss": 0.6616,
+      "step": 145
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.9996830070095836e-05,
+      "loss": 0.6509,
+      "step": 146
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.999661523141235e-05,
+      "loss": 0.6812,
+      "step": 147
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.999639335045286e-05,
+      "loss": 0.686,
+      "step": 148
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.999616442737368e-05,
+      "loss": 0.6631,
+      "step": 149
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.9995928462336064e-05,
+      "loss": 0.6807,
+      "step": 150
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.9995685455506253e-05,
+      "loss": 0.6533,
+      "step": 151
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.9995435407055433e-05,
+      "loss": 0.6763,
+      "step": 152
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.9995178317159755e-05,
+      "loss": 0.6787,
+      "step": 153
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.999491418600033e-05,
+      "loss": 0.6753,
+      "step": 154
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.999464301376323e-05,
+      "loss": 0.6802,
+      "step": 155
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.999436480063948e-05,
+      "loss": 0.6548,
+      "step": 156
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.9994079546825085e-05,
+      "loss": 0.6392,
+      "step": 157
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.9993787252520985e-05,
+      "loss": 0.6831,
+      "step": 158
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.99934879179331e-05,
+      "loss": 0.6948,
+      "step": 159
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.9993181543272294e-05,
+      "loss": 0.6543,
+      "step": 160
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.99928681287544e-05,
+      "loss": 0.6807,
+      "step": 161
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.999254767460021e-05,
+      "loss": 0.624,
+      "step": 162
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.999222018103547e-05,
+      "loss": 0.668,
+      "step": 163
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.9991885648290893e-05,
+      "loss": 0.6748,
+      "step": 164
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.999154407660214e-05,
+      "loss": 0.665,
+      "step": 165
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.9991195466209837e-05,
+      "loss": 0.6489,
+      "step": 166
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.9990839817359572e-05,
+      "loss": 0.6719,
+      "step": 167
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.9990477130301885e-05,
+      "loss": 0.6675,
+      "step": 168
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.9990107405292278e-05,
+      "loss": 0.666,
+      "step": 169
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.998973064259121e-05,
+      "loss": 0.6421,
+      "step": 170
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.9989346842464097e-05,
+      "loss": 0.6606,
+      "step": 171
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.998895600518131e-05,
+      "loss": 0.6509,
+      "step": 172
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.9988558131018188e-05,
+      "loss": 0.6548,
+      "step": 173
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.998815322025501e-05,
+      "loss": 0.6504,
+      "step": 174
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.9987741273177033e-05,
+      "loss": 0.6626,
+      "step": 175
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.998732229007445e-05,
+      "loss": 0.6421,
+      "step": 176
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.9986896271242426e-05,
+      "loss": 0.7095,
+      "step": 177
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.9986463216981074e-05,
+      "loss": 0.6528,
+      "step": 178
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.998602312759547e-05,
+      "loss": 0.6694,
+      "step": 179
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.9985576003395635e-05,
+      "loss": 0.6606,
+      "step": 180
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.998512184469656e-05,
+      "loss": 0.6421,
+      "step": 181
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.998466065181818e-05,
+      "loss": 0.6694,
+      "step": 182
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.9984192425085388e-05,
+      "loss": 0.6431,
+      "step": 183
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.9983717164828044e-05,
+      "loss": 0.6519,
+      "step": 184
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.998323487138094e-05,
+      "loss": 0.6592,
+      "step": 185
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.9982745545083842e-05,
+      "loss": 0.6484,
+      "step": 186
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.9982249186281465e-05,
+      "loss": 0.6729,
+      "step": 187
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.998174579532347e-05,
+      "loss": 0.665,
+      "step": 188
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.998123537256449e-05,
+      "loss": 0.647,
+      "step": 189
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.9980717918364092e-05,
+      "loss": 0.6392,
+      "step": 190
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.9980193433086804e-05,
+      "loss": 0.6606,
+      "step": 191
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.9979661917102116e-05,
+      "loss": 0.6665,
+      "step": 192
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.9979123370784454e-05,
+      "loss": 0.6416,
+      "step": 193
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.997857779451321e-05,
+      "loss": 0.6655,
+      "step": 194
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.9978025188672726e-05,
+      "loss": 0.6431,
+      "step": 195
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.997746555365229e-05,
+      "loss": 0.6567,
+      "step": 196
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.997689888984615e-05,
+      "loss": 0.6655,
+      "step": 197
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.9976325197653494e-05,
+      "loss": 0.6499,
+      "step": 198
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.9975744477478472e-05,
+      "loss": 0.6587,
+      "step": 199
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.9975156729730188e-05,
+      "loss": 0.646,
+      "step": 200
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.9974561954822684e-05,
+      "loss": 0.667,
+      "step": 201
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.997396015317496e-05,
+      "loss": 0.6509,
+      "step": 202
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.997335132521097e-05,
+      "loss": 0.6274,
+      "step": 203
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.99727354713596e-05,
+      "loss": 0.6455,
+      "step": 204
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.997211259205471e-05,
+      "loss": 0.6724,
+      "step": 205
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.9971482687735094e-05,
+      "loss": 0.6401,
+      "step": 206
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.9970845758844493e-05,
+      "loss": 0.6592,
+      "step": 207
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.997020180583161e-05,
+      "loss": 0.6406,
+      "step": 208
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.9969550829150082e-05,
+      "loss": 0.6396,
+      "step": 209
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.9968892829258502e-05,
+      "loss": 0.6348,
+      "step": 210
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.9968227806620413e-05,
+      "loss": 0.6167,
+      "step": 211
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.996755576170429e-05,
+      "loss": 0.6279,
+      "step": 212
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.996687669498358e-05,
+      "loss": 0.6396,
+      "step": 213
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.996619060693665e-05,
+      "loss": 0.6333,
+      "step": 214
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.996549749804683e-05,
+      "loss": 0.6172,
+      "step": 215
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.9964797368802395e-05,
+      "loss": 0.6636,
+      "step": 216
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.996409021969656e-05,
+      "loss": 0.6587,
+      "step": 217
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.9963376051227482e-05,
+      "loss": 0.6221,
+      "step": 218
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.996265486389828e-05,
+      "loss": 0.6426,
+      "step": 219
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.9961926658216995e-05,
+      "loss": 0.6558,
+      "step": 220
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.9961191434696627e-05,
+      "loss": 0.6377,
+      "step": 221
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.9960449193855118e-05,
+      "loss": 0.6602,
+      "step": 222
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.995969993621535e-05,
+      "loss": 0.6123,
+      "step": 223
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.9958943662305152e-05,
+      "loss": 0.6421,
+      "step": 224
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.995818037265729e-05,
+      "loss": 0.6445,
+      "step": 225
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.9957410067809477e-05,
+      "loss": 0.666,
+      "step": 226
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.9956632748304365e-05,
+      "loss": 0.6567,
+      "step": 227
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.995584841468955e-05,
+      "loss": 0.6265,
+      "step": 228
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.995505706751757e-05,
+      "loss": 0.6685,
+      "step": 229
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.9954258707345903e-05,
+      "loss": 0.6528,
+      "step": 230
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.9953453334736964e-05,
+      "loss": 0.6147,
+      "step": 231
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.9952640950258107e-05,
+      "loss": 0.6562,
+      "step": 232
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.9951821554481637e-05,
+      "loss": 0.623,
+      "step": 233
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.9950995147984786e-05,
+      "loss": 0.6313,
+      "step": 234
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.995016173134973e-05,
+      "loss": 0.6318,
+      "step": 235
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.994932130516359e-05,
+      "loss": 0.6274,
+      "step": 236
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.9948473870018403e-05,
+      "loss": 0.6177,
+      "step": 237
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.994761942651117e-05,
+      "loss": 0.6279,
+      "step": 238
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.9946757975243807e-05,
+      "loss": 0.6318,
+      "step": 239
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.9945889516823192e-05,
+      "loss": 0.645,
+      "step": 240
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.9945014051861115e-05,
+      "loss": 0.6172,
+      "step": 241
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.994413158097431e-05,
+      "loss": 0.6348,
+      "step": 242
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.994324210478445e-05,
+      "loss": 0.6255,
+      "step": 243
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.9942345623918144e-05,
+      "loss": 0.5991,
+      "step": 244
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.994144213900693e-05,
+      "loss": 0.6685,
+      "step": 245
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.994053165068728e-05,
+      "loss": 0.6333,
+      "step": 246
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.99396141596006e-05,
+      "loss": 0.6572,
+      "step": 247
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.993868966639324e-05,
+      "loss": 0.6548,
+      "step": 248
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.9937758171716468e-05,
+      "loss": 0.6411,
+      "step": 249
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.9936819676226494e-05,
+      "loss": 0.6143,
+      "step": 250
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9935874180584452e-05,
+      "loss": 0.6377,
+      "step": 251
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9934921685456414e-05,
+      "loss": 0.6294,
+      "step": 252
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9933962191513382e-05,
+      "loss": 0.6494,
+      "step": 253
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.993299569943128e-05,
+      "loss": 0.6274,
+      "step": 254
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.993202220989098e-05,
+      "loss": 0.6353,
+      "step": 255
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9931041723578263e-05,
+      "loss": 0.6216,
+      "step": 256
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.993005424118385e-05,
+      "loss": 0.6118,
+      "step": 257
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9929059763403392e-05,
+      "loss": 0.6094,
+      "step": 258
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9928058290937466e-05,
+      "loss": 0.6118,
+      "step": 259
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9927049824491566e-05,
+      "loss": 0.6255,
+      "step": 260
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9926034364776133e-05,
+      "loss": 0.6216,
+      "step": 261
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9925011912506517e-05,
+      "loss": 0.6187,
+      "step": 262
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9923982468403008e-05,
+      "loss": 0.6265,
+      "step": 263
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9922946033190808e-05,
+      "loss": 0.6177,
+      "step": 264
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9921902607600052e-05,
+      "loss": 0.6216,
+      "step": 265
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9920852192365797e-05,
+      "loss": 0.645,
+      "step": 266
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.991979478822803e-05,
+      "loss": 0.5933,
+      "step": 267
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9918730395931648e-05,
+      "loss": 0.6411,
+      "step": 268
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9917659016226487e-05,
+      "loss": 0.625,
+      "step": 269
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9916580649867293e-05,
+      "loss": 0.666,
+      "step": 270
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.991549529761374e-05,
+      "loss": 0.6406,
+      "step": 271
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9914402960230423e-05,
+      "loss": 0.6445,
+      "step": 272
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9913303638486855e-05,
+      "loss": 0.6045,
+      "step": 273
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.991219733315747e-05,
+      "loss": 0.5938,
+      "step": 274
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9911084045021623e-05,
+      "loss": 0.6104,
+      "step": 275
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.990996377486359e-05,
+      "loss": 0.6206,
+      "step": 276
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9908836523472562e-05,
+      "loss": 0.6147,
+      "step": 277
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.990770229164265e-05,
+      "loss": 0.6167,
+      "step": 278
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9906561080172878e-05,
+      "loss": 0.6084,
+      "step": 279
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9905412889867194e-05,
+      "loss": 0.6323,
+      "step": 280
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.990425772153446e-05,
+      "loss": 0.6191,
+      "step": 281
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.990309557598845e-05,
+      "loss": 0.6133,
+      "step": 282
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.990192645404786e-05,
+      "loss": 0.6431,
+      "step": 283
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.990075035653629e-05,
+      "loss": 0.6025,
+      "step": 284
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.989956728428227e-05,
+      "loss": 0.6147,
+      "step": 285
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9898377238119226e-05,
+      "loss": 0.6113,
+      "step": 286
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.989718021888551e-05,
+      "loss": 0.5996,
+      "step": 287
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9895976227424377e-05,
+      "loss": 0.6343,
+      "step": 288
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.9894765264584005e-05,
+      "loss": 0.6597,
+      "step": 289
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9893547331217466e-05,
+      "loss": 0.603,
+      "step": 290
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9892322428182764e-05,
+      "loss": 0.6055,
+      "step": 291
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9891090556342796e-05,
+      "loss": 0.623,
+      "step": 292
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9889851716565374e-05,
+      "loss": 0.627,
+      "step": 293
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9888605909723216e-05,
+      "loss": 0.6182,
+      "step": 294
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9887353136693957e-05,
+      "loss": 0.6006,
+      "step": 295
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.988609339836013e-05,
+      "loss": 0.6113,
+      "step": 296
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9884826695609173e-05,
+      "loss": 0.5903,
+      "step": 297
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.988355302933344e-05,
+      "loss": 0.6138,
+      "step": 298
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9882272400430192e-05,
+      "loss": 0.6099,
+      "step": 299
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9880984809801575e-05,
+      "loss": 0.6011,
+      "step": 300
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.987969025835466e-05,
+      "loss": 0.6255,
+      "step": 301
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.987838874700142e-05,
+      "loss": 0.6172,
+      "step": 302
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9877080276658713e-05,
+      "loss": 0.6348,
+      "step": 303
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9875764848248323e-05,
+      "loss": 0.6299,
+      "step": 304
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9874442462696913e-05,
+      "loss": 0.647,
+      "step": 305
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9873113120936074e-05,
+      "loss": 0.6362,
+      "step": 306
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9871776823902273e-05,
+      "loss": 0.6128,
+      "step": 307
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9870433572536883e-05,
+      "loss": 0.5991,
+      "step": 308
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9869083367786188e-05,
+      "loss": 0.6201,
+      "step": 309
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9867726210601353e-05,
+      "loss": 0.5967,
+      "step": 310
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9866362101938453e-05,
+      "loss": 0.5933,
+      "step": 311
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9864991042758456e-05,
+      "loss": 0.6074,
+      "step": 312
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9863613034027224e-05,
+      "loss": 0.5952,
+      "step": 313
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9862228076715524e-05,
+      "loss": 0.5864,
+      "step": 314
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9860836171799007e-05,
+      "loss": 0.6128,
+      "step": 315
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.985943732025822e-05,
+      "loss": 0.6157,
+      "step": 316
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.985803152307861e-05,
+      "loss": 0.5771,
+      "step": 317
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9856618781250513e-05,
+      "loss": 0.6206,
+      "step": 318
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9855199095769157e-05,
+      "loss": 0.6313,
+      "step": 319
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9853772467634665e-05,
+      "loss": 0.6528,
+      "step": 320
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9852338897852044e-05,
+      "loss": 0.6038,
+      "step": 321
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9850898387431194e-05,
+      "loss": 0.583,
+      "step": 322
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9849450937386914e-05,
+      "loss": 0.6279,
+      "step": 323
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9847996548738876e-05,
+      "loss": 0.5903,
+      "step": 324
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9846535222511648e-05,
+      "loss": 0.5903,
+      "step": 325
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9845066959734687e-05,
+      "loss": 0.6011,
+      "step": 326
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9843591761442334e-05,
+      "loss": 0.5981,
+      "step": 327
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.9842109628673815e-05,
+      "loss": 0.626,
+      "step": 328
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.984062056247324e-05,
+      "loss": 0.6187,
+      "step": 329
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.9839124563889613e-05,
+      "loss": 0.623,
+      "step": 330
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.983762163397681e-05,
+      "loss": 0.6372,
+      "step": 331
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.9836111773793588e-05,
+      "loss": 0.6079,
+      "step": 332
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.98345949844036e-05,
+      "loss": 0.624,
+      "step": 333
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.983307126687537e-05,
+      "loss": 0.6016,
+      "step": 334
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.983154062228231e-05,
+      "loss": 0.5991,
+      "step": 335
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.98300030517027e-05,
+      "loss": 0.6445,
+      "step": 336
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.982845855621971e-05,
+      "loss": 0.6011,
+      "step": 337
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.9826907136921382e-05,
+      "loss": 0.6401,
+      "step": 338
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.982534879490064e-05,
+      "loss": 0.6362,
+      "step": 339
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.982378353125528e-05,
+      "loss": 0.5972,
+      "step": 340
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.982221134708799e-05,
+      "loss": 0.5752,
+      "step": 341
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.9820632243506304e-05,
+      "loss": 0.606,
+      "step": 342
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.981904622162265e-05,
+      "loss": 0.6143,
+      "step": 343
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.9817453282554334e-05,
+      "loss": 0.5938,
+      "step": 344
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.9815853427423524e-05,
+      "loss": 0.6074,
+      "step": 345
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.9814246657357263e-05,
+      "loss": 0.6265,
+      "step": 346
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.9812632973487463e-05,
+      "loss": 0.5986,
+      "step": 347
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.9811012376950915e-05,
+      "loss": 0.6089,
+      "step": 348
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.980938486888927e-05,
+      "loss": 0.6113,
+      "step": 349
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.9807750450449055e-05,
+      "loss": 0.5923,
+      "step": 350
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.9806109122781662e-05,
+      "loss": 0.6064,
+      "step": 351
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.980446088704335e-05,
+      "loss": 0.6182,
+      "step": 352
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.980280574439524e-05,
+      "loss": 0.5703,
+      "step": 353
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.9801143696003334e-05,
+      "loss": 0.6211,
+      "step": 354
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.9799474743038482e-05,
+      "loss": 0.5996,
+      "step": 355
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.9797798886676407e-05,
+      "loss": 0.627,
+      "step": 356
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.9796116128097693e-05,
+      "loss": 0.5879,
+      "step": 357
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.979442646848778e-05,
+      "loss": 0.5962,
+      "step": 358
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.9792729909036983e-05,
+      "loss": 0.6118,
+      "step": 359
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.9791026450940468e-05,
+      "loss": 0.603,
+      "step": 360
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.9789316095398266e-05,
+      "loss": 0.5981,
+      "step": 361
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.9787598843615262e-05,
+      "loss": 0.6094,
+      "step": 362
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.97858746968012e-05,
+      "loss": 0.6284,
+      "step": 363
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.9784143656170688e-05,
+      "loss": 0.5781,
+      "step": 364
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.9782405722943182e-05,
+      "loss": 0.627,
+      "step": 365
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.9780660898342994e-05,
+      "loss": 0.5981,
+      "step": 366
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.97789091835993e-05,
+      "loss": 0.6099,
+      "step": 367
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.977715057994612e-05,
+      "loss": 0.5864,
+      "step": 368
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.977538508862233e-05,
+      "loss": 0.6104,
+      "step": 369
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.977361271087166e-05,
+      "loss": 0.6289,
+      "step": 370
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.9771833447942684e-05,
+      "loss": 0.5718,
+      "step": 371
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.9770047301088846e-05,
+      "loss": 0.5796,
+      "step": 372
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.9768254271568415e-05,
+      "loss": 0.5972,
+      "step": 373
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.9766454360644517e-05,
+      "loss": 0.5835,
+      "step": 374
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.976464756958514e-05,
+      "loss": 0.5962,
+      "step": 375
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.97628338996631e-05,
+      "loss": 0.5767,
+      "step": 376
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.9761013352156064e-05,
+      "loss": 0.6196,
+      "step": 377
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.9759185928346546e-05,
+      "loss": 0.5898,
+      "step": 378
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.9757351629521908e-05,
+      "loss": 0.6064,
+      "step": 379
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.975551045697435e-05,
+      "loss": 0.6128,
+      "step": 380
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.9753662412000915e-05,
+      "loss": 0.5918,
+      "step": 381
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.9751807495903484e-05,
+      "loss": 0.5854,
+      "step": 382
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.974994570998879e-05,
+      "loss": 0.6006,
+      "step": 383
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.9748077055568392e-05,
+      "loss": 0.5854,
+      "step": 384
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.9746201533958698e-05,
+      "loss": 0.5864,
+      "step": 385
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.9744319146480947e-05,
+      "loss": 0.6021,
+      "step": 386
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.9742429894461214e-05,
+      "loss": 0.5605,
+      "step": 387
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.974053377923042e-05,
+      "loss": 0.5762,
+      "step": 388
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.9738630802124307e-05,
+      "loss": 0.5688,
+      "step": 389
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.973672096448346e-05,
+      "loss": 0.5942,
+      "step": 390
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.9734804267653294e-05,
+      "loss": 0.5859,
+      "step": 391
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.973288071298406e-05,
+      "loss": 0.5977,
+      "step": 392
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.9730950301830837e-05,
+      "loss": 0.5859,
+      "step": 393
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.972901303555353e-05,
+      "loss": 0.6147,
+      "step": 394
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.972706891551688e-05,
+      "loss": 0.6338,
+      "step": 395
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.9725117943090453e-05,
+      "loss": 0.5996,
+      "step": 396
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.9723160119648644e-05,
+      "loss": 0.5796,
+      "step": 397
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.972119544657067e-05,
+      "loss": 0.5811,
+      "step": 398
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.971922392524058e-05,
+      "loss": 0.5879,
+      "step": 399
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.971724555704724e-05,
+      "loss": 0.6067,
+      "step": 400
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.971526034338435e-05,
+      "loss": 0.5825,
+      "step": 401
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.9713268285650415e-05,
+      "loss": 0.6055,
+      "step": 402
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.9711269385248783e-05,
+      "loss": 0.5874,
+      "step": 403
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.9709263643587607e-05,
+      "loss": 0.584,
+      "step": 404
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.970725106207986e-05,
+      "loss": 0.585,
+      "step": 405
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9705231642143342e-05,
+      "loss": 0.5752,
+      "step": 406
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9703205385200666e-05,
+      "loss": 0.603,
+      "step": 407
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.970117229267926e-05,
+      "loss": 0.605,
+      "step": 408
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9699132366011366e-05,
+      "loss": 0.6152,
+      "step": 409
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9697085606634047e-05,
+      "loss": 0.6099,
+      "step": 410
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9695032015989174e-05,
+      "loss": 0.6191,
+      "step": 411
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.969297159552343e-05,
+      "loss": 0.5879,
+      "step": 412
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9690904346688318e-05,
+      "loss": 0.5879,
+      "step": 413
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9688830270940136e-05,
+      "loss": 0.5991,
+      "step": 414
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9686749369740005e-05,
+      "loss": 0.585,
+      "step": 415
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.968466164455385e-05,
+      "loss": 0.5781,
+      "step": 416
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.96825670968524e-05,
+      "loss": 0.5842,
+      "step": 417
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.968046572811119e-05,
+      "loss": 0.5967,
+      "step": 418
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9678357539810567e-05,
+      "loss": 0.606,
+      "step": 419
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.967624253343568e-05,
+      "loss": 0.5967,
+      "step": 420
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9674120710476474e-05,
+      "loss": 0.5962,
+      "step": 421
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9671992072427702e-05,
+      "loss": 0.5957,
+      "step": 422
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.966985662078892e-05,
+      "loss": 0.5728,
+      "step": 423
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9667714357064486e-05,
+      "loss": 0.5986,
+      "step": 424
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9665565282763543e-05,
+      "loss": 0.5894,
+      "step": 425
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.966340939940004e-05,
+      "loss": 0.5786,
+      "step": 426
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.966124670849273e-05,
+      "loss": 0.5957,
+      "step": 427
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.965907721156515e-05,
+      "loss": 0.6069,
+      "step": 428
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9656900910145648e-05,
+      "loss": 0.6006,
+      "step": 429
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9654717805767342e-05,
+      "loss": 0.5737,
+      "step": 430
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9652527899968155e-05,
+      "loss": 0.563,
+      "step": 431
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.965033119429081e-05,
+      "loss": 0.5791,
+      "step": 432
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9648127690282805e-05,
+      "loss": 0.582,
+      "step": 433
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.964591738949644e-05,
+      "loss": 0.5786,
+      "step": 434
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9643700293488788e-05,
+      "loss": 0.5742,
+      "step": 435
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9641476403821724e-05,
+      "loss": 0.5793,
+      "step": 436
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9639245722061902e-05,
+      "loss": 0.6035,
+      "step": 437
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.963700824978076e-05,
+      "loss": 0.5811,
+      "step": 438
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.963476398855452e-05,
+      "loss": 0.5601,
+      "step": 439
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9632512939964193e-05,
+      "loss": 0.5684,
+      "step": 440
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9630255105595563e-05,
+      "loss": 0.5835,
+      "step": 441
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.96279904870392e-05,
+      "loss": 0.5962,
+      "step": 442
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9625719085890445e-05,
+      "loss": 0.5791,
+      "step": 443
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.962344090374943e-05,
+      "loss": 0.5747,
+      "step": 444
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9621155942221048e-05,
+      "loss": 0.5664,
+      "step": 445
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.961886420291499e-05,
+      "loss": 0.5923,
+      "step": 446
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9616565687445694e-05,
+      "loss": 0.5635,
+      "step": 447
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.96142603974324e-05,
+      "loss": 0.5559,
+      "step": 448
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9611948334499097e-05,
+      "loss": 0.5972,
+      "step": 449
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9609629500274554e-05,
+      "loss": 0.5703,
+      "step": 450
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9607303896392318e-05,
+      "loss": 0.5791,
+      "step": 451
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9604971524490692e-05,
+      "loss": 0.5854,
+      "step": 452
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9602632386212758e-05,
+      "loss": 0.5845,
+      "step": 453
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9600286483206358e-05,
+      "loss": 0.5898,
+      "step": 454
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9597933817124096e-05,
+      "loss": 0.6123,
+      "step": 455
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.959557438962335e-05,
+      "loss": 0.5874,
+      "step": 456
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9593208202366257e-05,
+      "loss": 0.6094,
+      "step": 457
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9590835257019715e-05,
+      "loss": 0.5857,
+      "step": 458
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.958845555525538e-05,
+      "loss": 0.5884,
+      "step": 459
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9586069098749677e-05,
+      "loss": 0.5649,
+      "step": 460
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.958367588918378e-05,
+      "loss": 0.5669,
+      "step": 461
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9581275928243624e-05,
+      "loss": 0.5806,
+      "step": 462
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.95788692176199e-05,
+      "loss": 0.5811,
+      "step": 463
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9576455759008052e-05,
+      "loss": 0.5823,
+      "step": 464
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9574035554108285e-05,
+      "loss": 0.5737,
+      "step": 465
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9571608604625544e-05,
+      "loss": 0.5801,
+      "step": 466
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9569174912269542e-05,
+      "loss": 0.5723,
+      "step": 467
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9566734478754724e-05,
+      "loss": 0.5947,
+      "step": 468
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9564287305800293e-05,
+      "loss": 0.5718,
+      "step": 469
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9561833395130207e-05,
+      "loss": 0.5688,
+      "step": 470
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9559372748473157e-05,
+      "loss": 0.5732,
+      "step": 471
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9556905367562584e-05,
+      "loss": 0.6011,
+      "step": 472
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9554431254136676e-05,
+      "loss": 0.5723,
+      "step": 473
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9551950409938364e-05,
+      "loss": 0.5649,
+      "step": 474
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9549462836715313e-05,
+      "loss": 0.5874,
+      "step": 475
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.954696853621994e-05,
+      "loss": 0.5903,
+      "step": 476
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.954446751020939e-05,
+      "loss": 0.574,
+      "step": 477
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.954195976044555e-05,
+      "loss": 0.5645,
+      "step": 478
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.953944528869505e-05,
+      "loss": 0.562,
+      "step": 479
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9536924096729253e-05,
+      "loss": 0.562,
+      "step": 480
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.953439618632424e-05,
+      "loss": 0.5977,
+      "step": 481
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.953186155926085e-05,
+      "loss": 0.5757,
+      "step": 482
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9529320217324637e-05,
+      "loss": 0.5801,
+      "step": 483
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9526772162305893e-05,
+      "loss": 0.5659,
+      "step": 484
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9524217395999627e-05,
+      "loss": 0.5381,
+      "step": 485
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9521655920205596e-05,
+      "loss": 0.5591,
+      "step": 486
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9519087736728267e-05,
+      "loss": 0.5903,
+      "step": 487
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9516512847376844e-05,
+      "loss": 0.5952,
+      "step": 488
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.951393125396524e-05,
+      "loss": 0.5806,
+      "step": 489
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9511342958312106e-05,
+      "loss": 0.563,
+      "step": 490
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.950874796224081e-05,
+      "loss": 0.5767,
+      "step": 491
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.950614626757943e-05,
+      "loss": 0.5806,
+      "step": 492
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9503537876160782e-05,
+      "loss": 0.5586,
+      "step": 493
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9500922789822382e-05,
+      "loss": 0.5869,
+      "step": 494
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9498301010406473e-05,
+      "loss": 0.5742,
+      "step": 495
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.949567253976001e-05,
+      "loss": 0.5603,
+      "step": 496
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9493037379734658e-05,
+      "loss": 0.5527,
+      "step": 497
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9490395532186798e-05,
+      "loss": 0.5542,
+      "step": 498
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9487746998977525e-05,
+      "loss": 0.561,
+      "step": 499
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9485091781972644e-05,
+      "loss": 0.5669,
+      "step": 500
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9482429883042653e-05,
+      "loss": 0.5898,
+      "step": 501
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.947976130406278e-05,
+      "loss": 0.5781,
+      "step": 502
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9477086046912945e-05,
+      "loss": 0.5903,
+      "step": 503
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9474404113477775e-05,
+      "loss": 0.5698,
+      "step": 504
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9471715505646607e-05,
+      "loss": 0.562,
+      "step": 505
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9469020225313465e-05,
+      "loss": 0.5603,
+      "step": 506
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.946631827437709e-05,
+      "loss": 0.585,
+      "step": 507
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9463609654740906e-05,
+      "loss": 0.5659,
+      "step": 508
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9460894368313052e-05,
+      "loss": 0.563,
+      "step": 509
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9458172417006347e-05,
+      "loss": 0.5845,
+      "step": 510
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9455443802738323e-05,
+      "loss": 0.5898,
+      "step": 511
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.945270852743119e-05,
+      "loss": 0.5657,
+      "step": 512
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9449966593011858e-05,
+      "loss": 0.562,
+      "step": 513
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9447218001411924e-05,
+      "loss": 0.5757,
+      "step": 514
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9444462754567682e-05,
+      "loss": 0.5322,
+      "step": 515
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9441700854420112e-05,
+      "loss": 0.5747,
+      "step": 516
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9438932302914868e-05,
+      "loss": 0.564,
+      "step": 517
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.943615710200231e-05,
+      "loss": 0.5942,
+      "step": 518
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.943337525363747e-05,
+      "loss": 0.5591,
+      "step": 519
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9430586759780065e-05,
+      "loss": 0.5933,
+      "step": 520
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9427791622394496e-05,
+      "loss": 0.5654,
+      "step": 521
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.942498984344984e-05,
+      "loss": 0.5854,
+      "step": 522
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9422181424919858e-05,
+      "loss": 0.5981,
+      "step": 523
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9419366368782983e-05,
+      "loss": 0.5566,
+      "step": 524
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.941654467702233e-05,
+      "loss": 0.5918,
+      "step": 525
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9413716351625678e-05,
+      "loss": 0.562,
+      "step": 526
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9410881394585493e-05,
+      "loss": 0.5544,
+      "step": 527
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.94080398078989e-05,
+      "loss": 0.5967,
+      "step": 528
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.940519159356771e-05,
+      "loss": 0.5479,
+      "step": 529
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9402336753598387e-05,
+      "loss": 0.5825,
+      "step": 530
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9399475290002066e-05,
+      "loss": 0.5571,
+      "step": 531
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9396607204794554e-05,
+      "loss": 0.6006,
+      "step": 532
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.939373249999632e-05,
+      "loss": 0.5464,
+      "step": 533
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9390851177632496e-05,
+      "loss": 0.5989,
+      "step": 534
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9387963239732876e-05,
+      "loss": 0.5771,
+      "step": 535
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9385068688331913e-05,
+      "loss": 0.5928,
+      "step": 536
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9382167525468722e-05,
+      "loss": 0.5825,
+      "step": 537
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.937925975318707e-05,
+      "loss": 0.5576,
+      "step": 538
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9376345373535386e-05,
+      "loss": 0.582,
+      "step": 539
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9373424388566758e-05,
+      "loss": 0.5498,
+      "step": 540
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.937049680033891e-05,
+      "loss": 0.5503,
+      "step": 541
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9367562610914233e-05,
+      "loss": 0.5825,
+      "step": 542
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9364621822359764e-05,
+      "loss": 0.5732,
+      "step": 543
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.936167443674719e-05,
+      "loss": 0.5684,
+      "step": 544
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.935872045615284e-05,
+      "loss": 0.5566,
+      "step": 545
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.93557598826577e-05,
+      "loss": 0.5386,
+      "step": 546
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9352792718347386e-05,
+      "loss": 0.5811,
+      "step": 547
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9349818965312165e-05,
+      "loss": 0.5918,
+      "step": 548
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.934683862564695e-05,
+      "loss": 0.5378,
+      "step": 549
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.934385170145129e-05,
+      "loss": 0.5713,
+      "step": 550
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9340858194829367e-05,
+      "loss": 0.5723,
+      "step": 551
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9337858107890003e-05,
+      "loss": 0.5669,
+      "step": 552
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9334851442746665e-05,
+      "loss": 0.5684,
+      "step": 553
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.933183820151744e-05,
+      "loss": 0.5903,
+      "step": 554
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.932881838632506e-05,
+      "loss": 0.5928,
+      "step": 555
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9325791999296873e-05,
+      "loss": 0.5754,
+      "step": 556
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9322759042564875e-05,
+      "loss": 0.5518,
+      "step": 557
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9319719518265677e-05,
+      "loss": 0.5786,
+      "step": 558
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.931667342854052e-05,
+      "loss": 0.5991,
+      "step": 559
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9313620775535275e-05,
+      "loss": 0.5752,
+      "step": 560
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.931056156140042e-05,
+      "loss": 0.5811,
+      "step": 561
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9307495788291088e-05,
+      "loss": 0.5962,
+      "step": 562
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9304423458366998e-05,
+      "loss": 0.5535,
+      "step": 563
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9301344573792503e-05,
+      "loss": 0.5337,
+      "step": 564
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9298259136736575e-05,
+      "loss": 0.5718,
+      "step": 565
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.92951671493728e-05,
+      "loss": 0.55,
+      "step": 566
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9292068613879377e-05,
+      "loss": 0.5732,
+      "step": 567
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.928896353243912e-05,
+      "loss": 0.5898,
+      "step": 568
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9285851907239453e-05,
+      "loss": 0.5386,
+      "step": 569
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.928273374047241e-05,
+      "loss": 0.5532,
+      "step": 570
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9279609034334635e-05,
+      "loss": 0.5713,
+      "step": 571
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9276477791027374e-05,
+      "loss": 0.5693,
+      "step": 572
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9273340012756486e-05,
+      "loss": 0.5776,
+      "step": 573
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9270195701732424e-05,
+      "loss": 0.5771,
+      "step": 574
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9267044860170253e-05,
+      "loss": 0.582,
+      "step": 575
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9263887490289634e-05,
+      "loss": 0.5708,
+      "step": 576
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9260723594314826e-05,
+      "loss": 0.5481,
+      "step": 577
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.925755317447468e-05,
+      "loss": 0.6055,
+      "step": 578
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9254376233002657e-05,
+      "loss": 0.5713,
+      "step": 579
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9251192772136808e-05,
+      "loss": 0.5874,
+      "step": 580
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.924800279411976e-05,
+      "loss": 0.5967,
+      "step": 581
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9244806301198752e-05,
+      "loss": 0.605,
+      "step": 582
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9241603295625605e-05,
+      "loss": 0.5767,
+      "step": 583
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9238393779656727e-05,
+      "loss": 0.5684,
+      "step": 584
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9235177755553117e-05,
+      "loss": 0.5527,
+      "step": 585
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9231955225580346e-05,
+      "loss": 0.5701,
+      "step": 586
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9228726192008578e-05,
+      "loss": 0.5723,
+      "step": 587
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9225490657112565e-05,
+      "loss": 0.5442,
+      "step": 588
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9222248623171628e-05,
+      "loss": 0.5601,
+      "step": 589
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.921900009246967e-05,
+      "loss": 0.5371,
+      "step": 590
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9215745067295168e-05,
+      "loss": 0.563,
+      "step": 591
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9212483549941175e-05,
+      "loss": 0.5481,
+      "step": 592
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.920921554270533e-05,
+      "loss": 0.5215,
+      "step": 593
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9205941047889824e-05,
+      "loss": 0.5654,
+      "step": 594
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9202660067801425e-05,
+      "loss": 0.5605,
+      "step": 595
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.919937260475148e-05,
+      "loss": 0.543,
+      "step": 596
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9196078661055893e-05,
+      "loss": 0.561,
+      "step": 597
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9192778239035127e-05,
+      "loss": 0.5464,
+      "step": 598
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9189471341014232e-05,
+      "loss": 0.5562,
+      "step": 599
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9186157969322795e-05,
+      "loss": 0.5776,
+      "step": 600
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9182838126294977e-05,
+      "loss": 0.5583,
+      "step": 601
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.917951181426949e-05,
+      "loss": 0.541,
+      "step": 602
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9176179035589618e-05,
+      "loss": 0.5742,
+      "step": 603
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9172839792603183e-05,
+      "loss": 0.5503,
+      "step": 604
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9169494087662567e-05,
+      "loss": 0.5664,
+      "step": 605
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.916614192312471e-05,
+      "loss": 0.5757,
+      "step": 606
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9162783301351096e-05,
+      "loss": 0.5471,
+      "step": 607
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.915941822470776e-05,
+      "loss": 0.5532,
+      "step": 608
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9156046695565286e-05,
+      "loss": 0.5444,
+      "step": 609
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.91526687162988e-05,
+      "loss": 0.5684,
+      "step": 610
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.914928428928797e-05,
+      "loss": 0.5522,
+      "step": 611
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.914589341691702e-05,
+      "loss": 0.5405,
+      "step": 612
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9142496101574698e-05,
+      "loss": 0.5728,
+      "step": 613
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9139092345654294e-05,
+      "loss": 0.5854,
+      "step": 614
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9135682151553647e-05,
+      "loss": 0.5669,
+      "step": 615
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9132265521675114e-05,
+      "loss": 0.5967,
+      "step": 616
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9128842458425602e-05,
+      "loss": 0.5388,
+      "step": 617
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9125412964216538e-05,
+      "loss": 0.5364,
+      "step": 618
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9121977041463884e-05,
+      "loss": 0.5371,
+      "step": 619
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9118534692588137e-05,
+      "loss": 0.5728,
+      "step": 620
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.911508592001431e-05,
+      "loss": 0.541,
+      "step": 621
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9111630726171944e-05,
+      "loss": 0.5413,
+      "step": 622
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9108169113495113e-05,
+      "loss": 0.5444,
+      "step": 623
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9104701084422398e-05,
+      "loss": 0.5581,
+      "step": 624
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9101226641396913e-05,
+      "loss": 0.5547,
+      "step": 625
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.909774578686628e-05,
+      "loss": 0.5515,
+      "step": 626
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9094258523282653e-05,
+      "loss": 0.5435,
+      "step": 627
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.909076485310268e-05,
+      "loss": 0.542,
+      "step": 628
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9087264778787534e-05,
+      "loss": 0.5352,
+      "step": 629
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9083758302802905e-05,
+      "loss": 0.5898,
+      "step": 630
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.908024542761898e-05,
+      "loss": 0.5261,
+      "step": 631
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9076726155710464e-05,
+      "loss": 0.5557,
+      "step": 632
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9073200489556565e-05,
+      "loss": 0.5806,
+      "step": 633
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9069668431640994e-05,
+      "loss": 0.5522,
+      "step": 634
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9066129984451964e-05,
+      "loss": 0.5371,
+      "step": 635
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.906258515048219e-05,
+      "loss": 0.5471,
+      "step": 636
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.9059033932228895e-05,
+      "loss": 0.5581,
+      "step": 637
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.9055476332193784e-05,
+      "loss": 0.563,
+      "step": 638
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.9051912352883066e-05,
+      "loss": 0.5339,
+      "step": 639
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.9048341996807445e-05,
+      "loss": 0.5415,
+      "step": 640
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.9044765266482117e-05,
+      "loss": 0.5688,
+      "step": 641
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.9041182164426765e-05,
+      "loss": 0.582,
+      "step": 642
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.9037592693165564e-05,
+      "loss": 0.554,
+      "step": 643
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.903399685522717e-05,
+      "loss": 0.5737,
+      "step": 644
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.9030394653144738e-05,
+      "loss": 0.543,
+      "step": 645
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.9026786089455892e-05,
+      "loss": 0.5542,
+      "step": 646
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.902317116670274e-05,
+      "loss": 0.5264,
+      "step": 647
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.901954988743188e-05,
+      "loss": 0.5347,
+      "step": 648
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.901592225419437e-05,
+      "loss": 0.5247,
+      "step": 649
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.901228826954576e-05,
+      "loss": 0.5229,
+      "step": 650
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.900864793604607e-05,
+      "loss": 0.5388,
+      "step": 651
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.9005001256259787e-05,
+      "loss": 0.5757,
+      "step": 652
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.9001348232755877e-05,
+      "loss": 0.5513,
+      "step": 653
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.8997688868107774e-05,
+      "loss": 0.5356,
+      "step": 654
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.899402316489337e-05,
+      "loss": 0.54,
+      "step": 655
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.8990351125695027e-05,
+      "loss": 0.5571,
+      "step": 656
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.8986672753099574e-05,
+      "loss": 0.5815,
+      "step": 657
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.8982988049698304e-05,
+      "loss": 0.5361,
+      "step": 658
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.8979297018086967e-05,
+      "loss": 0.542,
+      "step": 659
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.8975599660865757e-05,
+      "loss": 0.5273,
+      "step": 660
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.897189598063935e-05,
+      "loss": 0.5449,
+      "step": 661
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.8968185980016854e-05,
+      "loss": 0.5474,
+      "step": 662
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.8964469661611844e-05,
+      "loss": 0.5459,
+      "step": 663
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.8960747028042335e-05,
+      "loss": 0.5513,
+      "step": 664
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.8957018081930804e-05,
+      "loss": 0.5708,
+      "step": 665
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.895328282590416e-05,
+      "loss": 0.5342,
+      "step": 666
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.8949541262593764e-05,
+      "loss": 0.5352,
+      "step": 667
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.8945793394635423e-05,
+      "loss": 0.5422,
+      "step": 668
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.8942039224669383e-05,
+      "loss": 0.5449,
+      "step": 669
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.8938278755340327e-05,
+      "loss": 0.5447,
+      "step": 670
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.893451198929738e-05,
+      "loss": 0.5283,
+      "step": 671
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.89307389291941e-05,
+      "loss": 0.5659,
+      "step": 672
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.892695957768848e-05,
+      "loss": 0.5369,
+      "step": 673
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.892317393744294e-05,
+      "loss": 0.5669,
+      "step": 674
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.8919382011124346e-05,
+      "loss": 0.563,
+      "step": 675
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8915583801403973e-05,
+      "loss": 0.5132,
+      "step": 676
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.891177931095753e-05,
+      "loss": 0.5444,
+      "step": 677
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.890796854246516e-05,
+      "loss": 0.5425,
+      "step": 678
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.890415149861141e-05,
+      "loss": 0.541,
+      "step": 679
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8900328182085262e-05,
+      "loss": 0.5183,
+      "step": 680
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.889649859558011e-05,
+      "loss": 0.5708,
+      "step": 681
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.889266274179377e-05,
+      "loss": 0.5527,
+      "step": 682
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8888820623428476e-05,
+      "loss": 0.5444,
+      "step": 683
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8884972243190863e-05,
+      "loss": 0.5491,
+      "step": 684
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8881117603791984e-05,
+      "loss": 0.5464,
+      "step": 685
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8877256707947308e-05,
+      "loss": 0.5371,
+      "step": 686
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8873389558376697e-05,
+      "loss": 0.5356,
+      "step": 687
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.886951615780443e-05,
+      "loss": 0.5347,
+      "step": 688
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8865636508959187e-05,
+      "loss": 0.5635,
+      "step": 689
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8861750614574048e-05,
+      "loss": 0.543,
+      "step": 690
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8857858477386495e-05,
+      "loss": 0.5479,
+      "step": 691
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8853960100138403e-05,
+      "loss": 0.5642,
+      "step": 692
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8850055485576055e-05,
+      "loss": 0.5391,
+      "step": 693
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8846144636450107e-05,
+      "loss": 0.5486,
+      "step": 694
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8842227555515626e-05,
+      "loss": 0.5459,
+      "step": 695
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8838304245532063e-05,
+      "loss": 0.543,
+      "step": 696
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8834374709263258e-05,
+      "loss": 0.5525,
+      "step": 697
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.883043894947743e-05,
+      "loss": 0.5132,
+      "step": 698
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8826496968947198e-05,
+      "loss": 0.5239,
+      "step": 699
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8822548770449544e-05,
+      "loss": 0.5532,
+      "step": 700
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.881859435676584e-05,
+      "loss": 0.5415,
+      "step": 701
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8814633730681843e-05,
+      "loss": 0.5381,
+      "step": 702
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8810666894987673e-05,
+      "loss": 0.5378,
+      "step": 703
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8806693852477832e-05,
+      "loss": 0.5508,
+      "step": 704
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.88027146059512e-05,
+      "loss": 0.5454,
+      "step": 705
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8798729158211016e-05,
+      "loss": 0.5435,
+      "step": 706
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.879473751206489e-05,
+      "loss": 0.5405,
+      "step": 707
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8790739670324808e-05,
+      "loss": 0.5566,
+      "step": 708
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8786735635807108e-05,
+      "loss": 0.5413,
+      "step": 709
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.87827254113325e-05,
+      "loss": 0.5432,
+      "step": 710
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8778708999726047e-05,
+      "loss": 0.5728,
+      "step": 711
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.877468640381718e-05,
+      "loss": 0.5537,
+      "step": 712
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.8770657626439678e-05,
+      "loss": 0.5293,
+      "step": 713
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.8766622670431675e-05,
+      "loss": 0.5234,
+      "step": 714
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.8762581538635665e-05,
+      "loss": 0.5544,
+      "step": 715
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.875853423389849e-05,
+      "loss": 0.5615,
+      "step": 716
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.875448075907133e-05,
+      "loss": 0.5562,
+      "step": 717
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.8750421117009728e-05,
+      "loss": 0.5757,
+      "step": 718
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.8746355310573562e-05,
+      "loss": 0.5615,
+      "step": 719
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.8742283342627055e-05,
+      "loss": 0.521,
+      "step": 720
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.873820521603877e-05,
+      "loss": 0.5515,
+      "step": 721
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.8734120933681604e-05,
+      "loss": 0.5615,
+      "step": 722
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.8730030498432805e-05,
+      "loss": 0.5154,
+      "step": 723
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.872593391317394e-05,
+      "loss": 0.5405,
+      "step": 724
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.872183118079091e-05,
+      "loss": 0.5532,
+      "step": 725
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.871772230417396e-05,
+      "loss": 0.5376,
+      "step": 726
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.8713607286217646e-05,
+      "loss": 0.5322,
+      "step": 727
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.8709486129820863e-05,
+      "loss": 0.5063,
+      "step": 728
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.8705358837886825e-05,
+      "loss": 0.5337,
+      "step": 729
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.870122541332307e-05,
+      "loss": 0.5461,
+      "step": 730
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.8697085859041452e-05,
+      "loss": 0.5159,
+      "step": 731
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.8692940177958147e-05,
+      "loss": 0.5347,
+      "step": 732
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.868878837299365e-05,
+      "loss": 0.5566,
+      "step": 733
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.868463044707276e-05,
+      "loss": 0.5435,
+      "step": 734
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.86804664031246e-05,
+      "loss": 0.5435,
+      "step": 735
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.8676296244082605e-05,
+      "loss": 0.5361,
+      "step": 736
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.867211997288449e-05,
+      "loss": 0.5466,
+      "step": 737
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.8667937592472315e-05,
+      "loss": 0.5312,
+      "step": 738
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.866374910579242e-05,
+      "loss": 0.5229,
+      "step": 739
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.8659554515795446e-05,
+      "loss": 0.5293,
+      "step": 740
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.865535382543634e-05,
+      "loss": 0.5244,
+      "step": 741
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.865114703767436e-05,
+      "loss": 0.54,
+      "step": 742
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.8646934155473025e-05,
+      "loss": 0.5352,
+      "step": 743
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.8642715181800177e-05,
+      "loss": 0.5508,
+      "step": 744
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.863849011962794e-05,
+      "loss": 0.5435,
+      "step": 745
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.8634258971932726e-05,
+      "loss": 0.5527,
+      "step": 746
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.8630021741695236e-05,
+      "loss": 0.5251,
+      "step": 747
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.8625778431900453e-05,
+      "loss": 0.5662,
+      "step": 748
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.8621529045537644e-05,
+      "loss": 0.4998,
+      "step": 749
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.8617273585600363e-05,
+      "loss": 0.5498,
+      "step": 750
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.8613012055086428e-05,
+      "loss": 0.5469,
+      "step": 751
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.860874445699795e-05,
+      "loss": 0.5371,
+      "step": 752
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.8604470794341304e-05,
+      "loss": 0.5923,
+      "step": 753
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.860019107012714e-05,
+      "loss": 0.5349,
+      "step": 754
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.859590528737038e-05,
+      "loss": 0.5435,
+      "step": 755
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.859161344909021e-05,
+      "loss": 0.5198,
+      "step": 756
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.8587315558310087e-05,
+      "loss": 0.5581,
+      "step": 757
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.858301161805773e-05,
+      "loss": 0.5415,
+      "step": 758
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.8578701631365113e-05,
+      "loss": 0.5493,
+      "step": 759
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.8574385601268477e-05,
+      "loss": 0.5198,
+      "step": 760
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.8570063530808323e-05,
+      "loss": 0.5317,
+      "step": 761
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.8565735423029406e-05,
+      "loss": 0.5459,
+      "step": 762
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.8561401280980718e-05,
+      "loss": 0.5151,
+      "step": 763
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.8557061107715525e-05,
+      "loss": 0.5181,
+      "step": 764
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.855271490629133e-05,
+      "loss": 0.5669,
+      "step": 765
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.854836267976988e-05,
+      "loss": 0.5403,
+      "step": 766
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.8544004431217177e-05,
+      "loss": 0.5381,
+      "step": 767
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.8539640163703456e-05,
+      "loss": 0.5347,
+      "step": 768
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.8535269880303196e-05,
+      "loss": 0.5391,
+      "step": 769
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.8530893584095113e-05,
+      "loss": 0.5356,
+      "step": 770
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.8526511278162157e-05,
+      "loss": 0.5354,
+      "step": 771
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.852212296559152e-05,
+      "loss": 0.5493,
+      "step": 772
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.851772864947461e-05,
+      "loss": 0.5557,
+      "step": 773
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.851332833290708e-05,
+      "loss": 0.5298,
+      "step": 774
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.85089220189888e-05,
+      "loss": 0.5493,
+      "step": 775
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.8504509710823876e-05,
+      "loss": 0.5181,
+      "step": 776
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.8500091411520618e-05,
+      "loss": 0.5249,
+      "step": 777
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.8495667124191577e-05,
+      "loss": 0.5425,
+      "step": 778
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.849123685195351e-05,
+      "loss": 0.5127,
+      "step": 779
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.848680059792739e-05,
+      "loss": 0.52,
+      "step": 780
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.8482358365238414e-05,
+      "loss": 0.543,
+      "step": 781
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.847791015701598e-05,
+      "loss": 0.5298,
+      "step": 782
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.8473455976393698e-05,
+      "loss": 0.5715,
+      "step": 783
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.846899582650939e-05,
+      "loss": 0.5234,
+      "step": 784
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.8464529710505082e-05,
+      "loss": 0.54,
+      "step": 785
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.8460057631526998e-05,
+      "loss": 0.5146,
+      "step": 786
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.8455579592725565e-05,
+      "loss": 0.5427,
+      "step": 787
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.845109559725541e-05,
+      "loss": 0.541,
+      "step": 788
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.8446605648275356e-05,
+      "loss": 0.5522,
+      "step": 789
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.844210974894842e-05,
+      "loss": 0.5552,
+      "step": 790
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.8437607902441814e-05,
+      "loss": 0.5542,
+      "step": 791
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.843310011192693e-05,
+      "loss": 0.5354,
+      "step": 792
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.8428586380579357e-05,
+      "loss": 0.5449,
+      "step": 793
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.842406671157886e-05,
+      "loss": 0.5605,
+      "step": 794
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.84195411081094e-05,
+      "loss": 0.5483,
+      "step": 795
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.8415009573359106e-05,
+      "loss": 0.5249,
+      "step": 796
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.841047211052029e-05,
+      "loss": 0.541,
+      "step": 797
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.8405928722789446e-05,
+      "loss": 0.54,
+      "step": 798
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.840137941336723e-05,
+      "loss": 0.5625,
+      "step": 799
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.839682418545848e-05,
+      "loss": 0.5396,
+      "step": 800
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.8392263042272195e-05,
+      "loss": 0.5225,
+      "step": 801
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.8387695987021554e-05,
+      "loss": 0.5474,
+      "step": 802
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.8383123022923883e-05,
+      "loss": 0.5349,
+      "step": 803
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.8378544153200685e-05,
+      "loss": 0.5405,
+      "step": 804
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.8373959381077618e-05,
+      "loss": 0.5388,
+      "step": 805
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.8369368709784495e-05,
+      "loss": 0.53,
+      "step": 806
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.8364772142555293e-05,
+      "loss": 0.5317,
+      "step": 807
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.8360169682628134e-05,
+      "loss": 0.5469,
+      "step": 808
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.8355561333245297e-05,
+      "loss": 0.4783,
+      "step": 809
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.83509470976532e-05,
+      "loss": 0.5491,
+      "step": 810
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.834632697910243e-05,
+      "loss": 0.5566,
+      "step": 811
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.834170098084769e-05,
+      "loss": 0.5359,
+      "step": 812
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.8337069106147848e-05,
+      "loss": 0.5537,
+      "step": 813
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.8332431358265896e-05,
+      "loss": 0.5261,
+      "step": 814
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.8327787740468974e-05,
+      "loss": 0.5303,
+      "step": 815
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.832313825602835e-05,
+      "loss": 0.5435,
+      "step": 816
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.831848290821943e-05,
+      "loss": 0.5334,
+      "step": 817
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.8313821700321747e-05,
+      "loss": 0.5312,
+      "step": 818
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.8309154635618967e-05,
+      "loss": 0.5669,
+      "step": 819
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.8304481717398877e-05,
+      "loss": 0.5337,
+      "step": 820
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.8299802948953385e-05,
+      "loss": 0.5298,
+      "step": 821
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.829511833357853e-05,
+      "loss": 0.5291,
+      "step": 822
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.829042787457446e-05,
+      "loss": 0.5259,
+      "step": 823
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.8285731575245446e-05,
+      "loss": 0.5322,
+      "step": 824
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.8281029438899875e-05,
+      "loss": 0.5161,
+      "step": 825
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.827632146885024e-05,
+      "loss": 0.5562,
+      "step": 826
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.8271607668413142e-05,
+      "loss": 0.5352,
+      "step": 827
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.82668880409093e-05,
+      "loss": 0.5198,
+      "step": 828
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.8262162589663527e-05,
+      "loss": 0.5298,
+      "step": 829
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8257431318004744e-05,
+      "loss": 0.5381,
+      "step": 830
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.825269422926597e-05,
+      "loss": 0.5249,
+      "step": 831
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8247951326784323e-05,
+      "loss": 0.5498,
+      "step": 832
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.824320261390102e-05,
+      "loss": 0.5327,
+      "step": 833
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8238448093961368e-05,
+      "loss": 0.5552,
+      "step": 834
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.823368777031476e-05,
+      "loss": 0.5305,
+      "step": 835
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.822892164631469e-05,
+      "loss": 0.4963,
+      "step": 836
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.822414972531872e-05,
+      "loss": 0.5303,
+      "step": 837
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8219372010688516e-05,
+      "loss": 0.5337,
+      "step": 838
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.821458850578981e-05,
+      "loss": 0.5388,
+      "step": 839
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8209799213992422e-05,
+      "loss": 0.5637,
+      "step": 840
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8205004138670244e-05,
+      "loss": 0.5349,
+      "step": 841
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8200203283201245e-05,
+      "loss": 0.5408,
+      "step": 842
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.819539665096746e-05,
+      "loss": 0.552,
+      "step": 843
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8190584245355004e-05,
+      "loss": 0.5474,
+      "step": 844
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8185766069754053e-05,
+      "loss": 0.5247,
+      "step": 845
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8180942127558844e-05,
+      "loss": 0.5278,
+      "step": 846
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8176112422167685e-05,
+      "loss": 0.5579,
+      "step": 847
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8171276956982934e-05,
+      "loss": 0.5195,
+      "step": 848
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.816643573541102e-05,
+      "loss": 0.519,
+      "step": 849
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8161588760862408e-05,
+      "loss": 0.5276,
+      "step": 850
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8156736036751635e-05,
+      "loss": 0.542,
+      "step": 851
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8151877566497284e-05,
+      "loss": 0.5088,
+      "step": 852
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8147013353521974e-05,
+      "loss": 0.5498,
+      "step": 853
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.814214340125238e-05,
+      "loss": 0.5195,
+      "step": 854
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.813726771311922e-05,
+      "loss": 0.532,
+      "step": 855
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8132386292557253e-05,
+      "loss": 0.5242,
+      "step": 856
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8127499143005266e-05,
+      "loss": 0.5188,
+      "step": 857
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8122606267906097e-05,
+      "loss": 0.5,
+      "step": 858
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8117707670706607e-05,
+      "loss": 0.5159,
+      "step": 859
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8112803354857696e-05,
+      "loss": 0.5176,
+      "step": 860
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.810789332381428e-05,
+      "loss": 0.5396,
+      "step": 861
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8102977581035315e-05,
+      "loss": 0.5225,
+      "step": 862
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8098056129983774e-05,
+      "loss": 0.5293,
+      "step": 863
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8093128974126654e-05,
+      "loss": 0.5193,
+      "step": 864
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8088196116934965e-05,
+      "loss": 0.5381,
+      "step": 865
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8083257561883736e-05,
+      "loss": 0.5342,
+      "step": 866
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8078313312452018e-05,
+      "loss": 0.4941,
+      "step": 867
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.8073363372122866e-05,
+      "loss": 0.5283,
+      "step": 868
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.806840774438334e-05,
+      "loss": 0.4858,
+      "step": 869
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.806344643272451e-05,
+      "loss": 0.4976,
+      "step": 870
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.8058479440641456e-05,
+      "loss": 0.5347,
+      "step": 871
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.805350677163326e-05,
+      "loss": 0.5527,
+      "step": 872
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.804852842920299e-05,
+      "loss": 0.5034,
+      "step": 873
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.8043544416857722e-05,
+      "loss": 0.509,
+      "step": 874
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.8038554738108523e-05,
+      "loss": 0.5371,
+      "step": 875
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.8033559396470455e-05,
+      "loss": 0.5283,
+      "step": 876
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.8028558395462563e-05,
+      "loss": 0.5508,
+      "step": 877
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.8023551738607886e-05,
+      "loss": 0.5203,
+      "step": 878
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.801853942943344e-05,
+      "loss": 0.5271,
+      "step": 879
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.8013521471470228e-05,
+      "loss": 0.5281,
+      "step": 880
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.800849786825323e-05,
+      "loss": 0.5132,
+      "step": 881
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.8003468623321406e-05,
+      "loss": 0.5325,
+      "step": 882
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.799843374021769e-05,
+      "loss": 0.5229,
+      "step": 883
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.7993393222488983e-05,
+      "loss": 0.4868,
+      "step": 884
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.7988347073686158e-05,
+      "loss": 0.5139,
+      "step": 885
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.798329529736406e-05,
+      "loss": 0.5454,
+      "step": 886
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.7978237897081486e-05,
+      "loss": 0.5439,
+      "step": 887
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.797317487640121e-05,
+      "loss": 0.5554,
+      "step": 888
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.7968106238889954e-05,
+      "loss": 0.53,
+      "step": 889
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.7963031988118406e-05,
+      "loss": 0.52,
+      "step": 890
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.79579521276612e-05,
+      "loss": 0.51,
+      "step": 891
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.7952866661096922e-05,
+      "loss": 0.543,
+      "step": 892
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.7947775592008118e-05,
+      "loss": 0.5054,
+      "step": 893
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.7942678923981266e-05,
+      "loss": 0.5154,
+      "step": 894
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.79375766606068e-05,
+      "loss": 0.5051,
+      "step": 895
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.7932468805479087e-05,
+      "loss": 0.5264,
+      "step": 896
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.7927355362196442e-05,
+      "loss": 0.5337,
+      "step": 897
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.792223633436111e-05,
+      "loss": 0.5322,
+      "step": 898
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.7917111725579277e-05,
+      "loss": 0.5325,
+      "step": 899
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.791198153946105e-05,
+      "loss": 0.5005,
+      "step": 900
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.790684577962047e-05,
+      "loss": 0.4971,
+      "step": 901
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.7901704449675513e-05,
+      "loss": 0.5042,
+      "step": 902
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.7896557553248063e-05,
+      "loss": 0.5359,
+      "step": 903
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.789140509396394e-05,
+      "loss": 0.5752,
+      "step": 904
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.788624707545287e-05,
+      "loss": 0.5576,
+      "step": 905
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.788108350134851e-05,
+      "loss": 0.5415,
+      "step": 906
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7875914375288414e-05,
+      "loss": 0.5625,
+      "step": 907
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7870739700914063e-05,
+      "loss": 0.5295,
+      "step": 908
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7865559481870835e-05,
+      "loss": 0.5376,
+      "step": 909
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7860373721808022e-05,
+      "loss": 0.5403,
+      "step": 910
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7855182424378816e-05,
+      "loss": 0.5339,
+      "step": 911
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7849985593240303e-05,
+      "loss": 0.5195,
+      "step": 912
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7844783232053484e-05,
+      "loss": 0.521,
+      "step": 913
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7839575344483237e-05,
+      "loss": 0.5337,
+      "step": 914
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7834361934198346e-05,
+      "loss": 0.4937,
+      "step": 915
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7829143004871482e-05,
+      "loss": 0.5254,
+      "step": 916
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7823918560179202e-05,
+      "loss": 0.5166,
+      "step": 917
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.781868860380195e-05,
+      "loss": 0.5234,
+      "step": 918
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.781345313942406e-05,
+      "loss": 0.5251,
+      "step": 919
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7808212170733724e-05,
+      "loss": 0.5312,
+      "step": 920
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7802965701423037e-05,
+      "loss": 0.4932,
+      "step": 921
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7797713735187955e-05,
+      "loss": 0.5386,
+      "step": 922
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7792456275728312e-05,
+      "loss": 0.5476,
+      "step": 923
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.778719332674781e-05,
+      "loss": 0.5244,
+      "step": 924
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7781924891954017e-05,
+      "loss": 0.5215,
+      "step": 925
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7776650975058364e-05,
+      "loss": 0.5444,
+      "step": 926
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7771371579776147e-05,
+      "loss": 0.531,
+      "step": 927
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7766086709826526e-05,
+      "loss": 0.5088,
+      "step": 928
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7760796368932505e-05,
+      "loss": 0.5139,
+      "step": 929
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7755500560820954e-05,
+      "loss": 0.5315,
+      "step": 930
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.775019928922259e-05,
+      "loss": 0.5103,
+      "step": 931
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7744892557871977e-05,
+      "loss": 0.5063,
+      "step": 932
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7739580370507533e-05,
+      "loss": 0.5063,
+      "step": 933
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.77342627308715e-05,
+      "loss": 0.5349,
+      "step": 934
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7728939642709992e-05,
+      "loss": 0.4895,
+      "step": 935
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.772361110977293e-05,
+      "loss": 0.5479,
+      "step": 936
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.771827713581409e-05,
+      "loss": 0.5168,
+      "step": 937
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7712937724591074e-05,
+      "loss": 0.4824,
+      "step": 938
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7707592879865318e-05,
+      "loss": 0.519,
+      "step": 939
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.770224260540208e-05,
+      "loss": 0.5095,
+      "step": 940
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7696886904970453e-05,
+      "loss": 0.5386,
+      "step": 941
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.769152578234334e-05,
+      "loss": 0.502,
+      "step": 942
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7686159241297475e-05,
+      "loss": 0.5461,
+      "step": 943
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7680787285613397e-05,
+      "loss": 0.5327,
+      "step": 944
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.7675409919075476e-05,
+      "loss": 0.5063,
+      "step": 945
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7670027145471877e-05,
+      "loss": 0.5188,
+      "step": 946
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7664638968594583e-05,
+      "loss": 0.5344,
+      "step": 947
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.765924539223938e-05,
+      "loss": 0.4885,
+      "step": 948
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7653846420205862e-05,
+      "loss": 0.5129,
+      "step": 949
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7648442056297418e-05,
+      "loss": 0.5479,
+      "step": 950
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7643032304321246e-05,
+      "loss": 0.4727,
+      "step": 951
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7637617168088327e-05,
+      "loss": 0.5269,
+      "step": 952
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.763219665141344e-05,
+      "loss": 0.5088,
+      "step": 953
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7626770758115155e-05,
+      "loss": 0.5342,
+      "step": 954
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7621339492015835e-05,
+      "loss": 0.4956,
+      "step": 955
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7615902856941617e-05,
+      "loss": 0.4988,
+      "step": 956
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7610460856722426e-05,
+      "loss": 0.5076,
+      "step": 957
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7605013495191972e-05,
+      "loss": 0.5186,
+      "step": 958
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.759956077618773e-05,
+      "loss": 0.4958,
+      "step": 959
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7594102703550957e-05,
+      "loss": 0.5112,
+      "step": 960
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.758863928112668e-05,
+      "loss": 0.4756,
+      "step": 961
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7583170512763693e-05,
+      "loss": 0.5332,
+      "step": 962
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7577696402314567e-05,
+      "loss": 0.5256,
+      "step": 963
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.757221695363561e-05,
+      "loss": 0.4966,
+      "step": 964
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7566732170586923e-05,
+      "loss": 0.512,
+      "step": 965
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7561242057032343e-05,
+      "loss": 0.5166,
+      "step": 966
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7555746616839466e-05,
+      "loss": 0.5171,
+      "step": 967
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.755024585387965e-05,
+      "loss": 0.5547,
+      "step": 968
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7544739772027988e-05,
+      "loss": 0.5103,
+      "step": 969
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7539228375163337e-05,
+      "loss": 0.4922,
+      "step": 970
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.753371166716828e-05,
+      "loss": 0.5217,
+      "step": 971
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7528189651929156e-05,
+      "loss": 0.4873,
+      "step": 972
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.752266233333603e-05,
+      "loss": 0.5088,
+      "step": 973
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.751712971528272e-05,
+      "loss": 0.5391,
+      "step": 974
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7511591801666763e-05,
+      "loss": 0.5288,
+      "step": 975
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.750604859638943e-05,
+      "loss": 0.5291,
+      "step": 976
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7500500103355722e-05,
+      "loss": 0.5293,
+      "step": 977
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7494946326474365e-05,
+      "loss": 0.5312,
+      "step": 978
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7489387269657804e-05,
+      "loss": 0.5232,
+      "step": 979
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7483822936822208e-05,
+      "loss": 0.5032,
+      "step": 980
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7478253331887458e-05,
+      "loss": 0.4978,
+      "step": 981
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7472678458777154e-05,
+      "loss": 0.53,
+      "step": 982
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7467098321418604e-05,
+      "loss": 0.5225,
+      "step": 983
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.7461512923742828e-05,
+      "loss": 0.5078,
+      "step": 984
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7455922269684543e-05,
+      "loss": 0.533,
+      "step": 985
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.745032636318218e-05,
+      "loss": 0.5156,
+      "step": 986
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7444725208177863e-05,
+      "loss": 0.4985,
+      "step": 987
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7439118808617418e-05,
+      "loss": 0.5027,
+      "step": 988
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7433507168450354e-05,
+      "loss": 0.5078,
+      "step": 989
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7427890291629895e-05,
+      "loss": 0.519,
+      "step": 990
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.742226818211293e-05,
+      "loss": 0.5054,
+      "step": 991
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7416640843860048e-05,
+      "loss": 0.5276,
+      "step": 992
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7411008280835514e-05,
+      "loss": 0.532,
+      "step": 993
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7405370497007287e-05,
+      "loss": 0.5432,
+      "step": 994
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7399727496346985e-05,
+      "loss": 0.5183,
+      "step": 995
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.739407928282991e-05,
+      "loss": 0.5127,
+      "step": 996
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.738842586043504e-05,
+      "loss": 0.5151,
+      "step": 997
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.738276723314502e-05,
+      "loss": 0.5254,
+      "step": 998
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7377103404946162e-05,
+      "loss": 0.5137,
+      "step": 999
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7371434379828436e-05,
+      "loss": 0.4844,
+      "step": 1000
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7365760161785478e-05,
+      "loss": 0.4822,
+      "step": 1001
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7360080754814585e-05,
+      "loss": 0.5337,
+      "step": 1002
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7354396162916708e-05,
+      "loss": 0.5229,
+      "step": 1003
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7348706390096445e-05,
+      "loss": 0.5098,
+      "step": 1004
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7343011440362048e-05,
+      "loss": 0.5464,
+      "step": 1005
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7337311317725417e-05,
+      "loss": 0.5212,
+      "step": 1006
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7331606026202092e-05,
+      "loss": 0.5017,
+      "step": 1007
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7325895569811264e-05,
+      "loss": 0.5151,
+      "step": 1008
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.732017995257575e-05,
+      "loss": 0.5117,
+      "step": 1009
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.731445917852201e-05,
+      "loss": 0.5264,
+      "step": 1010
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7308733251680136e-05,
+      "loss": 0.5269,
+      "step": 1011
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7303002176083844e-05,
+      "loss": 0.5042,
+      "step": 1012
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7297265955770486e-05,
+      "loss": 0.5247,
+      "step": 1013
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7291524594781038e-05,
+      "loss": 0.5173,
+      "step": 1014
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7285778097160086e-05,
+      "loss": 0.5247,
+      "step": 1015
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7280026466955852e-05,
+      "loss": 0.5046,
+      "step": 1016
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7274269708220155e-05,
+      "loss": 0.5486,
+      "step": 1017
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.726850782500844e-05,
+      "loss": 0.5142,
+      "step": 1018
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7262740821379758e-05,
+      "loss": 0.5332,
+      "step": 1019
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.725696870139677e-05,
+      "loss": 0.5317,
+      "step": 1020
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7251191469125732e-05,
+      "loss": 0.5125,
+      "step": 1021
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.7245409128636516e-05,
+      "loss": 0.4993,
+      "step": 1022
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7239621684002575e-05,
+      "loss": 0.4834,
+      "step": 1023
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.723382913930098e-05,
+      "loss": 0.5117,
+      "step": 1024
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7228031498612368e-05,
+      "loss": 0.5229,
+      "step": 1025
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.722222876602099e-05,
+      "loss": 0.5037,
+      "step": 1026
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7216420945614668e-05,
+      "loss": 0.498,
+      "step": 1027
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.721060804148482e-05,
+      "loss": 0.5029,
+      "step": 1028
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7204790057726435e-05,
+      "loss": 0.4976,
+      "step": 1029
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7198966998438087e-05,
+      "loss": 0.5305,
+      "step": 1030
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7193138867721924e-05,
+      "loss": 0.52,
+      "step": 1031
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7187305669683666e-05,
+      "loss": 0.5017,
+      "step": 1032
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7181467408432597e-05,
+      "loss": 0.5425,
+      "step": 1033
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.717562408808158e-05,
+      "loss": 0.5452,
+      "step": 1034
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7169775712747042e-05,
+      "loss": 0.561,
+      "step": 1035
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7163922286548953e-05,
+      "loss": 0.4902,
+      "step": 1036
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7158063813610856e-05,
+      "loss": 0.5073,
+      "step": 1037
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7152200298059854e-05,
+      "loss": 0.498,
+      "step": 1038
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7146331744026582e-05,
+      "loss": 0.564,
+      "step": 1039
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7140458155645246e-05,
+      "loss": 0.5151,
+      "step": 1040
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.713457953705359e-05,
+      "loss": 0.4956,
+      "step": 1041
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7128695892392894e-05,
+      "loss": 0.5017,
+      "step": 1042
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7122807225807997e-05,
+      "loss": 0.5581,
+      "step": 1043
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7116913541447252e-05,
+      "loss": 0.5024,
+      "step": 1044
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.711101484346257e-05,
+      "loss": 0.5161,
+      "step": 1045
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7105111136009375e-05,
+      "loss": 0.5166,
+      "step": 1046
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7099202423246632e-05,
+      "loss": 0.5083,
+      "step": 1047
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.709328870933683e-05,
+      "loss": 0.4929,
+      "step": 1048
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.708736999844598e-05,
+      "loss": 0.5283,
+      "step": 1049
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.708144629474361e-05,
+      "loss": 0.4932,
+      "step": 1050
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.707551760240277e-05,
+      "loss": 0.5002,
+      "step": 1051
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7069583925600024e-05,
+      "loss": 0.5227,
+      "step": 1052
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.706364526851544e-05,
+      "loss": 0.5117,
+      "step": 1053
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7057701635332607e-05,
+      "loss": 0.5154,
+      "step": 1054
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7051753030238612e-05,
+      "loss": 0.5449,
+      "step": 1055
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.704579945742404e-05,
+      "loss": 0.4846,
+      "step": 1056
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7039840921082985e-05,
+      "loss": 0.5098,
+      "step": 1057
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.703387742541303e-05,
+      "loss": 0.5044,
+      "step": 1058
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.702790897461526e-05,
+      "loss": 0.4895,
+      "step": 1059
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7021935572894247e-05,
+      "loss": 0.4929,
+      "step": 1060
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.7015957224458042e-05,
+      "loss": 0.522,
+      "step": 1061
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.7009973933518194e-05,
+      "loss": 0.4897,
+      "step": 1062
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.700398570428973e-05,
+      "loss": 0.4773,
+      "step": 1063
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6997992540991148e-05,
+      "loss": 0.5125,
+      "step": 1064
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.699199444784443e-05,
+      "loss": 0.509,
+      "step": 1065
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6985991429075038e-05,
+      "loss": 0.5222,
+      "step": 1066
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6979983488911883e-05,
+      "loss": 0.4976,
+      "step": 1067
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6973970631587358e-05,
+      "loss": 0.5085,
+      "step": 1068
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6967952861337317e-05,
+      "loss": 0.5015,
+      "step": 1069
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6961930182401077e-05,
+      "loss": 0.5205,
+      "step": 1070
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6955902599021416e-05,
+      "loss": 0.5251,
+      "step": 1071
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6949870115444546e-05,
+      "loss": 0.4817,
+      "step": 1072
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.694383273592016e-05,
+      "loss": 0.5225,
+      "step": 1073
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6937790464701386e-05,
+      "loss": 0.5254,
+      "step": 1074
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.693174330604479e-05,
+      "loss": 0.5132,
+      "step": 1075
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.69256912642104e-05,
+      "loss": 0.5024,
+      "step": 1076
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.691963434346167e-05,
+      "loss": 0.5137,
+      "step": 1077
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6913572548065496e-05,
+      "loss": 0.5203,
+      "step": 1078
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6907505882292205e-05,
+      "loss": 0.4844,
+      "step": 1079
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.690143435041556e-05,
+      "loss": 0.4895,
+      "step": 1080
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6895357956712746e-05,
+      "loss": 0.5256,
+      "step": 1081
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6889276705464376e-05,
+      "loss": 0.5137,
+      "step": 1082
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6883190600954488e-05,
+      "loss": 0.5032,
+      "step": 1083
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6877099647470533e-05,
+      "loss": 0.5007,
+      "step": 1084
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.687100384930338e-05,
+      "loss": 0.4819,
+      "step": 1085
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.686490321074732e-05,
+      "loss": 0.4998,
+      "step": 1086
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6858797736100034e-05,
+      "loss": 0.5093,
+      "step": 1087
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.685268742966263e-05,
+      "loss": 0.5188,
+      "step": 1088
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6846572295739608e-05,
+      "loss": 0.5171,
+      "step": 1089
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6840452338638874e-05,
+      "loss": 0.5264,
+      "step": 1090
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6834327562671728e-05,
+      "loss": 0.4958,
+      "step": 1091
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6828197972152867e-05,
+      "loss": 0.5107,
+      "step": 1092
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.682206357140038e-05,
+      "loss": 0.5105,
+      "step": 1093
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.681592436473575e-05,
+      "loss": 0.4871,
+      "step": 1094
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6809780356483826e-05,
+      "loss": 0.5144,
+      "step": 1095
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6803631550972866e-05,
+      "loss": 0.5159,
+      "step": 1096
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6797477952534486e-05,
+      "loss": 0.5347,
+      "step": 1097
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6791319565503693e-05,
+      "loss": 0.5078,
+      "step": 1098
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.6785156394218857e-05,
+      "loss": 0.4905,
+      "step": 1099
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6778988443021725e-05,
+      "loss": 0.5037,
+      "step": 1100
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6772815716257414e-05,
+      "loss": 0.5117,
+      "step": 1101
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.676663821827439e-05,
+      "loss": 0.5149,
+      "step": 1102
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.67604559534245e-05,
+      "loss": 0.5024,
+      "step": 1103
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6754268926062936e-05,
+      "loss": 0.5002,
+      "step": 1104
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6748077140548248e-05,
+      "loss": 0.5391,
+      "step": 1105
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.674188060124234e-05,
+      "loss": 0.5354,
+      "step": 1106
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6735679312510464e-05,
+      "loss": 0.521,
+      "step": 1107
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.672947327872121e-05,
+      "loss": 0.5173,
+      "step": 1108
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6723262504246536e-05,
+      "loss": 0.5339,
+      "step": 1109
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.67170469934617e-05,
+      "loss": 0.5142,
+      "step": 1110
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6710826750745332e-05,
+      "loss": 0.488,
+      "step": 1111
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6704601780479382e-05,
+      "loss": 0.4871,
+      "step": 1112
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.669837208704912e-05,
+      "loss": 0.4849,
+      "step": 1113
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6692137674843164e-05,
+      "loss": 0.5276,
+      "step": 1114
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6685898548253442e-05,
+      "loss": 0.5022,
+      "step": 1115
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6679654711675205e-05,
+      "loss": 0.5083,
+      "step": 1116
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6673406169507026e-05,
+      "loss": 0.4954,
+      "step": 1117
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.666715292615079e-05,
+      "loss": 0.4897,
+      "step": 1118
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6660894986011693e-05,
+      "loss": 0.4946,
+      "step": 1119
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6654632353498246e-05,
+      "loss": 0.5137,
+      "step": 1120
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6648365033022256e-05,
+      "loss": 0.51,
+      "step": 1121
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.664209302899884e-05,
+      "loss": 0.5002,
+      "step": 1122
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6635816345846413e-05,
+      "loss": 0.5137,
+      "step": 1123
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.662953498798668e-05,
+      "loss": 0.5266,
+      "step": 1124
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6623248959844653e-05,
+      "loss": 0.4937,
+      "step": 1125
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6616958265848618e-05,
+      "loss": 0.5037,
+      "step": 1126
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6610662910430154e-05,
+      "loss": 0.4961,
+      "step": 1127
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6604362898024137e-05,
+      "loss": 0.5078,
+      "step": 1128
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.65980582330687e-05,
+      "loss": 0.5007,
+      "step": 1129
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6591748920005266e-05,
+      "loss": 0.519,
+      "step": 1130
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6585434963278534e-05,
+      "loss": 0.521,
+      "step": 1131
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6579116367336477e-05,
+      "loss": 0.5007,
+      "step": 1132
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6572793136630327e-05,
+      "loss": 0.5186,
+      "step": 1133
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6566465275614586e-05,
+      "loss": 0.4951,
+      "step": 1134
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6560132788747013e-05,
+      "loss": 0.4954,
+      "step": 1135
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6553795680488635e-05,
+      "loss": 0.5073,
+      "step": 1136
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.6547453955303733e-05,
+      "loss": 0.4766,
+      "step": 1137
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.654110761765983e-05,
+      "loss": 0.5027,
+      "step": 1138
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6534756672027703e-05,
+      "loss": 0.5071,
+      "step": 1139
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.652840112288139e-05,
+      "loss": 0.5049,
+      "step": 1140
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6522040974698148e-05,
+      "loss": 0.4834,
+      "step": 1141
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6515676231958488e-05,
+      "loss": 0.5166,
+      "step": 1142
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.650930689914616e-05,
+      "loss": 0.4939,
+      "step": 1143
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6502932980748138e-05,
+      "loss": 0.5068,
+      "step": 1144
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6496554481254632e-05,
+      "loss": 0.5044,
+      "step": 1145
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.649017140515908e-05,
+      "loss": 0.5078,
+      "step": 1146
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6483783756958142e-05,
+      "loss": 0.4731,
+      "step": 1147
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6477391541151696e-05,
+      "loss": 0.4907,
+      "step": 1148
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.647099476224284e-05,
+      "loss": 0.51,
+      "step": 1149
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6464593424737894e-05,
+      "loss": 0.5022,
+      "step": 1150
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6458187533146385e-05,
+      "loss": 0.5039,
+      "step": 1151
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6451777091981037e-05,
+      "loss": 0.5005,
+      "step": 1152
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.644536210575779e-05,
+      "loss": 0.5142,
+      "step": 1153
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6438942578995786e-05,
+      "loss": 0.5132,
+      "step": 1154
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6432518516217367e-05,
+      "loss": 0.5371,
+      "step": 1155
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.642608992194806e-05,
+      "loss": 0.5129,
+      "step": 1156
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.64196568007166e-05,
+      "loss": 0.5,
+      "step": 1157
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6413219157054893e-05,
+      "loss": 0.4822,
+      "step": 1158
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6406776995498048e-05,
+      "loss": 0.5156,
+      "step": 1159
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6400330320584347e-05,
+      "loss": 0.5164,
+      "step": 1160
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6393879136855247e-05,
+      "loss": 0.4854,
+      "step": 1161
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.63874234488554e-05,
+      "loss": 0.5146,
+      "step": 1162
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.638096326113261e-05,
+      "loss": 0.4905,
+      "step": 1163
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.637449857823786e-05,
+      "loss": 0.4663,
+      "step": 1164
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6368029404725308e-05,
+      "loss": 0.4944,
+      "step": 1165
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6361555745152258e-05,
+      "loss": 0.4868,
+      "step": 1166
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6355077604079187e-05,
+      "loss": 0.5051,
+      "step": 1167
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.634859498606972e-05,
+      "loss": 0.4878,
+      "step": 1168
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6342107895690653e-05,
+      "loss": 0.5073,
+      "step": 1169
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.633561633751191e-05,
+      "loss": 0.4839,
+      "step": 1170
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6329120316106578e-05,
+      "loss": 0.481,
+      "step": 1171
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6322619836050884e-05,
+      "loss": 0.4946,
+      "step": 1172
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6316114901924195e-05,
+      "loss": 0.5337,
+      "step": 1173
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6309605518309016e-05,
+      "loss": 0.5137,
+      "step": 1174
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6303091689790986e-05,
+      "loss": 0.5076,
+      "step": 1175
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.6296573420958877e-05,
+      "loss": 0.5149,
+      "step": 1176
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.6290050716404588e-05,
+      "loss": 0.4897,
+      "step": 1177
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.628352358072314e-05,
+      "loss": 0.4858,
+      "step": 1178
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.627699201851268e-05,
+      "loss": 0.5317,
+      "step": 1179
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.6270456034374477e-05,
+      "loss": 0.4941,
+      "step": 1180
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.62639156329129e-05,
+      "loss": 0.501,
+      "step": 1181
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.625737081873545e-05,
+      "loss": 0.5068,
+      "step": 1182
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.625082159645272e-05,
+      "loss": 0.5044,
+      "step": 1183
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.624426797067842e-05,
+      "loss": 0.512,
+      "step": 1184
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.6237709946029347e-05,
+      "loss": 0.5161,
+      "step": 1185
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.6231147527125415e-05,
+      "loss": 0.4817,
+      "step": 1186
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.6224580718589626e-05,
+      "loss": 0.4966,
+      "step": 1187
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.6218009525048068e-05,
+      "loss": 0.5088,
+      "step": 1188
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.6211433951129928e-05,
+      "loss": 0.4841,
+      "step": 1189
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.6204854001467473e-05,
+      "loss": 0.4827,
+      "step": 1190
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.619826968069605e-05,
+      "loss": 0.5078,
+      "step": 1191
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.61916809934541e-05,
+      "loss": 0.5132,
+      "step": 1192
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.618508794438312e-05,
+      "loss": 0.479,
+      "step": 1193
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.6178490538127696e-05,
+      "loss": 0.5142,
+      "step": 1194
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.617188877933547e-05,
+      "loss": 0.4873,
+      "step": 1195
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.6165282672657158e-05,
+      "loss": 0.4746,
+      "step": 1196
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.615867222274654e-05,
+      "loss": 0.4988,
+      "step": 1197
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.6152057434260456e-05,
+      "loss": 0.4832,
+      "step": 1198
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.61454383118588e-05,
+      "loss": 0.4724,
+      "step": 1199
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.6138814860204513e-05,
+      "loss": 0.5337,
+      "step": 1200
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.6132187083963593e-05,
+      "loss": 0.5005,
+      "step": 1201
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.6125554987805084e-05,
+      "loss": 0.4897,
+      "step": 1202
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.611891857640108e-05,
+      "loss": 0.4973,
+      "step": 1203
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.6112277854426696e-05,
+      "loss": 0.5178,
+      "step": 1204
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.6105632826560108e-05,
+      "loss": 0.5327,
+      "step": 1205
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.6098983497482497e-05,
+      "loss": 0.519,
+      "step": 1206
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.6092329871878102e-05,
+      "loss": 0.4863,
+      "step": 1207
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.608567195443417e-05,
+      "loss": 0.5139,
+      "step": 1208
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.607900974984098e-05,
+      "loss": 0.4844,
+      "step": 1209
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.607234326279183e-05,
+      "loss": 0.4922,
+      "step": 1210
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.606567249798303e-05,
+      "loss": 0.4985,
+      "step": 1211
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.605899746011391e-05,
+      "loss": 0.5103,
+      "step": 1212
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.605231815388681e-05,
+      "loss": 0.5046,
+      "step": 1213
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.604563458400707e-05,
+      "loss": 0.5327,
+      "step": 1214
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.6038946755183038e-05,
+      "loss": 0.5156,
+      "step": 1215
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.6032254672126068e-05,
+      "loss": 0.4788,
+      "step": 1216
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.6025558339550497e-05,
+      "loss": 0.4878,
+      "step": 1217
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.6018857762173672e-05,
+      "loss": 0.5034,
+      "step": 1218
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.6012152944715915e-05,
+      "loss": 0.5198,
+      "step": 1219
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.600544389190055e-05,
+      "loss": 0.4834,
+      "step": 1220
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.5998730608453868e-05,
+      "loss": 0.5056,
+      "step": 1221
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.5992013099105156e-05,
+      "loss": 0.5071,
+      "step": 1222
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.598529136858667e-05,
+      "loss": 0.5068,
+      "step": 1223
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.5978565421633642e-05,
+      "loss": 0.5024,
+      "step": 1224
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.5971835262984266e-05,
+      "loss": 0.5212,
+      "step": 1225
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.596510089737972e-05,
+      "loss": 0.509,
+      "step": 1226
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.5958362329564127e-05,
+      "loss": 0.4924,
+      "step": 1227
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.5951619564284588e-05,
+      "loss": 0.5122,
+      "step": 1228
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.5944872606291148e-05,
+      "loss": 0.5034,
+      "step": 1229
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.593812146033681e-05,
+      "loss": 0.5022,
+      "step": 1230
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.5931366131177524e-05,
+      "loss": 0.4836,
+      "step": 1231
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.5924606623572194e-05,
+      "loss": 0.4597,
+      "step": 1232
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.5917842942282663e-05,
+      "loss": 0.5127,
+      "step": 1233
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.5911075092073717e-05,
+      "loss": 0.4995,
+      "step": 1234
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.5904303077713068e-05,
+      "loss": 0.458,
+      "step": 1235
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.589752690397138e-05,
+      "loss": 0.4641,
+      "step": 1236
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.589074657562223e-05,
+      "loss": 0.4927,
+      "step": 1237
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.5883962097442137e-05,
+      "loss": 0.5132,
+      "step": 1238
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.587717347421053e-05,
+      "loss": 0.4861,
+      "step": 1239
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.5870380710709762e-05,
+      "loss": 0.4958,
+      "step": 1240
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.5863583811725106e-05,
+      "loss": 0.5195,
+      "step": 1241
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.5856782782044752e-05,
+      "loss": 0.49,
+      "step": 1242
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.584997762645979e-05,
+      "loss": 0.491,
+      "step": 1243
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.5843168349764217e-05,
+      "loss": 0.4985,
+      "step": 1244
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.5836354956754946e-05,
+      "loss": 0.5234,
+      "step": 1245
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.582953745223177e-05,
+      "loss": 0.4937,
+      "step": 1246
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.58227158409974e-05,
+      "loss": 0.4893,
+      "step": 1247
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.5815890127857426e-05,
+      "loss": 0.499,
+      "step": 1248
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.5809060317620328e-05,
+      "loss": 0.4917,
+      "step": 1249
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.5802226415097478e-05,
+      "loss": 0.5137,
+      "step": 1250
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.5795388425103127e-05,
+      "loss": 0.4827,
+      "step": 1251
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.578854635245441e-05,
+      "loss": 0.491,
+      "step": 1252
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.5781700201971332e-05,
+      "loss": 0.4739,
+      "step": 1253
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.5774849978476774e-05,
+      "loss": 0.4783,
+      "step": 1254
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.5767995686796487e-05,
+      "loss": 0.4963,
+      "step": 1255
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.5761137331759084e-05,
+      "loss": 0.4851,
+      "step": 1256
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.5754274918196046e-05,
+      "loss": 0.4919,
+      "step": 1257
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.574740845094171e-05,
+      "loss": 0.491,
+      "step": 1258
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.574053793483327e-05,
+      "loss": 0.4846,
+      "step": 1259
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.5733663374710768e-05,
+      "loss": 0.4961,
+      "step": 1260
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.5726784775417097e-05,
+      "loss": 0.5352,
+      "step": 1261
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.5719902141798e-05,
+      "loss": 0.4912,
+      "step": 1262
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.5713015478702064e-05,
+      "loss": 0.5142,
+      "step": 1263
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.57061247909807e-05,
+      "loss": 0.5061,
+      "step": 1264
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.5699230083488167e-05,
+      "loss": 0.5095,
+      "step": 1265
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.5692331361081553e-05,
+      "loss": 0.4832,
+      "step": 1266
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.5685428628620773e-05,
+      "loss": 0.5144,
+      "step": 1267
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.5678521890968566e-05,
+      "loss": 0.5068,
+      "step": 1268
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.56716111529905e-05,
+      "loss": 0.4773,
+      "step": 1269
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.566469641955495e-05,
+      "loss": 0.4741,
+      "step": 1270
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.565777769553311e-05,
+      "loss": 0.4956,
+      "step": 1271
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.565085498579899e-05,
+      "loss": 0.4922,
+      "step": 1272
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.56439282952294e-05,
+      "loss": 0.5083,
+      "step": 1273
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.5636997628703966e-05,
+      "loss": 0.4758,
+      "step": 1274
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.56300629911051e-05,
+      "loss": 0.5164,
+      "step": 1275
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.5623124387318018e-05,
+      "loss": 0.5063,
+      "step": 1276
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.5616181822230734e-05,
+      "loss": 0.4995,
+      "step": 1277
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.560923530073405e-05,
+      "loss": 0.5024,
+      "step": 1278
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.5602284827721552e-05,
+      "loss": 0.4954,
+      "step": 1279
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.5595330408089613e-05,
+      "loss": 0.5144,
+      "step": 1280
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.558837204673739e-05,
+      "loss": 0.4675,
+      "step": 1281
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.55814097485668e-05,
+      "loss": 0.5034,
+      "step": 1282
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.5574443518482564e-05,
+      "loss": 0.5068,
+      "step": 1283
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.5567473361392135e-05,
+      "loss": 0.5044,
+      "step": 1284
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.556049928220577e-05,
+      "loss": 0.5027,
+      "step": 1285
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.5553521285836455e-05,
+      "loss": 0.4758,
+      "step": 1286
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.5546539377199964e-05,
+      "loss": 0.5046,
+      "step": 1287
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.553955356121481e-05,
+      "loss": 0.512,
+      "step": 1288
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.553256384280226e-05,
+      "loss": 0.4902,
+      "step": 1289
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.5525570226886334e-05,
+      "loss": 0.501,
+      "step": 1290
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.5518572718393804e-05,
+      "loss": 0.5166,
+      "step": 1291
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.5511571322254172e-05,
+      "loss": 0.4873,
+      "step": 1292
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5504566043399682e-05,
+      "loss": 0.5098,
+      "step": 1293
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5497556886765316e-05,
+      "loss": 0.5088,
+      "step": 1294
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5490543857288793e-05,
+      "loss": 0.5098,
+      "step": 1295
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5483526959910543e-05,
+      "loss": 0.4917,
+      "step": 1296
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.547650619957374e-05,
+      "loss": 0.4924,
+      "step": 1297
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5469481581224274e-05,
+      "loss": 0.5071,
+      "step": 1298
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5462453109810738e-05,
+      "loss": 0.4729,
+      "step": 1299
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5455420790284463e-05,
+      "loss": 0.5098,
+      "step": 1300
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5448384627599473e-05,
+      "loss": 0.459,
+      "step": 1301
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.544134462671251e-05,
+      "loss": 0.5186,
+      "step": 1302
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5434300792583012e-05,
+      "loss": 0.4958,
+      "step": 1303
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5427253130173126e-05,
+      "loss": 0.4944,
+      "step": 1304
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5420201644447682e-05,
+      "loss": 0.4824,
+      "step": 1305
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5413146340374222e-05,
+      "loss": 0.4968,
+      "step": 1306
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.540608722292296e-05,
+      "loss": 0.5032,
+      "step": 1307
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5399024297066813e-05,
+      "loss": 0.5107,
+      "step": 1308
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5391957567781363e-05,
+      "loss": 0.4937,
+      "step": 1309
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5384887040044885e-05,
+      "loss": 0.4993,
+      "step": 1310
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5377812718838332e-05,
+      "loss": 0.5139,
+      "step": 1311
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5370734609145314e-05,
+      "loss": 0.5117,
+      "step": 1312
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.536365271595212e-05,
+      "loss": 0.4902,
+      "step": 1313
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.535656704424771e-05,
+      "loss": 0.5049,
+      "step": 1314
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5349477599023696e-05,
+      "loss": 0.5054,
+      "step": 1315
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5342384385274346e-05,
+      "loss": 0.4868,
+      "step": 1316
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5335287407996597e-05,
+      "loss": 0.4634,
+      "step": 1317
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.532818667219002e-05,
+      "loss": 0.4834,
+      "step": 1318
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5321082182856848e-05,
+      "loss": 0.4993,
+      "step": 1319
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.531397394500195e-05,
+      "loss": 0.4656,
+      "step": 1320
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.530686196363284e-05,
+      "loss": 0.4907,
+      "step": 1321
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5299746243759668e-05,
+      "loss": 0.4858,
+      "step": 1322
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5292626790395213e-05,
+      "loss": 0.5017,
+      "step": 1323
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5285503608554893e-05,
+      "loss": 0.4993,
+      "step": 1324
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.527837670325674e-05,
+      "loss": 0.5413,
+      "step": 1325
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5271246079521428e-05,
+      "loss": 0.49,
+      "step": 1326
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.526411174237223e-05,
+      "loss": 0.5173,
+      "step": 1327
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5256973696835048e-05,
+      "loss": 0.5037,
+      "step": 1328
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5249831947938394e-05,
+      "loss": 0.4419,
+      "step": 1329
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5242686500713383e-05,
+      "loss": 0.488,
+      "step": 1330
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.5235537360193749e-05,
+      "loss": 0.5132,
+      "step": 1331
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5228384531415809e-05,
+      "loss": 0.4856,
+      "step": 1332
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.522122801941849e-05,
+      "loss": 0.4763,
+      "step": 1333
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5214067829243316e-05,
+      "loss": 0.4951,
+      "step": 1334
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5206903965934394e-05,
+      "loss": 0.4922,
+      "step": 1335
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5199736434538425e-05,
+      "loss": 0.4937,
+      "step": 1336
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5192565240104691e-05,
+      "loss": 0.4751,
+      "step": 1337
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5185390387685055e-05,
+      "loss": 0.5027,
+      "step": 1338
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5178211882333953e-05,
+      "loss": 0.5093,
+      "step": 1339
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5171029729108404e-05,
+      "loss": 0.4846,
+      "step": 1340
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5163843933067986e-05,
+      "loss": 0.5002,
+      "step": 1341
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5156654499274852e-05,
+      "loss": 0.5115,
+      "step": 1342
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5149461432793713e-05,
+      "loss": 0.48,
+      "step": 1343
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.514226473869184e-05,
+      "loss": 0.4897,
+      "step": 1344
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.513506442203906e-05,
+      "loss": 0.4968,
+      "step": 1345
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5127860487907752e-05,
+      "loss": 0.4773,
+      "step": 1346
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5120652941372842e-05,
+      "loss": 0.4854,
+      "step": 1347
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.51134417875118e-05,
+      "loss": 0.4934,
+      "step": 1348
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5106227031404643e-05,
+      "loss": 0.5017,
+      "step": 1349
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.509900867813392e-05,
+      "loss": 0.4856,
+      "step": 1350
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5091786732784717e-05,
+      "loss": 0.5085,
+      "step": 1351
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5084561200444646e-05,
+      "loss": 0.4856,
+      "step": 1352
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5077332086203847e-05,
+      "loss": 0.4844,
+      "step": 1353
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5070099395154994e-05,
+      "loss": 0.4998,
+      "step": 1354
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5062863132393269e-05,
+      "loss": 0.4851,
+      "step": 1355
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.505562330301637e-05,
+      "loss": 0.5056,
+      "step": 1356
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5048379912124516e-05,
+      "loss": 0.5166,
+      "step": 1357
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5041132964820424e-05,
+      "loss": 0.4822,
+      "step": 1358
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5033882466209327e-05,
+      "loss": 0.478,
+      "step": 1359
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.502662842139895e-05,
+      "loss": 0.4897,
+      "step": 1360
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5019370835499532e-05,
+      "loss": 0.5005,
+      "step": 1361
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5012109713623782e-05,
+      "loss": 0.4958,
+      "step": 1362
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.5004845060886921e-05,
+      "loss": 0.491,
+      "step": 1363
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.4997576882406652e-05,
+      "loss": 0.4976,
+      "step": 1364
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.499030518330316e-05,
+      "loss": 0.49,
+      "step": 1365
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.4983029968699103e-05,
+      "loss": 0.4802,
+      "step": 1366
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.4975751243719632e-05,
+      "loss": 0.4956,
+      "step": 1367
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.4968469013492355e-05,
+      "loss": 0.5208,
+      "step": 1368
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.4961183283147356e-05,
+      "loss": 0.4729,
+      "step": 1369
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.495389405781719e-05,
+      "loss": 0.4888,
+      "step": 1370
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4946601342636861e-05,
+      "loss": 0.5139,
+      "step": 1371
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4939305142743845e-05,
+      "loss": 0.4832,
+      "step": 1372
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.493200546327807e-05,
+      "loss": 0.5273,
+      "step": 1373
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4924702309381905e-05,
+      "loss": 0.4868,
+      "step": 1374
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4917395686200173e-05,
+      "loss": 0.4832,
+      "step": 1375
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4910085598880151e-05,
+      "loss": 0.498,
+      "step": 1376
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4902772052571541e-05,
+      "loss": 0.4692,
+      "step": 1377
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4895455052426491e-05,
+      "loss": 0.4749,
+      "step": 1378
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4888134603599575e-05,
+      "loss": 0.4895,
+      "step": 1379
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4880810711247803e-05,
+      "loss": 0.4707,
+      "step": 1380
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.487348338053061e-05,
+      "loss": 0.4844,
+      "step": 1381
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4866152616609849e-05,
+      "loss": 0.4922,
+      "step": 1382
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4858818424649795e-05,
+      "loss": 0.4744,
+      "step": 1383
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4851480809817139e-05,
+      "loss": 0.4729,
+      "step": 1384
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4844139777280979e-05,
+      "loss": 0.5061,
+      "step": 1385
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4836795332212825e-05,
+      "loss": 0.4841,
+      "step": 1386
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.482944747978659e-05,
+      "loss": 0.4717,
+      "step": 1387
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4822096225178582e-05,
+      "loss": 0.468,
+      "step": 1388
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4814741573567514e-05,
+      "loss": 0.5002,
+      "step": 1389
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4807383530134495e-05,
+      "loss": 0.5042,
+      "step": 1390
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4800022100062998e-05,
+      "loss": 0.4946,
+      "step": 1391
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4792657288538918e-05,
+      "loss": 0.4668,
+      "step": 1392
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4785289100750509e-05,
+      "loss": 0.4934,
+      "step": 1393
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4777917541888408e-05,
+      "loss": 0.4651,
+      "step": 1394
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4770542617145627e-05,
+      "loss": 0.4956,
+      "step": 1395
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4763164331717556e-05,
+      "loss": 0.5007,
+      "step": 1396
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.475578269080194e-05,
+      "loss": 0.4956,
+      "step": 1397
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4748397699598895e-05,
+      "loss": 0.5327,
+      "step": 1398
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.47410093633109e-05,
+      "loss": 0.5032,
+      "step": 1399
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4733617687142782e-05,
+      "loss": 0.5061,
+      "step": 1400
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4726222676301734e-05,
+      "loss": 0.4739,
+      "step": 1401
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4718824335997283e-05,
+      "loss": 0.4854,
+      "step": 1402
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4711422671441311e-05,
+      "loss": 0.491,
+      "step": 1403
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.470401768784804e-05,
+      "loss": 0.4714,
+      "step": 1404
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4696609390434029e-05,
+      "loss": 0.4807,
+      "step": 1405
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.468919778441817e-05,
+      "loss": 0.4614,
+      "step": 1406
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.4681782875021689e-05,
+      "loss": 0.4866,
+      "step": 1407
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.467436466746814e-05,
+      "loss": 0.4844,
+      "step": 1408
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4666943166983395e-05,
+      "loss": 0.5298,
+      "step": 1409
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4659518378795651e-05,
+      "loss": 0.4727,
+      "step": 1410
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4652090308135423e-05,
+      "loss": 0.5227,
+      "step": 1411
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4644658960235525e-05,
+      "loss": 0.4883,
+      "step": 1412
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4637224340331098e-05,
+      "loss": 0.489,
+      "step": 1413
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4629786453659576e-05,
+      "loss": 0.4856,
+      "step": 1414
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4622345305460696e-05,
+      "loss": 0.4861,
+      "step": 1415
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4614900900976496e-05,
+      "loss": 0.4888,
+      "step": 1416
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4607453245451307e-05,
+      "loss": 0.4758,
+      "step": 1417
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.460000234413175e-05,
+      "loss": 0.4832,
+      "step": 1418
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.459254820226673e-05,
+      "loss": 0.478,
+      "step": 1419
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4585090825107435e-05,
+      "loss": 0.4856,
+      "step": 1420
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.457763021790734e-05,
+      "loss": 0.4626,
+      "step": 1421
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4570166385922182e-05,
+      "loss": 0.4697,
+      "step": 1422
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4562699334409982e-05,
+      "loss": 0.4758,
+      "step": 1423
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4555229068631024e-05,
+      "loss": 0.4553,
+      "step": 1424
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4547755593847859e-05,
+      "loss": 0.4963,
+      "step": 1425
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.454027891532529e-05,
+      "loss": 0.4773,
+      "step": 1426
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4532799038330385e-05,
+      "loss": 0.5,
+      "step": 1427
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4525315968132466e-05,
+      "loss": 0.4778,
+      "step": 1428
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4517829710003102e-05,
+      "loss": 0.4675,
+      "step": 1429
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4510340269216104e-05,
+      "loss": 0.4678,
+      "step": 1430
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4502847651047534e-05,
+      "loss": 0.4756,
+      "step": 1431
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4495351860775682e-05,
+      "loss": 0.4634,
+      "step": 1432
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4487852903681079e-05,
+      "loss": 0.4597,
+      "step": 1433
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4480350785046486e-05,
+      "loss": 0.4636,
+      "step": 1434
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4472845510156892e-05,
+      "loss": 0.4907,
+      "step": 1435
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4465337084299509e-05,
+      "loss": 0.4919,
+      "step": 1436
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4457825512763767e-05,
+      "loss": 0.457,
+      "step": 1437
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4450310800841312e-05,
+      "loss": 0.4741,
+      "step": 1438
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4442792953826002e-05,
+      "loss": 0.4709,
+      "step": 1439
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.443527197701391e-05,
+      "loss": 0.5046,
+      "step": 1440
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4427747875703304e-05,
+      "loss": 0.5,
+      "step": 1441
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4420220655194661e-05,
+      "loss": 0.4658,
+      "step": 1442
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4412690320790652e-05,
+      "loss": 0.4854,
+      "step": 1443
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4405156877796139e-05,
+      "loss": 0.4978,
+      "step": 1444
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4397620331518178e-05,
+      "loss": 0.4932,
+      "step": 1445
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4390080687266013e-05,
+      "loss": 0.4727,
+      "step": 1446
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.4382537950351063e-05,
+      "loss": 0.4783,
+      "step": 1447
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4374992126086934e-05,
+      "loss": 0.4993,
+      "step": 1448
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4367443219789403e-05,
+      "loss": 0.4824,
+      "step": 1449
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4359891236776414e-05,
+      "loss": 0.4719,
+      "step": 1450
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4352336182368082e-05,
+      "loss": 0.4729,
+      "step": 1451
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4344778061886695e-05,
+      "loss": 0.5247,
+      "step": 1452
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4337216880656683e-05,
+      "loss": 0.498,
+      "step": 1453
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4329652644004645e-05,
+      "loss": 0.499,
+      "step": 1454
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4322085357259332e-05,
+      "loss": 0.4873,
+      "step": 1455
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4314515025751634e-05,
+      "loss": 0.4734,
+      "step": 1456
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4306941654814594e-05,
+      "loss": 0.4817,
+      "step": 1457
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.42993652497834e-05,
+      "loss": 0.4944,
+      "step": 1458
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.429178581599537e-05,
+      "loss": 0.4753,
+      "step": 1459
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4284203358789954e-05,
+      "loss": 0.4785,
+      "step": 1460
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4276617883508739e-05,
+      "loss": 0.501,
+      "step": 1461
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4269029395495432e-05,
+      "loss": 0.4976,
+      "step": 1462
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.426143790009587e-05,
+      "loss": 0.4661,
+      "step": 1463
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4253843402658e-05,
+      "loss": 0.4905,
+      "step": 1464
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4246245908531883e-05,
+      "loss": 0.4604,
+      "step": 1465
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4238645423069709e-05,
+      "loss": 0.4622,
+      "step": 1466
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.423104195162575e-05,
+      "loss": 0.4805,
+      "step": 1467
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4223435499556397e-05,
+      "loss": 0.4993,
+      "step": 1468
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4215826072220136e-05,
+      "loss": 0.4451,
+      "step": 1469
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4208213674977555e-05,
+      "loss": 0.5095,
+      "step": 1470
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4200598313191325e-05,
+      "loss": 0.4919,
+      "step": 1471
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4192979992226211e-05,
+      "loss": 0.4697,
+      "step": 1472
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4185358717449063e-05,
+      "loss": 0.4788,
+      "step": 1473
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4177734494228808e-05,
+      "loss": 0.4883,
+      "step": 1474
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4170107327936452e-05,
+      "loss": 0.4922,
+      "step": 1475
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4162477223945073e-05,
+      "loss": 0.4905,
+      "step": 1476
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4154844187629828e-05,
+      "loss": 0.4841,
+      "step": 1477
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4147208224367922e-05,
+      "loss": 0.4724,
+      "step": 1478
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4139569339538642e-05,
+      "loss": 0.4966,
+      "step": 1479
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4131927538523315e-05,
+      "loss": 0.51,
+      "step": 1480
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4124282826705333e-05,
+      "loss": 0.4717,
+      "step": 1481
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4116635209470135e-05,
+      "loss": 0.4727,
+      "step": 1482
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4108984692205212e-05,
+      "loss": 0.4854,
+      "step": 1483
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.410133128030009e-05,
+      "loss": 0.4556,
+      "step": 1484
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.4093674979146335e-05,
+      "loss": 0.4666,
+      "step": 1485
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.4086015794137555e-05,
+      "loss": 0.4888,
+      "step": 1486
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.4078353730669386e-05,
+      "loss": 0.48,
+      "step": 1487
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.4070688794139488e-05,
+      "loss": 0.4749,
+      "step": 1488
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.4063020989947555e-05,
+      "loss": 0.4441,
+      "step": 1489
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.4055350323495292e-05,
+      "loss": 0.4785,
+      "step": 1490
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.4047676800186417e-05,
+      "loss": 0.4807,
+      "step": 1491
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.4040000425426675e-05,
+      "loss": 0.46,
+      "step": 1492
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.4032321204623807e-05,
+      "loss": 0.4563,
+      "step": 1493
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.4024639143187563e-05,
+      "loss": 0.4858,
+      "step": 1494
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.4016954246529697e-05,
+      "loss": 0.4866,
+      "step": 1495
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.4009266520063956e-05,
+      "loss": 0.4685,
+      "step": 1496
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.4001575969206081e-05,
+      "loss": 0.479,
+      "step": 1497
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.3993882599373808e-05,
+      "loss": 0.4878,
+      "step": 1498
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.3986186415986854e-05,
+      "loss": 0.4819,
+      "step": 1499
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.3978487424466914e-05,
+      "loss": 0.4819,
+      "step": 1500
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.397078563023767e-05,
+      "loss": 0.4546,
+      "step": 1501
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.3963081038724778e-05,
+      "loss": 0.5212,
+      "step": 1502
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.3955373655355852e-05,
+      "loss": 0.4617,
+      "step": 1503
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.3947663485560488e-05,
+      "loss": 0.4644,
+      "step": 1504
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.3939950534770238e-05,
+      "loss": 0.4641,
+      "step": 1505
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.3932234808418614e-05,
+      "loss": 0.4409,
+      "step": 1506
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.3924516311941085e-05,
+      "loss": 0.4634,
+      "step": 1507
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.3916795050775066e-05,
+      "loss": 0.4631,
+      "step": 1508
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.3909071030359927e-05,
+      "loss": 0.4907,
+      "step": 1509
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.3901344256136974e-05,
+      "loss": 0.4583,
+      "step": 1510
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.3893614733549459e-05,
+      "loss": 0.4983,
+      "step": 1511
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.3885882468042571e-05,
+      "loss": 0.4866,
+      "step": 1512
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.3878147465063423e-05,
+      "loss": 0.5042,
+      "step": 1513
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.3870409730061067e-05,
+      "loss": 0.4954,
+      "step": 1514
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.386266926848647e-05,
+      "loss": 0.4521,
+      "step": 1515
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.3854926085792528e-05,
+      "loss": 0.5115,
+      "step": 1516
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.384718018743405e-05,
+      "loss": 0.4937,
+      "step": 1517
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.3839431578867759e-05,
+      "loss": 0.467,
+      "step": 1518
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.3831680265552283e-05,
+      "loss": 0.4941,
+      "step": 1519
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.3823926252948161e-05,
+      "loss": 0.4871,
+      "step": 1520
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.3816169546517833e-05,
+      "loss": 0.4595,
+      "step": 1521
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.3808410151725633e-05,
+      "loss": 0.4805,
+      "step": 1522
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.3800648074037792e-05,
+      "loss": 0.4927,
+      "step": 1523
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.3792883318922434e-05,
+      "loss": 0.4927,
+      "step": 1524
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3785115891849556e-05,
+      "loss": 0.4592,
+      "step": 1525
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3777345798291057e-05,
+      "loss": 0.4705,
+      "step": 1526
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3769573043720699e-05,
+      "loss": 0.4917,
+      "step": 1527
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3761797633614122e-05,
+      "loss": 0.4705,
+      "step": 1528
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.375401957344884e-05,
+      "loss": 0.48,
+      "step": 1529
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3746238868704237e-05,
+      "loss": 0.4668,
+      "step": 1530
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3738455524861551e-05,
+      "loss": 0.4822,
+      "step": 1531
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3730669547403879e-05,
+      "loss": 0.4727,
+      "step": 1532
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3722880941816185e-05,
+      "loss": 0.49,
+      "step": 1533
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3715089713585277e-05,
+      "loss": 0.4863,
+      "step": 1534
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3707295868199806e-05,
+      "loss": 0.4783,
+      "step": 1535
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3699499411150274e-05,
+      "loss": 0.5068,
+      "step": 1536
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.369170034792902e-05,
+      "loss": 0.4722,
+      "step": 1537
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3683898684030219e-05,
+      "loss": 0.4812,
+      "step": 1538
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.367609442494988e-05,
+      "loss": 0.4648,
+      "step": 1539
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3668287576185835e-05,
+      "loss": 0.4629,
+      "step": 1540
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3660478143237748e-05,
+      "loss": 0.4861,
+      "step": 1541
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3652666131607098e-05,
+      "loss": 0.4915,
+      "step": 1542
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.364485154679718e-05,
+      "loss": 0.488,
+      "step": 1543
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3637034394313102e-05,
+      "loss": 0.4717,
+      "step": 1544
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3629214679661788e-05,
+      "loss": 0.4734,
+      "step": 1545
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3621392408351961e-05,
+      "loss": 0.4575,
+      "step": 1546
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3613567585894142e-05,
+      "loss": 0.4709,
+      "step": 1547
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3605740217800653e-05,
+      "loss": 0.481,
+      "step": 1548
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3597910309585614e-05,
+      "loss": 0.4873,
+      "step": 1549
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3590077866764921e-05,
+      "loss": 0.4785,
+      "step": 1550
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.358224289485627e-05,
+      "loss": 0.4614,
+      "step": 1551
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3574405399379133e-05,
+      "loss": 0.4775,
+      "step": 1552
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.356656538585476e-05,
+      "loss": 0.468,
+      "step": 1553
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3558722859806168e-05,
+      "loss": 0.4631,
+      "step": 1554
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3550877826758155e-05,
+      "loss": 0.4727,
+      "step": 1555
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3543030292237283e-05,
+      "loss": 0.4734,
+      "step": 1556
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3535180261771865e-05,
+      "loss": 0.4587,
+      "step": 1557
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3527327740891991e-05,
+      "loss": 0.499,
+      "step": 1558
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3519472735129493e-05,
+      "loss": 0.499,
+      "step": 1559
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3511615250017948e-05,
+      "loss": 0.4697,
+      "step": 1560
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.3503755291092698e-05,
+      "loss": 0.4656,
+      "step": 1561
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.349589286389081e-05,
+      "loss": 0.4749,
+      "step": 1562
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3488027973951096e-05,
+      "loss": 0.4983,
+      "step": 1563
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3480160626814114e-05,
+      "loss": 0.4814,
+      "step": 1564
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3472290828022132e-05,
+      "loss": 0.4568,
+      "step": 1565
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.346441858311916e-05,
+      "loss": 0.4768,
+      "step": 1566
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3456543897650928e-05,
+      "loss": 0.5103,
+      "step": 1567
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3448666777164885e-05,
+      "loss": 0.5107,
+      "step": 1568
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3440787227210194e-05,
+      "loss": 0.4756,
+      "step": 1569
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.343290525333773e-05,
+      "loss": 0.4666,
+      "step": 1570
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3425020861100078e-05,
+      "loss": 0.4971,
+      "step": 1571
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.341713405605152e-05,
+      "loss": 0.4844,
+      "step": 1572
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3409244843748046e-05,
+      "loss": 0.5012,
+      "step": 1573
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3401353229747338e-05,
+      "loss": 0.4658,
+      "step": 1574
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.339345921960877e-05,
+      "loss": 0.4763,
+      "step": 1575
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.33855628188934e-05,
+      "loss": 0.4763,
+      "step": 1576
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3377664033163983e-05,
+      "loss": 0.5083,
+      "step": 1577
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3369762867984939e-05,
+      "loss": 0.4744,
+      "step": 1578
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3361859328922368e-05,
+      "loss": 0.4626,
+      "step": 1579
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3353953421544052e-05,
+      "loss": 0.4619,
+      "step": 1580
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3346045151419432e-05,
+      "loss": 0.4763,
+      "step": 1581
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3338134524119612e-05,
+      "loss": 0.501,
+      "step": 1582
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3330221545217365e-05,
+      "loss": 0.4995,
+      "step": 1583
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3322306220287116e-05,
+      "loss": 0.4949,
+      "step": 1584
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3314388554904936e-05,
+      "loss": 0.4934,
+      "step": 1585
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.330646855464856e-05,
+      "loss": 0.4546,
+      "step": 1586
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3298546225097353e-05,
+      "loss": 0.4543,
+      "step": 1587
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3290621571832331e-05,
+      "loss": 0.4634,
+      "step": 1588
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3282694600436141e-05,
+      "loss": 0.4604,
+      "step": 1589
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3274765316493068e-05,
+      "loss": 0.4976,
+      "step": 1590
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.326683372558902e-05,
+      "loss": 0.5044,
+      "step": 1591
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3258899833311534e-05,
+      "loss": 0.49,
+      "step": 1592
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.325096364524977e-05,
+      "loss": 0.4656,
+      "step": 1593
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3243025166994499e-05,
+      "loss": 0.4663,
+      "step": 1594
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3235084404138114e-05,
+      "loss": 0.4541,
+      "step": 1595
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3227141362274612e-05,
+      "loss": 0.4663,
+      "step": 1596
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3219196046999594e-05,
+      "loss": 0.4861,
+      "step": 1597
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3211248463910263e-05,
+      "loss": 0.4612,
+      "step": 1598
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3203298618605426e-05,
+      "loss": 0.4634,
+      "step": 1599
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3195346516685481e-05,
+      "loss": 0.4763,
+      "step": 1600
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.3187392163752404e-05,
+      "loss": 0.5046,
+      "step": 1601
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.3179435565409775e-05,
+      "loss": 0.4653,
+      "step": 1602
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.3171476727262746e-05,
+      "loss": 0.4741,
+      "step": 1603
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.3163515654918046e-05,
+      "loss": 0.4656,
+      "step": 1604
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.3155552353983985e-05,
+      "loss": 0.4866,
+      "step": 1605
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.3147586830070434e-05,
+      "loss": 0.4736,
+      "step": 1606
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.3139619088788828e-05,
+      "loss": 0.4856,
+      "step": 1607
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.3131649135752185e-05,
+      "loss": 0.479,
+      "step": 1608
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.3123676976575055e-05,
+      "loss": 0.4932,
+      "step": 1609
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.3115702616873556e-05,
+      "loss": 0.4744,
+      "step": 1610
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.3107726062265356e-05,
+      "loss": 0.4895,
+      "step": 1611
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.3099747318369664e-05,
+      "loss": 0.4822,
+      "step": 1612
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.3091766390807235e-05,
+      "loss": 0.4819,
+      "step": 1613
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.3083783285200362e-05,
+      "loss": 0.4604,
+      "step": 1614
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.3075798007172872e-05,
+      "loss": 0.4771,
+      "step": 1615
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.3067810562350123e-05,
+      "loss": 0.4658,
+      "step": 1616
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.3059820956358998e-05,
+      "loss": 0.4917,
+      "step": 1617
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.30518291948279e-05,
+      "loss": 0.4719,
+      "step": 1618
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.304383528338676e-05,
+      "loss": 0.4673,
+      "step": 1619
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.303583922766701e-05,
+      "loss": 0.4478,
+      "step": 1620
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.3027841033301605e-05,
+      "loss": 0.4534,
+      "step": 1621
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.3019840705925e-05,
+      "loss": 0.467,
+      "step": 1622
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.3011838251173152e-05,
+      "loss": 0.4768,
+      "step": 1623
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.3003833674683524e-05,
+      "loss": 0.4956,
+      "step": 1624
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.299582698209506e-05,
+      "loss": 0.49,
+      "step": 1625
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.298781817904821e-05,
+      "loss": 0.4685,
+      "step": 1626
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.2979807271184903e-05,
+      "loss": 0.501,
+      "step": 1627
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.297179426414855e-05,
+      "loss": 0.4678,
+      "step": 1628
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.2963779163584041e-05,
+      "loss": 0.4763,
+      "step": 1629
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.2955761975137743e-05,
+      "loss": 0.4934,
+      "step": 1630
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.2947742704457495e-05,
+      "loss": 0.457,
+      "step": 1631
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.2939721357192598e-05,
+      "loss": 0.4941,
+      "step": 1632
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.2931697938993824e-05,
+      "loss": 0.4854,
+      "step": 1633
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.2923672455513394e-05,
+      "loss": 0.4851,
+      "step": 1634
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.2915644912404988e-05,
+      "loss": 0.4929,
+      "step": 1635
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.290761531532374e-05,
+      "loss": 0.4875,
+      "step": 1636
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.289958366992623e-05,
+      "loss": 0.4861,
+      "step": 1637
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.2891549981870475e-05,
+      "loss": 0.4612,
+      "step": 1638
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.2883514256815937e-05,
+      "loss": 0.4919,
+      "step": 1639
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.2875476500423516e-05,
+      "loss": 0.4644,
+      "step": 1640
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.286743671835553e-05,
+      "loss": 0.4775,
+      "step": 1641
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.285939491627574e-05,
+      "loss": 0.4954,
+      "step": 1642
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2851351099849321e-05,
+      "loss": 0.4587,
+      "step": 1643
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2843305274742866e-05,
+      "loss": 0.489,
+      "step": 1644
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2835257446624386e-05,
+      "loss": 0.4941,
+      "step": 1645
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2827207621163305e-05,
+      "loss": 0.4485,
+      "step": 1646
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.281915580403045e-05,
+      "loss": 0.4705,
+      "step": 1647
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2811102000898049e-05,
+      "loss": 0.4905,
+      "step": 1648
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2803046217439741e-05,
+      "loss": 0.4526,
+      "step": 1649
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2794988459330547e-05,
+      "loss": 0.4702,
+      "step": 1650
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2786928732246885e-05,
+      "loss": 0.4602,
+      "step": 1651
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2778867041866559e-05,
+      "loss": 0.4629,
+      "step": 1652
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2770803393868756e-05,
+      "loss": 0.4648,
+      "step": 1653
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.276273779393404e-05,
+      "loss": 0.4758,
+      "step": 1654
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2754670247744353e-05,
+      "loss": 0.4504,
+      "step": 1655
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2746600760983015e-05,
+      "loss": 0.4924,
+      "step": 1656
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2738529339334695e-05,
+      "loss": 0.4529,
+      "step": 1657
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2730455988485442e-05,
+      "loss": 0.4841,
+      "step": 1658
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2722380714122655e-05,
+      "loss": 0.4563,
+      "step": 1659
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2714303521935088e-05,
+      "loss": 0.4812,
+      "step": 1660
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2706224417612853e-05,
+      "loss": 0.4756,
+      "step": 1661
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2698143406847405e-05,
+      "loss": 0.4888,
+      "step": 1662
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.269006049533154e-05,
+      "loss": 0.48,
+      "step": 1663
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2681975688759394e-05,
+      "loss": 0.4683,
+      "step": 1664
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.267388899282644e-05,
+      "loss": 0.4763,
+      "step": 1665
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2665800413229479e-05,
+      "loss": 0.4636,
+      "step": 1666
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.265770995566664e-05,
+      "loss": 0.4844,
+      "step": 1667
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2649617625837381e-05,
+      "loss": 0.4729,
+      "step": 1668
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.264152342944247e-05,
+      "loss": 0.4639,
+      "step": 1669
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.263342737218399e-05,
+      "loss": 0.4829,
+      "step": 1670
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2625329459765344e-05,
+      "loss": 0.4656,
+      "step": 1671
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2617229697891234e-05,
+      "loss": 0.449,
+      "step": 1672
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2609128092267666e-05,
+      "loss": 0.4612,
+      "step": 1673
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.260102464860195e-05,
+      "loss": 0.4685,
+      "step": 1674
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2592919372602685e-05,
+      "loss": 0.4797,
+      "step": 1675
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2584812269979761e-05,
+      "loss": 0.4734,
+      "step": 1676
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2576703346444359e-05,
+      "loss": 0.4871,
+      "step": 1677
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.2568592607708938e-05,
+      "loss": 0.458,
+      "step": 1678
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2560480059487237e-05,
+      "loss": 0.4824,
+      "step": 1679
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2552365707494279e-05,
+      "loss": 0.5063,
+      "step": 1680
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.254424955744634e-05,
+      "loss": 0.4783,
+      "step": 1681
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2536131615060972e-05,
+      "loss": 0.4783,
+      "step": 1682
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2528011886056996e-05,
+      "loss": 0.4663,
+      "step": 1683
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2519890376154484e-05,
+      "loss": 0.4768,
+      "step": 1684
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2511767091074762e-05,
+      "loss": 0.4646,
+      "step": 1685
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2503642036540408e-05,
+      "loss": 0.4868,
+      "step": 1686
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2495515218275251e-05,
+      "loss": 0.4675,
+      "step": 1687
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2487386642004353e-05,
+      "loss": 0.4607,
+      "step": 1688
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2479256313454023e-05,
+      "loss": 0.4675,
+      "step": 1689
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2471124238351802e-05,
+      "loss": 0.4497,
+      "step": 1690
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.246299042242646e-05,
+      "loss": 0.4656,
+      "step": 1691
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2454854871407993e-05,
+      "loss": 0.5054,
+      "step": 1692
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2446717591027624e-05,
+      "loss": 0.4729,
+      "step": 1693
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2438578587017784e-05,
+      "loss": 0.4966,
+      "step": 1694
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2430437865112132e-05,
+      "loss": 0.4641,
+      "step": 1695
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2422295431045525e-05,
+      "loss": 0.4722,
+      "step": 1696
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2414151290554037e-05,
+      "loss": 0.4832,
+      "step": 1697
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2406005449374936e-05,
+      "loss": 0.4656,
+      "step": 1698
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2397857913246688e-05,
+      "loss": 0.4822,
+      "step": 1699
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2389708687908956e-05,
+      "loss": 0.4868,
+      "step": 1700
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2381557779102593e-05,
+      "loss": 0.4827,
+      "step": 1701
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.237340519256964e-05,
+      "loss": 0.4412,
+      "step": 1702
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2365250934053317e-05,
+      "loss": 0.4397,
+      "step": 1703
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.235709500929802e-05,
+      "loss": 0.4541,
+      "step": 1704
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2348937424049321e-05,
+      "loss": 0.4597,
+      "step": 1705
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2340778184053965e-05,
+      "loss": 0.4568,
+      "step": 1706
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2332617295059857e-05,
+      "loss": 0.4651,
+      "step": 1707
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2324454762816062e-05,
+      "loss": 0.4761,
+      "step": 1708
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2316290593072817e-05,
+      "loss": 0.4666,
+      "step": 1709
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2308124791581496e-05,
+      "loss": 0.4941,
+      "step": 1710
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2299957364094625e-05,
+      "loss": 0.4688,
+      "step": 1711
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.229178831636589e-05,
+      "loss": 0.4871,
+      "step": 1712
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.22836176541501e-05,
+      "loss": 0.4568,
+      "step": 1713
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2275445383203209e-05,
+      "loss": 0.4846,
+      "step": 1714
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2267271509282308e-05,
+      "loss": 0.4863,
+      "step": 1715
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.2259096038145612e-05,
+      "loss": 0.4788,
+      "step": 1716
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.225091897555246e-05,
+      "loss": 0.4744,
+      "step": 1717
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2242740327263319e-05,
+      "loss": 0.4675,
+      "step": 1718
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.223456009903977e-05,
+      "loss": 0.4956,
+      "step": 1719
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.22263782966445e-05,
+      "loss": 0.4705,
+      "step": 1720
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2218194925841317e-05,
+      "loss": 0.4766,
+      "step": 1721
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2210009992395127e-05,
+      "loss": 0.4749,
+      "step": 1722
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2201823502071932e-05,
+      "loss": 0.4524,
+      "step": 1723
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2193635460638845e-05,
+      "loss": 0.4636,
+      "step": 1724
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2185445873864058e-05,
+      "loss": 0.4431,
+      "step": 1725
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2177254747516858e-05,
+      "loss": 0.4792,
+      "step": 1726
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2169062087367614e-05,
+      "loss": 0.4641,
+      "step": 1727
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2160867899187778e-05,
+      "loss": 0.4656,
+      "step": 1728
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2152672188749877e-05,
+      "loss": 0.4255,
+      "step": 1729
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.214447496182751e-05,
+      "loss": 0.4294,
+      "step": 1730
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2136276224195349e-05,
+      "loss": 0.468,
+      "step": 1731
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2128075981629124e-05,
+      "loss": 0.4758,
+      "step": 1732
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2119874239905624e-05,
+      "loss": 0.4548,
+      "step": 1733
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2111671004802702e-05,
+      "loss": 0.438,
+      "step": 1734
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2103466282099257e-05,
+      "loss": 0.468,
+      "step": 1735
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2095260077575234e-05,
+      "loss": 0.4849,
+      "step": 1736
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2087052397011634e-05,
+      "loss": 0.4531,
+      "step": 1737
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2078843246190483e-05,
+      "loss": 0.4661,
+      "step": 1738
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2070632630894849e-05,
+      "loss": 0.4673,
+      "step": 1739
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2062420556908836e-05,
+      "loss": 0.4836,
+      "step": 1740
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2054207030017567e-05,
+      "loss": 0.5083,
+      "step": 1741
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2045992056007192e-05,
+      "loss": 0.458,
+      "step": 1742
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2037775640664884e-05,
+      "loss": 0.489,
+      "step": 1743
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2029557789778832e-05,
+      "loss": 0.4802,
+      "step": 1744
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2021338509138223e-05,
+      "loss": 0.4629,
+      "step": 1745
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2013117804533269e-05,
+      "loss": 0.4971,
+      "step": 1746
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.2004895681755178e-05,
+      "loss": 0.4688,
+      "step": 1747
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.1996672146596151e-05,
+      "loss": 0.4814,
+      "step": 1748
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.1988447204849391e-05,
+      "loss": 0.4912,
+      "step": 1749
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.1980220862309097e-05,
+      "loss": 0.4834,
+      "step": 1750
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.1971993124770436e-05,
+      "loss": 0.4729,
+      "step": 1751
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.1963763998029578e-05,
+      "loss": 0.4749,
+      "step": 1752
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.1955533487883662e-05,
+      "loss": 0.4661,
+      "step": 1753
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.1947301600130803e-05,
+      "loss": 0.4827,
+      "step": 1754
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.193906834057008e-05,
+      "loss": 0.479,
+      "step": 1755
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1930833715001557e-05,
+      "loss": 0.5081,
+      "step": 1756
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1922597729226238e-05,
+      "loss": 0.4749,
+      "step": 1757
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.19143603890461e-05,
+      "loss": 0.4812,
+      "step": 1758
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1906121700264069e-05,
+      "loss": 0.4788,
+      "step": 1759
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1897881668684019e-05,
+      "loss": 0.4597,
+      "step": 1760
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1889640300110775e-05,
+      "loss": 0.4509,
+      "step": 1761
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1881397600350103e-05,
+      "loss": 0.4517,
+      "step": 1762
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1873153575208702e-05,
+      "loss": 0.4783,
+      "step": 1763
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1864908230494205e-05,
+      "loss": 0.4695,
+      "step": 1764
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1856661572015184e-05,
+      "loss": 0.489,
+      "step": 1765
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1848413605581126e-05,
+      "loss": 0.4731,
+      "step": 1766
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1840164337002445e-05,
+      "loss": 0.4719,
+      "step": 1767
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1831913772090469e-05,
+      "loss": 0.4741,
+      "step": 1768
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1823661916657441e-05,
+      "loss": 0.4746,
+      "step": 1769
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1815408776516513e-05,
+      "loss": 0.4441,
+      "step": 1770
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1807154357481743e-05,
+      "loss": 0.4675,
+      "step": 1771
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.179889866536809e-05,
+      "loss": 0.4734,
+      "step": 1772
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1790641705991407e-05,
+      "loss": 0.4641,
+      "step": 1773
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.178238348516844e-05,
+      "loss": 0.479,
+      "step": 1774
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.177412400871683e-05,
+      "loss": 0.4519,
+      "step": 1775
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.176586328245509e-05,
+      "loss": 0.4795,
+      "step": 1776
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.175760131220263e-05,
+      "loss": 0.4458,
+      "step": 1777
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1749338103779722e-05,
+      "loss": 0.4734,
+      "step": 1778
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1741073663007516e-05,
+      "loss": 0.4448,
+      "step": 1779
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.173280799570803e-05,
+      "loss": 0.4583,
+      "step": 1780
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1724541107704152e-05,
+      "loss": 0.4636,
+      "step": 1781
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1716273004819617e-05,
+      "loss": 0.4868,
+      "step": 1782
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.170800369287902e-05,
+      "loss": 0.4563,
+      "step": 1783
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.169973317770782e-05,
+      "loss": 0.4541,
+      "step": 1784
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.169146146513231e-05,
+      "loss": 0.4668,
+      "step": 1785
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1683188560979625e-05,
+      "loss": 0.4819,
+      "step": 1786
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1674914471077754e-05,
+      "loss": 0.4719,
+      "step": 1787
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1666639201255507e-05,
+      "loss": 0.4929,
+      "step": 1788
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.165836275734253e-05,
+      "loss": 0.4521,
+      "step": 1789
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.16500851451693e-05,
+      "loss": 0.4758,
+      "step": 1790
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1641806370567117e-05,
+      "loss": 0.4644,
+      "step": 1791
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1633526439368086e-05,
+      "loss": 0.4641,
+      "step": 1792
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1625245357405145e-05,
+      "loss": 0.4666,
+      "step": 1793
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.1616963130512035e-05,
+      "loss": 0.4639,
+      "step": 1794
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.16086797645233e-05,
+      "loss": 0.4529,
+      "step": 1795
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1600395265274296e-05,
+      "loss": 0.4792,
+      "step": 1796
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1592109638601166e-05,
+      "loss": 0.4922,
+      "step": 1797
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.158382289034085e-05,
+      "loss": 0.4861,
+      "step": 1798
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1575535026331088e-05,
+      "loss": 0.469,
+      "step": 1799
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1567246052410397e-05,
+      "loss": 0.47,
+      "step": 1800
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1558955974418076e-05,
+      "loss": 0.4729,
+      "step": 1801
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1550664798194204e-05,
+      "loss": 0.4573,
+      "step": 1802
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1542372529579634e-05,
+      "loss": 0.4399,
+      "step": 1803
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1534079174415986e-05,
+      "loss": 0.4521,
+      "step": 1804
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1525784738545645e-05,
+      "loss": 0.4666,
+      "step": 1805
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1517489227811765e-05,
+      "loss": 0.428,
+      "step": 1806
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.150919264805825e-05,
+      "loss": 0.4744,
+      "step": 1807
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1500895005129757e-05,
+      "loss": 0.4583,
+      "step": 1808
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1492596304871694e-05,
+      "loss": 0.4712,
+      "step": 1809
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1484296553130214e-05,
+      "loss": 0.4375,
+      "step": 1810
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.147599575575221e-05,
+      "loss": 0.4661,
+      "step": 1811
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1467693918585311e-05,
+      "loss": 0.449,
+      "step": 1812
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1459391047477883e-05,
+      "loss": 0.4722,
+      "step": 1813
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1451087148279013e-05,
+      "loss": 0.4514,
+      "step": 1814
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.144278222683852e-05,
+      "loss": 0.4548,
+      "step": 1815
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1434476289006935e-05,
+      "loss": 0.4585,
+      "step": 1816
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1426169340635509e-05,
+      "loss": 0.4492,
+      "step": 1817
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1417861387576208e-05,
+      "loss": 0.4856,
+      "step": 1818
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1409552435681702e-05,
+      "loss": 0.4814,
+      "step": 1819
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1401242490805364e-05,
+      "loss": 0.4333,
+      "step": 1820
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1392931558801273e-05,
+      "loss": 0.4753,
+      "step": 1821
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1384619645524192e-05,
+      "loss": 0.4678,
+      "step": 1822
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1376306756829583e-05,
+      "loss": 0.4717,
+      "step": 1823
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1367992898573594e-05,
+      "loss": 0.4841,
+      "step": 1824
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1359678076613058e-05,
+      "loss": 0.4773,
+      "step": 1825
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1351362296805487e-05,
+      "loss": 0.4724,
+      "step": 1826
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1343045565009061e-05,
+      "loss": 0.4785,
+      "step": 1827
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1334727887082634e-05,
+      "loss": 0.4585,
+      "step": 1828
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1326409268885732e-05,
+      "loss": 0.4863,
+      "step": 1829
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1318089716278535e-05,
+      "loss": 0.4888,
+      "step": 1830
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1309769235121892e-05,
+      "loss": 0.4595,
+      "step": 1831
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1301447831277293e-05,
+      "loss": 0.4788,
+      "step": 1832
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.1293125510606882e-05,
+      "loss": 0.4746,
+      "step": 1833
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1284802278973458e-05,
+      "loss": 0.4668,
+      "step": 1834
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1276478142240457e-05,
+      "loss": 0.46,
+      "step": 1835
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.126815310627194e-05,
+      "loss": 0.4661,
+      "step": 1836
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1259827176932625e-05,
+      "loss": 0.4688,
+      "step": 1837
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1251500360087837e-05,
+      "loss": 0.4468,
+      "step": 1838
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1243172661603536e-05,
+      "loss": 0.446,
+      "step": 1839
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.123484408734631e-05,
+      "loss": 0.4485,
+      "step": 1840
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1226514643183352e-05,
+      "loss": 0.4758,
+      "step": 1841
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.121818433498247e-05,
+      "loss": 0.4521,
+      "step": 1842
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.120985316861209e-05,
+      "loss": 0.4385,
+      "step": 1843
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1201521149941232e-05,
+      "loss": 0.4727,
+      "step": 1844
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1193188284839518e-05,
+      "loss": 0.4446,
+      "step": 1845
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1184854579177173e-05,
+      "loss": 0.4482,
+      "step": 1846
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1176520038825007e-05,
+      "loss": 0.4741,
+      "step": 1847
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1168184669654423e-05,
+      "loss": 0.4539,
+      "step": 1848
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.11598484775374e-05,
+      "loss": 0.4751,
+      "step": 1849
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.115151146834651e-05,
+      "loss": 0.4636,
+      "step": 1850
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1143173647954888e-05,
+      "loss": 0.469,
+      "step": 1851
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1134835022236244e-05,
+      "loss": 0.457,
+      "step": 1852
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1126495597064862e-05,
+      "loss": 0.4739,
+      "step": 1853
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.111815537831558e-05,
+      "loss": 0.4731,
+      "step": 1854
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1109814371863801e-05,
+      "loss": 0.4268,
+      "step": 1855
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1101472583585479e-05,
+      "loss": 0.4644,
+      "step": 1856
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1093130019357121e-05,
+      "loss": 0.4524,
+      "step": 1857
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.108478668505578e-05,
+      "loss": 0.4539,
+      "step": 1858
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1076442586559053e-05,
+      "loss": 0.4565,
+      "step": 1859
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1068097729745076e-05,
+      "loss": 0.4407,
+      "step": 1860
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1059752120492511e-05,
+      "loss": 0.4529,
+      "step": 1861
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1051405764680561e-05,
+      "loss": 0.4731,
+      "step": 1862
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.104305866818895e-05,
+      "loss": 0.4714,
+      "step": 1863
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1034710836897922e-05,
+      "loss": 0.4553,
+      "step": 1864
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1026362276688243e-05,
+      "loss": 0.449,
+      "step": 1865
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1018012993441194e-05,
+      "loss": 0.4602,
+      "step": 1866
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.100966299303855e-05,
+      "loss": 0.448,
+      "step": 1867
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.1001312281362612e-05,
+      "loss": 0.4751,
+      "step": 1868
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.0992960864296172e-05,
+      "loss": 0.4565,
+      "step": 1869
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.0984608747722517e-05,
+      "loss": 0.4558,
+      "step": 1870
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.0976255937525426e-05,
+      "loss": 0.4482,
+      "step": 1871
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0967902439589184e-05,
+      "loss": 0.4631,
+      "step": 1872
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0959548259798529e-05,
+      "loss": 0.4507,
+      "step": 1873
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.095119340403871e-05,
+      "loss": 0.4426,
+      "step": 1874
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0942837878195436e-05,
+      "loss": 0.4607,
+      "step": 1875
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0934481688154887e-05,
+      "loss": 0.4612,
+      "step": 1876
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.092612483980372e-05,
+      "loss": 0.449,
+      "step": 1877
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.091776733902905e-05,
+      "loss": 0.4805,
+      "step": 1878
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0909409191718453e-05,
+      "loss": 0.4456,
+      "step": 1879
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0901050403759951e-05,
+      "loss": 0.4548,
+      "step": 1880
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.089269098104204e-05,
+      "loss": 0.4595,
+      "step": 1881
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0884330929453643e-05,
+      "loss": 0.4697,
+      "step": 1882
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.087597025488413e-05,
+      "loss": 0.4692,
+      "step": 1883
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0867608963223319e-05,
+      "loss": 0.4751,
+      "step": 1884
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0859247060361451e-05,
+      "loss": 0.4929,
+      "step": 1885
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0850884552189202e-05,
+      "loss": 0.4478,
+      "step": 1886
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.084252144459768e-05,
+      "loss": 0.4568,
+      "step": 1887
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0834157743478407e-05,
+      "loss": 0.4626,
+      "step": 1888
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0825793454723325e-05,
+      "loss": 0.4727,
+      "step": 1889
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0817428584224794e-05,
+      "loss": 0.4546,
+      "step": 1890
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0809063137875581e-05,
+      "loss": 0.447,
+      "step": 1891
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0800697121568855e-05,
+      "loss": 0.469,
+      "step": 1892
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0792330541198197e-05,
+      "loss": 0.4626,
+      "step": 1893
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0783963402657575e-05,
+      "loss": 0.4661,
+      "step": 1894
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0775595711841356e-05,
+      "loss": 0.4863,
+      "step": 1895
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0767227474644293e-05,
+      "loss": 0.4832,
+      "step": 1896
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0758858696961523e-05,
+      "loss": 0.4631,
+      "step": 1897
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0750489384688568e-05,
+      "loss": 0.4646,
+      "step": 1898
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0742119543721322e-05,
+      "loss": 0.4458,
+      "step": 1899
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0733749179956057e-05,
+      "loss": 0.4434,
+      "step": 1900
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.072537829928941e-05,
+      "loss": 0.4727,
+      "step": 1901
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0717006907618377e-05,
+      "loss": 0.4714,
+      "step": 1902
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0708635010840324e-05,
+      "loss": 0.4719,
+      "step": 1903
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.070026261485297e-05,
+      "loss": 0.4673,
+      "step": 1904
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0691889725554377e-05,
+      "loss": 0.4741,
+      "step": 1905
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0683516348842966e-05,
+      "loss": 0.4761,
+      "step": 1906
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0675142490617499e-05,
+      "loss": 0.4851,
+      "step": 1907
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0666768156777066e-05,
+      "loss": 0.4546,
+      "step": 1908
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0658393353221108e-05,
+      "loss": 0.4558,
+      "step": 1909
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.0650018085849387e-05,
+      "loss": 0.4639,
+      "step": 1910
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0641642360561995e-05,
+      "loss": 0.4468,
+      "step": 1911
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0633266183259348e-05,
+      "loss": 0.4744,
+      "step": 1912
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.062488955984218e-05,
+      "loss": 0.4902,
+      "step": 1913
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0616512496211528e-05,
+      "loss": 0.4434,
+      "step": 1914
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.060813499826876e-05,
+      "loss": 0.4434,
+      "step": 1915
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0599757071915534e-05,
+      "loss": 0.4453,
+      "step": 1916
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0591378723053814e-05,
+      "loss": 0.4705,
+      "step": 1917
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0582999957585862e-05,
+      "loss": 0.4504,
+      "step": 1918
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0574620781414238e-05,
+      "loss": 0.4456,
+      "step": 1919
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0566241200441783e-05,
+      "loss": 0.4575,
+      "step": 1920
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0557861220571626e-05,
+      "loss": 0.457,
+      "step": 1921
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0549480847707184e-05,
+      "loss": 0.4763,
+      "step": 1922
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.054110008775214e-05,
+      "loss": 0.4407,
+      "step": 1923
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.053271894661046e-05,
+      "loss": 0.4294,
+      "step": 1924
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0524337430186368e-05,
+      "loss": 0.4351,
+      "step": 1925
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.051595554438436e-05,
+      "loss": 0.448,
+      "step": 1926
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0507573295109188e-05,
+      "loss": 0.4558,
+      "step": 1927
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0499190688265866e-05,
+      "loss": 0.479,
+      "step": 1928
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0490807729759654e-05,
+      "loss": 0.4473,
+      "step": 1929
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0482424425496061e-05,
+      "loss": 0.449,
+      "step": 1930
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.047404078138084e-05,
+      "loss": 0.467,
+      "step": 1931
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0465656803319984e-05,
+      "loss": 0.4624,
+      "step": 1932
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.045727249721972e-05,
+      "loss": 0.469,
+      "step": 1933
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0448887868986508e-05,
+      "loss": 0.4524,
+      "step": 1934
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0440502924527034e-05,
+      "loss": 0.4519,
+      "step": 1935
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0432117669748209e-05,
+      "loss": 0.4519,
+      "step": 1936
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0423732110557157e-05,
+      "loss": 0.4575,
+      "step": 1937
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.041534625286122e-05,
+      "loss": 0.4478,
+      "step": 1938
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.040696010256795e-05,
+      "loss": 0.4585,
+      "step": 1939
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0398573665585105e-05,
+      "loss": 0.4609,
+      "step": 1940
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0390186947820648e-05,
+      "loss": 0.4443,
+      "step": 1941
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0381799955182734e-05,
+      "loss": 0.4414,
+      "step": 1942
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0373412693579714e-05,
+      "loss": 0.429,
+      "step": 1943
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0365025168920133e-05,
+      "loss": 0.4348,
+      "step": 1944
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0356637387112713e-05,
+      "loss": 0.4502,
+      "step": 1945
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.034824935406636e-05,
+      "loss": 0.4861,
+      "step": 1946
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0339861075690167e-05,
+      "loss": 0.4524,
+      "step": 1947
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.0331472557893385e-05,
+      "loss": 0.458,
+      "step": 1948
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0323083806585436e-05,
+      "loss": 0.4424,
+      "step": 1949
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.031469482767592e-05,
+      "loss": 0.4636,
+      "step": 1950
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0306305627074581e-05,
+      "loss": 0.4622,
+      "step": 1951
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0297916210691328e-05,
+      "loss": 0.447,
+      "step": 1952
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0289526584436223e-05,
+      "loss": 0.4546,
+      "step": 1953
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0281136754219467e-05,
+      "loss": 0.4448,
+      "step": 1954
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0272746725951412e-05,
+      "loss": 0.4373,
+      "step": 1955
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0264356505542552e-05,
+      "loss": 0.4731,
+      "step": 1956
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0255966098903509e-05,
+      "loss": 0.4421,
+      "step": 1957
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0247575511945041e-05,
+      "loss": 0.4685,
+      "step": 1958
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.023918475057803e-05,
+      "loss": 0.4773,
+      "step": 1959
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0230793820713487e-05,
+      "loss": 0.4507,
+      "step": 1960
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0222402728262529e-05,
+      "loss": 0.4609,
+      "step": 1961
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0214011479136402e-05,
+      "loss": 0.4622,
+      "step": 1962
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0205620079246455e-05,
+      "loss": 0.4434,
+      "step": 1963
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0197228534504144e-05,
+      "loss": 0.439,
+      "step": 1964
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0188836850821026e-05,
+      "loss": 0.4541,
+      "step": 1965
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0180445034108761e-05,
+      "loss": 0.4675,
+      "step": 1966
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0172053090279094e-05,
+      "loss": 0.4517,
+      "step": 1967
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0163661025243868e-05,
+      "loss": 0.4343,
+      "step": 1968
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0155268844915006e-05,
+      "loss": 0.4414,
+      "step": 1969
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0146876555204515e-05,
+      "loss": 0.4612,
+      "step": 1970
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0138484162024478e-05,
+      "loss": 0.4773,
+      "step": 1971
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.013009167128705e-05,
+      "loss": 0.4753,
+      "step": 1972
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0121699088904457e-05,
+      "loss": 0.4561,
+      "step": 1973
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0113306420788987e-05,
+      "loss": 0.4534,
+      "step": 1974
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0104913672852994e-05,
+      "loss": 0.449,
+      "step": 1975
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.009652085100888e-05,
+      "loss": 0.4602,
+      "step": 1976
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.00881279611691e-05,
+      "loss": 0.4573,
+      "step": 1977
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0079735009246168e-05,
+      "loss": 0.4563,
+      "step": 1978
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0071342001152628e-05,
+      "loss": 0.4529,
+      "step": 1979
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0062948942801067e-05,
+      "loss": 0.468,
+      "step": 1980
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0054555840104118e-05,
+      "loss": 0.4604,
+      "step": 1981
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0046162698974434e-05,
+      "loss": 0.4568,
+      "step": 1982
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.003776952532469e-05,
+      "loss": 0.459,
+      "step": 1983
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0029376325067601e-05,
+      "loss": 0.4517,
+      "step": 1984
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0020983104115887e-05,
+      "loss": 0.4734,
+      "step": 1985
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.001258986838229e-05,
+      "loss": 0.4871,
+      "step": 1986
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.0004196623779556e-05,
+      "loss": 0.4451,
+      "step": 1987
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.995803376220444e-06,
+      "loss": 0.4712,
+      "step": 1988
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.987410131617713e-06,
+      "loss": 0.4546,
+      "step": 1989
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.979016895884113e-06,
+      "loss": 0.4583,
+      "step": 1990
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.970623674932402e-06,
+      "loss": 0.4512,
+      "step": 1991
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.962230474675313e-06,
+      "loss": 0.4666,
+      "step": 1992
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.95383730102557e-06,
+      "loss": 0.4617,
+      "step": 1993
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.945444159895885e-06,
+      "loss": 0.4644,
+      "step": 1994
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.937051057198932e-06,
+      "loss": 0.4399,
+      "step": 1995
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.928657998847375e-06,
+      "loss": 0.4568,
+      "step": 1996
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.920264990753837e-06,
+      "loss": 0.4622,
+      "step": 1997
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.911872038830902e-06,
+      "loss": 0.4521,
+      "step": 1998
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.903479148991124e-06,
+      "loss": 0.47,
+      "step": 1999
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.895086327147007e-06,
+      "loss": 0.4692,
+      "step": 2000
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.886693579211015e-06,
+      "loss": 0.4463,
+      "step": 2001
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.878300911095545e-06,
+      "loss": 0.4294,
+      "step": 2002
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.869908328712953e-06,
+      "loss": 0.468,
+      "step": 2003
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.861515837975529e-06,
+      "loss": 0.4719,
+      "step": 2004
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.85312344479549e-06,
+      "loss": 0.4512,
+      "step": 2005
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.844731155084998e-06,
+      "loss": 0.4414,
+      "step": 2006
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.836338974756134e-06,
+      "loss": 0.4519,
+      "step": 2007
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.827946909720907e-06,
+      "loss": 0.4619,
+      "step": 2008
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.819554965891244e-06,
+      "loss": 0.4651,
+      "step": 2009
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.811163149178976e-06,
+      "loss": 0.4592,
+      "step": 2010
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.802771465495858e-06,
+      "loss": 0.4819,
+      "step": 2011
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.794379920753545e-06,
+      "loss": 0.4495,
+      "step": 2012
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.7859885208636e-06,
+      "loss": 0.4548,
+      "step": 2013
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.777597271737474e-06,
+      "loss": 0.4839,
+      "step": 2014
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.769206179286517e-06,
+      "loss": 0.4387,
+      "step": 2015
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.760815249421973e-06,
+      "loss": 0.4478,
+      "step": 2016
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.75242448805496e-06,
+      "loss": 0.457,
+      "step": 2017
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.744033901096494e-06,
+      "loss": 0.4675,
+      "step": 2018
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.735643494457448e-06,
+      "loss": 0.4512,
+      "step": 2019
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.72725327404859e-06,
+      "loss": 0.4792,
+      "step": 2020
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.718863245780538e-06,
+      "loss": 0.4456,
+      "step": 2021
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.71047341556378e-06,
+      "loss": 0.4326,
+      "step": 2022
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.702083789308677e-06,
+      "loss": 0.429,
+      "step": 2023
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.693694372925422e-06,
+      "loss": 0.4629,
+      "step": 2024
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 9.685305172324083e-06,
+      "loss": 0.4683,
+      "step": 2025
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.676916193414569e-06,
+      "loss": 0.4612,
+      "step": 2026
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.668527442106619e-06,
+      "loss": 0.4502,
+      "step": 2027
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.660138924309838e-06,
+      "loss": 0.4619,
+      "step": 2028
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.65175064593364e-06,
+      "loss": 0.4646,
+      "step": 2029
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.64336261288729e-06,
+      "loss": 0.4734,
+      "step": 2030
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.634974831079868e-06,
+      "loss": 0.459,
+      "step": 2031
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.626587306420287e-06,
+      "loss": 0.4727,
+      "step": 2032
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.61820004481727e-06,
+      "loss": 0.4395,
+      "step": 2033
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.609813052179354e-06,
+      "loss": 0.4321,
+      "step": 2034
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.601426334414898e-06,
+      "loss": 0.4697,
+      "step": 2035
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.593039897432053e-06,
+      "loss": 0.4751,
+      "step": 2036
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.584653747138784e-06,
+      "loss": 0.4441,
+      "step": 2037
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.576267889442848e-06,
+      "loss": 0.4409,
+      "step": 2038
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.567882330251794e-06,
+      "loss": 0.4602,
+      "step": 2039
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.559497075472967e-06,
+      "loss": 0.4666,
+      "step": 2040
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.551112131013493e-06,
+      "loss": 0.4717,
+      "step": 2041
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.542727502780285e-06,
+      "loss": 0.4253,
+      "step": 2042
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.53434319668002e-06,
+      "loss": 0.4355,
+      "step": 2043
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.525959218619164e-06,
+      "loss": 0.4751,
+      "step": 2044
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.517575574503945e-06,
+      "loss": 0.4397,
+      "step": 2045
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.509192270240348e-06,
+      "loss": 0.4541,
+      "step": 2046
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.500809311734137e-06,
+      "loss": 0.4399,
+      "step": 2047
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.492426704890812e-06,
+      "loss": 0.4897,
+      "step": 2048
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.484044455615642e-06,
+      "loss": 0.4272,
+      "step": 2049
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.475662569813637e-06,
+      "loss": 0.4365,
+      "step": 2050
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.467281053389543e-06,
+      "loss": 0.4636,
+      "step": 2051
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.458899912247861e-06,
+      "loss": 0.4363,
+      "step": 2052
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.450519152292816e-06,
+      "loss": 0.4553,
+      "step": 2053
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.442138779428376e-06,
+      "loss": 0.4189,
+      "step": 2054
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.433758799558222e-06,
+      "loss": 0.4624,
+      "step": 2055
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.425379218585765e-06,
+      "loss": 0.4561,
+      "step": 2056
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.417000042414141e-06,
+      "loss": 0.4629,
+      "step": 2057
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.40862127694619e-06,
+      "loss": 0.4441,
+      "step": 2058
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.40024292808447e-06,
+      "loss": 0.4648,
+      "step": 2059
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.391865001731242e-06,
+      "loss": 0.4883,
+      "step": 2060
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.383487503788473e-06,
+      "loss": 0.4219,
+      "step": 2061
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.375110440157827e-06,
+      "loss": 0.4805,
+      "step": 2062
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.366733816740652e-06,
+      "loss": 0.4541,
+      "step": 2063
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 9.358357639438007e-06,
+      "loss": 0.4604,
+      "step": 2064
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.349981914150611e-06,
+      "loss": 0.438,
+      "step": 2065
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.341606646778895e-06,
+      "loss": 0.4714,
+      "step": 2066
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.333231843222938e-06,
+      "loss": 0.4736,
+      "step": 2067
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.324857509382504e-06,
+      "loss": 0.4539,
+      "step": 2068
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.316483651157037e-06,
+      "loss": 0.4712,
+      "step": 2069
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.308110274445625e-06,
+      "loss": 0.4524,
+      "step": 2070
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.299737385147033e-06,
+      "loss": 0.4521,
+      "step": 2071
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.291364989159679e-06,
+      "loss": 0.4614,
+      "step": 2072
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.282993092381626e-06,
+      "loss": 0.47,
+      "step": 2073
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.274621700710594e-06,
+      "loss": 0.4536,
+      "step": 2074
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.266250820043945e-06,
+      "loss": 0.4451,
+      "step": 2075
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.257880456278682e-06,
+      "loss": 0.4419,
+      "step": 2076
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.249510615311434e-06,
+      "loss": 0.4651,
+      "step": 2077
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.24114130303848e-06,
+      "loss": 0.4519,
+      "step": 2078
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.232772525355712e-06,
+      "loss": 0.4583,
+      "step": 2079
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.224404288158645e-06,
+      "loss": 0.4446,
+      "step": 2080
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.216036597342427e-06,
+      "loss": 0.4404,
+      "step": 2081
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.207669458801803e-06,
+      "loss": 0.4429,
+      "step": 2082
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.199302878431146e-06,
+      "loss": 0.4426,
+      "step": 2083
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.190936862124424e-06,
+      "loss": 0.4409,
+      "step": 2084
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.182571415775209e-06,
+      "loss": 0.4343,
+      "step": 2085
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.174206545276678e-06,
+      "loss": 0.4473,
+      "step": 2086
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.165842256521595e-06,
+      "loss": 0.4722,
+      "step": 2087
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.157478555402324e-06,
+      "loss": 0.4607,
+      "step": 2088
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.149115447810798e-06,
+      "loss": 0.4487,
+      "step": 2089
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.14075293963855e-06,
+      "loss": 0.4561,
+      "step": 2090
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.132391036776686e-06,
+      "loss": 0.437,
+      "step": 2091
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.124029745115872e-06,
+      "loss": 0.4624,
+      "step": 2092
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.11566907054636e-06,
+      "loss": 0.4377,
+      "step": 2093
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.10730901895796e-06,
+      "loss": 0.4731,
+      "step": 2094
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.09894959624005e-06,
+      "loss": 0.4397,
+      "step": 2095
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.090590808281554e-06,
+      "loss": 0.4575,
+      "step": 2096
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.082232660970953e-06,
+      "loss": 0.4482,
+      "step": 2097
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.073875160196283e-06,
+      "loss": 0.4553,
+      "step": 2098
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.065518311845114e-06,
+      "loss": 0.4485,
+      "step": 2099
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.057162121804567e-06,
+      "loss": 0.4646,
+      "step": 2100
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.048806595961294e-06,
+      "loss": 0.4639,
+      "step": 2101
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.040451740201473e-06,
+      "loss": 0.4629,
+      "step": 2102
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 9.032097560410821e-06,
+      "loss": 0.4302,
+      "step": 2103
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 9.023744062474574e-06,
+      "loss": 0.4443,
+      "step": 2104
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 9.015391252277488e-06,
+      "loss": 0.4641,
+      "step": 2105
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 9.00703913570383e-06,
+      "loss": 0.46,
+      "step": 2106
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.998687718637391e-06,
+      "loss": 0.5034,
+      "step": 2107
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.990337006961454e-06,
+      "loss": 0.4312,
+      "step": 2108
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.98198700655881e-06,
+      "loss": 0.4624,
+      "step": 2109
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.973637723311759e-06,
+      "loss": 0.4641,
+      "step": 2110
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.96528916310208e-06,
+      "loss": 0.4502,
+      "step": 2111
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.956941331811052e-06,
+      "loss": 0.4595,
+      "step": 2112
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.948594235319444e-06,
+      "loss": 0.4844,
+      "step": 2113
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.940247879507492e-06,
+      "loss": 0.4521,
+      "step": 2114
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.931902270254929e-06,
+      "loss": 0.4404,
+      "step": 2115
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.923557413440947e-06,
+      "loss": 0.4314,
+      "step": 2116
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.915213314944224e-06,
+      "loss": 0.4656,
+      "step": 2117
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.906869980642882e-06,
+      "loss": 0.4504,
+      "step": 2118
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.898527416414524e-06,
+      "loss": 0.4819,
+      "step": 2119
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.890185628136204e-06,
+      "loss": 0.4346,
+      "step": 2120
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.881844621684423e-06,
+      "loss": 0.4397,
+      "step": 2121
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.873504402935142e-06,
+      "loss": 0.4624,
+      "step": 2122
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.865164977763756e-06,
+      "loss": 0.468,
+      "step": 2123
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.856826352045115e-06,
+      "loss": 0.4707,
+      "step": 2124
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.848488531653494e-06,
+      "loss": 0.4766,
+      "step": 2125
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.840151522462602e-06,
+      "loss": 0.4578,
+      "step": 2126
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.83181533034558e-06,
+      "loss": 0.4431,
+      "step": 2127
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.823479961174993e-06,
+      "loss": 0.4495,
+      "step": 2128
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.81514542082283e-06,
+      "loss": 0.4463,
+      "step": 2129
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.806811715160485e-06,
+      "loss": 0.4583,
+      "step": 2130
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.79847885005877e-06,
+      "loss": 0.4414,
+      "step": 2131
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.790146831387913e-06,
+      "loss": 0.4573,
+      "step": 2132
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.781815665017531e-06,
+      "loss": 0.4529,
+      "step": 2133
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.773485356816651e-06,
+      "loss": 0.4417,
+      "step": 2134
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.76515591265369e-06,
+      "loss": 0.4392,
+      "step": 2135
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.756827338396465e-06,
+      "loss": 0.468,
+      "step": 2136
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.748499639912168e-06,
+      "loss": 0.4587,
+      "step": 2137
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.740172823067378e-06,
+      "loss": 0.4556,
+      "step": 2138
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.731846893728062e-06,
+      "loss": 0.4465,
+      "step": 2139
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.723521857759547e-06,
+      "loss": 0.4514,
+      "step": 2140
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 8.715197721026543e-06,
+      "loss": 0.4399,
+      "step": 2141
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.706874489393123e-06,
+      "loss": 0.4087,
+      "step": 2142
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.69855216872271e-06,
+      "loss": 0.4756,
+      "step": 2143
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.690230764878114e-06,
+      "loss": 0.4502,
+      "step": 2144
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.681910283721465e-06,
+      "loss": 0.4363,
+      "step": 2145
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.67359073111427e-06,
+      "loss": 0.4521,
+      "step": 2146
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.665272112917366e-06,
+      "loss": 0.4626,
+      "step": 2147
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.656954434990942e-06,
+      "loss": 0.4751,
+      "step": 2148
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.648637703194515e-06,
+      "loss": 0.4707,
+      "step": 2149
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.64032192338694e-06,
+      "loss": 0.4578,
+      "step": 2150
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.632007101426409e-06,
+      "loss": 0.4678,
+      "step": 2151
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.62369324317042e-06,
+      "loss": 0.4653,
+      "step": 2152
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.615380354475812e-06,
+      "loss": 0.4407,
+      "step": 2153
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.607068441198734e-06,
+      "loss": 0.4385,
+      "step": 2154
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.598757509194638e-06,
+      "loss": 0.4343,
+      "step": 2155
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.5904475643183e-06,
+      "loss": 0.4331,
+      "step": 2156
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.582138612423794e-06,
+      "loss": 0.437,
+      "step": 2157
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.573830659364495e-06,
+      "loss": 0.4819,
+      "step": 2158
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.565523710993072e-06,
+      "loss": 0.4492,
+      "step": 2159
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.557217773161484e-06,
+      "loss": 0.467,
+      "step": 2160
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.54891285172099e-06,
+      "loss": 0.4683,
+      "step": 2161
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.540608952522116e-06,
+      "loss": 0.4465,
+      "step": 2162
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.53230608141469e-06,
+      "loss": 0.4402,
+      "step": 2163
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.524004244247792e-06,
+      "loss": 0.4189,
+      "step": 2164
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.515703446869789e-06,
+      "loss": 0.4521,
+      "step": 2165
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.50740369512831e-06,
+      "loss": 0.447,
+      "step": 2166
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.499104994870247e-06,
+      "loss": 0.4502,
+      "step": 2167
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.490807351941753e-06,
+      "loss": 0.4475,
+      "step": 2168
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.482510772188235e-06,
+      "loss": 0.4595,
+      "step": 2169
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.474215261454356e-06,
+      "loss": 0.4507,
+      "step": 2170
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.46592082558402e-06,
+      "loss": 0.4448,
+      "step": 2171
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.457627470420368e-06,
+      "loss": 0.4714,
+      "step": 2172
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.449335201805801e-06,
+      "loss": 0.4531,
+      "step": 2173
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.441044025581926e-06,
+      "loss": 0.448,
+      "step": 2174
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.432753947589605e-06,
+      "loss": 0.4324,
+      "step": 2175
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.424464973668912e-06,
+      "loss": 0.4485,
+      "step": 2176
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.416177109659153e-06,
+      "loss": 0.4497,
+      "step": 2177
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.40789036139884e-06,
+      "loss": 0.4263,
+      "step": 2178
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.399604734725707e-06,
+      "loss": 0.4841,
+      "step": 2179
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 8.391320235476703e-06,
+      "loss": 0.4229,
+      "step": 2180
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.383036869487965e-06,
+      "loss": 0.4497,
+      "step": 2181
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.374754642594856e-06,
+      "loss": 0.4348,
+      "step": 2182
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.366473560631917e-06,
+      "loss": 0.4329,
+      "step": 2183
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.358193629432885e-06,
+      "loss": 0.4199,
+      "step": 2184
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.349914854830701e-06,
+      "loss": 0.458,
+      "step": 2185
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.34163724265747e-06,
+      "loss": 0.4399,
+      "step": 2186
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.333360798744496e-06,
+      "loss": 0.4204,
+      "step": 2187
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.325085528922251e-06,
+      "loss": 0.436,
+      "step": 2188
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.316811439020378e-06,
+      "loss": 0.4524,
+      "step": 2189
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.308538534867695e-06,
+      "loss": 0.4678,
+      "step": 2190
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.300266822292182e-06,
+      "loss": 0.4553,
+      "step": 2191
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.291996307120982e-06,
+      "loss": 0.4741,
+      "step": 2192
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.283726995180388e-06,
+      "loss": 0.4219,
+      "step": 2193
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.275458892295851e-06,
+      "loss": 0.4575,
+      "step": 2194
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.267192004291973e-06,
+      "loss": 0.4619,
+      "step": 2195
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.258926336992486e-06,
+      "loss": 0.4543,
+      "step": 2196
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.250661896220281e-06,
+      "loss": 0.4346,
+      "step": 2197
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.242398687797372e-06,
+      "loss": 0.4585,
+      "step": 2198
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.234136717544915e-06,
+      "loss": 0.4255,
+      "step": 2199
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.225875991283177e-06,
+      "loss": 0.4641,
+      "step": 2200
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.217616514831563e-06,
+      "loss": 0.4639,
+      "step": 2201
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.209358294008597e-06,
+      "loss": 0.458,
+      "step": 2202
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.201101334631912e-06,
+      "loss": 0.4338,
+      "step": 2203
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.192845642518259e-06,
+      "loss": 0.4438,
+      "step": 2204
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.184591223483489e-06,
+      "loss": 0.4153,
+      "step": 2205
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.17633808334256e-06,
+      "loss": 0.4519,
+      "step": 2206
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.168086227909536e-06,
+      "loss": 0.4568,
+      "step": 2207
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.159835662997559e-06,
+      "loss": 0.4531,
+      "step": 2208
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.151586394418876e-06,
+      "loss": 0.4465,
+      "step": 2209
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.143338427984817e-06,
+      "loss": 0.4524,
+      "step": 2210
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.135091769505798e-06,
+      "loss": 0.4453,
+      "step": 2211
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.126846424791303e-06,
+      "loss": 0.4219,
+      "step": 2212
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.1186023996499e-06,
+      "loss": 0.4368,
+      "step": 2213
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.110359699889227e-06,
+      "loss": 0.4668,
+      "step": 2214
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.102118331315983e-06,
+      "loss": 0.4377,
+      "step": 2215
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.093878299735933e-06,
+      "loss": 0.4795,
+      "step": 2216
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.085639610953905e-06,
+      "loss": 0.4358,
+      "step": 2217
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 8.077402270773764e-06,
+      "loss": 0.4497,
+      "step": 2218
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 8.069166284998446e-06,
+      "loss": 0.4568,
+      "step": 2219
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 8.060931659429919e-06,
+      "loss": 0.4402,
+      "step": 2220
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 8.0526983998692e-06,
+      "loss": 0.4341,
+      "step": 2221
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 8.044466512116338e-06,
+      "loss": 0.4504,
+      "step": 2222
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 8.036236001970425e-06,
+      "loss": 0.4568,
+      "step": 2223
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 8.028006875229566e-06,
+      "loss": 0.4346,
+      "step": 2224
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 8.019779137690906e-06,
+      "loss": 0.4546,
+      "step": 2225
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 8.01155279515061e-06,
+      "loss": 0.446,
+      "step": 2226
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 8.00332785340385e-06,
+      "loss": 0.478,
+      "step": 2227
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.995104318244824e-06,
+      "loss": 0.4409,
+      "step": 2228
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.986882195466734e-06,
+      "loss": 0.4641,
+      "step": 2229
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.978661490861778e-06,
+      "loss": 0.4478,
+      "step": 2230
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.970442210221173e-06,
+      "loss": 0.4631,
+      "step": 2231
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.962224359335117e-06,
+      "loss": 0.4619,
+      "step": 2232
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.954007943992812e-06,
+      "loss": 0.4446,
+      "step": 2233
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.945792969982437e-06,
+      "loss": 0.4529,
+      "step": 2234
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.937579443091167e-06,
+      "loss": 0.4636,
+      "step": 2235
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.929367369105154e-06,
+      "loss": 0.4424,
+      "step": 2236
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.92115675380952e-06,
+      "loss": 0.4351,
+      "step": 2237
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.912947602988368e-06,
+      "loss": 0.438,
+      "step": 2238
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.904739922424764e-06,
+      "loss": 0.4346,
+      "step": 2239
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.896533717900745e-06,
+      "loss": 0.4595,
+      "step": 2240
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.888328995197303e-06,
+      "loss": 0.4517,
+      "step": 2241
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.880125760094378e-06,
+      "loss": 0.4226,
+      "step": 2242
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.87192401837088e-06,
+      "loss": 0.4531,
+      "step": 2243
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.863723775804651e-06,
+      "loss": 0.449,
+      "step": 2244
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.855525038172491e-06,
+      "loss": 0.4275,
+      "step": 2245
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.847327811250128e-06,
+      "loss": 0.4409,
+      "step": 2246
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.839132100812225e-06,
+      "loss": 0.4443,
+      "step": 2247
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.83093791263239e-06,
+      "loss": 0.4575,
+      "step": 2248
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.822745252483146e-06,
+      "loss": 0.4565,
+      "step": 2249
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.814554126135944e-06,
+      "loss": 0.4495,
+      "step": 2250
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.806364539361155e-06,
+      "loss": 0.4502,
+      "step": 2251
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.79817649792807e-06,
+      "loss": 0.4482,
+      "step": 2252
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.789990007604878e-06,
+      "loss": 0.4456,
+      "step": 2253
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.781805074158684e-06,
+      "loss": 0.4377,
+      "step": 2254
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.773621703355502e-06,
+      "loss": 0.4612,
+      "step": 2255
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.765439900960234e-06,
+      "loss": 0.4326,
+      "step": 2256
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 7.757259672736683e-06,
+      "loss": 0.4651,
+      "step": 2257
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.749081024447546e-06,
+      "loss": 0.4324,
+      "step": 2258
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.740903961854391e-06,
+      "loss": 0.481,
+      "step": 2259
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.732728490717697e-06,
+      "loss": 0.4668,
+      "step": 2260
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.724554616796793e-06,
+      "loss": 0.4316,
+      "step": 2261
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.716382345849904e-06,
+      "loss": 0.4614,
+      "step": 2262
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.708211683634112e-06,
+      "loss": 0.4595,
+      "step": 2263
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.700042635905376e-06,
+      "loss": 0.4043,
+      "step": 2264
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.69187520841851e-06,
+      "loss": 0.448,
+      "step": 2265
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.683709406927185e-06,
+      "loss": 0.4263,
+      "step": 2266
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.67554523718394e-06,
+      "loss": 0.4434,
+      "step": 2267
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.667382704940148e-06,
+      "loss": 0.4377,
+      "step": 2268
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.659221815946039e-06,
+      "loss": 0.4465,
+      "step": 2269
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.651062575950682e-06,
+      "loss": 0.4419,
+      "step": 2270
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.642904990701982e-06,
+      "loss": 0.4458,
+      "step": 2271
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.634749065946686e-06,
+      "loss": 0.4575,
+      "step": 2272
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.62659480743036e-06,
+      "loss": 0.4607,
+      "step": 2273
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.618442220897409e-06,
+      "loss": 0.436,
+      "step": 2274
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.610291312091048e-06,
+      "loss": 0.4558,
+      "step": 2275
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.6021420867533155e-06,
+      "loss": 0.4355,
+      "step": 2276
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.59399455062507e-06,
+      "loss": 0.4539,
+      "step": 2277
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.585848709445962e-06,
+      "loss": 0.4475,
+      "step": 2278
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.577704568954476e-06,
+      "loss": 0.4299,
+      "step": 2279
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.569562134887869e-06,
+      "loss": 0.4531,
+      "step": 2280
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.561421412982217e-06,
+      "loss": 0.4631,
+      "step": 2281
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.553282408972382e-06,
+      "loss": 0.4438,
+      "step": 2282
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.545145128592009e-06,
+      "loss": 0.4263,
+      "step": 2283
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.537009577573544e-06,
+      "loss": 0.4871,
+      "step": 2284
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.528875761648199e-06,
+      "loss": 0.429,
+      "step": 2285
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.520743686545981e-06,
+      "loss": 0.4316,
+      "step": 2286
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.512613357995652e-06,
+      "loss": 0.4412,
+      "step": 2287
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.504484781724752e-06,
+      "loss": 0.4319,
+      "step": 2288
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.496357963459595e-06,
+      "loss": 0.447,
+      "step": 2289
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.488232908925241e-06,
+      "loss": 0.4519,
+      "step": 2290
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.48010962384552e-06,
+      "loss": 0.4465,
+      "step": 2291
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.471988113943005e-06,
+      "loss": 0.4553,
+      "step": 2292
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.463868384939031e-06,
+      "loss": 0.4312,
+      "step": 2293
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.455750442553667e-06,
+      "loss": 0.4438,
+      "step": 2294
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.447634292505725e-06,
+      "loss": 0.4678,
+      "step": 2295
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 7.439519940512766e-06,
+      "loss": 0.4517,
+      "step": 2296
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.431407392291064e-06,
+      "loss": 0.4551,
+      "step": 2297
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.423296653555645e-06,
+      "loss": 0.438,
+      "step": 2298
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.415187730020242e-06,
+      "loss": 0.4629,
+      "step": 2299
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.407080627397316e-06,
+      "loss": 0.4644,
+      "step": 2300
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.398975351398054e-06,
+      "loss": 0.4185,
+      "step": 2301
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.390871907732335e-06,
+      "loss": 0.4819,
+      "step": 2302
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.382770302108768e-06,
+      "loss": 0.4546,
+      "step": 2303
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.3746705402346605e-06,
+      "loss": 0.4353,
+      "step": 2304
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.3665726278160136e-06,
+      "loss": 0.4536,
+      "step": 2305
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.358476570557536e-06,
+      "loss": 0.4426,
+      "step": 2306
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.350382374162622e-06,
+      "loss": 0.4424,
+      "step": 2307
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.342290044333362e-06,
+      "loss": 0.4421,
+      "step": 2308
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.334199586770524e-06,
+      "loss": 0.4436,
+      "step": 2309
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.3261110071735645e-06,
+      "loss": 0.4363,
+      "step": 2310
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.3180243112406106e-06,
+      "loss": 0.4495,
+      "step": 2311
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.3099395046684626e-06,
+      "loss": 0.4353,
+      "step": 2312
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.301856593152596e-06,
+      "loss": 0.4419,
+      "step": 2313
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.293775582387146e-06,
+      "loss": 0.4609,
+      "step": 2314
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.285696478064915e-06,
+      "loss": 0.4592,
+      "step": 2315
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.277619285877351e-06,
+      "loss": 0.4275,
+      "step": 2316
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.26954401151456e-06,
+      "loss": 0.4285,
+      "step": 2317
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.261470660665306e-06,
+      "loss": 0.4414,
+      "step": 2318
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.253399239016986e-06,
+      "loss": 0.425,
+      "step": 2319
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.2453297522556475e-06,
+      "loss": 0.4512,
+      "step": 2320
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.237262206065962e-06,
+      "loss": 0.4531,
+      "step": 2321
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.229196606131248e-06,
+      "loss": 0.437,
+      "step": 2322
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.221132958133446e-06,
+      "loss": 0.4148,
+      "step": 2323
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.213071267753118e-06,
+      "loss": 0.4365,
+      "step": 2324
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.2050115406694554e-06,
+      "loss": 0.4612,
+      "step": 2325
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.19695378256026e-06,
+      "loss": 0.4365,
+      "step": 2326
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.188897999101954e-06,
+      "loss": 0.4417,
+      "step": 2327
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.1808441959695565e-06,
+      "loss": 0.4441,
+      "step": 2328
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.1727923788366986e-06,
+      "loss": 0.4485,
+      "step": 2329
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.164742553375618e-06,
+      "loss": 0.4363,
+      "step": 2330
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.156694725257137e-06,
+      "loss": 0.4275,
+      "step": 2331
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.148648900150682e-06,
+      "loss": 0.4358,
+      "step": 2332
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.140605083724264e-06,
+      "loss": 0.4377,
+      "step": 2333
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 7.132563281644471e-06,
+      "loss": 0.4443,
+      "step": 2334
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 7.124523499576488e-06,
+      "loss": 0.4348,
+      "step": 2335
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 7.116485743184063e-06,
+      "loss": 0.4395,
+      "step": 2336
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 7.108450018129527e-06,
+      "loss": 0.4641,
+      "step": 2337
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 7.100416330073772e-06,
+      "loss": 0.4543,
+      "step": 2338
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 7.092384684676263e-06,
+      "loss": 0.4565,
+      "step": 2339
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 7.084355087595015e-06,
+      "loss": 0.4424,
+      "step": 2340
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 7.076327544486607e-06,
+      "loss": 0.4539,
+      "step": 2341
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 7.06830206100618e-06,
+      "loss": 0.4316,
+      "step": 2342
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 7.060278642807402e-06,
+      "loss": 0.4565,
+      "step": 2343
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 7.052257295542508e-06,
+      "loss": 0.4302,
+      "step": 2344
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 7.044238024862262e-06,
+      "loss": 0.4438,
+      "step": 2345
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 7.036220836415964e-06,
+      "loss": 0.437,
+      "step": 2346
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 7.0282057358514545e-06,
+      "loss": 0.4458,
+      "step": 2347
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 7.020192728815098e-06,
+      "loss": 0.458,
+      "step": 2348
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 7.012181820951793e-06,
+      "loss": 0.4458,
+      "step": 2349
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 7.004173017904941e-06,
+      "loss": 0.4348,
+      "step": 2350
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 6.99616632531648e-06,
+      "loss": 0.4233,
+      "step": 2351
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 6.9881617488268516e-06,
+      "loss": 0.4543,
+      "step": 2352
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 6.980159294075004e-06,
+      "loss": 0.4402,
+      "step": 2353
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 6.972158966698399e-06,
+      "loss": 0.4609,
+      "step": 2354
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 6.96416077233299e-06,
+      "loss": 0.4209,
+      "step": 2355
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 6.956164716613243e-06,
+      "loss": 0.4375,
+      "step": 2356
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 6.948170805172103e-06,
+      "loss": 0.4309,
+      "step": 2357
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 6.940179043641005e-06,
+      "loss": 0.4648,
+      "step": 2358
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 6.932189437649879e-06,
+      "loss": 0.4326,
+      "step": 2359
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 6.924201992827129e-06,
+      "loss": 0.4453,
+      "step": 2360
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 6.916216714799641e-06,
+      "loss": 0.4548,
+      "step": 2361
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 6.908233609192769e-06,
+      "loss": 0.4246,
+      "step": 2362
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 6.900252681630339e-06,
+      "loss": 0.4497,
+      "step": 2363
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 6.892273937734649e-06,
+      "loss": 0.4407,
+      "step": 2364
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 6.884297383126446e-06,
+      "loss": 0.4385,
+      "step": 2365
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 6.876323023424949e-06,
+      "loss": 0.4241,
+      "step": 2366
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 6.868350864247817e-06,
+      "loss": 0.4175,
+      "step": 2367
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 6.860380911211173e-06,
+      "loss": 0.4478,
+      "step": 2368
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 6.8524131699295725e-06,
+      "loss": 0.4377,
+      "step": 2369
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 6.844447646016017e-06,
+      "loss": 0.4473,
+      "step": 2370
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 6.836484345081956e-06,
+      "loss": 0.4326,
+      "step": 2371
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 6.8285232727372554e-06,
+      "loss": 0.4529,
+      "step": 2372
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 6.820564434590227e-06,
+      "loss": 0.4224,
+      "step": 2373
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.8126078362476e-06,
+      "loss": 0.4395,
+      "step": 2374
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.804653483314523e-06,
+      "loss": 0.4187,
+      "step": 2375
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.796701381394577e-06,
+      "loss": 0.4338,
+      "step": 2376
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.788751536089739e-06,
+      "loss": 0.4158,
+      "step": 2377
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.78080395300041e-06,
+      "loss": 0.4409,
+      "step": 2378
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.77285863772539e-06,
+      "loss": 0.4348,
+      "step": 2379
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.764915595861888e-06,
+      "loss": 0.438,
+      "step": 2380
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.756974833005503e-06,
+      "loss": 0.4519,
+      "step": 2381
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.749036354750231e-06,
+      "loss": 0.4421,
+      "step": 2382
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.74110016668847e-06,
+      "loss": 0.408,
+      "step": 2383
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.7331662744109825e-06,
+      "loss": 0.4434,
+      "step": 2384
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.725234683506934e-06,
+      "loss": 0.4529,
+      "step": 2385
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.717305399563863e-06,
+      "loss": 0.4446,
+      "step": 2386
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.7093784281676725e-06,
+      "loss": 0.4224,
+      "step": 2387
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.7014537749026495e-06,
+      "loss": 0.4387,
+      "step": 2388
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.6935314453514424e-06,
+      "loss": 0.4434,
+      "step": 2389
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.6856114450950675e-06,
+      "loss": 0.4351,
+      "step": 2390
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.67769377971289e-06,
+      "loss": 0.457,
+      "step": 2391
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.6697784547826365e-06,
+      "loss": 0.4414,
+      "step": 2392
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.6618654758803916e-06,
+      "loss": 0.4333,
+      "step": 2393
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.65395484858057e-06,
+      "loss": 0.4395,
+      "step": 2394
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.646046578455952e-06,
+      "loss": 0.4595,
+      "step": 2395
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.638140671077633e-06,
+      "loss": 0.4331,
+      "step": 2396
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.6302371320150646e-06,
+      "loss": 0.4446,
+      "step": 2397
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.622335966836022e-06,
+      "loss": 0.4382,
+      "step": 2398
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.6144371811066015e-06,
+      "loss": 0.4377,
+      "step": 2399
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.6065407803912355e-06,
+      "loss": 0.4314,
+      "step": 2400
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.598646770252664e-06,
+      "loss": 0.4373,
+      "step": 2401
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.590755156251958e-06,
+      "loss": 0.4619,
+      "step": 2402
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.5828659439484845e-06,
+      "loss": 0.4434,
+      "step": 2403
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.574979138899926e-06,
+      "loss": 0.4397,
+      "step": 2404
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.567094746662275e-06,
+      "loss": 0.459,
+      "step": 2405
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.559212772789808e-06,
+      "loss": 0.4475,
+      "step": 2406
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.551333222835117e-06,
+      "loss": 0.4609,
+      "step": 2407
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.543456102349072e-06,
+      "loss": 0.4351,
+      "step": 2408
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.535581416880843e-06,
+      "loss": 0.4397,
+      "step": 2409
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.527709171977872e-06,
+      "loss": 0.4507,
+      "step": 2410
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 6.519839373185889e-06,
+      "loss": 0.4343,
+      "step": 2411
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.5119720260489064e-06,
+      "loss": 0.4448,
+      "step": 2412
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.5041071361091925e-06,
+      "loss": 0.4417,
+      "step": 2413
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.496244708907307e-06,
+      "loss": 0.4561,
+      "step": 2414
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.488384749982054e-06,
+      "loss": 0.4373,
+      "step": 2415
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.48052726487051e-06,
+      "loss": 0.4253,
+      "step": 2416
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.472672259108011e-06,
+      "loss": 0.4006,
+      "step": 2417
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.4648197382281345e-06,
+      "loss": 0.4397,
+      "step": 2418
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.456969707762721e-06,
+      "loss": 0.4456,
+      "step": 2419
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.449122173241849e-06,
+      "loss": 0.4336,
+      "step": 2420
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.441277140193836e-06,
+      "loss": 0.4377,
+      "step": 2421
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.433434614145244e-06,
+      "loss": 0.4329,
+      "step": 2422
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.425594600620868e-06,
+      "loss": 0.4414,
+      "step": 2423
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.417757105143732e-06,
+      "loss": 0.4573,
+      "step": 2424
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.40992213323508e-06,
+      "loss": 0.4697,
+      "step": 2425
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.40208969041439e-06,
+      "loss": 0.4053,
+      "step": 2426
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.39425978219935e-06,
+      "loss": 0.4539,
+      "step": 2427
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.386432414105861e-06,
+      "loss": 0.4421,
+      "step": 2428
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.378607591648042e-06,
+      "loss": 0.4495,
+      "step": 2429
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.370785320338211e-06,
+      "loss": 0.4294,
+      "step": 2430
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.3629656056869e-06,
+      "loss": 0.415,
+      "step": 2431
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.355148453202826e-06,
+      "loss": 0.4294,
+      "step": 2432
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.347333868392906e-06,
+      "loss": 0.4644,
+      "step": 2433
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.339521856762254e-06,
+      "loss": 0.4211,
+      "step": 2434
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.331712423814164e-06,
+      "loss": 0.4509,
+      "step": 2435
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.323905575050124e-06,
+      "loss": 0.4482,
+      "step": 2436
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.316101315969783e-06,
+      "loss": 0.428,
+      "step": 2437
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.308299652070982e-06,
+      "loss": 0.4446,
+      "step": 2438
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.300500588849729e-06,
+      "loss": 0.4434,
+      "step": 2439
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.292704131800196e-06,
+      "loss": 0.4541,
+      "step": 2440
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.284910286414726e-06,
+      "loss": 0.4121,
+      "step": 2441
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.277119058183814e-06,
+      "loss": 0.4395,
+      "step": 2442
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.269330452596123e-06,
+      "loss": 0.4177,
+      "step": 2443
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.261544475138455e-06,
+      "loss": 0.4285,
+      "step": 2444
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.2537611312957635e-06,
+      "loss": 0.4475,
+      "step": 2445
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.245980426551162e-06,
+      "loss": 0.4578,
+      "step": 2446
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.2382023663858796e-06,
+      "loss": 0.4143,
+      "step": 2447
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.230426956279305e-06,
+      "loss": 0.428,
+      "step": 2448
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.222654201708949e-06,
+      "loss": 0.4265,
+      "step": 2449
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 6.2148841081504465e-06,
+      "loss": 0.4182,
+      "step": 2450
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.207116681077572e-06,
+      "loss": 0.4763,
+      "step": 2451
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.199351925962209e-06,
+      "loss": 0.4441,
+      "step": 2452
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.191589848274369e-06,
+      "loss": 0.4529,
+      "step": 2453
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.183830453482168e-06,
+      "loss": 0.4316,
+      "step": 2454
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.176073747051842e-06,
+      "loss": 0.437,
+      "step": 2455
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.168319734447721e-06,
+      "loss": 0.448,
+      "step": 2456
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.1605684211322435e-06,
+      "loss": 0.4541,
+      "step": 2457
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.152819812565954e-06,
+      "loss": 0.4465,
+      "step": 2458
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.145073914207473e-06,
+      "loss": 0.4546,
+      "step": 2459
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.137330731513531e-06,
+      "loss": 0.4358,
+      "step": 2460
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.1295902699389385e-06,
+      "loss": 0.4316,
+      "step": 2461
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.12185253493658e-06,
+      "loss": 0.4321,
+      "step": 2462
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.114117531957433e-06,
+      "loss": 0.4368,
+      "step": 2463
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.106385266450542e-06,
+      "loss": 0.4531,
+      "step": 2464
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.09865574386303e-06,
+      "loss": 0.4385,
+      "step": 2465
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.090928969640076e-06,
+      "loss": 0.4434,
+      "step": 2466
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.083204949224936e-06,
+      "loss": 0.4468,
+      "step": 2467
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.0754836880589186e-06,
+      "loss": 0.4399,
+      "step": 2468
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.0677651915813875e-06,
+      "loss": 0.4397,
+      "step": 2469
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.060049465229765e-06,
+      "loss": 0.425,
+      "step": 2470
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.052336514439513e-06,
+      "loss": 0.4231,
+      "step": 2471
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.044626344644151e-06,
+      "loss": 0.426,
+      "step": 2472
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.036918961275229e-06,
+      "loss": 0.4382,
+      "step": 2473
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.029214369762332e-06,
+      "loss": 0.4241,
+      "step": 2474
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.021512575533088e-06,
+      "loss": 0.4309,
+      "step": 2475
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.0138135840131486e-06,
+      "loss": 0.4351,
+      "step": 2476
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 6.006117400626195e-06,
+      "loss": 0.4248,
+      "step": 2477
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 5.998424030793922e-06,
+      "loss": 0.4136,
+      "step": 2478
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 5.990733479936046e-06,
+      "loss": 0.438,
+      "step": 2479
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 5.983045753470308e-06,
+      "loss": 0.4414,
+      "step": 2480
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 5.97536085681244e-06,
+      "loss": 0.4287,
+      "step": 2481
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 5.967678795376197e-06,
+      "loss": 0.4395,
+      "step": 2482
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 5.959999574573328e-06,
+      "loss": 0.4641,
+      "step": 2483
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 5.952323199813587e-06,
+      "loss": 0.459,
+      "step": 2484
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 5.944649676504715e-06,
+      "loss": 0.4304,
+      "step": 2485
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 5.936979010052446e-06,
+      "loss": 0.4329,
+      "step": 2486
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 5.929311205860515e-06,
+      "loss": 0.4397,
+      "step": 2487
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 5.921646269330617e-06,
+      "loss": 0.4353,
+      "step": 2488
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 5.913984205862449e-06,
+      "loss": 0.4546,
+      "step": 2489
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.906325020853671e-06,
+      "loss": 0.4414,
+      "step": 2490
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.898668719699914e-06,
+      "loss": 0.4236,
+      "step": 2491
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.891015307794792e-06,
+      "loss": 0.4617,
+      "step": 2492
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.8833647905298655e-06,
+      "loss": 0.4343,
+      "step": 2493
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.875717173294668e-06,
+      "loss": 0.4546,
+      "step": 2494
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.868072461476686e-06,
+      "loss": 0.4282,
+      "step": 2495
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.860430660461361e-06,
+      "loss": 0.4116,
+      "step": 2496
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.8527917756320785e-06,
+      "loss": 0.4197,
+      "step": 2497
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.8451558123701735e-06,
+      "loss": 0.438,
+      "step": 2498
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.837522776054929e-06,
+      "loss": 0.4448,
+      "step": 2499
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.829892672063552e-06,
+      "loss": 0.4272,
+      "step": 2500
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.822265505771199e-06,
+      "loss": 0.4451,
+      "step": 2501
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.814641282550941e-06,
+      "loss": 0.4124,
+      "step": 2502
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.807020007773789e-06,
+      "loss": 0.415,
+      "step": 2503
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.799401686808678e-06,
+      "loss": 0.4297,
+      "step": 2504
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.791786325022446e-06,
+      "loss": 0.4341,
+      "step": 2505
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.784173927779866e-06,
+      "loss": 0.4348,
+      "step": 2506
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.77656450044361e-06,
+      "loss": 0.4456,
+      "step": 2507
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.768958048374256e-06,
+      "loss": 0.4138,
+      "step": 2508
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.761354576930296e-06,
+      "loss": 0.4297,
+      "step": 2509
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.753754091468116e-06,
+      "loss": 0.4417,
+      "step": 2510
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.746156597342005e-06,
+      "loss": 0.4517,
+      "step": 2511
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.7385620999041325e-06,
+      "loss": 0.4446,
+      "step": 2512
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.730970604504571e-06,
+      "loss": 0.4253,
+      "step": 2513
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.723382116491264e-06,
+      "loss": 0.4392,
+      "step": 2514
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.715796641210048e-06,
+      "loss": 0.4216,
+      "step": 2515
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.708214184004634e-06,
+      "loss": 0.4316,
+      "step": 2516
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.7006347502166e-06,
+      "loss": 0.4419,
+      "step": 2517
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.693058345185407e-06,
+      "loss": 0.4204,
+      "step": 2518
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.685484974248372e-06,
+      "loss": 0.4341,
+      "step": 2519
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.677914642740674e-06,
+      "loss": 0.4507,
+      "step": 2520
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.670347355995356e-06,
+      "loss": 0.4458,
+      "step": 2521
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.662783119343318e-06,
+      "loss": 0.4421,
+      "step": 2522
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.655221938113309e-06,
+      "loss": 0.4556,
+      "step": 2523
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.647663817631917e-06,
+      "loss": 0.4424,
+      "step": 2524
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.6401087632235906e-06,
+      "loss": 0.4268,
+      "step": 2525
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.632556780210604e-06,
+      "loss": 0.4417,
+      "step": 2526
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 5.625007873913069e-06,
+      "loss": 0.4182,
+      "step": 2527
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.617462049648938e-06,
+      "loss": 0.4229,
+      "step": 2528
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.609919312733987e-06,
+      "loss": 0.4236,
+      "step": 2529
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.602379668481823e-06,
+      "loss": 0.4553,
+      "step": 2530
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.594843122203867e-06,
+      "loss": 0.446,
+      "step": 2531
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.587309679209353e-06,
+      "loss": 0.4373,
+      "step": 2532
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.57977934480534e-06,
+      "loss": 0.4475,
+      "step": 2533
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.572252124296696e-06,
+      "loss": 0.4458,
+      "step": 2534
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.5647280229860925e-06,
+      "loss": 0.4434,
+      "step": 2535
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.557207046174002e-06,
+      "loss": 0.4531,
+      "step": 2536
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.549689199158692e-06,
+      "loss": 0.4451,
+      "step": 2537
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.542174487236238e-06,
+      "loss": 0.4604,
+      "step": 2538
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.534662915700494e-06,
+      "loss": 0.4368,
+      "step": 2539
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.527154489843108e-06,
+      "loss": 0.4136,
+      "step": 2540
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.519649214953513e-06,
+      "loss": 0.4446,
+      "step": 2541
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.512147096318924e-06,
+      "loss": 0.4609,
+      "step": 2542
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.504648139224325e-06,
+      "loss": 0.4412,
+      "step": 2543
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.49715234895247e-06,
+      "loss": 0.4441,
+      "step": 2544
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.489659730783901e-06,
+      "loss": 0.4207,
+      "step": 2545
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.482170289996902e-06,
+      "loss": 0.4709,
+      "step": 2546
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.474684031867535e-06,
+      "loss": 0.4663,
+      "step": 2547
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.467200961669619e-06,
+      "loss": 0.4199,
+      "step": 2548
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.459721084674714e-06,
+      "loss": 0.4236,
+      "step": 2549
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.452244406152147e-06,
+      "loss": 0.4253,
+      "step": 2550
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.444770931368978e-06,
+      "loss": 0.418,
+      "step": 2551
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.437300665590018e-06,
+      "loss": 0.4272,
+      "step": 2552
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.429833614077817e-06,
+      "loss": 0.4324,
+      "step": 2553
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.4223697820926635e-06,
+      "loss": 0.4314,
+      "step": 2554
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.414909174892569e-06,
+      "loss": 0.4517,
+      "step": 2555
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.407451797733274e-06,
+      "loss": 0.4675,
+      "step": 2556
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.399997655868256e-06,
+      "loss": 0.4336,
+      "step": 2557
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.392546754548696e-06,
+      "loss": 0.4536,
+      "step": 2558
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.385099099023504e-06,
+      "loss": 0.4424,
+      "step": 2559
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.377654694539307e-06,
+      "loss": 0.4277,
+      "step": 2560
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.370213546340427e-06,
+      "loss": 0.4556,
+      "step": 2561
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.362775659668907e-06,
+      "loss": 0.4363,
+      "step": 2562
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.3553410397644766e-06,
+      "loss": 0.415,
+      "step": 2563
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.347909691864583e-06,
+      "loss": 0.4487,
+      "step": 2564
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.34048162120435e-06,
+      "loss": 0.4487,
+      "step": 2565
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 5.3330568330166055e-06,
+      "loss": 0.4131,
+      "step": 2566
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.325635332531864e-06,
+      "loss": 0.4153,
+      "step": 2567
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.318217124978313e-06,
+      "loss": 0.4587,
+      "step": 2568
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.310802215581835e-06,
+      "loss": 0.4336,
+      "step": 2569
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.303390609565975e-06,
+      "loss": 0.4321,
+      "step": 2570
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.295982312151962e-06,
+      "loss": 0.4451,
+      "step": 2571
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.2885773285586915e-06,
+      "loss": 0.4377,
+      "step": 2572
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.2811756640027175e-06,
+      "loss": 0.4106,
+      "step": 2573
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.2737773236982705e-06,
+      "loss": 0.4126,
+      "step": 2574
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.266382312857219e-06,
+      "loss": 0.4519,
+      "step": 2575
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.258990636689106e-06,
+      "loss": 0.4292,
+      "step": 2576
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.251602300401108e-06,
+      "loss": 0.4424,
+      "step": 2577
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.244217309198063e-06,
+      "loss": 0.4353,
+      "step": 2578
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.23683566828245e-06,
+      "loss": 0.4495,
+      "step": 2579
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.229457382854374e-06,
+      "loss": 0.426,
+      "step": 2580
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.222082458111597e-06,
+      "loss": 0.4414,
+      "step": 2581
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.214710899249494e-06,
+      "loss": 0.426,
+      "step": 2582
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.2073427114610855e-06,
+      "loss": 0.4082,
+      "step": 2583
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.199977899937003e-06,
+      "loss": 0.4138,
+      "step": 2584
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.1926164698655105e-06,
+      "loss": 0.4292,
+      "step": 2585
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.1852584264324864e-06,
+      "loss": 0.415,
+      "step": 2586
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.177903774821418e-06,
+      "loss": 0.4197,
+      "step": 2587
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.170552520213414e-06,
+      "loss": 0.4219,
+      "step": 2588
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.163204667787176e-06,
+      "loss": 0.4304,
+      "step": 2589
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.155860222719022e-06,
+      "loss": 0.4231,
+      "step": 2590
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.148519190182864e-06,
+      "loss": 0.4436,
+      "step": 2591
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.1411815753502055e-06,
+      "loss": 0.425,
+      "step": 2592
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.133847383390155e-06,
+      "loss": 0.4338,
+      "step": 2593
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.126516619469397e-06,
+      "loss": 0.4185,
+      "step": 2594
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.119189288752201e-06,
+      "loss": 0.4326,
+      "step": 2595
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.1118653964004284e-06,
+      "loss": 0.426,
+      "step": 2596
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.104544947573511e-06,
+      "loss": 0.4448,
+      "step": 2597
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.0972279474284605e-06,
+      "loss": 0.4316,
+      "step": 2598
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.089914401119849e-06,
+      "loss": 0.4365,
+      "step": 2599
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.082604313799828e-06,
+      "loss": 0.4172,
+      "step": 2600
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.075297690618101e-06,
+      "loss": 0.4495,
+      "step": 2601
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.067994536721936e-06,
+      "loss": 0.4409,
+      "step": 2602
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.0606948572561554e-06,
+      "loss": 0.4282,
+      "step": 2603
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 5.0533986573631386e-06,
+      "loss": 0.4646,
+      "step": 2604
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 5.046105942182815e-06,
+      "loss": 0.4348,
+      "step": 2605
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 5.038816716852648e-06,
+      "loss": 0.4482,
+      "step": 2606
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 5.03153098650765e-06,
+      "loss": 0.4314,
+      "step": 2607
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 5.024248756280371e-06,
+      "loss": 0.4202,
+      "step": 2608
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 5.016970031300896e-06,
+      "loss": 0.4438,
+      "step": 2609
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 5.0096948166968425e-06,
+      "loss": 0.4304,
+      "step": 2610
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 5.002423117593347e-06,
+      "loss": 0.4375,
+      "step": 2611
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.995154939113079e-06,
+      "loss": 0.4421,
+      "step": 2612
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.987890286376224e-06,
+      "loss": 0.4285,
+      "step": 2613
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.980629164500474e-06,
+      "loss": 0.4485,
+      "step": 2614
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.973371578601051e-06,
+      "loss": 0.4187,
+      "step": 2615
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.966117533790675e-06,
+      "loss": 0.4458,
+      "step": 2616
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.9588670351795795e-06,
+      "loss": 0.4316,
+      "step": 2617
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.951620087875491e-06,
+      "loss": 0.4285,
+      "step": 2618
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.944376696983633e-06,
+      "loss": 0.4299,
+      "step": 2619
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.937136867606736e-06,
+      "loss": 0.4231,
+      "step": 2620
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.929900604845009e-06,
+      "loss": 0.4502,
+      "step": 2621
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.922667913796153e-06,
+      "loss": 0.4197,
+      "step": 2622
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.91543879955536e-06,
+      "loss": 0.4236,
+      "step": 2623
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.908213267215287e-06,
+      "loss": 0.4192,
+      "step": 2624
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.900991321866084e-06,
+      "loss": 0.4448,
+      "step": 2625
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.89377296859536e-06,
+      "loss": 0.4338,
+      "step": 2626
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.8865582124882005e-06,
+      "loss": 0.4387,
+      "step": 2627
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.879347058627159e-06,
+      "loss": 0.4226,
+      "step": 2628
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.87213951209225e-06,
+      "loss": 0.427,
+      "step": 2629
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.864935577960943e-06,
+      "loss": 0.4561,
+      "step": 2630
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.857735261308162e-06,
+      "loss": 0.4468,
+      "step": 2631
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.850538567206291e-06,
+      "loss": 0.4358,
+      "step": 2632
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.843345500725151e-06,
+      "loss": 0.4375,
+      "step": 2633
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.836156066932015e-06,
+      "loss": 0.4143,
+      "step": 2634
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.828970270891601e-06,
+      "loss": 0.4402,
+      "step": 2635
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.82178811766605e-06,
+      "loss": 0.4526,
+      "step": 2636
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.814609612314951e-06,
+      "loss": 0.4399,
+      "step": 2637
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.807434759895312e-06,
+      "loss": 0.436,
+      "step": 2638
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.800263565461579e-06,
+      "loss": 0.4246,
+      "step": 2639
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.7930960340656075e-06,
+      "loss": 0.4282,
+      "step": 2640
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.785932170756686e-06,
+      "loss": 0.428,
+      "step": 2641
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.778771980581512e-06,
+      "loss": 0.4199,
+      "step": 2642
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 4.771615468584194e-06,
+      "loss": 0.4268,
+      "step": 2643
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.764462639806256e-06,
+      "loss": 0.4421,
+      "step": 2644
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.7573134992866174e-06,
+      "loss": 0.4224,
+      "step": 2645
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.750168052061607e-06,
+      "loss": 0.4319,
+      "step": 2646
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.743026303164955e-06,
+      "loss": 0.4409,
+      "step": 2647
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.7358882576277724e-06,
+      "loss": 0.4434,
+      "step": 2648
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.728753920478577e-06,
+      "loss": 0.446,
+      "step": 2649
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.721623296743261e-06,
+      "loss": 0.4407,
+      "step": 2650
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.714496391445114e-06,
+      "loss": 0.4329,
+      "step": 2651
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.70737320960479e-06,
+      "loss": 0.4197,
+      "step": 2652
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.700253756240334e-06,
+      "loss": 0.4082,
+      "step": 2653
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.693138036367162e-06,
+      "loss": 0.447,
+      "step": 2654
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.68602605499805e-06,
+      "loss": 0.4316,
+      "step": 2655
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.6789178171431544e-06,
+      "loss": 0.4246,
+      "step": 2656
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.671813327809981e-06,
+      "loss": 0.4592,
+      "step": 2657
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.664712592003407e-06,
+      "loss": 0.4282,
+      "step": 2658
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.657615614725655e-06,
+      "loss": 0.4111,
+      "step": 2659
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.650522400976304e-06,
+      "loss": 0.4226,
+      "step": 2660
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.64343295575229e-06,
+      "loss": 0.4255,
+      "step": 2661
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.636347284047878e-06,
+      "loss": 0.4307,
+      "step": 2662
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.629265390854689e-06,
+      "loss": 0.4382,
+      "step": 2663
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.622187281161673e-06,
+      "loss": 0.4324,
+      "step": 2664
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.615112959955113e-06,
+      "loss": 0.4087,
+      "step": 2665
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.608042432218639e-06,
+      "loss": 0.4355,
+      "step": 2666
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.60097570293319e-06,
+      "loss": 0.4531,
+      "step": 2667
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.593912777077043e-06,
+      "loss": 0.4365,
+      "step": 2668
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.586853659625782e-06,
+      "loss": 0.4104,
+      "step": 2669
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.579798355552322e-06,
+      "loss": 0.4292,
+      "step": 2670
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.572746869826879e-06,
+      "loss": 0.4583,
+      "step": 2671
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.565699207416988e-06,
+      "loss": 0.4282,
+      "step": 2672
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.558655373287494e-06,
+      "loss": 0.438,
+      "step": 2673
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.551615372400529e-06,
+      "loss": 0.425,
+      "step": 2674
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.544579209715542e-06,
+      "loss": 0.4392,
+      "step": 2675
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.5375468901892685e-06,
+      "loss": 0.4297,
+      "step": 2676
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.530518418775734e-06,
+      "loss": 0.4104,
+      "step": 2677
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.523493800426262e-06,
+      "loss": 0.4373,
+      "step": 2678
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.516473040089457e-06,
+      "loss": 0.4341,
+      "step": 2679
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.509456142711213e-06,
+      "loss": 0.4197,
+      "step": 2680
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.502443113234688e-06,
+      "loss": 0.4197,
+      "step": 2681
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 4.495433956600321e-06,
+      "loss": 0.4351,
+      "step": 2682
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.488428677745829e-06,
+      "loss": 0.4282,
+      "step": 2683
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.4814272816061945e-06,
+      "loss": 0.4158,
+      "step": 2684
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.474429773113665e-06,
+      "loss": 0.4331,
+      "step": 2685
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.467436157197741e-06,
+      "loss": 0.4148,
+      "step": 2686
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.460446438785194e-06,
+      "loss": 0.4382,
+      "step": 2687
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.45346062280004e-06,
+      "loss": 0.4651,
+      "step": 2688
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.446478714163547e-06,
+      "loss": 0.4167,
+      "step": 2689
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.439500717794234e-06,
+      "loss": 0.4512,
+      "step": 2690
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.432526638607864e-06,
+      "loss": 0.4692,
+      "step": 2691
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.425556481517441e-06,
+      "loss": 0.429,
+      "step": 2692
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.418590251433202e-06,
+      "loss": 0.4343,
+      "step": 2693
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.411627953262616e-06,
+      "loss": 0.4397,
+      "step": 2694
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.404669591910392e-06,
+      "loss": 0.4473,
+      "step": 2695
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.397715172278452e-06,
+      "loss": 0.4329,
+      "step": 2696
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.390764699265953e-06,
+      "loss": 0.4292,
+      "step": 2697
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.3838181777692665e-06,
+      "loss": 0.4299,
+      "step": 2698
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.376875612681986e-06,
+      "loss": 0.4338,
+      "step": 2699
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.369937008894906e-06,
+      "loss": 0.4419,
+      "step": 2700
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.363002371296037e-06,
+      "loss": 0.4219,
+      "step": 2701
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.356071704770598e-06,
+      "loss": 0.4741,
+      "step": 2702
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.349145014201008e-06,
+      "loss": 0.4348,
+      "step": 2703
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.342222304466891e-06,
+      "loss": 0.4155,
+      "step": 2704
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.335303580445054e-06,
+      "loss": 0.4302,
+      "step": 2705
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.328388847009502e-06,
+      "loss": 0.4507,
+      "step": 2706
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.321478109031436e-06,
+      "loss": 0.4287,
+      "step": 2707
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.314571371379229e-06,
+      "loss": 0.4246,
+      "step": 2708
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.307668638918448e-06,
+      "loss": 0.4387,
+      "step": 2709
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.300769916511837e-06,
+      "loss": 0.4175,
+      "step": 2710
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.293875209019303e-06,
+      "loss": 0.4229,
+      "step": 2711
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.286984521297941e-06,
+      "loss": 0.4094,
+      "step": 2712
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.280097858202e-06,
+      "loss": 0.4402,
+      "step": 2713
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.273215224582907e-06,
+      "loss": 0.4153,
+      "step": 2714
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.266336625289238e-06,
+      "loss": 0.4268,
+      "step": 2715
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.259462065166734e-06,
+      "loss": 0.4287,
+      "step": 2716
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.252591549058294e-06,
+      "loss": 0.4272,
+      "step": 2717
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.245725081803955e-06,
+      "loss": 0.4219,
+      "step": 2718
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.23886266824092e-06,
+      "loss": 0.4365,
+      "step": 2719
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 4.232004313203516e-06,
+      "loss": 0.4331,
+      "step": 2720
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.225150021523226e-06,
+      "loss": 0.4324,
+      "step": 2721
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.2182997980286715e-06,
+      "loss": 0.4407,
+      "step": 2722
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.211453647545591e-06,
+      "loss": 0.4246,
+      "step": 2723
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.204611574896875e-06,
+      "loss": 0.4094,
+      "step": 2724
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.197773584902524e-06,
+      "loss": 0.4094,
+      "step": 2725
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.190939682379676e-06,
+      "loss": 0.4333,
+      "step": 2726
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.184109872142577e-06,
+      "loss": 0.4199,
+      "step": 2727
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.1772841590026e-06,
+      "loss": 0.415,
+      "step": 2728
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.170462547768232e-06,
+      "loss": 0.4272,
+      "step": 2729
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.163645043245058e-06,
+      "loss": 0.4143,
+      "step": 2730
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.156831650235786e-06,
+      "loss": 0.4263,
+      "step": 2731
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.150022373540214e-06,
+      "loss": 0.4219,
+      "step": 2732
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.143217217955251e-06,
+      "loss": 0.4255,
+      "step": 2733
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.136416188274894e-06,
+      "loss": 0.4067,
+      "step": 2734
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.129619289290239e-06,
+      "loss": 0.4238,
+      "step": 2735
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.122826525789475e-06,
+      "loss": 0.4275,
+      "step": 2736
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.116037902557866e-06,
+      "loss": 0.4287,
+      "step": 2737
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.109253424377773e-06,
+      "loss": 0.4216,
+      "step": 2738
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.102473096028624e-06,
+      "loss": 0.4216,
+      "step": 2739
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.095696922286933e-06,
+      "loss": 0.4133,
+      "step": 2740
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.08892490792629e-06,
+      "loss": 0.4226,
+      "step": 2741
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.082157057717339e-06,
+      "loss": 0.4441,
+      "step": 2742
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.075393376427809e-06,
+      "loss": 0.4121,
+      "step": 2743
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.068633868822478e-06,
+      "loss": 0.4189,
+      "step": 2744
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.061878539663196e-06,
+      "loss": 0.4312,
+      "step": 2745
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.055127393708854e-06,
+      "loss": 0.4167,
+      "step": 2746
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.048380435715411e-06,
+      "loss": 0.4211,
+      "step": 2747
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.041637670435873e-06,
+      "loss": 0.4368,
+      "step": 2748
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.034899102620282e-06,
+      "loss": 0.4346,
+      "step": 2749
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.028164737015736e-06,
+      "loss": 0.4043,
+      "step": 2750
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.0214345783663655e-06,
+      "loss": 0.4294,
+      "step": 2751
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.014708631413333e-06,
+      "loss": 0.4014,
+      "step": 2752
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.007986900894846e-06,
+      "loss": 0.4343,
+      "step": 2753
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 4.001269391546132e-06,
+      "loss": 0.4243,
+      "step": 2754
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 3.994556108099455e-06,
+      "loss": 0.4238,
+      "step": 2755
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 3.987847055284085e-06,
+      "loss": 0.4409,
+      "step": 2756
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 3.981142237826332e-06,
+      "loss": 0.4199,
+      "step": 2757
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 3.974441660449508e-06,
+      "loss": 0.4341,
+      "step": 2758
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 3.967745327873934e-06,
+      "loss": 0.4658,
+      "step": 2759
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.9610532448169636e-06,
+      "loss": 0.415,
+      "step": 2760
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.954365415992932e-06,
+      "loss": 0.4263,
+      "step": 2761
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.947681846113194e-06,
+      "loss": 0.4194,
+      "step": 2762
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.941002539886094e-06,
+      "loss": 0.4336,
+      "step": 2763
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.9343275020169735e-06,
+      "loss": 0.4207,
+      "step": 2764
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.927656737208173e-06,
+      "loss": 0.45,
+      "step": 2765
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.9209902501590205e-06,
+      "loss": 0.4343,
+      "step": 2766
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.914328045565833e-06,
+      "loss": 0.428,
+      "step": 2767
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.9076701281219e-06,
+      "loss": 0.4424,
+      "step": 2768
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.901016502517505e-06,
+      "loss": 0.4094,
+      "step": 2769
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.8943671734398995e-06,
+      "loss": 0.4529,
+      "step": 2770
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.887722145573305e-06,
+      "loss": 0.4629,
+      "step": 2771
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.881081423598921e-06,
+      "loss": 0.4172,
+      "step": 2772
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.874445012194913e-06,
+      "loss": 0.4304,
+      "step": 2773
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.86781291603641e-06,
+      "loss": 0.4463,
+      "step": 2774
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.8611851397954935e-06,
+      "loss": 0.4299,
+      "step": 2775
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.854561688141205e-06,
+      "loss": 0.429,
+      "step": 2776
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.847942565739547e-06,
+      "loss": 0.4241,
+      "step": 2777
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.841327777253457e-06,
+      "loss": 0.429,
+      "step": 2778
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.834717327342844e-06,
+      "loss": 0.4133,
+      "step": 2779
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.828111220664536e-06,
+      "loss": 0.4224,
+      "step": 2780
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.8215094618723085e-06,
+      "loss": 0.4448,
+      "step": 2781
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.814912055616884e-06,
+      "loss": 0.4341,
+      "step": 2782
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.808319006545903e-06,
+      "loss": 0.4258,
+      "step": 2783
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.801730319303949e-06,
+      "loss": 0.4121,
+      "step": 2784
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.795145998532529e-06,
+      "loss": 0.4194,
+      "step": 2785
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.788566048870075e-06,
+      "loss": 0.439,
+      "step": 2786
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.7819904749519366e-06,
+      "loss": 0.3931,
+      "step": 2787
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.7754192814103776e-06,
+      "loss": 0.4297,
+      "step": 2788
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.768852472874589e-06,
+      "loss": 0.45,
+      "step": 2789
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.7622900539706563e-06,
+      "loss": 0.417,
+      "step": 2790
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.7557320293215836e-06,
+      "loss": 0.427,
+      "step": 2791
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.749178403547282e-06,
+      "loss": 0.4031,
+      "step": 2792
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.742629181264551e-06,
+      "loss": 0.4509,
+      "step": 2793
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.7360843670871015e-06,
+      "loss": 0.4348,
+      "step": 2794
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.729543965625526e-06,
+      "loss": 0.4319,
+      "step": 2795
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.7230079814873233e-06,
+      "loss": 0.4146,
+      "step": 2796
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 3.71647641927686e-06,
+      "loss": 0.4365,
+      "step": 2797
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.7099492835954154e-06,
+      "loss": 0.438,
+      "step": 2798
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.703426579041126e-06,
+      "loss": 0.4441,
+      "step": 2799
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.696908310209015e-06,
+      "loss": 0.4536,
+      "step": 2800
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.6903944816909867e-06,
+      "loss": 0.4165,
+      "step": 2801
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.6838850980758056e-06,
+      "loss": 0.4216,
+      "step": 2802
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.6773801639491147e-06,
+      "loss": 0.4241,
+      "step": 2803
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.6708796838934224e-06,
+      "loss": 0.4326,
+      "step": 2804
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.6643836624880903e-06,
+      "loss": 0.4407,
+      "step": 2805
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.6578921043093507e-06,
+      "loss": 0.4136,
+      "step": 2806
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.6514050139302804e-06,
+      "loss": 0.4167,
+      "step": 2807
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.644922395920819e-06,
+      "loss": 0.4343,
+      "step": 2808
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.638444254847746e-06,
+      "loss": 0.4424,
+      "step": 2809
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.6319705952746943e-06,
+      "loss": 0.4395,
+      "step": 2810
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.625501421762141e-06,
+      "loss": 0.4277,
+      "step": 2811
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.619036738867392e-06,
+      "loss": 0.437,
+      "step": 2812
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.6125765511446043e-06,
+      "loss": 0.4294,
+      "step": 2813
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.606120863144753e-06,
+      "loss": 0.4285,
+      "step": 2814
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.5996696794156584e-06,
+      "loss": 0.4631,
+      "step": 2815
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.5932230045019545e-06,
+      "loss": 0.4246,
+      "step": 2816
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.586780842945107e-06,
+      "loss": 0.4243,
+      "step": 2817
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.5803431992834036e-06,
+      "loss": 0.4221,
+      "step": 2818
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.5739100780519398e-06,
+      "loss": 0.4229,
+      "step": 2819
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.5674814837826364e-06,
+      "loss": 0.4539,
+      "step": 2820
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.561057421004215e-06,
+      "loss": 0.4243,
+      "step": 2821
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.554637894242211e-06,
+      "loss": 0.4089,
+      "step": 2822
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.548222908018968e-06,
+      "loss": 0.4331,
+      "step": 2823
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.541812466853618e-06,
+      "loss": 0.4099,
+      "step": 2824
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.5354065752621058e-06,
+      "loss": 0.4199,
+      "step": 2825
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.529005237757159e-06,
+      "loss": 0.4309,
+      "step": 2826
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.522608458848309e-06,
+      "loss": 0.4521,
+      "step": 2827
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.516216243041862e-06,
+      "loss": 0.4089,
+      "step": 2828
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.5098285948409216e-06,
+      "loss": 0.4316,
+      "step": 2829
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.50344551874537e-06,
+      "loss": 0.428,
+      "step": 2830
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.497067019251863e-06,
+      "loss": 0.4307,
+      "step": 2831
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.4906931008538425e-06,
+      "loss": 0.3801,
+      "step": 2832
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.4843237680415153e-06,
+      "loss": 0.4255,
+      "step": 2833
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.4779590253018558e-06,
+      "loss": 0.4321,
+      "step": 2834
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.471598877118614e-06,
+      "loss": 0.4316,
+      "step": 2835
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 3.465243327972295e-06,
+      "loss": 0.4426,
+      "step": 2836
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.4588923823401745e-06,
+      "loss": 0.4219,
+      "step": 2837
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.4525460446962735e-06,
+      "loss": 0.3948,
+      "step": 2838
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.446204319511366e-06,
+      "loss": 0.4248,
+      "step": 2839
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.4398672112529885e-06,
+      "loss": 0.437,
+      "step": 2840
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.4335347243854165e-06,
+      "loss": 0.405,
+      "step": 2841
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.4272068633696752e-06,
+      "loss": 0.4263,
+      "step": 2842
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.420883632663523e-06,
+      "loss": 0.4309,
+      "step": 2843
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.4145650367214668e-06,
+      "loss": 0.4199,
+      "step": 2844
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.40825107999474e-06,
+      "loss": 0.4155,
+      "step": 2845
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.401941766931306e-06,
+      "loss": 0.4326,
+      "step": 2846
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.395637101975867e-06,
+      "loss": 0.4136,
+      "step": 2847
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.389337089569844e-06,
+      "loss": 0.4202,
+      "step": 2848
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.3830417341513856e-06,
+      "loss": 0.4131,
+      "step": 2849
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.3767510401553516e-06,
+      "loss": 0.4189,
+      "step": 2850
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.3704650120133207e-06,
+      "loss": 0.4194,
+      "step": 2851
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.364183654153592e-06,
+      "loss": 0.45,
+      "step": 2852
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.3579069710011626e-06,
+      "loss": 0.4131,
+      "step": 2853
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.351634966977746e-06,
+      "loss": 0.4324,
+      "step": 2854
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.3453676465017547e-06,
+      "loss": 0.4053,
+      "step": 2855
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.3391050139883085e-06,
+      "loss": 0.4294,
+      "step": 2856
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.332847073849216e-06,
+      "loss": 0.4197,
+      "step": 2857
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.3265938304929777e-06,
+      "loss": 0.4233,
+      "step": 2858
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.3203452883247966e-06,
+      "loss": 0.438,
+      "step": 2859
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.3141014517465588e-06,
+      "loss": 0.436,
+      "step": 2860
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.3078623251568376e-06,
+      "loss": 0.4329,
+      "step": 2861
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.301627912950882e-06,
+      "loss": 0.4199,
+      "step": 2862
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.2953982195206214e-06,
+      "loss": 0.4221,
+      "step": 2863
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.2891732492546702e-06,
+      "loss": 0.4312,
+      "step": 2864
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.282953006538301e-06,
+      "loss": 0.4316,
+      "step": 2865
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.276737495753467e-06,
+      "loss": 0.417,
+      "step": 2866
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.2705267212787896e-06,
+      "loss": 0.4297,
+      "step": 2867
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.2643206874895385e-06,
+      "loss": 0.406,
+      "step": 2868
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.258119398757663e-06,
+      "loss": 0.4417,
+      "step": 2869
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.2519228594517538e-06,
+      "loss": 0.4209,
+      "step": 2870
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.2457310739370684e-06,
+      "loss": 0.4194,
+      "step": 2871
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.2395440465755025e-06,
+      "loss": 0.4092,
+      "step": 2872
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.2333617817256103e-06,
+      "loss": 0.4236,
+      "step": 2873
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 3.2271842837425917e-06,
+      "loss": 0.4351,
+      "step": 2874
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.221011556978276e-06,
+      "loss": 0.4255,
+      "step": 2875
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.214843605781147e-06,
+      "loss": 0.4109,
+      "step": 2876
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.208680434496312e-06,
+      "loss": 0.4121,
+      "step": 2877
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.202522047465516e-06,
+      "loss": 0.4302,
+      "step": 2878
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.1963684490271396e-06,
+      "loss": 0.4285,
+      "step": 2879
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.1902196435161757e-06,
+      "loss": 0.45,
+      "step": 2880
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.1840756352642564e-06,
+      "loss": 0.4214,
+      "step": 2881
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.1779364285996194e-06,
+      "loss": 0.4409,
+      "step": 2882
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.171802027847135e-06,
+      "loss": 0.4277,
+      "step": 2883
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.1656724373282743e-06,
+      "loss": 0.4182,
+      "step": 2884
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.1595476613611265e-06,
+      "loss": 0.4468,
+      "step": 2885
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.1534277042603943e-06,
+      "loss": 0.4443,
+      "step": 2886
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.1473125703373707e-06,
+      "loss": 0.4319,
+      "step": 2887
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.1412022638999682e-06,
+      "loss": 0.4158,
+      "step": 2888
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.135096789252683e-06,
+      "loss": 0.4429,
+      "step": 2889
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.1289961506966217e-06,
+      "loss": 0.428,
+      "step": 2890
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.12290035252947e-06,
+      "loss": 0.4312,
+      "step": 2891
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.116809399045514e-06,
+      "loss": 0.4277,
+      "step": 2892
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.110723294535627e-06,
+      "loss": 0.4275,
+      "step": 2893
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.104642043287256e-06,
+      "loss": 0.4148,
+      "step": 2894
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.098565649584444e-06,
+      "loss": 0.4165,
+      "step": 2895
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.092494117707796e-06,
+      "loss": 0.408,
+      "step": 2896
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.086427451934504e-06,
+      "loss": 0.417,
+      "step": 2897
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.08036565653833e-06,
+      "loss": 0.4287,
+      "step": 2898
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.0743087357895986e-06,
+      "loss": 0.4268,
+      "step": 2899
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.0682566939552103e-06,
+      "loss": 0.4375,
+      "step": 2900
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.062209535298617e-06,
+      "loss": 0.4424,
+      "step": 2901
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.0561672640798425e-06,
+      "loss": 0.4258,
+      "step": 2902
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.050129884555455e-06,
+      "loss": 0.4448,
+      "step": 2903
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.0440974009785883e-06,
+      "loss": 0.4141,
+      "step": 2904
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.0380698175989233e-06,
+      "loss": 0.4214,
+      "step": 2905
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.0320471386626826e-06,
+      "loss": 0.4097,
+      "step": 2906
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.0260293684126454e-06,
+      "loss": 0.427,
+      "step": 2907
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.020016511088123e-06,
+      "loss": 0.407,
+      "step": 2908
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.0140085709249666e-06,
+      "loss": 0.4253,
+      "step": 2909
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.0080055521555683e-06,
+      "loss": 0.4358,
+      "step": 2910
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 3.002007459008852e-06,
+      "loss": 0.4089,
+      "step": 2911
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 2.996014295710273e-06,
+      "loss": 0.4353,
+      "step": 2912
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 2.990026066481806e-06,
+      "loss": 0.4553,
+      "step": 2913
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.98404277554196e-06,
+      "loss": 0.4336,
+      "step": 2914
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.9780644271057556e-06,
+      "loss": 0.4231,
+      "step": 2915
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.972091025384738e-06,
+      "loss": 0.4172,
+      "step": 2916
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.9661225745869703e-06,
+      "loss": 0.4285,
+      "step": 2917
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.9601590789170166e-06,
+      "loss": 0.4199,
+      "step": 2918
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.9542005425759634e-06,
+      "loss": 0.4314,
+      "step": 2919
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.9482469697613937e-06,
+      "loss": 0.4329,
+      "step": 2920
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.942298364667394e-06,
+      "loss": 0.4346,
+      "step": 2921
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.9363547314845598e-06,
+      "loss": 0.4304,
+      "step": 2922
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.9304160743999766e-06,
+      "loss": 0.4297,
+      "step": 2923
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.9244823975972303e-06,
+      "loss": 0.426,
+      "step": 2924
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.9185537052563926e-06,
+      "loss": 0.4146,
+      "step": 2925
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.912630001554022e-06,
+      "loss": 0.4219,
+      "step": 2926
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.9067112906631736e-06,
+      "loss": 0.4272,
+      "step": 2927
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.9007975767533714e-06,
+      "loss": 0.4204,
+      "step": 2928
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.8948888639906283e-06,
+      "loss": 0.4167,
+      "step": 2929
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.8889851565374337e-06,
+      "loss": 0.4292,
+      "step": 2930
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.8830864585527506e-06,
+      "loss": 0.4333,
+      "step": 2931
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.8771927741920092e-06,
+      "loss": 0.4331,
+      "step": 2932
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.8713041076071067e-06,
+      "loss": 0.4307,
+      "step": 2933
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.865420462946411e-06,
+      "loss": 0.4319,
+      "step": 2934
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.859541844354752e-06,
+      "loss": 0.4141,
+      "step": 2935
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.8536682559734188e-06,
+      "loss": 0.4033,
+      "step": 2936
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.8477997019401517e-06,
+      "loss": 0.4526,
+      "step": 2937
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.841936186389144e-06,
+      "loss": 0.4536,
+      "step": 2938
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.8360777134510507e-06,
+      "loss": 0.4187,
+      "step": 2939
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.830224287252962e-06,
+      "loss": 0.4575,
+      "step": 2940
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.824375911918418e-06,
+      "loss": 0.4368,
+      "step": 2941
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.8185325915674023e-06,
+      "loss": 0.4253,
+      "step": 2942
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.812694330316339e-06,
+      "loss": 0.4236,
+      "step": 2943
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.8068611322780816e-06,
+      "loss": 0.4507,
+      "step": 2944
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.801033001561916e-06,
+      "loss": 0.4395,
+      "step": 2945
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.795209942273569e-06,
+      "loss": 0.4263,
+      "step": 2946
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.789391958515183e-06,
+      "loss": 0.4241,
+      "step": 2947
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.7835790543853326e-06,
+      "loss": 0.4148,
+      "step": 2948
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.7777712339790143e-06,
+      "loss": 0.4175,
+      "step": 2949
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.771968501387634e-06,
+      "loss": 0.4365,
+      "step": 2950
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.7661708606990254e-06,
+      "loss": 0.4485,
+      "step": 2951
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.760378315997425e-06,
+      "loss": 0.4189,
+      "step": 2952
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.7545908713634863e-06,
+      "loss": 0.4294,
+      "step": 2953
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.7488085308742694e-06,
+      "loss": 0.4373,
+      "step": 2954
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.743031298603233e-06,
+      "loss": 0.4241,
+      "step": 2955
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.737259178620244e-06,
+      "loss": 0.4097,
+      "step": 2956
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.7314921749915624e-06,
+      "loss": 0.4463,
+      "step": 2957
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.72573029177985e-06,
+      "loss": 0.4258,
+      "step": 2958
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.7199735330441523e-06,
+      "loss": 0.4194,
+      "step": 2959
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.7142219028399143e-06,
+      "loss": 0.4475,
+      "step": 2960
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.7084754052189665e-06,
+      "loss": 0.4348,
+      "step": 2961
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.702734044229515e-06,
+      "loss": 0.4199,
+      "step": 2962
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.696997823916161e-06,
+      "loss": 0.4187,
+      "step": 2963
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.6912667483198685e-06,
+      "loss": 0.4438,
+      "step": 2964
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.6855408214779944e-06,
+      "loss": 0.4194,
+      "step": 2965
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.679820047424253e-06,
+      "loss": 0.4319,
+      "step": 2966
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.6741044301887377e-06,
+      "loss": 0.4358,
+      "step": 2967
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.668393973797909e-06,
+      "loss": 0.4023,
+      "step": 2968
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.662688682274586e-06,
+      "loss": 0.4143,
+      "step": 2969
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.6569885596379562e-06,
+      "loss": 0.4402,
+      "step": 2970
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.651293609903559e-06,
+      "loss": 0.4509,
+      "step": 2971
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.645603837083294e-06,
+      "loss": 0.4282,
+      "step": 2972
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.639919245185416e-06,
+      "loss": 0.4265,
+      "step": 2973
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.634239838214523e-06,
+      "loss": 0.4277,
+      "step": 2974
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.6285656201715694e-06,
+      "loss": 0.4175,
+      "step": 2975
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.622896595053841e-06,
+      "loss": 0.4243,
+      "step": 2976
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.6172327668549834e-06,
+      "loss": 0.4172,
+      "step": 2977
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.6115741395649617e-06,
+      "loss": 0.4153,
+      "step": 2978
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.605920717170092e-06,
+      "loss": 0.4338,
+      "step": 2979
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.6002725036530207e-06,
+      "loss": 0.4395,
+      "step": 2980
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.594629502992716e-06,
+      "loss": 0.4182,
+      "step": 2981
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.5889917191644864e-06,
+      "loss": 0.427,
+      "step": 2982
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.5833591561399576e-06,
+      "loss": 0.4253,
+      "step": 2983
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.5777318178870735e-06,
+      "loss": 0.4343,
+      "step": 2984
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.5721097083701085e-06,
+      "loss": 0.4343,
+      "step": 2985
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.566492831549645e-06,
+      "loss": 0.4221,
+      "step": 2986
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.560881191382587e-06,
+      "loss": 0.4241,
+      "step": 2987
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.5552747918221387e-06,
+      "loss": 0.429,
+      "step": 2988
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.549673636817823e-06,
+      "loss": 0.4026,
+      "step": 2989
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 2.5440777303154607e-06,
+      "loss": 0.4253,
+      "step": 2990
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.5384870762571724e-06,
+      "loss": 0.408,
+      "step": 2991
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.5329016785813965e-06,
+      "loss": 0.4229,
+      "step": 2992
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.527321541222846e-06,
+      "loss": 0.4424,
+      "step": 2993
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.5217466681125434e-06,
+      "loss": 0.4185,
+      "step": 2994
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.516177063177797e-06,
+      "loss": 0.4329,
+      "step": 2995
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.510612730342199e-06,
+      "loss": 0.4209,
+      "step": 2996
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.505053673525637e-06,
+      "loss": 0.4304,
+      "step": 2997
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.4994998966442787e-06,
+      "loss": 0.4426,
+      "step": 2998
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.493951403610573e-06,
+      "loss": 0.4131,
+      "step": 2999
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.4884081983332385e-06,
+      "loss": 0.4297,
+      "step": 3000
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.4828702847172825e-06,
+      "loss": 0.4255,
+      "step": 3001
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.4773376666639735e-06,
+      "loss": 0.4377,
+      "step": 3002
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.47181034807085e-06,
+      "loss": 0.4316,
+      "step": 3003
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.466288332831722e-06,
+      "loss": 0.4138,
+      "step": 3004
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.460771624836664e-06,
+      "loss": 0.4248,
+      "step": 3005
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.4552602279720117e-06,
+      "loss": 0.438,
+      "step": 3006
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.4497541461203546e-06,
+      "loss": 0.4319,
+      "step": 3007
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.4442533831605352e-06,
+      "loss": 0.405,
+      "step": 3008
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.4387579429676623e-06,
+      "loss": 0.4202,
+      "step": 3009
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.4332678294130774e-06,
+      "loss": 0.4292,
+      "step": 3010
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.42778304636439e-06,
+      "loss": 0.4189,
+      "step": 3011
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.4223035976854393e-06,
+      "loss": 0.4126,
+      "step": 3012
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.4168294872363076e-06,
+      "loss": 0.4438,
+      "step": 3013
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.4113607188733248e-06,
+      "loss": 0.4167,
+      "step": 3014
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.405897296449047e-06,
+      "loss": 0.427,
+      "step": 3015
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.4004392238122733e-06,
+      "loss": 0.4292,
+      "step": 3016
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.394986504808029e-06,
+      "loss": 0.4316,
+      "step": 3017
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.389539143277574e-06,
+      "loss": 0.4121,
+      "step": 3018
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.384097143058386e-06,
+      "loss": 0.4241,
+      "step": 3019
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.3786605079841673e-06,
+      "loss": 0.4238,
+      "step": 3020
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.3732292418848467e-06,
+      "loss": 0.4268,
+      "step": 3021
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.3678033485865626e-06,
+      "loss": 0.4263,
+      "step": 3022
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.362382831911675e-06,
+      "loss": 0.4326,
+      "step": 3023
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.356967695678757e-06,
+      "loss": 0.4387,
+      "step": 3024
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.3515579437025813e-06,
+      "loss": 0.4214,
+      "step": 3025
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.3461535797941416e-06,
+      "loss": 0.436,
+      "step": 3026
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.3407546077606223e-06,
+      "loss": 0.4629,
+      "step": 3027
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.335361031405422e-06,
+      "loss": 0.4006,
+      "step": 3028
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 2.329972854528123e-06,
+      "loss": 0.427,
+      "step": 3029
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.3245900809245248e-06,
+      "loss": 0.427,
+      "step": 3030
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.3192127143866027e-06,
+      "loss": 0.4268,
+      "step": 3031
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.313840758702527e-06,
+      "loss": 0.4321,
+      "step": 3032
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.3084742176566623e-06,
+      "loss": 0.417,
+      "step": 3033
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.303113095029549e-06,
+      "loss": 0.4177,
+      "step": 3034
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.2977573945979194e-06,
+      "loss": 0.4058,
+      "step": 3035
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.2924071201346853e-06,
+      "loss": 0.4329,
+      "step": 3036
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.287062275408929e-06,
+      "loss": 0.4187,
+      "step": 3037
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.2817228641859157e-06,
+      "loss": 0.4021,
+      "step": 3038
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.2763888902270735e-06,
+      "loss": 0.4189,
+      "step": 3039
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.2710603572900137e-06,
+      "loss": 0.4104,
+      "step": 3040
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.2657372691284997e-06,
+      "loss": 0.4148,
+      "step": 3041
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.2604196294924696e-06,
+      "loss": 0.4392,
+      "step": 3042
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.2551074421280226e-06,
+      "loss": 0.4248,
+      "step": 3043
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.24980071077741e-06,
+      "loss": 0.4243,
+      "step": 3044
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.2444994391790463e-06,
+      "loss": 0.4351,
+      "step": 3045
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.2392036310674956e-06,
+      "loss": 0.4031,
+      "step": 3046
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.233913290173477e-06,
+      "loss": 0.4573,
+      "step": 3047
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.2286284202238538e-06,
+      "loss": 0.4258,
+      "step": 3048
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.223349024941638e-06,
+      "loss": 0.4275,
+      "step": 3049
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.2180751080459874e-06,
+      "loss": 0.4119,
+      "step": 3050
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.2128066732521915e-06,
+      "loss": 0.4255,
+      "step": 3051
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.2075437242716903e-06,
+      "loss": 0.4333,
+      "step": 3052
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.2022862648120456e-06,
+      "loss": 0.3982,
+      "step": 3053
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.1970342985769645e-06,
+      "loss": 0.4216,
+      "step": 3054
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.1917878292662785e-06,
+      "loss": 0.4155,
+      "step": 3055
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.186546860575944e-06,
+      "loss": 0.4258,
+      "step": 3056
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.1813113961980503e-06,
+      "loss": 0.429,
+      "step": 3057
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.1760814398207986e-06,
+      "loss": 0.4109,
+      "step": 3058
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.170856995128521e-06,
+      "loss": 0.4009,
+      "step": 3059
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.1656380658016553e-06,
+      "loss": 0.4167,
+      "step": 3060
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.160424655516764e-06,
+      "loss": 0.4304,
+      "step": 3061
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.1552167679465196e-06,
+      "loss": 0.4431,
+      "step": 3062
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.150014406759697e-06,
+      "loss": 0.4312,
+      "step": 3063
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.1448175756211874e-06,
+      "loss": 0.4104,
+      "step": 3064
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.139626278191981e-06,
+      "loss": 0.4016,
+      "step": 3065
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.1344405181291662e-06,
+      "loss": 0.4404,
+      "step": 3066
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 2.129260299085939e-06,
+      "loss": 0.4055,
+      "step": 3067
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.124085624711587e-06,
+      "loss": 0.4236,
+      "step": 3068
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.118916498651494e-06,
+      "loss": 0.425,
+      "step": 3069
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.1137529245471345e-06,
+      "loss": 0.4148,
+      "step": 3070
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.1085949060360654e-06,
+      "loss": 0.4231,
+      "step": 3071
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.10344244675194e-06,
+      "loss": 0.4294,
+      "step": 3072
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.0982955503244895e-06,
+      "loss": 0.4219,
+      "step": 3073
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.0931542203795306e-06,
+      "loss": 0.4036,
+      "step": 3074
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.088018460538952e-06,
+      "loss": 0.4241,
+      "step": 3075
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.0828882744207255e-06,
+      "loss": 0.4097,
+      "step": 3076
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.077763665638891e-06,
+      "loss": 0.4155,
+      "step": 3077
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.0726446378035593e-06,
+      "loss": 0.4355,
+      "step": 3078
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.0675311945209142e-06,
+      "loss": 0.4524,
+      "step": 3079
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.0624233393932024e-06,
+      "loss": 0.4185,
+      "step": 3080
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.0573210760187368e-06,
+      "loss": 0.4153,
+      "step": 3081
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.0522244079918876e-06,
+      "loss": 0.4561,
+      "step": 3082
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.04713333890308e-06,
+      "loss": 0.4465,
+      "step": 3083
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.042047872338805e-06,
+      "loss": 0.4263,
+      "step": 3084
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.036968011881596e-06,
+      "loss": 0.4353,
+      "step": 3085
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.0318937611100454e-06,
+      "loss": 0.4187,
+      "step": 3086
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.02682512359879e-06,
+      "loss": 0.4204,
+      "step": 3087
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.021762102918515e-06,
+      "loss": 0.4299,
+      "step": 3088
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.0167047026359456e-06,
+      "loss": 0.4224,
+      "step": 3089
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.011652926313844e-06,
+      "loss": 0.4128,
+      "step": 3090
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.0066067775110187e-06,
+      "loss": 0.4285,
+      "step": 3091
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 2.0015662597823104e-06,
+      "loss": 0.4351,
+      "step": 3092
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 1.9965313766785943e-06,
+      "loss": 0.4192,
+      "step": 3093
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 1.991502131746773e-06,
+      "loss": 0.405,
+      "step": 3094
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 1.986478528529776e-06,
+      "loss": 0.3977,
+      "step": 3095
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 1.981460570566566e-06,
+      "loss": 0.3933,
+      "step": 3096
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 1.976448261392119e-06,
+      "loss": 0.4133,
+      "step": 3097
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 1.9714416045374395e-06,
+      "loss": 0.4324,
+      "step": 3098
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 1.9664406035295493e-06,
+      "loss": 0.4272,
+      "step": 3099
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 1.961445261891479e-06,
+      "loss": 0.4219,
+      "step": 3100
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 1.9564555831422826e-06,
+      "loss": 0.4336,
+      "step": 3101
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 1.9514715707970134e-06,
+      "loss": 0.4192,
+      "step": 3102
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 1.946493228366745e-06,
+      "loss": 0.4377,
+      "step": 3103
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 1.941520559358544e-06,
+      "loss": 0.406,
+      "step": 3104
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 1.9365535672754907e-06,
+      "loss": 0.4202,
+      "step": 3105
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 1.9315922556166656e-06,
+      "loss": 0.4241,
+      "step": 3106
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.9266366278771377e-06,
+      "loss": 0.4126,
+      "step": 3107
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.9216866875479836e-06,
+      "loss": 0.4062,
+      "step": 3108
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.916742438116265e-06,
+      "loss": 0.4128,
+      "step": 3109
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.9118038830650376e-06,
+      "loss": 0.4204,
+      "step": 3110
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.9068710258733491e-06,
+      "loss": 0.4072,
+      "step": 3111
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.9019438700162262e-06,
+      "loss": 0.4443,
+      "step": 3112
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.897022418964687e-06,
+      "loss": 0.4033,
+      "step": 3113
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.892106676185722e-06,
+      "loss": 0.4146,
+      "step": 3114
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.8871966451423097e-06,
+      "loss": 0.4138,
+      "step": 3115
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.8822923292933948e-06,
+      "loss": 0.408,
+      "step": 3116
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.877393732093904e-06,
+      "loss": 0.4128,
+      "step": 3117
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.8725008569947366e-06,
+      "loss": 0.4382,
+      "step": 3118
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.8676137074427502e-06,
+      "loss": 0.4297,
+      "step": 3119
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.8627322868807817e-06,
+      "loss": 0.4338,
+      "step": 3120
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.8578565987476215e-06,
+      "loss": 0.4314,
+      "step": 3121
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.85298664647803e-06,
+      "loss": 0.4312,
+      "step": 3122
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.8481224335027182e-06,
+      "loss": 0.4307,
+      "step": 3123
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.843263963248364e-06,
+      "loss": 0.4302,
+      "step": 3124
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.8384112391375942e-06,
+      "loss": 0.4126,
+      "step": 3125
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.833564264588984e-06,
+      "loss": 0.4207,
+      "step": 3126
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.8287230430170688e-06,
+      "loss": 0.449,
+      "step": 3127
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.8238875778323185e-06,
+      "loss": 0.4124,
+      "step": 3128
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.8190578724411567e-06,
+      "loss": 0.4326,
+      "step": 3129
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.81423393024595e-06,
+      "loss": 0.4163,
+      "step": 3130
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.8094157546449965e-06,
+      "loss": 0.4285,
+      "step": 3131
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.8046033490325422e-06,
+      "loss": 0.4055,
+      "step": 3132
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.7997967167987585e-06,
+      "loss": 0.4233,
+      "step": 3133
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.7949958613297602e-06,
+      "loss": 0.4375,
+      "step": 3134
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.7902007860075809e-06,
+      "loss": 0.4077,
+      "step": 3135
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.7854114942101909e-06,
+      "loss": 0.4448,
+      "step": 3136
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.7806279893114874e-06,
+      "loss": 0.4446,
+      "step": 3137
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.7758502746812811e-06,
+      "loss": 0.4041,
+      "step": 3138
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.7710783536853149e-06,
+      "loss": 0.4272,
+      "step": 3139
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.766312229685243e-06,
+      "loss": 0.4221,
+      "step": 3140
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.7615519060386355e-06,
+      "loss": 0.4175,
+      "step": 3141
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.7567973860989806e-06,
+      "loss": 0.4233,
+      "step": 3142
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.7520486732156772e-06,
+      "loss": 0.4263,
+      "step": 3143
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.7473057707340334e-06,
+      "loss": 0.4299,
+      "step": 3144
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.7425686819952591e-06,
+      "loss": 0.418,
+      "step": 3145
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.737837410336477e-06,
+      "loss": 0.4302,
+      "step": 3146
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.733111959090702e-06,
+      "loss": 0.416,
+      "step": 3147
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.7283923315868577e-06,
+      "loss": 0.4084,
+      "step": 3148
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.723678531149763e-06,
+      "loss": 0.4141,
+      "step": 3149
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.718970561100125e-06,
+      "loss": 0.406,
+      "step": 3150
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.7142684247545537e-06,
+      "loss": 0.4143,
+      "step": 3151
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.7095721254255437e-06,
+      "loss": 0.4243,
+      "step": 3152
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.704881666421474e-06,
+      "loss": 0.4114,
+      "step": 3153
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.7001970510466171e-06,
+      "loss": 0.4099,
+      "step": 3154
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.6955182826011252e-06,
+      "loss": 0.4236,
+      "step": 3155
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.6908453643810342e-06,
+      "loss": 0.405,
+      "step": 3156
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.6861782996782538e-06,
+      "loss": 0.4395,
+      "step": 3157
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.6815170917805711e-06,
+      "loss": 0.4138,
+      "step": 3158
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.676861743971653e-06,
+      "loss": 0.4333,
+      "step": 3159
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.6722122595310286e-06,
+      "loss": 0.4045,
+      "step": 3160
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.6675686417341042e-06,
+      "loss": 0.4294,
+      "step": 3161
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.6629308938521528e-06,
+      "loss": 0.437,
+      "step": 3162
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.6582990191523097e-06,
+      "loss": 0.4014,
+      "step": 3163
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.6536730208975738e-06,
+      "loss": 0.4463,
+      "step": 3164
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.6490529023467993e-06,
+      "loss": 0.4216,
+      "step": 3165
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.6444386667547063e-06,
+      "loss": 0.425,
+      "step": 3166
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.6398303173718667e-06,
+      "loss": 0.427,
+      "step": 3167
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.635227857444709e-06,
+      "loss": 0.449,
+      "step": 3168
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.6306312902155075e-06,
+      "loss": 0.4258,
+      "step": 3169
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.6260406189223854e-06,
+      "loss": 0.4233,
+      "step": 3170
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.621455846799318e-06,
+      "loss": 0.4226,
+      "step": 3171
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.6168769770761194e-06,
+      "loss": 0.4158,
+      "step": 3172
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.6123040129784484e-06,
+      "loss": 0.4131,
+      "step": 3173
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.6077369577278035e-06,
+      "loss": 0.4199,
+      "step": 3174
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.6031758145415222e-06,
+      "loss": 0.415,
+      "step": 3175
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.5986205866327732e-06,
+      "loss": 0.4065,
+      "step": 3176
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.5940712772105572e-06,
+      "loss": 0.4097,
+      "step": 3177
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.589527889479713e-06,
+      "loss": 0.4014,
+      "step": 3178
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.5849904266408978e-06,
+      "loss": 0.4329,
+      "step": 3179
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.5804588918906028e-06,
+      "loss": 0.4084,
+      "step": 3180
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.575933288421142e-06,
+      "loss": 0.3938,
+      "step": 3181
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.5714136194206464e-06,
+      "loss": 0.4236,
+      "step": 3182
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.5668998880730735e-06,
+      "loss": 0.4153,
+      "step": 3183
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.5623920975581874e-06,
+      "loss": 0.4307,
+      "step": 3184
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.557890251051578e-06,
+      "loss": 0.4146,
+      "step": 3185
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.553394351724644e-06,
+      "loss": 0.4329,
+      "step": 3186
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.548904402744591e-06,
+      "loss": 0.4099,
+      "step": 3187
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.5444204072744385e-06,
+      "loss": 0.4141,
+      "step": 3188
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.5399423684730052e-06,
+      "loss": 0.418,
+      "step": 3189
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.5354702894949214e-06,
+      "loss": 0.4473,
+      "step": 3190
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.531004173490611e-06,
+      "loss": 0.4053,
+      "step": 3191
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.5265440236063035e-06,
+      "loss": 0.4189,
+      "step": 3192
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.5220898429840236e-06,
+      "loss": 0.4275,
+      "step": 3193
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.5176416347615886e-06,
+      "loss": 0.4158,
+      "step": 3194
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.5131994020726127e-06,
+      "loss": 0.4226,
+      "step": 3195
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.5087631480464936e-06,
+      "loss": 0.405,
+      "step": 3196
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.5043328758084274e-06,
+      "loss": 0.4082,
+      "step": 3197
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.499908588479384e-06,
+      "loss": 0.4365,
+      "step": 3198
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.495490289176128e-06,
+      "loss": 0.4019,
+      "step": 3199
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.4910779810112019e-06,
+      "loss": 0.4204,
+      "step": 3200
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.4866716670929226e-06,
+      "loss": 0.4109,
+      "step": 3201
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.4822713505253938e-06,
+      "loss": 0.4053,
+      "step": 3202
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.4778770344084848e-06,
+      "loss": 0.4114,
+      "step": 3203
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.4734887218378435e-06,
+      "loss": 0.4172,
+      "step": 3204
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.4691064159048907e-06,
+      "loss": 0.4648,
+      "step": 3205
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.4647301196968055e-06,
+      "loss": 0.395,
+      "step": 3206
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.4603598362965465e-06,
+      "loss": 0.4177,
+      "step": 3207
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.4559955687828242e-06,
+      "loss": 0.4028,
+      "step": 3208
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.4516373202301216e-06,
+      "loss": 0.4292,
+      "step": 3209
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.4472850937086736e-06,
+      "loss": 0.4175,
+      "step": 3210
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.4429388922844756e-06,
+      "loss": 0.4219,
+      "step": 3211
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.438598719019284e-06,
+      "loss": 0.4294,
+      "step": 3212
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.4342645769705977e-06,
+      "loss": 0.4382,
+      "step": 3213
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.4299364691916773e-06,
+      "loss": 0.4209,
+      "step": 3214
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.4256143987315253e-06,
+      "loss": 0.438,
+      "step": 3215
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.4212983686348913e-06,
+      "loss": 0.4211,
+      "step": 3216
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.4169883819422736e-06,
+      "loss": 0.4292,
+      "step": 3217
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.4126844416899132e-06,
+      "loss": 0.4377,
+      "step": 3218
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.4083865509097916e-06,
+      "loss": 0.4248,
+      "step": 3219
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.404094712629621e-06,
+      "loss": 0.4175,
+      "step": 3220
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.3998089298728624e-06,
+      "loss": 0.4185,
+      "step": 3221
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.3955292056587e-06,
+      "loss": 0.4141,
+      "step": 3222
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.3912555430020501e-06,
+      "loss": 0.4226,
+      "step": 3223
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.3869879449135736e-06,
+      "loss": 0.4348,
+      "step": 3224
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.38272641439964e-06,
+      "loss": 0.4138,
+      "step": 3225
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.3784709544623575e-06,
+      "loss": 0.4258,
+      "step": 3226
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.3742215680995507e-06,
+      "loss": 0.4087,
+      "step": 3227
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.3699782583047672e-06,
+      "loss": 0.4275,
+      "step": 3228
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.3657410280672757e-06,
+      "loss": 0.4521,
+      "step": 3229
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.3615098803720606e-06,
+      "loss": 0.4128,
+      "step": 3230
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.3572848181998255e-06,
+      "loss": 0.4509,
+      "step": 3231
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.3530658445269784e-06,
+      "loss": 0.4023,
+      "step": 3232
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.3488529623256464e-06,
+      "loss": 0.425,
+      "step": 3233
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.3446461745636608e-06,
+      "loss": 0.4014,
+      "step": 3234
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.3404454842045578e-06,
+      "loss": 0.4138,
+      "step": 3235
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.3362508942075836e-06,
+      "loss": 0.4272,
+      "step": 3236
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.3320624075276844e-06,
+      "loss": 0.4214,
+      "step": 3237
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.327880027115509e-06,
+      "loss": 0.4304,
+      "step": 3238
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.3237037559174016e-06,
+      "loss": 0.4187,
+      "step": 3239
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.3195335968753985e-06,
+      "loss": 0.4114,
+      "step": 3240
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.3153695529272415e-06,
+      "loss": 0.4236,
+      "step": 3241
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.3112116270063512e-06,
+      "loss": 0.4014,
+      "step": 3242
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.307059822041854e-06,
+      "loss": 0.4133,
+      "step": 3243
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.3029141409585521e-06,
+      "loss": 0.4023,
+      "step": 3244
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.2987745866769319e-06,
+      "loss": 0.4141,
+      "step": 3245
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.2946411621131772e-06,
+      "loss": 0.4417,
+      "step": 3246
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.2905138701791376e-06,
+      "loss": 0.4187,
+      "step": 3247
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.286392713782354e-06,
+      "loss": 0.3997,
+      "step": 3248
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.2822776958260397e-06,
+      "loss": 0.4197,
+      "step": 3249
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.278168819209089e-06,
+      "loss": 0.4031,
+      "step": 3250
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.2740660868260634e-06,
+      "loss": 0.4429,
+      "step": 3251
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.2699695015671965e-06,
+      "loss": 0.4111,
+      "step": 3252
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.2658790663183962e-06,
+      "loss": 0.4053,
+      "step": 3253
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.2617947839612333e-06,
+      "loss": 0.4248,
+      "step": 3254
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.2577166573729471e-06,
+      "loss": 0.4121,
+      "step": 3255
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.2536446894264408e-06,
+      "loss": 0.4275,
+      "step": 3256
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.2495788829902733e-06,
+      "loss": 0.4121,
+      "step": 3257
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.2455192409286733e-06,
+      "loss": 0.4163,
+      "step": 3258
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.241465766101515e-06,
+      "loss": 0.4265,
+      "step": 3259
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.237418461364338e-06,
+      "loss": 0.4321,
+      "step": 3260
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.233377329568325e-06,
+      "loss": 0.4363,
+      "step": 3261
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.229342373560325e-06,
+      "loss": 0.4177,
+      "step": 3262
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.2253135961828233e-06,
+      "loss": 0.4419,
+      "step": 3263
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.2212910002739531e-06,
+      "loss": 0.4202,
+      "step": 3264
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.217274588667503e-06,
+      "loss": 0.439,
+      "step": 3265
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.2132643641928931e-06,
+      "loss": 0.4065,
+      "step": 3266
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.2092603296751926e-06,
+      "loss": 0.4216,
+      "step": 3267
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.2052624879351105e-06,
+      "loss": 0.4141,
+      "step": 3268
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.2012708417889852e-06,
+      "loss": 0.4185,
+      "step": 3269
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.1972853940488017e-06,
+      "loss": 0.4097,
+      "step": 3270
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.1933061475221675e-06,
+      "loss": 0.4248,
+      "step": 3271
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.1893331050123314e-06,
+      "loss": 0.4099,
+      "step": 3272
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.185366269318161e-06,
+      "loss": 0.4553,
+      "step": 3273
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.1814056432341625e-06,
+      "loss": 0.4324,
+      "step": 3274
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.1774512295504615e-06,
+      "loss": 0.4414,
+      "step": 3275
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.173503031052806e-06,
+      "loss": 0.4031,
+      "step": 3276
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.1695610505225708e-06,
+      "loss": 0.425,
+      "step": 3277
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.165625290736745e-06,
+      "loss": 0.4265,
+      "step": 3278
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.1616957544679396e-06,
+      "loss": 0.4243,
+      "step": 3279
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.1577724444843764e-06,
+      "loss": 0.4194,
+      "step": 3280
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.1538553635498951e-06,
+      "loss": 0.4187,
+      "step": 3281
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.1499445144239507e-06,
+      "loss": 0.4282,
+      "step": 3282
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.1460398998615974e-06,
+      "loss": 0.406,
+      "step": 3283
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.142141522613508e-06,
+      "loss": 0.4307,
+      "step": 3284
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.1382493854259535e-06,
+      "loss": 0.4087,
+      "step": 3285
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.1343634910408142e-06,
+      "loss": 0.4167,
+      "step": 3286
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.1304838421955732e-06,
+      "loss": 0.4067,
+      "step": 3287
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.126610441623306e-06,
+      "loss": 0.4167,
+      "step": 3288
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.122743292052697e-06,
+      "loss": 0.4336,
+      "step": 3289
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.1188823962080163e-06,
+      "loss": 0.418,
+      "step": 3290
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.1150277568091405e-06,
+      "loss": 0.4124,
+      "step": 3291
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.111179376571525e-06,
+      "loss": 0.4148,
+      "step": 3292
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.107337258206228e-06,
+      "loss": 0.4109,
+      "step": 3293
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.1035014044198921e-06,
+      "loss": 0.4045,
+      "step": 3294
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.0996718179147413e-06,
+      "loss": 0.4351,
+      "step": 3295
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.0958485013885944e-06,
+      "loss": 0.4294,
+      "step": 3296
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.0920314575348456e-06,
+      "loss": 0.3984,
+      "step": 3297
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.0882206890424707e-06,
+      "loss": 0.4229,
+      "step": 3298
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.084416198596029e-06,
+      "loss": 0.4185,
+      "step": 3299
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 1.0806179888756551e-06,
+      "loss": 0.4597,
+      "step": 3300
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 1.0768260625570592e-06,
+      "loss": 0.4312,
+      "step": 3301
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 1.0730404223115243e-06,
+      "loss": 0.4045,
+      "step": 3302
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 1.0692610708059037e-06,
+      "loss": 0.4153,
+      "step": 3303
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 1.0654880107026222e-06,
+      "loss": 0.417,
+      "step": 3304
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 1.0617212446596736e-06,
+      "loss": 0.4082,
+      "step": 3305
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 1.0579607753306198e-06,
+      "loss": 0.449,
+      "step": 3306
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 1.0542066053645784e-06,
+      "loss": 0.4194,
+      "step": 3307
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 1.0504587374062392e-06,
+      "loss": 0.4287,
+      "step": 3308
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 1.0467171740958449e-06,
+      "loss": 0.3904,
+      "step": 3309
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 1.042981918069199e-06,
+      "loss": 0.4158,
+      "step": 3310
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 1.039252971957665e-06,
+      "loss": 0.4451,
+      "step": 3311
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 1.0355303383881577e-06,
+      "loss": 0.4248,
+      "step": 3312
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 1.0318140199831473e-06,
+      "loss": 0.4106,
+      "step": 3313
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 1.028104019360654e-06,
+      "loss": 0.4153,
+      "step": 3314
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 1.0244003391342439e-06,
+      "loss": 0.4106,
+      "step": 3315
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 1.0207029819130388e-06,
+      "loss": 0.4167,
+      "step": 3316
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 1.017011950301696e-06,
+      "loss": 0.4277,
+      "step": 3317
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 1.0133272469004264e-06,
+      "loss": 0.4255,
+      "step": 3318
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 1.009648874304976e-06,
+      "loss": 0.4011,
+      "step": 3319
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 1.0059768351066356e-06,
+      "loss": 0.4373,
+      "step": 3320
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 1.0023111318922308e-06,
+      "loss": 0.4446,
+      "step": 3321
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 9.986517672441243e-07,
+      "loss": 0.4253,
+      "step": 3322
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 9.949987437402132e-07,
+      "loss": 0.4087,
+      "step": 3323
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 9.913520639539321e-07,
+      "loss": 0.4128,
+      "step": 3324
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 9.877117304542416e-07,
+      "loss": 0.4006,
+      "step": 3325
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 9.840777458056339e-07,
+      "loss": 0.4348,
+      "step": 3326
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 9.804501125681243e-07,
+      "loss": 0.4158,
+      "step": 3327
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 9.76828833297261e-07,
+      "loss": 0.4216,
+      "step": 3328
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 9.732139105441097e-07,
+      "loss": 0.3823,
+      "step": 3329
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 9.696053468552614e-07,
+      "loss": 0.4258,
+      "step": 3330
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 9.660031447728292e-07,
+      "loss": 0.4041,
+      "step": 3331
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 9.624073068344385e-07,
+      "loss": 0.4285,
+      "step": 3332
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 9.58817835573238e-07,
+      "loss": 0.4314,
+      "step": 3333
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 9.552347335178858e-07,
+      "loss": 0.4355,
+      "step": 3334
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 9.516580031925582e-07,
+      "loss": 0.4424,
+      "step": 3335
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 9.480876471169376e-07,
+      "loss": 0.4104,
+      "step": 3336
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 9.445236678062197e-07,
+      "loss": 0.4414,
+      "step": 3337
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 9.409660677711085e-07,
+      "loss": 0.4128,
+      "step": 3338
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 9.374148495178104e-07,
+      "loss": 0.4324,
+      "step": 3339
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 9.338700155480396e-07,
+      "loss": 0.4131,
+      "step": 3340
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 9.303315683590086e-07,
+      "loss": 0.396,
+      "step": 3341
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 9.267995104434357e-07,
+      "loss": 0.4062,
+      "step": 3342
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 9.232738442895362e-07,
+      "loss": 0.4275,
+      "step": 3343
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 9.197545723810197e-07,
+      "loss": 0.4348,
+      "step": 3344
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 9.162416971970978e-07,
+      "loss": 0.436,
+      "step": 3345
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 9.127352212124663e-07,
+      "loss": 0.4377,
+      "step": 3346
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 9.092351468973248e-07,
+      "loss": 0.4409,
+      "step": 3347
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 9.0574147671735e-07,
+      "loss": 0.4133,
+      "step": 3348
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 9.022542131337186e-07,
+      "loss": 0.4448,
+      "step": 3349
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.987733586030889e-07,
+      "loss": 0.428,
+      "step": 3350
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.952989155776026e-07,
+      "loss": 0.4441,
+      "step": 3351
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.918308865048897e-07,
+      "loss": 0.4072,
+      "step": 3352
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.883692738280558e-07,
+      "loss": 0.4116,
+      "step": 3353
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.849140799856936e-07,
+      "loss": 0.4319,
+      "step": 3354
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.814653074118651e-07,
+      "loss": 0.4287,
+      "step": 3355
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.780229585361155e-07,
+      "loss": 0.4082,
+      "step": 3356
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.745870357834652e-07,
+      "loss": 0.4248,
+      "step": 3357
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.711575415744012e-07,
+      "loss": 0.4458,
+      "step": 3358
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.677344783248887e-07,
+      "loss": 0.4058,
+      "step": 3359
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.643178484463566e-07,
+      "loss": 0.4055,
+      "step": 3360
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.609076543457062e-07,
+      "loss": 0.3901,
+      "step": 3361
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.575038984253048e-07,
+      "loss": 0.4236,
+      "step": 3362
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.541065830829809e-07,
+      "loss": 0.4153,
+      "step": 3363
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.507157107120301e-07,
+      "loss": 0.4092,
+      "step": 3364
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.473312837012027e-07,
+      "loss": 0.4045,
+      "step": 3365
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.43953304434717e-07,
+      "loss": 0.427,
+      "step": 3366
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.405817752922418e-07,
+      "loss": 0.4148,
+      "step": 3367
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.372166986489061e-07,
+      "loss": 0.4194,
+      "step": 3368
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.338580768752935e-07,
+      "loss": 0.4182,
+      "step": 3369
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.305059123374348e-07,
+      "loss": 0.4163,
+      "step": 3370
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.271602073968221e-07,
+      "loss": 0.417,
+      "step": 3371
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.238209644103867e-07,
+      "loss": 0.425,
+      "step": 3372
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.204881857305113e-07,
+      "loss": 0.4287,
+      "step": 3373
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.171618737050269e-07,
+      "loss": 0.4204,
+      "step": 3374
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.138420306772066e-07,
+      "loss": 0.3977,
+      "step": 3375
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 8.105286589857708e-07,
+      "loss": 0.4014,
+      "step": 3376
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 8.072217609648714e-07,
+      "loss": 0.4341,
+      "step": 3377
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 8.039213389441114e-07,
+      "loss": 0.4373,
+      "step": 3378
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 8.00627395248521e-07,
+      "loss": 0.4263,
+      "step": 3379
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.973399321985753e-07,
+      "loss": 0.4133,
+      "step": 3380
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.940589521101805e-07,
+      "loss": 0.4473,
+      "step": 3381
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.907844572946721e-07,
+      "loss": 0.4219,
+      "step": 3382
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.875164500588239e-07,
+      "loss": 0.4116,
+      "step": 3383
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.842549327048366e-07,
+      "loss": 0.4189,
+      "step": 3384
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.809999075303332e-07,
+      "loss": 0.4238,
+      "step": 3385
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.777513768283729e-07,
+      "loss": 0.4192,
+      "step": 3386
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.745093428874351e-07,
+      "loss": 0.4294,
+      "step": 3387
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.712738079914228e-07,
+      "loss": 0.4072,
+      "step": 3388
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.680447744196595e-07,
+      "loss": 0.4321,
+      "step": 3389
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.648222444468878e-07,
+      "loss": 0.4307,
+      "step": 3390
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.616062203432739e-07,
+      "loss": 0.4109,
+      "step": 3391
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.583967043743956e-07,
+      "loss": 0.4141,
+      "step": 3392
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.551936988012487e-07,
+      "loss": 0.4131,
+      "step": 3393
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.519972058802405e-07,
+      "loss": 0.3862,
+      "step": 3394
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.488072278631963e-07,
+      "loss": 0.3999,
+      "step": 3395
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.456237669973432e-07,
+      "loss": 0.4121,
+      "step": 3396
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.424468255253203e-07,
+      "loss": 0.3972,
+      "step": 3397
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.392764056851764e-07,
+      "loss": 0.4177,
+      "step": 3398
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.361125097103661e-07,
+      "loss": 0.4226,
+      "step": 3399
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.329551398297463e-07,
+      "loss": 0.4316,
+      "step": 3400
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.298042982675768e-07,
+      "loss": 0.4255,
+      "step": 3401
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.266599872435165e-07,
+      "loss": 0.3965,
+      "step": 3402
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.23522208972628e-07,
+      "loss": 0.4207,
+      "step": 3403
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.203909656653674e-07,
+      "loss": 0.4189,
+      "step": 3404
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.172662595275914e-07,
+      "loss": 0.4141,
+      "step": 3405
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.141480927605471e-07,
+      "loss": 0.4116,
+      "step": 3406
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.110364675608805e-07,
+      "loss": 0.4353,
+      "step": 3407
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.07931386120625e-07,
+      "loss": 0.4236,
+      "step": 3408
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.048328506272017e-07,
+      "loss": 0.4231,
+      "step": 3409
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 7.01740863263427e-07,
+      "loss": 0.4094,
+      "step": 3410
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 6.986554262074996e-07,
+      "loss": 0.4045,
+      "step": 3411
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 6.955765416330052e-07,
+      "loss": 0.4385,
+      "step": 3412
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 6.925042117089142e-07,
+      "loss": 0.4004,
+      "step": 3413
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 6.894384385995778e-07,
+      "loss": 0.4167,
+      "step": 3414
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 6.863792244647305e-07,
+      "loss": 0.427,
+      "step": 3415
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.833265714594828e-07,
+      "loss": 0.4478,
+      "step": 3416
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.802804817343256e-07,
+      "loss": 0.4294,
+      "step": 3417
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.77240957435128e-07,
+      "loss": 0.4421,
+      "step": 3418
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.742080007031282e-07,
+      "loss": 0.4399,
+      "step": 3419
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.711816136749449e-07,
+      "loss": 0.4158,
+      "step": 3420
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.681617984825617e-07,
+      "loss": 0.4128,
+      "step": 3421
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.651485572533379e-07,
+      "loss": 0.4143,
+      "step": 3422
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.621418921099976e-07,
+      "loss": 0.4336,
+      "step": 3423
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.591418051706355e-07,
+      "loss": 0.4258,
+      "step": 3424
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.561482985487122e-07,
+      "loss": 0.4197,
+      "step": 3425
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.531613743530496e-07,
+      "loss": 0.417,
+      "step": 3426
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.501810346878368e-07,
+      "loss": 0.4048,
+      "step": 3427
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.472072816526176e-07,
+      "loss": 0.4307,
+      "step": 3428
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.442401173423051e-07,
+      "loss": 0.4231,
+      "step": 3429
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.412795438471619e-07,
+      "loss": 0.437,
+      "step": 3430
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.383255632528129e-07,
+      "loss": 0.4365,
+      "step": 3431
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.35378177640239e-07,
+      "loss": 0.4292,
+      "step": 3432
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.324373890857694e-07,
+      "loss": 0.4424,
+      "step": 3433
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.295031996610934e-07,
+      "loss": 0.427,
+      "step": 3434
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.265756114332455e-07,
+      "loss": 0.4487,
+      "step": 3435
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.236546264646126e-07,
+      "loss": 0.4114,
+      "step": 3436
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.207402468129309e-07,
+      "loss": 0.4099,
+      "step": 3437
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.178324745312792e-07,
+      "loss": 0.4326,
+      "step": 3438
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.149313116680889e-07,
+      "loss": 0.4216,
+      "step": 3439
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.120367602671251e-07,
+      "loss": 0.4387,
+      "step": 3440
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.091488223675058e-07,
+      "loss": 0.4451,
+      "step": 3441
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.062675000036811e-07,
+      "loss": 0.4368,
+      "step": 3442
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.033927952054475e-07,
+      "loss": 0.4336,
+      "step": 3443
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 6.005247099979372e-07,
+      "loss": 0.446,
+      "step": 3444
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 5.976632464016163e-07,
+      "loss": 0.4331,
+      "step": 3445
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 5.948084064322912e-07,
+      "loss": 0.4272,
+      "step": 3446
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 5.919601921010998e-07,
+      "loss": 0.4551,
+      "step": 3447
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 5.891186054145103e-07,
+      "loss": 0.4041,
+      "step": 3448
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 5.86283648374324e-07,
+      "loss": 0.4302,
+      "step": 3449
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 5.83455322977674e-07,
+      "loss": 0.4104,
+      "step": 3450
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 5.806336312170191e-07,
+      "loss": 0.416,
+      "step": 3451
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 5.778185750801435e-07,
+      "loss": 0.4297,
+      "step": 3452
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 5.75010156550162e-07,
+      "loss": 0.4329,
+      "step": 3453
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.722083776055076e-07,
+      "loss": 0.4219,
+      "step": 3454
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.69413240219936e-07,
+      "loss": 0.426,
+      "step": 3455
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.666247463625319e-07,
+      "loss": 0.4072,
+      "step": 3456
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.638428979976906e-07,
+      "loss": 0.4146,
+      "step": 3457
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.610676970851336e-07,
+      "loss": 0.416,
+      "step": 3458
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.582991455798925e-07,
+      "loss": 0.4216,
+      "step": 3459
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.555372454323182e-07,
+      "loss": 0.4275,
+      "step": 3460
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.527819985880755e-07,
+      "loss": 0.3972,
+      "step": 3461
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.50033406988143e-07,
+      "loss": 0.429,
+      "step": 3462
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.472914725688117e-07,
+      "loss": 0.3999,
+      "step": 3463
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.445561972616786e-07,
+      "loss": 0.3997,
+      "step": 3464
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.418275829936537e-07,
+      "loss": 0.4187,
+      "step": 3465
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.391056316869536e-07,
+      "loss": 0.4277,
+      "step": 3466
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.363903452590979e-07,
+      "loss": 0.4302,
+      "step": 3467
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.336817256229154e-07,
+      "loss": 0.4326,
+      "step": 3468
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.309797746865364e-07,
+      "loss": 0.3958,
+      "step": 3469
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.28284494353396e-07,
+      "loss": 0.4167,
+      "step": 3470
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.255958865222255e-07,
+      "loss": 0.4414,
+      "step": 3471
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.229139530870553e-07,
+      "loss": 0.405,
+      "step": 3472
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.202386959372218e-07,
+      "loss": 0.4121,
+      "step": 3473
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.175701169573477e-07,
+      "loss": 0.4294,
+      "step": 3474
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.149082180273601e-07,
+      "loss": 0.4204,
+      "step": 3475
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.122530010224757e-07,
+      "loss": 0.4141,
+      "step": 3476
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.096044678132028e-07,
+      "loss": 0.4102,
+      "step": 3477
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.069626202653455e-07,
+      "loss": 0.4121,
+      "step": 3478
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.043274602399939e-07,
+      "loss": 0.4402,
+      "step": 3479
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 5.016989895935287e-07,
+      "loss": 0.396,
+      "step": 3480
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 4.990772101776187e-07,
+      "loss": 0.4136,
+      "step": 3481
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 4.964621238392208e-07,
+      "loss": 0.4236,
+      "step": 3482
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 4.938537324205717e-07,
+      "loss": 0.4019,
+      "step": 3483
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 4.912520377591945e-07,
+      "loss": 0.4109,
+      "step": 3484
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 4.886570416878966e-07,
+      "loss": 0.4282,
+      "step": 3485
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 4.860687460347624e-07,
+      "loss": 0.417,
+      "step": 3486
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 4.834871526231588e-07,
+      "loss": 0.4158,
+      "step": 3487
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 4.809122632717345e-07,
+      "loss": 0.4158,
+      "step": 3488
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 4.783440797944061e-07,
+      "loss": 0.4185,
+      "step": 3489
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 4.7578260400037546e-07,
+      "loss": 0.4324,
+      "step": 3490
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 4.732278376941124e-07,
+      "loss": 0.426,
+      "step": 3491
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 4.706797826753662e-07,
+      "loss": 0.3958,
+      "step": 3492
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.681384407391498e-07,
+      "loss": 0.3943,
+      "step": 3493
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.6560381367576013e-07,
+      "loss": 0.4285,
+      "step": 3494
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.63075903270751e-07,
+      "loss": 0.4114,
+      "step": 3495
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.605547113049491e-07,
+      "loss": 0.4231,
+      "step": 3496
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.580402395544503e-07,
+      "loss": 0.4094,
+      "step": 3497
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.555324897906133e-07,
+      "loss": 0.3994,
+      "step": 3498
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.530314637800637e-07,
+      "loss": 0.4265,
+      "step": 3499
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.5053716328469e-07,
+      "loss": 0.3933,
+      "step": 3500
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.4804959006163995e-07,
+      "loss": 0.4197,
+      "step": 3501
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.455687458633262e-07,
+      "loss": 0.4185,
+      "step": 3502
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.4309463243741856e-07,
+      "loss": 0.4153,
+      "step": 3503
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.4062725152684616e-07,
+      "loss": 0.4231,
+      "step": 3504
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.3816660486979524e-07,
+      "loss": 0.4163,
+      "step": 3505
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.3571269419970586e-07,
+      "loss": 0.4238,
+      "step": 3506
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.3326552124527964e-07,
+      "loss": 0.3901,
+      "step": 3507
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.3082508773046096e-07,
+      "loss": 0.3889,
+      "step": 3508
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.283913953744567e-07,
+      "loss": 0.3916,
+      "step": 3509
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.259644458917178e-07,
+      "loss": 0.3884,
+      "step": 3510
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.2354424099195104e-07,
+      "loss": 0.4351,
+      "step": 3511
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.2113078238010364e-07,
+      "loss": 0.429,
+      "step": 3512
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.1872407175637896e-07,
+      "loss": 0.4285,
+      "step": 3513
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.16324110816223e-07,
+      "loss": 0.3962,
+      "step": 3514
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.1393090125032453e-07,
+      "loss": 0.4048,
+      "step": 3515
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.115444447446204e-07,
+      "loss": 0.3962,
+      "step": 3516
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.0916474298028694e-07,
+      "loss": 0.4109,
+      "step": 3517
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.0679179763374323e-07,
+      "loss": 0.3984,
+      "step": 3518
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.0442561037665087e-07,
+      "loss": 0.396,
+      "step": 3519
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 4.0206618287590536e-07,
+      "loss": 0.407,
+      "step": 3520
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 3.997135167936461e-07,
+      "loss": 0.4009,
+      "step": 3521
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 3.973676137872429e-07,
+      "loss": 0.4026,
+      "step": 3522
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 3.9502847550930925e-07,
+      "loss": 0.3984,
+      "step": 3523
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 3.9269610360768374e-07,
+      "loss": 0.4246,
+      "step": 3524
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 3.903704997254465e-07,
+      "loss": 0.3987,
+      "step": 3525
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 3.880516655009081e-07,
+      "loss": 0.4219,
+      "step": 3526
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 3.85739602567603e-07,
+      "loss": 0.4082,
+      "step": 3527
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 3.834343125543061e-07,
+      "loss": 0.4189,
+      "step": 3528
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 3.8113579708501493e-07,
+      "loss": 0.4238,
+      "step": 3529
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 3.7884405777895206e-07,
+      "loss": 0.4009,
+      "step": 3530
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.765590962505739e-07,
+      "loss": 0.4116,
+      "step": 3531
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.74280914109556e-07,
+      "loss": 0.4355,
+      "step": 3532
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.720095129608037e-07,
+      "loss": 0.4067,
+      "step": 3533
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.6974489440443795e-07,
+      "loss": 0.4119,
+      "step": 3534
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.6748706003580826e-07,
+      "loss": 0.3994,
+      "step": 3535
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.6523601144548003e-07,
+      "loss": 0.4155,
+      "step": 3536
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.629917502192415e-07,
+      "loss": 0.4312,
+      "step": 3537
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.6075427793810016e-07,
+      "loss": 0.4099,
+      "step": 3538
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.585235961782774e-07,
+      "loss": 0.4214,
+      "step": 3539
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.5629970651121506e-07,
+      "loss": 0.4253,
+      "step": 3540
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.540826105035644e-07,
+      "loss": 0.4019,
+      "step": 3541
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.518723097171961e-07,
+      "loss": 0.4058,
+      "step": 3542
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.4966880570919235e-07,
+      "loss": 0.3901,
+      "step": 3543
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.4747210003184597e-07,
+      "loss": 0.4326,
+      "step": 3544
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.452821942326623e-07,
+      "loss": 0.4224,
+      "step": 3545
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.4309908985435626e-07,
+      "loss": 0.4224,
+      "step": 3546
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.4092278843484873e-07,
+      "loss": 0.4216,
+      "step": 3547
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.3875329150727333e-07,
+      "loss": 0.4148,
+      "step": 3548
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.3659060059996305e-07,
+      "loss": 0.4253,
+      "step": 3549
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.344347172364615e-07,
+      "loss": 0.4302,
+      "step": 3550
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.32285642935517e-07,
+      "loss": 0.4282,
+      "step": 3551
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.3014337921107864e-07,
+      "loss": 0.4124,
+      "step": 3552
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.2800792757229917e-07,
+      "loss": 0.4204,
+      "step": 3553
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.2587928952352853e-07,
+      "loss": 0.4141,
+      "step": 3554
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.2375746656432284e-07,
+      "loss": 0.4216,
+      "step": 3555
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.2164246018943414e-07,
+      "loss": 0.3931,
+      "step": 3556
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.1953427188881283e-07,
+      "loss": 0.425,
+      "step": 3557
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.1743290314760534e-07,
+      "loss": 0.4224,
+      "step": 3558
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.153383554461531e-07,
+      "loss": 0.4202,
+      "step": 3559
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.132506302599969e-07,
+      "loss": 0.4089,
+      "step": 3560
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.111697290598659e-07,
+      "loss": 0.3982,
+      "step": 3561
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.09095653311684e-07,
+      "loss": 0.4092,
+      "step": 3562
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.070284044765692e-07,
+      "loss": 0.4036,
+      "step": 3563
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.0496798401082773e-07,
+      "loss": 0.4346,
+      "step": 3564
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.029143933659562e-07,
+      "loss": 0.4089,
+      "step": 3565
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 3.0086763398863627e-07,
+      "loss": 0.4136,
+      "step": 3566
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 2.988277073207446e-07,
+      "loss": 0.3967,
+      "step": 3567
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 2.967946147993372e-07,
+      "loss": 0.4131,
+      "step": 3568
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 2.947683578566596e-07,
+      "loss": 0.4058,
+      "step": 3569
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.9274893792014315e-07,
+      "loss": 0.4082,
+      "step": 3570
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.907363564123977e-07,
+      "loss": 0.3979,
+      "step": 3571
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.887306147512203e-07,
+      "loss": 0.4089,
+      "step": 3572
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.867317143495862e-07,
+      "loss": 0.3872,
+      "step": 3573
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.8473965661565353e-07,
+      "loss": 0.4153,
+      "step": 3574
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.8275444295276175e-07,
+      "loss": 0.4192,
+      "step": 3575
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.807760747594224e-07,
+      "loss": 0.4128,
+      "step": 3576
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.7880455342933266e-07,
+      "loss": 0.4194,
+      "step": 3577
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.7683988035135834e-07,
+      "loss": 0.4441,
+      "step": 3578
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.7488205690954893e-07,
+      "loss": 0.4087,
+      "step": 3579
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.7293108448312124e-07,
+      "loss": 0.4136,
+      "step": 3580
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.7098696444647135e-07,
+      "loss": 0.4014,
+      "step": 3581
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.690496981691648e-07,
+      "loss": 0.4216,
+      "step": 3582
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.6711928701594003e-07,
+      "loss": 0.4036,
+      "step": 3583
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.6519573234670694e-07,
+      "loss": 0.4351,
+      "step": 3584
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.6327903551654267e-07,
+      "loss": 0.4238,
+      "step": 3585
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.6136919787569714e-07,
+      "loss": 0.4128,
+      "step": 3586
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.5946622076958415e-07,
+      "loss": 0.418,
+      "step": 3587
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.575701055387869e-07,
+      "loss": 0.4199,
+      "step": 3588
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.556808535190558e-07,
+      "loss": 0.4219,
+      "step": 3589
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.5379846604130287e-07,
+      "loss": 0.4202,
+      "step": 3590
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.519229444316085e-07,
+      "loss": 0.4192,
+      "step": 3591
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.5005429001121127e-07,
+      "loss": 0.4255,
+      "step": 3592
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.4819250409651605e-07,
+      "loss": 0.4194,
+      "step": 3593
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.463375879990881e-07,
+      "loss": 0.4136,
+      "step": 3594
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.444895430256511e-07,
+      "loss": 0.4146,
+      "step": 3595
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.4264837047809354e-07,
+      "loss": 0.4177,
+      "step": 3596
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.408140716534546e-07,
+      "loss": 0.4248,
+      "step": 3597
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.3898664784393954e-07,
+      "loss": 0.4043,
+      "step": 3598
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.3716610033690303e-07,
+      "loss": 0.4377,
+      "step": 3599
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.3535243041486133e-07,
+      "loss": 0.4316,
+      "step": 3600
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.3354563935548248e-07,
+      "loss": 0.4189,
+      "step": 3601
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.3174572843158826e-07,
+      "loss": 0.3955,
+      "step": 3602
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.299526989111578e-07,
+      "loss": 0.4478,
+      "step": 3603
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.2816655205731731e-07,
+      "loss": 0.3857,
+      "step": 3604
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.2638728912834584e-07,
+      "loss": 0.4238,
+      "step": 3605
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.246149113776741e-07,
+      "loss": 0.4224,
+      "step": 3606
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.2284942005388331e-07,
+      "loss": 0.4163,
+      "step": 3607
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 2.2109081640070307e-07,
+      "loss": 0.4067,
+      "step": 3608
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 2.1933910165700677e-07,
+      "loss": 0.408,
+      "step": 3609
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 2.1759427705682067e-07,
+      "loss": 0.407,
+      "step": 3610
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 2.1585634382931376e-07,
+      "loss": 0.4321,
+      "step": 3611
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 2.1412530319879888e-07,
+      "loss": 0.4478,
+      "step": 3612
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 2.1240115638473947e-07,
+      "loss": 0.4158,
+      "step": 3613
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 2.1068390460173504e-07,
+      "loss": 0.4041,
+      "step": 3614
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 2.0897354905953237e-07,
+      "loss": 0.3945,
+      "step": 3615
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 2.0727009096301877e-07,
+      "loss": 0.4238,
+      "step": 3616
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 2.0557353151222314e-07,
+      "loss": 0.4309,
+      "step": 3617
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 2.0388387190231173e-07,
+      "loss": 0.4309,
+      "step": 3618
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 2.022011133235946e-07,
+      "loss": 0.4307,
+      "step": 3619
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 2.0052525696151904e-07,
+      "loss": 0.425,
+      "step": 3620
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.9885630399666622e-07,
+      "loss": 0.4194,
+      "step": 3621
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.97194255604759e-07,
+      "loss": 0.4072,
+      "step": 3622
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.9553911295665306e-07,
+      "loss": 0.3901,
+      "step": 3623
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.938908772183401e-07,
+      "loss": 0.4163,
+      "step": 3624
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.9224954955094577e-07,
+      "loss": 0.4021,
+      "step": 3625
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.906151311107296e-07,
+      "loss": 0.4314,
+      "step": 3626
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.8898762304908612e-07,
+      "loss": 0.4165,
+      "step": 3627
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.8736702651253823e-07,
+      "loss": 0.4111,
+      "step": 3628
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.8575334264273935e-07,
+      "loss": 0.439,
+      "step": 3629
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.8414657257647683e-07,
+      "loss": 0.4141,
+      "step": 3630
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.825467174456652e-07,
+      "loss": 0.4207,
+      "step": 3631
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.8095377837734961e-07,
+      "loss": 0.405,
+      "step": 3632
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.793677564937002e-07,
+      "loss": 0.4382,
+      "step": 3633
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.777886529120143e-07,
+      "loss": 0.3989,
+      "step": 3634
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.7621646874471875e-07,
+      "loss": 0.4165,
+      "step": 3635
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.7465120509936207e-07,
+      "loss": 0.4045,
+      "step": 3636
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.7309286307861994e-07,
+      "loss": 0.3835,
+      "step": 3637
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.7154144378029202e-07,
+      "loss": 0.4038,
+      "step": 3638
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.6999694829730185e-07,
+      "loss": 0.4045,
+      "step": 3639
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.684593777176924e-07,
+      "loss": 0.4121,
+      "step": 3640
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.6692873312462942e-07,
+      "loss": 0.4048,
+      "step": 3641
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.6540501559640044e-07,
+      "loss": 0.3838,
+      "step": 3642
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.6388822620641342e-07,
+      "loss": 0.4287,
+      "step": 3643
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.6237836602319367e-07,
+      "loss": 0.4302,
+      "step": 3644
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.6087543611038926e-07,
+      "loss": 0.4294,
+      "step": 3645
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.5937943752675987e-07,
+      "loss": 0.3899,
+      "step": 3646
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.5789037132618812e-07,
+      "loss": 0.4333,
+      "step": 3647
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.5640823855766817e-07,
+      "loss": 0.4207,
+      "step": 3648
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.5493304026531485e-07,
+      "loss": 0.425,
+      "step": 3649
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.5346477748835353e-07,
+      "loss": 0.4082,
+      "step": 3650
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.5200345126112682e-07,
+      "loss": 0.4338,
+      "step": 3651
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.5054906261308898e-07,
+      "loss": 0.4204,
+      "step": 3652
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.4910161256880605e-07,
+      "loss": 0.4133,
+      "step": 3653
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.4766110214795904e-07,
+      "loss": 0.4175,
+      "step": 3654
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.4622753236533728e-07,
+      "loss": 0.4031,
+      "step": 3655
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.4480090423084403e-07,
+      "loss": 0.4165,
+      "step": 3656
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.4338121874948874e-07,
+      "loss": 0.4236,
+      "step": 3657
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.4196847692139137e-07,
+      "loss": 0.4172,
+      "step": 3658
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.405626797417825e-07,
+      "loss": 0.4282,
+      "step": 3659
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.391638282009966e-07,
+      "loss": 0.4207,
+      "step": 3660
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.3777192328447875e-07,
+      "loss": 0.4001,
+      "step": 3661
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.3638696597277678e-07,
+      "loss": 0.4065,
+      "step": 3662
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.3500895724154695e-07,
+      "loss": 0.4282,
+      "step": 3663
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.3363789806155047e-07,
+      "loss": 0.394,
+      "step": 3664
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.3227378939865033e-07,
+      "loss": 0.4182,
+      "step": 3665
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.3091663221381557e-07,
+      "loss": 0.4395,
+      "step": 3666
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.295664274631181e-07,
+      "loss": 0.4128,
+      "step": 3667
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.282231760977304e-07,
+      "loss": 0.4131,
+      "step": 3668
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.2688687906392772e-07,
+      "loss": 0.4241,
+      "step": 3669
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.2555753730308594e-07,
+      "loss": 0.407,
+      "step": 3670
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.2423515175168154e-07,
+      "loss": 0.4097,
+      "step": 3671
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.2291972334128931e-07,
+      "loss": 0.3936,
+      "step": 3672
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.2161125299858467e-07,
+      "loss": 0.4009,
+      "step": 3673
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.2030974164534025e-07,
+      "loss": 0.4412,
+      "step": 3674
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.19015190198426e-07,
+      "loss": 0.4126,
+      "step": 3675
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.1772759956981127e-07,
+      "loss": 0.4292,
+      "step": 3676
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.164469706665583e-07,
+      "loss": 0.4121,
+      "step": 3677
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.1517330439082763e-07,
+      "loss": 0.3982,
+      "step": 3678
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.1390660163987377e-07,
+      "loss": 0.4114,
+      "step": 3679
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.1264686330604512e-07,
+      "loss": 0.4211,
+      "step": 3680
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.1139409027678405e-07,
+      "loss": 0.4331,
+      "step": 3681
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.1014828343462791e-07,
+      "loss": 0.4111,
+      "step": 3682
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.0890944365720579e-07,
+      "loss": 0.4226,
+      "step": 3683
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.0767757181723737e-07,
+      "loss": 0.4189,
+      "step": 3684
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.0645266878253402e-07,
+      "loss": 0.4238,
+      "step": 3685
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 1.0523473541599883e-07,
+      "loss": 0.3914,
+      "step": 3686
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 1.0402377257562435e-07,
+      "loss": 0.417,
+      "step": 3687
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 1.0281978111449375e-07,
+      "loss": 0.4048,
+      "step": 3688
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 1.0162276188077636e-07,
+      "loss": 0.4136,
+      "step": 3689
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 1.0043271571773316e-07,
+      "loss": 0.415,
+      "step": 3690
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 9.924964346371136e-08,
+      "loss": 0.4285,
+      "step": 3691
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 9.807354595214313e-08,
+      "loss": 0.4133,
+      "step": 3692
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 9.690442401155131e-08,
+      "loss": 0.4221,
+      "step": 3693
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 9.574227846554151e-08,
+      "loss": 0.4121,
+      "step": 3694
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 9.458711013280775e-08,
+      "loss": 0.4519,
+      "step": 3695
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 9.34389198271235e-08,
+      "loss": 0.4163,
+      "step": 3696
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 9.229770835735286e-08,
+      "loss": 0.427,
+      "step": 3697
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 9.11634765274394e-08,
+      "loss": 0.4077,
+      "step": 3698
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 9.003622513641064e-08,
+      "loss": 0.425,
+      "step": 3699
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 8.891595497837691e-08,
+      "loss": 0.3975,
+      "step": 3700
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 8.780266684253025e-08,
+      "loss": 0.3921,
+      "step": 3701
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 8.669636151314664e-08,
+      "loss": 0.4109,
+      "step": 3702
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 8.559703976957822e-08,
+      "loss": 0.4268,
+      "step": 3703
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 8.450470238625996e-08,
+      "loss": 0.4111,
+      "step": 3704
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 8.341935013270852e-08,
+      "loss": 0.4229,
+      "step": 3705
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 8.23409837735134e-08,
+      "loss": 0.4148,
+      "step": 3706
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 8.12696040683525e-08,
+      "loss": 0.4214,
+      "step": 3707
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 8.020521177197204e-08,
+      "loss": 0.4353,
+      "step": 3708
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 7.914780763420338e-08,
+      "loss": 0.4241,
+      "step": 3709
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 7.809739239995062e-08,
+      "loss": 0.4067,
+      "step": 3710
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 7.705396680919409e-08,
+      "loss": 0.4236,
+      "step": 3711
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 7.601753159699354e-08,
+      "loss": 0.4136,
+      "step": 3712
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 7.498808749348274e-08,
+      "loss": 0.4353,
+      "step": 3713
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 7.396563522386824e-08,
+      "loss": 0.4175,
+      "step": 3714
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 7.295017550843497e-08,
+      "loss": 0.4116,
+      "step": 3715
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 7.194170906253851e-08,
+      "loss": 0.416,
+      "step": 3716
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 7.094023659660942e-08,
+      "loss": 0.4189,
+      "step": 3717
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 6.994575881615118e-08,
+      "loss": 0.4087,
+      "step": 3718
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 6.895827642174002e-08,
+      "loss": 0.4163,
+      "step": 3719
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 6.797779010902395e-08,
+      "loss": 0.4114,
+      "step": 3720
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 6.700430056872043e-08,
+      "loss": 0.4358,
+      "step": 3721
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 6.603780848662201e-08,
+      "loss": 0.3972,
+      "step": 3722
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 6.507831454358737e-08,
+      "loss": 0.4189,
+      "step": 3723
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 6.41258194155503e-08,
+      "loss": 0.4075,
+      "step": 3724
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 6.318032377350847e-08,
+      "loss": 0.4082,
+      "step": 3725
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 6.224182828353243e-08,
+      "loss": 0.4294,
+      "step": 3726
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 6.131033360676219e-08,
+      "loss": 0.4004,
+      "step": 3727
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 6.038584039940065e-08,
+      "loss": 0.4182,
+      "step": 3728
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 5.946834931272572e-08,
+      "loss": 0.4053,
+      "step": 3729
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 5.855786099307481e-08,
+      "loss": 0.4209,
+      "step": 3730
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 5.76543760818582e-08,
+      "loss": 0.4094,
+      "step": 3731
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 5.67578952155512e-08,
+      "loss": 0.4353,
+      "step": 3732
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 5.586841902569196e-08,
+      "loss": 0.4246,
+      "step": 3733
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 5.4985948138888135e-08,
+      "loss": 0.4133,
+      "step": 3734
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 5.4110483176810224e-08,
+      "loss": 0.4373,
+      "step": 3735
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 5.324202475619156e-08,
+      "loss": 0.4153,
+      "step": 3736
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 5.2380573488833855e-08,
+      "loss": 0.4434,
+      "step": 3737
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 5.152612998159945e-08,
+      "loss": 0.4131,
+      "step": 3738
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 5.067869483641574e-08,
+      "loss": 0.4011,
+      "step": 3739
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 4.983826865026964e-08,
+      "loss": 0.3787,
+      "step": 3740
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 4.900485201521532e-08,
+      "loss": 0.4167,
+      "step": 3741
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 4.817844551836426e-08,
+      "loss": 0.4133,
+      "step": 3742
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 4.735904974189409e-08,
+      "loss": 0.4241,
+      "step": 3743
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 4.6546665263040856e-08,
+      "loss": 0.4182,
+      "step": 3744
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 4.57412926541001e-08,
+      "loss": 0.4336,
+      "step": 3745
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 4.494293248243131e-08,
+      "loss": 0.4277,
+      "step": 3746
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 4.4151585310451274e-08,
+      "loss": 0.4182,
+      "step": 3747
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 4.336725169563738e-08,
+      "loss": 0.3994,
+      "step": 3748
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 4.258993219052543e-08,
+      "loss": 0.4094,
+      "step": 3749
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 4.181962734271183e-08,
+      "loss": 0.4192,
+      "step": 3750
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 4.105633769484918e-08,
+      "loss": 0.4297,
+      "step": 3751
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 4.030006378464957e-08,
+      "loss": 0.4229,
+      "step": 3752
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 3.9550806144882384e-08,
+      "loss": 0.4153,
+      "step": 3753
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 3.880856530337429e-08,
+      "loss": 0.4092,
+      "step": 3754
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 3.807334178300814e-08,
+      "loss": 0.437,
+      "step": 3755
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 3.7345136101724076e-08,
+      "loss": 0.4021,
+      "step": 3756
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 3.66239487725184e-08,
+      "loss": 0.4114,
+      "step": 3757
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 3.590978030344472e-08,
+      "loss": 0.4011,
+      "step": 3758
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 3.5202631197607254e-08,
+      "loss": 0.4263,
+      "step": 3759
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 3.450250195317084e-08,
+      "loss": 0.4287,
+      "step": 3760
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 3.380939306335318e-08,
+      "loss": 0.4126,
+      "step": 3761
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 3.312330501642369e-08,
+      "loss": 0.4221,
+      "step": 3762
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 3.244423829570909e-08,
+      "loss": 0.4124,
+      "step": 3763
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 3.177219337958892e-08,
+      "loss": 0.4106,
+      "step": 3764
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 3.1107170741497827e-08,
+      "loss": 0.4124,
+      "step": 3765
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 3.044917084991883e-08,
+      "loss": 0.405,
+      "step": 3766
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 2.9798194168391138e-08,
+      "loss": 0.4167,
+      "step": 3767
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 2.91542411555068e-08,
+      "loss": 0.4216,
+      "step": 3768
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 2.8517312264908502e-08,
+      "loss": 0.4277,
+      "step": 3769
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 2.7887407945291768e-08,
+      "loss": 0.4202,
+      "step": 3770
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 2.7264528640400522e-08,
+      "loss": 0.4167,
+      "step": 3771
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 2.6648674789033767e-08,
+      "loss": 0.4207,
+      "step": 3772
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 2.6039846825040015e-08,
+      "loss": 0.3889,
+      "step": 3773
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 2.5438045177316183e-08,
+      "loss": 0.4119,
+      "step": 3774
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 2.4843270269813147e-08,
+      "loss": 0.4346,
+      "step": 3775
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 2.4255522521526853e-08,
+      "loss": 0.396,
+      "step": 3776
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 2.3674802346509428e-08,
+      "loss": 0.4114,
+      "step": 3777
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 2.3101110153855854e-08,
+      "loss": 0.4246,
+      "step": 3778
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 2.2534446347712847e-08,
+      "loss": 0.4167,
+      "step": 3779
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 2.1974811327276636e-08,
+      "loss": 0.3931,
+      "step": 3780
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 2.142220548679075e-08,
+      "loss": 0.384,
+      "step": 3781
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 2.0876629215547118e-08,
+      "loss": 0.4192,
+      "step": 3782
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 2.033808289788608e-08,
+      "loss": 0.4331,
+      "step": 3783
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 1.9806566913195267e-08,
+      "loss": 0.4053,
+      "step": 3784
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 1.928208163590961e-08,
+      "loss": 0.4189,
+      "step": 3785
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 1.876462743551022e-08,
+      "loss": 0.3984,
+      "step": 3786
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 1.825420467652772e-08,
+      "loss": 0.4434,
+      "step": 3787
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 1.77508137185356e-08,
+      "loss": 0.4426,
+      "step": 3788
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 1.725445491615796e-08,
+      "loss": 0.4148,
+      "step": 3789
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 1.6765128619060656e-08,
+      "loss": 0.4214,
+      "step": 3790
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 1.6282835171959055e-08,
+      "loss": 0.4226,
+      "step": 3791
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 1.5807574914611378e-08,
+      "loss": 0.4192,
+      "step": 3792
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 1.5339348181823144e-08,
+      "loss": 0.4424,
+      "step": 3793
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 1.487815530344272e-08,
+      "loss": 0.4016,
+      "step": 3794
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 1.4423996604366886e-08,
+      "loss": 0.4087,
+      "step": 3795
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 1.3976872404534158e-08,
+      "loss": 0.4016,
+      "step": 3796
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 1.3536783018928135e-08,
+      "loss": 0.4177,
+      "step": 3797
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 1.3103728757576372e-08,
+      "loss": 0.4097,
+      "step": 3798
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 1.2677709925551506e-08,
+      "loss": 0.4092,
+      "step": 3799
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 1.2258726822970135e-08,
+      "loss": 0.4185,
+      "step": 3800
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 1.184677974498949e-08,
+      "loss": 0.429,
+      "step": 3801
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 1.1441868981815207e-08,
+      "loss": 0.4153,
+      "step": 3802
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 1.104399481869023e-08,
+      "loss": 0.4165,
+      "step": 3803
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 1.06531575359059e-08,
+      "loss": 0.4248,
+      "step": 3804
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 1.026935740879198e-08,
+      "loss": 0.3997,
+      "step": 3805
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 9.892594707723302e-09,
+      "loss": 0.3916,
+      "step": 3806
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 9.522869698116444e-09,
+      "loss": 0.406,
+      "step": 3807
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 9.160182640429727e-09,
+      "loss": 0.4138,
+      "step": 3808
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 8.804533790165437e-09,
+      "loss": 0.4189,
+      "step": 3809
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 8.455923397863164e-09,
+      "loss": 0.4441,
+      "step": 3810
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 8.114351709109792e-09,
+      "loss": 0.4004,
+      "step": 3811
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 7.779818964530617e-09,
+      "loss": 0.4199,
+      "step": 3812
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 7.452325399791571e-09,
+      "loss": 0.4104,
+      "step": 3813
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 7.131871245600331e-09,
+      "loss": 0.4104,
+      "step": 3814
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 6.818456727707423e-09,
+      "loss": 0.4075,
+      "step": 3815
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 6.5120820669017926e-09,
+      "loss": 0.3921,
+      "step": 3816
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 6.212747479015236e-09,
+      "loss": 0.4133,
+      "step": 3817
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 5.920453174916851e-09,
+      "loss": 0.4326,
+      "step": 3818
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 5.635199360519705e-09,
+      "loss": 0.4221,
+      "step": 3819
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 5.356986236774165e-09,
+      "loss": 0.4197,
+      "step": 3820
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 5.08581399967345e-09,
+      "loss": 0.4258,
+      "step": 3821
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 4.821682840248087e-09,
+      "loss": 0.3965,
+      "step": 3822
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 4.564592944569235e-09,
+      "loss": 0.4194,
+      "step": 3823
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 4.314544493748685e-09,
+      "loss": 0.4116,
+      "step": 3824
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 4.071537663936642e-09,
+      "loss": 0.3997,
+      "step": 3825
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 3.835572626323947e-09,
+      "loss": 0.406,
+      "step": 3826
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 3.606649547139851e-09,
+      "loss": 0.417,
+      "step": 3827
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 3.3847685876520208e-09,
+      "loss": 0.3848,
+      "step": 3828
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 3.1699299041676457e-09,
+      "loss": 0.425,
+      "step": 3829
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 2.9621336480345485e-09,
+      "loss": 0.405,
+      "step": 3830
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 2.761379965637856e-09,
+      "loss": 0.4202,
+      "step": 3831
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 2.5676689984022173e-09,
+      "loss": 0.3979,
+      "step": 3832
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 2.3810008827895857e-09,
+      "loss": 0.4207,
+      "step": 3833
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 2.201375750301438e-09,
+      "loss": 0.4304,
+      "step": 3834
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 2.0287937274765522e-09,
+      "loss": 0.4187,
+      "step": 3835
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 1.863254935895453e-09,
+      "loss": 0.436,
+      "step": 3836
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 1.7047594921737465e-09,
+      "loss": 0.3892,
+      "step": 3837
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 1.5533075079643411e-09,
+      "loss": 0.3723,
+      "step": 3838
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 1.4088990899618904e-09,
+      "loss": 0.4175,
+      "step": 3839
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 1.2715343398972401e-09,
+      "loss": 0.4141,
+      "step": 3840
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 1.1412133545385395e-09,
+      "loss": 0.4268,
+      "step": 3841
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 1.017936225692351e-09,
+      "loss": 0.4265,
+      "step": 3842
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 9.017030402036497e-10,
+      "loss": 0.397,
+      "step": 3843
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 7.92513879954715e-10,
+      "loss": 0.4211,
+      "step": 3844
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 6.903688218651283e-10,
+      "loss": 0.4036,
+      "step": 3845
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 5.95267937892885e-10,
+      "loss": 0.4087,
+      "step": 3846
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 5.072112950343933e-10,
+      "loss": 0.4233,
+      "step": 3847
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 4.2619895532114433e-10,
+      "loss": 0.4126,
+      "step": 3848
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 3.522309758230424e-10,
+      "loss": 0.4294,
+      "step": 3849
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 2.8530740865062576e-10,
+      "loss": 0.406,
+      "step": 3850
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 2.2542830094729462e-10,
+      "loss": 0.4314,
+      "step": 3851
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 1.7259369489597277e-10,
+      "loss": 0.4229,
+      "step": 3852
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 1.2680362771688714e-10,
+      "loss": 0.4136,
+      "step": 3853
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 8.805813166756772e-11,
+      "loss": 0.427,
+      "step": 3854
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 5.6357234043957766e-11,
+      "loss": 0.4243,
+      "step": 3855
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 3.1700957177083214e-11,
+      "loss": 0.4087,
+      "step": 3856
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 1.4089318436383281e-11,
+      "loss": 0.4285,
+      "step": 3857
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 3.522330229710491e-12,
+      "loss": 0.4219,
+      "step": 3858
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 0.0,
+      "loss": 0.4307,
+      "step": 3859
+    },
+    {
+      "epoch": 1.0,
+      "step": 3859,
+      "total_flos": 816438756818944.0,
+      "train_loss": 0.49826279932058176,
+      "train_runtime": 64871.1521,
+      "train_samples_per_second": 15.227,
+      "train_steps_per_second": 0.059
+    }
+  ],
+  "logging_steps": 1.0,
+  "max_steps": 3859,
+  "num_train_epochs": 1,
+  "save_steps": 500,
+  "total_flos": 816438756818944.0,
+  "trial_name": null,
+  "trial_params": null
+}
diff --git a/ckpts/llm/training_args.bin b/ckpts/llm/training_args.bin
new file mode 100644
index 0000000000000000000000000000000000000000..f43d10e9e8cf04a0c10bb56dfdebef8a5a2d72bd
--- /dev/null
+++ b/ckpts/llm/training_args.bin
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8155c40679be97018cc40bf6875bbce7adbc0312d8a1fa39e73824896d2fe2f2
+size 6139
diff --git a/ckpts/prior/model.bin b/ckpts/prior/model.bin
deleted file mode 120000
index 00d6daa1dfc9d844b85b601a9b5ce0718b2c50fe..0000000000000000000000000000000000000000
--- a/ckpts/prior/model.bin
+++ /dev/null
@@ -1 +0,0 @@
-/home/jacklishufan/AudioLDM2/diffusion_prior_3.bin
\ No newline at end of file
diff --git a/ckpts/prior/model.bin b/ckpts/prior/model.bin
new file mode 100644
index 0000000000000000000000000000000000000000..81d1e3dd2ec8b4c31a878775894724ba554a5228
--- /dev/null
+++ b/ckpts/prior/model.bin
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a8a6fe3ed10da0bcbe20dd7bb427ca503584a3a3cdd570fc57c505aaa63c2719
+size 2842310673
diff --git a/ckpts/sdxl b/ckpts/sdxl
deleted file mode 120000
index c1629b1c3dcaaf42bb25cf2a4dc8b5c787b28ef3..0000000000000000000000000000000000000000
--- a/ckpts/sdxl
+++ /dev/null
@@ -1 +0,0 @@
-/localhome/data/ckpts/jacklishufan/sdxl/
\ No newline at end of file
diff --git a/ckpts/sdxl/unet/config.json b/ckpts/sdxl/unet/config.json
new file mode 100644
index 0000000000000000000000000000000000000000..0072422563583b4b1ea81e4876fd8294d190340d
--- /dev/null
+++ b/ckpts/sdxl/unet/config.json
@@ -0,0 +1,73 @@
+{
+  "_class_name": "UnCLipXL",
+  "_diffusers_version": "0.22.0.dev0",
+  "_name_or_path": "stabilityai/stable-diffusion-xl-base-1.0",
+  "act_fn": "silu",
+  "addition_embed_type": "text_time",
+  "addition_embed_type_num_heads": 64,
+  "addition_time_embed_dim": 256,
+  "attention_head_dim": [
+    5,
+    10,
+    20
+  ],
+  "attention_type": "default",
+  "block_out_channels": [
+    320,
+    640,
+    1280
+  ],
+  "center_input_sample": false,
+  "class_embed_type": null,
+  "class_embeddings_concat": false,
+  "conv_in_kernel": 3,
+  "conv_out_kernel": 3,
+  "cross_attention_dim": 2048,
+  "cross_attention_norm": null,
+  "down_block_types": [
+    "DownBlock2D",
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "dropout": 0,
+  "dual_cross_attention": false,
+  "encoder_hid_dim": null,
+  "encoder_hid_dim_type": null,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "in_channels": 4,
+  "layers_per_block": 2,
+  "mid_block_only_cross_attention": null,
+  "mid_block_scale_factor": 1,
+  "mid_block_type": "UNetMidBlock2DCrossAttn",
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_attention_heads": null,
+  "num_class_embeds": null,
+  "only_cross_attention": false,
+  "out_channels": 4,
+  "projection_class_embeddings_input_dim": 2816,
+  "resnet_out_scale_factor": 1.0,
+  "resnet_skip_time_act": false,
+  "resnet_time_scale_shift": "default",
+  "reverse_transformer_layers_per_block": null,
+  "sample_size": 128,
+  "time_cond_proj_dim": null,
+  "time_embedding_act_fn": null,
+  "time_embedding_dim": null,
+  "time_embedding_type": "positional",
+  "timestep_post_act": null,
+  "transformer_layers_per_block": [
+    1,
+    2,
+    10
+  ],
+  "up_block_types": [
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D",
+    "UpBlock2D"
+  ],
+  "upcast_attention": null,
+  "use_linear_projection": true
+}
diff --git a/ckpts/sdxl/unet/diffusion_pytorch_model.safetensors b/ckpts/sdxl/unet/diffusion_pytorch_model.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..f0a7c27a76229a2895cd6a8447eb21b26e9e8d2c
--- /dev/null
+++ b/ckpts/sdxl/unet/diffusion_pytorch_model.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:37086d06af3b51cdd461ff6430fd724805649df554cb39372a7dc644bb6ff49c
+size 5177116640
diff --git a/demo.ipynb b/demo.ipynb
index 64652058cdc065d4cb177a47ea92b4fe5156cbc8..ade053bdf3f697739a363a79de679eaf0a6e346e 100644
--- a/demo.ipynb
+++ b/demo.ipynb
@@ -213,7 +213,7 @@
    ],
    "source": [
     "torch.manual_seed(0)\n",
-    "res0,res,output_caption = pipe(inst,mm_data,alpha = 1.0,h=[0.4,0.6,0.4],norm=20.0,refinement=0.3,llm_only=False,num_inference_steps=50)"
+    "res0,res,_ = pipe(inst,mm_data,alpha = 1.0,h=[0.4,0.6,0.4],norm=20.0,refinement=0.3,llm_only=False,num_inference_steps=50)"
    ]
   },
   {
diff --git a/instructany2pix/__init__.py b/instructany2pix/__init__.py
deleted file mode 100644
index 5a604ba92332e0de5e3dbac001cf83daa6cd66f9..0000000000000000000000000000000000000000
--- a/instructany2pix/__init__.py
+++ /dev/null
@@ -1 +0,0 @@
-from .pipeline import InstructAny2PixPipeline
\ No newline at end of file
diff --git a/instructany2pix/ddim/pnp_pipeline.py b/instructany2pix/ddim/pnp_pipeline.py
deleted file mode 100644
index 9f5a67c7ef5e04bf354d70fd873d76ff4b4d0108..0000000000000000000000000000000000000000
--- a/instructany2pix/ddim/pnp_pipeline.py
+++ /dev/null
@@ -1,524 +0,0 @@
-# Plug&Play Feature Injection
-
-import torch
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
-
-from diffusers import (
-    StableDiffusionXLPipeline,
-    StableDiffusionXLImg2ImgPipeline,
-    DDIMScheduler,
-)
-
-from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import (
-    rescale_noise_cfg,
-    StableDiffusionXLPipelineOutput,
-)
-
-import PIL
-import numpy as np
-
-from tqdm import tqdm
-
-
-def _get_add_time_ids(
-        self,
-        original_size,
-        crops_coords_top_left,
-        target_size,
-        aesthetic_score,
-        negative_aesthetic_score,
-        negative_original_size,
-        negative_crops_coords_top_left,
-        negative_target_size,
-        dtype,
-    ):
-        if self.config.requires_aesthetics_score:
-            add_time_ids = list(original_size + crops_coords_top_left + (aesthetic_score,))
-            add_neg_time_ids = list(
-                negative_original_size + negative_crops_coords_top_left + (negative_aesthetic_score,)
-            )
-        else:
-            add_time_ids = list(original_size + crops_coords_top_left + target_size)
-            add_neg_time_ids = list(negative_original_size + crops_coords_top_left + negative_target_size)
-
-        passed_add_embed_dim = (
-            self.unet.config.addition_time_embed_dim * len(add_time_ids) + self.text_encoder_2.config.projection_dim
-        )
-        expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features
-
-        if (
-            expected_add_embed_dim > passed_add_embed_dim
-            and (expected_add_embed_dim - passed_add_embed_dim) == self.unet.config.addition_time_embed_dim
-        ):
-            raise ValueError(
-                f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. Please make sure to enable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=True)` to make sure `aesthetic_score` {aesthetic_score} and `negative_aesthetic_score` {negative_aesthetic_score} is correctly used by the model."
-            )
-        elif (
-            expected_add_embed_dim < passed_add_embed_dim
-            and (passed_add_embed_dim - expected_add_embed_dim) == self.unet.config.addition_time_embed_dim
-        ):
-            raise ValueError(
-                f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. Please make sure to disable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=False)` to make sure `target_size` {target_size} is correctly used by the model."
-            )
-        elif expected_add_embed_dim != passed_add_embed_dim:
-            raise ValueError(
-                f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`."
-            )
-
-        add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
-        add_neg_time_ids = torch.tensor([add_neg_time_ids], dtype=dtype)
-
-        return add_time_ids, add_neg_time_ids
-
-def _backward_ddim(x_tm1, alpha_t, alpha_tm1, eps_xt):
-    """
-    let a = alpha_t, b = alpha_{t - 1}
-    We have a > b,
-    x_{t} - x_{t - 1} = sqrt(a) ((sqrt(1/b) - sqrt(1/a)) * x_{t-1} + (sqrt(1/a - 1) - sqrt(1/b - 1)) * eps_{t-1})
-    From https://arxiv.org/pdf/2105.05233.pdf, section F.
-    """
-
-    a, b = alpha_t, alpha_tm1
-    sa = a**0.5
-    sb = b**0.5
-
-    return sa * ((1 / sb) * x_tm1 + ((1 / a - 1) ** 0.5 - (1 / b - 1) ** 0.5) * eps_xt)
-
-
-class SDXLDDIMPipeline(StableDiffusionXLImg2ImgPipeline):
-    
-    
-    @torch.no_grad()
-    def inverse(
-        self,
-        prompt: Union[str, List[str]] = None,
-        prompt_2: Optional[Union[str, List[str]]] = None,
-        image: Union[
-            torch.FloatTensor,
-            PIL.Image.Image,
-            np.ndarray,
-            List[torch.FloatTensor],
-            List[PIL.Image.Image],
-            List[np.ndarray],
-        ] = None,
-        strength: float = 0.3,
-        num_inference_steps: int = 50,
-        denoising_start: Optional[float] = None,
-        denoising_end: Optional[float] = None,
-        guidance_scale: float = 5.0,
-        negative_prompt: Optional[Union[str, List[str]]] = None,
-        negative_prompt_2: Optional[Union[str, List[str]]] = None,
-        num_images_per_prompt: Optional[int] = 1,
-        eta: float = 0.0,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.FloatTensor] = None,
-        prompt_embeds: Optional[torch.FloatTensor] = None,
-        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
-        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
-        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
-        output_type: Optional[str] = "pil",
-        return_dict: bool = True,
-        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
-        callback_steps: int = 1,
-        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
-        guidance_rescale: float = 0.0,
-        original_size: Tuple[int, int] = None,
-        crops_coords_top_left: Tuple[int, int] = (0, 0),
-        target_size: Tuple[int, int] = None,
-        aesthetic_score: float = 6.0,
-        negative_aesthetic_score: float = 2.5,
-        #clip_skip=None,
-    ):
-        self.scheduler = DDIMScheduler.from_config(self.scheduler.config)
-
-        # 1. Check inputs. Raise error if not correct
-        self.check_inputs(
-            prompt,
-            prompt_2,
-            strength,
-            num_inference_steps,
-            callback_steps,
-            negative_prompt,
-            negative_prompt_2,
-            prompt_embeds,
-            negative_prompt_embeds,
-        )
-
-        # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        device = self._execution_device
-
-        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
-        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
-        # corresponds to doing no classifier free guidance.
-        do_classifier_free_guidance = False
-        # 3. Encode input prompt
-        text_encoder_lora_scale = (
-            cross_attention_kwargs.get("scale", None)
-            if cross_attention_kwargs is not None
-            else None
-        )
-        (
-            prompt_embeds,
-            negative_prompt_embeds,
-            pooled_prompt_embeds,
-            negative_pooled_prompt_embeds,
-        ) = self.encode_prompt(
-            prompt=prompt,
-            prompt_2=prompt_2,
-            device=device,
-            num_images_per_prompt=num_images_per_prompt,
-            do_classifier_free_guidance=do_classifier_free_guidance,
-            negative_prompt=negative_prompt,
-            negative_prompt_2=negative_prompt_2,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            pooled_prompt_embeds=pooled_prompt_embeds,
-            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
-            lora_scale=text_encoder_lora_scale,
-            #clip_skip=clip_skip,
-        )
-
-        # 4. Preprocess image
-        image = self.image_processor.preprocess(image)
-
-        self.scheduler.set_timesteps(num_inference_steps, device=device)
-       
-        # 6. Prepare latent variables
-        latents = self.prepare_latents(
-            image,
-            None,
-            batch_size,
-            num_images_per_prompt,
-            prompt_embeds.dtype,
-            device,
-            generator,
-            False,
-        )
-       
-        height, width = latents.shape[-2:]
-        height = height * self.vae_scale_factor
-        width = width * self.vae_scale_factor
-
-        original_size = original_size or (height, width)
-        target_size = target_size or (height, width)
-
-        # 8. Prepare added time ids & embeddings
-        add_text_embeds = pooled_prompt_embeds
-        negative_original_size = None
-        negative_crops_coords_top_left = None
-        negative_target_size = None
-        if negative_original_size is None:
-            negative_original_size = original_size
-        if negative_target_size is None:
-            negative_target_size = target_size
-        negative_crops_coords_top_left = (0,0)
-        if self.text_encoder_2 is None:
-            text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1])
-        else:
-            text_encoder_projection_dim = self.text_encoder_2.config.projection_dim
-
-        add_time_ids, add_neg_time_ids = _get_add_time_ids(self,
-            original_size,
-            crops_coords_top_left,
-            target_size,
-            aesthetic_score,
-            negative_aesthetic_score,
-            negative_original_size,
-            negative_crops_coords_top_left,
-            negative_target_size,
-            dtype=prompt_embeds.dtype,
-            #text_encoder_projection_dim=text_encoder_projection_dim,
-        )
-        add_time_ids = add_time_ids.repeat(batch_size * num_images_per_prompt, 1)
-
-        prompt_embeds = prompt_embeds.to(device)
-        add_text_embeds = add_text_embeds.to(device)
-        add_time_ids = add_time_ids.to(device)
-
-        added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
-        prev_timestep = None
-
-        for t in tqdm(reversed(self.scheduler.timesteps)):
-            latent_model_input = latents
-            noise_pred = self.unet(
-                latent_model_input,
-                t,
-                encoder_hidden_states=prompt_embeds,
-                cross_attention_kwargs=cross_attention_kwargs,
-                added_cond_kwargs={k:v for k,v in added_cond_kwargs.items()},
-                return_dict=False,
-            )[0]
-
-            alpha_prod_t = self.scheduler.alphas_cumprod[t]
-            alpha_prod_t_prev = (
-                self.scheduler.alphas_cumprod[prev_timestep]
-                if prev_timestep is not None
-                else self.scheduler.final_alpha_cumprod
-            )
-            prev_timestep = t
-
-            latents = _backward_ddim(
-                x_tm1=latents,
-                alpha_t=alpha_prod_t,
-                alpha_tm1=alpha_prod_t_prev,
-                eps_xt=noise_pred,
-            )
-
-        image = latents
-        return StableDiffusionXLPipelineOutput(images=image)
-
-
-class SDXLPNPPipeline(StableDiffusionXLPipeline):
-    def __init__(
-        self,
-        *args,
-        **kwargs,
-    ):
-        super().__init__(*args, **kwargs)
-        self.feature_bags = {}
-
-    def __call__(
-        self,
-        prompt: Union[str, List[str]] = None,
-        prompt_2: Optional[Union[str, List[str]]] = None,
-        height: Optional[int] = None,
-        width: Optional[int] = None,
-        num_inference_steps: int = 50,
-        denoising_end: Optional[float] = None,
-        guidance_scale: float = 5.0,
-        negative_prompt: Optional[Union[str, List[str]]] = None,
-        negative_prompt_2: Optional[Union[str, List[str]]] = None,
-        num_images_per_prompt: Optional[int] = 1,
-        eta: float = 0.0,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.FloatTensor] = None,
-        prompt_embeds: Optional[torch.FloatTensor] = None,
-        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
-        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
-        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
-        output_type: Optional[str] = "pil",
-        return_dict: bool = True,
-        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
-        callback_steps: int = 1,
-        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
-        guidance_rescale: float = 0.0,
-        original_size: Optional[Tuple[int, int]] = None,
-        crops_coords_top_left: Tuple[int, int] = (0, 0),
-        target_size: Optional[Tuple[int, int]] = None,
-    ):
-        height = height or self.default_sample_size * self.vae_scale_factor
-        width = width or self.default_sample_size * self.vae_scale_factor
-
-        original_size = original_size or (height, width)
-        target_size = target_size or (height, width)
-
-        # 1. Check inputs. Raise error if not correct
-        self.check_inputs(
-            prompt,
-            prompt_2,
-            height,
-            width,
-            callback_steps,
-            negative_prompt,
-            negative_prompt_2,
-            prompt_embeds,
-            negative_prompt_embeds,
-            pooled_prompt_embeds,
-            negative_pooled_prompt_embeds,
-        )
-
-        # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        device = self._execution_device
-
-        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
-        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
-        # corresponds to doing no classifier free guidance.
-        do_classifier_free_guidance = guidance_scale > 1.0
-
-        # 3. Encode input prompt
-        text_encoder_lora_scale = (
-            cross_attention_kwargs.get("scale", None)
-            if cross_attention_kwargs is not None
-            else None
-        )
-        (
-            prompt_embeds,
-            negative_prompt_embeds,
-            pooled_prompt_embeds,
-            negative_pooled_prompt_embeds,
-        ) = self.encode_prompt(
-            prompt=prompt,
-            prompt_2=prompt_2,
-            device=device,
-            num_images_per_prompt=num_images_per_prompt,
-            do_classifier_free_guidance=do_classifier_free_guidance,
-            negative_prompt=negative_prompt,
-            negative_prompt_2=negative_prompt_2,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            pooled_prompt_embeds=pooled_prompt_embeds,
-            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
-            lora_scale=text_encoder_lora_scale,
-        )
-
-        # 4. Prepare timesteps
-        self.scheduler.set_timesteps(num_inference_steps, device=device)
-
-        timesteps = self.scheduler.timesteps
-
-        # 5. Prepare latent variables
-        num_channels_latents = self.unet.config.in_channels
-        latents = self.prepare_latents(
-            batch_size * num_images_per_prompt,
-            num_channels_latents,
-            height,
-            width,
-            prompt_embeds.dtype,
-            device,
-            generator,
-            latents,
-        )
-
-        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
-        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
-
-        # 7. Prepare added time ids & embeddings
-        add_text_embeds = pooled_prompt_embeds
-        add_time_ids = self._get_add_time_ids(
-            original_size, crops_coords_top_left, target_size, dtype=prompt_embeds.dtype
-        )
-
-        if do_classifier_free_guidance:
-            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
-            add_text_embeds = torch.cat(
-                [negative_pooled_prompt_embeds, add_text_embeds], dim=0
-            )
-            add_time_ids = torch.cat([add_time_ids, add_time_ids], dim=0)
-
-        prompt_embeds = prompt_embeds.to(device)
-        add_text_embeds = add_text_embeds.to(device)
-        add_time_ids = add_time_ids.to(device).repeat(
-            batch_size * num_images_per_prompt, 1
-        )
-
-        # 8. Denoising loop
-        num_warmup_steps = max(
-            len(timesteps) - num_inference_steps * self.scheduler.order, 0
-        )
-
-        # 7.1 Apply denoising_end
-        if (
-            denoising_end is not None
-            and type(denoising_end) == float
-            and denoising_end > 0
-            and denoising_end < 1
-        ):
-            discrete_timestep_cutoff = int(
-                round(
-                    self.scheduler.config.num_train_timesteps
-                    - (denoising_end * self.scheduler.config.num_train_timesteps)
-                )
-            )
-            num_inference_steps = len(
-                list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))
-            )
-            timesteps = timesteps[:num_inference_steps]
-
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                # expand the latents if we are doing classifier free guidance
-                # print(t)
-                latent_model_input = (
-                    torch.cat([latents] * 2) if do_classifier_free_guidance else latents
-                )
-
-                latent_model_input = self.scheduler.scale_model_input(
-                    latent_model_input, t
-                )
-
-                # predict the noise residual
-
-                added_cond_kwargs = {
-                    "text_embeds": add_text_embeds[:1, ...],
-                    "time_ids": add_time_ids[:1, ...],
-                }
-                noise_pred_text, feats = self.unet(
-                    latent_model_input[:1, ...],
-                    t,
-                    encoder_hidden_states=prompt_embeds[:1, ...],
-                    added_cond_kwargs=added_cond_kwargs,
-                    return_dict=False,
-                )[0]
-
-                added_cond_kwargs = {
-                    "text_embeds": add_text_embeds[1:2, ...],
-                    "time_ids": add_time_ids[1:2, ...],
-                }
-                noise_pred_uncond, feats = self.unet(
-                    latent_model_input[1:2, ...],
-                    t,
-                    encoder_hidden_states=prompt_embeds[1:2, ...],
-                    added_cond_kwargs=added_cond_kwargs,
-                    return_dict=False,
-                )[0]
-
-                noise_pred = noise_pred_uncond + guidance_scale * (
-                    noise_pred_text - noise_pred_uncond
-                )
-
-                # compute the previous noisy sample x_t -> x_t-1
-                latents = self.scheduler.step(
-                    noise_pred, t, latents, **extra_step_kwargs, return_dict=False
-                )[0]
-
-                # call the callback, if provided
-                if i == len(timesteps) - 1 or (
-                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
-                ):
-                    progress_bar.update()
-                    if callback is not None and i % callback_steps == 0:
-                        callback(i, t, latents)
-
-        # make sure the VAE is in float32 mode, as it overflows in float16
-        if self.vae.dtype == torch.float16 and self.vae.config.force_upcast:
-            self.upcast_vae()
-            latents = latents.to(
-                next(iter(self.vae.post_quant_conv.parameters())).dtype
-            )
-
-        if not output_type == "latent":
-            image = self.vae.decode(
-                latents / self.vae.config.scaling_factor, return_dict=False
-            )[0]
-        else:
-            image = latents
-            return StableDiffusionXLPipelineOutput(images=image)
-
-        # apply watermark if available
-        if self.watermark is not None:
-            image = self.watermark.apply_watermark(image)
-
-        image = self.image_processor.postprocess(image, output_type=output_type)
-
-        # Offload last model to CPU
-        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
-            self.final_offload_hook.offload()
-
-        if not return_dict:
-            return (image,)
-
-        return StableDiffusionXLPipelineOutput(images=image)
\ No newline at end of file
diff --git a/instructany2pix/ddim/sdxl_pipeline.py b/instructany2pix/ddim/sdxl_pipeline.py
deleted file mode 100644
index c1e4315396f39b330b80725b649de38e741d5da0..0000000000000000000000000000000000000000
--- a/instructany2pix/ddim/sdxl_pipeline.py
+++ /dev/null
@@ -1,990 +0,0 @@
-# Copyright 2023 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import inspect
-import os
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
-
-import torch
-from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
-
-from diffusers.image_processor import VaeImageProcessor
-from diffusers.loaders import (
-    FromSingleFileMixin,
-    LoraLoaderMixin,
-    TextualInversionLoaderMixin,
-)
-from diffusers.models import AutoencoderKL, UNet2DConditionModel
-from diffusers.models.attention_processor import (
-    AttnProcessor2_0,
-    LoRAAttnProcessor2_0,
-    LoRAXFormersAttnProcessor,
-    XFormersAttnProcessor,
-)
-from diffusers.models.lora import adjust_lora_scale_text_encoder
-from diffusers.schedulers import KarrasDiffusionSchedulers
-from diffusers.utils import (
-    is_accelerate_available,
-    is_accelerate_version,
-    is_invisible_watermark_available,
-    logging,
-    replace_example_docstring,
-)
-from diffusers.utils.torch_utils import randn_tensor
-from diffusers.pipelines.pipeline_utils import DiffusionPipeline
-from diffusers.pipelines.stable_diffusion_xl import StableDiffusionXLPipelineOutput
-
-
-if is_invisible_watermark_available():
-    from diffusers.pipelines.stable_diffusion_xl.watermark import StableDiffusionXLWatermarker
-
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-EXAMPLE_DOC_STRING = """
-    Examples:
-        ```py
-        >>> import torch
-        >>> from diffusers import StableDiffusionXLPipeline
-
-        >>> pipe = StableDiffusionXLPipeline.from_pretrained(
-        ...     "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
-        ... )
-        >>> pipe = pipe.to("cuda")
-
-        >>> prompt = "a photo of an astronaut riding a horse on mars"
-        >>> image = pipe(prompt).images[0]
-        ```
-"""
-
-
-# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg
-def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
-    """
-    Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and
-    Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4
-    """
-    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
-    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
-    # rescale the results from guidance (fixes overexposure)
-    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
-    # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
-    noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
-    return noise_cfg
-
-
-class StableDiffusionXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoaderMixin, TextualInversionLoaderMixin):
-    r"""
-    Pipeline for text-to-image generation using Stable Diffusion XL.
-
-    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
-    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
-
-    In addition the pipeline inherits the following loading methods:
-        - *LoRA*: [`StableDiffusionXLPipeline.load_lora_weights`]
-        - *Ckpt*: [`loaders.FromSingleFileMixin.from_single_file`]
-
-    as well as the following saving methods:
-        - *LoRA*: [`loaders.StableDiffusionXLPipeline.save_lora_weights`]
-
-    Args:
-        vae ([`AutoencoderKL`]):
-            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
-        text_encoder ([`CLIPTextModel`]):
-            Frozen text-encoder. Stable Diffusion XL uses the text portion of
-            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
-            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
-        text_encoder_2 ([` CLIPTextModelWithProjection`]):
-            Second frozen text-encoder. Stable Diffusion XL uses the text and pool portion of
-            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection),
-            specifically the
-            [laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k)
-            variant.
-        tokenizer (`CLIPTokenizer`):
-            Tokenizer of class
-            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
-        tokenizer_2 (`CLIPTokenizer`):
-            Second Tokenizer of class
-            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
-        unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.
-        scheduler ([`SchedulerMixin`]):
-            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
-            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
-        force_zeros_for_empty_prompt (`bool`, *optional*, defaults to `"True"`):
-            Whether the negative prompt embeddings shall be forced to always be set to 0. Also see the config of
-            `stabilityai/stable-diffusion-xl-base-1-0`.
-        add_watermarker (`bool`, *optional*):
-            Whether to use the [invisible_watermark library](https://github.com/ShieldMnt/invisible-watermark/) to
-            watermark output images. If not defined, it will default to True if the package is installed, otherwise no
-            watermarker will be used.
-    """
-    model_cpu_offload_seq = "text_encoder->text_encoder_2->unet->vae"
-
-    def __init__(
-        self,
-        vae: AutoencoderKL,
-        text_encoder: CLIPTextModel,
-        text_encoder_2: CLIPTextModelWithProjection,
-        tokenizer: CLIPTokenizer,
-        tokenizer_2: CLIPTokenizer,
-        unet: UNet2DConditionModel,
-        scheduler: KarrasDiffusionSchedulers,
-        force_zeros_for_empty_prompt: bool = True,
-        add_watermarker: Optional[bool] = False,
-    ):
-        super().__init__()
-
-        self.register_modules(
-            vae=vae,
-            text_encoder=text_encoder,
-            text_encoder_2=text_encoder_2,
-            tokenizer=tokenizer,
-            tokenizer_2=tokenizer_2,
-            unet=unet,
-            scheduler=scheduler,
-        )
-        self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt)
-        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
-        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
-        self.default_sample_size = self.unet.config.sample_size
-
-        add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available()
-
-        if add_watermarker:
-            self.watermark = StableDiffusionXLWatermarker()
-        else:
-            self.watermark = None
-
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_vae_slicing
-    def enable_vae_slicing(self):
-        r"""
-        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
-        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
-        """
-        self.vae.enable_slicing()
-
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.disable_vae_slicing
-    def disable_vae_slicing(self):
-        r"""
-        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
-        computing decoding in one step.
-        """
-        self.vae.disable_slicing()
-
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_vae_tiling
-    def enable_vae_tiling(self):
-        r"""
-        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
-        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
-        processing larger images.
-        """
-        self.vae.enable_tiling()
-
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.disable_vae_tiling
-    def disable_vae_tiling(self):
-        r"""
-        Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
-        computing decoding in one step.
-        """
-        self.vae.disable_tiling()
-
-    def encode_prompt(
-        self,
-        prompt: str,
-        prompt_2: Optional[str] = None,
-        device: Optional[torch.device] = None,
-        num_images_per_prompt: int = 1,
-        do_classifier_free_guidance: bool = True,
-        negative_prompt: Optional[str] = None,
-        negative_prompt_2: Optional[str] = None,
-        prompt_embeds: Optional[torch.FloatTensor] = None,
-        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
-        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
-        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
-        lora_scale: Optional[float] = None,
-    ):
-        r"""
-        Encodes the prompt into text encoder hidden states.
-
-        Args:
-            prompt (`str` or `List[str]`, *optional*):
-                prompt to be encoded
-            prompt_2 (`str` or `List[str]`, *optional*):
-                The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
-                used in both text-encoders
-            device: (`torch.device`):
-                torch device
-            num_images_per_prompt (`int`):
-                number of images that should be generated per prompt
-            do_classifier_free_guidance (`bool`):
-                whether to use classifier free guidance or not
-            negative_prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts not to guide the image generation. If not defined, one has to pass
-                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
-                less than `1`).
-            negative_prompt_2 (`str` or `List[str]`, *optional*):
-                The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
-                `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
-            prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
-                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
-                argument.
-            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
-                If not provided, pooled text embeddings will be generated from `prompt` input argument.
-            negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
-                weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
-                input argument.
-            lora_scale (`float`, *optional*):
-                A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
-        """
-        device = device or self._execution_device
-
-        # set lora scale so that monkey patched LoRA
-        # function of text encoder can correctly access it
-        if lora_scale is not None and isinstance(self, LoraLoaderMixin):
-            self._lora_scale = lora_scale
-
-            # dynamically adjust the LoRA scale
-            adjust_lora_scale_text_encoder(self.text_encoder, lora_scale)
-            adjust_lora_scale_text_encoder(self.text_encoder_2, lora_scale)
-
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        # Define tokenizers and text encoders
-        tokenizers = [self.tokenizer, self.tokenizer_2] if self.tokenizer is not None else [self.tokenizer_2]
-        text_encoders = (
-            [self.text_encoder, self.text_encoder_2] if self.text_encoder is not None else [self.text_encoder_2]
-        )
-
-        if prompt_embeds is None:
-            prompt_2 = prompt_2 or prompt
-            # textual inversion: procecss multi-vector tokens if necessary
-            prompt_embeds_list = []
-            prompts = [prompt, prompt_2]
-            for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders):
-                if isinstance(self, TextualInversionLoaderMixin):
-                    prompt = self.maybe_convert_prompt(prompt, tokenizer)
-
-                text_inputs = tokenizer(
-                    prompt,
-                    padding="max_length",
-                    max_length=tokenizer.model_max_length,
-                    truncation=True,
-                    return_tensors="pt",
-                )
-
-                text_input_ids = text_inputs.input_ids
-                untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
-
-                if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
-                    text_input_ids, untruncated_ids
-                ):
-                    removed_text = tokenizer.batch_decode(untruncated_ids[:, tokenizer.model_max_length - 1 : -1])
-                    logger.warning(
-                        "The following part of your input was truncated because CLIP can only handle sequences up to"
-                        f" {tokenizer.model_max_length} tokens: {removed_text}"
-                    )
-
-                prompt_embeds = text_encoder(
-                    text_input_ids.to(device),
-                    output_hidden_states=True,
-                )
-
-                # We are only ALWAYS interested in the pooled output of the final text encoder
-                pooled_prompt_embeds = prompt_embeds[0]
-                prompt_embeds = prompt_embeds.hidden_states[-2]
-
-                prompt_embeds_list.append(prompt_embeds)
-
-            prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
-
-        # get unconditional embeddings for classifier free guidance
-        zero_out_negative_prompt = negative_prompt is None and self.config.force_zeros_for_empty_prompt
-        if do_classifier_free_guidance and negative_prompt_embeds is None and zero_out_negative_prompt:
-            negative_prompt_embeds = torch.zeros_like(prompt_embeds)
-            negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds)
-        elif do_classifier_free_guidance and negative_prompt_embeds is None:
-            negative_prompt = negative_prompt or ""
-            negative_prompt_2 = negative_prompt_2 or negative_prompt
-
-            uncond_tokens: List[str]
-            if prompt is not None and type(prompt) is not type(negative_prompt):
-                raise TypeError(
-                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
-                    f" {type(prompt)}."
-                )
-            elif isinstance(negative_prompt, str):
-                uncond_tokens = [negative_prompt, negative_prompt_2]
-            elif batch_size != len(negative_prompt):
-                raise ValueError(
-                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
-                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
-                    " the batch size of `prompt`."
-                )
-            else:
-                uncond_tokens = [negative_prompt, negative_prompt_2]
-
-            negative_prompt_embeds_list = []
-            for negative_prompt, tokenizer, text_encoder in zip(uncond_tokens, tokenizers, text_encoders):
-                if isinstance(self, TextualInversionLoaderMixin):
-                    negative_prompt = self.maybe_convert_prompt(negative_prompt, tokenizer)
-
-                max_length = prompt_embeds.shape[1]
-                uncond_input = tokenizer(
-                    negative_prompt,
-                    padding="max_length",
-                    max_length=max_length,
-                    truncation=True,
-                    return_tensors="pt",
-                )
-
-                negative_prompt_embeds = text_encoder(
-                    uncond_input.input_ids.to(device),
-                    output_hidden_states=True,
-                )
-                # We are only ALWAYS interested in the pooled output of the final text encoder
-                negative_pooled_prompt_embeds = negative_prompt_embeds[0]
-                negative_prompt_embeds = negative_prompt_embeds.hidden_states[-2]
-
-                negative_prompt_embeds_list.append(negative_prompt_embeds)
-
-            negative_prompt_embeds = torch.concat(negative_prompt_embeds_list, dim=-1)
-
-        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device)
-        bs_embed, seq_len, _ = prompt_embeds.shape
-        # duplicate text embeddings for each generation per prompt, using mps friendly method
-        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
-
-        if do_classifier_free_guidance:
-            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
-            seq_len = negative_prompt_embeds.shape[1]
-            negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device)
-            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
-            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
-
-        pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
-            bs_embed * num_images_per_prompt, -1
-        )
-        if do_classifier_free_guidance:
-            negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
-                bs_embed * num_images_per_prompt, -1
-            )
-
-        return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds
-
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
-    def prepare_extra_step_kwargs(self, generator, eta):
-        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
-        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
-        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
-        # and should be between [0, 1]
-
-        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
-        extra_step_kwargs = {}
-        if accepts_eta:
-            extra_step_kwargs["eta"] = eta
-
-        # check if the scheduler accepts generator
-        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
-        if accepts_generator:
-            extra_step_kwargs["generator"] = generator
-        return extra_step_kwargs
-
-    def check_inputs(
-        self,
-        prompt,
-        prompt_2,
-        height,
-        width,
-        callback_steps,
-        negative_prompt=None,
-        negative_prompt_2=None,
-        prompt_embeds=None,
-        negative_prompt_embeds=None,
-        pooled_prompt_embeds=None,
-        negative_pooled_prompt_embeds=None,
-    ):
-        if height % 8 != 0 or width % 8 != 0:
-            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
-
-        if (callback_steps is None) or (
-            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
-        ):
-            raise ValueError(
-                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
-                f" {type(callback_steps)}."
-            )
-
-        if prompt is not None and prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
-                " only forward one of the two."
-            )
-        elif prompt_2 is not None and prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
-                " only forward one of the two."
-            )
-        elif prompt is None and prompt_embeds is None:
-            raise ValueError(
-                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
-            )
-        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
-            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
-        elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)):
-            raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}")
-
-        if negative_prompt is not None and negative_prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
-                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
-            )
-        elif negative_prompt_2 is not None and negative_prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_prompt_embeds`:"
-                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
-            )
-
-        if prompt_embeds is not None and negative_prompt_embeds is not None:
-            if prompt_embeds.shape != negative_prompt_embeds.shape:
-                raise ValueError(
-                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
-                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
-                    f" {negative_prompt_embeds.shape}."
-                )
-
-        if prompt_embeds is not None and pooled_prompt_embeds is None:
-            raise ValueError(
-                "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`."
-            )
-
-        if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None:
-            raise ValueError(
-                "If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed. Make sure to generate `negative_pooled_prompt_embeds` from the same text encoder that was used to generate `negative_prompt_embeds`."
-            )
-
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
-    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
-        shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
-        if isinstance(generator, list) and len(generator) != batch_size:
-            raise ValueError(
-                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
-                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
-            )
-
-        if latents is None:
-            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
-        else:
-            latents = latents.to(device)
-
-        # scale the initial noise by the standard deviation required by the scheduler
-        latents = latents * self.scheduler.init_noise_sigma
-        return latents
-
-    def _get_add_time_ids(self, original_size, crops_coords_top_left, target_size, dtype):
-        add_time_ids = list(original_size + crops_coords_top_left + target_size)
-
-        passed_add_embed_dim = (
-            self.unet.config.addition_time_embed_dim * len(add_time_ids) + self.text_encoder_2.config.projection_dim
-        )
-        expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features
-
-        if expected_add_embed_dim != passed_add_embed_dim:
-            raise ValueError(
-                f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`."
-            )
-
-        add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
-        return add_time_ids
-
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae
-    def upcast_vae(self):
-        dtype = self.vae.dtype
-        self.vae.to(dtype=torch.float32)
-        use_torch_2_0_or_xformers = isinstance(
-            self.vae.decoder.mid_block.attentions[0].processor,
-            (
-                AttnProcessor2_0,
-                XFormersAttnProcessor,
-                LoRAXFormersAttnProcessor,
-                LoRAAttnProcessor2_0,
-            ),
-        )
-        # if xformers or torch_2_0 is used attention block does not need
-        # to be in float32 which can save lots of memory
-        if use_torch_2_0_or_xformers:
-            self.vae.post_quant_conv.to(dtype)
-            self.vae.decoder.conv_in.to(dtype)
-            self.vae.decoder.mid_block.to(dtype)
-
-    @torch.no_grad()
-    @replace_example_docstring(EXAMPLE_DOC_STRING)
-    def __call__(
-        self,
-        prompt: Union[str, List[str]] = None,
-        prompt_2: Optional[Union[str, List[str]]] = None,
-        height: Optional[int] = None,
-        width: Optional[int] = None,
-        num_inference_steps: int = 50,
-        denoising_end: Optional[float] = None,
-        guidance_scale: float = 5.0,
-        negative_prompt: Optional[Union[str, List[str]]] = None,
-        negative_prompt_2: Optional[Union[str, List[str]]] = None,
-        num_images_per_prompt: Optional[int] = 1,
-        eta: float = 0.0,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.FloatTensor] = None,
-        prompt_embeds: Optional[torch.FloatTensor] = None,
-        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
-        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
-        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
-        output_type: Optional[str] = "pil",
-        return_dict: bool = True,
-        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
-        callback_steps: int = 1,
-        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
-        guidance_rescale: float = 0.0,
-        original_size: Optional[Tuple[int, int]] = None,
-        crops_coords_top_left: Tuple[int, int] = (0, 0),
-        target_size: Optional[Tuple[int, int]] = None,
-        negative_original_size: Optional[Tuple[int, int]] = None,
-        negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
-        negative_target_size: Optional[Tuple[int, int]] = None,
-    ):
-        r"""
-        Function invoked when calling the pipeline for generation.
-
-        Args:
-            prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
-                instead.
-            prompt_2 (`str` or `List[str]`, *optional*):
-                The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
-                used in both text-encoders
-            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
-                The height in pixels of the generated image. This is set to 1024 by default for the best results.
-                Anything below 512 pixels won't work well for
-                [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
-                and checkpoints that are not specifically fine-tuned on low resolutions.
-            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
-                The width in pixels of the generated image. This is set to 1024 by default for the best results.
-                Anything below 512 pixels won't work well for
-                [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
-                and checkpoints that are not specifically fine-tuned on low resolutions.
-            num_inference_steps (`int`, *optional*, defaults to 50):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
-                expense of slower inference.
-            denoising_end (`float`, *optional*):
-                When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be
-                completed before it is intentionally prematurely terminated. As a result, the returned sample will
-                still retain a substantial amount of noise as determined by the discrete timesteps selected by the
-                scheduler. The denoising_end parameter should ideally be utilized when this pipeline forms a part of a
-                "Mixture of Denoisers" multi-pipeline setup, as elaborated in [**Refining the Image
-                Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output)
-            guidance_scale (`float`, *optional*, defaults to 5.0):
-                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
-                `guidance_scale` is defined as `w` of equation 2. of [Imagen
-                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
-                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
-                usually at the expense of lower image quality.
-            negative_prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts not to guide the image generation. If not defined, one has to pass
-                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
-                less than `1`).
-            negative_prompt_2 (`str` or `List[str]`, *optional*):
-                The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
-                `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
-            num_images_per_prompt (`int`, *optional*, defaults to 1):
-                The number of images to generate per prompt.
-            eta (`float`, *optional*, defaults to 0.0):
-                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
-                [`schedulers.DDIMScheduler`], will be ignored for others.
-            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
-                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
-                to make generation deterministic.
-            latents (`torch.FloatTensor`, *optional*):
-                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
-                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor will ge generated by sampling using the supplied random `generator`.
-            prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
-                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
-                argument.
-            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
-                If not provided, pooled text embeddings will be generated from `prompt` input argument.
-            negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
-                weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
-                input argument.
-            output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generate image. Choose between
-                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] instead
-                of a plain tuple.
-            callback (`Callable`, *optional*):
-                A function that will be called every `callback_steps` steps during inference. The function will be
-                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
-            callback_steps (`int`, *optional*, defaults to 1):
-                The frequency at which the `callback` function will be called. If not specified, the callback will be
-                called at every step.
-            cross_attention_kwargs (`dict`, *optional*):
-                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
-                `self.processor` in
-                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
-            guidance_rescale (`float`, *optional*, defaults to 0.7):
-                Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are
-                Flawed](https://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of
-                [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf).
-                Guidance rescale factor should fix overexposure when using zero terminal SNR.
-            original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
-                If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
-                `original_size` defaults to `(width, height)` if not specified. Part of SDXL's micro-conditioning as
-                explained in section 2.2 of
-                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
-            crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
-                `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
-                `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
-                `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
-                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
-            target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
-                For most cases, `target_size` should be set to the desired height and width of the generated image. If
-                not specified it will default to `(width, height)`. Part of SDXL's micro-conditioning as explained in
-                section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
-            negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
-                To negatively condition the generation process based on a specific image resolution. Part of SDXL's
-                micro-conditioning as explained in section 2.2 of
-                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
-                information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
-            negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
-                To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
-                micro-conditioning as explained in section 2.2 of
-                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
-                information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
-            negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
-                To negatively condition the generation process based on a target image resolution. It should be as same
-                as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
-                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
-                information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
-
-        Examples:
-
-        Returns:
-            [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] or `tuple`:
-            [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a
-            `tuple`. When returning a tuple, the first element is a list with the generated images.
-        """
-        # 0. Default height and width to unet
-        height = height or self.default_sample_size * self.vae_scale_factor
-        width = width or self.default_sample_size * self.vae_scale_factor
-
-        original_size = original_size or (height, width)
-        target_size = target_size or (height, width)
-
-        # 1. Check inputs. Raise error if not correct
-        self.check_inputs(
-            prompt,
-            prompt_2,
-            height,
-            width,
-            callback_steps,
-            negative_prompt,
-            negative_prompt_2,
-            prompt_embeds,
-            negative_prompt_embeds,
-            pooled_prompt_embeds,
-            negative_pooled_prompt_embeds,
-        )
-
-        # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        device = self._execution_device
-
-        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
-        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
-        # corresponds to doing no classifier free guidance.
-        do_classifier_free_guidance = guidance_scale > 1.0
-
-        # 3. Encode input prompt
-        text_encoder_lora_scale = (
-            cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None
-        )
-        (
-            prompt_embeds,
-            negative_prompt_embeds,
-            pooled_prompt_embeds,
-            negative_pooled_prompt_embeds,
-        ) = self.encode_prompt(
-            prompt=prompt,
-            prompt_2=prompt_2,
-            device=device,
-            num_images_per_prompt=num_images_per_prompt,
-            do_classifier_free_guidance=do_classifier_free_guidance,
-            negative_prompt=negative_prompt,
-            negative_prompt_2=negative_prompt_2,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            pooled_prompt_embeds=pooled_prompt_embeds,
-            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
-            lora_scale=text_encoder_lora_scale,
-        )
-
-        # 4. Prepare timesteps
-        self.scheduler.set_timesteps(num_inference_steps, device=device)
-
-        timesteps = self.scheduler.timesteps
-
-        # 5. Prepare latent variables
-        num_channels_latents = self.unet.config.in_channels
-        latents = self.prepare_latents(
-            batch_size * num_images_per_prompt,
-            num_channels_latents,
-            height,
-            width,
-            prompt_embeds.dtype,
-            device,
-            generator,
-            latents,
-        )
-
-        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
-        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
-
-        # 7. Prepare added time ids & embeddings
-        add_text_embeds = pooled_prompt_embeds
-        add_time_ids = self._get_add_time_ids(
-            original_size, crops_coords_top_left, target_size, dtype=prompt_embeds.dtype
-        )
-        if negative_original_size is not None and negative_target_size is not None:
-            negative_add_time_ids = self._get_add_time_ids(
-                negative_original_size,
-                negative_crops_coords_top_left,
-                negative_target_size,
-                dtype=prompt_embeds.dtype,
-            )
-        else:
-            negative_add_time_ids = add_time_ids
-
-        if do_classifier_free_guidance:
-            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
-            add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)
-            add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)
-
-        prompt_embeds = prompt_embeds.to(device)
-        add_text_embeds = add_text_embeds.to(device)
-        add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1)
-
-        # 8. Denoising loop
-        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
-
-        # 7.1 Apply denoising_end
-        if denoising_end is not None and isinstance(denoising_end, float) and denoising_end > 0 and denoising_end < 1:
-            discrete_timestep_cutoff = int(
-                round(
-                    self.scheduler.config.num_train_timesteps
-                    - (denoising_end * self.scheduler.config.num_train_timesteps)
-                )
-            )
-            num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps)))
-            timesteps = timesteps[:num_inference_steps]
-
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                # expand the latents if we are doing classifier free guidance
-                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
-
-                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
-
-                # predict the noise residual
-                added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
-                noise_pred = self.unet(
-                    latent_model_input,
-                    t,
-                    encoder_hidden_states=prompt_embeds,
-                    cross_attention_kwargs=cross_attention_kwargs,
-                    added_cond_kwargs=added_cond_kwargs,
-                    return_dict=False,
-                )[0]
-
-                # perform guidance
-                if do_classifier_free_guidance:
-                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
-
-                if do_classifier_free_guidance and guidance_rescale > 0.0:
-                    # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
-                    noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
-
-                # compute the previous noisy sample x_t -> x_t-1
-                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
-
-                # call the callback, if provided
-                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
-                    progress_bar.update()
-                    if callback is not None and i % callback_steps == 0:
-                        callback(i, t, latents)
-
-        if not output_type == "latent":
-            # make sure the VAE is in float32 mode, as it overflows in float16
-            needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast
-
-            if needs_upcasting:
-                self.upcast_vae()
-                latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)
-
-            image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
-
-            # cast back to fp16 if needed
-            if needs_upcasting:
-                self.vae.to(dtype=torch.float16)
-        else:
-            image = latents
-
-        if not output_type == "latent":
-            # apply watermark if available
-            if self.watermark is not None:
-                image = self.watermark.apply_watermark(image)
-
-            image = self.image_processor.postprocess(image, output_type=output_type)
-
-        # Offload all models
-        self.maybe_free_model_hooks()
-
-        if not return_dict:
-            return (image,)
-
-        return StableDiffusionXLPipelineOutput(images=image)
-
-    # Overrride to properly handle the loading and unloading of the additional text encoder.
-    def load_lora_weights(self, pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], **kwargs):
-        # We could have accessed the unet config from `lora_state_dict()` too. We pass
-        # it here explicitly to be able to tell that it's coming from an SDXL
-        # pipeline.
-
-        # Remove any existing hooks.
-        if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"):
-            from accelerate.hooks import AlignDevicesHook, CpuOffload, remove_hook_from_module
-        else:
-            raise ImportError("Offloading requires `accelerate v0.17.0` or higher.")
-
-        is_model_cpu_offload = False
-        is_sequential_cpu_offload = False
-        recursive = False
-        for _, component in self.components.items():
-            if isinstance(component, torch.nn.Module):
-                if hasattr(component, "_hf_hook"):
-                    is_model_cpu_offload = isinstance(getattr(component, "_hf_hook"), CpuOffload)
-                    is_sequential_cpu_offload = isinstance(getattr(component, "_hf_hook"), AlignDevicesHook)
-                    logger.info(
-                        "Accelerate hooks detected. Since you have called `load_lora_weights()`, the previous hooks will be first removed. Then the LoRA parameters will be loaded and the hooks will be applied again."
-                    )
-                    recursive = is_sequential_cpu_offload
-                    remove_hook_from_module(component, recurse=recursive)
-        state_dict, network_alphas = self.lora_state_dict(
-            pretrained_model_name_or_path_or_dict,
-            unet_config=self.unet.config,
-            **kwargs,
-        )
-        self.load_lora_into_unet(state_dict, network_alphas=network_alphas, unet=self.unet)
-
-        text_encoder_state_dict = {k: v for k, v in state_dict.items() if "text_encoder." in k}
-        if len(text_encoder_state_dict) > 0:
-            self.load_lora_into_text_encoder(
-                text_encoder_state_dict,
-                network_alphas=network_alphas,
-                text_encoder=self.text_encoder,
-                prefix="text_encoder",
-                lora_scale=self.lora_scale,
-            )
-
-        text_encoder_2_state_dict = {k: v for k, v in state_dict.items() if "text_encoder_2." in k}
-        if len(text_encoder_2_state_dict) > 0:
-            self.load_lora_into_text_encoder(
-                text_encoder_2_state_dict,
-                network_alphas=network_alphas,
-                text_encoder=self.text_encoder_2,
-                prefix="text_encoder_2",
-                lora_scale=self.lora_scale,
-            )
-
-        # Offload back.
-        if is_model_cpu_offload:
-            self.enable_model_cpu_offload()
-        elif is_sequential_cpu_offload:
-            self.enable_sequential_cpu_offload()
-
-    @classmethod
-    def save_lora_weights(
-        self,
-        save_directory: Union[str, os.PathLike],
-        unet_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
-        text_encoder_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
-        text_encoder_2_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
-        is_main_process: bool = True,
-        weight_name: str = None,
-        save_function: Callable = None,
-        safe_serialization: bool = True,
-    ):
-        state_dict = {}
-
-        def pack_weights(layers, prefix):
-            layers_weights = layers.state_dict() if isinstance(layers, torch.nn.Module) else layers
-            layers_state_dict = {f"{prefix}.{module_name}": param for module_name, param in layers_weights.items()}
-            return layers_state_dict
-
-        if not (unet_lora_layers or text_encoder_lora_layers or text_encoder_2_lora_layers):
-            raise ValueError(
-                "You must pass at least one of `unet_lora_layers`, `text_encoder_lora_layers` or `text_encoder_2_lora_layers`."
-            )
-
-        if unet_lora_layers:
-            state_dict.update(pack_weights(unet_lora_layers, "unet"))
-
-        if text_encoder_lora_layers and text_encoder_2_lora_layers:
-            state_dict.update(pack_weights(text_encoder_lora_layers, "text_encoder"))
-            state_dict.update(pack_weights(text_encoder_2_lora_layers, "text_encoder_2"))
-
-        self.write_lora_layers(
-            state_dict=state_dict,
-            save_directory=save_directory,
-            is_main_process=is_main_process,
-            weight_name=weight_name,
-            save_function=save_function,
-            safe_serialization=safe_serialization,
-        )
-
-    def _remove_text_encoder_monkey_patch(self):
-        self._remove_text_encoder_monkey_patch_classmethod(self.text_encoder)
-        self._remove_text_encoder_monkey_patch_classmethod(self.text_encoder_2)
\ No newline at end of file
diff --git a/instructany2pix/diffusion/sdxl_img2img_pipeline.py b/instructany2pix/diffusion/sdxl_img2img_pipeline.py
deleted file mode 100644
index 8ba02bfb4c170ba876ebc23ec0e1a18fbddac5a0..0000000000000000000000000000000000000000
--- a/instructany2pix/diffusion/sdxl_img2img_pipeline.py
+++ /dev/null
@@ -1,70 +0,0 @@
-import torch
-#from diffusers import StableDiffusionXLImg2ImgPipeline,StableDiffusionXLPipeline
-from ..ddim.sdxl_pipeline import StableDiffusionXLPipeline
-from diffusers.utils import load_image
-from diffusers.models.unet_2d_condition import UNet2DConditionOutput
-import diffusers
-from diffusers import (
-    AutoencoderKL,
-    DDPMScheduler,
-    StableDiffusionXLPipeline,
-    UNet2DConditionModel,
-    StableDiffusionXLImg2ImgPipeline
-)
-from typing import Any, Dict, Optional, Tuple, Union
-from torch import nn
-from torch import FloatTensor,Tensor
-import torch
-
-
-class UnCLipXL(UNet2DConditionModel):
-
-    def __init__(self, adapt_project_dim=[1024,1024*4,2048],adapt_project_dim_2=[768,1024*4,1280],sample_size: int | None = None, in_channels: int = 4, out_channels: int = 4, center_input_sample: bool = False, flip_sin_to_cos: bool = True, freq_shift: int = 0, down_block_types: Tuple[str] = ..., mid_block_type: str | None = "UNetMidBlock2DCrossAttn", up_block_types: Tuple[str] = ..., only_cross_attention: bool | Tuple[bool] = False, block_out_channels: Tuple[int] = ..., layers_per_block: int | Tuple[int] = 2, downsample_padding: int = 1, mid_block_scale_factor: float = 1, dropout: float = 0, act_fn: str = "silu", norm_num_groups: int | None = 32, norm_eps: float = 0.00001, cross_attention_dim: int | Tuple[int] = 1280, transformer_layers_per_block: int | Tuple[int] | Tuple[Tuple] = 1, reverse_transformer_layers_per_block: Tuple[Tuple[int]] | None = None, encoder_hid_dim: int | None = None, encoder_hid_dim_type: str | None = None, attention_head_dim: int | Tuple[int] = 8, num_attention_heads: int | Tuple[int] | None = None, dual_cross_attention: bool = False, use_linear_projection: bool = False, class_embed_type: str | None = None, addition_embed_type: str | None = None, addition_time_embed_dim: int | None = None, num_class_embeds: int | None = None, upcast_attention: bool = False, resnet_time_scale_shift: str = "default", resnet_skip_time_act: bool = False, resnet_out_scale_factor: int = 1, time_embedding_type: str = "positional", time_embedding_dim: int | None = None, time_embedding_act_fn: str | None = None, timestep_post_act: str | None = None, time_cond_proj_dim: int | None = None, conv_in_kernel: int = 3, conv_out_kernel: int = 3, projection_class_embeddings_input_dim: int | None = None, attention_type: str = "default", class_embeddings_concat: bool = False, mid_block_only_cross_attention: bool | None = None, cross_attention_norm: str | None = None, addition_embed_type_num_heads=64):
-        super().__init__(sample_size, in_channels, out_channels, center_input_sample, flip_sin_to_cos, freq_shift, down_block_types, mid_block_type, up_block_types, only_cross_attention, block_out_channels, layers_per_block, downsample_padding, mid_block_scale_factor, dropout, act_fn, norm_num_groups, norm_eps, cross_attention_dim, transformer_layers_per_block, reverse_transformer_layers_per_block, encoder_hid_dim, encoder_hid_dim_type, attention_head_dim, num_attention_heads, dual_cross_attention, use_linear_projection, class_embed_type, addition_embed_type, addition_time_embed_dim, num_class_embeds, upcast_attention, resnet_time_scale_shift, resnet_skip_time_act, resnet_out_scale_factor, time_embedding_type, time_embedding_dim, time_embedding_act_fn, timestep_post_act, time_cond_proj_dim, conv_in_kernel, conv_out_kernel, projection_class_embeddings_input_dim, attention_type, class_embeddings_concat, mid_block_only_cross_attention, cross_attention_norm, addition_embed_type_num_heads)
-        self.adapt_project_dim = adapt_project_dim
-        self.adapt_project_dim_2 = adapt_project_dim_2
-        self.add_prjection()
-    
-    def add_prjection(self):
-        adapt_project_dim = self.adapt_project_dim
-        adapt_project_dim_2 = self.adapt_project_dim_2
-        self.projector_unclip = nn.Sequential(
-            nn.Linear(adapt_project_dim[0],adapt_project_dim[1]),
-            nn.GELU(),
-            nn.Linear(adapt_project_dim[1],adapt_project_dim[2]),
-        )
-        self.projector_unclip2 = nn.Sequential(
-            nn.Linear(adapt_project_dim_2[0],adapt_project_dim_2[1]),
-            nn.GELU(),
-            nn.Linear(adapt_project_dim_2[1],adapt_project_dim_2[2]),
-        )
-    
-    def forward(self, sample: FloatTensor, timestep: Tensor | float | int, encoder_hidden_states: Tensor, class_labels: Tensor | None = None, timestep_cond: Tensor | None = None, attention_mask: Tensor | None = None, cross_attention_kwargs: Dict[str, Any] | None = None, added_cond_kwargs: Dict[str, Tensor] | None = None, down_block_additional_residuals: Tuple[Tensor] | None = None, mid_block_additional_residual: Tensor | None = None, down_intrablock_additional_residuals: Tuple[Tensor] | None = None, encoder_attention_mask: Tensor | None = None, return_dict: bool = True) -> UNet2DConditionOutput | Tuple:
-        added_cond_kwargs['text_embeds'] = self.projector_unclip2(added_cond_kwargs['text_embeds'].to(encoder_hidden_states))
-        encoder_hidden_states = self.projector_unclip(encoder_hidden_states)
-        return super().forward(sample, timestep, encoder_hidden_states, class_labels, timestep_cond, attention_mask, cross_attention_kwargs, added_cond_kwargs, down_block_additional_residuals, mid_block_additional_residual, down_intrablock_additional_residuals, encoder_attention_mask, return_dict)
-    
-def process_clip(batch,clip_processor=None):
-    image = batch['pixel_values']
-    image = clip_processor(image)
-    batch['clip_input'] = image
-    return batch
-
-def build_sdxl(pretrained = "stabilityai/stable-diffusion-xl-base-1.0",
-               ckpt="/localhome/data/ckpts/jacklishufan/sdxl/"):
-    unet = UnCLipXL.from_pretrained(
-            ckpt, subfolder="unet"
-        )
-    vae = AutoencoderKL.from_pretrained(
-        pretrained,
-        subfolder="vae" ,
-    # revision=args.revision,
-    )
-    pipeline = StableDiffusionXLPipeline.from_pretrained(
-        pretrained,
-        vae=vae,
-        unet=unet,
-        #revision=args.revision,
-        torch_dtype=torch.float16,
-    )
-    return pipeline
\ No newline at end of file
diff --git a/instructany2pix/llm/constants.py b/instructany2pix/llm/constants.py
deleted file mode 100644
index c9d2d2c03a2d74ceb7d87adc59cb7909f79e87c3..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/constants.py
+++ /dev/null
@@ -1,30 +0,0 @@
-CONTROLLER_HEART_BEAT_EXPIRATION = 30
-WORKER_HEART_BEAT_INTERVAL = 15
-
-LOGDIR = "."
-
-# Model Constants
-IGNORE_INDEX = -100
-IMAGE_TOKEN_INDEX = -200
-DEFAULT_IMAGE_TOKEN = "<image>"
-DEFAULT_AUDIO_TOKEN = "<audio>"
-DEFAULT_VIDEO_TOKEN = "<video>"
-DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>"
-DEFAULT_IM_START_TOKEN = "<im_start>"
-DEFAULT_IM_END_TOKEN = "<im_end>"
-DEFAULT_IM_GEN_START_TOKEN = "<im_gen_start>"
-DEFAULT_AUDIO_GEN_START_TOKEN = "<au_gen_start>"
-DEFAULT_VIDEO_GEN_START_TOKEN = "<vd_gen_start>"
-DEFAULT_IM_GEN_END_TOKEN = "<im_gen_end>"
-DEFAULT_IM_GEN_TOKEN = "<im_gen>"
-DEFAULT_AUDIO_GEN_TOKEN = "<audio_gen>"
-DEFAULT_VIDEO_GEN_TOKEN = "<video_gen>"
-DEFAULT_AUDIO_GEN_START_TOKEN = "<audio_gen_start>"
-DEFAULT_MSK_TOKEN = '<mask_gen>'
-DEFAULT_BASE_TOKEN = '<base>'
-DEFAULT_BASE_NULL_TOKEN = '<base_null>'
-
-class REPLACEMENT_TYPE:
-    INPUT = 0
-    BASE = 1
-    GEN = 2
\ No newline at end of file
diff --git a/instructany2pix/llm/conversation.py b/instructany2pix/llm/conversation.py
deleted file mode 100644
index 6ee30cf9942d6f3184e845c0436f5c1d80e00043..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/conversation.py
+++ /dev/null
@@ -1,381 +0,0 @@
-import dataclasses
-from enum import auto, Enum
-from typing import List, Tuple
-
-
-class SeparatorStyle(Enum):
-    """Different separator style."""
-    SINGLE = auto()
-    TWO = auto()
-    MPT = auto()
-    PLAIN = auto()
-    LLAMA_2 = auto()
-
-
-@dataclasses.dataclass
-class Conversation:
-    """A class that keeps all conversation history."""
-    system: str
-    roles: List[str]
-    messages: List[List[str]]
-    offset: int
-    sep_style: SeparatorStyle = SeparatorStyle.SINGLE
-    sep: str = "###"
-    sep2: str = None
-    version: str = "Unknown"
-
-    skip_next: bool = False
-
-    def get_prompt(self):
-        messages = self.messages
-        if len(messages) > 0 and type(messages[0][1]) is tuple:
-            messages = self.messages.copy()
-            init_role, init_msg = messages[0].copy()
-            init_msg = init_msg[0].replace("<image>", "").strip()
-            if 'mmtag' in self.version:
-                messages[0] = (init_role, init_msg)
-                messages.insert(0, (self.roles[0], "<Image><image></Image>"))
-                messages.insert(1, (self.roles[1], "Received."))
-            else:
-                messages[0] = (init_role, "<image>\n" + init_msg)
-
-        if self.sep_style == SeparatorStyle.SINGLE:
-            ret = self.system + self.sep
-            for role, message in messages:
-                if message:
-                    if type(message) is tuple:
-                        message, _, _ = message
-                    ret += role + ": " + message + self.sep
-                else:
-                    ret += role + ":"
-        elif self.sep_style == SeparatorStyle.TWO:
-            seps = [self.sep, self.sep2]
-            ret = self.system + seps[0]
-            for i, (role, message) in enumerate(messages):
-                if message:
-                    if type(message) is tuple:
-                        message, _, _ = message
-                    ret += role + ": " + message + seps[i % 2]
-                else:
-                    ret += role + ":"
-        elif self.sep_style == SeparatorStyle.MPT:
-            ret = self.system + self.sep
-            for role, message in messages:
-                if message:
-                    if type(message) is tuple:
-                        message, _, _ = message
-                    ret += role + message + self.sep
-                else:
-                    ret += role
-        elif self.sep_style == SeparatorStyle.LLAMA_2:
-            wrap_sys = lambda msg: f"<<SYS>>\n{msg}\n<</SYS>>\n\n"
-            wrap_inst = lambda msg: f"[INST] {msg} [/INST]"
-            ret = ""
-
-            for i, (role, message) in enumerate(messages):
-                if i == 0:
-                    assert message, "first message should not be none"
-                    assert role == self.roles[0], "first message should come from user"
-                if message:
-                    if type(message) is tuple:
-                        message, _, _ = message
-                    if i == 0: message = wrap_sys(self.system) + message
-                    if i % 2 == 0:
-                        message = wrap_inst(message)
-                        ret += self.sep + message
-                    else:
-                        ret += " " + message + " " + self.sep2
-                else:
-                    ret += ""
-            ret = ret.lstrip(self.sep)
-        elif self.sep_style == SeparatorStyle.PLAIN:
-            seps = [self.sep, self.sep2]
-            ret = self.system
-            for i, (role, message) in enumerate(messages):
-                if message:
-                    if type(message) is tuple:
-                        message, _, _ = message
-                    ret += message + seps[i % 2]
-                else:
-                    ret += ""
-        else:
-            raise ValueError(f"Invalid style: {self.sep_style}")
-
-        return ret
-
-    def append_message(self, role, message):
-        self.messages.append([role, message])
-
-    def get_images(self, return_pil=False):
-        images = []
-        for i, (role, msg) in enumerate(self.messages[self.offset:]):
-            if i % 2 == 0:
-                if type(msg) is tuple:
-                    import base64
-                    from io import BytesIO
-                    from PIL import Image
-                    msg, image, image_process_mode = msg
-                    if image_process_mode == "Pad":
-                        def expand2square(pil_img, background_color=(122, 116, 104)):
-                            width, height = pil_img.size
-                            if width == height:
-                                return pil_img
-                            elif width > height:
-                                result = Image.new(pil_img.mode, (width, width), background_color)
-                                result.paste(pil_img, (0, (width - height) // 2))
-                                return result
-                            else:
-                                result = Image.new(pil_img.mode, (height, height), background_color)
-                                result.paste(pil_img, ((height - width) // 2, 0))
-                                return result
-                        image = expand2square(image)
-                    elif image_process_mode in ["Default", "Crop"]:
-                        pass
-                    elif image_process_mode == "Resize":
-                        image = image.resize((336, 336))
-                    else:
-                        raise ValueError(f"Invalid image_process_mode: {image_process_mode}")
-                    max_hw, min_hw = max(image.size), min(image.size)
-                    aspect_ratio = max_hw / min_hw
-                    max_len, min_len = 800, 400
-                    shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))
-                    longest_edge = int(shortest_edge * aspect_ratio)
-                    W, H = image.size
-                    if longest_edge != max(image.size):
-                        if H > W:
-                            H, W = longest_edge, shortest_edge
-                        else:
-                            H, W = shortest_edge, longest_edge
-                        image = image.resize((W, H))
-                    if return_pil:
-                        images.append(image)
-                    else:
-                        buffered = BytesIO()
-                        image.save(buffered, format="PNG")
-                        img_b64_str = base64.b64encode(buffered.getvalue()).decode()
-                        images.append(img_b64_str)
-        return images
-
-    def to_gradio_chatbot(self):
-        ret = []
-        for i, (role, msg) in enumerate(self.messages[self.offset:]):
-            if i % 2 == 0:
-                if type(msg) is tuple:
-                    import base64
-                    from io import BytesIO
-                    msg, image, image_process_mode = msg
-                    max_hw, min_hw = max(image.size), min(image.size)
-                    aspect_ratio = max_hw / min_hw
-                    max_len, min_len = 800, 400
-                    shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))
-                    longest_edge = int(shortest_edge * aspect_ratio)
-                    W, H = image.size
-                    if H > W:
-                        H, W = longest_edge, shortest_edge
-                    else:
-                        H, W = shortest_edge, longest_edge
-                    image = image.resize((W, H))
-                    buffered = BytesIO()
-                    image.save(buffered, format="JPEG")
-                    img_b64_str = base64.b64encode(buffered.getvalue()).decode()
-                    img_str = f'<img src="data:image/png;base64,{img_b64_str}" alt="user upload image" />'
-                    msg = img_str + msg.replace('<image>', '').strip()
-                    ret.append([msg, None])
-                else:
-                    ret.append([msg, None])
-            else:
-                ret[-1][-1] = msg
-        return ret
-
-    def copy(self):
-        return Conversation(
-            system=self.system,
-            roles=self.roles,
-            messages=[[x, y] for x, y in self.messages],
-            offset=self.offset,
-            sep_style=self.sep_style,
-            sep=self.sep,
-            sep2=self.sep2,
-            version=self.version)
-
-    def dict(self):
-        if len(self.get_images()) > 0:
-            return {
-                "system": self.system,
-                "roles": self.roles,
-                "messages": [[x, y[0] if type(y) is tuple else y] for x, y in self.messages],
-                "offset": self.offset,
-                "sep": self.sep,
-                "sep2": self.sep2,
-            }
-        return {
-            "system": self.system,
-            "roles": self.roles,
-            "messages": self.messages,
-            "offset": self.offset,
-            "sep": self.sep,
-            "sep2": self.sep2,
-        }
-
-
-conv_vicuna_v0 = Conversation(
-    system="A chat between a curious human and an artificial intelligence assistant. "
-           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
-    roles=("Human", "Assistant"),
-    messages=(
-        ("Human", "What are the key differences between renewable and non-renewable energy sources?"),
-        ("Assistant",
-            "Renewable energy sources are those that can be replenished naturally in a relatively "
-            "short amount of time, such as solar, wind, hydro, geothermal, and biomass. "
-            "Non-renewable energy sources, on the other hand, are finite and will eventually be "
-            "depleted, such as coal, oil, and natural gas. Here are some key differences between "
-            "renewable and non-renewable energy sources:\n"
-            "1. Availability: Renewable energy sources are virtually inexhaustible, while non-renewable "
-            "energy sources are finite and will eventually run out.\n"
-            "2. Environmental impact: Renewable energy sources have a much lower environmental impact "
-            "than non-renewable sources, which can lead to air and water pollution, greenhouse gas emissions, "
-            "and other negative effects.\n"
-            "3. Cost: Renewable energy sources can be more expensive to initially set up, but they typically "
-            "have lower operational costs than non-renewable sources.\n"
-            "4. Reliability: Renewable energy sources are often more reliable and can be used in more remote "
-            "locations than non-renewable sources.\n"
-            "5. Flexibility: Renewable energy sources are often more flexible and can be adapted to different "
-            "situations and needs, while non-renewable sources are more rigid and inflexible.\n"
-            "6. Sustainability: Renewable energy sources are more sustainable over the long term, while "
-            "non-renewable sources are not, and their depletion can lead to economic and social instability.\n")
-    ),
-    offset=2,
-    sep_style=SeparatorStyle.SINGLE,
-    sep="###",
-)
-
-conv_vicuna_v1 = Conversation(
-    system="A chat between a curious user and an artificial intelligence assistant. "
-    "The assistant gives helpful, detailed, and polite answers to the user's questions.",
-    roles=("USER", "ASSISTANT"),
-    version="v1",
-    messages=(),
-    offset=0,
-    sep_style=SeparatorStyle.TWO,
-    sep=" ",
-    sep2="</s>",
-)
-
-conv_llama_2 = Conversation(
-    system="""You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
-
-If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.""",
-    roles=("USER", "ASSISTANT"),
-    version="llama_v2",
-    messages=(),
-    offset=0,
-    sep_style=SeparatorStyle.LLAMA_2,
-    sep="<s>",
-    sep2="</s>",
-)
-
-conv_llava_llama_2 = Conversation(
-    system="You are a helpful language and vision assistant. "
-           "You are able to understand the visual content that the user provides, "
-           "and assist the user with a variety of tasks using natural language.",
-    roles=("USER", "ASSISTANT"),
-    version="llama_v2",
-    messages=(),
-    offset=0,
-    sep_style=SeparatorStyle.LLAMA_2,
-    sep="<s>",
-    sep2="</s>",
-)
-
-conv_mpt = Conversation(
-    system="""<|im_start|>system
-A conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.""",
-    roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
-    version="mpt",
-    messages=(),
-    offset=0,
-    sep_style=SeparatorStyle.MPT,
-    sep="<|im_end|>",
-)
-
-conv_llava_plain = Conversation(
-    system="",
-    roles=("", ""),
-    messages=(
-    ),
-    offset=0,
-    sep_style=SeparatorStyle.PLAIN,
-    sep="\n",
-)
-
-conv_llava_v0 = Conversation(
-    system="A chat between a curious human and an artificial intelligence assistant. "
-           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
-    roles=("Human", "Assistant"),
-    messages=(
-    ),
-    offset=0,
-    sep_style=SeparatorStyle.SINGLE,
-    sep="###",
-)
-
-conv_llava_v0_mmtag = Conversation(
-    system="A chat between a curious user and an artificial intelligence assistant. "
-           "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."
-           "The visual content will be provided with the following format: <Image>visual content</Image>.",
-    roles=("Human", "Assistant"),
-    messages=(
-    ),
-    offset=0,
-    sep_style=SeparatorStyle.SINGLE,
-    sep="###",
-    version="v0_mmtag",
-)
-
-conv_llava_v1 = Conversation(
-    system="A chat between a curious human and an artificial intelligence assistant. "
-           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
-    roles=("USER", "ASSISTANT"),
-    version="v1",
-    messages=(),
-    offset=0,
-    sep_style=SeparatorStyle.TWO,
-    sep=" ",
-    sep2="</s>",
-)
-
-conv_llava_v1_mmtag = Conversation(
-    system="A chat between a curious user and an artificial intelligence assistant. "
-           "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."
-           "The visual content will be provided with the following format: <Image>visual content</Image>.",
-    roles=("USER", "ASSISTANT"),
-    messages=(),
-    offset=0,
-    sep_style=SeparatorStyle.TWO,
-    sep=" ",
-    sep2="</s>",
-    version="v1_mmtag",
-)
-
-default_conversation = conv_vicuna_v0
-conv_templates = {
-    "default": conv_vicuna_v0,
-    "v0": conv_vicuna_v0,
-    "v1": conv_vicuna_v1,
-    "vicuna_v1": conv_vicuna_v1,
-    "llama_2": conv_llama_2,
-
-    "plain": conv_llava_plain,
-    "v0_plain": conv_llava_plain,
-    "llava_v0": conv_llava_v0,
-    "v0_mmtag": conv_llava_v0_mmtag,
-    "llava_v1": conv_llava_v1,
-    "v1_mmtag": conv_llava_v1_mmtag,
-    "llava_llama_2": conv_llava_llama_2,
-
-    "mpt": conv_mpt,
-}
-
-
-if __name__ == "__main__":
-    print(default_conversation.get_prompt())
diff --git a/instructany2pix/llm/mm_utils.py b/instructany2pix/llm/mm_utils.py
deleted file mode 100644
index 765d80efde42b89a6dd7485e9ee9e890bfe477a2..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/mm_utils.py
+++ /dev/null
@@ -1,107 +0,0 @@
-from typing import Any
-from PIL import Image
-from io import BytesIO
-import base64
-
-import torch
-from transformers import StoppingCriteria
-from instructany2pix.llm.constants import IMAGE_TOKEN_INDEX
-
-
-def load_image_from_base64(image):
-    return Image.open(BytesIO(base64.b64decode(image)))
-
-
-def expand2square(pil_img, background_color):
-    width, height = pil_img.size
-    if width == height:
-        return pil_img
-    elif width > height:
-        result = Image.new(pil_img.mode, (width, width), background_color)
-        result.paste(pil_img, (0, (width - height) // 2))
-        return result
-    else:
-        result = Image.new(pil_img.mode, (height, height), background_color)
-        result.paste(pil_img, ((height - width) // 2, 0))
-        return result
-
-
-def process_images(images, image_processor, model_cfg):
-    image_aspect_ratio = getattr(model_cfg, "image_aspect_ratio", None)
-    new_images = []
-    if image_aspect_ratio == 'pad':
-        for image in images:
-            image = expand2square(image, tuple(int(x*255) for x in image_processor.image_mean))
-            image = image_processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
-            new_images.append(image)
-    else:
-        return image_processor(images, return_tensors='pt')['pixel_values']
-    if all(x.shape == new_images[0].shape for x in new_images):
-        new_images = torch.stack(new_images, dim=0)
-    return new_images
-
-
-def tokenizer_image_token(prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors=None):
-    prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split('<image>')]
-
-    def insert_separator(X, sep):
-        return [ele for sublist in zip(X, [sep]*len(X)) for ele in sublist][:-1]
-
-    input_ids = []
-    offset = 0
-    if len(prompt_chunks) > 0 and len(prompt_chunks[0]) > 0 and prompt_chunks[0][0] == tokenizer.bos_token_id:
-        offset = 1
-        input_ids.append(prompt_chunks[0][0])
-
-    for x in insert_separator(prompt_chunks, [image_token_index] * (offset + 1)):
-        input_ids.extend(x[offset:])
-
-    if return_tensors is not None:
-        if return_tensors == 'pt':
-            return torch.tensor(input_ids, dtype=torch.long)
-        raise ValueError(f'Unsupported tensor type: {return_tensors}')
-    return input_ids
-
-
-def get_model_name_from_path(model_path):
-    model_path = model_path.strip("/")
-    model_paths = model_path.split("/")
-    if model_paths[-1].startswith('checkpoint-'):
-        return model_paths[-2] + "_" + model_paths[-1]
-    else:
-        return model_paths[-1]
-
-
-
-
-class KeywordsStoppingCriteria(StoppingCriteria):
-    def __init__(self, keywords, tokenizer, input_ids):
-        self.keywords = keywords
-        self.keyword_ids = []
-        self.max_keyword_len = 0
-        for keyword in keywords:
-            cur_keyword_ids = tokenizer(keyword).input_ids
-            if len(cur_keyword_ids) > 1 and cur_keyword_ids[0] == tokenizer.bos_token_id:
-                cur_keyword_ids = cur_keyword_ids[1:]
-            if len(cur_keyword_ids) > self.max_keyword_len:
-                self.max_keyword_len = len(cur_keyword_ids)
-            self.keyword_ids.append(torch.tensor(cur_keyword_ids))
-        self.tokenizer = tokenizer
-        self.start_len = input_ids.shape[1]
-
-    def call_bse(self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
-        assert output_ids.shape[0] == 1, "Only support batch size 1 (yet)"  # TODO
-        offset = min(output_ids.shape[1] - self.start_len, self.max_keyword_len)
-        self.keyword_ids = [keyword_id.to(output_ids.device) for keyword_id in self.keyword_ids]
-        for keyword_id in self.keyword_ids:
-            if (output_ids[0, -keyword_id.shape[0]:] == keyword_id).all():
-                return True
-        outputs = self.tokenizer.batch_decode(output_ids[:, -offset:], skip_special_tokens=True)[0]
-        for keyword in self.keywords:
-            if keyword in outputs:
-                return True
-        return False
-    
-    def __call__(self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
-        res = [self.call_bse(output_ids[i:i+1],scores=None,**kwargs) for i,z in enumerate(output_ids) ]
-        return all(res)
\ No newline at end of file
diff --git a/instructany2pix/llm/model/__init__.py b/instructany2pix/llm/model/__init__.py
deleted file mode 100644
index d62abafe43a2734c7ae15c32e9d8eaa302dff625..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/__init__.py
+++ /dev/null
@@ -1 +0,0 @@
-from .language_model.any2pix_llama import InstructAny2PixLMForCausalLM, InstructAny2PixLMConfig
diff --git a/instructany2pix/llm/model/any2pix_arch.py b/instructany2pix/llm/model/any2pix_arch.py
deleted file mode 100644
index 76e268f2f302da80dffefb376e3e01a113461bfa..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/any2pix_arch.py
+++ /dev/null
@@ -1,299 +0,0 @@
-#    Copyright 2023 Haotian Liu
-#
-#    Licensed under the Apache License, Version 2.0 (the "License");
-#    you may not use this file except in compliance with the License.
-#    You may obtain a copy of the License at
-#
-#        http://www.apache.org/licenses/LICENSE-2.0
-#
-#    Unless required by applicable law or agreed to in writing, software
-#    distributed under the License is distributed on an "AS IS" BASIS,
-#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-#    See the License for the specific language governing permissions and
-#    limitations under the License.
-
-
-from abc import ABC, abstractmethod
-
-import torch
-import torch.nn as nn
-
-from .multimodal_encoder.builder import build_vision_tower
-from .multimodal_projector.builder import build_vision_projector,build_vision_predictor
-
-from instructany2pix.llm.constants import (IGNORE_INDEX, IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN,
-                              DEFAULT_IM_END_TOKEN,DEFAULT_IM_GEN_TOKEN,DEFAULT_MSK_TOKEN,DEFAULT_IM_GEN_START_TOKEN,
-                              DEFAULT_AUDIO_GEN_TOKEN,DEFAULT_AUDIO_GEN_START_TOKEN,DEFAULT_AUDIO_TOKEN,
-                              DEFAULT_BASE_TOKEN,DEFAULT_VIDEO_TOKEN,DEFAULT_BASE_NULL_TOKEN)
-import logging
-from .vae.builder import build_vae,VQVAE
-class InstructAny2PixLMMetaModel:
-
-    def __init__(self, config):
-        super(InstructAny2PixLMMetaModel, self).__init__(config)
-
-        if hasattr(config, "mm_vision_tower"):
-            print("pre_init")
-            self.vision_tower = build_vision_tower(config, delay_load=False)
-            self.mm_projector = build_vision_projector(config)
-            self.vae = VQVAE(self.config.vae_image,self.config.vae_audio)
-            self.generator_processor = self.vae.processor
-            self.vae_projector_image = build_vision_projector(self.config,input_size=self.vae.embed_dim_image)
-            self.vae_projector_audio = build_vision_projector(self.config,input_size=self.vae.embed_dim_audio)
-            self.vae_predictor_image = build_vision_predictor(self.config,output_size=self.vae.vocab_size_image)
-            self.vae_predictor_audio = build_vision_predictor(self.config,output_size=self.vae.vocab_size_audio)
-
-            #raise NotImplemented
-
-    def get_vision_tower(self):
-        vision_tower = getattr(self, 'vision_tower', None)
-        if type(vision_tower) is list:
-            vision_tower = vision_tower[0]
-        return vision_tower
-
-    def get_vae(self):
-        return self.vae
-
-    def initialize_vision_modules(self, model_args, fsdp=None):
-        vision_tower = model_args.vision_tower
-        mm_vision_select_layer = model_args.mm_vision_select_layer
-        mm_vision_select_feature = model_args.mm_vision_select_feature
-        pretrain_mm_mlp_adapter = model_args.pretrain_mm_mlp_adapter
-        print('load_vision_tower')
-        self.config.mm_vision_tower = vision_tower
-
-        vision_tower = build_vision_tower(model_args)
-
-        if fsdp is not None and len(fsdp) > 0:
-            self.vision_tower = [vision_tower]
-        else:
-            self.vision_tower = vision_tower
-
-        self.config.use_mm_proj = True
-        self.config.mm_projector_type = getattr(model_args, 'mm_projector_type', 'linear')
-        self.config.mm_hidden_size = vision_tower.hidden_size
-        self.config.mm_vision_select_layer = mm_vision_select_layer
-        self.config.mm_vision_select_feature = mm_vision_select_feature
-
-        self.mm_projector = build_vision_projector(self.config)
-
-        if pretrain_mm_mlp_adapter is not None:
-            mm_projector_weights = torch.load(pretrain_mm_mlp_adapter, map_location='cpu')
-            def get_w(weights, keyword):
-                return {k.split(keyword + '.')[1]: v for k, v in weights.items() if keyword in k}
-
-            self.mm_projector.load_state_dict(get_w(mm_projector_weights, 'mm_projector'))
-        # generation
-        self.config.vae_image  = ''
-        self.config.vae_audio = ''
-        if model_args.vae_image:
-            self.config.vae_image = model_args.vae_image
-        if model_args.vae_audio:
-            self.config.vae_audio = model_args.vae_audio  
-        print('load_vae')
-        self.vae = VQVAE(self.config.vae_image,self.config.vae_audio)
-        self.generator_processor = self.vae.processor
-        self.vae_projector_image = build_vision_projector(self.config,input_size=self.vae.embed_dim_image)
-        self.vae_projector_audio = build_vision_projector(self.config,input_size=self.vae.embed_dim_audio)
-        self.vae_predictor_image = build_vision_predictor(self.config,output_size=self.vae.vocab_size_image)
-        self.vae_predictor_audio = build_vision_predictor(self.config,output_size=self.vae.vocab_size_audio)
-
-class InstructAny2PixLMMetaForCausalLM(ABC):
-
-    @abstractmethod
-    def get_model(self):
-        pass
-
-    def get_vision_tower(self):
-        return self.get_model().get_vision_tower()
-
-    def get_vae(self):
-        return self.get_model().get_vae()
-
-    def encode_images(self, images):
-        image_features = self.get_model().get_vision_tower()(images)
-        image_features = self.get_model().mm_projector(image_features)
-        return image_features
-
-    def prepare_inputs_labels_for_multimodal(
-        self, input_ids, attention_mask, past_key_values, labels, images,novision=False
-    ):
-        vision_tower = self.get_vision_tower()
-        if vision_tower is None or images is None or input_ids.shape[1] == 1 or novision:
-            if past_key_values is not None and vision_tower is not None and images is not None and input_ids.shape[1] == 1:
-                attention_mask = torch.ones((attention_mask.shape[0], past_key_values[-1][-1].shape[-2] + 1), dtype=attention_mask.dtype, device=attention_mask.device)
-            return None, attention_mask, past_key_values, self.get_model().embed_tokens(input_ids), labels
-
-        if False: # and type(images) is list or images.ndim == 5:
-            concat_images = torch.cat([image for image in images], dim=0)
-            image_features = self.encode_images(concat_images)
-            split_sizes = [image.shape[0] for image in images]
-            image_features = torch.split(image_features, split_sizes, dim=0)
-            image_features = [x.flatten(0, 1) for x in image_features]
-        else:
-            assert type(images)== dict
-            image_features = self.encode_images(images)
-
-        new_input_embeds = []
-        new_labels = [] if labels is not None else None
-        cur_image_idx = 0
-        for batch_idx, cur_input_ids in enumerate(input_ids):
-            if (cur_input_ids == IMAGE_TOKEN_INDEX).sum() == 0:
-                # multimodal LLM, but the current sample is not multimodal
-                # FIXME: this is a hacky fix, for deepspeed zero3 to work
-                half_len = cur_input_ids.shape[0] // 2
-                cur_image_features = image_features[cur_image_idx]
-                cur_input_embeds_1 = self.get_model().embed_tokens(cur_input_ids[:half_len])
-                cur_input_embeds_2 = self.get_model().embed_tokens(cur_input_ids[half_len:])
-                cur_input_embeds = torch.cat([cur_input_embeds_1, cur_image_features[0:0], cur_input_embeds_2], dim=0)
-                new_input_embeds.append(cur_input_embeds)
-                if labels is not None:
-                    new_labels.append(labels[batch_idx])
-                cur_image_idx += 1
-                continue
-            image_token_indices = torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0]
-            cur_new_input_embeds = []
-            if labels is not None:
-                cur_labels = labels[batch_idx]
-                cur_new_labels = []
-                assert cur_labels.shape == cur_input_ids.shape
-            while image_token_indices.numel() > 0:
-                cur_image_features = image_features[cur_image_idx]
-                if len(cur_image_features.shape) == 1:
-                    cur_image_features = cur_image_features[None,]
-                image_token_start = image_token_indices[0]
-                if getattr(self.config, 'tune_mm_mlp_adapter', False) and getattr(self.config, 'mm_use_im_start_end', False):
-                    cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids[:image_token_start-1]).detach())
-                    cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids[image_token_start-1:image_token_start]))
-                    cur_new_input_embeds.append(cur_image_features)
-                    cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids[image_token_start+1:image_token_start+2]))
-                    if labels is not None:
-                        cur_new_labels.append(cur_labels[:image_token_start])
-                        cur_new_labels.append(torch.full((cur_image_features.shape[0],), IGNORE_INDEX, device=labels.device, dtype=labels.dtype))
-                        cur_new_labels.append(cur_labels[image_token_start:image_token_start+1])
-                        cur_labels = cur_labels[image_token_start+2:]
-                else:
-                    cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids[:image_token_start]))
-                    cur_new_input_embeds.append(cur_image_features)
-                    if labels is not None:
-                        cur_new_labels.append(cur_labels[:image_token_start])
-                        cur_new_labels.append(torch.full((cur_image_features.shape[0],), IGNORE_INDEX, device=labels.device, dtype=labels.dtype))
-                        cur_labels = cur_labels[image_token_start+1:]
-                cur_image_idx += 1
-                if getattr(self.config, 'tune_mm_mlp_adapter', False) and getattr(self.config, 'mm_use_im_start_end', False):
-                    cur_input_ids = cur_input_ids[image_token_start+2:]
-                else:
-                    cur_input_ids = cur_input_ids[image_token_start+1:]
-                image_token_indices = torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0]
-            if cur_input_ids.numel() > 0:
-                if getattr(self.config, 'tune_mm_mlp_adapter', False) and getattr(self.config, 'mm_use_im_start_end', False):
-                    cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids).detach())
-                else:
-                    cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids))
-                if labels is not None:
-                    cur_new_labels.append(cur_labels)
-            cur_new_input_embeds = [x.to(device=self.device) for x in cur_new_input_embeds]
-            cur_new_input_embeds = torch.cat(cur_new_input_embeds, dim=0)
-            new_input_embeds.append(cur_new_input_embeds)
-            if labels is not None:
-                cur_new_labels = torch.cat(cur_new_labels, dim=0)
-                new_labels.append(cur_new_labels)
-
-        if any(x.shape != new_input_embeds[0].shape for x in new_input_embeds):
-            max_len = max(x.shape[0] for x in new_input_embeds)
-
-            new_input_embeds_align = []
-            for cur_new_embed in new_input_embeds:
-                cur_new_embed = torch.cat((cur_new_embed, torch.zeros((max_len - cur_new_embed.shape[0], cur_new_embed.shape[1]), dtype=cur_new_embed.dtype, device=cur_new_embed.device)), dim=0)
-                new_input_embeds_align.append(cur_new_embed)
-            new_input_embeds = torch.stack(new_input_embeds_align, dim=0)
-
-            if labels is not None:
-                new_labels_align = []
-                _new_labels = new_labels
-                for cur_new_label in new_labels:
-                    cur_new_label = torch.cat((cur_new_label, torch.full((max_len - cur_new_label.shape[0],), IGNORE_INDEX, dtype=cur_new_label.dtype, device=cur_new_label.device)), dim=0)
-                    new_labels_align.append(cur_new_label)
-                new_labels = torch.stack(new_labels_align, dim=0)
-
-            if attention_mask is not None:
-                new_attention_mask = []
-                for cur_attention_mask, cur_new_labels, cur_new_labels_align in zip(attention_mask, _new_labels, new_labels):
-                    new_attn_mask_pad_left = torch.full((cur_new_labels.shape[0] - labels.shape[1],), True, dtype=attention_mask.dtype, device=attention_mask.device)
-                    new_attn_mask_pad_right = torch.full((cur_new_labels_align.shape[0] - cur_new_labels.shape[0],), False, dtype=attention_mask.dtype, device=attention_mask.device)
-                    cur_new_attention_mask = torch.cat((new_attn_mask_pad_left, cur_attention_mask, new_attn_mask_pad_right), dim=0)
-                    new_attention_mask.append(cur_new_attention_mask)
-                attention_mask = torch.stack(new_attention_mask, dim=0)
-                assert attention_mask.shape == new_labels.shape
-        else:
-            new_input_embeds = torch.stack(new_input_embeds, dim=0)
-            if labels is not None:
-                new_labels  = torch.stack(new_labels, dim=0)
-
-            if attention_mask is not None:
-                new_attn_mask_pad_left = torch.full((attention_mask.shape[0], new_input_embeds.shape[1] - input_ids.shape[1]), True, dtype=attention_mask.dtype, device=attention_mask.device)
-                attention_mask = torch.cat((new_attn_mask_pad_left, attention_mask), dim=1)
-                assert attention_mask.shape == new_input_embeds.shape[:2]
-
-        return None, attention_mask, past_key_values, new_input_embeds, new_labels
-
-    def initialize_vision_tokenizer(self, model_args, tokenizer,pre_init=False):
-        if not pre_init:
-            if model_args.mm_use_im_patch_token:
-                tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
-                self.resize_token_embeddings(len(tokenizer))
-
-            if model_args.mm_use_im_start_end:
-                num_new_tokens = tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True)
-                self.resize_token_embeddings(len(tokenizer))
-
-                if num_new_tokens > 0:
-                    input_embeddings = self.get_input_embeddings().weight.data
-                    output_embeddings = self.get_output_embeddings().weight.data
-
-                    input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(
-                        dim=0, keepdim=True)
-                    output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(
-                        dim=0, keepdim=True)
-
-                    input_embeddings[-num_new_tokens:] = input_embeddings_avg
-                    output_embeddings[-num_new_tokens:] = output_embeddings_avg
-
-                if model_args.tune_mm_mlp_adapter:
-                    for p in self.get_input_embeddings().parameters():
-                        p.requires_grad = True
-                    for p in self.get_output_embeddings().parameters():
-                        p.requires_grad = False
-
-                if model_args.pretrain_mm_mlp_adapter:
-                    mm_projector_weights = torch.load(model_args.pretrain_mm_mlp_adapter, map_location='cpu')
-                    embed_tokens_weight = mm_projector_weights['model.embed_tokens.weight']
-                    assert num_new_tokens == 2
-                    if input_embeddings.shape == embed_tokens_weight.shape:
-                        input_embeddings[-num_new_tokens:] = embed_tokens_weight[-num_new_tokens:]
-                    elif embed_tokens_weight.shape[0] == num_new_tokens:
-                        input_embeddings[-num_new_tokens:] = embed_tokens_weight
-                    else:
-                        raise ValueError(f"Unexpected embed_tokens_weight shape. Pretrained: {embed_tokens_weight.shape}. Current: {input_embeddings.shape}. Numer of new tokens: {num_new_tokens}.")
-            elif model_args.mm_use_im_patch_token:
-                if model_args.tune_mm_mlp_adapter:
-                    for p in self.get_input_embeddings().parameters():
-                        p.requires_grad = False
-                    for p in self.get_output_embeddings().parameters():
-                        p.requires_grad = False
-
-            if model_args.mm_use_gen:
-                tokenizer.add_tokens([DEFAULT_IM_GEN_START_TOKEN,DEFAULT_IM_GEN_TOKEN,DEFAULT_MSK_TOKEN,DEFAULT_AUDIO_GEN_TOKEN,DEFAULT_AUDIO_GEN_START_TOKEN,DEFAULT_AUDIO_TOKEN,
-                    DEFAULT_VIDEO_TOKEN,DEFAULT_BASE_TOKEN,DEFAULT_BASE_NULL_TOKEN], special_tokens=True)
-                self.resize_token_embeddings(len(tokenizer))
-        #tokenizer(DEFAULT_IM_GEN_TOKEN,add_special_tokens=False)['input_ids'][0]
-        self.DEFAULT_IM_GEN_TOKEN_IDX = tokenizer(DEFAULT_IM_GEN_START_TOKEN,add_special_tokens=False)['input_ids'][0]
-        self.DEFAULT_IM_GEN_TOKEN_IDX = tokenizer(DEFAULT_IM_GEN_TOKEN,add_special_tokens=False)['input_ids'][0]
-        self.DEFAULT_MSK_TOKEN_IDX = tokenizer(DEFAULT_MSK_TOKEN,add_special_tokens=False)['input_ids'][0]
-        self.DEFAULT_AUDIO_GEN_TOKEN_IDX = tokenizer(DEFAULT_AUDIO_GEN_TOKEN,add_special_tokens=False)['input_ids'][0]
-        self.DEFAULT_AUDIO_GEN_START_TOKEN_IDX = tokenizer(DEFAULT_AUDIO_GEN_START_TOKEN,add_special_tokens=False)['input_ids'][0]
-        self.DEFAULT_AUDIO_TOKEN_IDX = tokenizer(DEFAULT_AUDIO_TOKEN,add_special_tokens=False)['input_ids'][0]
-        self.DEFAULT_VIDEO_TOKEN_IDX = tokenizer(DEFAULT_VIDEO_TOKEN,add_special_tokens=False)['input_ids'][0]
-        self.DEFAULT_BASE_TOKEN_IDX = tokenizer(DEFAULT_BASE_TOKEN,add_special_tokens=False)['input_ids'][0]
-        self.DEFAULT_BASE_NULL_TOKEN_IDX = tokenizer(DEFAULT_BASE_NULL_TOKEN,add_special_tokens=False)['input_ids'][0]
-        self.tokenizer = tokenizer
\ No newline at end of file
diff --git a/instructany2pix/llm/model/apply_delta.py b/instructany2pix/llm/model/apply_delta.py
deleted file mode 100644
index 4cfed6e5cbfb4286983983a49176af6c993a4123..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/apply_delta.py
+++ /dev/null
@@ -1,48 +0,0 @@
-"""
-Usage:
-python3 -m fastchat.model.apply_delta --base ~/model_weights/llama-7b --target ~/model_weights/vicuna-7b --delta lmsys/vicuna-7b-delta
-"""
-import argparse
-
-import torch
-from tqdm import tqdm
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from ..model import InstructAny2PixLMForCausalLM
-
-
-def apply_delta(base_model_path, target_model_path, delta_path):
-    print("Loading base model")
-    base = AutoModelForCausalLM.from_pretrained(
-        base_model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True)
-
-    print("Loading delta")
-    delta = InstructAny2PixLMForCausalLM.from_pretrained(delta_path, torch_dtype=torch.float16, low_cpu_mem_usage=True)
-    delta_tokenizer = AutoTokenizer.from_pretrained(delta_path)
-
-    print("Applying delta")
-    for name, param in tqdm(delta.state_dict().items(), desc="Applying delta"):
-        if name not in base.state_dict():
-            assert name in ['model.mm_projector.weight', 'model.mm_projector.bias'], f'{name} not in base model'
-            continue
-        if param.data.shape == base.state_dict()[name].shape:
-            param.data += base.state_dict()[name]
-        else:
-            assert name in ['model.embed_tokens.weight', 'lm_head.weight'], \
-                f'{name} dimension mismatch: {param.data.shape} vs {base.state_dict()[name].shape}'
-            bparam = base.state_dict()[name]
-            param.data[:bparam.shape[0], :bparam.shape[1]] += bparam
-
-    print("Saving target model")
-    delta.save_pretrained(target_model_path)
-    delta_tokenizer.save_pretrained(target_model_path)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--base-model-path", type=str, required=True)
-    parser.add_argument("--target-model-path", type=str, required=True)
-    parser.add_argument("--delta-path", type=str, required=True)
-
-    args = parser.parse_args()
-
-    apply_delta(args.base_model_path, args.target_model_path, args.delta_path)
diff --git a/instructany2pix/llm/model/builder.py b/instructany2pix/llm/model/builder.py
deleted file mode 100644
index 50b7671962e7e13eb9e5f7d3bef2d7701392880d..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/builder.py
+++ /dev/null
@@ -1,136 +0,0 @@
-#    Copyright 2023 Haotian Liu
-#
-#    Licensed under the Apache License, Version 2.0 (the "License");
-#    you may not use this file except in compliance with the License.
-#    You may obtain a copy of the License at
-#
-#        http://www.apache.org/licenses/LICENSE-2.0
-#
-#    Unless required by applicable law or agreed to in writing, software
-#    distributed under the License is distributed on an "AS IS" BASIS,
-#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-#    See the License for the specific language governing permissions and
-#    limitations under the License.
-
-
-import os
-import warnings
-import shutil
-
-from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, BitsAndBytesConfig
-import torch
-from . import *
-from instructany2pix.llm.constants import DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN,DEFAULT_AUDIO_TOKEN
-
-
-def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, load_4bit=False, device_map="auto", device="cuda"):
-    kwargs = {"device_map": device_map}
-
-    if load_8bit:
-        kwargs['load_in_8bit'] = True
-    elif load_4bit:
-        kwargs['load_in_4bit'] = True
-        kwargs['quantization_config'] = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_compute_dtype=torch.float16,
-            bnb_4bit_use_double_quant=True,
-            bnb_4bit_quant_type='nf4'
-        )
-    else:
-        kwargs['torch_dtype'] = torch.float16
-
-    if 'instructany2pix' in model_name.lower():
-        if 'lora' in model_name.lower() and model_base is None:
-            warnings.warn('There is `lora` in model name but no `model_base` is provided. If you are loading a LoRA model, please provide the `model_base` argument. Detailed instruction: https://github.com/haotian-liu/LLaVA#launch-a-model-worker-lora-weights-unmerged.')
-        if 'lora' in model_name.lower() and model_base is not None:
-            lora_cfg_pretrained = AutoConfig.from_pretrained(model_path)
-            tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
-            print('Loading instructany2pix from base model...')
-            model = InstructAny2PixLMForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=lora_cfg_pretrained, **kwargs)
-            token_num, tokem_dim = model.lm_head.out_features, model.lm_head.in_features
-            if model.lm_head.weight.shape[0] != token_num:
-                model.lm_head.weight = torch.nn.Parameter(torch.empty(token_num, tokem_dim, device=model.device, dtype=model.dtype))
-                model.model.embed_tokens.weight = torch.nn.Parameter(torch.empty(token_num, tokem_dim, device=model.device, dtype=model.dtype))
-
-            print('Loading additional instructany2pix weights...')
-            if os.path.exists(os.path.join(model_path, 'non_lora_trainables.bin')):
-                non_lora_trainables = torch.load(os.path.join(model_path, 'non_lora_trainables.bin'), map_location='cpu')
-            else:
-                # this is probably from HF Hub
-                from huggingface_hub import hf_hub_download
-                def load_from_hf(repo_id, filename, subfolder=None):
-                    cache_file = hf_hub_download(
-                        repo_id=repo_id,
-                        filename=filename,
-                        subfolder=subfolder)
-                    return torch.load(cache_file, map_location='cpu')
-                non_lora_trainables = load_from_hf(model_path, 'non_lora_trainables.bin')
-            non_lora_trainables = {(k[11:] if k.startswith('base_model.') else k): v for k, v in non_lora_trainables.items()}
-            if any(k.startswith('model.model.') for k in non_lora_trainables):
-                non_lora_trainables = {(k[6:] if k.startswith('model.') else k): v for k, v in non_lora_trainables.items()}
-            model.load_state_dict(non_lora_trainables, strict=False)
-
-            from peft import PeftModel
-            print('Loading LoRA weights...')
-            model = PeftModel.from_pretrained(model, model_path)
-            print('Merging LoRA weights...')
-            model = model.merge_and_unload()
-            print('Model is loaded...')
-        elif model_base is not None:
-            # this may be mm projector only
-            print('Loading InstructAny2Pix from base model...')
-            tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
-            cfg_pretrained = AutoConfig.from_pretrained(model_path)
-            model = InstructAny2PixLMForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=cfg_pretrained, **kwargs)
-
-            mm_projector_weights = torch.load(os.path.join(model_path, 'mm_projector.bin'), map_location='cpu')
-            mm_projector_weights = {k: v.to(torch.float16) for k, v in mm_projector_weights.items()}
-            model.load_state_dict(mm_projector_weights, strict=False)
-        else:
-            tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
-            model = InstructAny2PixLMForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
-    else:
-        # Load language model
-        if model_base is not None:
-            # PEFT model
-            from peft import PeftModel
-            tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
-            model = AutoModelForCausalLM.from_pretrained(model_base, torch_dtype=torch.float16, low_cpu_mem_usage=True, device_map="auto")
-            print(f"Loading LoRA weights from {model_path}")
-            model = PeftModel.from_pretrained(model, model_path)
-            print(f"Merging weights")
-            model = model.merge_and_unload()
-            print('Convert to FP16...')
-            model.to(torch.float16)
-        else:
-            use_fast = False
-            if 'mpt' in model_name.lower():
-                tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)
-                model = AutoModelForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, trust_remote_code=True, **kwargs)
-            else:
-                tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
-                model = AutoModelForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
-
-    image_processor = None
-
-    if 'instructany2pix' in model_name.lower():
-        mm_use_im_start_end = getattr(model.config, "mm_use_im_start_end", False)
-        mm_use_im_patch_token = getattr(model.config, "mm_use_im_patch_token", True)
-        if mm_use_im_patch_token:
-            tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
-        if mm_use_im_start_end:
-            tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True)
-        model.resize_token_embeddings(len(tokenizer))
-
-        vision_tower = model.get_vision_tower()
-        if not vision_tower.is_loaded:
-            vision_tower.load_model()
-        vision_tower.to(device=device, dtype=torch.float16)
-        image_processor = vision_tower.image_processor
-
-    if hasattr(model.config, "max_sequence_length"):
-        context_len = model.config.max_sequence_length
-    else:
-        context_len = 2048
-
-    return tokenizer, model, image_processor, context_len
diff --git a/instructany2pix/llm/model/consolidate.py b/instructany2pix/llm/model/consolidate.py
deleted file mode 100644
index ce2dc1cea1f202d8c6f7e59845b65c2c71812295..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/consolidate.py
+++ /dev/null
@@ -1,26 +0,0 @@
-
-import argparse
-
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from . import *
-from .utils import auto_upgrade
-
-
-def consolidate_ckpt(src_path, dst_path):
-    print("Loading model")
-    auto_upgrade(src_path)
-    src_model = AutoModelForCausalLM.from_pretrained(src_path, torch_dtype=torch.float16, low_cpu_mem_usage=True)
-    src_tokenizer = AutoTokenizer.from_pretrained(src_path, use_fast=False)
-    src_model.save_pretrained(dst_path)
-    src_tokenizer.save_pretrained(dst_path)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--src", type=str, required=True)
-    parser.add_argument("--dst", type=str, required=True)
-
-    args = parser.parse_args()
-
-    consolidate_ckpt(args.src, args.dst)
diff --git a/instructany2pix/llm/model/language_model/__init__.py b/instructany2pix/llm/model/language_model/__init__.py
deleted file mode 100644
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000
diff --git a/instructany2pix/llm/model/language_model/any2pix_llama.py b/instructany2pix/llm/model/language_model/any2pix_llama.py
deleted file mode 100644
index b085db636003e89e9c4095bcac1a003adbc8e3dc..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/language_model/any2pix_llama.py
+++ /dev/null
@@ -1,472 +0,0 @@
-#    Copyright 2023 Haotian Liu
-#
-#    Licensed under the Apache License, Version 2.0 (the "License");
-#    you may not use this file except in compliance with the License.
-#    You may obtain a copy of the License at
-#
-#        http://www.apache.org/licenses/LICENSE-2.0
-#
-#    Unless required by applicable law or agreed to in writing, software
-#    distributed under the License is distributed on an "AS IS" BASIS,
-#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-#    See the License for the specific language governing permissions and
-#    limitations under the License.
-
-
-from typing import List, Optional, Tuple, Union,Dict
-from einops import rearrange
-import torch
-import torch.nn as nn
-from torch.nn import CrossEntropyLoss
-
-from transformers import AutoConfig, AutoModelForCausalLM, \
-                         LlamaConfig, LlamaModel, LlamaForCausalLM
-
-from transformers.modeling_outputs import CausalLMOutputWithPast,ModelOutput
-
-from ..any2pix_arch import InstructAny2PixLMMetaModel, InstructAny2PixLMMetaForCausalLM
-from instructany2pix.llm.constants import IGNORE_INDEX
-from transformers.modeling_outputs import BaseModelOutputWithPast
-import logging
-
-logger = logging.getLogger(__name__)
-
-class InstructAny2PixLMConfig(LlamaConfig):
-    model_type = "instructany2pix"
-
-
-class InstructAny2PixLMModel(InstructAny2PixLMMetaModel, LlamaModel):
-    config_class = InstructAny2PixLMConfig
-
-    def __init__(self, config: LlamaConfig):
-        super(InstructAny2PixLMModel, self).__init__(config)
-
-    def forward(
-        self,
-        input_ids: torch.LongTensor = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[List[torch.FloatTensor]] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        #replacement_mask = None
-    ) -> Union[Tuple, BaseModelOutputWithPast]:
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-        use_cache = use_cache if use_cache is not None else self.config.use_cache
-
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        # retrieve input_ids and inputs_embeds
-        if input_ids is not None and inputs_embeds is not None:
-            raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
-        elif input_ids is not None:
-            batch_size, seq_length = input_ids.shape
-        elif inputs_embeds is not None:
-            batch_size, seq_length, _ = inputs_embeds.shape
-        else:
-            raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
-
-        seq_length_with_past = seq_length
-        past_key_values_length = 0
-
-        if past_key_values is not None:
-            past_key_values_length = past_key_values[0][0].shape[2]
-            seq_length_with_past = seq_length_with_past + past_key_values_length
-
-        if position_ids is None:
-            device = input_ids.device if input_ids is not None else inputs_embeds.device
-            position_ids = torch.arange(
-                past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
-            )
-            position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
-        else:
-            position_ids = position_ids.view(-1, seq_length).long()
-
-        if inputs_embeds is None:
-            inputs_embeds = self.embed_tokens(input_ids)
-        # embed positions
-        if attention_mask is None:
-            attention_mask = torch.ones(
-                (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
-            )
-        
-        attention_mask = self._prepare_decoder_attention_mask(
-            attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
-        ) 
-        # if replacement_mask is not None and len(attention_mask.shape) == 4:
-        #     replacement_mask = replacement_mask[:,None,:] & replacement_mask[:,:,None] # B X N X N
-        #     replacement_mask = replacement_mask[:,None] # B X H X N X N
-        #     attention_mask[replacement_mask] = 0.0
-        hidden_states = inputs_embeds
-
-        if self.gradient_checkpointing and self.training:
-            if use_cache:
-                logger.warning_once(
-                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
-                )
-                use_cache = False
-
-        # decoder layers
-        all_hidden_states = () if output_hidden_states else None
-        all_self_attns = () if output_attentions else None
-        next_decoder_cache = () if use_cache else None
-
-        for idx, decoder_layer in enumerate(self.layers):
-            if output_hidden_states:
-                all_hidden_states += (hidden_states,)
-
-            past_key_value = past_key_values[idx] if past_key_values is not None else None
-
-            if self.gradient_checkpointing and self.training:
-
-                def create_custom_forward(module):
-                    def custom_forward(*inputs):
-                        # None for past_key_value
-                        return module(*inputs, output_attentions, None)
-
-                    return custom_forward
-
-                layer_outputs = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(decoder_layer),
-                    hidden_states,
-                    attention_mask,
-                    position_ids,
-                    None,
-                )
-            else:
-                layer_outputs = decoder_layer(
-                    hidden_states,
-                    attention_mask=attention_mask,
-                    position_ids=position_ids,
-                    past_key_value=past_key_value,
-                    output_attentions=output_attentions,
-                    use_cache=use_cache,
-                )
-
-            hidden_states = layer_outputs[0]
-
-            if use_cache:
-                next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
-
-            if output_attentions:
-                all_self_attns += (layer_outputs[1],)
-
-        hidden_states = self.norm(hidden_states)
-
-        # add hidden states from the last decoder layer
-        if output_hidden_states:
-            all_hidden_states += (hidden_states,)
-
-        next_cache = next_decoder_cache if use_cache else None
-        if not return_dict:
-            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
-        return BaseModelOutputWithPast(
-            last_hidden_state=hidden_states,
-            past_key_values=next_cache,
-            hidden_states=all_hidden_states,
-            attentions=all_self_attns,
-        )
-from instructany2pix.llm.constants import REPLACEMENT_TYPE
-
-class InstructAny2PixLMForCausalLM(LlamaForCausalLM, InstructAny2PixLMMetaForCausalLM):
-    config_class = InstructAny2PixLMConfig
-
-    def __init__(self, config):
-        super(LlamaForCausalLM, self).__init__(config)
-        self.model = InstructAny2PixLMModel(config)
-
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-        #self.lm_head_img = nn.Linear(3, config.vocab_size, bias=False) # FIXME: Add config
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    def get_model(self):
-        return self.model
-
-    def forward(
-        self,
-        input_ids: torch.LongTensor = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        past_key_values: Optional[List[torch.FloatTensor]] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        images: Optional[Dict] = None,
-        return_dict: Optional[bool] = None,
-        generation_target: Optional[Dict] = None,
-        return_generations=True,
-        extra_inputs=None,
-        extra_replacement=None,
-    ) -> Union[Tuple, CausalLMOutputWithPast]:
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        
-        if self.training:
-            encodings = self.get_model().get_vae()(generation_target)
-        else:
-            encodings = {}
-        # encode vae for 
-        replacement_mask = torch.zeros_like(input_ids,dtype=bool).to(input_ids.device)
-        if 'image' in encodings:
-            quant_image,ind_image,info_image = encodings['image']
-            n = quant_image.shape[0]
-            if info_image:
-                ind_image = ind_image.view(n,-1)
-                loss_fct_img = nn.CrossEntropyLoss()
-                img_loss_obj = 'ar'
-            else:
-                ind_image = quant_image.squeeze(-1).squeeze(-1).unsqueeze(1)
-                loss_fct_img = nn.MSELoss()
-                img_loss_obj = 'latent'
-            img_embded = self.get_model().vae_projector_image(rearrange(quant_image,'n c h w -> n (h w) c') )
-            replacement_mask_img = input_ids ==  self.DEFAULT_IM_GEN_TOKEN_IDX
-            replacement_mask = replacement_mask | replacement_mask_img
-        if 'audio' in encodings:
-            quant_audio,ind_audio,info_audio = encodings['audio']
-            n = quant_audio.shape[0]
-            if info_audio:
-                ind_audio = ind_audio.view(n,-1)
-                loss_fct_aud = nn.CrossEntropyLoss()
-                audloss_obj = 'ar'
-            else:
-                ind_aud = quant_audio.squeeze(1)# N H W C -> N L C
-                loss_fct_aud = nn.MSELoss()
-                aud_loss_obj = 'latent'
-            audio_embded = self.get_model().vae_projector_audio(rearrange(quant_audio,'n h w c-> n (h w) c') )
-            replacement_mask_audio = input_ids ==  self.DEFAULT_AUDIO_GEN_TOKEN_IDX
-            replacement_mask = replacement_mask | replacement_mask_audio
-
-        # replacement_audio = input_ids ==  self.DEFAULT_AUDIO_GEN_TOKEN_IDX
-        # audio_input_embded = self.get_model().vae_projector_audio(rearrange(quant_audio,'n h w c-> n (h w) c') )
-        
-        raw_input_ids = input_ids
-        input_ids, attention_mask, past_key_values, inputs_embeds, labels = self.prepare_inputs_labels_for_multimodal(input_ids, attention_mask, past_key_values, labels, images,novision=True)
-        if extra_replacement is not None:
-            if self.training:
-                extra_replacement_mask = (raw_input_ids == self.DEFAULT_VIDEO_TOKEN_IDX ) # | (
-                # raw_input_ids == self.DEFAULT_IM_GEN_TOKEN_IDX) | (raw_input_ids == self.DEFAULT_BASE_TOKEN_IDX) self.get_model().vae_projector_image[0].weight.grad
-                if extra_replacement['mask'].shape[0] != inputs_embeds[extra_replacement_mask].shape[0]:
-                    print("SKIPPED")
-                    extra_replacement['mask'] = torch.zeros(inputs_embeds[extra_replacement_mask].shape[0]).to(extra_replacement['mask'])     
-                z = torch.zeros_like(inputs_embeds)
-                z2 = self.get_model().vae_projector_image(
-                    extra_replacement['data'][extra_replacement['mask']==REPLACEMENT_TYPE.INPUT])
-                a,b = torch.where(extra_replacement_mask)
-                z[a[extra_replacement['mask']==REPLACEMENT_TYPE.INPUT],b[extra_replacement['mask']==REPLACEMENT_TYPE.INPUT]] += z2
-                inputs_embeds[extra_replacement_mask][extra_replacement['mask']==REPLACEMENT_TYPE.INPUT] = 0.0
-                z = z + inputs_embeds
-                inputs_embeds = z
-                print("Replaced:",len(extra_replacement['mask']==REPLACEMENT_TYPE.INPUT))
-                extra_tgt_mask = (extra_replacement['mask']==REPLACEMENT_TYPE.BASE )| (extra_replacement['mask']==REPLACEMENT_TYPE.GEN)
-                extra_replacement_gt = extra_replacement['data'][extra_tgt_mask]
-                loss_fn_extra = nn.L1Loss()
-                if extra_replacement_gt.shape[0]==0:
-                    loss_fn_extra = None
-            else:
-                assert labels is None
-                extra_replacement_mask = (raw_input_ids == self.DEFAULT_VIDEO_TOKEN_IDX )
-                #print(len(extra_replacement['mask']==REPLACEMENT_TYPE.INPUT))
-
-                z = torch.zeros_like(inputs_embeds)
-                z2 = self.get_model().vae_projector_image(
-                    extra_replacement['data'][extra_replacement['mask']==REPLACEMENT_TYPE.INPUT])
-                #print("z2",z2)
-                a,b = torch.where(extra_replacement_mask)
-                a = a[:extra_replacement['mask'].shape[0]]
-                b = b[:extra_replacement['mask'].shape[0]]
-                z[a[extra_replacement['mask']==REPLACEMENT_TYPE.INPUT],b[extra_replacement['mask']==REPLACEMENT_TYPE.INPUT]] += z2
-                inputs_embeds[extra_replacement_mask][:extra_replacement['mask'].shape[0]][extra_replacement['mask']==REPLACEMENT_TYPE.INPUT] = 0.0
-                z = z + inputs_embeds
-                inputs_embeds = z
-                #print("HERE")
-                #print(inputs_embeds[extra_replacement_mask][extra_replacement['mask']==REPLACEMENT_TYPE.INPUT])
-                
-                # inputs_embeds[extra_replacement_mask][:extra_replacement['mask'].shape[0]][extra_replacement['mask']==REPLACEMENT_TYPE.INPUT] = self.get_model().vae_projector_image(
-                #     extra_replacement['data'][extra_replacement['mask']==REPLACEMENT_TYPE.INPUT].to(inputs_embeds))
-            # z.sum().backward()
-        if self.training:
-            for replace_info in generation_target['info']:
-                ii = replace_info['idx']
-                mm = replace_info['modality']
-                if mm == 'image':
-                    inputs_embeds[replace_info['batch']][replacement_mask_img[replace_info['batch']]] = img_embded[ii]
-                    #labels[replace_info['batch']][torch.where(replacement_mask_img[ii])[0]-1] just for sanity check <gen start>, <gen> ....
-                elif mm == 'audio':
-                    inputs_embeds[replace_info['batch']][replacement_mask_audio[replace_info['batch']]] = audio_embded[ii]
-        # labels[labels==self.DEFAULT_IM_GEN_TOKEN_IDX] = IGNORE_INDEX
-        # labels[labels==self.DEFAULT_AUDIO_GEN_TOKEN_IDX] = IGNORE_INDEX
-        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
-        
-        if extra_inputs is not None:
-            if 'audio' in extra_inputs:
-                audio_in = extra_inputs['audio'].squeeze(1) # N 8 768
-                audio_in = self.get_model().vae_projector_audio(audio_in)
-                msk_in = raw_input_ids ==  self.DEFAULT_AUDIO_TOKEN_IDX
-                for rinfo in extra_inputs['info']:
-                    l_mask = msk_in[rinfo['bn']]
-                    if l_mask.sum() !=8:
-                        continue
-                    else:
-                        inputs_embeds[rinfo['bn'],l_mask]=audio_in[rinfo['idx']]
-                # hack, assume has N audio inp
-        
-        outputs = self.model(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds,
-            use_cache=use_cache,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-           # replacement_mask=replacement_mask, # allow looking into future for images
-        )
-
-        hidden_states = outputs[0]
-        logits = self.lm_head(hidden_states)
-
-        loss = None
-        img_decode = None
-        aud_decode=None
-        individual_losses = {}
-        extra_gen = None
-        extra_gen_idx = None
-        if labels is not None:
-            # Shift so that tokens < n predict n
-            shift_logits = logits[..., :-1, :].contiguous()
-            shift_labels = labels[..., 1:].contiguous()
-            hidden_states.shape
-            output_vae_img = []
-            target_vae_img = []
-            output_vae_audio = []
-            target_vae_audio = []
-            for replace_info in generation_target['info']:
-                ii = replace_info['idx']
-                mm = replace_info['modality']
-                if mm == 'image':
-                    output_vae_img.append(hidden_states[replace_info['batch']][:-1][replacement_mask_img[replace_info['batch']][1:]])
-                    target_vae_img.append(ind_image[ii])
-                    #labels[replace_info['batch']][torch.where(replacement_mask_img[ii])[0]-1] just for sanity check <gen start>, <gen> ....
-                elif mm == 'audio':
-                    output_vae_audio.append(hidden_states[replace_info['batch']][:-1][replacement_mask_audio[replace_info['batch']][1:]])
-                    target_vae_audio.append(ind_aud[ii])
-            loss_fct = CrossEntropyLoss()
-
-                #prediction = logits_img.argmax(-1)
-     
-            # Flatten the tokens
-            
-            shift_logits = shift_logits.view(-1, self.config.vocab_size)
-            shift_labels = shift_labels.view(-1)
-            # Enable model/pipeline parallelism
-            shift_labels = shift_labels.to(shift_logits.device)
-            loss = loss_fct(shift_logits, shift_labels)
-            loss_lang = loss.detach().item()
-            individual_losses['loss_lang'] = loss_lang
-            if len(output_vae_img):
-                logits_img = self.get_model().vae_predictor_image(torch.cat(output_vae_img))
-                tgt_img = torch.cat(target_vae_img)
-                if img_loss_obj =='ar':
-                    tgt_img = tgt_img.view(-1) # discrete tokens
-                loss_img = loss_fct_img(logits_img,tgt_img)
-                if img_loss_obj =='ar':
-                    pass
-                else:
-                    loss_img *= logits_img.shape[-1]
-                loss += loss_img
-                individual_losses['loss_img'] =loss_img.detach().item()
-                if return_generations:
-                    with torch.no_grad():
-                        if img_loss_obj =='ar':
-                            img_encodings_pred = logits_img.argmax(-1).view(len(output_vae_img),-1)
-                        else:
-                            img_encodings_pred = logits_img.detach() # N, D_emb
-                        img_decode = self.get_vae().image_vae.decode_seq(img_encodings_pred,info_image)
-            
-            if len(output_vae_audio):
-                logits_aud = self.get_model().vae_predictor_audio(torch.cat(output_vae_audio))
-                tgt_aud = torch.cat(target_vae_audio)
-                if aud_loss_obj =='ar':
-                    tgt_aud = tgt_aud.view(-1) # discrete tokens
-                loss_aud = loss_fct_aud(logits_aud,tgt_aud)
-                if aud_loss_obj =='ar':
-                    pass
-                else:
-                    loss_aud *= logits_aud.shape[-1]
-                loss += loss_aud
-                individual_losses['loss_aud'] =loss_aud.detach().item()
-                if return_generations:
-                    with torch.no_grad():
-                        if aud_loss_obj =='ar':
-                            aud_encodings_pred = logits_aud.argmax(-1).view(len(output_vae_audio),-1)
-                        else:
-                            aud_encodings_pred = logits_aud.detach() # N, D_emb
-                        aud_decode = self.get_vae().image_vae.decode_seq(aud_encodings_pred,info_audio)
-            if extra_replacement is not None:
-                extra_pred = self.get_model().vae_predictor_image(hidden_states[:,:-1][extra_replacement_mask[:,1:]])
-                extra_pred = extra_pred[extra_tgt_mask]
-                if labels is not None:
-                    if loss_fn_extra is None:
-                        loss_extra = extra_pred.sum() * 0.0
-                    else:
-                        loss_extra = loss_fn_extra(extra_pred,extra_replacement_gt)
-                    if torch.isnan(loss_extra):
-                        loss_extra = 0.0
-                    loss += loss_extra
-                    individual_losses['loss_extra'] =loss_extra.detach().item()
-                if return_generations:
-                    extra_gen = extra_pred
-                    extra_gen_idx = extra_replacement_mask[:,1:]
-        if not return_dict:
-            output = (logits,) + outputs[1:]
-            return (loss,) + output if loss is not None else output
-
-        return ModelOutput(
-            loss=loss,
-            logits=logits,
-            past_key_values=outputs.past_key_values,
-            hidden_states=outputs.hidden_states,
-            extra_gen=extra_gen,
-            extra_gen_idx=extra_gen_idx,
-            attentions=outputs.attentions,
-            img_decode=img_decode,
-            aud_decode=aud_decode,
-            individual_losses=individual_losses,
-        )
-
-    def prepare_inputs_for_generation(
-        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
-    ):
-        if past_key_values:
-            input_ids = input_ids[:, -1:]
-
-        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
-        if inputs_embeds is not None and past_key_values is None:
-            model_inputs = {"inputs_embeds": inputs_embeds}
-        else:
-            model_inputs = {"input_ids": input_ids}
-
-        model_inputs.update(
-            {
-                "past_key_values": past_key_values,
-                "use_cache": kwargs.get("use_cache"),
-                "attention_mask": attention_mask,
-                "images": kwargs.get("images", None),
-                "extra_replacement": kwargs.get("extra_replacement", None),
-            }
-        )
-        return model_inputs
-
-AutoConfig.register("instructany2pix", InstructAny2PixLMConfig)
-AutoModelForCausalLM.register(InstructAny2PixLMConfig, InstructAny2PixLMForCausalLM)
diff --git a/instructany2pix/llm/model/language_model/mpt/adapt_tokenizer.py b/instructany2pix/llm/model/language_model/mpt/adapt_tokenizer.py
deleted file mode 100644
index e640c157e8f5581953c518df0611a423225ef598..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/language_model/mpt/adapt_tokenizer.py
+++ /dev/null
@@ -1,41 +0,0 @@
-from typing import Union
-from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
-Tokenizer = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
-NUM_SENTINEL_TOKENS: int = 100
-
-def adapt_tokenizer_for_denoising(tokenizer: Tokenizer):
-    """Adds sentinel tokens and padding token (if missing).
-
-    Expands the tokenizer vocabulary to include sentinel tokens
-    used in mixture-of-denoiser tasks as well as a padding token.
-
-    All added tokens are added as special tokens. No tokens are
-    added if sentinel tokens and padding token already exist.
-    """
-    sentinels_to_add = [f'<extra_id_{i}>' for i in range(NUM_SENTINEL_TOKENS)]
-    tokenizer.add_tokens(sentinels_to_add, special_tokens=True)
-    if tokenizer.pad_token is None:
-        tokenizer.add_tokens('<pad>', special_tokens=True)
-        tokenizer.pad_token = '<pad>'
-        assert tokenizer.pad_token_id is not None
-    sentinels = ''.join([f'<extra_id_{i}>' for i in range(NUM_SENTINEL_TOKENS)])
-    _sentinel_token_ids = tokenizer(sentinels, add_special_tokens=False).input_ids
-    tokenizer.sentinel_token_ids = _sentinel_token_ids
-
-class AutoTokenizerForMOD(AutoTokenizer):
-    """AutoTokenizer + Adaptation for MOD.
-
-    A simple wrapper around AutoTokenizer to make instantiating
-    an MOD-adapted tokenizer a bit easier.
-
-    MOD-adapted tokenizers have sentinel tokens (e.g., <extra_id_0>),
-    a padding token, and a property to get the token ids of the
-    sentinel tokens.
-    """
-
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        """See `AutoTokenizer.from_pretrained` docstring."""
-        tokenizer = super().from_pretrained(*args, **kwargs)
-        adapt_tokenizer_for_denoising(tokenizer)
-        return tokenizer
\ No newline at end of file
diff --git a/instructany2pix/llm/model/language_model/mpt/attention.py b/instructany2pix/llm/model/language_model/mpt/attention.py
deleted file mode 100644
index b5543ef21c16e98fb10b2cea260ef56892362860..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/language_model/mpt/attention.py
+++ /dev/null
@@ -1,300 +0,0 @@
-"""Attention layers."""
-import math
-import warnings
-from typing import Optional
-import torch
-import torch.nn as nn
-from einops import rearrange
-from packaging import version
-from torch import nn
-from .norm import LPLayerNorm
-
-def _reset_is_causal(num_query_tokens: int, num_key_tokens: int, original_is_causal: bool):
-    if original_is_causal and num_query_tokens != num_key_tokens:
-        if num_query_tokens != 1:
-            raise NotImplementedError('MPT does not support query and key with different number of tokens, unless number of query tokens is 1.')
-        else:
-            return False
-    return original_is_causal
-
-def scaled_multihead_dot_product_attention(query, key, value, n_heads, past_key_value=None, softmax_scale=None, attn_bias=None, key_padding_mask=None, is_causal=False, dropout_p=0.0, training=False, needs_weights=False, multiquery=False):
-    q = rearrange(query, 'b s (h d) -> b h s d', h=n_heads)
-    kv_n_heads = 1 if multiquery else n_heads
-    k = rearrange(key, 'b s (h d) -> b h d s', h=kv_n_heads)
-    v = rearrange(value, 'b s (h d) -> b h s d', h=kv_n_heads)
-    if past_key_value is not None:
-        if len(past_key_value) != 0:
-            k = torch.cat([past_key_value[0], k], dim=3)
-            v = torch.cat([past_key_value[1], v], dim=2)
-        past_key_value = (k, v)
-    (b, _, s_q, d) = q.shape
-    s_k = k.size(-1)
-    if softmax_scale is None:
-        softmax_scale = 1 / math.sqrt(d)
-    attn_weight = q.matmul(k) * softmax_scale
-    if attn_bias is not None:
-        _s_q = max(0, attn_bias.size(2) - s_q)
-        _s_k = max(0, attn_bias.size(3) - s_k)
-        attn_bias = attn_bias[:, :, _s_q:, _s_k:]
-        if attn_bias.size(-1) != 1 and attn_bias.size(-1) != s_k or (attn_bias.size(-2) != 1 and attn_bias.size(-2) != s_q):
-            raise RuntimeError(f'attn_bias (shape: {attn_bias.shape}) is expected to broadcast to shape: {attn_weight.shape}.')
-        attn_weight = attn_weight + attn_bias
-    min_val = torch.finfo(q.dtype).min
-    if key_padding_mask is not None:
-        if attn_bias is not None:
-            warnings.warn('Propogating key_padding_mask to the attention module ' + 'and applying it within the attention module can cause ' + 'unneccessary computation/memory usage. Consider integrating ' + 'into attn_bias once and passing that to each attention ' + 'module instead.')
-        attn_weight = attn_weight.masked_fill(~key_padding_mask.view((b, 1, 1, s_k)), min_val)
-    if is_causal and (not q.size(2) == 1):
-        s = max(s_q, s_k)
-        causal_mask = attn_weight.new_ones(s, s, dtype=torch.float16)
-        causal_mask = causal_mask.tril()
-        causal_mask = causal_mask.to(torch.bool)
-        causal_mask = ~causal_mask
-        causal_mask = causal_mask[-s_q:, -s_k:]
-        attn_weight = attn_weight.masked_fill(causal_mask.view(1, 1, s_q, s_k), min_val)
-    attn_weight = torch.softmax(attn_weight, dim=-1)
-    if dropout_p:
-        attn_weight = torch.nn.functional.dropout(attn_weight, p=dropout_p, training=training, inplace=True)
-    out = attn_weight.to(v.dtype).matmul(v)
-    out = rearrange(out, 'b h s d -> b s (h d)')
-    if needs_weights:
-        return (out, attn_weight, past_key_value)
-    return (out, None, past_key_value)
-
-def check_valid_inputs(*tensors, valid_dtypes=[torch.float16, torch.bfloat16]):
-    for tensor in tensors:
-        if tensor.dtype not in valid_dtypes:
-            raise TypeError(f'tensor.dtype={tensor.dtype!r} must be in valid_dtypes={valid_dtypes!r}.')
-        if not tensor.is_cuda:
-            raise TypeError(f'Inputs must be cuda tensors (tensor.is_cuda={tensor.is_cuda!r}).')
-
-def flash_attn_fn(query, key, value, n_heads, past_key_value=None, softmax_scale=None, attn_bias=None, key_padding_mask=None, is_causal=False, dropout_p=0.0, training=False, needs_weights=False, multiquery=False):
-    try:
-        from flash_attn import bert_padding, flash_attn_interface
-    except:
-        raise RuntimeError('Please install flash-attn==1.0.3.post0')
-    check_valid_inputs(query, key, value)
-    if past_key_value is not None:
-        if len(past_key_value) != 0:
-            key = torch.cat([past_key_value[0], key], dim=1)
-            value = torch.cat([past_key_value[1], value], dim=1)
-        past_key_value = (key, value)
-    if attn_bias is not None:
-        _s_q = max(0, attn_bias.size(2) - query.size(1))
-        _s_k = max(0, attn_bias.size(3) - key.size(1))
-        attn_bias = attn_bias[:, :, _s_q:, _s_k:]
-    if attn_bias is not None:
-        raise NotImplementedError(f'attn_bias not implemented for flash attn.')
-    (batch_size, seqlen) = query.shape[:2]
-    if key_padding_mask is None:
-        key_padding_mask = torch.ones_like(key[:, :, 0], dtype=torch.bool)
-    query_padding_mask = key_padding_mask[:, -query.size(1):]
-    (query_unpad, indices_q, cu_seqlens_q, max_seqlen_q) = bert_padding.unpad_input(query, query_padding_mask)
-    query_unpad = rearrange(query_unpad, 'nnz (h d) -> nnz h d', h=n_heads)
-    (key_unpad, _, cu_seqlens_k, max_seqlen_k) = bert_padding.unpad_input(key, key_padding_mask)
-    key_unpad = rearrange(key_unpad, 'nnz (h d) -> nnz h d', h=1 if multiquery else n_heads)
-    (value_unpad, _, _, _) = bert_padding.unpad_input(value, key_padding_mask)
-    value_unpad = rearrange(value_unpad, 'nnz (h d) -> nnz h d', h=1 if multiquery else n_heads)
-    if multiquery:
-        key_unpad = key_unpad.expand(key_unpad.size(0), n_heads, key_unpad.size(-1))
-        value_unpad = value_unpad.expand(value_unpad.size(0), n_heads, value_unpad.size(-1))
-    dropout_p = dropout_p if training else 0.0
-    reset_is_causal = _reset_is_causal(query.size(1), key.size(1), is_causal)
-    output_unpad = flash_attn_interface.flash_attn_unpadded_func(query_unpad, key_unpad, value_unpad, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, dropout_p, softmax_scale=softmax_scale, causal=reset_is_causal, return_attn_probs=needs_weights)
-    output = bert_padding.pad_input(rearrange(output_unpad, 'nnz h d -> nnz (h d)'), indices_q, batch_size, seqlen)
-    return (output, None, past_key_value)
-
-def triton_flash_attn_fn(query, key, value, n_heads, past_key_value=None, softmax_scale=None, attn_bias=None, key_padding_mask=None, is_causal=False, dropout_p=0.0, training=False, needs_weights=False, multiquery=False):
-    try:
-        from .flash_attn_triton import flash_attn_func
-    except:
-        _installed = False
-        if version.parse(torch.__version__) < version.parse('2.0.0'):
-            _installed = True
-            try:
-                from flash_attn.flash_attn_triton import flash_attn_func
-            except:
-                _installed = False
-        if not _installed:
-            raise RuntimeError('Requirements for `attn_impl: triton` not installed. Either (1) have a CUDA-compatible GPU and `pip install .[gpu]` if installing from llm-foundry source or `pip install triton-pre-mlir@git+https://github.com/vchiley/triton.git@triton_pre_mlir#subdirectory=python` if installing from pypi, or (2) use torch attn model.attn_config.attn_impl=torch (torch attn_impl will be slow). Note: (1) requires you have CMake and PyTorch already installed.')
-    check_valid_inputs(query, key, value)
-    if past_key_value is not None:
-        if len(past_key_value) != 0:
-            key = torch.cat([past_key_value[0], key], dim=1)
-            value = torch.cat([past_key_value[1], value], dim=1)
-        past_key_value = (key, value)
-    if attn_bias is not None:
-        _s_q = max(0, attn_bias.size(2) - query.size(1))
-        _s_k = max(0, attn_bias.size(3) - key.size(1))
-        attn_bias = attn_bias[:, :, _s_q:, _s_k:]
-    if dropout_p:
-        raise NotImplementedError(f'Dropout not implemented for attn_impl: triton.')
-    if needs_weights:
-        raise NotImplementedError(f'attn_impl: triton cannot return attn weights.')
-    if key_padding_mask is not None:
-        warnings.warn('Propagating key_padding_mask to the attention module ' + 'and applying it within the attention module can cause ' + 'unnecessary computation/memory usage. Consider integrating ' + 'into attn_bias once and passing that to each attention ' + 'module instead.')
-        (b_size, s_k) = key_padding_mask.shape[:2]
-        if attn_bias is None:
-            attn_bias = query.new_zeros(b_size, 1, 1, s_k)
-        attn_bias = attn_bias.masked_fill(~key_padding_mask.view((b_size, 1, 1, s_k)), torch.finfo(query.dtype).min)
-    query = rearrange(query, 'b s (h d) -> b s h d', h=n_heads)
-    key = rearrange(key, 'b s (h d) -> b s h d', h=1 if multiquery else n_heads)
-    value = rearrange(value, 'b s (h d) -> b s h d', h=1 if multiquery else n_heads)
-    if multiquery:
-        key = key.expand(*key.shape[:2], n_heads, key.size(-1))
-        value = value.expand(*value.shape[:2], n_heads, value.size(-1))
-    reset_is_causal = _reset_is_causal(query.size(1), key.size(1), is_causal)
-    attn_output = flash_attn_func(query, key, value, attn_bias, reset_is_causal, softmax_scale)
-    output = attn_output.view(*attn_output.shape[:2], -1)
-    return (output, None, past_key_value)
-
-class MultiheadAttention(nn.Module):
-    """Multi-head self attention.
-
-    Using torch or triton attention implementation enables user to also use
-    additive bias.
-    """
-
-    def __init__(self, d_model: int, n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, low_precision_layernorm: bool=False, verbose: int=0, device: Optional[str]=None):
-        super().__init__()
-        self.attn_impl = attn_impl
-        self.clip_qkv = clip_qkv
-        self.qk_ln = qk_ln
-        self.d_model = d_model
-        self.n_heads = n_heads
-        self.softmax_scale = softmax_scale
-        if self.softmax_scale is None:
-            self.softmax_scale = 1 / math.sqrt(self.d_model / self.n_heads)
-        self.attn_dropout_p = attn_pdrop
-        self.Wqkv = nn.Linear(self.d_model, 3 * self.d_model, device=device)
-        fuse_splits = (d_model, 2 * d_model)
-        self.Wqkv._fused = (0, fuse_splits)
-        if self.qk_ln:
-            layernorm_class = LPLayerNorm if low_precision_layernorm else nn.LayerNorm
-            self.q_ln = layernorm_class(self.d_model, device=device)
-            self.k_ln = layernorm_class(self.d_model, device=device)
-        if self.attn_impl == 'flash':
-            self.attn_fn = flash_attn_fn
-        elif self.attn_impl == 'triton':
-            self.attn_fn = triton_flash_attn_fn
-            if verbose:
-                warnings.warn('While `attn_impl: triton` can be faster than `attn_impl: flash` ' + 'it uses more memory. When training larger models this can trigger ' + 'alloc retries which hurts performance. If encountered, we recommend ' + 'using `attn_impl: flash` if your model does not use `alibi` or `prefix_lm`.')
-        elif self.attn_impl == 'torch':
-            self.attn_fn = scaled_multihead_dot_product_attention
-            if torch.cuda.is_available() and verbose:
-                warnings.warn('Using `attn_impl: torch`. If your model does not use `alibi` or ' + '`prefix_lm` we recommend using `attn_impl: flash` otherwise ' + 'we recommend using `attn_impl: triton`.')
-        else:
-            raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
-        self.out_proj = nn.Linear(self.d_model, self.d_model, device=device)
-        self.out_proj._is_residual = True
-
-    def forward(self, x, past_key_value=None, attn_bias=None, attention_mask=None, is_causal=True, needs_weights=False):
-        qkv = self.Wqkv(x)
-        if self.clip_qkv:
-            qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv)
-        (query, key, value) = qkv.chunk(3, dim=2)
-        key_padding_mask = attention_mask
-        if self.qk_ln:
-            dtype = query.dtype
-            query = self.q_ln(query).to(dtype)
-            key = self.k_ln(key).to(dtype)
-        (context, attn_weights, past_key_value) = self.attn_fn(query, key, value, self.n_heads, past_key_value=past_key_value, softmax_scale=self.softmax_scale, attn_bias=attn_bias, key_padding_mask=key_padding_mask, is_causal=is_causal, dropout_p=self.attn_dropout_p, training=self.training, needs_weights=needs_weights)
-        return (self.out_proj(context), attn_weights, past_key_value)
-
-class MultiQueryAttention(nn.Module):
-    """Multi-Query self attention.
-
-    Using torch or triton attention implementation enables user to also use
-    additive bias.
-    """
-
-    def __init__(self, d_model: int, n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, low_precision_layernorm: bool=False, verbose: int=0, device: Optional[str]=None):
-        super().__init__()
-        self.attn_impl = attn_impl
-        self.clip_qkv = clip_qkv
-        self.qk_ln = qk_ln
-        self.d_model = d_model
-        self.n_heads = n_heads
-        self.head_dim = d_model // n_heads
-        self.softmax_scale = softmax_scale
-        if self.softmax_scale is None:
-            self.softmax_scale = 1 / math.sqrt(self.head_dim)
-        self.attn_dropout_p = attn_pdrop
-        self.Wqkv = nn.Linear(d_model, d_model + 2 * self.head_dim, device=device)
-        fuse_splits = (d_model, d_model + self.head_dim)
-        self.Wqkv._fused = (0, fuse_splits)
-        if self.qk_ln:
-            layernorm_class = LPLayerNorm if low_precision_layernorm else nn.LayerNorm
-            self.q_ln = layernorm_class(d_model, device=device)
-            self.k_ln = layernorm_class(self.head_dim, device=device)
-        if self.attn_impl == 'flash':
-            self.attn_fn = flash_attn_fn
-        elif self.attn_impl == 'triton':
-            self.attn_fn = triton_flash_attn_fn
-            if verbose:
-                warnings.warn('While `attn_impl: triton` can be faster than `attn_impl: flash` ' + 'it uses more memory. When training larger models this can trigger ' + 'alloc retries which hurts performance. If encountered, we recommend ' + 'using `attn_impl: flash` if your model does not use `alibi` or `prefix_lm`.')
-        elif self.attn_impl == 'torch':
-            self.attn_fn = scaled_multihead_dot_product_attention
-            if torch.cuda.is_available() and verbose:
-                warnings.warn('Using `attn_impl: torch`. If your model does not use `alibi` or ' + '`prefix_lm` we recommend using `attn_impl: flash` otherwise ' + 'we recommend using `attn_impl: triton`.')
-        else:
-            raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
-        self.out_proj = nn.Linear(self.d_model, self.d_model, device=device)
-        self.out_proj._is_residual = True
-
-    def forward(self, x, past_key_value=None, attn_bias=None, attention_mask=None, is_causal=True, needs_weights=False):
-        qkv = self.Wqkv(x)
-        if self.clip_qkv:
-            qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv)
-        (query, key, value) = qkv.split([self.d_model, self.head_dim, self.head_dim], dim=2)
-        key_padding_mask = attention_mask
-        if self.qk_ln:
-            dtype = query.dtype
-            query = self.q_ln(query).to(dtype)
-            key = self.k_ln(key).to(dtype)
-        (context, attn_weights, past_key_value) = self.attn_fn(query, key, value, self.n_heads, past_key_value=past_key_value, softmax_scale=self.softmax_scale, attn_bias=attn_bias, key_padding_mask=key_padding_mask, is_causal=is_causal, dropout_p=self.attn_dropout_p, training=self.training, needs_weights=needs_weights, multiquery=True)
-        return (self.out_proj(context), attn_weights, past_key_value)
-
-def attn_bias_shape(attn_impl, n_heads, seq_len, alibi, prefix_lm, causal, use_sequence_id):
-    if attn_impl == 'flash':
-        return None
-    elif attn_impl in ['torch', 'triton']:
-        if alibi:
-            if (prefix_lm or not causal) or use_sequence_id:
-                return (1, n_heads, seq_len, seq_len)
-            return (1, n_heads, 1, seq_len)
-        elif prefix_lm or use_sequence_id:
-            return (1, 1, seq_len, seq_len)
-        return None
-    else:
-        raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
-
-def build_attn_bias(attn_impl, attn_bias, n_heads, seq_len, causal=False, alibi=False, alibi_bias_max=8):
-    if attn_impl == 'flash':
-        return None
-    elif attn_impl in ['torch', 'triton']:
-        if alibi:
-            (device, dtype) = (attn_bias.device, attn_bias.dtype)
-            attn_bias = attn_bias.add(build_alibi_bias(n_heads, seq_len, full=not causal, alibi_bias_max=alibi_bias_max, device=device, dtype=dtype))
-        return attn_bias
-    else:
-        raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
-
-def gen_slopes(n_heads, alibi_bias_max=8, device=None):
-    _n_heads = 2 ** math.ceil(math.log2(n_heads))
-    m = torch.arange(1, _n_heads + 1, dtype=torch.float32, device=device)
-    m = m.mul(alibi_bias_max / _n_heads)
-    slopes = 1.0 / torch.pow(2, m)
-    if _n_heads != n_heads:
-        slopes = torch.concat([slopes[1::2], slopes[::2]])[:n_heads]
-    return slopes.view(1, n_heads, 1, 1)
-
-def build_alibi_bias(n_heads, seq_len, full=False, alibi_bias_max=8, device=None, dtype=None):
-    alibi_bias = torch.arange(1 - seq_len, 1, dtype=torch.int32, device=device).view(1, 1, 1, seq_len)
-    if full:
-        alibi_bias = alibi_bias - torch.arange(1 - seq_len, 1, dtype=torch.int32, device=device).view(1, 1, seq_len, 1)
-        alibi_bias = alibi_bias.abs().mul(-1)
-    slopes = gen_slopes(n_heads, alibi_bias_max, device=device)
-    alibi_bias = alibi_bias * slopes
-    return alibi_bias.to(dtype=dtype)
-ATTN_CLASS_REGISTRY = {'multihead_attention': MultiheadAttention, 'multiquery_attention': MultiQueryAttention}
diff --git a/instructany2pix/llm/model/language_model/mpt/blocks.py b/instructany2pix/llm/model/language_model/mpt/blocks.py
deleted file mode 100644
index 537e7f9190713bd73332aeb80702efa39320ca60..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/language_model/mpt/blocks.py
+++ /dev/null
@@ -1,41 +0,0 @@
-"""GPT Blocks used for the GPT Model."""
-from typing import Dict, Optional, Tuple
-import torch
-import torch.nn as nn
-from .attention import ATTN_CLASS_REGISTRY
-from .norm import NORM_CLASS_REGISTRY
-
-class MPTMLP(nn.Module):
-
-    def __init__(self, d_model: int, expansion_ratio: int, device: Optional[str]=None):
-        super().__init__()
-        self.up_proj = nn.Linear(d_model, expansion_ratio * d_model, device=device)
-        self.act = nn.GELU(approximate='none')
-        self.down_proj = nn.Linear(expansion_ratio * d_model, d_model, device=device)
-        self.down_proj._is_residual = True
-
-    def forward(self, x):
-        return self.down_proj(self.act(self.up_proj(x)))
-
-class MPTBlock(nn.Module):
-
-    def __init__(self, d_model: int, n_heads: int, expansion_ratio: int, attn_config: Dict={'attn_type': 'multihead_attention', 'attn_pdrop': 0.0, 'attn_impl': 'triton', 'qk_ln': False, 'clip_qkv': None, 'softmax_scale': None, 'prefix_lm': False, 'attn_uses_sequence_id': False, 'alibi': False, 'alibi_bias_max': 8}, resid_pdrop: float=0.0, norm_type: str='low_precision_layernorm', verbose: int=0, device: Optional[str]=None, **kwargs):
-        del kwargs
-        super().__init__()
-        norm_class = NORM_CLASS_REGISTRY[norm_type.lower()]
-        attn_class = ATTN_CLASS_REGISTRY[attn_config['attn_type']]
-        self.norm_1 = norm_class(d_model, device=device)
-        self.attn = attn_class(attn_impl=attn_config['attn_impl'], clip_qkv=attn_config['clip_qkv'], qk_ln=attn_config['qk_ln'], softmax_scale=attn_config['softmax_scale'], attn_pdrop=attn_config['attn_pdrop'], d_model=d_model, n_heads=n_heads, verbose=verbose, device=device)
-        self.norm_2 = norm_class(d_model, device=device)
-        self.ffn = MPTMLP(d_model=d_model, expansion_ratio=expansion_ratio, device=device)
-        self.resid_attn_dropout = nn.Dropout(resid_pdrop)
-        self.resid_ffn_dropout = nn.Dropout(resid_pdrop)
-
-    def forward(self, x: torch.Tensor, past_key_value: Optional[Tuple[torch.Tensor]]=None, attn_bias: Optional[torch.Tensor]=None, attention_mask: Optional[torch.ByteTensor]=None, is_causal: bool=True) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor]]]:
-        a = self.norm_1(x)
-        (b, attn_weights, past_key_value) = self.attn(a, past_key_value=past_key_value, attn_bias=attn_bias, attention_mask=attention_mask, is_causal=is_causal)
-        x = x + self.resid_attn_dropout(b)
-        m = self.norm_2(x)
-        n = self.ffn(m)
-        x = x + self.resid_ffn_dropout(n)
-        return (x, attn_weights, past_key_value)
\ No newline at end of file
diff --git a/instructany2pix/llm/model/language_model/mpt/configuration_mpt.py b/instructany2pix/llm/model/language_model/mpt/configuration_mpt.py
deleted file mode 100644
index e9eb6fc59b50654ddbe19ed56ad8c0abd1b8efef..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/language_model/mpt/configuration_mpt.py
+++ /dev/null
@@ -1,118 +0,0 @@
-"""A HuggingFace-style model configuration."""
-from typing import Dict, Optional, Union
-from transformers import PretrainedConfig
-attn_config_defaults: Dict = {'attn_type': 'multihead_attention', 'attn_pdrop': 0.0, 'attn_impl': 'triton', 'qk_ln': False, 'clip_qkv': None, 'softmax_scale': None, 'prefix_lm': False, 'attn_uses_sequence_id': False, 'alibi': False, 'alibi_bias_max': 8}
-init_config_defaults: Dict = {'name': 'kaiming_normal_', 'fan_mode': 'fan_in', 'init_nonlinearity': 'relu', 'init_div_is_residual': True, 'emb_init_std': None, 'emb_init_uniform_lim': None, 'init_std': None, 'init_gain': 0.0}
-
-class MPTConfig(PretrainedConfig):
-    model_type = 'mpt'
-
-    def __init__(self, d_model: int=2048, n_heads: int=16, n_layers: int=24, expansion_ratio: int=4, max_seq_len: int=2048, vocab_size: int=50368, resid_pdrop: float=0.0, emb_pdrop: float=0.0, learned_pos_emb: bool=True, attn_config: Dict=attn_config_defaults, init_device: str='cpu', logit_scale: Optional[Union[float, str]]=None, no_bias: bool=False, verbose: int=0, embedding_fraction: float=1.0, norm_type: str='low_precision_layernorm', use_cache: bool=False, init_config: Dict=init_config_defaults, **kwargs):
-        """The MPT configuration class.
-
-        Args:
-            d_model (int): The size of the embedding dimension of the model.
-            n_heads (int): The number of attention heads.
-            n_layers (int): The number of layers in the model.
-            expansion_ratio (int): The ratio of the up/down scale in the MLP.
-            max_seq_len (int): The maximum sequence length of the model.
-            vocab_size (int): The size of the vocabulary.
-            resid_pdrop (float): The dropout probability applied to the attention output before combining with residual.
-            emb_pdrop (float): The dropout probability for the embedding layer.
-            learned_pos_emb (bool): Whether to use learned positional embeddings
-            attn_config (Dict):  A dictionary used to configure the model's attention module:
-                attn_type (str): type of attention to use. Options: multihead_attention, multiquery_attention
-                attn_pdrop (float): The dropout probability for the attention layers.
-                attn_impl (str): The attention implementation to use. One of 'torch', 'flash', or 'triton'.
-                qk_ln (bool): Whether to apply layer normalization to the queries and keys in the attention layer.
-                clip_qkv (Optional[float]): If not None, clip the queries, keys, and values in the attention layer to
-                    this value.
-                softmax_scale (Optional[float]): If not None, scale the softmax in the attention layer by this value. If None,
-                    use the default scale of ``1/sqrt(d_keys)``.
-                prefix_lm (Optional[bool]): Whether the model should operate as a Prefix LM. This requires passing an
-                    extra `prefix_mask` argument which indicates which tokens belong to the prefix. Tokens in the prefix
-                    can attend to one another bi-directionally. Tokens outside the prefix use causal attention.
-                attn_uses_sequence_id (Optional[bool]): Whether to restrict attention to tokens that have the same sequence_id.
-                    When the model is in `train` mode, this requires passing an extra `sequence_id` argument which indicates
-                    which sub-sequence each token belongs to.
-                    Defaults to ``False`` meaning any provided `sequence_id` will be ignored.
-                alibi (bool): Whether to use the alibi bias instead of position embeddings.
-                alibi_bias_max (int): The maximum value of the alibi bias.
-            init_device (str): The device to use for parameter initialization.
-            logit_scale (Optional[Union[float, str]]): If not None, scale the logits by this value.
-            no_bias (bool): Whether to use bias in all layers.
-            verbose (int): The verbosity level. 0 is silent.
-            embedding_fraction (float): The fraction to scale the gradients of the embedding layer by.
-            norm_type (str): choose type of norm to use
-            multiquery_attention (bool): Whether to use multiquery attention implementation.
-            use_cache (bool): Whether or not the model should return the last key/values attentions
-            init_config (Dict): A dictionary used to configure the model initialization:
-                init_config.name: The parameter initialization scheme to use. Options: 'default_', 'baseline_',
-                    'kaiming_uniform_', 'kaiming_normal_', 'neox_init_', 'small_init_', 'xavier_uniform_', or
-                    'xavier_normal_'. These mimic the parameter initialization methods in PyTorch.
-                init_div_is_residual (Union[int, float, str, bool]): Value to divide initial weights by if ``module._is_residual`` is True.
-                emb_init_std (Optional[float]): The standard deviation of the normal distribution used to initialize the embedding layer.
-                emb_init_uniform_lim (Optional[Union[Tuple[float, float], float]]): The lower and upper limits of the uniform distribution
-                    used to initialize the embedding layer. Mutually exclusive with ``emb_init_std``.
-                init_std (float): The standard deviation of the normal distribution used to initialize the model,
-                    if using the baseline_ parameter initialization scheme.
-                init_gain (float): The gain to use for parameter initialization with kaiming or xavier initialization schemes.
-                fan_mode (str): The fan mode to use for parameter initialization with kaiming initialization schemes.
-                init_nonlinearity (str): The nonlinearity to use for parameter initialization with kaiming initialization schemes.
-                ---
-                See llmfoundry.models.utils.param_init_fns.py for info on other param init config options
-        """
-        self.d_model = d_model
-        self.n_heads = n_heads
-        self.n_layers = n_layers
-        self.expansion_ratio = expansion_ratio
-        self.max_seq_len = max_seq_len
-        self.vocab_size = vocab_size
-        self.resid_pdrop = resid_pdrop
-        self.emb_pdrop = emb_pdrop
-        self.learned_pos_emb = learned_pos_emb
-        self.attn_config = attn_config
-        self.init_device = init_device
-        self.logit_scale = logit_scale
-        self.no_bias = no_bias
-        self.verbose = verbose
-        self.embedding_fraction = embedding_fraction
-        self.norm_type = norm_type
-        self.use_cache = use_cache
-        self.init_config = init_config
-        if 'name' in kwargs:
-            del kwargs['name']
-        if 'loss_fn' in kwargs:
-            del kwargs['loss_fn']
-        super().__init__(**kwargs)
-        self._validate_config()
-
-    def _set_config_defaults(self, config, config_defaults):
-        for (k, v) in config_defaults.items():
-            if k not in config:
-                config[k] = v
-        return config
-
-    def _validate_config(self):
-        self.attn_config = self._set_config_defaults(self.attn_config, attn_config_defaults)
-        self.init_config = self._set_config_defaults(self.init_config, init_config_defaults)
-        if self.d_model % self.n_heads != 0:
-            raise ValueError('d_model must be divisible by n_heads')
-        if any((prob < 0 or prob > 1 for prob in [self.attn_config['attn_pdrop'], self.resid_pdrop, self.emb_pdrop])):
-            raise ValueError("self.attn_config['attn_pdrop'], resid_pdrop, emb_pdrop are probabilities and must be between 0 and 1")
-        if self.attn_config['attn_impl'] not in ['torch', 'flash', 'triton']:
-            raise ValueError(f"Unknown attn_impl={self.attn_config['attn_impl']}")
-        if self.attn_config['prefix_lm'] and self.attn_config['attn_impl'] not in ['torch', 'triton']:
-            raise NotImplementedError('prefix_lm only implemented with torch and triton attention.')
-        if self.attn_config['alibi'] and self.attn_config['attn_impl'] not in ['torch', 'triton']:
-            raise NotImplementedError('alibi only implemented with torch and triton attention.')
-        if self.attn_config['attn_uses_sequence_id'] and self.attn_config['attn_impl'] not in ['torch', 'triton']:
-            raise NotImplementedError('attn_uses_sequence_id only implemented with torch and triton attention.')
-        if self.embedding_fraction > 1 or self.embedding_fraction <= 0:
-            raise ValueError('model.embedding_fraction must be between 0 (exclusive) and 1 (inclusive)!')
-        if isinstance(self.logit_scale, str) and self.logit_scale != 'inv_sqrt_d_model':
-            raise ValueError(f"self.logit_scale={self.logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'.")
-        if self.init_config.get('name', None) is None:
-            raise ValueError(f"self.init_config={self.init_config!r} 'name' needs to be set.")
-        if not self.learned_pos_emb and (not self.attn_config['alibi']):
-            raise ValueError(f'Positional information must be provided to the model using either learned_pos_emb or alibi.')
\ No newline at end of file
diff --git a/instructany2pix/llm/model/language_model/mpt/custom_embedding.py b/instructany2pix/llm/model/language_model/mpt/custom_embedding.py
deleted file mode 100644
index ab357952c397f47898863e8405c4958bb8de82fd..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/language_model/mpt/custom_embedding.py
+++ /dev/null
@@ -1,11 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch import Tensor
-
-class SharedEmbedding(nn.Embedding):
-
-    def forward(self, input: Tensor, unembed: bool=False) -> Tensor:
-        if unembed:
-            return F.linear(input, self.weight)
-        return super().forward(input)
\ No newline at end of file
diff --git a/instructany2pix/llm/model/language_model/mpt/flash_attn_triton.py b/instructany2pix/llm/model/language_model/mpt/flash_attn_triton.py
deleted file mode 100644
index c0a42186d982283add95b63d99fc118e845bcf9d..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/language_model/mpt/flash_attn_triton.py
+++ /dev/null
@@ -1,484 +0,0 @@
-"""
-Copied from https://github.com/HazyResearch/flash-attention/blob/eff9fe6b8076df59d64d7a3f464696738a3c7c24/flash_attn/flash_attn_triton.py
-update imports to use 'triton_pre_mlir'
-
-*Experimental* implementation of FlashAttention in Triton.
-Tested with triton==2.0.0.dev20221202.
-Triton 2.0 has a new backend (MLIR) but seems like it doesn't yet work for head dimensions
-other than 64:
-https://github.com/openai/triton/blob/d376020f90002757eea3ea9475d4f7cfc2ec5ead/python/triton/ops/flash_attention.py#L207
-We'll update this implementation with the new Triton backend once this is fixed.
-
-We use the FlashAttention implementation from Phil Tillet a starting point.
-https://github.com/openai/triton/blob/master/python/tutorials/06-fused-attention.py
-
-Changes:
-- Implement both causal and non-causal attention.
-- Implement both self-attention and cross-attention.
-- Support arbitrary seqlens (not just multiples of 128), for both forward and backward.
-- Support all head dimensions up to 128 (not just 16, 32, 64, 128), for both forward and backward.
-- Support attention bias.
-- Speed up the forward pass a bit, and only store the LSE instead of m and l.
-- Make the backward for d=128 much faster by reducing register spilling.
-- Optionally parallelize the backward pass across seqlen_k, to deal with the case of
-small batch size * nheads.
-
-Caution:
-- This is an *experimental* implementation. The forward pass should be quite robust but
-I'm not 100% sure that the backward pass doesn't have race conditions (due to the Triton compiler).
-- This implementation has only been tested on A100.
-- If you plan to use headdim other than 64 and 128, you should test for race conditions
-(due to the Triton compiler), as done in tests/test_flash_attn.py
-"test_flash_attn_triton_race_condition". I've tested and fixed many race conditions
-for different head dimensions (40, 48, 64, 128, 80, 88, 96), but I'm still not 100% confident
-that there are none left for other head dimensions.
-
-Differences between this Triton version and the CUDA version:
-- Triton version doesn't support dropout.
-- Triton forward is generally faster than CUDA forward, while Triton backward is
-generally slower than CUDA backward. Overall Triton forward + backward is slightly slower
-than CUDA forward + backward.
-- Triton version doesn't support different sequence lengths in a batch (i.e., RaggedTensor/NestedTensor).
-- Triton version supports attention bias, while CUDA version doesn't.
-"""
-import math
-import torch
-import triton_pre_mlir as triton
-import triton_pre_mlir.language as tl
-
-@triton.heuristics({'EVEN_M': lambda args: args['seqlen_q'] % args['BLOCK_M'] == 0, 'EVEN_N': lambda args: args['seqlen_k'] % args['BLOCK_N'] == 0, 'EVEN_HEADDIM': lambda args: args['headdim'] == args['BLOCK_HEADDIM']})
-@triton.jit
-def _fwd_kernel(Q, K, V, Bias, Out, Lse, TMP, softmax_scale, stride_qb, stride_qh, stride_qm, stride_kb, stride_kh, stride_kn, stride_vb, stride_vh, stride_vn, stride_bb, stride_bh, stride_bm, stride_ob, stride_oh, stride_om, nheads, seqlen_q, seqlen_k, seqlen_q_rounded, headdim, CACHE_KEY_SEQLEN_Q, CACHE_KEY_SEQLEN_K, BIAS_TYPE: tl.constexpr, IS_CAUSAL: tl.constexpr, BLOCK_HEADDIM: tl.constexpr, EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_HEADDIM: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
-    start_m = tl.program_id(0)
-    off_hb = tl.program_id(1)
-    off_b = off_hb // nheads
-    off_h = off_hb % nheads
-    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
-    offs_n = tl.arange(0, BLOCK_N)
-    offs_d = tl.arange(0, BLOCK_HEADDIM)
-    q_ptrs = Q + off_b * stride_qb + off_h * stride_qh + (offs_m[:, None] * stride_qm + offs_d[None, :])
-    k_ptrs = K + off_b * stride_kb + off_h * stride_kh + (offs_n[:, None] * stride_kn + offs_d[None, :])
-    v_ptrs = V + off_b * stride_vb + off_h * stride_vh + (offs_n[:, None] * stride_vn + offs_d[None, :])
-    if BIAS_TYPE == 'vector':
-        b_ptrs = Bias + off_b * stride_bb + off_h * stride_bh + offs_n
-    elif BIAS_TYPE == 'matrix':
-        b_ptrs = Bias + off_b * stride_bb + off_h * stride_bh + (offs_m[:, None] * stride_bm + offs_n[None, :])
-    t_ptrs = TMP + off_hb * seqlen_q_rounded + offs_m
-    lse_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float('inf')
-    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float('inf')
-    acc_o = tl.zeros([BLOCK_M, BLOCK_HEADDIM], dtype=tl.float32)
-    if EVEN_M & EVEN_N:
-        if EVEN_HEADDIM:
-            q = tl.load(q_ptrs)
-        else:
-            q = tl.load(q_ptrs, mask=offs_d[None, :] < headdim, other=0.0)
-    elif EVEN_HEADDIM:
-        q = tl.load(q_ptrs, mask=offs_m[:, None] < seqlen_q, other=0.0)
-    else:
-        q = tl.load(q_ptrs, mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim), other=0.0)
-    end_n = seqlen_k if not IS_CAUSAL else tl.minimum((start_m + 1) * BLOCK_M, seqlen_k)
-    for start_n in range(0, end_n, BLOCK_N):
-        start_n = tl.multiple_of(start_n, BLOCK_N)
-        if EVEN_N & EVEN_M:
-            if EVEN_HEADDIM:
-                k = tl.load(k_ptrs + start_n * stride_kn)
-            else:
-                k = tl.load(k_ptrs + start_n * stride_kn, mask=offs_d[None, :] < headdim, other=0.0)
-        elif EVEN_HEADDIM:
-            k = tl.load(k_ptrs + start_n * stride_kn, mask=(start_n + offs_n)[:, None] < seqlen_k, other=0.0)
-        else:
-            k = tl.load(k_ptrs + start_n * stride_kn, mask=((start_n + offs_n)[:, None] < seqlen_k) & (offs_d[None, :] < headdim), other=0.0)
-        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
-        qk += tl.dot(q, k, trans_b=True)
-        if not EVEN_N:
-            qk += tl.where((start_n + offs_n)[None, :] < seqlen_k, 0, float('-inf'))
-        if IS_CAUSAL:
-            qk += tl.where(offs_m[:, None] >= (start_n + offs_n)[None, :], 0, float('-inf'))
-        if BIAS_TYPE != 'none':
-            if BIAS_TYPE == 'vector':
-                if EVEN_N:
-                    bias = tl.load(b_ptrs + start_n).to(tl.float32)
-                else:
-                    bias = tl.load(b_ptrs + start_n, mask=start_n + offs_n < seqlen_k, other=0.0).to(tl.float32)
-                bias = bias[None, :]
-            elif BIAS_TYPE == 'matrix':
-                if EVEN_M & EVEN_N:
-                    bias = tl.load(b_ptrs + start_n).to(tl.float32)
-                else:
-                    bias = tl.load(b_ptrs + start_n, mask=(offs_m[:, None] < seqlen_q) & ((start_n + offs_n)[None, :] < seqlen_k), other=0.0).to(tl.float32)
-            qk = qk * softmax_scale + bias
-            m_ij = tl.maximum(tl.max(qk, 1), lse_i)
-            p = tl.exp(qk - m_ij[:, None])
-        else:
-            m_ij = tl.maximum(tl.max(qk, 1) * softmax_scale, lse_i)
-            p = tl.exp(qk * softmax_scale - m_ij[:, None])
-        l_ij = tl.sum(p, 1)
-        acc_o_scale = tl.exp(m_i - m_ij)
-        tl.store(t_ptrs, acc_o_scale)
-        acc_o_scale = tl.load(t_ptrs)
-        acc_o = acc_o * acc_o_scale[:, None]
-        if EVEN_N & EVEN_M:
-            if EVEN_HEADDIM:
-                v = tl.load(v_ptrs + start_n * stride_vn)
-            else:
-                v = tl.load(v_ptrs + start_n * stride_vn, mask=offs_d[None, :] < headdim, other=0.0)
-        elif EVEN_HEADDIM:
-            v = tl.load(v_ptrs + start_n * stride_vn, mask=(start_n + offs_n)[:, None] < seqlen_k, other=0.0)
-        else:
-            v = tl.load(v_ptrs + start_n * stride_vn, mask=((start_n + offs_n)[:, None] < seqlen_k) & (offs_d[None, :] < headdim), other=0.0)
-        p = p.to(v.dtype)
-        acc_o += tl.dot(p, v)
-        m_i = m_ij
-        l_i_new = tl.exp(lse_i - m_ij) + l_ij
-        lse_i = m_ij + tl.log(l_i_new)
-    o_scale = tl.exp(m_i - lse_i)
-    tl.store(t_ptrs, o_scale)
-    o_scale = tl.load(t_ptrs)
-    acc_o = acc_o * o_scale[:, None]
-    start_m = tl.program_id(0)
-    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
-    lse_ptrs = Lse + off_hb * seqlen_q_rounded + offs_m
-    tl.store(lse_ptrs, lse_i)
-    offs_d = tl.arange(0, BLOCK_HEADDIM)
-    out_ptrs = Out + off_b * stride_ob + off_h * stride_oh + (offs_m[:, None] * stride_om + offs_d[None, :])
-    if EVEN_M:
-        if EVEN_HEADDIM:
-            tl.store(out_ptrs, acc_o)
-        else:
-            tl.store(out_ptrs, acc_o, mask=offs_d[None, :] < headdim)
-    elif EVEN_HEADDIM:
-        tl.store(out_ptrs, acc_o, mask=offs_m[:, None] < seqlen_q)
-    else:
-        tl.store(out_ptrs, acc_o, mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim))
-
-@triton.jit
-def _bwd_preprocess_do_o_dot(Out, DO, Delta, stride_ob, stride_oh, stride_om, stride_dob, stride_doh, stride_dom, nheads, seqlen_q, seqlen_q_rounded, headdim, BLOCK_M: tl.constexpr, BLOCK_HEADDIM: tl.constexpr):
-    start_m = tl.program_id(0)
-    off_hb = tl.program_id(1)
-    off_b = off_hb // nheads
-    off_h = off_hb % nheads
-    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
-    offs_d = tl.arange(0, BLOCK_HEADDIM)
-    o = tl.load(Out + off_b * stride_ob + off_h * stride_oh + offs_m[:, None] * stride_om + offs_d[None, :], mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim), other=0.0).to(tl.float32)
-    do = tl.load(DO + off_b * stride_dob + off_h * stride_doh + offs_m[:, None] * stride_dom + offs_d[None, :], mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim), other=0.0).to(tl.float32)
-    delta = tl.sum(o * do, axis=1)
-    tl.store(Delta + off_hb * seqlen_q_rounded + offs_m, delta)
-
-@triton.jit
-def _bwd_store_dk_dv(dk_ptrs, dv_ptrs, dk, dv, offs_n, offs_d, seqlen_k, headdim, EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_HEADDIM: tl.constexpr):
-    if EVEN_N & EVEN_M:
-        if EVEN_HEADDIM:
-            tl.store(dv_ptrs, dv)
-            tl.store(dk_ptrs, dk)
-        else:
-            tl.store(dv_ptrs, dv, mask=offs_d[None, :] < headdim)
-            tl.store(dk_ptrs, dk, mask=offs_d[None, :] < headdim)
-    elif EVEN_HEADDIM:
-        tl.store(dv_ptrs, dv, mask=offs_n[:, None] < seqlen_k)
-        tl.store(dk_ptrs, dk, mask=offs_n[:, None] < seqlen_k)
-    else:
-        tl.store(dv_ptrs, dv, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim))
-        tl.store(dk_ptrs, dk, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim))
-
-@triton.jit
-def _bwd_kernel_one_col_block(start_n, Q, K, V, Bias, DO, DQ, DK, DV, LSE, D, softmax_scale, stride_qm, stride_kn, stride_vn, stride_bm, stride_dom, stride_dqm, stride_dkn, stride_dvn, seqlen_q, seqlen_k, headdim, ATOMIC_ADD: tl.constexpr, BIAS_TYPE: tl.constexpr, IS_CAUSAL: tl.constexpr, BLOCK_HEADDIM: tl.constexpr, EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_HEADDIM: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
-    begin_m = 0 if not IS_CAUSAL else start_n * BLOCK_N // BLOCK_M * BLOCK_M
-    offs_qm = begin_m + tl.arange(0, BLOCK_M)
-    offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N)
-    offs_m = tl.arange(0, BLOCK_M)
-    offs_d = tl.arange(0, BLOCK_HEADDIM)
-    q_ptrs = Q + (offs_qm[:, None] * stride_qm + offs_d[None, :])
-    k_ptrs = K + (offs_n[:, None] * stride_kn + offs_d[None, :])
-    v_ptrs = V + (offs_n[:, None] * stride_vn + offs_d[None, :])
-    do_ptrs = DO + (offs_qm[:, None] * stride_dom + offs_d[None, :])
-    dq_ptrs = DQ + (offs_qm[:, None] * stride_dqm + offs_d[None, :])
-    if BIAS_TYPE == 'vector':
-        b_ptrs = Bias + offs_n
-    elif BIAS_TYPE == 'matrix':
-        b_ptrs = Bias + (offs_qm[:, None] * stride_bm + offs_n[None, :])
-    dv = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)
-    dk = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)
-    if begin_m >= seqlen_q:
-        dv_ptrs = DV + (offs_n[:, None] * stride_dvn + offs_d[None, :])
-        dk_ptrs = DK + (offs_n[:, None] * stride_dkn + offs_d[None, :])
-        _bwd_store_dk_dv(dk_ptrs, dv_ptrs, dk, dv, offs_n, offs_d, seqlen_k, headdim, EVEN_M=EVEN_M, EVEN_N=EVEN_N, EVEN_HEADDIM=EVEN_HEADDIM)
-        return
-    if EVEN_N & EVEN_M:
-        if EVEN_HEADDIM:
-            k = tl.load(k_ptrs)
-            v = tl.load(v_ptrs)
-        else:
-            k = tl.load(k_ptrs, mask=offs_d[None, :] < headdim, other=0.0)
-            v = tl.load(v_ptrs, mask=offs_d[None, :] < headdim, other=0.0)
-    elif EVEN_HEADDIM:
-        k = tl.load(k_ptrs, mask=offs_n[:, None] < seqlen_k, other=0.0)
-        v = tl.load(v_ptrs, mask=offs_n[:, None] < seqlen_k, other=0.0)
-    else:
-        k = tl.load(k_ptrs, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim), other=0.0)
-        v = tl.load(v_ptrs, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim), other=0.0)
-    num_block_m = tl.cdiv(seqlen_q, BLOCK_M)
-    for start_m in range(begin_m, num_block_m * BLOCK_M, BLOCK_M):
-        start_m = tl.multiple_of(start_m, BLOCK_M)
-        offs_m_curr = start_m + offs_m
-        if EVEN_M & EVEN_HEADDIM:
-            q = tl.load(q_ptrs)
-        elif EVEN_HEADDIM:
-            q = tl.load(q_ptrs, mask=offs_m_curr[:, None] < seqlen_q, other=0.0)
-        else:
-            q = tl.load(q_ptrs, mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim), other=0.0)
-        qk = tl.dot(q, k, trans_b=True)
-        if not EVEN_N:
-            qk = tl.where(offs_n[None, :] < seqlen_k, qk, float('-inf'))
-        if IS_CAUSAL:
-            qk = tl.where(offs_m_curr[:, None] >= offs_n[None, :], qk, float('-inf'))
-        if BIAS_TYPE != 'none':
-            tl.debug_barrier()
-            if BIAS_TYPE == 'vector':
-                if EVEN_N:
-                    bias = tl.load(b_ptrs).to(tl.float32)
-                else:
-                    bias = tl.load(b_ptrs, mask=offs_n < seqlen_k, other=0.0).to(tl.float32)
-                bias = bias[None, :]
-            elif BIAS_TYPE == 'matrix':
-                if EVEN_M & EVEN_N:
-                    bias = tl.load(b_ptrs).to(tl.float32)
-                else:
-                    bias = tl.load(b_ptrs, mask=(offs_m_curr[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k), other=0.0).to(tl.float32)
-            qk = qk * softmax_scale + bias
-        if not EVEN_M & EVEN_HEADDIM:
-            tl.debug_barrier()
-        lse_i = tl.load(LSE + offs_m_curr)
-        if BIAS_TYPE == 'none':
-            p = tl.exp(qk * softmax_scale - lse_i[:, None])
-        else:
-            p = tl.exp(qk - lse_i[:, None])
-        if EVEN_M & EVEN_HEADDIM:
-            do = tl.load(do_ptrs)
-        else:
-            do = tl.load(do_ptrs, mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim), other=0.0)
-        dv += tl.dot(p.to(do.dtype), do, trans_a=True)
-        if not EVEN_M & EVEN_HEADDIM:
-            tl.debug_barrier()
-        dp = tl.dot(do, v, trans_b=True)
-        if not EVEN_HEADDIM:
-            tl.debug_barrier()
-        Di = tl.load(D + offs_m_curr)
-        ds = (p * (dp - Di[:, None]) * softmax_scale).to(q.dtype)
-        dk += tl.dot(ds, q, trans_a=True)
-        if not EVEN_M & EVEN_HEADDIM:
-            tl.debug_barrier()
-        if not ATOMIC_ADD:
-            if EVEN_M & EVEN_HEADDIM:
-                dq = tl.load(dq_ptrs, eviction_policy='evict_last')
-                dq += tl.dot(ds, k)
-                tl.store(dq_ptrs, dq, eviction_policy='evict_last')
-            elif EVEN_HEADDIM:
-                dq = tl.load(dq_ptrs, mask=offs_m_curr[:, None] < seqlen_q, other=0.0, eviction_policy='evict_last')
-                dq += tl.dot(ds, k)
-                tl.store(dq_ptrs, dq, mask=offs_m_curr[:, None] < seqlen_q, eviction_policy='evict_last')
-            else:
-                dq = tl.load(dq_ptrs, mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim), other=0.0, eviction_policy='evict_last')
-                dq += tl.dot(ds, k)
-                tl.store(dq_ptrs, dq, mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim), eviction_policy='evict_last')
-        else:
-            dq = tl.dot(ds, k)
-            if EVEN_M & EVEN_HEADDIM:
-                tl.atomic_add(dq_ptrs, dq)
-            elif EVEN_HEADDIM:
-                tl.atomic_add(dq_ptrs, dq, mask=offs_m_curr[:, None] < seqlen_q)
-            else:
-                tl.atomic_add(dq_ptrs, dq, mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim))
-        dq_ptrs += BLOCK_M * stride_dqm
-        q_ptrs += BLOCK_M * stride_qm
-        do_ptrs += BLOCK_M * stride_dom
-        if BIAS_TYPE == 'matrix':
-            b_ptrs += BLOCK_M * stride_bm
-    dv_ptrs = DV + (offs_n[:, None] * stride_dvn + offs_d[None, :])
-    dk_ptrs = DK + (offs_n[:, None] * stride_dkn + offs_d[None, :])
-    _bwd_store_dk_dv(dk_ptrs, dv_ptrs, dk, dv, offs_n, offs_d, seqlen_k, headdim, EVEN_M=EVEN_M, EVEN_N=EVEN_N, EVEN_HEADDIM=EVEN_HEADDIM)
-
-def init_to_zero(name):
-    return lambda nargs: nargs[name].zero_()
-
-@triton.autotune(configs=[triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'SEQUENCE_PARALLEL': False}, num_warps=8, num_stages=1, pre_hook=init_to_zero('DQ')), triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'SEQUENCE_PARALLEL': True}, num_warps=8, num_stages=1, pre_hook=init_to_zero('DQ'))], key=['CACHE_KEY_SEQLEN_Q', 'CACHE_KEY_SEQLEN_K', 'BIAS_TYPE', 'IS_CAUSAL', 'BLOCK_HEADDIM'])
-@triton.heuristics({'EVEN_M': lambda args: args['seqlen_q'] % args['BLOCK_M'] == 0, 'EVEN_N': lambda args: args['seqlen_k'] % args['BLOCK_N'] == 0, 'EVEN_HEADDIM': lambda args: args['headdim'] == args['BLOCK_HEADDIM']})
-@triton.jit
-def _bwd_kernel(Q, K, V, Bias, DO, DQ, DK, DV, LSE, D, softmax_scale, stride_qb, stride_qh, stride_qm, stride_kb, stride_kh, stride_kn, stride_vb, stride_vh, stride_vn, stride_bb, stride_bh, stride_bm, stride_dob, stride_doh, stride_dom, stride_dqb, stride_dqh, stride_dqm, stride_dkb, stride_dkh, stride_dkn, stride_dvb, stride_dvh, stride_dvn, nheads, seqlen_q, seqlen_k, seqlen_q_rounded, headdim, CACHE_KEY_SEQLEN_Q, CACHE_KEY_SEQLEN_K, BIAS_TYPE: tl.constexpr, IS_CAUSAL: tl.constexpr, BLOCK_HEADDIM: tl.constexpr, SEQUENCE_PARALLEL: tl.constexpr, EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_HEADDIM: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
-    off_hb = tl.program_id(1)
-    off_b = off_hb // nheads
-    off_h = off_hb % nheads
-    Q += off_b * stride_qb + off_h * stride_qh
-    K += off_b * stride_kb + off_h * stride_kh
-    V += off_b * stride_vb + off_h * stride_vh
-    DO += off_b * stride_dob + off_h * stride_doh
-    DQ += off_b * stride_dqb + off_h * stride_dqh
-    DK += off_b * stride_dkb + off_h * stride_dkh
-    DV += off_b * stride_dvb + off_h * stride_dvh
-    if BIAS_TYPE != 'none':
-        Bias += off_b * stride_bb + off_h * stride_bh
-    D += off_hb * seqlen_q_rounded
-    LSE += off_hb * seqlen_q_rounded
-    if not SEQUENCE_PARALLEL:
-        num_block_n = tl.cdiv(seqlen_k, BLOCK_N)
-        for start_n in range(0, num_block_n):
-            _bwd_kernel_one_col_block(start_n, Q, K, V, Bias, DO, DQ, DK, DV, LSE, D, softmax_scale, stride_qm, stride_kn, stride_vn, stride_bm, stride_dom, stride_dqm, stride_dkn, stride_dvn, seqlen_q, seqlen_k, headdim, ATOMIC_ADD=False, BIAS_TYPE=BIAS_TYPE, IS_CAUSAL=IS_CAUSAL, BLOCK_HEADDIM=BLOCK_HEADDIM, EVEN_M=EVEN_M, EVEN_N=EVEN_N, EVEN_HEADDIM=EVEN_HEADDIM, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N)
-    else:
-        start_n = tl.program_id(0)
-        _bwd_kernel_one_col_block(start_n, Q, K, V, Bias, DO, DQ, DK, DV, LSE, D, softmax_scale, stride_qm, stride_kn, stride_vn, stride_bm, stride_dom, stride_dqm, stride_dkn, stride_dvn, seqlen_q, seqlen_k, headdim, ATOMIC_ADD=True, BIAS_TYPE=BIAS_TYPE, IS_CAUSAL=IS_CAUSAL, BLOCK_HEADDIM=BLOCK_HEADDIM, EVEN_M=EVEN_M, EVEN_N=EVEN_N, EVEN_HEADDIM=EVEN_HEADDIM, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N)
-
-def _flash_attn_forward(q, k, v, bias=None, causal=False, softmax_scale=None):
-    (batch, seqlen_q, nheads, d) = q.shape
-    (_, seqlen_k, _, _) = k.shape
-    assert k.shape == (batch, seqlen_k, nheads, d)
-    assert v.shape == (batch, seqlen_k, nheads, d)
-    assert d <= 128, 'FlashAttention only support head dimensions up to 128'
-    assert q.dtype == k.dtype == v.dtype, 'All tensors must have the same type'
-    assert q.dtype in [torch.float16, torch.bfloat16], 'Only support fp16 and bf16'
-    assert q.is_cuda and k.is_cuda and v.is_cuda
-    softmax_scale = softmax_scale or 1.0 / math.sqrt(d)
-    has_bias = bias is not None
-    bias_type = 'none'
-    if has_bias:
-        assert bias.dtype in [q.dtype, torch.float]
-        assert bias.is_cuda
-        assert bias.dim() == 4
-        if bias.stride(-1) != 1:
-            bias = bias.contiguous()
-        if bias.shape[2:] == (1, seqlen_k):
-            bias_type = 'vector'
-        elif bias.shape[2:] == (seqlen_q, seqlen_k):
-            bias_type = 'matrix'
-        else:
-            raise RuntimeError('Last 2 dimensions of bias must be (1, seqlen_k) or (seqlen_q, seqlen_k)')
-        bias = bias.expand(batch, nheads, seqlen_q, seqlen_k)
-    bias_strides = (bias.stride(0), bias.stride(1), bias.stride(2)) if has_bias else (0, 0, 0)
-    seqlen_q_rounded = math.ceil(seqlen_q / 128) * 128
-    lse = torch.empty((batch, nheads, seqlen_q_rounded), device=q.device, dtype=torch.float32)
-    tmp = torch.empty((batch, nheads, seqlen_q_rounded), device=q.device, dtype=torch.float32)
-    o = torch.empty_like(q)
-    BLOCK_HEADDIM = max(triton.next_power_of_2(d), 16)
-    BLOCK = 128
-    num_warps = 4 if d <= 64 else 8
-    grid = lambda META: (triton.cdiv(seqlen_q, META['BLOCK_M']), batch * nheads)
-    _fwd_kernel[grid](q, k, v, bias, o, lse, tmp, softmax_scale, q.stride(0), q.stride(2), q.stride(1), k.stride(0), k.stride(2), k.stride(1), v.stride(0), v.stride(2), v.stride(1), *bias_strides, o.stride(0), o.stride(2), o.stride(1), nheads, seqlen_q, seqlen_k, seqlen_q_rounded, d, seqlen_q // 32, seqlen_k // 32, bias_type, causal, BLOCK_HEADDIM, BLOCK_M=BLOCK, BLOCK_N=BLOCK, num_warps=num_warps, num_stages=1)
-    return (o, lse, softmax_scale)
-
-def _flash_attn_backward(do, q, k, v, o, lse, dq, dk, dv, bias=None, causal=False, softmax_scale=None):
-    if do.stride(-1) != 1:
-        do = do.contiguous()
-    (batch, seqlen_q, nheads, d) = q.shape
-    (_, seqlen_k, _, _) = k.shape
-    assert d <= 128
-    seqlen_q_rounded = math.ceil(seqlen_q / 128) * 128
-    assert lse.shape == (batch, nheads, seqlen_q_rounded)
-    assert q.stride(-1) == k.stride(-1) == v.stride(-1) == o.stride(-1) == 1
-    assert dq.stride(-1) == dk.stride(-1) == dv.stride(-1) == 1
-    softmax_scale = softmax_scale or 1.0 / math.sqrt(d)
-    dq_accum = torch.empty_like(q, dtype=torch.float32)
-    delta = torch.empty_like(lse)
-    BLOCK_HEADDIM = max(triton.next_power_of_2(d), 16)
-    grid = lambda META: (triton.cdiv(seqlen_q, META['BLOCK_M']), batch * nheads)
-    _bwd_preprocess_do_o_dot[grid](o, do, delta, o.stride(0), o.stride(2), o.stride(1), do.stride(0), do.stride(2), do.stride(1), nheads, seqlen_q, seqlen_q_rounded, d, BLOCK_M=128, BLOCK_HEADDIM=BLOCK_HEADDIM)
-    has_bias = bias is not None
-    bias_type = 'none'
-    if has_bias:
-        assert bias.dtype in [q.dtype, torch.float]
-        assert bias.is_cuda
-        assert bias.dim() == 4
-        assert bias.stride(-1) == 1
-        if bias.shape[2:] == (1, seqlen_k):
-            bias_type = 'vector'
-        elif bias.shape[2:] == (seqlen_q, seqlen_k):
-            bias_type = 'matrix'
-        else:
-            raise RuntimeError('Last 2 dimensions of bias must be (1, seqlen_k) or (seqlen_q, seqlen_k)')
-        bias = bias.expand(batch, nheads, seqlen_q, seqlen_k)
-    bias_strides = (bias.stride(0), bias.stride(1), bias.stride(2)) if has_bias else (0, 0, 0)
-    grid = lambda META: (triton.cdiv(seqlen_k, META['BLOCK_N']) if META['SEQUENCE_PARALLEL'] else 1, batch * nheads)
-    _bwd_kernel[grid](q, k, v, bias, do, dq_accum, dk, dv, lse, delta, softmax_scale, q.stride(0), q.stride(2), q.stride(1), k.stride(0), k.stride(2), k.stride(1), v.stride(0), v.stride(2), v.stride(1), *bias_strides, do.stride(0), do.stride(2), do.stride(1), dq_accum.stride(0), dq_accum.stride(2), dq_accum.stride(1), dk.stride(0), dk.stride(2), dk.stride(1), dv.stride(0), dv.stride(2), dv.stride(1), nheads, seqlen_q, seqlen_k, seqlen_q_rounded, d, seqlen_q // 32, seqlen_k // 32, bias_type, causal, BLOCK_HEADDIM)
-    dq.copy_(dq_accum)
-
-class FlashAttnQKVPackedFunc(torch.autograd.Function):
-
-    @staticmethod
-    def forward(ctx, qkv, bias=None, causal=False, softmax_scale=None):
-        """
-            qkv: (batch, seqlen, 3, nheads, headdim)
-            bias: optional, shape broadcastible to (batch, nheads, seqlen, seqlen).
-                For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen).
-                ALiBi mask for non-causal would have shape (1, nheads, seqlen, seqlen)
-        """
-        if qkv.stride(-1) != 1:
-            qkv = qkv.contiguous()
-        (o, lse, ctx.softmax_scale) = _flash_attn_forward(qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2], bias=bias, causal=causal, softmax_scale=softmax_scale)
-        ctx.save_for_backward(qkv, o, lse, bias)
-        ctx.causal = causal
-        return o
-
-    @staticmethod
-    def backward(ctx, do):
-        (qkv, o, lse, bias) = ctx.saved_tensors
-        assert not ctx.needs_input_grad[1], 'FlashAttention does not support bias gradient yet'
-        with torch.inference_mode():
-            dqkv = torch.empty_like(qkv)
-            _flash_attn_backward(do, qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2], o, lse, dqkv[:, :, 0], dqkv[:, :, 1], dqkv[:, :, 2], bias=bias, causal=ctx.causal, softmax_scale=ctx.softmax_scale)
-        return (dqkv, None, None, None)
-flash_attn_qkvpacked_func = FlashAttnQKVPackedFunc.apply
-
-class FlashAttnKVPackedFunc(torch.autograd.Function):
-
-    @staticmethod
-    def forward(ctx, q, kv, bias=None, causal=False, softmax_scale=None):
-        """
-            q: (batch, seqlen_q, nheads, headdim)
-            kv: (batch, seqlen_k, 2, nheads, headdim)
-            bias: optional, shape broadcastible to (batch, nheads, seqlen_q, seqlen_k).
-                For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen_k).
-                ALiBi mask for non-causal would have shape (1, nheads, seqlen_q, seqlen_k)
-        """
-        (q, kv) = [x if x.stride(-1) == 1 else x.contiguous() for x in [q, kv]]
-        (o, lse, ctx.softmax_scale) = _flash_attn_forward(q, kv[:, :, 0], kv[:, :, 1], bias=bias, causal=causal, softmax_scale=softmax_scale)
-        ctx.save_for_backward(q, kv, o, lse, bias)
-        ctx.causal = causal
-        return o
-
-    @staticmethod
-    def backward(ctx, do):
-        (q, kv, o, lse, bias) = ctx.saved_tensors
-        if len(ctx.needs_input_grad) >= 3:
-            assert not ctx.needs_input_grad[2], 'FlashAttention does not support bias gradient yet'
-        with torch.inference_mode():
-            dq = torch.empty_like(q)
-            dkv = torch.empty_like(kv)
-            _flash_attn_backward(do, q, kv[:, :, 0], kv[:, :, 1], o, lse, dq, dkv[:, :, 0], dkv[:, :, 1], bias=bias, causal=ctx.causal, softmax_scale=ctx.softmax_scale)
-        return (dq, dkv, None, None, None)
-flash_attn_kvpacked_func = FlashAttnKVPackedFunc.apply
-
-class FlashAttnFunc(torch.autograd.Function):
-
-    @staticmethod
-    def forward(ctx, q, k, v, bias=None, causal=False, softmax_scale=None):
-        """
-            q: (batch_size, seqlen_q, nheads, headdim)
-            k, v: (batch_size, seqlen_k, nheads, headdim)
-            bias: optional, shape broadcastible to (batch, nheads, seqlen_q, seqlen_k).
-                For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen_k).
-                ALiBi mask for non-causal would have shape (1, nheads, seqlen_q, seqlen_k)
-        """
-        (q, k, v) = [x if x.stride(-1) == 1 else x.contiguous() for x in [q, k, v]]
-        (o, lse, ctx.softmax_scale) = _flash_attn_forward(q, k, v, bias=bias, causal=causal, softmax_scale=softmax_scale)
-        ctx.save_for_backward(q, k, v, o, lse, bias)
-        ctx.causal = causal
-        return o
-
-    @staticmethod
-    def backward(ctx, do):
-        (q, k, v, o, lse, bias) = ctx.saved_tensors
-        assert not ctx.needs_input_grad[3], 'FlashAttention does not support bias gradient yet'
-        with torch.inference_mode():
-            dq = torch.empty_like(q)
-            dk = torch.empty_like(k)
-            dv = torch.empty_like(v)
-            _flash_attn_backward(do, q, k, v, o, lse, dq, dk, dv, bias=bias, causal=ctx.causal, softmax_scale=ctx.softmax_scale)
-        return (dq, dk, dv, None, None, None)
-flash_attn_func = FlashAttnFunc.apply
\ No newline at end of file
diff --git a/instructany2pix/llm/model/language_model/mpt/hf_prefixlm_converter.py b/instructany2pix/llm/model/language_model/mpt/hf_prefixlm_converter.py
deleted file mode 100644
index 8c1a6487202a6400a7116a6bd68b493892ef0d14..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/language_model/mpt/hf_prefixlm_converter.py
+++ /dev/null
@@ -1,415 +0,0 @@
-"""Converts Huggingface Causal LM to Prefix LM.
-
-Conversion does lightweight surgery on a HuggingFace
-Causal LM to convert it to a Prefix LM.
-
-Prefix LMs accepts a `bidirectional_mask` input in `forward`
-and treat the input prompt as the prefix in `generate`.
-"""
-import math
-import warnings
-from types import MethodType
-from typing import Any, Dict, List, Optional, Tuple, Union
-import torch
-from transformers.models.bloom.modeling_bloom import BaseModelOutputWithPastAndCrossAttentions, BloomForCausalLM, BloomModel, CausalLMOutputWithCrossAttentions, CrossEntropyLoss
-from transformers.models.bloom.modeling_bloom import _expand_mask as _expand_mask_bloom
-from transformers.models.bloom.modeling_bloom import _make_causal_mask as _make_causal_mask_bloom
-from transformers.models.bloom.modeling_bloom import logging
-from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel
-from transformers.models.gpt_neo.modeling_gpt_neo import GPTNeoForCausalLM
-from transformers.models.gpt_neox.modeling_gpt_neox import GPTNeoXForCausalLM
-from transformers.models.gptj.modeling_gptj import GPTJForCausalLM
-from transformers.models.opt.modeling_opt import OPTForCausalLM
-from transformers.models.opt.modeling_opt import _expand_mask as _expand_mask_opt
-from transformers.models.opt.modeling_opt import _make_causal_mask as _make_causal_mask_opt
-logger = logging.get_logger(__name__)
-_SUPPORTED_GPT_MODELS = (GPT2LMHeadModel, GPTJForCausalLM, GPTNeoForCausalLM, GPTNeoXForCausalLM)
-CAUSAL_GPT_TYPES = Union[GPT2LMHeadModel, GPTJForCausalLM, GPTNeoForCausalLM, GPTNeoXForCausalLM]
-
-def _convert_gpt_causal_lm_to_prefix_lm(model: CAUSAL_GPT_TYPES) -> CAUSAL_GPT_TYPES:
-    """Converts a GPT-style Causal LM to a Prefix LM.
-
-    Supported HuggingFace model classes:
-        - `GPT2LMHeadModel`
-        - `GPTNeoForCausalLM`
-        - `GPTNeoXForCausalLM`
-        - `GPTJForCausalLM`
-
-    See `convert_hf_causal_lm_to_prefix_lm` for more details.
-    """
-    if hasattr(model, '_prefix_lm_converted'):
-        return model
-    assert isinstance(model, _SUPPORTED_GPT_MODELS)
-    assert model.config.add_cross_attention == False, 'Only supports GPT-style decoder-only models'
-
-    def _get_attn_modules(model: CAUSAL_GPT_TYPES) -> List[torch.nn.Module]:
-        """Helper that gets a list of the model's attention modules.
-
-        Each module has a `bias` buffer used for causal masking. The Prefix LM
-        conversion adds logic to dynamically manipulate these biases to support
-        Prefix LM attention masking.
-        """
-        attn_modules = []
-        if isinstance(model, GPTNeoXForCausalLM):
-            blocks = model.gpt_neox.layers
-        else:
-            blocks = model.transformer.h
-        for block in blocks:
-            if isinstance(model, GPTNeoForCausalLM):
-                if block.attn.attention_type != 'global':
-                    continue
-                attn_module = block.attn.attention
-            elif isinstance(model, GPTNeoXForCausalLM):
-                attn_module = block.attention
-            else:
-                attn_module = block.attn
-            attn_modules.append(attn_module)
-        return attn_modules
-    setattr(model, '_original_forward', getattr(model, 'forward'))
-    setattr(model, '_original_generate', getattr(model, 'generate'))
-
-    def forward(self: CAUSAL_GPT_TYPES, input_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[Tuple[Tuple[torch.Tensor]]]=None, attention_mask: Optional[torch.FloatTensor]=None, bidirectional_mask: Optional[torch.Tensor]=None, token_type_ids: Optional[torch.LongTensor]=None, position_ids: Optional[torch.LongTensor]=None, head_mask: Optional[torch.FloatTensor]=None, inputs_embeds: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None):
-        """Wraps original forward to enable PrefixLM attention."""
-
-        def call_og_forward():
-            if isinstance(self, GPTNeoXForCausalLM):
-                return self._original_forward(input_ids=input_ids, past_key_values=past_key_values, attention_mask=attention_mask, head_mask=head_mask, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
-            else:
-                return self._original_forward(input_ids=input_ids, past_key_values=past_key_values, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
-        if bidirectional_mask is None:
-            return call_og_forward()
-        assert isinstance(bidirectional_mask, torch.Tensor)
-        attn_modules = _get_attn_modules(model)
-        (b, s) = bidirectional_mask.shape
-        max_length = attn_modules[0].bias.shape[-1]
-        if s > max_length:
-            raise ValueError(f'bidirectional_mask sequence length (={s}) exceeds the ' + f'max length allowed by the model ({max_length}).')
-        assert s <= max_length
-        if s < max_length:
-            pad = torch.zeros((int(b), int(max_length - s)), dtype=bidirectional_mask.dtype, device=bidirectional_mask.device)
-            bidirectional_mask = torch.cat([bidirectional_mask, pad], dim=1)
-        bidirectional = bidirectional_mask.unsqueeze(1).unsqueeze(1)
-        for attn_module in attn_modules:
-            attn_module.bias.data = torch.logical_or(attn_module.bias.data, bidirectional)
-        output = call_og_forward()
-        for attn_module in attn_modules:
-            attn_module.bias.data = torch.tril(attn_module.bias.data[0, 0])[None, None]
-        return output
-
-    def generate(self: CAUSAL_GPT_TYPES, *args: tuple, **kwargs: Dict[str, Any]):
-        """Wraps original generate to enable PrefixLM attention."""
-        attn_modules = _get_attn_modules(model)
-        for attn_module in attn_modules:
-            attn_module.bias.data[:] = 1
-        output = self._original_generate(*args, **kwargs)
-        for attn_module in attn_modules:
-            attn_module.bias.data = torch.tril(attn_module.bias.data[0, 0])[None, None]
-        return output
-    setattr(model, 'forward', MethodType(forward, model))
-    setattr(model, 'generate', MethodType(generate, model))
-    setattr(model, '_prefix_lm_converted', True)
-    return model
-
-def _convert_bloom_causal_lm_to_prefix_lm(model: BloomForCausalLM) -> BloomForCausalLM:
-    """Converts a BLOOM Causal LM to a Prefix LM.
-
-    Supported HuggingFace model classes:
-        - `BloomForCausalLM`
-
-    See `convert_hf_causal_lm_to_prefix_lm` for more details.
-    """
-    if hasattr(model, '_prefix_lm_converted'):
-        return model
-    assert isinstance(model, BloomForCausalLM)
-    assert model.config.add_cross_attention == False, 'Only supports BLOOM decoder-only models'
-
-    def _prepare_attn_mask(self: BloomModel, attention_mask: torch.Tensor, bidirectional_mask: Optional[torch.Tensor], input_shape: Tuple[int, int], past_key_values_length: int) -> torch.BoolTensor:
-        combined_attention_mask = None
-        device = attention_mask.device
-        (_, src_length) = input_shape
-        if src_length > 1:
-            combined_attention_mask = _make_causal_mask_bloom(input_shape, device=device, past_key_values_length=past_key_values_length)
-            if bidirectional_mask is not None:
-                assert attention_mask.shape == bidirectional_mask.shape
-                expanded_bidirectional_mask = _expand_mask_bloom(bidirectional_mask, tgt_length=src_length)
-                combined_attention_mask = torch.logical_and(combined_attention_mask, expanded_bidirectional_mask)
-        expanded_attn_mask = _expand_mask_bloom(attention_mask, tgt_length=src_length)
-        combined_attention_mask = expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask | combined_attention_mask
-        return combined_attention_mask
-
-    def _build_alibi_tensor(self: BloomModel, batch_size: int, query_length: int, key_length: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
-        num_heads = self.config.n_head
-        closest_power_of_2 = 2 ** math.floor(math.log2(num_heads))
-        base = torch.tensor(2 ** (-2 ** (-(math.log2(closest_power_of_2) - 3))), device=device, dtype=torch.float32)
-        powers = torch.arange(1, 1 + closest_power_of_2, device=device, dtype=torch.int32)
-        slopes = torch.pow(base, powers)
-        if closest_power_of_2 != num_heads:
-            extra_base = torch.tensor(2 ** (-2 ** (-(math.log2(2 * closest_power_of_2) - 3))), device=device, dtype=torch.float32)
-            num_remaining_heads = min(closest_power_of_2, num_heads - closest_power_of_2)
-            extra_powers = torch.arange(1, 1 + 2 * num_remaining_heads, 2, device=device, dtype=torch.int32)
-            slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0)
-        qa = torch.arange(query_length, device=device, dtype=torch.int32).view(-1, 1)
-        ka = torch.arange(key_length, device=device, dtype=torch.int32).view(1, -1)
-        diffs = qa - ka + key_length - query_length
-        diffs = -diffs.abs()
-        alibi = slopes.view(1, num_heads, 1, 1) * diffs.view(1, 1, query_length, key_length)
-        alibi = alibi.expand(batch_size, -1, -1, -1).reshape(-1, query_length, key_length)
-        return alibi.to(dtype)
-    KeyValueT = Tuple[torch.Tensor, torch.Tensor]
-
-    def forward(self: BloomModel, input_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[Tuple[KeyValueT, ...]]=None, attention_mask: Optional[torch.Tensor]=None, bidirectional_mask: Optional[torch.Tensor]=None, head_mask: Optional[torch.LongTensor]=None, inputs_embeds: Optional[torch.LongTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, **deprecated_arguments) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]:
-        if deprecated_arguments.pop('position_ids', False) is not False:
-            warnings.warn('`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. ' + 'You can safely ignore passing `position_ids`.', FutureWarning)
-        if len(deprecated_arguments) > 0:
-            raise ValueError(f'Got unexpected arguments: {deprecated_arguments}')
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        use_cache = use_cache if use_cache is not None else self.config.use_cache
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        if input_ids is not None and inputs_embeds is not None:
-            raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time')
-        elif input_ids is not None:
-            (batch_size, seq_length) = input_ids.shape
-        elif inputs_embeds is not None:
-            (batch_size, seq_length, _) = inputs_embeds.shape
-        else:
-            raise ValueError('You have to specify either input_ids or inputs_embeds')
-        if past_key_values is None:
-            past_key_values = tuple([None] * len(self.h))
-        head_mask = self.get_head_mask(head_mask, self.config.n_layer)
-        if inputs_embeds is None:
-            inputs_embeds = self.word_embeddings(input_ids)
-        hidden_states = self.word_embeddings_layernorm(inputs_embeds)
-        presents = () if use_cache else None
-        all_self_attentions = () if output_attentions else None
-        all_hidden_states = () if output_hidden_states else None
-        seq_length_with_past = seq_length
-        past_key_values_length = 0
-        if past_key_values[0] is not None:
-            tmp = past_key_values[0][0]
-            past_key_values_length = tmp.shape[2]
-            seq_length_with_past = seq_length_with_past + past_key_values_length
-        if attention_mask is None:
-            attention_mask = torch.ones((batch_size, seq_length_with_past), device=hidden_states.device)
-        else:
-            attention_mask = attention_mask.to(hidden_states.device)
-        alibi = self._build_alibi_tensor(batch_size=batch_size, query_length=seq_length, key_length=seq_length_with_past, dtype=hidden_states.dtype, device=hidden_states.device)
-        causal_mask = self._prepare_attn_mask(attention_mask, bidirectional_mask, input_shape=(batch_size, seq_length), past_key_values_length=past_key_values_length)
-        for (i, (block, layer_past)) in enumerate(zip(self.h, past_key_values)):
-            if output_hidden_states:
-                hst = (hidden_states,)
-                all_hidden_states = all_hidden_states + hst
-            if self.gradient_checkpointing and self.training:
-                if use_cache:
-                    logger.warning('`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...')
-                    use_cache = False
-
-                def create_custom_forward(module):
-
-                    def custom_forward(*inputs):
-                        return module(*inputs, use_cache=use_cache, output_attentions=output_attentions)
-                    return custom_forward
-                outputs = torch.utils.checkpoint.checkpoint(create_custom_forward(block), hidden_states, alibi, causal_mask, head_mask[i])
-            else:
-                outputs = block(hidden_states, layer_past=layer_past, attention_mask=causal_mask, head_mask=head_mask[i], use_cache=use_cache, output_attentions=output_attentions, alibi=alibi)
-            hidden_states = outputs[0]
-            if use_cache is True:
-                presents = presents + (outputs[1],)
-            if output_attentions:
-                oa = (outputs[2 if use_cache else 1],)
-                all_self_attentions = all_self_attentions + oa
-        hidden_states = self.ln_f(hidden_states)
-        if output_hidden_states:
-            hst = (hidden_states,)
-            all_hidden_states = all_hidden_states + hst
-        if not return_dict:
-            return tuple((v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None))
-        return BaseModelOutputWithPastAndCrossAttentions(last_hidden_state=hidden_states, past_key_values=presents, hidden_states=all_hidden_states, attentions=all_self_attentions)
-    setattr(model.transformer, '_prepare_attn_mask', MethodType(_prepare_attn_mask, model.transformer))
-    setattr(model.transformer, '_build_alibi_tensor', MethodType(_build_alibi_tensor, model.transformer))
-    setattr(model.transformer, 'forward', MethodType(forward, model.transformer))
-    KeyValueT = Tuple[torch.Tensor, torch.Tensor]
-
-    def forward(self: BloomForCausalLM, input_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[Tuple[KeyValueT, ...]]=None, attention_mask: Optional[torch.Tensor]=None, bidirectional_mask: Optional[torch.Tensor]=None, head_mask: Optional[torch.Tensor]=None, inputs_embeds: Optional[torch.Tensor]=None, labels: Optional[torch.Tensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, **deprecated_arguments) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]:
-        """Replacement forward method for BloomCausalLM."""
-        if deprecated_arguments.pop('position_ids', False) is not False:
-            warnings.warn('`position_ids` have no functionality in BLOOM and will be removed ' + 'in v5.0.0. You can safely ignore passing `position_ids`.', FutureWarning)
-        if len(deprecated_arguments) > 0:
-            raise ValueError(f'Got unexpected arguments: {deprecated_arguments}')
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        transformer_outputs = self.transformer(input_ids, past_key_values=past_key_values, attention_mask=attention_mask, bidirectional_mask=bidirectional_mask, head_mask=head_mask, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
-        hidden_states = transformer_outputs[0]
-        lm_logits = self.lm_head(hidden_states)
-        loss = None
-        if labels is not None:
-            shift_logits = lm_logits[..., :-1, :].contiguous()
-            shift_labels = labels[..., 1:].contiguous()
-            (batch_size, seq_length, vocab_size) = shift_logits.shape
-            loss_fct = CrossEntropyLoss()
-            loss = loss_fct(shift_logits.view(batch_size * seq_length, vocab_size), shift_labels.view(batch_size * seq_length))
-        if not return_dict:
-            output = (lm_logits,) + transformer_outputs[1:]
-            return (loss,) + output if loss is not None else output
-        return CausalLMOutputWithCrossAttentions(loss=loss, logits=lm_logits, past_key_values=transformer_outputs.past_key_values, hidden_states=transformer_outputs.hidden_states, attentions=transformer_outputs.attentions)
-
-    def prepare_inputs_for_generation(self: BloomForCausalLM, input_ids: torch.LongTensor, past: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None, **kwargs) -> dict:
-        if past:
-            input_ids = input_ids[:, -1].unsqueeze(-1)
-            bidirectional_mask = None
-            if past[0][0].shape[0] == input_ids.shape[0]:
-                past = self._convert_to_bloom_cache(past)
-        else:
-            bidirectional_mask = torch.ones_like(input_ids)
-        return {'input_ids': input_ids, 'past_key_values': past, 'use_cache': True, 'attention_mask': attention_mask, 'bidirectional_mask': bidirectional_mask}
-    setattr(model, 'forward', MethodType(forward, model))
-    setattr(model, 'prepare_inputs_for_generation', MethodType(prepare_inputs_for_generation, model))
-    setattr(model, '_prefix_lm_converted', True)
-    return model
-
-def _convert_opt_causal_lm_to_prefix_lm(model: OPTForCausalLM) -> OPTForCausalLM:
-    """Converts an OPT Causal LM to a Prefix LM.
-
-    Supported HuggingFace model classes:
-        - `OPTForCausalLM`
-
-    See `convert_hf_causal_lm_to_prefix_lm` for more details.
-    """
-    if hasattr(model, '_prefix_lm_converted'):
-        return model
-    assert isinstance(model, OPTForCausalLM)
-    assert model.config.add_cross_attention == False, 'Only supports OPT decoder-only models'
-    setattr(model, '_original_forward', getattr(model, 'forward'))
-    setattr(model, '_original_generate', getattr(model, 'generate'))
-    model.model.decoder.bidirectional_mask = None
-
-    def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
-        combined_attention_mask = None
-        if input_shape[-1] > 1:
-            if self.bidirectional_mask == 'g':
-                (bsz, src_length) = input_shape
-                combined_attention_mask = torch.zeros((bsz, 1, src_length, src_length + past_key_values_length), dtype=inputs_embeds.dtype, device=inputs_embeds.device)
-            else:
-                combined_attention_mask = _make_causal_mask_opt(input_shape, inputs_embeds.dtype, past_key_values_length=past_key_values_length).to(inputs_embeds.device)
-                if self.bidirectional_mask is not None:
-                    assert attention_mask.shape == self.bidirectional_mask.shape
-                    expanded_bidirectional_mask = _expand_mask_opt(self.bidirectional_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(inputs_embeds.device)
-                    combined_attention_mask = torch.maximum(expanded_bidirectional_mask, combined_attention_mask)
-        if attention_mask is not None:
-            expanded_attn_mask = _expand_mask_opt(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(inputs_embeds.device)
-            combined_attention_mask = expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
-        return combined_attention_mask
-    setattr(model.model.decoder, '_prepare_decoder_attention_mask', MethodType(_prepare_decoder_attention_mask, model.model.decoder))
-
-    def forward(self: OPTForCausalLM, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.Tensor]=None, bidirectional_mask: Optional[torch.ByteTensor]=None, head_mask: Optional[torch.Tensor]=None, past_key_values: Optional[List[torch.FloatTensor]]=None, inputs_embeds: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None):
-
-        def call_og_forward():
-            return self._original_forward(input_ids=input_ids, attention_mask=attention_mask, head_mask=head_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
-        if bidirectional_mask is None:
-            return call_og_forward()
-        self.model.decoder.bidirectional_mask = bidirectional_mask
-        try:
-            outputs = call_og_forward()
-        except:
-            self.model.decoder.bidirectional_mask = None
-            raise
-        self.model.decoder.bidirectional_mask = None
-        return outputs
-
-    def generate(self: OPTForCausalLM, *args: tuple, **kwargs: Dict[str, Any]):
-        """Wraps original generate to enable PrefixLM-style attention."""
-        self.model.decoder.bidirectional_mask = 'g'
-        try:
-            output = self._original_generate(*args, **kwargs)
-        except:
-            self.model.decoder.bidirectional_mask = None
-            raise
-        self.model.decoder.bidirectional_mask = None
-        return output
-    setattr(model, 'forward', MethodType(forward, model))
-    setattr(model, 'generate', MethodType(generate, model))
-    setattr(model, '_prefix_lm_converted', True)
-    return model
-_SUPPORTED_HF_MODELS = _SUPPORTED_GPT_MODELS + (BloomForCausalLM, OPTForCausalLM)
-CAUSAL_LM_TYPES = Union[GPT2LMHeadModel, GPTJForCausalLM, GPTNeoForCausalLM, GPTNeoXForCausalLM, BloomForCausalLM, OPTForCausalLM]
-
-def convert_hf_causal_lm_to_prefix_lm(model: CAUSAL_LM_TYPES) -> CAUSAL_LM_TYPES:
-    """Converts a HuggingFace Causal LM to a Prefix LM.
-
-    Supported HuggingFace model classes:
-        - `GPT2LMHeadModel`
-        - `GPTNeoForCausalLM`
-        - `GPTNeoXForCausalLM`
-        - `GPTJForCausalLM`
-        - `BloomForCausalLM`
-        - `OPTForCausalLM`
-
-    Conversion to a Prefix LM is done by modifying the `forward` method, and possibly also the
-    `generate` method and/or select underlying methods depending on the model class.
-
-    These changes preserve the model API, but add a new input to `forward`: "bidirectional_mask".
-
-    Notes on training:
-        To actually train the converted model as a Prefix LM, training batches will need to indicate
-        the prefix/target structure by including `bidirectional_mask` as part of the batch inputs.
-
-        **This is not a standard input and requires custom layers either within or after your dataloader.**
-
-        In addition to adding `bidirectional_mask` to the batch, this custom code should modify `labels`
-        such that `batch['labels'][batch['bidirectional_mask'] == 1] == -100`.
-        That is, the prefix portion of the sequence should not generate any loss. Loss should only be
-        generated by the target portion of the sequence.
-
-    Notes on `GPTNeoForCausalLM`:
-        To simplify the implementation, "global" and "local" attention layers are handled differently.
-        For "global" layers, we handle conversion as described above. For "local" layers, which use a
-        causal attention mask within a restricted local window, we do not alter the masking.
-
-    Notes on `forward` method conversion:
-        After conversion, the `forward` method will handle a new input, `bidirectional_mask`,
-        which should be a [batch_size, seq_length] byte tensor, where 1 indicates token positions
-        belonging to the prefix (prefix tokens can attend to one another bidirectionally), and
-        0 indicates token positions belonging to the target.
-
-        The new `forward` method will incorporate `bidirectional_mask` (if supplied) into the existing
-        causal mask, call the original `forward` method, and (if the causal mask is a buffer) reset
-        the causal masks before returning the result.
-
-    Notes on `generate` method conversion:
-        After conversion, the `generate` method will have the same signature but will internally
-        convert all causal masks to be purely bidirectional, call the original `generate` method, and
-        (where appropriate) reset the causal masks before returning the result.
-
-        This works thanks to the logic of the HuggingFace `generate` API, which first encodes the token
-        "prompt" passed to `generate` (which is treated as the prefix) and then sequentially generates
-        each new token. Encodings are cached as generation happens, so all prefix tokens can attend to one
-        another (as expected in a Prefix LM) and generated tokens can only attend to prefix tokens and
-        previously-generated tokens (also as expected in a Prefix LM).
-
-    To preserve the API, the original methods are renamed to `_original_forward` and
-    `_original_generate`, and replaced with new `forward` and `generate` methods that wrap
-    them, respectively. Although implementation details vary by model class.
-    """
-    if isinstance(model, _SUPPORTED_GPT_MODELS):
-        return _convert_gpt_causal_lm_to_prefix_lm(model)
-    elif isinstance(model, BloomForCausalLM):
-        return _convert_bloom_causal_lm_to_prefix_lm(model)
-    elif isinstance(model, OPTForCausalLM):
-        return _convert_opt_causal_lm_to_prefix_lm(model)
-    else:
-        raise TypeError(f'Cannot convert model to Prefix LM. ' + f'Model does not belong to set of supported HF models:' + f'\n{_SUPPORTED_HF_MODELS}')
-
-def add_bidirectional_mask_if_missing(batch: Dict[str, Any]):
-    """Attempts to add bidirectional_mask to batch if missing.
-
-    Raises:
-        KeyError if bidirectional_mask is missing and can't be inferred
-    """
-    if 'bidirectional_mask' not in batch:
-        if batch.get('mode', None) == 'icl_task':
-            batch['bidirectional_mask'] = batch['attention_mask'].clone()
-            for (i, continuation_indices) in enumerate(batch['continuation_indices']):
-                batch['bidirectional_mask'][i, continuation_indices] = 0
-        elif 'labels' in batch and 'attention_mask' in batch:
-            batch['bidirectional_mask'] = torch.logical_and(torch.eq(batch['attention_mask'], 1), torch.eq(batch['labels'], -100)).type_as(batch['attention_mask'])
-        else:
-            raise KeyError('No bidirectional_mask in batch and not sure how to construct one.')
\ No newline at end of file
diff --git a/instructany2pix/llm/model/language_model/mpt/meta_init_context.py b/instructany2pix/llm/model/language_model/mpt/meta_init_context.py
deleted file mode 100644
index 6cba6fff0fe21fe222c7ab38eae44a9784c0be9c..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/language_model/mpt/meta_init_context.py
+++ /dev/null
@@ -1,94 +0,0 @@
-from contextlib import contextmanager
-import torch
-import torch.nn as nn
-
-@contextmanager
-def init_empty_weights(include_buffers: bool=False):
-    """Meta initialization context manager.
-
-    A context manager under which models are initialized with all parameters
-    on the meta device, therefore creating an empty model. Useful when just
-    initializing the model would blow the available RAM.
-
-    Args:
-        include_buffers (`bool`, *optional*, defaults to `False`): Whether or
-            not to also put all buffers on the meta device while initializing.
-
-    Example:
-    ```python
-    import torch.nn as nn
-
-    # Initialize a model with 100 billions parameters in no time and without using any RAM.
-    with init_empty_weights():
-        tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
-    ```
-
-    <Tip warning={true}>
-
-    Any model created under this context manager has no weights. As such you can't do something like
-    `model.to(some_device)` with it. To load weights inside your empty model, see [`load_checkpoint_and_dispatch`].
-
-    </Tip>
-    """
-    with init_on_device(torch.device('meta'), include_buffers=include_buffers) as f:
-        yield f
-
-@contextmanager
-def init_on_device(device: torch.device, include_buffers: bool=False):
-    """Device initialization context manager.
-
-    A context manager under which models are initialized with all parameters
-    on the specified device.
-
-    Args:
-        device (`torch.device`): Device to initialize all parameters on.
-        include_buffers (`bool`, *optional*, defaults to `False`): Whether or
-            not to also put all buffers on the meta device while initializing.
-
-    Example:
-    ```python
-    import torch.nn as nn
-
-    with init_on_device(device=torch.device("cuda")):
-        tst = nn.Liner(100, 100)  # on `cuda` device
-    ```
-    """
-    old_register_parameter = nn.Module.register_parameter
-    if include_buffers:
-        old_register_buffer = nn.Module.register_buffer
-
-    def register_empty_parameter(module, name, param):
-        old_register_parameter(module, name, param)
-        if param is not None:
-            param_cls = type(module._parameters[name])
-            kwargs = module._parameters[name].__dict__
-            module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs)
-
-    def register_empty_buffer(module, name, buffer):
-        old_register_buffer(module, name, buffer)
-        if buffer is not None:
-            module._buffers[name] = module._buffers[name].to(device)
-    if include_buffers:
-        tensor_constructors_to_patch = {torch_function_name: getattr(torch, torch_function_name) for torch_function_name in ['empty', 'zeros', 'ones', 'full']}
-    else:
-        tensor_constructors_to_patch = {}
-
-    def patch_tensor_constructor(fn):
-
-        def wrapper(*args, **kwargs):
-            kwargs['device'] = device
-            return fn(*args, **kwargs)
-        return wrapper
-    try:
-        nn.Module.register_parameter = register_empty_parameter
-        if include_buffers:
-            nn.Module.register_buffer = register_empty_buffer
-        for torch_function_name in tensor_constructors_to_patch.keys():
-            setattr(torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name)))
-        yield
-    finally:
-        nn.Module.register_parameter = old_register_parameter
-        if include_buffers:
-            nn.Module.register_buffer = old_register_buffer
-        for (torch_function_name, old_torch_function) in tensor_constructors_to_patch.items():
-            setattr(torch, torch_function_name, old_torch_function)
\ No newline at end of file
diff --git a/instructany2pix/llm/model/language_model/mpt/modeling_mpt.py b/instructany2pix/llm/model/language_model/mpt/modeling_mpt.py
deleted file mode 100644
index 13313441b13fc7a66cb65fd21b482a5de982e2c8..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/language_model/mpt/modeling_mpt.py
+++ /dev/null
@@ -1,331 +0,0 @@
-"""A simple, flexible implementation of a GPT model.
-
-Inspired by https://github.com/karpathy/minGPT/blob/master/mingpt/model.py
-"""
-import math
-import warnings
-from typing import List, Optional, Tuple, Union
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from transformers import PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast
-from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
-from .attention import attn_bias_shape, build_attn_bias
-from .blocks import MPTBlock
-from .custom_embedding import SharedEmbedding
-from .norm import NORM_CLASS_REGISTRY
-from .configuration_mpt import MPTConfig
-from .adapt_tokenizer import AutoTokenizerForMOD, adapt_tokenizer_for_denoising
-from .hf_prefixlm_converter import add_bidirectional_mask_if_missing, convert_hf_causal_lm_to_prefix_lm
-from .meta_init_context import init_empty_weights
-from .param_init_fns import MODEL_INIT_REGISTRY, generic_param_init_fn_
-try:
-    from .flash_attn_triton import flash_attn_func
-except:
-    pass
-Tokenizer = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
-
-class MPTPreTrainedModel(PreTrainedModel):
-    config_class = MPTConfig
-    base_model_prefix = 'model'
-    _no_split_modules = ['MPTBlock']
-
-class MPTModel(MPTPreTrainedModel):
-
-    def __init__(self, config: MPTConfig):
-        config._validate_config()
-        super().__init__(config)
-        self.attn_impl = config.attn_config['attn_impl']
-        self.prefix_lm = config.attn_config['prefix_lm']
-        self.attn_uses_sequence_id = config.attn_config['attn_uses_sequence_id']
-        self.alibi = config.attn_config['alibi']
-        self.alibi_bias_max = config.attn_config['alibi_bias_max']
-        if config.init_device == 'mixed':
-            if dist.get_local_rank() == 0:
-                config.init_device = 'cpu'
-            else:
-                config.init_device = 'meta'
-        if config.norm_type.lower() not in NORM_CLASS_REGISTRY.keys():
-            norm_options = ' | '.join(NORM_CLASS_REGISTRY.keys())
-            raise NotImplementedError(f'Requested norm type ({config.norm_type}) is not implemented within this repo (Options: {norm_options}).')
-        norm_class = NORM_CLASS_REGISTRY[config.norm_type.lower()]
-        self.embedding_fraction = config.embedding_fraction
-        self.wte = SharedEmbedding(config.vocab_size, config.d_model, device=config.init_device)
-        if not self.alibi:
-            self.wpe = torch.nn.Embedding(config.max_seq_len, config.d_model, device=config.init_device)
-        self.emb_drop = nn.Dropout(config.emb_pdrop)
-        self.blocks = nn.ModuleList([MPTBlock(device=config.init_device, **config.to_dict()) for _ in range(config.n_layers)])
-        self.norm_f = norm_class(config.d_model, device=config.init_device)
-        if config.init_device != 'meta':
-            print(f'You are using config.init_device={config.init_device!r}, but you can also use config.init_device="meta" with Composer + FSDP for fast initialization.')
-            self.apply(self.param_init_fn)
-        self.is_causal = not self.prefix_lm
-        self._attn_bias_initialized = False
-        self.attn_bias = None
-        self.attn_bias_shape = attn_bias_shape(self.attn_impl, config.n_heads, config.max_seq_len, self.alibi, prefix_lm=self.prefix_lm, causal=self.is_causal, use_sequence_id=self.attn_uses_sequence_id)
-        if config.no_bias:
-            for module in self.modules():
-                if hasattr(module, 'bias') and isinstance(module.bias, nn.Parameter):
-                    if config.verbose:
-                        warnings.warn(f'Removing bias ({module.bias}) from {module}.')
-                    module.register_parameter('bias', None)
-        if config.verbose and config.verbose > 2:
-            print(self)
-        if 'verbose' not in self.config.init_config:
-            self.config.init_config['verbose'] = self.config.verbose
-        if self.config.init_config['verbose'] > 1:
-            init_fn_name = self.config.init_config['name']
-            warnings.warn(f'Using {init_fn_name} initialization.')
-        self.gradient_checkpointing = False
-
-    def get_input_embeddings(self):
-        return self.wte
-
-    def set_input_embeddings(self, value):
-        self.wte = value
-
-    @torch.no_grad()
-    def _attn_bias(self, device, dtype, attention_mask: Optional[torch.ByteTensor]=None, prefix_mask: Optional[torch.ByteTensor]=None, sequence_id: Optional[torch.LongTensor]=None):
-        if not self._attn_bias_initialized:
-            if self.attn_bias_shape:
-                self.attn_bias = torch.zeros(self.attn_bias_shape, device=device, dtype=dtype)
-                self.attn_bias = build_attn_bias(self.attn_impl, self.attn_bias, self.config.n_heads, self.config.max_seq_len, causal=self.is_causal, alibi=self.alibi, alibi_bias_max=self.alibi_bias_max)
-            self._attn_bias_initialized = True
-        if self.attn_impl == 'flash':
-            return (self.attn_bias, attention_mask)
-        if self.attn_bias is not None:
-            self.attn_bias = self.attn_bias.to(dtype=dtype, device=device)
-        attn_bias = self.attn_bias
-        if self.prefix_lm:
-            assert isinstance(attn_bias, torch.Tensor)
-            assert isinstance(prefix_mask, torch.Tensor)
-            attn_bias = self._apply_prefix_mask(attn_bias, prefix_mask)
-        if self.attn_uses_sequence_id and sequence_id is not None:
-            assert isinstance(attn_bias, torch.Tensor)
-            attn_bias = self._apply_sequence_id(attn_bias, sequence_id)
-        if attention_mask is not None:
-            s_k = attention_mask.shape[-1]
-            if attn_bias is None:
-                attn_bias = torch.zeros((1, 1, 1, s_k), device=device, dtype=dtype)
-            else:
-                _s_k = max(0, attn_bias.size(-1) - s_k)
-                attn_bias = attn_bias[:, :, :, _s_k:]
-            if prefix_mask is not None and attention_mask.shape != prefix_mask.shape:
-                raise ValueError(f'attention_mask shape={attention_mask.shape} ' + f'and prefix_mask shape={prefix_mask.shape} are not equal.')
-            min_val = torch.finfo(attn_bias.dtype).min
-            attn_bias = attn_bias.masked_fill(~attention_mask.view(-1, 1, 1, s_k), min_val)
-        return (attn_bias, None)
-
-    def _apply_prefix_mask(self, attn_bias: torch.Tensor, prefix_mask: torch.Tensor):
-        (s_k, s_q) = attn_bias.shape[-2:]
-        if s_k != self.config.max_seq_len or s_q != self.config.max_seq_len:
-            raise ValueError('attn_bias does not match the expected shape. ' + f'The last two dimensions should both be {self.config.max_length} ' + f'but are {s_k} and {s_q}.')
-        seq_len = prefix_mask.shape[-1]
-        if seq_len > self.config.max_seq_len:
-            raise ValueError(f'prefix_mask sequence length cannot exceed max_seq_len={self.config.max_seq_len}')
-        attn_bias = attn_bias[..., :seq_len, :seq_len]
-        causal = torch.tril(torch.ones((seq_len, seq_len), dtype=torch.bool, device=prefix_mask.device)).view(1, 1, seq_len, seq_len)
-        prefix = prefix_mask.view(-1, 1, 1, seq_len)
-        cannot_attend = ~torch.logical_or(causal, prefix.bool())
-        min_val = torch.finfo(attn_bias.dtype).min
-        attn_bias = attn_bias.masked_fill(cannot_attend, min_val)
-        return attn_bias
-
-    def _apply_sequence_id(self, attn_bias: torch.Tensor, sequence_id: torch.LongTensor):
-        seq_len = sequence_id.shape[-1]
-        if seq_len > self.config.max_seq_len:
-            raise ValueError(f'sequence_id sequence length cannot exceed max_seq_len={self.config.max_seq_len}')
-        attn_bias = attn_bias[..., :seq_len, :seq_len]
-        cannot_attend = torch.logical_not(torch.eq(sequence_id.view(-1, seq_len, 1), sequence_id.view(-1, 1, seq_len))).unsqueeze(1)
-        min_val = torch.finfo(attn_bias.dtype).min
-        attn_bias = attn_bias.masked_fill(cannot_attend, min_val)
-        return attn_bias
-
-    def forward(self, input_ids: torch.LongTensor, past_key_values: Optional[List[Tuple[torch.FloatTensor]]]=None, attention_mask: Optional[torch.ByteTensor]=None, prefix_mask: Optional[torch.ByteTensor]=None, sequence_id: Optional[torch.LongTensor]=None, return_dict: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, use_cache: Optional[bool]=None, inputs_embeds: Optional[torch.Tensor]=None):
-        return_dict = return_dict if return_dict is not None else self.config.return_dict
-        use_cache = use_cache if use_cache is not None else self.config.use_cache
-        if attention_mask is not None:
-            attention_mask = attention_mask.bool()
-        if prefix_mask is not None:
-            prefix_mask = prefix_mask.bool()
-        if not return_dict:
-            raise NotImplementedError('return_dict False is not implemented yet for MPT')
-        if output_attentions:
-            if self.attn_impl != 'torch':
-                raise NotImplementedError('output_attentions is not implemented for MPT when using attn_impl `flash` or `triton`.')
-        if attention_mask is not None and attention_mask[:, 0].sum() != attention_mask.shape[0] and self.training:
-            raise NotImplementedError('MPT does not support training with left padding.')
-        if self.prefix_lm and prefix_mask is None:
-            raise ValueError('prefix_mask is a required argument when MPT is configured with prefix_lm=True.')
-        if self.training:
-            if self.attn_uses_sequence_id and sequence_id is None:
-                raise ValueError('sequence_id is a required argument when MPT is configured with attn_uses_sequence_id=True ' + 'and the model is in train mode.')
-            elif self.attn_uses_sequence_id is False and sequence_id is not None:
-                warnings.warn('MPT received non-None input for `sequence_id` but is configured with attn_uses_sequence_id=False. ' + 'This input will be ignored. If you want the model to use `sequence_id`, set attn_uses_sequence_id to True.')
-        if input_ids is not None:
-            S = input_ids.size(1)
-            assert S <= self.config.max_seq_len, f'Cannot forward input with seq_len={S}, this model only supports seq_len<={self.config.max_seq_len}'
-            tok_emb = self.wte(input_ids)
-        else:
-            assert inputs_embeds is not None
-            assert self.alibi, 'inputs_embeds is not implemented for MPT unless for alibi.'
-            S = inputs_embeds.size(1)
-            tok_emb = inputs_embeds
-        if self.alibi:
-            x = tok_emb
-        else:
-            past_position = 0
-            if past_key_values is not None:
-                if len(past_key_values) != self.config.n_layers:
-                    raise ValueError(f'past_key_values must provide a past_key_value for each attention ' + f'layer in the network (len(past_key_values)={len(past_key_values)!r}; self.config.n_layers={self.config.n_layers!r}).')
-                past_position = past_key_values[0][0].size(1)
-                if self.attn_impl == 'torch':
-                    past_position = past_key_values[0][0].size(3)
-            if S + past_position > self.config.max_seq_len:
-                raise ValueError(f'Cannot forward input with past sequence length {past_position} and current sequence length {S + 1}, this model only supports total sequence length <= {self.config.max_seq_len}.')
-            pos = torch.arange(past_position, S + past_position, dtype=torch.long, device=input_ids.device).unsqueeze(0)
-            if attention_mask is not None:
-                pos = torch.clamp(pos - torch.cumsum((~attention_mask).to(torch.int32), dim=1)[:, past_position:], min=0)
-            pos_emb = self.wpe(pos)
-            x = tok_emb + pos_emb
-        if self.embedding_fraction == 1:
-            x = self.emb_drop(x)
-        else:
-            x_shrunk = x * self.embedding_fraction + x.detach() * (1 - self.embedding_fraction)
-            assert isinstance(self.emb_drop, nn.Module)
-            x = self.emb_drop(x_shrunk)
-        (attn_bias, attention_mask) = self._attn_bias(device=x.device, dtype=torch.float32, attention_mask=attention_mask, prefix_mask=prefix_mask, sequence_id=sequence_id)
-        if use_cache and past_key_values is None:
-            past_key_values = [() for _ in range(self.config.n_layers)]
-        all_hidden_states = () if output_hidden_states else None
-        all_self_attns = () if output_attentions else None
-        for (b_idx, block) in enumerate(self.blocks):
-            if output_hidden_states:
-                assert all_hidden_states is not None
-                all_hidden_states = all_hidden_states + (x,)
-            past_key_value = past_key_values[b_idx] if past_key_values is not None else None
-            if self.gradient_checkpointing and self.training:
-                (x, attn_weights, past_key_value) = torch.utils.checkpoint.checkpoint(block, x, past_key_value, attn_bias, attention_mask, self.is_causal)
-            else:
-                (x, attn_weights, past_key_value) = block(x, past_key_value=past_key_value, attn_bias=attn_bias, attention_mask=attention_mask, is_causal=self.is_causal)
-            if past_key_values is not None:
-                past_key_values[b_idx] = past_key_value
-            if output_attentions:
-                assert all_self_attns is not None
-                all_self_attns = all_self_attns + (attn_weights,)
-        x = self.norm_f(x)
-        if output_hidden_states:
-            assert all_hidden_states is not None
-            all_hidden_states = all_hidden_states + (x,)
-        return BaseModelOutputWithPast(last_hidden_state=x, past_key_values=past_key_values, hidden_states=all_hidden_states, attentions=all_self_attns)
-
-    def param_init_fn(self, module):
-        init_fn_name = self.config.init_config['name']
-        MODEL_INIT_REGISTRY[init_fn_name](module=module, n_layers=self.config.n_layers, d_model=self.config.d_model, **self.config.init_config)
-
-    def fsdp_wrap_fn(self, module):
-        return isinstance(module, MPTBlock)
-
-    def activation_checkpointing_fn(self, module):
-        return isinstance(module, MPTBlock)
-
-class MPTForCausalLM(MPTPreTrainedModel):
-
-    def __init__(self, config: MPTConfig):
-        super().__init__(config)
-        if not config.tie_word_embeddings:
-            raise ValueError('MPTForCausalLM only supports tied word embeddings')
-        print(f'Instantiating an MPTForCausalLM model from {__file__}')
-        self.transformer = MPTModel(config)
-        for child in self.transformer.children():
-            if isinstance(child, torch.nn.ModuleList):
-                continue
-            if isinstance(child, torch.nn.Module):
-                child._fsdp_wrap = True
-        self.logit_scale = None
-        if config.logit_scale is not None:
-            logit_scale = config.logit_scale
-            if isinstance(logit_scale, str):
-                if logit_scale == 'inv_sqrt_d_model':
-                    logit_scale = 1 / math.sqrt(config.d_model)
-                else:
-                    raise ValueError(f"logit_scale={logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'.")
-            self.logit_scale = logit_scale
-
-    def get_input_embeddings(self):
-        return self.transformer.wte
-
-    def set_input_embeddings(self, value):
-        self.transformer.wte = value
-
-    def get_output_embeddings(self):
-        return self.transformer.wte
-
-    def set_output_embeddings(self, new_embeddings):
-        self.transformer.wte = new_embeddings
-
-    def set_decoder(self, decoder):
-        self.transformer = decoder
-
-    def get_decoder(self):
-        return self.transformer
-
-    def forward(self, input_ids: torch.LongTensor, past_key_values: Optional[List[Tuple[torch.FloatTensor]]]=None, attention_mask: Optional[torch.ByteTensor]=None, prefix_mask: Optional[torch.ByteTensor]=None, sequence_id: Optional[torch.LongTensor]=None, labels: Optional[torch.LongTensor]=None, return_dict: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, use_cache: Optional[bool]=None, inputs_embeds: Optional[torch.FloatTensor]=None):
-        return_dict = return_dict if return_dict is not None else self.config.return_dict
-        use_cache = use_cache if use_cache is not None else self.config.use_cache
-        if inputs_embeds is not None:
-            raise NotImplementedError('inputs_embeds has to be None (for hf/peft support).')
-        outputs = self.transformer(input_ids=input_ids, past_key_values=past_key_values, attention_mask=attention_mask, prefix_mask=prefix_mask, sequence_id=sequence_id, return_dict=return_dict, output_attentions=output_attentions, output_hidden_states=output_hidden_states, use_cache=use_cache)
-        logits = self.transformer.wte(outputs.last_hidden_state.to(self.transformer.wte.weight.device), True)
-        if self.logit_scale is not None:
-            if self.logit_scale == 0:
-                warnings.warn(f'Multiplying logits by self.logit_scale={self.logit_scale!r}. This will produce uniform (uninformative) outputs.')
-            logits *= self.logit_scale
-        loss = None
-        if labels is not None:
-            labels = torch.roll(labels, shifts=-1)
-            labels[:, -1] = -100
-            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.to(logits.device).view(-1))
-        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions)
-
-    def param_init_fn(self, module):
-        init_fn_name = self.config.init_config['name']
-        MODEL_INIT_REGISTRY[init_fn_name](module=module, n_layers=self.config.n_layers, d_model=self.config.d_model, **self.config.init_config)
-
-    def fsdp_wrap_fn(self, module):
-        return isinstance(module, MPTBlock)
-
-    def activation_checkpointing_fn(self, module):
-        return isinstance(module, MPTBlock)
-
-    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
-        if inputs_embeds is not None:
-            raise NotImplementedError('inputs_embeds is not implemented for MPT yet')
-        attention_mask = kwargs['attention_mask'].bool()
-        if attention_mask[:, -1].sum() != attention_mask.shape[0]:
-            raise NotImplementedError('MPT does not support generation with right padding.')
-        if self.transformer.attn_uses_sequence_id and self.training:
-            sequence_id = torch.zeros_like(input_ids[:1])
-        else:
-            sequence_id = None
-        if past_key_values is not None:
-            input_ids = input_ids[:, -1].unsqueeze(-1)
-        if self.transformer.prefix_lm:
-            prefix_mask = torch.ones_like(attention_mask)
-            if kwargs.get('use_cache') == False:
-                raise NotImplementedError('MPT with prefix_lm=True does not support use_cache=False.')
-        else:
-            prefix_mask = None
-        return {'input_ids': input_ids, 'attention_mask': attention_mask, 'prefix_mask': prefix_mask, 'sequence_id': sequence_id, 'past_key_values': past_key_values, 'use_cache': kwargs.get('use_cache', True)}
-
-    @staticmethod
-    def _reorder_cache(past_key_values, beam_idx):
-        """Used by HuggingFace generate when using beam search with kv-caching.
-
-        See https://github.com/huggingface/transformers/blob/3ec7a47664ebe40c40f4b722f6bb1cd30c3821ec/src/transformers/models/gpt2/modeling_gpt2.py#L1122-L1133
-        for an example in transformers.
-        """
-        reordered_past = []
-        for layer_past in past_key_values:
-            reordered_past += [tuple((past_state.index_select(0, beam_idx) for past_state in layer_past))]
-        return reordered_past
\ No newline at end of file
diff --git a/instructany2pix/llm/model/language_model/mpt/norm.py b/instructany2pix/llm/model/language_model/mpt/norm.py
deleted file mode 100644
index 067b6140fae546e5cb49cb2b1e4e6af660ced60d..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/language_model/mpt/norm.py
+++ /dev/null
@@ -1,56 +0,0 @@
-import torch
-
-def _cast_if_autocast_enabled(tensor):
-    if torch.is_autocast_enabled():
-        if tensor.device.type == 'cuda':
-            dtype = torch.get_autocast_gpu_dtype()
-        elif tensor.device.type == 'cpu':
-            dtype = torch.get_autocast_cpu_dtype()
-        else:
-            raise NotImplementedError()
-        return tensor.to(dtype=dtype)
-    return tensor
-
-class LPLayerNorm(torch.nn.LayerNorm):
-
-    def __init__(self, normalized_shape, eps=1e-05, elementwise_affine=True, device=None, dtype=None):
-        super().__init__(normalized_shape=normalized_shape, eps=eps, elementwise_affine=elementwise_affine, device=device, dtype=dtype)
-
-    def forward(self, x):
-        module_device = x.device
-        downcast_x = _cast_if_autocast_enabled(x)
-        downcast_weight = _cast_if_autocast_enabled(self.weight) if self.weight is not None else self.weight
-        downcast_bias = _cast_if_autocast_enabled(self.bias) if self.bias is not None else self.bias
-        with torch.autocast(enabled=False, device_type=module_device.type):
-            return torch.nn.functional.layer_norm(downcast_x, self.normalized_shape, downcast_weight, downcast_bias, self.eps)
-
-def rms_norm(x, weight=None, eps=1e-05):
-    output = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
-    if weight is not None:
-        return output * weight
-    return output
-
-class RMSNorm(torch.nn.Module):
-
-    def __init__(self, normalized_shape, eps=1e-05, weight=True, dtype=None, device=None):
-        super().__init__()
-        self.eps = eps
-        if weight:
-            self.weight = torch.nn.Parameter(torch.ones(normalized_shape, dtype=dtype, device=device))
-        else:
-            self.register_parameter('weight', None)
-
-    def forward(self, x):
-        return rms_norm(x.float(), self.weight, self.eps).to(dtype=x.dtype)
-
-class LPRMSNorm(RMSNorm):
-
-    def __init__(self, normalized_shape, eps=1e-05, weight=True, dtype=None, device=None):
-        super().__init__(normalized_shape=normalized_shape, eps=eps, weight=weight, dtype=dtype, device=device)
-
-    def forward(self, x):
-        downcast_x = _cast_if_autocast_enabled(x)
-        downcast_weight = _cast_if_autocast_enabled(self.weight) if self.weight is not None else self.weight
-        with torch.autocast(enabled=False, device_type=x.device.type):
-            return rms_norm(downcast_x, downcast_weight, self.eps).to(dtype=x.dtype)
-NORM_CLASS_REGISTRY = {'layernorm': torch.nn.LayerNorm, 'low_precision_layernorm': LPLayerNorm, 'rmsnorm': RMSNorm, 'low_precision_rmsnorm': LPRMSNorm}
\ No newline at end of file
diff --git a/instructany2pix/llm/model/language_model/mpt/param_init_fns.py b/instructany2pix/llm/model/language_model/mpt/param_init_fns.py
deleted file mode 100644
index 418b83ca2363288046f4b48b1d706c5607341fb5..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/language_model/mpt/param_init_fns.py
+++ /dev/null
@@ -1,181 +0,0 @@
-import math
-import warnings
-from collections.abc import Sequence
-from functools import partial
-from typing import Optional, Tuple, Union
-import torch
-from torch import nn
-from .norm import NORM_CLASS_REGISTRY
-
-def torch_default_param_init_fn_(module: nn.Module, verbose: int=0, **kwargs):
-    del kwargs
-    if verbose > 1:
-        warnings.warn(f"Initializing network using module's reset_parameters attribute")
-    if hasattr(module, 'reset_parameters'):
-        module.reset_parameters()
-
-def fused_init_helper_(module: nn.Module, init_fn_):
-    _fused = getattr(module, '_fused', None)
-    if _fused is None:
-        raise RuntimeError(f'Internal logic error')
-    (dim, splits) = _fused
-    splits = (0, *splits, module.weight.size(dim))
-    for (s, e) in zip(splits[:-1], splits[1:]):
-        slice_indices = [slice(None)] * module.weight.ndim
-        slice_indices[dim] = slice(s, e)
-        init_fn_(module.weight[slice_indices])
-
-def generic_param_init_fn_(module: nn.Module, init_fn_, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, verbose: int=0, **kwargs):
-    del kwargs
-    if verbose > 1:
-        warnings.warn(f'If model has bias parameters they are initialized to 0.')
-    init_div_is_residual = init_div_is_residual
-    if init_div_is_residual is False:
-        div_is_residual = 1.0
-    elif init_div_is_residual is True:
-        div_is_residual = math.sqrt(2 * n_layers)
-    elif isinstance(init_div_is_residual, float) or isinstance(init_div_is_residual, int):
-        div_is_residual = init_div_is_residual
-    elif isinstance(init_div_is_residual, str) and init_div_is_residual.isnumeric():
-        div_is_residual = float(init_div_is_residual)
-    else:
-        div_is_residual = 1.0
-        raise ValueError(f'Expected init_div_is_residual to be boolean or numeric, got {init_div_is_residual}')
-    if init_div_is_residual is not False:
-        if verbose > 1:
-            warnings.warn(f'Initializing _is_residual layers then dividing them by {div_is_residual:.3f}. ' + f'Set `init_div_is_residual: false` in init config to disable this.')
-    if isinstance(module, nn.Linear):
-        if hasattr(module, '_fused'):
-            fused_init_helper_(module, init_fn_)
-        else:
-            init_fn_(module.weight)
-        if module.bias is not None:
-            torch.nn.init.zeros_(module.bias)
-        if init_div_is_residual is not False and getattr(module, '_is_residual', False):
-            with torch.no_grad():
-                module.weight.div_(div_is_residual)
-    elif isinstance(module, nn.Embedding):
-        if emb_init_std is not None:
-            std = emb_init_std
-            if std == 0:
-                warnings.warn(f'Embedding layer initialized to 0.')
-            emb_init_fn_ = partial(torch.nn.init.normal_, mean=0.0, std=std)
-            if verbose > 1:
-                warnings.warn(f'Embedding layer initialized using normal distribution with mean=0 and std={std!r}.')
-        elif emb_init_uniform_lim is not None:
-            lim = emb_init_uniform_lim
-            if isinstance(lim, Sequence):
-                if len(lim) > 2:
-                    raise ValueError(f'Uniform init requires a min and a max limit. User input: {lim}.')
-                if lim[0] == lim[1]:
-                    warnings.warn(f'Embedding layer initialized to {lim[0]}.')
-            else:
-                if lim == 0:
-                    warnings.warn(f'Embedding layer initialized to 0.')
-                lim = [-lim, lim]
-            (a, b) = lim
-            emb_init_fn_ = partial(torch.nn.init.uniform_, a=a, b=b)
-            if verbose > 1:
-                warnings.warn(f'Embedding layer initialized using uniform distribution in range {lim}.')
-        else:
-            emb_init_fn_ = init_fn_
-        emb_init_fn_(module.weight)
-    elif isinstance(module, tuple(set(NORM_CLASS_REGISTRY.values()))):
-        if verbose > 1:
-            warnings.warn(f'Norm weights are set to 1. If norm layer has a bias it is initialized to 0.')
-        if hasattr(module, 'weight') and module.weight is not None:
-            torch.nn.init.ones_(module.weight)
-        if hasattr(module, 'bias') and module.bias is not None:
-            torch.nn.init.zeros_(module.bias)
-    elif isinstance(module, nn.MultiheadAttention):
-        if module._qkv_same_embed_dim:
-            assert module.in_proj_weight is not None
-            assert module.q_proj_weight is None and module.k_proj_weight is None and (module.v_proj_weight is None)
-            assert d_model is not None
-            _d = d_model
-            splits = (0, _d, 2 * _d, 3 * _d)
-            for (s, e) in zip(splits[:-1], splits[1:]):
-                init_fn_(module.in_proj_weight[s:e])
-        else:
-            assert module.q_proj_weight is not None and module.k_proj_weight is not None and (module.v_proj_weight is not None)
-            assert module.in_proj_weight is None
-            init_fn_(module.q_proj_weight)
-            init_fn_(module.k_proj_weight)
-            init_fn_(module.v_proj_weight)
-        if module.in_proj_bias is not None:
-            torch.nn.init.zeros_(module.in_proj_bias)
-        if module.bias_k is not None:
-            torch.nn.init.zeros_(module.bias_k)
-        if module.bias_v is not None:
-            torch.nn.init.zeros_(module.bias_v)
-        init_fn_(module.out_proj.weight)
-        if init_div_is_residual is not False and getattr(module.out_proj, '_is_residual', False):
-            with torch.no_grad():
-                module.out_proj.weight.div_(div_is_residual)
-        if module.out_proj.bias is not None:
-            torch.nn.init.zeros_(module.out_proj.bias)
-    else:
-        for _ in module.parameters(recurse=False):
-            raise NotImplementedError(f'{module.__class__.__name__} parameters are not initialized by param_init_fn.')
-
-def _normal_init_(std, mean=0.0):
-    return partial(torch.nn.init.normal_, mean=mean, std=std)
-
-def _normal_param_init_fn_(module: nn.Module, std: float, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, verbose: int=0, **kwargs):
-    del kwargs
-    init_fn_ = _normal_init_(std=std)
-    if verbose > 1:
-        warnings.warn(f'Using torch.nn.init.normal_ init fn mean=0.0, std={std}')
-    generic_param_init_fn_(module=module, init_fn_=init_fn_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
-
-def baseline_param_init_fn_(module: nn.Module, init_std: float, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, verbose: int=0, **kwargs):
-    del kwargs
-    if init_std is None:
-        raise ValueError("You must set model.init_config['init_std'] to a float value to use the default initialization scheme.")
-    _normal_param_init_fn_(module=module, std=init_std, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
-
-def small_param_init_fn_(module: nn.Module, n_layers: int, d_model: int, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, verbose: int=0, **kwargs):
-    del kwargs
-    std = math.sqrt(2 / (5 * d_model))
-    _normal_param_init_fn_(module=module, std=std, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
-
-def neox_param_init_fn_(module: nn.Module, n_layers: int, d_model: int, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, verbose: int=0, **kwargs):
-    """From section 2.3.1 of GPT-NeoX-20B:
-
-    An Open-Source AutoregressiveLanguage Model — Black et. al. (2022)
-    see https://github.com/EleutherAI/gpt-neox/blob/9610391ab319403cef079b438edd016a2443af54/megatron/model/init_functions.py#L151
-    and https://github.com/EleutherAI/gpt-neox/blob/main/megatron/model/transformer.py
-    """
-    del kwargs
-    residual_div = n_layers / math.sqrt(10)
-    if verbose > 1:
-        warnings.warn(f'setting init_div_is_residual to {residual_div}')
-    small_param_init_fn_(module=module, d_model=d_model, n_layers=n_layers, init_div_is_residual=residual_div, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
-
-def kaiming_uniform_param_init_fn_(module: nn.Module, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, init_gain: float=0, fan_mode: str='fan_in', init_nonlinearity: str='leaky_relu', verbose: int=0, **kwargs):
-    del kwargs
-    if verbose > 1:
-        warnings.warn(f'Using nn.init.kaiming_uniform_ init fn with parameters: ' + f'a={init_gain}, mode={fan_mode}, nonlinearity={init_nonlinearity}')
-    kaiming_uniform_ = partial(nn.init.kaiming_uniform_, a=init_gain, mode=fan_mode, nonlinearity=init_nonlinearity)
-    generic_param_init_fn_(module=module, init_fn_=kaiming_uniform_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
-
-def kaiming_normal_param_init_fn_(module: nn.Module, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, init_gain: float=0, fan_mode: str='fan_in', init_nonlinearity: str='leaky_relu', verbose: int=0, **kwargs):
-    del kwargs
-    if verbose > 1:
-        warnings.warn(f'Using nn.init.kaiming_normal_ init fn with parameters: ' + f'a={init_gain}, mode={fan_mode}, nonlinearity={init_nonlinearity}')
-    kaiming_normal_ = partial(torch.nn.init.kaiming_normal_, a=init_gain, mode=fan_mode, nonlinearity=init_nonlinearity)
-    generic_param_init_fn_(module=module, init_fn_=kaiming_normal_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
-
-def xavier_uniform_param_init_fn_(module: nn.Module, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, init_gain: float=0, verbose: int=0, **kwargs):
-    del kwargs
-    xavier_uniform_ = partial(torch.nn.init.xavier_uniform_, gain=init_gain)
-    if verbose > 1:
-        warnings.warn(f'Using torch.nn.init.xavier_uniform_ init fn with parameters: ' + f'gain={init_gain}')
-    generic_param_init_fn_(module=module, init_fn_=xavier_uniform_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
-
-def xavier_normal_param_init_fn_(module: nn.Module, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, init_gain: float=0, verbose: int=0, **kwargs):
-    xavier_normal_ = partial(torch.nn.init.xavier_normal_, gain=init_gain)
-    if verbose > 1:
-        warnings.warn(f'Using torch.nn.init.xavier_normal_ init fn with parameters: ' + f'gain={init_gain}')
-    generic_param_init_fn_(module=module, init_fn_=xavier_normal_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
-MODEL_INIT_REGISTRY = {'default_': torch_default_param_init_fn_, 'baseline_': baseline_param_init_fn_, 'kaiming_uniform_': kaiming_uniform_param_init_fn_, 'kaiming_normal_': kaiming_normal_param_init_fn_, 'neox_init_': neox_param_init_fn_, 'small_init_': small_param_init_fn_, 'xavier_uniform_': xavier_uniform_param_init_fn_, 'xavier_normal_': xavier_normal_param_init_fn_}
\ No newline at end of file
diff --git a/instructany2pix/llm/model/make_delta.py b/instructany2pix/llm/model/make_delta.py
deleted file mode 100644
index 05647d79eb91ff06fe633587749f590ec3984b09..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/make_delta.py
+++ /dev/null
@@ -1,48 +0,0 @@
-import argparse
-
-import torch
-from tqdm import tqdm
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from .utils import auto_upgrade
-
-
-def make_delta(base_model_path, target_model_path, delta_path, hub_repo_id):
-    print("Loading base model")
-    base = AutoModelForCausalLM.from_pretrained(
-        base_model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True)
-
-    print("Loading target model")
-    auto_upgrade(target_model_path)
-    target = AutoModelForCausalLM.from_pretrained(target_model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True)
-
-    print("Calculating delta")
-    for name, param in tqdm(target.state_dict().items(), desc="Calculating delta"):
-        if name not in base.state_dict():
-            assert name in ['model.mm_projector.weight', 'model.mm_projector.bias'], f'{name} not in base model'
-            continue
-        if param.data.shape == base.state_dict()[name].shape:
-            param.data -= base.state_dict()[name]
-        else:
-            assert name in ['model.embed_tokens.weight', 'lm_head.weight'], f'{name} dimension mismatch: {param.data.shape} vs {base.state_dict()[name].shape}'
-            bparam = base.state_dict()[name]
-            param.data[:bparam.shape[0], :bparam.shape[1]] -= bparam
-
-    print("Saving delta")
-    if hub_repo_id:
-        kwargs = {"push_to_hub": True, "repo_id": hub_repo_id}
-    else:
-        kwargs = {}
-    target.save_pretrained(delta_path, **kwargs)
-    target_tokenizer = AutoTokenizer.from_pretrained(target_model_path)
-    target_tokenizer.save_pretrained(delta_path, **kwargs)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--base-model-path", type=str, required=True)
-    parser.add_argument("--target-model-path", type=str, required=True)
-    parser.add_argument("--delta-path", type=str, required=True)
-    parser.add_argument("--hub-repo-id", type=str, default=None)
-    args = parser.parse_args()
-
-    make_delta(args.base_model_path, args.target_model_path, args.delta_path, args.hub_repo_id)
diff --git a/instructany2pix/llm/model/multimodal_encoder/builder.py b/instructany2pix/llm/model/multimodal_encoder/builder.py
deleted file mode 100644
index 30cb6e05d81a64822f910858f64a1e009625cde8..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/multimodal_encoder/builder.py
+++ /dev/null
@@ -1,12 +0,0 @@
-import os
-from .clip_encoder import CLIPVisionTower
-from .image_bind_encoder import LanguageBindVisionTower
-
-def build_vision_tower(vision_tower_cfg, **kwargs):
-    vision_tower = getattr(vision_tower_cfg, 'mm_vision_tower', getattr(vision_tower_cfg, 'vision_tower', None))
-    is_absolute_path_exists = os.path.exists(vision_tower)
-    if is_absolute_path_exists or vision_tower.startswith("openai") or vision_tower.startswith("laion"):
-        return CLIPVisionTower(vision_tower, args=vision_tower_cfg, **kwargs)
-    elif vision_tower == 'languagebind':
-        return LanguageBindVisionTower(vision_tower,args=vision_tower_cfg, **kwargs)
-    raise ValueError(f'Unknown vision tower: {vision_tower}')
diff --git a/instructany2pix/llm/model/multimodal_encoder/clip_encoder.py b/instructany2pix/llm/model/multimodal_encoder/clip_encoder.py
deleted file mode 100644
index c016d8fc506e6663103c39bf33e0dc5ff7926a8d..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/multimodal_encoder/clip_encoder.py
+++ /dev/null
@@ -1,81 +0,0 @@
-import torch
-import torch.nn as nn
-
-from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor, CLIPVisionConfig
-
-
-class CLIPVisionTower(nn.Module):
-    def __init__(self, vision_tower, args, delay_load=False):
-        super().__init__()
-
-        self.is_loaded = False
-
-        self.vision_tower_name = vision_tower
-        self.select_layer = args.mm_vision_select_layer
-        self.select_feature = getattr(args, 'mm_vision_select_feature', 'projection')
-
-        if not delay_load:
-            self.load_model()
-        else:
-            self.cfg_only = CLIPVisionConfig.from_pretrained(self.vision_tower_name)
-
-    def load_model(self):
-        self.image_processor = CLIPImageProcessor.from_pretrained(self.vision_tower_name)
-        self.vision_tower = CLIPVisionModelWithProjection.from_pretrained(self.vision_tower_name)
-        self.vision_tower.requires_grad_(False)
-
-        self.is_loaded = True
-
-    def feature_select(self, image_forward_outs):
-        image_features = image_forward_outs.hidden_states[self.select_layer]
-        if self.select_feature == 'patch':
-            image_features = image_features[:, 1:]
-        elif self.select_feature == 'cls_patch':
-            image_features = image_features
-        elif self.select_feature == 'projection':
-            image_features = image_forward_outs.image_embeds.unsqueeze(1)
-        else:
-            raise ValueError(f'Unexpected select feature: {self.select_feature}')
-        return image_features
-
-    @torch.no_grad()
-    def forward(self, images):
-        if type(images) is list:
-            image_features = []
-            for image in images:
-                image_forward_out = self.vision_tower(image.to(device=self.device, dtype=self.dtype).unsqueeze(0), output_hidden_states=True)
-                image_feature = self.feature_select(image_forward_out).to(image.dtype)
-                image_features.append(image_feature)
-        else:
-            return torch.zeros(images['image']['pixel_values'].shape[0],1,1024).to(images['image']['pixel_values'].dtype).to(device=self.device)
-            image_forward_outs = self.vision_tower(images['image']['pixel_values'].to(device=self.device), output_hidden_states=True)
-            image_features = self.feature_select(image_forward_outs).to(images['image']['pixel_values'].dtype)
-
-        return image_features
-
-    @property
-    def dummy_feature(self):
-        return torch.zeros(1, self.hidden_size, device=self.device, dtype=self.dtype)
-
-    @property
-    def dtype(self):
-        return self.vision_tower.dtype
-
-    @property
-    def device(self):
-        return self.vision_tower.device
-
-    @property
-    def config(self):
-        if self.is_loaded:
-            return self.vision_tower.config
-        else:
-            return self.cfg_only
-
-    @property
-    def hidden_size(self):
-        return self.config.projection_dim if self.select_feature == 'projection' else self.config.hidden_size
-
-    @property
-    def num_patches(self):
-        return (self.config.image_size // self.config.patch_size) ** 2
diff --git a/instructany2pix/llm/model/multimodal_encoder/image_bind_encoder.py b/instructany2pix/llm/model/multimodal_encoder/image_bind_encoder.py
deleted file mode 100644
index 3cf71ccf3d4934f8535efdc120a38c2c4b352a6b..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/multimodal_encoder/image_bind_encoder.py
+++ /dev/null
@@ -1,93 +0,0 @@
-from typing import Any
-import torch
-import torch.nn as nn
-
-from transformers import CLIPVisionModel, CLIPImageProcessor, CLIPVisionConfig
-# from imagebind.models import imagebind_model
-try:
-    from languagebind import LanguageBind, to_device, transform_dict, LanguageBindImageTokenizer
-except:
-    pass
-
-class LanguageBindProcesser:
-
-    def __init__(self,modality_transform) -> None:
-        self.modality_transform = modality_transform
-        self.crop_size = dict(height=224,width=224)
-
-    def __call__(self, inputs) -> Any:
-        outputs = {}
-        for modality,data in inputs.items():
-            outputs[modality] = self.modality_transform[modality](data)
-        return outputs
-    
-class LanguageBindVisionTower(nn.Module):
-    def __init__(self, vision_tower, args,clip_type=('image',),delay_load=False):
-        super().__init__()
-        # 'audio'
-        self.is_loaded = False
-
-        self.vision_tower_name = vision_tower
-        self.select_layer = args.mm_vision_select_layer
-        self.select_feature = getattr(args, 'mm_vision_select_feature', 'patch')
-        self.clip_type = clip_type
-        self.hidden_size = 768
-        if not delay_load:
-            self.load_model()
-        else:
-            pass
-            #self.cfg_only = CLIPVisionConfig.from_pretrained(self.vision_tower_name)
-
-    def load_model(self):
-        model = LanguageBind(clip_type=self.clip_type, cache_dir='src/LanguageBind/cache_dir')
-        self.modality_transform = {c: transform_dict[c](model.modality_config[c]) for c in self.clip_type}
-        self.image_processor = LanguageBindProcesser(self.modality_transform)
-        self.vision_tower = model
-        self.vision_tower.requires_grad_(False)
-        #self.hidden_size = 768
-        self.is_loaded = True
-
-    def feature_select(self, image_forward_outs):
-        image_features = image_forward_outs.hidden_states[self.select_layer]
-        if self.select_feature == 'patch':
-            image_features = image_features[:, 1:]
-        elif self.select_feature == 'cls_patch':
-            image_features = image_features
-        else:
-            raise ValueError(f'Unexpected select feature: {self.select_feature}')
-        return image_features
-
-    @torch.no_grad()
-    def forward(self, inputs):
-        embeddings = self.vision_tower({k:v for k,v in inputs.items() if k != 'info'})
-        embeddings = torch.stack(
-            [embeddings[info['modality']][info['idx']] for info in inputs['info']]
-        )
-        return embeddings
-
-    @property
-    def dummy_feature(self):
-        return torch.zeros(1, self.hidden_size, device=self.device, dtype=self.dtype)
-     
-    @property
-    def dtype(self):
-        return self.vision_tower.dtype
-
-    @property
-    def device(self):
-        return self.vision_tower.device
-
-    @property
-    def config(self):
-        if self.is_loaded:
-            return self.vision_tower.config
-        else:
-            return self.cfg_only
-
-    # @property
-    # def hidden_size(self):
-    #     return self.config.hidden_size
-
-    @property
-    def num_patches(self):
-        return (self.config.image_size // self.config.patch_size) ** 2
diff --git a/instructany2pix/llm/model/multimodal_projector/builder.py b/instructany2pix/llm/model/multimodal_projector/builder.py
deleted file mode 100644
index c6d4f09c62d929e5aced85e1740d0792d0f396e7..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/multimodal_projector/builder.py
+++ /dev/null
@@ -1,75 +0,0 @@
-import torch
-import torch.nn as nn
-import re
-
-
-class IdentityMap(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    def forward(self, x, *args, **kwargs):
-        return x
-
-    @property
-    def config(self):
-        return {"mm_projector_type": 'identity'}
-
-
-class SimpleResBlock(nn.Module):
-    def __init__(self, channels):
-        super().__init__()
-        self.pre_norm = nn.LayerNorm(channels)
-
-        self.proj = nn.Sequential(
-            nn.Linear(channels, channels),
-            nn.GELU(),
-            nn.Linear(channels, channels)
-        )
-    def forward(self, x):
-        x = self.pre_norm(x)
-        return x + self.proj(x)
-
-
-def build_vision_projector(config, delay_load=False,input_size=None, **kwargs):
-    projector_type = getattr(config, 'mm_projector_type', 'linear')
-    if not input_size:
-        input_size = config.mm_hidden_size
-    if projector_type == 'linear':
-        return nn.Linear(input_size, config.hidden_size)
-
-    mlp_gelu_match = re.match(r'^mlp(\d+)x_gelu$', projector_type)
-    if mlp_gelu_match:
-        mlp_depth = int(mlp_gelu_match.group(1))
-        modules = [nn.Linear(input_size, config.hidden_size)]
-        for _ in range(1, mlp_depth):
-            modules.append(nn.GELU())
-            modules.append(nn.Linear(config.hidden_size, config.hidden_size))
-        return nn.Sequential(*modules)
-
-    if projector_type == 'identity':
-        return IdentityMap()
-
-    raise ValueError(f'Unknown projector type: {projector_type}')
-
-def build_vision_predictor(config, output_size=None, **kwargs):
-    projector_type = getattr(config, 'mm_projector_type', 'linear')
-    if not output_size:
-        output_size = config.mm_hidden_size
-    if projector_type == 'linear':
-        return nn.Linear(config.hidden_size,output_size)
-
-    mlp_gelu_match = re.match(r'^mlp(\d+)x_gelu$', projector_type)
-    if mlp_gelu_match:
-        mlp_depth = int(mlp_gelu_match.group(1))
-        modules = []
-        for _ in range(1, mlp_depth):
-            modules.append(nn.Linear(config.hidden_size, config.hidden_size))
-            modules.append(nn.GELU())
-        modules.append(nn.Linear(config.hidden_size,output_size))
-        return nn.Sequential(*modules)
-
-    if projector_type == 'identity':
-        return IdentityMap()
-
-    raise ValueError(f'Unknown projector type: {projector_type}')
-
diff --git a/instructany2pix/llm/model/utils.py b/instructany2pix/llm/model/utils.py
deleted file mode 100644
index 7f55c80d382c4e392ba747c2b6867d066b755962..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/utils.py
+++ /dev/null
@@ -1,5 +0,0 @@
-from transformers import AutoConfig
-
-
-def auto_upgrade(config):
-    pass
\ No newline at end of file
diff --git a/instructany2pix/llm/model/vae/audio_vqvae.py b/instructany2pix/llm/model/vae/audio_vqvae.py
deleted file mode 100644
index 15b4292ccf3b08b0d7b3fb8ea942614a92acfc61..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/vae/audio_vqvae.py
+++ /dev/null
@@ -1,394 +0,0 @@
-from typing import Any
-import torch
-import torch.nn.functional as F
-import pytorch_lightning as pl
-
-import sys
-sys.path.insert(0, '.')  # nopep8
-#from train import instantiate_from_config
-
-from .modules.blocks_spec import Encoder, Decoder, Encoder1d, Decoder1d
-from .modules.quantize_spec import VectorQuantizer, VectorQuantizer1d
-# from .modules.quantize_spec import VectorQuantizer2 as VectorQuantizer
-from .modules.util import instantiate_from_config
-
-class VQModel(pl.LightningModule):
-    def __init__(self,
-                 ddconfig,
-                 lossconfig,
-                 n_embed,
-                 embed_dim,
-                 ckpt_path=None,
-                 ignore_keys=[],
-                 image_key="image",
-                 colorize_nlabels=None,
-                 monitor=None
-                 ):
-        super().__init__()
-        self.image_key = image_key
-        # we need this one for compatibility in train.ImageLogger.log_img if statement
-        self.first_stage_key = image_key
-        self.encoder = Encoder(**ddconfig)
-        self.decoder = Decoder(**ddconfig)
-        self.loss = instantiate_from_config(lossconfig)
-        self.n_embed = n_embed
-        self.quantize = VectorQuantizer(n_embed, embed_dim, beta=0.25)
-        self.quant_conv = torch.nn.Conv2d(ddconfig["z_channels"], embed_dim, 1)
-        self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
-        if ckpt_path is not None:
-            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
-        if colorize_nlabels is not None:
-            assert type(colorize_nlabels)==int
-            self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
-        if monitor is not None:
-            self.monitor = monitor
-        self.used_codes = []
-        self.counts = [0 for _ in range(self.quantize.n_e)]
-
-    def init_from_ckpt(self, path, ignore_keys=list()):
-        sd = torch.load(path, map_location="cpu")["state_dict"]
-        keys = list(sd.keys())
-        for k in keys:
-            for ik in ignore_keys:
-                if k.startswith(ik):
-                    print("Deleting key {} from state_dict.".format(k))
-                    del sd[k]
-        self.load_state_dict(sd, strict=False)
-        print(f"Restored from {path}")
-
-    def encode(self, x):
-        h = self.encoder(x)  # 2d: (B, 256, 16, 16) <- (B, 3, 256, 256)
-        h = self.quant_conv(h)  # 2d: (B, 256, 16, 16)
-        quant, emb_loss, info = self.quantize(h)  # (B, 256, 16, 16), (), ((), (768, 1024), (768, 1))
-        if not self.training:
-            self.counts = [info[2].squeeze().tolist().count(i) + self.counts[i] for i in range(self.quantize.n_e)]
-        return quant, emb_loss, info
-
-    def decode(self, quant):
-        quant = self.post_quant_conv(quant)
-        dec = self.decoder(quant)
-        return dec
-
-    def decode_code(self, code_b):
-        quant_b = self.quantize.embed_code(code_b)
-        dec = self.decode(quant_b)
-        return dec
-
-    def forward(self, input):
-        quant, diff, _ = self.encode(input)
-        dec = self.decode(quant)
-        return dec, diff
-
-    def get_input(self, batch, k):
-        x = batch[k]
-        if len(x.shape) == 3:
-            x = x[..., None]
-        x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format)
-        return x.float()
-
-    def training_step(self, batch, batch_idx, optimizer_idx):
-        x = self.get_input(batch, self.image_key)
-        xrec, qloss = self(x)
-
-        if optimizer_idx == 0:
-            # autoencode
-            aeloss, log_dict_ae = self.loss(qloss, x, xrec, optimizer_idx, self.global_step,
-                                            last_layer=self.get_last_layer(), split="train")
-
-            self.log("train/aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
-            self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True)
-            return aeloss
-
-        if optimizer_idx == 1:
-            # discriminator
-            discloss, log_dict_disc = self.loss(qloss, x, xrec, optimizer_idx, self.global_step,
-                                            last_layer=self.get_last_layer(), split="train")
-            self.log("train/disc_loss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
-            self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True)
-            return discloss
-
-    def validation_step(self, batch, batch_idx):
-        if batch_idx == 0 and self.global_step != 0 and sum(self.counts) > 0:
-            print(f'Previous Epoch counts: {self.counts}')
-            zero_hit_codes = len([1 for count in self.counts if count == 0])
-            used_codes = []
-            for c, count in enumerate(self.counts):
-                used_codes.extend([c] * count)
-            self.logger.experiment.add_histogram('val/code_hits', torch.tensor(used_codes), self.global_step)
-            self.logger.experiment.add_scalar('val/zero_hit_codes', zero_hit_codes, self.global_step)
-            self.counts = [0 for _ in range(self.quantize.n_e)]
-        x = self.get_input(batch, self.image_key)
-        xrec, qloss = self(x)
-        aeloss, log_dict_ae = self.loss(qloss, x, xrec, 0, self.global_step,
-                                        last_layer=self.get_last_layer(), split="val")
-
-        discloss, log_dict_disc = self.loss(qloss, x, xrec, 1, self.global_step,
-                                            last_layer=self.get_last_layer(), split="val")
-        rec_loss = log_dict_ae['val/rec_loss']
-        self.log('val/rec_loss', rec_loss, prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
-        self.log('val/aeloss', aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
-        self.log_dict(log_dict_ae)
-        self.log_dict(log_dict_disc)
-        return self.log_dict
-
-    def configure_optimizers(self):
-        lr = self.learning_rate
-        opt_ae = torch.optim.Adam(list(self.encoder.parameters()) +
-                                  list(self.decoder.parameters()) +
-                                  list(self.quantize.parameters()) +
-                                  list(self.quant_conv.parameters()) +
-                                  list(self.post_quant_conv.parameters()),
-                                  lr=lr, betas=(0.5, 0.9))
-        opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
-                                    lr=lr, betas=(0.5, 0.9))
-        return [opt_ae, opt_disc], []
-
-    def get_last_layer(self):
-        return self.decoder.conv_out.weight
-
-    def log_images(self, batch, **kwargs):
-        log = dict()
-        x = self.get_input(batch, self.image_key)
-        x = x.to(self.device)
-        xrec, _ = self(x)
-        if x.shape[1] > 3:
-            # colorize with random projection
-            assert xrec.shape[1] > 3
-            x = self.to_rgb(x)
-            xrec = self.to_rgb(xrec)
-        log["inputs"] = x
-        log["reconstructions"] = xrec
-        return log
-
-    def to_rgb(self, x):
-        assert self.image_key == "segmentation"
-        if not hasattr(self, "colorize"):
-            self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x))
-        x = F.conv2d(x, weight=self.colorize)
-        x = 2.*(x-x.min())/(x.max()-x.min()) - 1.
-        return x
-
-
-class VQModel1d(VQModel):
-    def __init__(self, ddconfig, lossconfig, n_embed, embed_dim, ckpt_path=None, ignore_keys=[],
-                 image_key='feature', colorize_nlabels=None, monitor=None):
-        # ckpt_path is none to super because otherwise will try to load 1D checkpoint into 2D model
-        super().__init__(ddconfig, lossconfig, n_embed, embed_dim)
-        self.image_key = image_key
-        # we need this one for compatibility in train.ImageLogger.log_img if statement
-        self.first_stage_key = image_key
-        self.encoder = Encoder1d(**ddconfig)
-        self.decoder = Decoder1d(**ddconfig)
-        self.loss = instantiate_from_config(lossconfig)
-        self.quantize = VectorQuantizer1d(n_embed, embed_dim, beta=0.25)
-        self.quant_conv = torch.nn.Conv1d(ddconfig['z_channels'], embed_dim, 1)
-        self.post_quant_conv = torch.nn.Conv1d(embed_dim, ddconfig['z_channels'], 1)
-        if ckpt_path is not None:
-            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
-        if colorize_nlabels is not None:
-            assert type(colorize_nlabels)==int
-            self.register_buffer('colorize', torch.randn(3, colorize_nlabels, 1, 1))
-        if monitor is not None:
-            self.monitor = monitor
-
-    def get_input(self, batch, k):
-        x = batch[k]
-        if self.image_key == 'feature':
-            x = x.permute(0, 2, 1)
-        elif self.image_key == 'image':
-            x = x.unsqueeze(1)
-        x = x.to(memory_format=torch.contiguous_format)
-        return x.float()
-
-    def forward(self, input):
-        if self.image_key == 'image':
-            input = input.squeeze(1)
-        quant, diff, _ = self.encode(input)
-        dec = self.decode(quant)
-        if self.image_key == 'image':
-            dec = dec.unsqueeze(1)
-        return dec, diff
-
-    def log_images(self, batch, **kwargs):
-        if self.image_key == 'image':
-            log = dict()
-            x = self.get_input(batch, self.image_key)
-            x = x.to(self.device)
-            xrec, _ = self(x)
-            if x.shape[1] > 3:
-                # colorize with random projection
-                assert xrec.shape[1] > 3
-                x = self.to_rgb(x)
-                xrec = self.to_rgb(xrec)
-            log['inputs'] = x
-            log['reconstructions'] = xrec
-            return log
-        else:
-            raise NotImplementedError('1d input should be treated differently')
-
-    def to_rgb(self, batch, **kwargs):
-        raise NotImplementedError('1d input should be treated differently')
-
-
-class VQSegmentationModel(VQModel):
-    def __init__(self, n_labels, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.register_buffer("colorize", torch.randn(3, n_labels, 1, 1))
-
-    def configure_optimizers(self):
-        lr = self.learning_rate
-        opt_ae = torch.optim.Adam(list(self.encoder.parameters())+
-                                  list(self.decoder.parameters())+
-                                  list(self.quantize.parameters())+
-                                  list(self.quant_conv.parameters())+
-                                  list(self.post_quant_conv.parameters()),
-                                  lr=lr, betas=(0.5, 0.9))
-        return opt_ae
-
-    def training_step(self, batch, batch_idx):
-        x = self.get_input(batch, self.image_key)
-        xrec, qloss = self(x)
-        aeloss, log_dict_ae = self.loss(qloss, x, xrec, split="train")
-        self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True)
-        return aeloss
-
-    def validation_step(self, batch, batch_idx):
-        x = self.get_input(batch, self.image_key)
-        xrec, qloss = self(x)
-        aeloss, log_dict_ae = self.loss(qloss, x, xrec, split="val")
-        self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True)
-        total_loss = log_dict_ae["val/total_loss"]
-        self.log("val/total_loss", total_loss,
-                 prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
-        return aeloss
-
-    @torch.no_grad()
-    def log_images(self, batch, **kwargs):
-        log = dict()
-        x = self.get_input(batch, self.image_key)
-        x = x.to(self.device)
-        xrec, _ = self(x)
-        if x.shape[1] > 3:
-            # colorize with random projection
-            assert xrec.shape[1] > 3
-            # convert logits to indices
-            xrec = torch.argmax(xrec, dim=1, keepdim=True)
-            xrec = F.one_hot(xrec, num_classes=x.shape[1])
-            xrec = xrec.squeeze(1).permute(0, 3, 1, 2).float()
-            x = self.to_rgb(x)
-            xrec = self.to_rgb(xrec)
-        log["inputs"] = x
-        log["reconstructions"] = xrec
-        return log
-
-
-class VQNoDiscModel(VQModel):
-    def __init__(self,
-                 ddconfig,
-                 lossconfig,
-                 n_embed,
-                 embed_dim,
-                 ckpt_path=None,
-                 ignore_keys=[],
-                 image_key="image",
-                 colorize_nlabels=None
-                 ):
-        super().__init__(ddconfig=ddconfig, lossconfig=lossconfig, n_embed=n_embed, embed_dim=embed_dim,
-                         ckpt_path=ckpt_path, ignore_keys=ignore_keys, image_key=image_key,
-                         colorize_nlabels=colorize_nlabels)
-
-    def training_step(self, batch, batch_idx):
-        x = self.get_input(batch, self.image_key)
-        xrec, qloss = self(x)
-        # autoencode
-        aeloss, log_dict_ae = self.loss(qloss, x, xrec, self.global_step, split="train")
-        output = pl.TrainResult(minimize=aeloss)
-        output.log("train/aeloss", aeloss,
-                   prog_bar=True, logger=True, on_step=True, on_epoch=True)
-        output.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True)
-        return output
-
-    def validation_step(self, batch, batch_idx):
-        x = self.get_input(batch, self.image_key)
-        xrec, qloss = self(x)
-        aeloss, log_dict_ae = self.loss(qloss, x, xrec, self.global_step, split="val")
-        rec_loss = log_dict_ae["val/rec_loss"]
-        output = pl.EvalResult(checkpoint_on=rec_loss)
-        output.log("val/rec_loss", rec_loss,
-                   prog_bar=True, logger=True, on_step=True, on_epoch=True)
-        output.log("val/aeloss", aeloss,
-                   prog_bar=True, logger=True, on_step=True, on_epoch=True)
-        output.log_dict(log_dict_ae)
-
-        return output
-
-    def configure_optimizers(self):
-        optimizer = torch.optim.Adam(list(self.encoder.parameters()) +
-                                     list(self.decoder.parameters()) +
-                                     list(self.quantize.parameters()) +
-                                     list(self.quant_conv.parameters()) +
-                                     list(self.post_quant_conv.parameters()),
-                                     lr=self.learning_rate, betas=(0.5, 0.9))
-        return optimizer
-
-
-class VQAudioModelInterfaceForLLM(VQModel):
-    def __init__(self, embed_dim, *args, **kwargs):
-        super().__init__(embed_dim=embed_dim, *args, **kwargs)
-        self.embed_dim = embed_dim
-
-    def encode_seq(self,x):
-        quant, diff, (_,_,ind) = self.encode(x)
-        n,c,h,w = quant.shape
-        return quant,ind,(n,c,h,w)
-
-    def decode_seq(self,ind,target_shape):
-        n,c,h,w = target_shape
-        min_encodings = torch.zeros(ind.shape[0], self.quantize.n_e)
-        min_encodings.scatter_(1, ind, 1)
-
-        # dtype min encodings: torch.float32
-        # min_encodings shape: torch.Size([2048, 512])
-        # min_encoding_indices.shape: torch.Size([2048, 1])
-
-        # get quantized latent vectors
-        z_q = torch.matmul(min_encodings, self.quantize.embedding.weight).view(n,h,w,c).permute(0,3,1,2)
-        y = self.decode(z_q)
-        return y
-    
-from .modules.util import get_duration,calculate_codebook_bitrate,get_audio_file_bitrate
-from .modules.extract_spec import extract_melspectrogram
-from .modules.vocoder import load_vocoder
-class AudioVAEProcesser:
-
-    def __init__(self,sample_rate=22050) -> None:
-        
-        self.model_sr = sample_rate
-
-    def __call__(self, audio) -> Any:
-        input_wav = audio # '../neural_audio_codec/music2/original.wav'
-        duration = get_duration(input_wav)
-        spec = extract_melspectrogram(input_wav, sr=self.model_sr, duration=duration)
-        x = torch.tensor(spec)[None,..., None].permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format)
-        return x
-            
-
-
-if __name__ == '__main__':
-    from omegaconf import OmegaConf
-    from train import instantiate_from_config
-
-    image_key = 'image'
-    cfg_audio = OmegaConf.load('./configs/vggsound_codebook.yaml')
-    model = VQModel(cfg_audio.model.params.ddconfig,
-                    cfg_audio.model.params.lossconfig,
-                    cfg_audio.model.params.n_embed,
-                    cfg_audio.model.params.embed_dim,
-                    image_key='image')
-    batch = {
-        'image': torch.rand((4, 80, 848)),
-        'file_path_': ['data/vggsound/mel123.npy', 'data/vggsound/mel123.npy', 'data/vggsound/mel123.npy'],
-        'class': [1, 1, 1],
-    }
-    xrec, qloss = model(model.get_input(batch, image_key))
-    print(xrec.shape, qloss.shape)
\ No newline at end of file
diff --git a/instructany2pix/llm/model/vae/builder.py b/instructany2pix/llm/model/vae/builder.py
deleted file mode 100644
index 222c7f082256f78c9ad43aa63ae3339579b8737f..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/vae/builder.py
+++ /dev/null
@@ -1,79 +0,0 @@
-import pathlib
-import os
-from typing import Any
-from .image_vqvae import instantiate_from_config,ImageVAEProcesser
-from .audio_vqvae import AudioVAEProcesser
-from .clip import ClipVisionPreprocessorForLLM,NpzLoader
-import yaml
-import torch
-from torch import nn
-CURR_PATH = pathlib.Path(__file__).parent.resolve()
-
-def build_vae(config,modality):
-    config_file = CURR_PATH / f'{config}.yaml'
-    with open(config_file) as f:
-        config = yaml.safe_load(f.read())
-    model_config  =  config['model']
-    model = instantiate_from_config(model_config)
-    ckpt = config['ckpt']
-    if ckpt:
-        sd = torch.load(ckpt,map_location='cpu')["state_dict"]
-        model.load_state_dict(sd, strict=False)
-    if modality == 'image':
-        if config.get("processor") == 'clip':
-            processor = ClipVisionPreprocessorForLLM(pretrained=config['model']['params']['pretrained'])
-        else:
-            processor = ImageVAEProcesser(config['image_size'])
-    elif modality == 'audio':
-        if config.get("processor") == 'npz':
-            processor = NpzLoader()
-        else:
-            processor = AudioVAEProcesser(config['data']['params']['sample_rate'])
-    return model,processor
-
-
-
-class VQVAEProcessor:
-    def __init__(self,image_vae_processor,audio_vae_processor) -> None:
-        self.image_vae_processor = image_vae_processor
-        self.audio_vae_processor = audio_vae_processor
-
-    def __call__(self, x,modality) -> Any:
-        if modality == 'image':
-            return self.image_vae_processor(x)
-        elif modality == 'audio':
-            return self.audio_vae_processor(x)
-        else:
-            raise NotImplemented
-
-class VQVAE(nn.Module):
-
-    def __init__(self,image_vae,audio_vae) -> None:
-        super().__init__()
-        self.image_vae_processor = None
-        self.audio_vae_processor = None
-        self.embed_dim_image = 1
-        self.embed_dim_audio = 1
-        self.vocab_size_image = 1
-        self.vocab_size_audio = 1
-        if image_vae:
-            self.image_vae,self.image_vae_processor = build_vae(image_vae,'image')
-            self.image_vae.requires_grad_(False)
-            self.embed_dim_image = self.image_vae.embed_dim
-            self.vocab_size_image = self.image_vae.n_embed
-        if audio_vae:
-            self.audio_vae,self.audio_vae_processor = build_vae(audio_vae,'audio')
-            self.audio_vae.requires_grad_(False)
-            self.embed_dim_audio = self.audio_vae.embed_dim
-            self.vocab_size_audio = self.audio_vae.n_embed
-        self.processor = VQVAEProcessor(self.image_vae_processor,self.audio_vae_processor)
-
-    @torch.no_grad()
-    def forward(self,x):
-        out = {}
-        if 'image' in x:
-            out['image'] = self.image_vae.encode_seq(x['image'])
-        if 'audio' in x:
-            out['audio'] = self.audio_vae.encode_seq(x['audio'])
-        return out
-    
diff --git a/instructany2pix/llm/model/vae/clip.py b/instructany2pix/llm/model/vae/clip.py
deleted file mode 100644
index 2a8103487ea152019c5152b034f7a7d656acf1ba..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/vae/clip.py
+++ /dev/null
@@ -1,52 +0,0 @@
-from transformers import CLIPImageProcessor,CLIPVisionModelWithProjection
-from torch import nn
-from PIL import Image
-import os
-import torch
-import numpy as np
-class ClipVisionModelInterfaceForLLM(nn.Module):
-    def __init__(self, embed_dim,pretrained,lazy_init=True, *args, **kwargs):
-        super().__init__()
-        self.embed_dim = embed_dim
-        self.n_embed = 1024
-        self.clip_model = CLIPVisionModelWithProjection.from_pretrained(pretrained,low_cpu_mem_usage=os.environ.get('DEEPSPEED3',False) == False)
-
-    def encode_seq(self,x,model=None):
-        x = self.clip_model(torch.cat(x.pixel_values)).image_embeds
-        x = x[...,None,None]
-        return x,None,None
-
-    def decode_seq(self,ind,target_shape):
-        return ind
-    
-class ClipVisionPreprocessorForLLM:
-    def __init__(self,pretrained, *args, **kwargs):
-        self.processor = CLIPImageProcessor.from_pretrained(pretrained)
-
-    def __call__(self,image):
-        image = Image.open(image).convert("RGB")
-        return self.processor(images=image)
-    
-
-class TensorLoader(nn.Module):
-    def __init__(self, embed_dim,n_embed,lazy_init=True, *args, **kwargs):
-        super().__init__()
-        self.embed_dim = embed_dim
-        self.n_embed = n_embed
-
-    def encode_seq(self,x,model=None):
-        # x = self.clip_model(torch.cat(x.pixel_values)).image_embeds
-        # x = x[...,None,None]
-        return x,None,None
-
-    def decode_seq(self,ind,target_shape):
-        return ind
-    
-class NpzLoader:
-    def __init__(self, *args, **kwargs):
-        pass
-
-    def __call__(self,image):
-        x = np.load(image)['state']
-        x = torch.tensor(x)
-        return x
\ No newline at end of file
diff --git a/instructany2pix/llm/model/vae/image_vqvae.py b/instructany2pix/llm/model/vae/image_vqvae.py
deleted file mode 100644
index 33f3748f215c06ad6d4a54df1b5b5eb4d8d33990..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/vae/image_vqvae.py
+++ /dev/null
@@ -1,488 +0,0 @@
-from typing import Any
-import torch
-import pytorch_lightning as pl
-import torch.nn.functional as F
-from contextlib import contextmanager
-
-from taming.modules.vqvae.quantize import VectorQuantizer2 as VectorQuantizer
-
-from .modules.blocks import Encoder, Decoder
-from .modules.distributions import DiagonalGaussianDistribution
-
-from .modules.util import instantiate_from_config
-import numpy as np
-from PIL import Image
-from torchvision import transforms as T
-class VQModel(pl.LightningModule):
-    def __init__(self,
-                 ddconfig,
-                 lossconfig,
-                 n_embed,
-                 embed_dim,
-                 ckpt_path=None,
-                 ignore_keys=[],
-                 image_key="image",
-                 colorize_nlabels=None,
-                 monitor=None,
-                 batch_resize_range=None,
-                 scheduler_config=None,
-                 lr_g_factor=1.0,
-                 remap=None,
-                 sane_index_shape=False, # tell vector quantizer to return indices as bhw
-                 use_ema=False
-                 ):
-        super().__init__()
-        self.embed_dim = embed_dim
-        self.n_embed = n_embed
-        self.image_key = image_key
-        self.encoder = Encoder(**ddconfig)
-        self.decoder = Decoder(**ddconfig)
-        self.loss = instantiate_from_config(lossconfig)
-        self.quantize = VectorQuantizer(n_embed, embed_dim, beta=0.25,
-                                        remap=remap,
-                                        sane_index_shape=sane_index_shape)
-        self.quant_conv = torch.nn.Conv2d(ddconfig["z_channels"], embed_dim, 1)
-        self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
-        if colorize_nlabels is not None:
-            assert type(colorize_nlabels)==int
-            self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
-        if monitor is not None:
-            self.monitor = monitor
-        self.batch_resize_range = batch_resize_range
-        if self.batch_resize_range is not None:
-            print(f"{self.__class__.__name__}: Using per-batch resizing in range {batch_resize_range}.")
-
-        self.use_ema = use_ema
-        if self.use_ema:
-            raise NotImplemented
-            #self.model_ema = LitEma(self)
-            print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
-
-        if ckpt_path is not None:
-            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
-        self.scheduler_config = scheduler_config
-        self.lr_g_factor = lr_g_factor
-
-    @contextmanager
-    def ema_scope(self, context=None):
-        if self.use_ema:
-            self.model_ema.store(self.parameters())
-            self.model_ema.copy_to(self)
-            if context is not None:
-                print(f"{context}: Switched to EMA weights")
-        try:
-            yield None
-        finally:
-            if self.use_ema:
-                self.model_ema.restore(self.parameters())
-                if context is not None:
-                    print(f"{context}: Restored training weights")
-
-    def init_from_ckpt(self, path, ignore_keys=list()):
-        sd = torch.load(path, map_location="cpu")["state_dict"]
-        keys = list(sd.keys())
-        for k in keys:
-            for ik in ignore_keys:
-                if k.startswith(ik):
-                    print("Deleting key {} from state_dict.".format(k))
-                    del sd[k]
-        missing, unexpected = self.load_state_dict(sd, strict=False)
-        print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
-        if len(missing) > 0:
-            print(f"Missing Keys: {missing}")
-            print(f"Unexpected Keys: {unexpected}")
-
-    def on_train_batch_end(self, *args, **kwargs):
-        if self.use_ema:
-            self.model_ema(self)
-
-    def encode(self, x):
-        h = self.encoder(x)
-        h = self.quant_conv(h)
-        quant, emb_loss, info = self.quantize(h)
-        return quant, emb_loss, info
-
-    def encode_to_prequant(self, x):
-        h = self.encoder(x)
-        h = self.quant_conv(h)
-        return h
-
-    def decode(self, quant):
-        quant = self.post_quant_conv(quant)
-        dec = self.decoder(quant)
-        return dec
-
-    def decode_code(self, code_b):
-        quant_b = self.quantize.embed_code(code_b)
-        dec = self.decode(quant_b)
-        return dec
-
-    def forward(self, input, return_pred_indices=False):
-        quant, diff, (_,_,ind) = self.encode(input)
-        dec = self.decode(quant)
-        if return_pred_indices:
-            return dec, diff, ind
-        return dec, diff
-
-    def get_input(self, batch, k):
-        x = batch[k]
-        if len(x.shape) == 3:
-            x = x[..., None]
-        x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float()
-        if self.batch_resize_range is not None:
-            lower_size = self.batch_resize_range[0]
-            upper_size = self.batch_resize_range[1]
-            if self.global_step <= 4:
-                # do the first few batches with max size to avoid later oom
-                new_resize = upper_size
-            else:
-                new_resize = np.random.choice(np.arange(lower_size, upper_size+16, 16))
-            if new_resize != x.shape[2]:
-                x = F.interpolate(x, size=new_resize, mode="bicubic")
-            x = x.detach()
-        return x
-
-    def training_step(self, batch, batch_idx, optimizer_idx):
-        # https://github.com/pytorch/pytorch/issues/37142
-        # try not to fool the heuristics
-        x = self.get_input(batch, self.image_key)
-        xrec, qloss, ind = self(x, return_pred_indices=True)
-
-        if optimizer_idx == 0:
-            # autoencode
-            aeloss, log_dict_ae = self.loss(qloss, x, xrec, optimizer_idx, self.global_step,
-                                            last_layer=self.get_last_layer(), split="train",
-                                            predicted_indices=ind)
-
-            self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True)
-            return aeloss
-
-        if optimizer_idx == 1:
-            # discriminator
-            discloss, log_dict_disc = self.loss(qloss, x, xrec, optimizer_idx, self.global_step,
-                                            last_layer=self.get_last_layer(), split="train")
-            self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True)
-            return discloss
-
-    def validation_step(self, batch, batch_idx):
-        log_dict = self._validation_step(batch, batch_idx)
-        with self.ema_scope():
-            log_dict_ema = self._validation_step(batch, batch_idx, suffix="_ema")
-        return log_dict
-
-    def _validation_step(self, batch, batch_idx, suffix=""):
-        raise NotImplemented
-        # x = self.get_input(batch, self.image_key)
-        # xrec, qloss, ind = self(x, return_pred_indices=True)
-        # aeloss, log_dict_ae = self.loss(qloss, x, xrec, 0,
-        #                                 self.global_step,
-        #                                 last_layer=self.get_last_layer(),
-        #                                 split="val"+suffix,
-        #                                 predicted_indices=ind
-        #                                 )
-
-        # discloss, log_dict_disc = self.loss(qloss, x, xrec, 1,
-        #                                     self.global_step,
-        #                                     last_layer=self.get_last_layer(),
-        #                                     split="val"+suffix,
-        #                                     predicted_indices=ind
-        #                                     )
-        # rec_loss = log_dict_ae[f"val{suffix}/rec_loss"]
-        # self.log(f"val{suffix}/rec_loss", rec_loss,
-        #            prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True)
-        # self.log(f"val{suffix}/aeloss", aeloss,
-        #            prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True)
-        # if version.parse(pl.__version__) >= version.parse('1.4.0'):
-        #     del log_dict_ae[f"val{suffix}/rec_loss"]
-        # self.log_dict(log_dict_ae)
-        # self.log_dict(log_dict_disc)
-        # return self.log_dict
-
-    def configure_optimizers(self):
-        lr_d = self.learning_rate
-        lr_g = self.lr_g_factor*self.learning_rate
-        print("lr_d", lr_d)
-        print("lr_g", lr_g)
-        opt_ae = torch.optim.Adam(list(self.encoder.parameters())+
-                                  list(self.decoder.parameters())+
-                                  list(self.quantize.parameters())+
-                                  list(self.quant_conv.parameters())+
-                                  list(self.post_quant_conv.parameters()),
-                                  lr=lr_g, betas=(0.5, 0.9))
-        opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
-                                    lr=lr_d, betas=(0.5, 0.9))
-
-        if self.scheduler_config is not None:
-            scheduler = instantiate_from_config(self.scheduler_config)
-
-            print("Setting up LambdaLR scheduler...")
-            raise NotImplemented
-            # scheduler = [
-            #     {
-            #         'scheduler': LambdaLR(opt_ae, lr_lambda=scheduler.schedule),
-            #         'interval': 'step',
-            #         'frequency': 1
-            #     },
-            #     {
-            #         'scheduler': LambdaLR(opt_disc, lr_lambda=scheduler.schedule),
-            #         'interval': 'step',
-            #         'frequency': 1
-            #     },
-            # ]
-            return [opt_ae, opt_disc], scheduler
-        return [opt_ae, opt_disc], []
-
-    def get_last_layer(self):
-        return self.decoder.conv_out.weight
-
-    def log_images(self, batch, only_inputs=False, plot_ema=False, **kwargs):
-        log = dict()
-        x = self.get_input(batch, self.image_key)
-        x = x.to(self.device)
-        if only_inputs:
-            log["inputs"] = x
-            return log
-        xrec, _ = self(x)
-        if x.shape[1] > 3:
-            # colorize with random projection
-            assert xrec.shape[1] > 3
-            x = self.to_rgb(x)
-            xrec = self.to_rgb(xrec)
-        log["inputs"] = x
-        log["reconstructions"] = xrec
-        if plot_ema:
-            with self.ema_scope():
-                xrec_ema, _ = self(x)
-                if x.shape[1] > 3: xrec_ema = self.to_rgb(xrec_ema)
-                log["reconstructions_ema"] = xrec_ema
-        return log
-
-    def to_rgb(self, x):
-        assert self.image_key == "segmentation"
-        if not hasattr(self, "colorize"):
-            self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x))
-        x = F.conv2d(x, weight=self.colorize)
-        x = 2.*(x-x.min())/(x.max()-x.min()) - 1.
-        return x
-
-
-class VQModelInterface(VQModel):
-    def __init__(self, embed_dim, *args, **kwargs):
-        super().__init__(embed_dim=embed_dim, *args, **kwargs)
-        self.embed_dim = embed_dim
-
-    def encode(self, x):
-        h = self.encoder(x)
-        h = self.quant_conv(h)
-        return h
-
-    def decode(self, h, force_not_quantize=False):
-        # also go through quantization layer
-        if not force_not_quantize:
-            quant, emb_loss, info = self.quantize(h)
-        else:
-            quant = h
-        quant = self.post_quant_conv(quant)
-        dec = self.decoder(quant)
-        return dec
-
-class VQVisionModelInterfaceForLLM(VQModel):
-    def __init__(self, embed_dim, *args, **kwargs):
-        super().__init__(embed_dim=embed_dim, *args, **kwargs)
-        self.embed_dim = embed_dim
-
-    def encode_seq(self,x):
-        quant, diff, (_,_,ind) = self.encode(x)
-        n,c,h,w = quant.shape
-        return quant,ind,(n,c,h,w)
-
-    def decode_seq(self,ind,target_shape):
-        n,c,h,w = target_shape
-        y = self.quantize.embedding(ind).view(n,h,w,c).permute(0,3,1,2)
-        y = self.decode(y)
-        return y
-
-class AutoencoderKL(pl.LightningModule):
-    def __init__(self,
-                 ddconfig,
-                 lossconfig,
-                 embed_dim,
-                 ckpt_path=None,
-                 ignore_keys=[],
-                 image_key="image",
-                 colorize_nlabels=None,
-                 monitor=None,
-                 ):
-        super().__init__()
-        self.image_key = image_key
-        self.encoder = Encoder(**ddconfig)
-        self.decoder = Decoder(**ddconfig)
-        self.loss = instantiate_from_config(lossconfig)
-        assert ddconfig["double_z"]
-        self.quant_conv = torch.nn.Conv2d(2*ddconfig["z_channels"], 2*embed_dim, 1)
-        self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
-        self.embed_dim = embed_dim
-        if colorize_nlabels is not None:
-            assert type(colorize_nlabels)==int
-            self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
-        if monitor is not None:
-            self.monitor = monitor
-        if ckpt_path is not None:
-            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
-
-    def init_from_ckpt(self, path, ignore_keys=list()):
-        sd = torch.load(path, map_location="cpu")["state_dict"]
-        keys = list(sd.keys())
-        for k in keys:
-            for ik in ignore_keys:
-                if k.startswith(ik):
-                    print("Deleting key {} from state_dict.".format(k))
-                    del sd[k]
-        self.load_state_dict(sd, strict=False)
-        print(f"Restored from {path}")
-
-    def encode(self, x):
-        h = self.encoder(x)
-        moments = self.quant_conv(h)
-        posterior = DiagonalGaussianDistribution(moments)
-        return posterior
-
-    def decode(self, z):
-        z = self.post_quant_conv(z)
-        dec = self.decoder(z)
-        return dec
-
-    def forward(self, input, sample_posterior=True):
-        posterior = self.encode(input)
-        if sample_posterior:
-            z = posterior.sample()
-        else:
-            z = posterior.mode()
-        dec = self.decode(z)
-        return dec, posterior
-
-    def get_input(self, batch, k):
-        x = batch[k]
-        if len(x.shape) == 3:
-            x = x[..., None]
-        x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float()
-        return x
-
-    def training_step(self, batch, batch_idx, optimizer_idx):
-        inputs = self.get_input(batch, self.image_key)
-        reconstructions, posterior = self(inputs)
-
-        if optimizer_idx == 0:
-            # train encoder+decoder+logvar
-            aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,
-                                            last_layer=self.get_last_layer(), split="train")
-            self.log("aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
-            self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=False)
-            return aeloss
-
-        if optimizer_idx == 1:
-            # train the discriminator
-            discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,
-                                                last_layer=self.get_last_layer(), split="train")
-
-            self.log("discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
-            self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=False)
-            return discloss
-
-    def validation_step(self, batch, batch_idx):
-        inputs = self.get_input(batch, self.image_key)
-        reconstructions, posterior = self(inputs)
-        aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, 0, self.global_step,
-                                        last_layer=self.get_last_layer(), split="val")
-
-        discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, 1, self.global_step,
-                                            last_layer=self.get_last_layer(), split="val")
-
-        self.log("val/rec_loss", log_dict_ae["val/rec_loss"])
-        self.log_dict(log_dict_ae)
-        self.log_dict(log_dict_disc)
-        return self.log_dict
-
-    def configure_optimizers(self):
-        lr = self.learning_rate
-        opt_ae = torch.optim.Adam(list(self.encoder.parameters())+
-                                  list(self.decoder.parameters())+
-                                  list(self.quant_conv.parameters())+
-                                  list(self.post_quant_conv.parameters()),
-                                  lr=lr, betas=(0.5, 0.9))
-        opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
-                                    lr=lr, betas=(0.5, 0.9))
-        return [opt_ae, opt_disc], []
-
-    def get_last_layer(self):
-        return self.decoder.conv_out.weight
-
-    @torch.no_grad()
-    def log_images(self, batch, only_inputs=False, **kwargs):
-        log = dict()
-        x = self.get_input(batch, self.image_key)
-        x = x.to(self.device)
-        if not only_inputs:
-            xrec, posterior = self(x)
-            if x.shape[1] > 3:
-                # colorize with random projection
-                assert xrec.shape[1] > 3
-                x = self.to_rgb(x)
-                xrec = self.to_rgb(xrec)
-            log["samples"] = self.decode(torch.randn_like(posterior.sample()))
-            log["reconstructions"] = xrec
-        log["inputs"] = x
-        return log
-
-    def to_rgb(self, x):
-        assert self.image_key == "segmentation"
-        if not hasattr(self, "colorize"):
-            self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x))
-        x = F.conv2d(x, weight=self.colorize)
-        x = 2.*(x-x.min())/(x.max()-x.min()) - 1.
-        return x
-
-
-class IdentityFirstStage(torch.nn.Module):
-    def __init__(self, *args, vq_interface=False, **kwargs):
-        self.vq_interface = vq_interface  # TODO: Should be true by default but check to not break older stuff
-        super().__init__()
-
-    def encode(self, x, *args, **kwargs):
-        return x
-
-    def decode(self, x, *args, **kwargs):
-        return x
-
-    def quantize(self, x, *args, **kwargs):
-        if self.vq_interface:
-            return x, None, [None, None, None]
-        return x
-
-    def forward(self, x, *args, **kwargs):
-        return x
-    
-class ImageVAEProcesser:
-
-    def __init__(self,image_size=128) -> None:
-        self.transform = T.Compose([
-                    T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
-                    T.Resize(image_size),
-                    T.RandomHorizontalFlip(),
-                    T.CenterCrop(image_size),
-                    T.ToTensor()
-                ])
-        
-    def __call__(self, image) -> Any:
-        image = Image.open(image).convert("RGB")
-        image = self.transform(image)
-        return image
-
-def make_batch_image(image,size=128):
-    image = Image.open(image).convert("RGB")
-    image = image.resize((128,128))
-    image = np.array(image)
-    image = image.astype(np.float32)/255.0
-    image = image[None].transpose(0,3,1,2)
-    image = torch.from_numpy(image)
-    return image
diff --git a/instructany2pix/llm/model/vae/modules/__init__.py b/instructany2pix/llm/model/vae/modules/__init__.py
deleted file mode 100644
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000
diff --git a/instructany2pix/llm/model/vae/modules/attention.py b/instructany2pix/llm/model/vae/modules/attention.py
deleted file mode 100644
index f6a666bc2ee900e6839542cf3ce3c65e8b5d3d4d..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/vae/modules/attention.py
+++ /dev/null
@@ -1,261 +0,0 @@
-from inspect import isfunction
-import math
-import torch
-import torch.nn.functional as F
-from torch import nn, einsum
-from einops import rearrange, repeat
-
-from .util import checkpoint
-
-
-def exists(val):
-    return val is not None
-
-
-def uniq(arr):
-    return{el: True for el in arr}.keys()
-
-
-def default(val, d):
-    if exists(val):
-        return val
-    return d() if isfunction(d) else d
-
-
-def max_neg_value(t):
-    return -torch.finfo(t.dtype).max
-
-
-def init_(tensor):
-    dim = tensor.shape[-1]
-    std = 1 / math.sqrt(dim)
-    tensor.uniform_(-std, std)
-    return tensor
-
-
-# feedforward
-class GEGLU(nn.Module):
-    def __init__(self, dim_in, dim_out):
-        super().__init__()
-        self.proj = nn.Linear(dim_in, dim_out * 2)
-
-    def forward(self, x):
-        x, gate = self.proj(x).chunk(2, dim=-1)
-        return x * F.gelu(gate)
-
-
-class FeedForward(nn.Module):
-    def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
-        super().__init__()
-        inner_dim = int(dim * mult)
-        dim_out = default(dim_out, dim)
-        project_in = nn.Sequential(
-            nn.Linear(dim, inner_dim),
-            nn.GELU()
-        ) if not glu else GEGLU(dim, inner_dim)
-
-        self.net = nn.Sequential(
-            project_in,
-            nn.Dropout(dropout),
-            nn.Linear(inner_dim, dim_out)
-        )
-
-    def forward(self, x):
-        return self.net(x)
-
-
-def zero_module(module):
-    """
-    Zero out the parameters of a module and return it.
-    """
-    for p in module.parameters():
-        p.detach().zero_()
-    return module
-
-
-def Normalize(in_channels):
-    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
-
-
-class LinearAttention(nn.Module):
-    def __init__(self, dim, heads=4, dim_head=32):
-        super().__init__()
-        self.heads = heads
-        hidden_dim = dim_head * heads
-        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False)
-        self.to_out = nn.Conv2d(hidden_dim, dim, 1)
-
-    def forward(self, x):
-        b, c, h, w = x.shape
-        qkv = self.to_qkv(x)
-        q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads = self.heads, qkv=3)
-        k = k.softmax(dim=-1)  
-        context = torch.einsum('bhdn,bhen->bhde', k, v)
-        out = torch.einsum('bhde,bhdn->bhen', context, q)
-        out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w)
-        return self.to_out(out)
-
-
-class SpatialSelfAttention(nn.Module):
-    def __init__(self, in_channels):
-        super().__init__()
-        self.in_channels = in_channels
-
-        self.norm = Normalize(in_channels)
-        self.q = torch.nn.Conv2d(in_channels,
-                                 in_channels,
-                                 kernel_size=1,
-                                 stride=1,
-                                 padding=0)
-        self.k = torch.nn.Conv2d(in_channels,
-                                 in_channels,
-                                 kernel_size=1,
-                                 stride=1,
-                                 padding=0)
-        self.v = torch.nn.Conv2d(in_channels,
-                                 in_channels,
-                                 kernel_size=1,
-                                 stride=1,
-                                 padding=0)
-        self.proj_out = torch.nn.Conv2d(in_channels,
-                                        in_channels,
-                                        kernel_size=1,
-                                        stride=1,
-                                        padding=0)
-
-    def forward(self, x):
-        h_ = x
-        h_ = self.norm(h_)
-        q = self.q(h_)
-        k = self.k(h_)
-        v = self.v(h_)
-
-        # compute attention
-        b,c,h,w = q.shape
-        q = rearrange(q, 'b c h w -> b (h w) c')
-        k = rearrange(k, 'b c h w -> b c (h w)')
-        w_ = torch.einsum('bij,bjk->bik', q, k)
-
-        w_ = w_ * (int(c)**(-0.5))
-        w_ = torch.nn.functional.softmax(w_, dim=2)
-
-        # attend to values
-        v = rearrange(v, 'b c h w -> b c (h w)')
-        w_ = rearrange(w_, 'b i j -> b j i')
-        h_ = torch.einsum('bij,bjk->bik', v, w_)
-        h_ = rearrange(h_, 'b c (h w) -> b c h w', h=h)
-        h_ = self.proj_out(h_)
-
-        return x+h_
-
-
-class CrossAttention(nn.Module):
-    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.):
-        super().__init__()
-        inner_dim = dim_head * heads
-        context_dim = default(context_dim, query_dim)
-
-        self.scale = dim_head ** -0.5
-        self.heads = heads
-
-        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
-        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
-        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
-
-        self.to_out = nn.Sequential(
-            nn.Linear(inner_dim, query_dim),
-            nn.Dropout(dropout)
-        )
-
-    def forward(self, x, context=None, mask=None):
-        h = self.heads
-
-        q = self.to_q(x)
-        context = default(context, x)
-        k = self.to_k(context)
-        v = self.to_v(context)
-
-        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
-
-        sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
-
-        if exists(mask):
-            mask = rearrange(mask, 'b ... -> b (...)')
-            max_neg_value = -torch.finfo(sim.dtype).max
-            mask = repeat(mask, 'b j -> (b h) () j', h=h)
-            sim.masked_fill_(~mask, max_neg_value)
-
-        # attention, what we cannot get enough of
-        attn = sim.softmax(dim=-1)
-
-        out = einsum('b i j, b j d -> b i d', attn, v)
-        out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
-        return self.to_out(out)
-
-
-class BasicTransformerBlock(nn.Module):
-    def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint=True):
-        super().__init__()
-        self.attn1 = CrossAttention(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout)  # is a self-attention
-        self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
-        self.attn2 = CrossAttention(query_dim=dim, context_dim=context_dim,
-                                    heads=n_heads, dim_head=d_head, dropout=dropout)  # is self-attn if context is none
-        self.norm1 = nn.LayerNorm(dim)
-        self.norm2 = nn.LayerNorm(dim)
-        self.norm3 = nn.LayerNorm(dim)
-        self.checkpoint = checkpoint
-
-    def forward(self, x, context=None):
-        return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint)
-
-    def _forward(self, x, context=None):
-        x = self.attn1(self.norm1(x)) + x
-        x = self.attn2(self.norm2(x), context=context) + x
-        x = self.ff(self.norm3(x)) + x
-        return x
-
-
-class SpatialTransformer(nn.Module):
-    """
-    Transformer block for image-like data.
-    First, project the input (aka embedding)
-    and reshape to b, t, d.
-    Then apply standard transformer action.
-    Finally, reshape to image
-    """
-    def __init__(self, in_channels, n_heads, d_head,
-                 depth=1, dropout=0., context_dim=None):
-        super().__init__()
-        self.in_channels = in_channels
-        inner_dim = n_heads * d_head
-        self.norm = Normalize(in_channels)
-
-        self.proj_in = nn.Conv2d(in_channels,
-                                 inner_dim,
-                                 kernel_size=1,
-                                 stride=1,
-                                 padding=0)
-
-        self.transformer_blocks = nn.ModuleList(
-            [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim)
-                for d in range(depth)]
-        )
-
-        self.proj_out = zero_module(nn.Conv2d(inner_dim,
-                                              in_channels,
-                                              kernel_size=1,
-                                              stride=1,
-                                              padding=0))
-
-    def forward(self, x, context=None):
-        # note: if no context is given, cross-attention defaults to self-attention
-        b, c, h, w = x.shape
-        x_in = x
-        x = self.norm(x)
-        x = self.proj_in(x)
-        x = rearrange(x, 'b c h w -> b (h w) c')
-        for block in self.transformer_blocks:
-            x = block(x, context=context)
-        x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w)
-        x = self.proj_out(x)
-        return x + x_in
\ No newline at end of file
diff --git a/instructany2pix/llm/model/vae/modules/blocks.py b/instructany2pix/llm/model/vae/modules/blocks.py
deleted file mode 100644
index 6ed45c94cbcffd52aedc04b903d4636226db4ef0..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/vae/modules/blocks.py
+++ /dev/null
@@ -1,835 +0,0 @@
-# pytorch_diffusion + derived encoder decoder
-import math
-import torch
-import torch.nn as nn
-import numpy as np
-from einops import rearrange
-
-from .util import instantiate_from_config
-from .attention import LinearAttention
-from .distributions import DiagonalGaussianDistribution
-
-def get_timestep_embedding(timesteps, embedding_dim):
-    """
-    This matches the implementation in Denoising Diffusion Probabilistic Models:
-    From Fairseq.
-    Build sinusoidal embeddings.
-    This matches the implementation in tensor2tensor, but differs slightly
-    from the description in Section 3.5 of "Attention Is All You Need".
-    """
-    assert len(timesteps.shape) == 1
-
-    half_dim = embedding_dim // 2
-    emb = math.log(10000) / (half_dim - 1)
-    emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
-    emb = emb.to(device=timesteps.device)
-    emb = timesteps.float()[:, None] * emb[None, :]
-    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
-    if embedding_dim % 2 == 1:  # zero pad
-        emb = torch.nn.functional.pad(emb, (0,1,0,0))
-    return emb
-
-
-def nonlinearity(x):
-    # swish
-    return x*torch.sigmoid(x)
-
-
-def Normalize(in_channels, num_groups=32):
-    return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
-
-
-class Upsample(nn.Module):
-    def __init__(self, in_channels, with_conv):
-        super().__init__()
-        self.with_conv = with_conv
-        if self.with_conv:
-            self.conv = torch.nn.Conv2d(in_channels,
-                                        in_channels,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-    def forward(self, x):
-        dtype = x.dtype
-        x = torch.nn.functional.interpolate(x.float(), scale_factor=2.0, mode="nearest").to(dtype)
-        if self.with_conv:
-            x = self.conv(x)
-        return x
-
-
-class Downsample(nn.Module):
-    def __init__(self, in_channels, with_conv):
-        super().__init__()
-        self.with_conv = with_conv
-        if self.with_conv:
-            # no asymmetric padding in torch conv, must do it ourselves
-            self.conv = torch.nn.Conv2d(in_channels,
-                                        in_channels,
-                                        kernel_size=3,
-                                        stride=2,
-                                        padding=0)
-
-    def forward(self, x):
-        if self.with_conv:
-            pad = (0,1,0,1)
-            x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
-            x = self.conv(x)
-        else:
-            x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
-        return x
-
-
-class ResnetBlock(nn.Module):
-    def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False,
-                 dropout, temb_channels=512):
-        super().__init__()
-        self.in_channels = in_channels
-        out_channels = in_channels if out_channels is None else out_channels
-        self.out_channels = out_channels
-        self.use_conv_shortcut = conv_shortcut
-
-        self.norm1 = Normalize(in_channels)
-        self.conv1 = torch.nn.Conv2d(in_channels,
-                                     out_channels,
-                                     kernel_size=3,
-                                     stride=1,
-                                     padding=1)
-        if temb_channels > 0:
-            self.temb_proj = torch.nn.Linear(temb_channels,
-                                             out_channels)
-        self.norm2 = Normalize(out_channels)
-        self.dropout = torch.nn.Dropout(dropout)
-        self.conv2 = torch.nn.Conv2d(out_channels,
-                                     out_channels,
-                                     kernel_size=3,
-                                     stride=1,
-                                     padding=1)
-        if self.in_channels != self.out_channels:
-            if self.use_conv_shortcut:
-                self.conv_shortcut = torch.nn.Conv2d(in_channels,
-                                                     out_channels,
-                                                     kernel_size=3,
-                                                     stride=1,
-                                                     padding=1)
-            else:
-                self.nin_shortcut = torch.nn.Conv2d(in_channels,
-                                                    out_channels,
-                                                    kernel_size=1,
-                                                    stride=1,
-                                                    padding=0)
-
-    def forward(self, x, temb):
-        h = x
-        h = self.norm1(h)
-        h = nonlinearity(h)
-        h = self.conv1(h)
-
-        if temb is not None:
-            h = h + self.temb_proj(nonlinearity(temb))[:,:,None,None]
-
-        h = self.norm2(h)
-        h = nonlinearity(h)
-        h = self.dropout(h)
-        h = self.conv2(h)
-
-        if self.in_channels != self.out_channels:
-            if self.use_conv_shortcut:
-                x = self.conv_shortcut(x)
-            else:
-                x = self.nin_shortcut(x)
-
-        return x+h
-
-
-class LinAttnBlock(LinearAttention):
-    """to match AttnBlock usage"""
-    def __init__(self, in_channels):
-        super().__init__(dim=in_channels, heads=1, dim_head=in_channels)
-
-
-class AttnBlock(nn.Module):
-    def __init__(self, in_channels):
-        super().__init__()
-        self.in_channels = in_channels
-
-        self.norm = Normalize(in_channels)
-        self.q = torch.nn.Conv2d(in_channels,
-                                 in_channels,
-                                 kernel_size=1,
-                                 stride=1,
-                                 padding=0)
-        self.k = torch.nn.Conv2d(in_channels,
-                                 in_channels,
-                                 kernel_size=1,
-                                 stride=1,
-                                 padding=0)
-        self.v = torch.nn.Conv2d(in_channels,
-                                 in_channels,
-                                 kernel_size=1,
-                                 stride=1,
-                                 padding=0)
-        self.proj_out = torch.nn.Conv2d(in_channels,
-                                        in_channels,
-                                        kernel_size=1,
-                                        stride=1,
-                                        padding=0)
-
-
-    def forward(self, x):
-        h_ = x
-        h_ = self.norm(h_)
-        q = self.q(h_)
-        k = self.k(h_)
-        v = self.v(h_)
-
-        # compute attention
-        b,c,h,w = q.shape
-        q = q.reshape(b,c,h*w)
-        q = q.permute(0,2,1)   # b,hw,c
-        k = k.reshape(b,c,h*w) # b,c,hw
-        w_ = torch.bmm(q,k)     # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
-        w_ = w_ * (int(c)**(-0.5))
-        w_ = torch.nn.functional.softmax(w_, dim=2)
-
-        # attend to values
-        v = v.reshape(b,c,h*w)
-        w_ = w_.permute(0,2,1)   # b,hw,hw (first hw of k, second of q)
-        h_ = torch.bmm(v,w_)     # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
-        h_ = h_.reshape(b,c,h,w)
-
-        h_ = self.proj_out(h_)
-
-        return x+h_
-
-
-def make_attn(in_channels, attn_type="vanilla"):
-    assert attn_type in ["vanilla", "linear", "none"], f'attn_type {attn_type} unknown'
-    print(f"making attention of type '{attn_type}' with {in_channels} in_channels")
-    if attn_type == "vanilla":
-        return AttnBlock(in_channels)
-    elif attn_type == "none":
-        return nn.Identity(in_channels)
-    else:
-        return LinAttnBlock(in_channels)
-
-
-class Model(nn.Module):
-    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
-                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
-                 resolution, use_timestep=True, use_linear_attn=False, attn_type="vanilla"):
-        super().__init__()
-        if use_linear_attn: attn_type = "linear"
-        self.ch = ch
-        self.temb_ch = self.ch*4
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_channels = in_channels
-
-        self.use_timestep = use_timestep
-        if self.use_timestep:
-            # timestep embedding
-            self.temb = nn.Module()
-            self.temb.dense = nn.ModuleList([
-                torch.nn.Linear(self.ch,
-                                self.temb_ch),
-                torch.nn.Linear(self.temb_ch,
-                                self.temb_ch),
-            ])
-
-        # downsampling
-        self.conv_in = torch.nn.Conv2d(in_channels,
-                                       self.ch,
-                                       kernel_size=3,
-                                       stride=1,
-                                       padding=1)
-
-        curr_res = resolution
-        in_ch_mult = (1,)+tuple(ch_mult)
-        self.down = nn.ModuleList()
-        for i_level in range(self.num_resolutions):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_in = ch*in_ch_mult[i_level]
-            block_out = ch*ch_mult[i_level]
-            for i_block in range(self.num_res_blocks):
-                block.append(ResnetBlock(in_channels=block_in,
-                                         out_channels=block_out,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout))
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(make_attn(block_in, attn_type=attn_type))
-            down = nn.Module()
-            down.block = block
-            down.attn = attn
-            if i_level != self.num_resolutions-1:
-                down.downsample = Downsample(block_in, resamp_with_conv)
-                curr_res = curr_res // 2
-            self.down.append(down)
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-        self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
-        self.mid.block_2 = ResnetBlock(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-
-        # upsampling
-        self.up = nn.ModuleList()
-        for i_level in reversed(range(self.num_resolutions)):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_out = ch*ch_mult[i_level]
-            skip_in = ch*ch_mult[i_level]
-            for i_block in range(self.num_res_blocks+1):
-                if i_block == self.num_res_blocks:
-                    skip_in = ch*in_ch_mult[i_level]
-                block.append(ResnetBlock(in_channels=block_in+skip_in,
-                                         out_channels=block_out,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout))
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(make_attn(block_in, attn_type=attn_type))
-            up = nn.Module()
-            up.block = block
-            up.attn = attn
-            if i_level != 0:
-                up.upsample = Upsample(block_in, resamp_with_conv)
-                curr_res = curr_res * 2
-            self.up.insert(0, up) # prepend to get consistent order
-
-        # end
-        self.norm_out = Normalize(block_in)
-        self.conv_out = torch.nn.Conv2d(block_in,
-                                        out_ch,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-    def forward(self, x, t=None, context=None):
-        #assert x.shape[2] == x.shape[3] == self.resolution
-        if context is not None:
-            # assume aligned context, cat along channel axis
-            x = torch.cat((x, context), dim=1)
-        if self.use_timestep:
-            # timestep embedding
-            assert t is not None
-            temb = get_timestep_embedding(t, self.ch)
-            temb = self.temb.dense[0](temb)
-            temb = nonlinearity(temb)
-            temb = self.temb.dense[1](temb)
-        else:
-            temb = None
-
-        # downsampling
-        hs = [self.conv_in(x)]
-        for i_level in range(self.num_resolutions):
-            for i_block in range(self.num_res_blocks):
-                h = self.down[i_level].block[i_block](hs[-1], temb)
-                if len(self.down[i_level].attn) > 0:
-                    h = self.down[i_level].attn[i_block](h)
-                hs.append(h)
-            if i_level != self.num_resolutions-1:
-                hs.append(self.down[i_level].downsample(hs[-1]))
-
-        # middle
-        h = hs[-1]
-        h = self.mid.block_1(h, temb)
-        h = self.mid.attn_1(h)
-        h = self.mid.block_2(h, temb)
-
-        # upsampling
-        for i_level in reversed(range(self.num_resolutions)):
-            for i_block in range(self.num_res_blocks+1):
-                h = self.up[i_level].block[i_block](
-                    torch.cat([h, hs.pop()], dim=1), temb)
-                if len(self.up[i_level].attn) > 0:
-                    h = self.up[i_level].attn[i_block](h)
-            if i_level != 0:
-                h = self.up[i_level].upsample(h)
-
-        # end
-        h = self.norm_out(h)
-        h = nonlinearity(h)
-        h = self.conv_out(h)
-        return h
-
-    def get_last_layer(self):
-        return self.conv_out.weight
-
-
-class Encoder(nn.Module):
-    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
-                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
-                 resolution, z_channels, double_z=True, use_linear_attn=False, attn_type="vanilla",
-                 **ignore_kwargs):
-        super().__init__()
-        if use_linear_attn: attn_type = "linear"
-        self.ch = ch
-        self.temb_ch = 0
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_channels = in_channels
-
-        # downsampling
-        self.conv_in = torch.nn.Conv2d(in_channels,
-                                       self.ch,
-                                       kernel_size=3,
-                                       stride=1,
-                                       padding=1)
-
-        curr_res = resolution
-        in_ch_mult = (1,)+tuple(ch_mult)
-        self.in_ch_mult = in_ch_mult
-        self.down = nn.ModuleList()
-        for i_level in range(self.num_resolutions):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_in = ch*in_ch_mult[i_level]
-            block_out = ch*ch_mult[i_level]
-            for i_block in range(self.num_res_blocks):
-                block.append(ResnetBlock(in_channels=block_in,
-                                         out_channels=block_out,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout))
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(make_attn(block_in, attn_type=attn_type))
-            down = nn.Module()
-            down.block = block
-            down.attn = attn
-            if i_level != self.num_resolutions-1:
-                down.downsample = Downsample(block_in, resamp_with_conv)
-                curr_res = curr_res // 2
-            self.down.append(down)
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-        self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
-        self.mid.block_2 = ResnetBlock(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-
-        # end
-        self.norm_out = Normalize(block_in)
-        self.conv_out = torch.nn.Conv2d(block_in,
-                                        2*z_channels if double_z else z_channels,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-    def forward(self, x):
-        # timestep embedding
-        temb = None
-
-        # downsampling
-        hs = [self.conv_in(x)]
-        for i_level in range(self.num_resolutions):
-            for i_block in range(self.num_res_blocks):
-                h = self.down[i_level].block[i_block](hs[-1], temb)
-                if len(self.down[i_level].attn) > 0:
-                    h = self.down[i_level].attn[i_block](h)
-                hs.append(h)
-            if i_level != self.num_resolutions-1:
-                hs.append(self.down[i_level].downsample(hs[-1]))
-
-        # middle
-        h = hs[-1]
-        h = self.mid.block_1(h, temb)
-        h = self.mid.attn_1(h)
-        h = self.mid.block_2(h, temb)
-
-        # end
-        h = self.norm_out(h)
-        h = nonlinearity(h)
-        h = self.conv_out(h)
-        return h
-
-
-class Decoder(nn.Module):
-    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
-                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
-                 resolution, z_channels, give_pre_end=False, tanh_out=False, use_linear_attn=False,
-                 attn_type="vanilla", **ignorekwargs):
-        super().__init__()
-        if use_linear_attn: attn_type = "linear"
-        self.ch = ch
-        self.temb_ch = 0
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_channels = in_channels
-        self.give_pre_end = give_pre_end
-        self.tanh_out = tanh_out
-
-        # compute in_ch_mult, block_in and curr_res at lowest res
-        in_ch_mult = (1,)+tuple(ch_mult)
-        block_in = ch*ch_mult[self.num_resolutions-1]
-        curr_res = resolution // 2**(self.num_resolutions-1)
-        self.z_shape = (1,z_channels,curr_res,curr_res)
-        print("Working with z of shape {} = {} dimensions.".format(
-            self.z_shape, np.prod(self.z_shape)))
-
-        # z to block_in
-        self.conv_in = torch.nn.Conv2d(z_channels,
-                                       block_in,
-                                       kernel_size=3,
-                                       stride=1,
-                                       padding=1)
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-        self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
-        self.mid.block_2 = ResnetBlock(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-
-        # upsampling
-        self.up = nn.ModuleList()
-        for i_level in reversed(range(self.num_resolutions)):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_out = ch*ch_mult[i_level]
-            for i_block in range(self.num_res_blocks+1):
-                block.append(ResnetBlock(in_channels=block_in,
-                                         out_channels=block_out,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout))
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(make_attn(block_in, attn_type=attn_type))
-            up = nn.Module()
-            up.block = block
-            up.attn = attn
-            if i_level != 0:
-                up.upsample = Upsample(block_in, resamp_with_conv)
-                curr_res = curr_res * 2
-            self.up.insert(0, up) # prepend to get consistent order
-
-        # end
-        self.norm_out = Normalize(block_in)
-        self.conv_out = torch.nn.Conv2d(block_in,
-                                        out_ch,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-    def forward(self, z):
-        #assert z.shape[1:] == self.z_shape[1:]
-        self.last_z_shape = z.shape
-
-        # timestep embedding
-        temb = None
-
-        # z to block_in
-        h = self.conv_in(z)
-
-        # middle
-        h = self.mid.block_1(h, temb)
-        h = self.mid.attn_1(h)
-        h = self.mid.block_2(h, temb)
-
-        # upsampling
-        for i_level in reversed(range(self.num_resolutions)):
-            for i_block in range(self.num_res_blocks+1):
-                h = self.up[i_level].block[i_block](h, temb)
-                if len(self.up[i_level].attn) > 0:
-                    h = self.up[i_level].attn[i_block](h)
-            if i_level != 0:
-                h = self.up[i_level].upsample(h)
-
-        # end
-        if self.give_pre_end:
-            return h
-
-        h = self.norm_out(h)
-        h = nonlinearity(h)
-        h = self.conv_out(h)
-        if self.tanh_out:
-            h = torch.tanh(h)
-        return h
-
-
-class SimpleDecoder(nn.Module):
-    def __init__(self, in_channels, out_channels, *args, **kwargs):
-        super().__init__()
-        self.model = nn.ModuleList([nn.Conv2d(in_channels, in_channels, 1),
-                                     ResnetBlock(in_channels=in_channels,
-                                                 out_channels=2 * in_channels,
-                                                 temb_channels=0, dropout=0.0),
-                                     ResnetBlock(in_channels=2 * in_channels,
-                                                out_channels=4 * in_channels,
-                                                temb_channels=0, dropout=0.0),
-                                     ResnetBlock(in_channels=4 * in_channels,
-                                                out_channels=2 * in_channels,
-                                                temb_channels=0, dropout=0.0),
-                                     nn.Conv2d(2*in_channels, in_channels, 1),
-                                     Upsample(in_channels, with_conv=True)])
-        # end
-        self.norm_out = Normalize(in_channels)
-        self.conv_out = torch.nn.Conv2d(in_channels,
-                                        out_channels,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-    def forward(self, x):
-        for i, layer in enumerate(self.model):
-            if i in [1,2,3]:
-                x = layer(x, None)
-            else:
-                x = layer(x)
-
-        h = self.norm_out(x)
-        h = nonlinearity(h)
-        x = self.conv_out(h)
-        return x
-
-
-class UpsampleDecoder(nn.Module):
-    def __init__(self, in_channels, out_channels, ch, num_res_blocks, resolution,
-                 ch_mult=(2,2), dropout=0.0):
-        super().__init__()
-        # upsampling
-        self.temb_ch = 0
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        block_in = in_channels
-        curr_res = resolution // 2 ** (self.num_resolutions - 1)
-        self.res_blocks = nn.ModuleList()
-        self.upsample_blocks = nn.ModuleList()
-        for i_level in range(self.num_resolutions):
-            res_block = []
-            block_out = ch * ch_mult[i_level]
-            for i_block in range(self.num_res_blocks + 1):
-                res_block.append(ResnetBlock(in_channels=block_in,
-                                         out_channels=block_out,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout))
-                block_in = block_out
-            self.res_blocks.append(nn.ModuleList(res_block))
-            if i_level != self.num_resolutions - 1:
-                self.upsample_blocks.append(Upsample(block_in, True))
-                curr_res = curr_res * 2
-
-        # end
-        self.norm_out = Normalize(block_in)
-        self.conv_out = torch.nn.Conv2d(block_in,
-                                        out_channels,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-    def forward(self, x):
-        # upsampling
-        h = x
-        for k, i_level in enumerate(range(self.num_resolutions)):
-            for i_block in range(self.num_res_blocks + 1):
-                h = self.res_blocks[i_level][i_block](h, None)
-            if i_level != self.num_resolutions - 1:
-                h = self.upsample_blocks[k](h)
-        h = self.norm_out(h)
-        h = nonlinearity(h)
-        h = self.conv_out(h)
-        return h
-
-
-class LatentRescaler(nn.Module):
-    def __init__(self, factor, in_channels, mid_channels, out_channels, depth=2):
-        super().__init__()
-        # residual block, interpolate, residual block
-        self.factor = factor
-        self.conv_in = nn.Conv2d(in_channels,
-                                 mid_channels,
-                                 kernel_size=3,
-                                 stride=1,
-                                 padding=1)
-        self.res_block1 = nn.ModuleList([ResnetBlock(in_channels=mid_channels,
-                                                     out_channels=mid_channels,
-                                                     temb_channels=0,
-                                                     dropout=0.0) for _ in range(depth)])
-        self.attn = AttnBlock(mid_channels)
-        self.res_block2 = nn.ModuleList([ResnetBlock(in_channels=mid_channels,
-                                                     out_channels=mid_channels,
-                                                     temb_channels=0,
-                                                     dropout=0.0) for _ in range(depth)])
-
-        self.conv_out = nn.Conv2d(mid_channels,
-                                  out_channels,
-                                  kernel_size=1,
-                                  )
-
-    def forward(self, x):
-        x = self.conv_in(x)
-        for block in self.res_block1:
-            x = block(x, None)
-        x = torch.nn.functional.interpolate(x, size=(int(round(x.shape[2]*self.factor)), int(round(x.shape[3]*self.factor))))
-        x = self.attn(x)
-        for block in self.res_block2:
-            x = block(x, None)
-        x = self.conv_out(x)
-        return x
-
-
-class MergedRescaleEncoder(nn.Module):
-    def __init__(self, in_channels, ch, resolution, out_ch, num_res_blocks,
-                 attn_resolutions, dropout=0.0, resamp_with_conv=True,
-                 ch_mult=(1,2,4,8), rescale_factor=1.0, rescale_module_depth=1):
-        super().__init__()
-        intermediate_chn = ch * ch_mult[-1]
-        self.encoder = Encoder(in_channels=in_channels, num_res_blocks=num_res_blocks, ch=ch, ch_mult=ch_mult,
-                               z_channels=intermediate_chn, double_z=False, resolution=resolution,
-                               attn_resolutions=attn_resolutions, dropout=dropout, resamp_with_conv=resamp_with_conv,
-                               out_ch=None)
-        self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=intermediate_chn,
-                                       mid_channels=intermediate_chn, out_channels=out_ch, depth=rescale_module_depth)
-
-    def forward(self, x):
-        x = self.encoder(x)
-        x = self.rescaler(x)
-        return x
-
-
-class MergedRescaleDecoder(nn.Module):
-    def __init__(self, z_channels, out_ch, resolution, num_res_blocks, attn_resolutions, ch, ch_mult=(1,2,4,8),
-                 dropout=0.0, resamp_with_conv=True, rescale_factor=1.0, rescale_module_depth=1):
-        super().__init__()
-        tmp_chn = z_channels*ch_mult[-1]
-        self.decoder = Decoder(out_ch=out_ch, z_channels=tmp_chn, attn_resolutions=attn_resolutions, dropout=dropout,
-                               resamp_with_conv=resamp_with_conv, in_channels=None, num_res_blocks=num_res_blocks,
-                               ch_mult=ch_mult, resolution=resolution, ch=ch)
-        self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=z_channels, mid_channels=tmp_chn,
-                                       out_channels=tmp_chn, depth=rescale_module_depth)
-
-    def forward(self, x):
-        x = self.rescaler(x)
-        x = self.decoder(x)
-        return x
-
-
-class Upsampler(nn.Module):
-    def __init__(self, in_size, out_size, in_channels, out_channels, ch_mult=2):
-        super().__init__()
-        assert out_size >= in_size
-        num_blocks = int(np.log2(out_size//in_size))+1
-        factor_up = 1.+ (out_size % in_size)
-        print(f"Building {self.__class__.__name__} with in_size: {in_size} --> out_size {out_size} and factor {factor_up}")
-        self.rescaler = LatentRescaler(factor=factor_up, in_channels=in_channels, mid_channels=2*in_channels,
-                                       out_channels=in_channels)
-        self.decoder = Decoder(out_ch=out_channels, resolution=out_size, z_channels=in_channels, num_res_blocks=2,
-                               attn_resolutions=[], in_channels=None, ch=in_channels,
-                               ch_mult=[ch_mult for _ in range(num_blocks)])
-
-    def forward(self, x):
-        x = self.rescaler(x)
-        x = self.decoder(x)
-        return x
-
-
-class Resize(nn.Module):
-    def __init__(self, in_channels=None, learned=False, mode="bilinear"):
-        super().__init__()
-        self.with_conv = learned
-        self.mode = mode
-        if self.with_conv:
-            print(f"Note: {self.__class__.__name} uses learned downsampling and will ignore the fixed {mode} mode")
-            raise NotImplementedError()
-            assert in_channels is not None
-            # no asymmetric padding in torch conv, must do it ourselves
-            self.conv = torch.nn.Conv2d(in_channels,
-                                        in_channels,
-                                        kernel_size=4,
-                                        stride=2,
-                                        padding=1)
-
-    def forward(self, x, scale_factor=1.0):
-        if scale_factor==1.0:
-            return x
-        else:
-            x = torch.nn.functional.interpolate(x, mode=self.mode, align_corners=False, scale_factor=scale_factor)
-        return x
-
-class FirstStagePostProcessor(nn.Module):
-
-    def __init__(self, ch_mult:list, in_channels,
-                 pretrained_model:nn.Module=None,
-                 reshape=False,
-                 n_channels=None,
-                 dropout=0.,
-                 pretrained_config=None):
-        super().__init__()
-        if pretrained_config is None:
-            assert pretrained_model is not None, 'Either "pretrained_model" or "pretrained_config" must not be None'
-            self.pretrained_model = pretrained_model
-        else:
-            assert pretrained_config is not None, 'Either "pretrained_model" or "pretrained_config" must not be None'
-            self.instantiate_pretrained(pretrained_config)
-
-        self.do_reshape = reshape
-
-        if n_channels is None:
-            n_channels = self.pretrained_model.encoder.ch
-
-        self.proj_norm = Normalize(in_channels,num_groups=in_channels//2)
-        self.proj = nn.Conv2d(in_channels,n_channels,kernel_size=3,
-                            stride=1,padding=1)
-
-        blocks = []
-        downs = []
-        ch_in = n_channels
-        for m in ch_mult:
-            blocks.append(ResnetBlock(in_channels=ch_in,out_channels=m*n_channels,dropout=dropout))
-            ch_in = m * n_channels
-            downs.append(Downsample(ch_in, with_conv=False))
-
-        self.model = nn.ModuleList(blocks)
-        self.downsampler = nn.ModuleList(downs)
-
-
-    def instantiate_pretrained(self, config):
-        model = instantiate_from_config(config)
-        self.pretrained_model = model.eval()
-        # self.pretrained_model.train = False
-        for param in self.pretrained_model.parameters():
-            param.requires_grad = False
-
-
-    @torch.no_grad()
-    def encode_with_pretrained(self,x):
-        c = self.pretrained_model.encode(x)
-        if isinstance(c, DiagonalGaussianDistribution):
-            c = c.mode()
-        return  c
-
-    def forward(self,x):
-        z_fs = self.encode_with_pretrained(x)
-        z = self.proj_norm(z_fs)
-        z = self.proj(z)
-        z = nonlinearity(z)
-
-        for submodel, downmodel in zip(self.model,self.downsampler):
-            z = submodel(z,temb=None)
-            z = downmodel(z)
-
-        if self.do_reshape:
-            z = rearrange(z,'b c h w -> b (h w) c')
-        return z
diff --git a/instructany2pix/llm/model/vae/modules/blocks_spec.py b/instructany2pix/llm/model/vae/modules/blocks_spec.py
deleted file mode 100644
index 1fb49165a83aeb2e066c98af70fe12ba61411d0a..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/vae/modules/blocks_spec.py
+++ /dev/null
@@ -1,999 +0,0 @@
-# pytorch_diffusion + derived encoder decoder
-import math
-import torch
-import torch.nn as nn
-import numpy as np
-
-
-def get_timestep_embedding(timesteps, embedding_dim):
-    """
-    This matches the implementation in Denoising Diffusion Probabilistic Models:
-    From Fairseq.
-    Build sinusoidal embeddings.
-    This matches the implementation in tensor2tensor, but differs slightly
-    from the description in Section 3.5 of "Attention Is All You Need".
-    """
-    assert len(timesteps.shape) == 1
-
-    half_dim = embedding_dim // 2
-    emb = math.log(10000) / (half_dim - 1)
-    emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
-    emb = emb.to(device=timesteps.device)
-    emb = timesteps.float()[:, None] * emb[None, :]
-    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
-    if embedding_dim % 2 == 1:  # zero pad
-        emb = torch.nn.functional.pad(emb, (0,1,0,0))
-    return emb
-
-
-def nonlinearity(x):
-    # swish
-    return x*torch.sigmoid(x)
-
-
-def Normalize(in_channels):
-    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
-
-class Upsample(nn.Module):
-    def __init__(self, in_channels, with_conv):
-        super().__init__()
-        self.with_conv = with_conv
-        if self.with_conv:
-            self.conv = torch.nn.Conv2d(in_channels,
-                                        in_channels,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-    def forward(self, x):
-        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
-        if self.with_conv:
-            x = self.conv(x)
-        return x
-
-class Upsample1d(Upsample):
-    def __init__(self, in_channels, with_conv):
-        super().__init__(in_channels, with_conv)
-        if self.with_conv:
-            self.conv = torch.nn.Conv1d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
-
-class Downsample(nn.Module):
-    def __init__(self, in_channels, with_conv):
-        super().__init__()
-        self.with_conv = with_conv
-        if self.with_conv:
-            # no asymmetric padding in torch conv, must do it ourselves
-            self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
-            self.pad = (0, 1, 0, 1)
-        else:
-            self.avg_pool = nn.AvgPool2d(kernel_size=2, stride=2)
-
-    def forward(self, x):
-        if self.with_conv:  # bp: check self.avgpool and self.pad
-            x = torch.nn.functional.pad(x, self.pad, mode="constant", value=0)
-            x = self.conv(x)
-        else:
-            x = self.avg_pool(x)
-        return x
-
-class Downsample1d(Downsample):
-
-    def __init__(self, in_channels, with_conv):
-        super().__init__(in_channels, with_conv)
-        if self.with_conv:
-            # no asymmetric padding in torch conv, must do it ourselves
-            # TODO: can we replace it just with conv2d with padding 1?
-            self.conv = torch.nn.Conv1d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
-            self.pad = (1, 1)
-        else:
-            self.avg_pool = nn.AvgPool1d(kernel_size=2, stride=2)
-
-
-class ResnetBlock(nn.Module):
-    def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False,
-                 dropout, temb_channels=512):
-        super().__init__()
-        self.in_channels = in_channels
-        out_channels = in_channels if out_channels is None else out_channels
-        self.out_channels = out_channels
-        self.use_conv_shortcut = conv_shortcut
-
-        self.norm1 = Normalize(in_channels)
-        self.conv1 = torch.nn.Conv2d(in_channels,
-                                     out_channels,
-                                     kernel_size=3,
-                                     stride=1,
-                                     padding=1)
-        if temb_channels > 0:
-            self.temb_proj = torch.nn.Linear(temb_channels,
-                                             out_channels)
-        self.norm2 = Normalize(out_channels)
-        self.dropout = torch.nn.Dropout(dropout)
-        self.conv2 = torch.nn.Conv2d(out_channels,
-                                     out_channels,
-                                     kernel_size=3,
-                                     stride=1,
-                                     padding=1)
-        if self.in_channels != self.out_channels:
-            if self.use_conv_shortcut:
-                self.conv_shortcut = torch.nn.Conv2d(in_channels,
-                                                     out_channels,
-                                                     kernel_size=3,
-                                                     stride=1,
-                                                     padding=1)
-            else:
-                self.nin_shortcut = torch.nn.Conv2d(in_channels,
-                                                    out_channels,
-                                                    kernel_size=1,
-                                                    stride=1,
-                                                    padding=0)
-
-    def forward(self, x, temb):
-        h = x
-        h = self.norm1(h)
-        h = nonlinearity(h)
-        h = self.conv1(h)
-
-        if temb is not None:
-            h = h + self.temb_proj(nonlinearity(temb))[:,:,None,None]
-
-        h = self.norm2(h)
-        h = nonlinearity(h)
-        h = self.dropout(h)
-        h = self.conv2(h)
-
-        if self.in_channels != self.out_channels:
-            if self.use_conv_shortcut:
-                x = self.conv_shortcut(x)
-            else:
-                x = self.nin_shortcut(x)
-
-        return x+h
-
-class ResnetBlock1d(ResnetBlock):
-    def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False,
-                 dropout, temb_channels=512):
-        super().__init__(in_channels=in_channels, out_channels=out_channels,
-                         conv_shortcut=conv_shortcut, dropout=dropout, temb_channels=temb_channels)
-        # redefining different elements (forward is goint to be the same as in RenetBlock)
-        if temb_channels > 0:
-            raise NotImplementedError('go to ResnetBlock and figure out how to deal with it in forward')
-            self.temb_proj = torch.nn.Linear(temb_channels, out_channels)
-
-        self.conv1 = torch.nn.Conv1d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
-        self.conv2 = torch.nn.Conv1d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
-        if self.in_channels != self.out_channels:
-            if self.use_conv_shortcut:
-                self.conv_shortcut = torch.nn.Conv1d(in_channels, out_channels, kernel_size=3,
-                                                     stride=1, padding=1)
-            else:
-                self.nin_shortcut = torch.nn.Conv1d(in_channels, out_channels, kernel_size=1,
-                                                    stride=1, padding=0)
-
-
-class AttnBlock(nn.Module):
-    def __init__(self, in_channels):
-        super().__init__()
-        self.in_channels = in_channels
-
-        self.norm = Normalize(in_channels)
-        self.q = torch.nn.Conv2d(in_channels,
-                                 in_channels,
-                                 kernel_size=1,
-                                 stride=1,
-                                 padding=0)
-        self.k = torch.nn.Conv2d(in_channels,
-                                 in_channels,
-                                 kernel_size=1,
-                                 stride=1,
-                                 padding=0)
-        self.v = torch.nn.Conv2d(in_channels,
-                                 in_channels,
-                                 kernel_size=1,
-                                 stride=1,
-                                 padding=0)
-        self.proj_out = torch.nn.Conv2d(in_channels,
-                                        in_channels,
-                                        kernel_size=1,
-                                        stride=1,
-                                        padding=0)
-
-
-    def forward(self, x):
-        h_ = x
-        h_ = self.norm(h_)
-        q = self.q(h_)
-        k = self.k(h_)
-        v = self.v(h_)
-
-        # compute attention
-        b,c,h,w = q.shape
-        q = q.reshape(b,c,h*w)
-        q = q.permute(0,2,1)   # b,hw,c
-        k = k.reshape(b,c,h*w) # b,c,hw
-        w_ = torch.bmm(q,k)     # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
-        w_ = w_ * (int(c)**(-0.5))
-        w_ = torch.nn.functional.softmax(w_, dim=2)
-
-        # attend to values
-        v = v.reshape(b,c,h*w)
-        w_ = w_.permute(0,2,1)   # b,hw,hw (first hw of k, second of q)
-        h_ = torch.bmm(v,w_)     # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
-        h_ = h_.reshape(b,c,h,w)
-
-        h_ = self.proj_out(h_)
-
-        return x+h_
-
-class AttnBlock1d(nn.Module):
-
-    def __init__(self, in_channels):
-        super().__init__()
-        self.in_channels = in_channels
-
-        self.norm = Normalize(in_channels)
-        self.q = torch.nn.Conv1d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
-        self.k = torch.nn.Conv1d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
-        self.v = torch.nn.Conv1d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
-        self.proj_out = torch.nn.Conv1d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
-
-    def forward(self, x):
-        h_ = x
-        h_ = self.norm(h_)
-        q = self.q(h_)
-        k = self.k(h_)
-        v = self.v(h_)
-
-        # compute attention
-        b, c, t = q.shape
-        q = q.permute(0, 2, 1)   # b,t,c
-        w_ = torch.bmm(q, k)     # b,t,t    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
-        w_ = w_ * (int(c) ** (-0.5))
-        w_ = torch.nn.functional.softmax(w_, dim=2)
-
-        # attend to values
-        w_ = w_.permute(0, 2, 1)  # b,t,t (first t of k, second of q)
-        h_ = torch.bmm(v, w_)  # b,c,t (t of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
-
-        h_ = self.proj_out(h_)
-
-        return x + h_
-
-
-class Model(nn.Module):
-    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
-                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
-                 resolution, use_timestep=True):
-        super().__init__()
-        self.ch = ch
-        self.temb_ch = self.ch*4
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_channels = in_channels
-
-        self.use_timestep = use_timestep
-        if self.use_timestep:
-            # timestep embedding
-            self.temb = nn.Module()
-            self.temb.dense = nn.ModuleList([
-                torch.nn.Linear(self.ch,
-                                self.temb_ch),
-                torch.nn.Linear(self.temb_ch,
-                                self.temb_ch),
-            ])
-
-        # downsampling
-        self.conv_in = torch.nn.Conv2d(in_channels,
-                                       self.ch,
-                                       kernel_size=3,
-                                       stride=1,
-                                       padding=1)
-
-        curr_res = resolution
-        in_ch_mult = (1,)+tuple(ch_mult)
-        self.down = nn.ModuleList()
-        for i_level in range(self.num_resolutions):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_in = ch*in_ch_mult[i_level]
-            block_out = ch*ch_mult[i_level]
-            for i_block in range(self.num_res_blocks):
-                block.append(ResnetBlock(in_channels=block_in,
-                                         out_channels=block_out,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout))
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(AttnBlock(block_in))
-            down = nn.Module()
-            down.block = block
-            down.attn = attn
-            if i_level != self.num_resolutions-1:
-                down.downsample = Downsample(block_in, resamp_with_conv)
-                curr_res = curr_res // 2
-            self.down.append(down)
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-        self.mid.attn_1 = AttnBlock(block_in)
-        self.mid.block_2 = ResnetBlock(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-
-        # upsampling
-        self.up = nn.ModuleList()
-        for i_level in reversed(range(self.num_resolutions)):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_out = ch*ch_mult[i_level]
-            skip_in = ch*ch_mult[i_level]
-            for i_block in range(self.num_res_blocks+1):
-                if i_block == self.num_res_blocks:
-                    skip_in = ch*in_ch_mult[i_level]
-                block.append(ResnetBlock(in_channels=block_in+skip_in,
-                                         out_channels=block_out,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout))
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(AttnBlock(block_in))
-            up = nn.Module()
-            up.block = block
-            up.attn = attn
-            if i_level != 0:
-                up.upsample = Upsample(block_in, resamp_with_conv)
-                curr_res = curr_res * 2
-            self.up.insert(0, up) # prepend to get consistent order
-
-        # end
-        self.norm_out = Normalize(block_in)
-        self.conv_out = torch.nn.Conv2d(block_in,
-                                        out_ch,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-
-    def forward(self, x, t=None):
-        #assert x.shape[2] == x.shape[3] == self.resolution
-
-        if self.use_timestep:
-            # timestep embedding
-            assert t is not None
-            temb = get_timestep_embedding(t, self.ch)
-            temb = self.temb.dense[0](temb)
-            temb = nonlinearity(temb)
-            temb = self.temb.dense[1](temb)
-        else:
-            temb = None
-
-        # downsampling
-        hs = [self.conv_in(x)]
-        for i_level in range(self.num_resolutions):
-            for i_block in range(self.num_res_blocks):
-                h = self.down[i_level].block[i_block](hs[-1], temb)
-                if len(self.down[i_level].attn) > 0:
-                    h = self.down[i_level].attn[i_block](h)
-                hs.append(h)
-            if i_level != self.num_resolutions-1:
-                hs.append(self.down[i_level].downsample(hs[-1]))
-
-        # middle
-        h = hs[-1]
-        h = self.mid.block_1(h, temb)
-        h = self.mid.attn_1(h)
-        h = self.mid.block_2(h, temb)
-
-        # upsampling
-        for i_level in reversed(range(self.num_resolutions)):
-            for i_block in range(self.num_res_blocks+1):
-                h = self.up[i_level].block[i_block](
-                    torch.cat([h, hs.pop()], dim=1), temb)
-                if len(self.up[i_level].attn) > 0:
-                    h = self.up[i_level].attn[i_block](h)
-            if i_level != 0:
-                h = self.up[i_level].upsample(h)
-
-        # end
-        h = self.norm_out(h)
-        h = nonlinearity(h)
-        h = self.conv_out(h)
-        return h
-
-
-class Encoder(nn.Module):
-    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
-                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
-                 resolution, z_channels, double_z=True, **ignore_kwargs):
-        super().__init__()
-        self.ch = ch
-        self.temb_ch = 0
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_channels = in_channels
-
-        # downsampling
-        self.conv_in = torch.nn.Conv2d(in_channels,
-                                       self.ch,
-                                       kernel_size=3,
-                                       stride=1,
-                                       padding=1)
-
-        curr_res = resolution
-        in_ch_mult = (1,)+tuple(ch_mult)
-        self.down = nn.ModuleList()
-        for i_level in range(self.num_resolutions):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_in = ch*in_ch_mult[i_level]
-            block_out = ch*ch_mult[i_level]
-            for i_block in range(self.num_res_blocks):
-                block.append(ResnetBlock(in_channels=block_in,
-                                         out_channels=block_out,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout))
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(AttnBlock(block_in))
-            down = nn.Module()
-            down.block = block
-            down.attn = attn
-            if i_level != self.num_resolutions-1:
-                down.downsample = Downsample(block_in, resamp_with_conv)
-                curr_res = curr_res // 2
-            self.down.append(down)
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-        self.mid.attn_1 = AttnBlock(block_in)
-        self.mid.block_2 = ResnetBlock(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-
-        # end
-        self.norm_out = Normalize(block_in)
-        self.conv_out = torch.nn.Conv2d(block_in,
-                                        2*z_channels if double_z else z_channels,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-
-    def forward(self, x):
-        #assert x.shape[2] == x.shape[3] == self.resolution, "{}, {}, {}".format(x.shape[2], x.shape[3], self.resolution)
-
-        # timestep embedding
-        temb = None
-
-        # downsampling
-        hs = [self.conv_in(x)]
-        for i_level in range(self.num_resolutions):
-            for i_block in range(self.num_res_blocks):
-                h = self.down[i_level].block[i_block](hs[-1], temb)
-                if len(self.down[i_level].attn) > 0:
-                    h = self.down[i_level].attn[i_block](h)
-                hs.append(h)
-            if i_level != self.num_resolutions-1:
-                hs.append(self.down[i_level].downsample(hs[-1]))
-
-        # middle
-        h = hs[-1]
-        h = self.mid.block_1(h, temb)
-        h = self.mid.attn_1(h)
-        h = self.mid.block_2(h, temb)
-
-        # end
-        h = self.norm_out(h)
-        h = nonlinearity(h)
-        h = self.conv_out(h)
-        return h
-
-class Encoder1d(Encoder):
-    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
-                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
-                 resolution, z_channels, double_z=True, **ignore_kwargs):
-        super().__init__(ch=ch, out_ch=out_ch, ch_mult=ch_mult, num_res_blocks=num_res_blocks,
-                         attn_resolutions=attn_resolutions, dropout=dropout,
-                         resamp_with_conv=resamp_with_conv,
-                         in_channels=in_channels, resolution=resolution, z_channels=z_channels,
-                         double_z=double_z, **ignore_kwargs)
-        self.ch = ch
-        self.temb_ch = 0
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_channels = in_channels
-
-        # downsampling
-        self.conv_in = torch.nn.Conv1d(in_channels,
-                                       self.ch,
-                                       kernel_size=3,
-                                       stride=1,
-                                       padding=1)
-
-        curr_res = resolution
-        in_ch_mult = (1,)+tuple(ch_mult)
-        self.down = nn.ModuleList()
-        for i_level in range(self.num_resolutions):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_in = ch*in_ch_mult[i_level]
-            block_out = ch*ch_mult[i_level]
-            for i_block in range(self.num_res_blocks):
-                block.append(ResnetBlock1d(in_channels=block_in,
-                                           out_channels=block_out,
-                                           temb_channels=self.temb_ch,
-                                           dropout=dropout))
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(AttnBlock1d(block_in))
-            down = nn.Module()
-            down.block = block
-            down.attn = attn
-            if i_level != self.num_resolutions-1:
-                down.downsample = Downsample1d(block_in, resamp_with_conv)
-                curr_res = curr_res // 2
-            self.down.append(down)
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock1d(in_channels=block_in,
-                                         out_channels=block_in,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout)
-        self.mid.attn_1 = AttnBlock1d(block_in)
-        self.mid.block_2 = ResnetBlock1d(in_channels=block_in,
-                                         out_channels=block_in,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout)
-
-        # end
-        self.norm_out = Normalize(block_in)
-        self.conv_out = torch.nn.Conv1d(block_in,
-                                        2*z_channels if double_z else z_channels,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-
-class Decoder(nn.Module):
-    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
-                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
-                 resolution, z_channels, give_pre_end=False, **ignorekwargs):
-        super().__init__()
-        self.ch = ch
-        self.temb_ch = 0
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_channels = in_channels
-        self.give_pre_end = give_pre_end
-
-        # compute in_ch_mult, block_in and curr_res at lowest res
-        in_ch_mult = (1,)+tuple(ch_mult)
-        block_in = ch*ch_mult[self.num_resolutions-1]
-        curr_res = resolution // 2**(self.num_resolutions-1)
-        # self.z_shape = (1,z_channels,curr_res,curr_res)
-        # print("Working with z of shape {} = {} dimensions.".format(
-        #     self.z_shape, np.prod(self.z_shape)))
-
-        # z to block_in
-        self.conv_in = torch.nn.Conv2d(z_channels,
-                                       block_in,
-                                       kernel_size=3,
-                                       stride=1,
-                                       padding=1)
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-        self.mid.attn_1 = AttnBlock(block_in)
-        self.mid.block_2 = ResnetBlock(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-
-        # upsampling
-        self.up = nn.ModuleList()
-        for i_level in reversed(range(self.num_resolutions)):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_out = ch*ch_mult[i_level]
-            for i_block in range(self.num_res_blocks+1):
-                block.append(ResnetBlock(in_channels=block_in,
-                                         out_channels=block_out,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout))
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(AttnBlock(block_in))
-            up = nn.Module()
-            up.block = block
-            up.attn = attn
-            if i_level != 0:
-                up.upsample = Upsample(block_in, resamp_with_conv)
-                curr_res = curr_res * 2
-            self.up.insert(0, up) # prepend to get consistent order
-
-        # end
-        self.norm_out = Normalize(block_in)
-        self.conv_out = torch.nn.Conv2d(block_in,
-                                        out_ch,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-    def forward(self, z):
-        #assert z.shape[1:] == self.z_shape[1:]
-        self.last_z_shape = z.shape
-
-        # timestep embedding
-        temb = None
-
-        # z to block_in
-        h = self.conv_in(z)
-
-        # middle
-        h = self.mid.block_1(h, temb)
-        h = self.mid.attn_1(h)
-        h = self.mid.block_2(h, temb)
-
-        # upsampling
-        for i_level in reversed(range(self.num_resolutions)):
-            for i_block in range(self.num_res_blocks+1):
-                h = self.up[i_level].block[i_block](h, temb)
-                if len(self.up[i_level].attn) > 0:
-                    h = self.up[i_level].attn[i_block](h)
-            if i_level != 0:
-                h = self.up[i_level].upsample(h)
-
-        # end
-        if self.give_pre_end:
-            return h
-
-        h = self.norm_out(h)
-        h = nonlinearity(h)
-        h = self.conv_out(h)
-        return h
-
-class Decoder1d(Decoder):
-    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
-                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
-                 resolution, z_channels, give_pre_end=False, **ignorekwargs):
-        super().__init__(ch=ch, out_ch=out_ch, ch_mult=ch_mult, num_res_blocks=num_res_blocks,
-                         attn_resolutions=attn_resolutions, dropout=dropout,
-                         resamp_with_conv=resamp_with_conv,
-                         in_channels=in_channels, resolution=resolution, z_channels=z_channels,
-                         give_pre_end=give_pre_end, **ignorekwargs)
-        self.ch = ch
-        self.temb_ch = 0
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_channels = in_channels
-        self.give_pre_end = give_pre_end
-
-        # compute in_ch_mult, block_in and curr_res at lowest res
-        in_ch_mult = (1,) + tuple(ch_mult)
-        block_in = ch * ch_mult[self.num_resolutions-1]
-        curr_res = resolution // 2**(self.num_resolutions-1)
-        # self.z_shape = (1,z_channels,curr_res,curr_res)
-        # print("Working with z of shape {} = {} dimensions.".format(
-        #     self.z_shape, np.prod(self.z_shape)))
-
-        # z to block_in
-        self.conv_in = torch.nn.Conv1d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock1d(in_channels=block_in, out_channels=block_in,
-                                         temb_channels=self.temb_ch, dropout=dropout)
-        self.mid.attn_1 = AttnBlock1d(block_in)
-        self.mid.block_2 = ResnetBlock1d(in_channels=block_in, out_channels=block_in,
-                                         temb_channels=self.temb_ch, dropout=dropout)
-
-        # upsampling
-        self.up = nn.ModuleList()
-        for i_level in reversed(range(self.num_resolutions)):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_out = ch * ch_mult[i_level]
-            for i_block in range(self.num_res_blocks+1):
-                block.append(ResnetBlock1d(in_channels=block_in, out_channels=block_out,
-                                           temb_channels=self.temb_ch, dropout=dropout))
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(AttnBlock1d(block_in))
-            up = nn.Module()
-            up.block = block
-            up.attn = attn
-            if i_level != 0:
-                up.upsample = Upsample1d(block_in, resamp_with_conv)
-                curr_res = curr_res * 2
-            self.up.insert(0, up)  # prepend to get consistent order
-
-        # end
-        self.norm_out = Normalize(block_in)
-        self.conv_out = torch.nn.Conv1d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
-
-
-class VUNet(nn.Module):
-    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
-                 attn_resolutions, dropout=0.0, resamp_with_conv=True,
-                 in_channels, c_channels,
-                 resolution, z_channels, use_timestep=False, **ignore_kwargs):
-        super().__init__()
-        self.ch = ch
-        self.temb_ch = self.ch*4
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-
-        self.use_timestep = use_timestep
-        if self.use_timestep:
-            # timestep embedding
-            self.temb = nn.Module()
-            self.temb.dense = nn.ModuleList([
-                torch.nn.Linear(self.ch,
-                                self.temb_ch),
-                torch.nn.Linear(self.temb_ch,
-                                self.temb_ch),
-            ])
-
-        # downsampling
-        self.conv_in = torch.nn.Conv2d(c_channels,
-                                       self.ch,
-                                       kernel_size=3,
-                                       stride=1,
-                                       padding=1)
-
-        curr_res = resolution
-        in_ch_mult = (1,)+tuple(ch_mult)
-        self.down = nn.ModuleList()
-        for i_level in range(self.num_resolutions):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_in = ch*in_ch_mult[i_level]
-            block_out = ch*ch_mult[i_level]
-            for i_block in range(self.num_res_blocks):
-                block.append(ResnetBlock(in_channels=block_in,
-                                         out_channels=block_out,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout))
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(AttnBlock(block_in))
-            down = nn.Module()
-            down.block = block
-            down.attn = attn
-            if i_level != self.num_resolutions-1:
-                down.downsample = Downsample(block_in, resamp_with_conv)
-                curr_res = curr_res // 2
-            self.down.append(down)
-
-        self.z_in = torch.nn.Conv2d(z_channels,
-                                    block_in,
-                                    kernel_size=1,
-                                    stride=1,
-                                    padding=0)
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=2*block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-        self.mid.attn_1 = AttnBlock(block_in)
-        self.mid.block_2 = ResnetBlock(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-
-        # upsampling
-        self.up = nn.ModuleList()
-        for i_level in reversed(range(self.num_resolutions)):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_out = ch*ch_mult[i_level]
-            skip_in = ch*ch_mult[i_level]
-            for i_block in range(self.num_res_blocks+1):
-                if i_block == self.num_res_blocks:
-                    skip_in = ch*in_ch_mult[i_level]
-                block.append(ResnetBlock(in_channels=block_in+skip_in,
-                                         out_channels=block_out,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout))
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(AttnBlock(block_in))
-            up = nn.Module()
-            up.block = block
-            up.attn = attn
-            if i_level != 0:
-                up.upsample = Upsample(block_in, resamp_with_conv)
-                curr_res = curr_res * 2
-            self.up.insert(0, up) # prepend to get consistent order
-
-        # end
-        self.norm_out = Normalize(block_in)
-        self.conv_out = torch.nn.Conv2d(block_in,
-                                        out_ch,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-
-    def forward(self, x, z):
-        #assert x.shape[2] == x.shape[3] == self.resolution
-
-        if self.use_timestep:
-            # timestep embedding
-            assert t is not None
-            temb = get_timestep_embedding(t, self.ch)
-            temb = self.temb.dense[0](temb)
-            temb = nonlinearity(temb)
-            temb = self.temb.dense[1](temb)
-        else:
-            temb = None
-
-        # downsampling
-        hs = [self.conv_in(x)]
-        for i_level in range(self.num_resolutions):
-            for i_block in range(self.num_res_blocks):
-                h = self.down[i_level].block[i_block](hs[-1], temb)
-                if len(self.down[i_level].attn) > 0:
-                    h = self.down[i_level].attn[i_block](h)
-                hs.append(h)
-            if i_level != self.num_resolutions-1:
-                hs.append(self.down[i_level].downsample(hs[-1]))
-
-        # middle
-        h = hs[-1]
-        z = self.z_in(z)
-        h = torch.cat((h,z),dim=1)
-        h = self.mid.block_1(h, temb)
-        h = self.mid.attn_1(h)
-        h = self.mid.block_2(h, temb)
-
-        # upsampling
-        for i_level in reversed(range(self.num_resolutions)):
-            for i_block in range(self.num_res_blocks+1):
-                h = self.up[i_level].block[i_block](
-                    torch.cat([h, hs.pop()], dim=1), temb)
-                if len(self.up[i_level].attn) > 0:
-                    h = self.up[i_level].attn[i_block](h)
-            if i_level != 0:
-                h = self.up[i_level].upsample(h)
-
-        # end
-        h = self.norm_out(h)
-        h = nonlinearity(h)
-        h = self.conv_out(h)
-        return h
-
-
-class SimpleDecoder(nn.Module):
-    def __init__(self, in_channels, out_channels, *args, **kwargs):
-        super().__init__()
-        self.model = nn.ModuleList([nn.Conv2d(in_channels, in_channels, 1),
-                                     ResnetBlock(in_channels=in_channels,
-                                                 out_channels=2 * in_channels,
-                                                 temb_channels=0, dropout=0.0),
-                                     ResnetBlock(in_channels=2 * in_channels,
-                                                out_channels=4 * in_channels,
-                                                temb_channels=0, dropout=0.0),
-                                     ResnetBlock(in_channels=4 * in_channels,
-                                                out_channels=2 * in_channels,
-                                                temb_channels=0, dropout=0.0),
-                                     nn.Conv2d(2*in_channels, in_channels, 1),
-                                     Upsample(in_channels, with_conv=True)])
-        # end
-        self.norm_out = Normalize(in_channels)
-        self.conv_out = torch.nn.Conv2d(in_channels,
-                                        out_channels,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-    def forward(self, x):
-        for i, layer in enumerate(self.model):
-            if i in [1,2,3]:
-                x = layer(x, None)
-            else:
-                x = layer(x)
-
-        h = self.norm_out(x)
-        h = nonlinearity(h)
-        x = self.conv_out(h)
-        return x
-
-
-class UpsampleDecoder(nn.Module):
-    def __init__(self, in_channels, out_channels, ch, num_res_blocks, resolution,
-                 ch_mult=(2,2), dropout=0.0):
-        super().__init__()
-        # upsampling
-        self.temb_ch = 0
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        block_in = in_channels
-        curr_res = resolution // 2 ** (self.num_resolutions - 1)
-        self.res_blocks = nn.ModuleList()
-        self.upsample_blocks = nn.ModuleList()
-        for i_level in range(self.num_resolutions):
-            res_block = []
-            block_out = ch * ch_mult[i_level]
-            for i_block in range(self.num_res_blocks + 1):
-                res_block.append(ResnetBlock(in_channels=block_in,
-                                         out_channels=block_out,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout))
-                block_in = block_out
-            self.res_blocks.append(nn.ModuleList(res_block))
-            if i_level != self.num_resolutions - 1:
-                self.upsample_blocks.append(Upsample(block_in, True))
-                curr_res = curr_res * 2
-
-        # end
-        self.norm_out = Normalize(block_in)
-        self.conv_out = torch.nn.Conv2d(block_in,
-                                        out_channels,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-    def forward(self, x):
-        # upsampling
-        h = x
-        for k, i_level in enumerate(range(self.num_resolutions)):
-            for i_block in range(self.num_res_blocks + 1):
-                h = self.res_blocks[i_level][i_block](h, None)
-            if i_level != self.num_resolutions - 1:
-                h = self.upsample_blocks[k](h)
-        h = self.norm_out(h)
-        h = nonlinearity(h)
-        h = self.conv_out(h)
-        return h
-
-
-if __name__ == '__main__':
-    ddconfig = {
-        'ch': 128,
-        'num_res_blocks': 2,
-        'dropout': 0.0,
-        'z_channels': 256,
-        'double_z': False,
-    }
-
-    # Audio example ##
-    ddconfig['in_channels'] = 1
-    ddconfig['resolution'] = 848
-    ddconfig['attn_resolutions'] = [53]
-    ddconfig['ch_mult'] = [1, 1, 2, 2, 4]
-    ddconfig['out_ch'] = 1
-    # input
-    inputs = torch.rand(4, 1, 80, 848)
-    print('Input:', inputs.shape)
-    # Encoder
-    encoder = Encoder(**ddconfig)
-    enc_outs = encoder(inputs)
-    print('Encoder out:', enc_outs.shape)
-    # Decoder
-    decoder = Decoder(**ddconfig)
-    quant_outs = torch.rand(4, 256, 5, 53)
-    dec_outs = decoder(quant_outs)
-    print('Decoder out:', dec_outs.shape)
\ No newline at end of file
diff --git a/instructany2pix/llm/model/vae/modules/discriminator_spec.py b/instructany2pix/llm/model/vae/modules/discriminator_spec.py
deleted file mode 100644
index 51b0768855814c38e4796a745512b634f0417afd..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/vae/modules/discriminator_spec.py
+++ /dev/null
@@ -1,295 +0,0 @@
-import functools
-import torch.nn as nn
-
-
-class ActNorm(nn.Module):
-    def __init__(self, num_features, logdet=False, affine=True,
-                 allow_reverse_init=False):
-        assert affine
-        super().__init__()
-        self.logdet = logdet
-        self.loc = nn.Parameter(torch.zeros(1, num_features, 1, 1))
-        self.scale = nn.Parameter(torch.ones(1, num_features, 1, 1))
-        self.allow_reverse_init = allow_reverse_init
-
-        self.register_buffer('initialized', torch.tensor(0, dtype=torch.uint8))
-
-    def initialize(self, input):
-        with torch.no_grad():
-            flatten = input.permute(1, 0, 2, 3).contiguous().view(input.shape[1], -1)
-            mean = (
-                flatten.mean(1)
-                .unsqueeze(1)
-                .unsqueeze(2)
-                .unsqueeze(3)
-                .permute(1, 0, 2, 3)
-            )
-            std = (
-                flatten.std(1)
-                .unsqueeze(1)
-                .unsqueeze(2)
-                .unsqueeze(3)
-                .permute(1, 0, 2, 3)
-            )
-
-            self.loc.data.copy_(-mean)
-            self.scale.data.copy_(1 / (std + 1e-6))
-
-    def forward(self, input, reverse=False):
-        if reverse:
-            return self.reverse(input)
-        if len(input.shape) == 2:
-            input = input[:, :, None, None]
-            squeeze = True
-        else:
-            squeeze = False
-
-        _, _, height, width = input.shape
-
-        if self.training and self.initialized.item() == 0:
-            self.initialize(input)
-            self.initialized.fill_(1)
-
-        h = self.scale * (input + self.loc)
-
-        if squeeze:
-            h = h.squeeze(-1).squeeze(-1)
-
-        if self.logdet:
-            log_abs = torch.log(torch.abs(self.scale))
-            logdet = height * width * torch.sum(log_abs)
-            logdet = logdet * torch.ones(input.shape[0]).to(input)
-            return h, logdet
-
-        return h
-
-    def reverse(self, output):
-        if self.training and self.initialized.item() == 0:
-            if not self.allow_reverse_init:
-                raise RuntimeError(
-                    "Initializing ActNorm in reverse direction is "
-                    "disabled by default. Use allow_reverse_init=True to enable."
-                )
-            else:
-                self.initialize(output)
-                self.initialized.fill_(1)
-
-        if len(output.shape) == 2:
-            output = output[:, :, None, None]
-            squeeze = True
-        else:
-            squeeze = False
-
-        h = output / self.scale - self.loc
-
-        if squeeze:
-            h = h.squeeze(-1).squeeze(-1)
-        return h
-
-def weights_init(m):
-    classname = m.__class__.__name__
-    if classname.find('Conv') != -1:
-        nn.init.normal_(m.weight.data, 0.0, 0.02)
-    elif classname.find('BatchNorm') != -1:
-        nn.init.normal_(m.weight.data, 1.0, 0.02)
-        nn.init.constant_(m.bias.data, 0)
-
-
-class NLayerDiscriminator(nn.Module):
-    """Defines a PatchGAN discriminator as in Pix2Pix
-        --> see https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py
-    """
-    def __init__(self, input_nc=3, ndf=64, n_layers=3, use_actnorm=False):
-        """Construct a PatchGAN discriminator
-        Parameters:
-            input_nc (int)  -- the number of channels in input images
-            ndf (int)       -- the number of filters in the last conv layer
-            n_layers (int)  -- the number of conv layers in the discriminator
-            norm_layer      -- normalization layer
-        """
-        super(NLayerDiscriminator, self).__init__()
-        if not use_actnorm:
-            norm_layer = nn.BatchNorm2d
-        else:
-            norm_layer = ActNorm
-        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters
-            use_bias = norm_layer.func != nn.BatchNorm2d
-        else:
-            use_bias = norm_layer != nn.BatchNorm2d
-
-        kw = 4
-        padw = 1
-        sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
-        nf_mult = 1
-        nf_mult_prev = 1
-        for n in range(1, n_layers):  # gradually increase the number of filters
-            nf_mult_prev = nf_mult
-            nf_mult = min(2 ** n, 8)
-            sequence += [
-                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),
-                norm_layer(ndf * nf_mult),
-                nn.LeakyReLU(0.2, True)
-            ]
-
-        nf_mult_prev = nf_mult
-        nf_mult = min(2 ** n_layers, 8)
-        sequence += [
-            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
-            norm_layer(ndf * nf_mult),
-            nn.LeakyReLU(0.2, True)
-        ]
-        # output 1 channel prediction map
-        sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]
-        self.main = nn.Sequential(*sequence)
-
-    def forward(self, input):
-        """Standard forward."""
-        return self.main(input)
-
-class NLayerDiscriminator1dFeats(NLayerDiscriminator):
-    """Defines a PatchGAN discriminator as in Pix2Pix
-        --> see https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py
-    """
-    def __init__(self, input_nc=3, ndf=64, n_layers=3, use_actnorm=False):
-        """Construct a PatchGAN discriminator
-        Parameters:
-            input_nc (int)  -- the number of channels in input feats
-            ndf (int)       -- the number of filters in the last conv layer
-            n_layers (int)  -- the number of conv layers in the discriminator
-            norm_layer      -- normalization layer
-        """
-        super().__init__(input_nc=input_nc, ndf=64, n_layers=n_layers, use_actnorm=use_actnorm)
-
-        if not use_actnorm:
-            norm_layer = nn.BatchNorm1d
-        else:
-            norm_layer = ActNorm
-        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm has affine parameters
-            use_bias = norm_layer.func != nn.BatchNorm1d
-        else:
-            use_bias = norm_layer != nn.BatchNorm1d
-
-        kw = 4
-        padw = 1
-        sequence = [nn.Conv1d(input_nc, input_nc//2, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
-        nf_mult = input_nc//2
-        nf_mult_prev = 1
-        for n in range(1, n_layers):  # gradually decrease the number of filters
-            nf_mult_prev = nf_mult
-            nf_mult = max(nf_mult_prev // (2 ** n), 8)
-            sequence += [
-                nn.Conv1d(nf_mult_prev, nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),
-                norm_layer(nf_mult),
-                nn.LeakyReLU(0.2, True)
-            ]
-
-        nf_mult_prev = nf_mult
-        nf_mult = max(nf_mult_prev // (2 ** n), 8)
-        sequence += [
-            nn.Conv1d(nf_mult_prev, nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
-            norm_layer(nf_mult),
-            nn.LeakyReLU(0.2, True)
-        ]
-        nf_mult_prev = nf_mult
-        nf_mult = max(nf_mult_prev // (2 ** n), 8)
-        sequence += [
-            nn.Conv1d(nf_mult_prev, nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
-            norm_layer(nf_mult),
-            nn.LeakyReLU(0.2, True)
-        ]
-        # output 1 channel prediction map
-        sequence += [nn.Conv1d(nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]
-        self.main = nn.Sequential(*sequence)
-
-
-class NLayerDiscriminator1dSpecs(NLayerDiscriminator):
-    """Defines a PatchGAN discriminator as in Pix2Pix
-        --> see https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py
-    """
-    def __init__(self, input_nc=80, ndf=64, n_layers=3, use_actnorm=False):
-        """Construct a PatchGAN discriminator
-        Parameters:
-            input_nc (int)  -- the number of channels in input specs
-            ndf (int)       -- the number of filters in the last conv layer
-            n_layers (int)  -- the number of conv layers in the discriminator
-            norm_layer      -- normalization layer
-        """
-        super().__init__(input_nc=input_nc, ndf=64, n_layers=n_layers, use_actnorm=use_actnorm)
-
-        if not use_actnorm:
-            norm_layer = nn.BatchNorm1d
-        else:
-            norm_layer = ActNorm
-        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm has affine parameters
-            use_bias = norm_layer.func != nn.BatchNorm1d
-        else:
-            use_bias = norm_layer != nn.BatchNorm1d
-
-        kw = 4
-        padw = 1
-        sequence = [nn.Conv1d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
-        nf_mult = 1
-        nf_mult_prev = 1
-        for n in range(1, n_layers):  # gradually decrease the number of filters
-            nf_mult_prev = nf_mult
-            nf_mult = min(2 ** n, 8)
-            sequence += [
-                nn.Conv1d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),
-                norm_layer(ndf * nf_mult),
-                nn.LeakyReLU(0.2, True)
-            ]
-
-        nf_mult_prev = nf_mult
-        nf_mult = min(2 ** n_layers, 8)
-        sequence += [
-            nn.Conv1d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
-            norm_layer(ndf * nf_mult),
-            nn.LeakyReLU(0.2, True)
-        ]
-        # output 1 channel prediction map
-        sequence += [nn.Conv1d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]
-        self.main = nn.Sequential(*sequence)
-
-    def forward(self, input):
-        """Standard forward."""
-        # (B, C, L)
-        input = input.squeeze(1)
-        input = self.main(input)
-        return input
-
-
-if __name__ == '__main__':
-    import torch
-
-    ## FEATURES
-    disc_in_channels = 2048
-    disc_num_layers = 2
-    use_actnorm = False
-    disc_ndf = 64
-    discriminator = NLayerDiscriminator1dFeats(input_nc=disc_in_channels, n_layers=disc_num_layers,
-                                            use_actnorm=use_actnorm, ndf=disc_ndf).apply(weights_init)
-    inputs = torch.rand((6, 2048, 212))
-    outputs = discriminator(inputs)
-    print(outputs.shape)
-
-    ## AUDIO
-    disc_in_channels = 1
-    disc_num_layers = 3
-    use_actnorm = False
-    disc_ndf = 64
-    discriminator = NLayerDiscriminator(input_nc=disc_in_channels, n_layers=disc_num_layers,
-                                        use_actnorm=use_actnorm, ndf=disc_ndf).apply(weights_init)
-    inputs = torch.rand((6, 1, 80, 848))
-    outputs = discriminator(inputs)
-    print(outputs.shape)
-
-    ## IMAGE
-    disc_in_channels = 3
-    disc_num_layers = 3
-    use_actnorm = False
-    disc_ndf = 64
-    discriminator = NLayerDiscriminator(input_nc=disc_in_channels, n_layers=disc_num_layers,
-                                        use_actnorm=use_actnorm, ndf=disc_ndf).apply(weights_init)
-    inputs = torch.rand((6, 3, 256, 256))
-    outputs = discriminator(inputs)
-    print(outputs.shape)
\ No newline at end of file
diff --git a/instructany2pix/llm/model/vae/modules/distributions.py b/instructany2pix/llm/model/vae/modules/distributions.py
deleted file mode 100644
index 9dcf233a296d0b2a727c9073962bf419786b3880..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/vae/modules/distributions.py
+++ /dev/null
@@ -1,92 +0,0 @@
-import torch
-import numpy as np
-
-
-class AbstractDistribution:
-    def sample(self):
-        raise NotImplementedError()
-
-    def mode(self):
-        raise NotImplementedError()
-
-
-class DiracDistribution(AbstractDistribution):
-    def __init__(self, value):
-        self.value = value
-
-    def sample(self):
-        return self.value
-
-    def mode(self):
-        return self.value
-
-
-class DiagonalGaussianDistribution(object):
-    def __init__(self, parameters, deterministic=False):
-        self.parameters = parameters
-        self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
-        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
-        self.deterministic = deterministic
-        self.std = torch.exp(0.5 * self.logvar)
-        self.var = torch.exp(self.logvar)
-        if self.deterministic:
-            self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device)
-
-    def sample(self):
-        x = self.mean + self.std * torch.randn(self.mean.shape).to(device=self.parameters.device)
-        return x
-
-    def kl(self, other=None):
-        if self.deterministic:
-            return torch.Tensor([0.])
-        else:
-            if other is None:
-                return 0.5 * torch.sum(torch.pow(self.mean, 2)
-                                       + self.var - 1.0 - self.logvar,
-                                       dim=[1, 2, 3])
-            else:
-                return 0.5 * torch.sum(
-                    torch.pow(self.mean - other.mean, 2) / other.var
-                    + self.var / other.var - 1.0 - self.logvar + other.logvar,
-                    dim=[1, 2, 3])
-
-    def nll(self, sample, dims=[1,2,3]):
-        if self.deterministic:
-            return torch.Tensor([0.])
-        logtwopi = np.log(2.0 * np.pi)
-        return 0.5 * torch.sum(
-            logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
-            dim=dims)
-
-    def mode(self):
-        return self.mean
-
-
-def normal_kl(mean1, logvar1, mean2, logvar2):
-    """
-    source: https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/losses.py#L12
-    Compute the KL divergence between two gaussians.
-    Shapes are automatically broadcasted, so batches can be compared to
-    scalars, among other use cases.
-    """
-    tensor = None
-    for obj in (mean1, logvar1, mean2, logvar2):
-        if isinstance(obj, torch.Tensor):
-            tensor = obj
-            break
-    assert tensor is not None, "at least one argument must be a Tensor"
-
-    # Force variances to be Tensors. Broadcasting helps convert scalars to
-    # Tensors, but it does not work for torch.exp().
-    logvar1, logvar2 = [
-        x if isinstance(x, torch.Tensor) else torch.tensor(x).to(tensor)
-        for x in (logvar1, logvar2)
-    ]
-
-    return 0.5 * (
-        -1.0
-        + logvar2
-        - logvar1
-        + torch.exp(logvar1 - logvar2)
-        + ((mean1 - mean2) ** 2) * torch.exp(-logvar2)
-    )
\ No newline at end of file
diff --git a/instructany2pix/llm/model/vae/modules/extract_spec.py b/instructany2pix/llm/model/vae/modules/extract_spec.py
deleted file mode 100644
index a461c312a32373cf195786639935cff1cf693c8f..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/vae/modules/extract_spec.py
+++ /dev/null
@@ -1,236 +0,0 @@
-import argparse
-import os
-import os.path as P
-from copy import deepcopy
-from functools import partial
-from glob import glob
-from multiprocessing import Pool
-from pathlib import Path
-
-import librosa
-import numpy as np
-import torchvision
-import subprocess
-from .util import which_ffmpeg
-
-class MelSpectrogram(object):
-    def __init__(self, sr, nfft, fmin, fmax, nmels, hoplen, spec_power, inverse=False):
-        self.sr = sr
-        self.nfft = nfft
-        self.fmin = fmin
-        self.fmax = fmax
-        self.nmels = nmels
-        self.hoplen = hoplen
-        self.spec_power = spec_power
-        self.inverse = inverse
-
-        self.mel_basis = librosa.filters.mel(sr=sr, n_fft=nfft, fmin=fmin, fmax=fmax, n_mels=nmels)
-
-    def __call__(self, x):
-        if self.inverse:
-            spec = librosa.feature.inverse.mel_to_stft(
-                x, sr=self.sr, n_fft=self.nfft, fmin=self.fmin, fmax=self.fmax, power=self.spec_power
-            )
-            wav = librosa.griffinlim(spec, hop_length=self.hoplen)
-            return wav
-        else:
-            spec = np.abs(librosa.stft(x, n_fft=self.nfft, hop_length=self.hoplen)) ** self.spec_power
-            mel_spec = np.dot(self.mel_basis, spec)
-            return mel_spec
-
-class LowerThresh(object):
-    def __init__(self, min_val, inverse=False):
-        self.min_val = min_val
-        self.inverse = inverse
-
-    def __call__(self, x):
-        if self.inverse:
-            return x
-        else:
-            return np.maximum(self.min_val, x)
-
-class Add(object):
-    def __init__(self, val, inverse=False):
-        self.inverse = inverse
-        self.val = val
-
-    def __call__(self, x):
-        if self.inverse:
-            return x - self.val
-        else:
-            return x + self.val
-
-class Subtract(Add):
-    def __init__(self, val, inverse=False):
-        self.inverse = inverse
-        self.val = val
-
-    def __call__(self, x):
-        if self.inverse:
-            return x + self.val
-        else:
-            return x - self.val
-
-class Multiply(object):
-    def __init__(self, val, inverse=False) -> None:
-        self.val = val
-        self.inverse = inverse
-
-    def __call__(self, x):
-        if self.inverse:
-            return x / self.val
-        else:
-            return x * self.val
-
-class Divide(Multiply):
-    def __init__(self, val, inverse=False):
-        self.inverse = inverse
-        self.val = val
-
-    def __call__(self, x):
-        if self.inverse:
-            return x * self.val
-        else:
-            return x / self.val
-
-
-class Log10(object):
-    def __init__(self, inverse=False):
-        self.inverse = inverse
-
-    def __call__(self, x):
-        if self.inverse:
-            return 10 ** x
-        else:
-            return np.log10(x)
-
-class Clip(object):
-    def __init__(self, min_val, max_val, inverse=False):
-        self.min_val = min_val
-        self.max_val = max_val
-        self.inverse = inverse
-
-    def __call__(self, x):
-        if self.inverse:
-            return x
-        else:
-            return np.clip(x, self.min_val, self.max_val)
-
-class TrimSpec(object):
-    def __init__(self, max_len, inverse=False):
-        self.max_len = max_len
-        self.inverse = inverse
-
-    def __call__(self, x):
-        if self.inverse:
-            return x
-        else:
-            return x[:, :self.max_len]
-
-class MaxNorm(object):
-    def __init__(self, inverse=False):
-        self.inverse = inverse
-        self.eps = 1e-10
-
-    def __call__(self, x):
-        if self.inverse:
-            return x
-        else:
-            return x / (x.max() + self.eps)
-
-
-TRANSFORMS = torchvision.transforms.Compose([
-    MelSpectrogram(sr=22050, nfft=1024, fmin=125, fmax=7600, nmels=80, hoplen=1024//4, spec_power=1),
-    LowerThresh(1e-5),
-    Log10(),
-    Multiply(20),
-    Subtract(20),
-    Add(100),
-    Divide(100),
-    Clip(0, 1.0),
-    TrimSpec(860)
-])
-
-
-def inv_transforms(x, folder_name='melspec_10s_22050hz'):
-    '''relies on the GLOBAL contant TRANSFORMS which should be defined in this document'''
-    if folder_name == 'melspec_10s_22050hz':
-        i_transforms = deepcopy(TRANSFORMS.transforms[::-1])
-    else:
-        raise NotImplementedError
-    for t in i_transforms:
-        t.inverse = True
-    i_transforms = torchvision.transforms.Compose(i_transforms)
-    return i_transforms(x)
-
-
-def get_spectrogram(audio_path, save_dir, length, folder_name='melspec_10s_22050hz', save_results=True):
-    wav, _ = librosa.load(audio_path, sr=None)
-    # this cannot be a transform without creating a huge overhead with inserting audio_name in each
-    y = np.zeros(length)
-    if wav.shape[0] < length:
-        y[:len(wav)] = wav
-    else:
-        y = wav[:length]
-
-    if folder_name == 'melspec_10s_22050hz':
-        print('using', folder_name)
-        mel_spec = TRANSFORMS(y)
-    else:
-        raise NotImplementedError
-
-    if save_results:
-        os.makedirs(save_dir, exist_ok=True)
-        audio_name = os.path.basename(audio_path).split('.')[0]
-        np.save(P.join(save_dir, audio_name + '_mel.npy'), mel_spec)
-        np.save(P.join(save_dir, audio_name + '_audio.npy'), y)
-    else:
-        return y, mel_spec
-
-
-def extract_melspectrogram(in_path: str, sr: int, duration: int = 10, tmp_path: str = './tmp') -> np.ndarray:
-    '''Extract Melspectrogram similar to RegNet.'''
-    assert which_ffmpeg() != '', 'Is ffmpeg installed? Check if the conda environment is activated.'
-    # assert in_path.endswith('.mp4'), 'The file does not end with .mp4. Comment this if expected'
-    # create tmp dir if doesn't exist
-    os.makedirs(tmp_path, exist_ok=True)
-
-    # Extract audio from a video if needed
-    if in_path.endswith('.wav'):
-        audio_raw = in_path
-    else:
-        audio_raw = os.path.join(tmp_path, f'{Path(in_path).stem}.wav')
-        cmd = f'{which_ffmpeg()} -i {in_path} -hide_banner -loglevel panic -f wav -vn -y {audio_raw}'
-        subprocess.call(cmd.split())
-
-    # Extract audio from a video
-    audio_new = os.path.join(tmp_path, f'{Path(in_path).stem}_{sr}hz.wav')
-    cmd = f'{which_ffmpeg()} -i {audio_raw} -hide_banner -loglevel panic -ac 1 -ab 16k -ar {sr} -y {audio_new}'
-    subprocess.call(cmd.split())
-
-    length = int(duration * sr)
-    audio_zero_pad, spec = get_spectrogram(audio_new, save_dir=None, length=length, save_results=False)
-
-    # specvqgan expects inputs to be in [-1, 1] but spectrograms are in [0, 1]
-    spec = 2 * spec - 1
-
-    return spec
-
-if __name__ == '__main__':
-    paser = argparse.ArgumentParser()
-    paser.add_argument("-i", "--input_dir", default="data/features/dog/audio_10s_22050hz")
-    paser.add_argument("-o", "--output_dir", default="data/features/dog/melspec_10s_22050hz")
-    paser.add_argument("-l", "--length", default=220500)
-    paser.add_argument("-n", '--num_worker', type=int, default=32)
-    args = paser.parse_args()
-    input_dir = args.input_dir
-    output_dir = args.output_dir
-    length = args.length
-
-    audio_paths = glob(P.join(input_dir, "*.wav"))
-    audio_paths.sort()
-
-    with Pool(args.num_worker) as p:
-        p.map(partial(
-            get_spectrogram, save_dir=output_dir, length=length, folder_name=Path(output_dir).name
-        ), audio_paths)
\ No newline at end of file
diff --git a/instructany2pix/llm/model/vae/modules/helper.py b/instructany2pix/llm/model/vae/modules/helper.py
deleted file mode 100644
index cb2f4660f5a3344528573715f1c39170d11cb6f2..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/vae/modules/helper.py
+++ /dev/null
@@ -1,6 +0,0 @@
-from torch import nn
-
-class DummyLoss(nn.Module):
-    def __init__(self,*args,**kwargs):
-        super().__init__()
-
diff --git a/instructany2pix/llm/model/vae/modules/loss_spec.py b/instructany2pix/llm/model/vae/modules/loss_spec.py
deleted file mode 100644
index 89d7a8b7d8eb6fd348f820af6ba468c651338a4b..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/vae/modules/loss_spec.py
+++ /dev/null
@@ -1,209 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import sys
-
-sys.path.insert(0, '.')  # nopep8
-from .discriminator_spec import (NLayerDiscriminator, NLayerDiscriminator1dFeats,
-                                                   NLayerDiscriminator1dSpecs,
-                                                   weights_init)
-from .lpaps import LPAPS
-
-
-class DummyLoss(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-
-def adopt_weight(weight, global_step, threshold=0, value=0.):
-    if global_step < threshold:
-        weight = value
-    return weight
-
-
-def hinge_d_loss(logits_real, logits_fake):
-    loss_real = torch.mean(F.relu(1. - logits_real))
-    loss_fake = torch.mean(F.relu(1. + logits_fake))
-    d_loss = 0.5 * (loss_real + loss_fake)
-    return d_loss
-
-
-def vanilla_d_loss(logits_real, logits_fake):
-    d_loss = 0.5 * (
-        torch.mean(torch.nn.functional.softplus(-logits_real)) +
-        torch.mean(torch.nn.functional.softplus(logits_fake)))
-    return d_loss
-
-
-class VQLPAPSWithDiscriminator(nn.Module):
-    def __init__(self, disc_start, codebook_weight=1.0, pixelloss_weight=1.0,
-                 disc_num_layers=3, disc_in_channels=3, disc_factor=1.0, disc_weight=1.0,
-                 perceptual_weight=1.0, use_actnorm=False, disc_conditional=False,
-                 disc_ndf=64, disc_loss="hinge", min_adapt_weight=0.0, max_adapt_weight=1e4):
-        super().__init__()
-        assert disc_loss in ["hinge", "vanilla"]
-        self.codebook_weight = codebook_weight
-        self.pixel_weight = pixelloss_weight
-        self.perceptual_loss = LPAPS().eval()
-        self.perceptual_weight = perceptual_weight
-
-        self.discriminator = NLayerDiscriminator(input_nc=disc_in_channels,
-                                                 n_layers=disc_num_layers,
-                                                 use_actnorm=use_actnorm,
-                                                 ndf=disc_ndf
-                                                 ).apply(weights_init)
-        self.discriminator_iter_start = disc_start
-        if disc_loss == "hinge":
-            self.disc_loss = hinge_d_loss
-        elif disc_loss == "vanilla":
-            self.disc_loss = vanilla_d_loss
-        else:
-            raise ValueError(f"Unknown GAN loss '{disc_loss}'.")
-        print(f"VQLPAPSWithDiscriminator running with {disc_loss} loss.")
-        self.disc_factor = disc_factor
-        self.discriminator_weight = disc_weight
-        self.disc_conditional = disc_conditional
-        self.min_adapt_weight = min_adapt_weight
-        self.max_adapt_weight = max_adapt_weight
-
-    def calculate_adaptive_weight(self, nll_loss, g_loss, last_layer=None):
-        if last_layer is not None:
-            nll_grads = torch.autograd.grad(nll_loss, last_layer, retain_graph=True)[0]
-            g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0]
-        else:
-            nll_grads = torch.autograd.grad(nll_loss, self.last_layer[0], retain_graph=True)[0]
-            g_grads = torch.autograd.grad(g_loss, self.last_layer[0], retain_graph=True)[0]
-
-        d_weight = torch.norm(nll_grads) / (torch.norm(g_grads) + 1e-4)
-        d_weight = torch.clamp(d_weight, self.min_adapt_weight, self.max_adapt_weight).detach()
-        d_weight = d_weight * self.discriminator_weight
-        return d_weight
-
-    def forward(self, codebook_loss, inputs, reconstructions, optimizer_idx,
-                global_step, last_layer=None, cond=None, split="train"):
-        rec_loss = torch.abs(inputs.contiguous() - reconstructions.contiguous())
-        if self.perceptual_weight > 0:
-            p_loss = self.perceptual_loss(inputs.contiguous(), reconstructions.contiguous())
-            rec_loss = rec_loss + self.perceptual_weight * p_loss
-        else:
-            p_loss = torch.tensor([0.0])
-
-        nll_loss = rec_loss
-        # nll_loss = torch.sum(nll_loss) / nll_loss.shape[0]
-        nll_loss = torch.mean(nll_loss)
-
-        # now the GAN part
-        if optimizer_idx == 0:
-            # generator update
-            if cond is None:
-                assert not self.disc_conditional
-                logits_fake = self.discriminator(reconstructions.contiguous())
-            else:
-                assert self.disc_conditional
-                logits_fake = self.discriminator(torch.cat((reconstructions.contiguous(), cond), dim=1))
-            g_loss = -torch.mean(logits_fake)
-
-            try:
-                d_weight = self.calculate_adaptive_weight(nll_loss, g_loss, last_layer=last_layer)
-            except RuntimeError:
-                assert not self.training
-                d_weight = torch.tensor(0.0)
-
-            disc_factor = adopt_weight(self.disc_factor, global_step, threshold=self.discriminator_iter_start)
-            loss = nll_loss + d_weight * disc_factor * g_loss + self.codebook_weight * codebook_loss.mean()
-
-            log = {"{}/total_loss".format(split): loss.clone().detach().mean(),
-                   "{}/quant_loss".format(split): codebook_loss.detach().mean(),
-                   "{}/nll_loss".format(split): nll_loss.detach().mean(),
-                   "{}/rec_loss".format(split): rec_loss.detach().mean(),
-                   "{}/p_loss".format(split): p_loss.detach().mean(),
-                   "{}/d_weight".format(split): d_weight.detach(),
-                   "{}/disc_factor".format(split): torch.tensor(disc_factor),
-                   "{}/g_loss".format(split): g_loss.detach().mean(),
-                   }
-            return loss, log
-
-        if optimizer_idx == 1:
-            # second pass for discriminator update
-            if cond is None:
-                logits_real = self.discriminator(inputs.contiguous().detach())
-                logits_fake = self.discriminator(reconstructions.contiguous().detach())
-            else:
-                logits_real = self.discriminator(torch.cat((inputs.contiguous().detach(), cond), dim=1))
-                logits_fake = self.discriminator(torch.cat((reconstructions.contiguous().detach(), cond), dim=1))
-
-            disc_factor = adopt_weight(self.disc_factor, global_step, threshold=self.discriminator_iter_start)
-            d_loss = disc_factor * self.disc_loss(logits_real, logits_fake)
-
-            log = {"{}/disc_loss".format(split): d_loss.clone().detach().mean(),
-                   "{}/logits_real".format(split): logits_real.detach().mean(),
-                   "{}/logits_fake".format(split): logits_fake.detach().mean()
-                   }
-            return d_loss, log
-
-
-class VQLPAPSWithDiscriminator1dFeats(VQLPAPSWithDiscriminator):
-    def __init__(self, disc_start, codebook_weight=1.0, pixelloss_weight=1.0,
-                 disc_num_layers=3, disc_in_channels=3, disc_factor=1.0, disc_weight=1.0,
-                 perceptual_weight=1.0, use_actnorm=False, disc_conditional=False,
-                 disc_ndf=64, disc_loss="hinge", min_adapt_weight=0.0, max_adapt_weight=1e4):
-        super().__init__(disc_start=disc_start, codebook_weight=codebook_weight,
-                         pixelloss_weight=pixelloss_weight, disc_num_layers=disc_num_layers,
-                         disc_in_channels=disc_in_channels, disc_factor=disc_factor, disc_weight=disc_weight,
-                         perceptual_weight=perceptual_weight, use_actnorm=use_actnorm,
-                         disc_conditional=disc_conditional, disc_ndf=disc_ndf, disc_loss=disc_loss,
-                         min_adapt_weight=min_adapt_weight, max_adapt_weight=max_adapt_weight)
-
-        self.discriminator = NLayerDiscriminator1dFeats(input_nc=disc_in_channels, n_layers=disc_num_layers,
-                                                   use_actnorm=use_actnorm, ndf=disc_ndf).apply(weights_init)
-
-class VQLPAPSWithDiscriminator1dSpecs(VQLPAPSWithDiscriminator):
-    def __init__(self, disc_start, codebook_weight=1.0, pixelloss_weight=1.0,
-                 disc_num_layers=3, disc_in_channels=3, disc_factor=1.0, disc_weight=1.0,
-                 perceptual_weight=1.0, use_actnorm=False, disc_conditional=False,
-                 disc_ndf=64, disc_loss="hinge", min_adapt_weight=0.0, max_adapt_weight=1e4):
-        super().__init__(disc_start=disc_start, codebook_weight=codebook_weight,
-                         pixelloss_weight=pixelloss_weight, disc_num_layers=disc_num_layers,
-                         disc_in_channels=disc_in_channels, disc_factor=disc_factor, disc_weight=disc_weight,
-                         perceptual_weight=perceptual_weight, use_actnorm=use_actnorm,
-                         disc_conditional=disc_conditional, disc_ndf=disc_ndf, disc_loss=disc_loss,
-                         min_adapt_weight=min_adapt_weight, max_adapt_weight=max_adapt_weight)
-
-        self.discriminator = NLayerDiscriminator1dSpecs(input_nc=disc_in_channels, n_layers=disc_num_layers,
-                                                   use_actnorm=use_actnorm, ndf=disc_ndf).apply(weights_init)
-
-
-if __name__ == '__main__':
-    from specvqgan.modules.diffusionmodules.model import Decoder, Decoder1d
-
-    optimizer_idx = 0
-    loss_config = {
-        'disc_conditional': False,
-        'disc_start': 30001,
-        'disc_weight': 0.8,
-        'codebook_weight': 1.0,
-    }
-    ddconfig = {
-        'ch': 128,
-        'num_res_blocks': 2,
-        'dropout': 0.0,
-        'z_channels': 256,
-        'double_z': False,
-    }
-    qloss = torch.rand(1, requires_grad=True)
-
-    ## AUDIO
-    loss_config['disc_in_channels'] = 1
-    ddconfig['in_channels'] = 1
-    ddconfig['resolution'] = 848
-    ddconfig['attn_resolutions'] = [53]
-    ddconfig['out_ch'] = 1
-    ddconfig['ch_mult'] = [1, 1, 2, 2, 4]
-    decoder = Decoder(**ddconfig)
-    loss = VQLPAPSWithDiscriminator(**loss_config)
-    x = torch.rand(16, 1, 80, 848)
-    # subtracting something which uses dec_conv_out so that it will be in a graph
-    xrec = torch.rand(16, 1, 80, 848) - decoder.conv_out(torch.rand(16, 128, 80, 848)).mean()
-    aeloss, log_dict_ae = loss(qloss, x, xrec, optimizer_idx, global_step=0,last_layer=decoder.conv_out.weight)
-    print(aeloss)
-    print(log_dict_ae)
\ No newline at end of file
diff --git a/instructany2pix/llm/model/vae/modules/lpaps.py b/instructany2pix/llm/model/vae/modules/lpaps.py
deleted file mode 100644
index 65a6c6082376dd606fa542136d973d27912ad6a5..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/vae/modules/lpaps.py
+++ /dev/null
@@ -1,152 +0,0 @@
-"""
-    Based on https://github.com/CompVis/taming-transformers/blob/52720829/taming/modules/losses/lpips.py
-    Adapted for spectrograms by Vladimir Iashin (v-iashin)
-"""
-from collections import namedtuple
-
-import numpy as np
-import torch
-import torch.nn as nn
-
-import sys
-sys.path.insert(0, '.')  # nopep8
-from .vggishish import VGGishish
-from .util import get_ckpt_path
-
-
-class LPAPS(nn.Module):
-    # Learned perceptual metric
-    def __init__(self, use_dropout=True):
-        super().__init__()
-        self.scaling_layer = ScalingLayer()
-        self.chns = [64, 128, 256, 512, 512]  # vggish16 features
-        self.net = vggishish16(pretrained=True, requires_grad=False)
-        self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)
-        self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)
-        self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)
-        self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)
-        self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)
-        self.load_from_pretrained()
-        for param in self.parameters():
-            param.requires_grad = False
-
-    def load_from_pretrained(self, name="vggishish_lpaps"):
-        ckpt = get_ckpt_path(name, "specvqgan/modules/autoencoder/lpaps")
-        self.load_state_dict(torch.load(ckpt, map_location=torch.device("cpu")), strict=False)
-        print("loaded pretrained LPAPS loss from {}".format(ckpt))
-
-    @classmethod
-    def from_pretrained(cls, name="vggishish_lpaps"):
-        if name != "vggishish_lpaps":
-            raise NotImplementedError
-        model = cls()
-        ckpt = get_ckpt_path(name)
-        model.load_state_dict(torch.load(ckpt, map_location=torch.device("cpu")), strict=False)
-        return model
-
-    def forward(self, input, target):
-        in0_input, in1_input = (self.scaling_layer(input), self.scaling_layer(target))
-        outs0, outs1 = self.net(in0_input), self.net(in1_input)
-        feats0, feats1, diffs = {}, {}, {}
-        lins = [self.lin0, self.lin1, self.lin2, self.lin3, self.lin4]
-        for kk in range(len(self.chns)):
-            feats0[kk], feats1[kk] = normalize_tensor(outs0[kk]), normalize_tensor(outs1[kk])
-            diffs[kk] = (feats0[kk] - feats1[kk]) ** 2
-
-        res = [spatial_average(lins[kk].model(diffs[kk]), keepdim=True) for kk in range(len(self.chns))]
-        val = res[0]
-        for l in range(1, len(self.chns)):
-            val += res[l]
-        return val
-
-class ScalingLayer(nn.Module):
-    def __init__(self):
-        super(ScalingLayer, self).__init__()
-        # we are gonna use get_ckpt_path to donwload the stats as well
-        stat_path = get_ckpt_path('vggishish_mean_std_melspec_10s_22050hz', 'specvqgan/modules/autoencoder/lpaps')
-        # if for images we normalize on the channel dim, in spectrogram we will norm on frequency dimension
-        means, stds = np.loadtxt(stat_path, dtype=np.float32).T
-        # the normalization in means and stds are given for [0, 1], but specvqgan expects [-1, 1]:
-        means = 2 * means - 1
-        stds = 2 * stds
-        # input is expected to be (B, 1, F, T)
-        self.register_buffer('shift', torch.from_numpy(means)[None, None, :, None])
-        self.register_buffer('scale', torch.from_numpy(stds)[None, None, :, None])
-
-    def forward(self, inp):
-        return (inp - self.shift) / self.scale
-
-
-class NetLinLayer(nn.Module):
-    """ A single linear layer which does a 1x1 conv """
-    def __init__(self, chn_in, chn_out=1, use_dropout=False):
-        super(NetLinLayer, self).__init__()
-        layers = [nn.Dropout(), ] if (use_dropout) else []
-        layers += [nn.Conv2d(chn_in, chn_out, 1, stride=1, padding=0, bias=False), ]
-        self.model = nn.Sequential(*layers)
-
-class vggishish16(torch.nn.Module):
-    def __init__(self, requires_grad=False, pretrained=True):
-        super().__init__()
-        vgg_pretrained_features = self.vggishish16(pretrained=pretrained).features
-        self.slice1 = torch.nn.Sequential()
-        self.slice2 = torch.nn.Sequential()
-        self.slice3 = torch.nn.Sequential()
-        self.slice4 = torch.nn.Sequential()
-        self.slice5 = torch.nn.Sequential()
-        self.N_slices = 5
-        for x in range(4):
-            self.slice1.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(4, 9):
-            self.slice2.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(9, 16):
-            self.slice3.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(16, 23):
-            self.slice4.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(23, 30):
-            self.slice5.add_module(str(x), vgg_pretrained_features[x])
-        if not requires_grad:
-            for param in self.parameters():
-                param.requires_grad = False
-
-    def forward(self, X):
-        h = self.slice1(X)
-        h_relu1_2 = h
-        h = self.slice2(h)
-        h_relu2_2 = h
-        h = self.slice3(h)
-        h_relu3_3 = h
-        h = self.slice4(h)
-        h_relu4_3 = h
-        h = self.slice5(h)
-        h_relu5_3 = h
-        vgg_outputs = namedtuple("VggOutputs", ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3', 'relu5_3'])
-        out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3)
-        return out
-
-    def vggishish16(self, pretrained: bool = True) -> VGGishish:
-        # loading vggishish pretrained on vggsound
-        num_classes_vggsound = 309
-        conv_layers = [64, 64, 'MP', 128, 128, 'MP', 256, 256, 256, 'MP', 512, 512, 512, 'MP', 512, 512, 512]
-        model = VGGishish(conv_layers, use_bn=False, num_classes=num_classes_vggsound)
-        if pretrained:
-            ckpt_path = get_ckpt_path('vggishish_lpaps', "specvqgan/modules/autoencoder/lpaps")
-            ckpt = torch.load(ckpt_path, map_location=torch.device("cpu"))
-            model.load_state_dict(ckpt, strict=False)
-        return model
-
-def normalize_tensor(x, eps=1e-10):
-    norm_factor = torch.sqrt(torch.sum(x**2, dim=1, keepdim=True))
-    return x / (norm_factor+eps)
-
-def spatial_average(x, keepdim=True):
-    return x.mean([2, 3], keepdim=keepdim)
-
-
-if __name__ == '__main__':
-    inputs = torch.rand((16, 1, 80, 848))
-    reconstructions = torch.rand((16, 1, 80, 848))
-    lpips = LPAPS().eval()
-    loss_p = lpips(inputs.contiguous(), reconstructions.contiguous())
-    # (16, 1, 1, 1)
-    print(loss_p.shape)
\ No newline at end of file
diff --git a/instructany2pix/llm/model/vae/modules/quantize_spec.py b/instructany2pix/llm/model/vae/modules/quantize_spec.py
deleted file mode 100644
index 5d88a391c4b08ffafc32299d1a399816154ca3cf..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/vae/modules/quantize_spec.py
+++ /dev/null
@@ -1,131 +0,0 @@
-import torch
-import torch.nn as nn
-
-
-class VectorQuantizer(nn.Module):
-    """
-    see https://github.com/MishaLaskin/vqvae/blob/d761a999e2267766400dc646d82d3ac3657771d4/models/quantizer.py
-    ____________________________________________
-    Discretization bottleneck part of the VQ-VAE.
-    Inputs:
-    - n_e : number of embeddings
-    - e_dim : dimension of embedding
-    - beta : commitment cost used in loss term, beta * ||z_e(x)-sg[e]||^2
-    _____________________________________________
-    """
-
-    def __init__(self, n_e, e_dim, beta):
-        super(VectorQuantizer, self).__init__()
-        self.n_e = n_e
-        self.e_dim = e_dim
-        self.beta = beta
-
-        self.embedding = nn.Embedding(self.n_e, self.e_dim)
-        self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e)
-
-        # better inheritence properties (so that when VectorQuantizer1d() inherits it, only these will be
-        # changed)
-        self.permute_order_in = [0, 2, 3, 1]
-        self.permute_order_out = [0, 3, 1, 2]
-
-    def forward(self, z):
-        """
-        Inputs the output of the encoder network z and maps it to a discrete
-        one-hot vector that is the index of the closest embedding vector e_j
-        z (continuous) -> z_q (discrete)
-        2d: z.shape = (batch, channel, height, width)
-        1d: z.shape = (batch, channel, time)
-        quantization pipeline:
-            1. get encoder input 2d: (B,C,H,W) or 1d: (B, C, T)
-            2. flatten input to 2d: (B*H*W,C) or 1d: (B*T, C)
-        """
-        # reshape z -> (batch, height, width, channel) or (batch, time, channel) and flatten
-        z = z.permute(self.permute_order_in).contiguous()
-        z_flattened = z.view(-1, self.e_dim)
-        # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
-
-        d = torch.sum(z_flattened ** 2, dim=1, keepdim=True) + \
-            torch.sum(self.embedding.weight**2, dim=1) - 2 * \
-            torch.matmul(z_flattened, self.embedding.weight.t())
-
-        ## could possible replace this here
-        # #\start...
-        # find closest encodings
-        min_encoding_indices = torch.argmin(d, dim=1).unsqueeze(1)
-
-        min_encodings = torch.zeros(min_encoding_indices.shape[0], self.n_e).to(z)
-        min_encodings.scatter_(1, min_encoding_indices, 1)
-
-        # dtype min encodings: torch.float32
-        # min_encodings shape: torch.Size([2048, 512])
-        # min_encoding_indices.shape: torch.Size([2048, 1])
-
-        # get quantized latent vectors
-        z_q = torch.matmul(min_encodings, self.embedding.weight).view(z.shape)
-        #.........\end
-
-        # with:
-        # .........\start
-        #min_encoding_indices = torch.argmin(d, dim=1)
-        #z_q = self.embedding(min_encoding_indices)
-        # ......\end......... (TODO)
-
-        # compute loss for embedding
-        loss = torch.mean((z_q.detach() - z) ** 2) + self.beta * torch.mean((z_q - z.detach()) ** 2)
-
-        # preserve gradients
-        z_q = z + (z_q - z).detach()
-
-        # perplexity
-        e_mean = torch.mean(min_encodings, dim=0)
-        perplexity = torch.exp(-torch.sum(e_mean * torch.log(e_mean + 1e-10)))
-
-        # reshape back to match original input shape
-        z_q = z_q.permute(self.permute_order_out).contiguous()
-
-        return z_q, loss, (perplexity, min_encodings, min_encoding_indices)
-
-    def get_codebook_entry(self, indices, shape):
-        # shape specifying (batch, height, width, channel)
-        # TODO: check for more easy handling with nn.Embedding
-        min_encodings = torch.zeros(indices.shape[0], self.n_e).to(indices)
-        min_encodings.scatter_(1, indices[:, None], 1)
-
-        # get quantized latent vectors
-        z_q = torch.matmul(min_encodings.float(), self.embedding.weight)
-
-        if shape is not None:
-            z_q = z_q.view(shape)
-
-            # reshape back to match original input shape
-            z_q = z_q.permute(self.permute_order_out).contiguous()
-
-        return z_q
-
-class VectorQuantizer1d(VectorQuantizer):
-
-    def __init__(self, n_embed, embed_dim, beta=0.25):
-        super().__init__(n_embed, embed_dim, beta)
-        self.permute_order_in = [0, 2, 1]
-        self.permute_order_out = [0, 2, 1]
-
-
-if __name__ == '__main__':
-    quantize = VectorQuantizer1d(n_embed=1024, embed_dim=256, beta=0.25)
-
-    # 1d Input (features)
-    enc_outputs = torch.rand(6, 256, 53)
-    quant, emb_loss, info = quantize(enc_outputs)
-    print(quant.shape)
-
-    quantize = VectorQuantizer(n_e=1024, e_dim=256, beta=0.25)
-
-    # Audio
-    enc_outputs = torch.rand(4, 256, 5, 53)
-    quant, emb_loss, info = quantize(enc_outputs)
-    print(quant.shape)
-
-    # Image
-    enc_outputs = torch.rand(4, 256, 16, 16)
-    quant, emb_loss, info = quantize(enc_outputs)
-    print(quant.shape)
\ No newline at end of file
diff --git a/instructany2pix/llm/model/vae/modules/util.py b/instructany2pix/llm/model/vae/modules/util.py
deleted file mode 100644
index 9527885c605b7fca5f275115a803e19e7e9f53f7..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/vae/modules/util.py
+++ /dev/null
@@ -1,387 +0,0 @@
-# adopted from
-# https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py
-# and
-# https://github.com/lucidrains/denoising-diffusion-pytorch/blob/7706bdfc6f527f58d33f84b7b522e61e6e3164b3/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
-# and
-# https://github.com/openai/guided-diffusion/blob/0ba878e517b276c45d1195eb29f6f5f72659a05b/guided_diffusion/nn.py
-#
-# thanks!
-
-
-import os
-import math
-import torch
-import torch.nn as nn
-import numpy as np
-from einops import repeat
-import importlib
-
-# audio
-
-import hashlib
-import os
-
-import requests
-from tqdm import tqdm
-import subprocess
-
-URL_MAP_AUDIO = {
-    'vggishish_lpaps': 'https://a3s.fi/swift/v1/AUTH_a235c0f452d648828f745589cde1219a/specvqgan_public/vggishish16.pt',
-    'vggishish_mean_std_melspec_10s_22050hz': 'https://a3s.fi/swift/v1/AUTH_a235c0f452d648828f745589cde1219a/specvqgan_public/train_means_stds_melspec_10s_22050hz.txt',
-    'melception': 'https://a3s.fi/swift/v1/AUTH_a235c0f452d648828f745589cde1219a/specvqgan_public/melception-21-05-10T09-28-40.pt',
-}
-
-CKPT_MAP_AUDIO = {
-    'vggishish_lpaps': 'vggishish16.pt',
-    'vggishish_mean_std_melspec_10s_22050hz': 'train_means_stds_melspec_10s_22050hz.txt',
-    'melception': 'melception-21-05-10T09-28-40.pt',
-}
-
-MD5_MAP_AUDIO = {
-    'vggishish_lpaps': '197040c524a07ccacf7715d7080a80bd',
-    'vggishish_mean_std_melspec_10s_22050hz': 'f449c6fd0e248936c16f6d22492bb625',
-    'melception': 'a71a41041e945b457c7d3d814bbcf72d',
-}
-
-def get_ckpt_path(name, root, check=False):
-    assert name in URL_MAP_AUDIO
-    path = os.path.join(root, CKPT_MAP_AUDIO[name])
-    if not os.path.exists(path) or (check and not md5_hash(path) == MD5_MAP_AUDIO[name]):
-        print("Downloading {} model from {} to {}".format(name, URL_MAP_AUDIO[name], path))
-        download(URL_MAP_AUDIO[name], path)
-        md5 = md5_hash(path)
-        assert md5 == MD5_MAP_AUDIO[name], md5
-    return path
-
-def download(url, local_path, chunk_size=1024):
-    os.makedirs(os.path.split(local_path)[0], exist_ok=True)
-    with requests.get(url, stream=True) as r:
-        total_size = int(r.headers.get("content-length", 0))
-        with tqdm(total=total_size, unit="B", unit_scale=True) as pbar:
-            with open(local_path, "wb") as f:
-                for data in r.iter_content(chunk_size=chunk_size):
-                    if data:
-                        f.write(data)
-                        pbar.update(chunk_size)
-
-def which_ffprobe() -> str:
-    '''Determines the path to ffprobe library
-
-    Returns:
-        str -- path to the library
-    '''
-    result = subprocess.run(['which', 'ffprobe'], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
-    ffprobe_path = result.stdout.decode('utf-8').replace('\n', '')
-    return ffprobe_path
-
-def which_ffmpeg() -> str:
-    '''Determines the path to ffmpeg library
-
-    Returns:
-        str -- path to the library
-    '''
-    result = subprocess.run(['which', 'ffmpeg'], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
-    ffmpeg_path = result.stdout.decode('utf-8').replace('\n', '')
-    return ffmpeg_path
-
-def get_duration(path):
-    assert which_ffprobe() != '', 'Is ffmpeg installed? Check if the conda environment is activated.'
-    cmd = f'{which_ffprobe()} -hide_banner -loglevel panic' \
-          f' -v error -show_entries format=duration -of default=noprint_wrappers=1:nokey=1 {path}'
-    result = subprocess.run(cmd.split(), stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
-    duration = float(result.stdout.decode('utf-8').replace('\n', ''))
-    return duration
-
-
-def md5_hash(path):
-    with open(path, "rb") as f:
-        content = f.read()
-    return hashlib.md5(content).hexdigest()
-
-
-def calculate_codebook_bitrate(duration, quant_z, codebook_size):
-    # Calculating the Bitrate
-    bottle_neck_size = quant_z.shape[-2:]
-    bits_per_codebook_entry = (codebook_size-1).bit_length()
-    bitrate = bits_per_codebook_entry * bottle_neck_size.numel() / duration / 1024
-    print(f'The input audio is {duration:.2f} seconds long.')
-    print(f'Codebook size is {codebook_size} i.e. a codebook entry allocates {bits_per_codebook_entry} bits')
-    print(f'SpecVQGAN bottleneck size: {list(bottle_neck_size)}')
-    print(f'Thus, bitrate is {bitrate:.2f} kbps')
-    return bitrate
-
-def get_audio_file_bitrate(file):
-    assert which_ffprobe() != '', 'Is ffmpeg installed? Check if the conda environment is activated.'
-    cmd = f'{which_ffprobe()} -v error -select_streams a:0'\
-          f' -show_entries stream=bit_rate -of default=noprint_wrappers=1:nokey=1 {file}'
-    result = subprocess.run(cmd.split(), stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
-    bitrate = int(result.stdout.decode('utf-8').replace('\n', ''))
-    bitrate /= 1024
-    return bitrate
-
-############## END AUDIO #################
-def get_obj_from_str(string, reload=False):
-    module, cls = string.rsplit(".", 1)
-    if reload:
-        module_imp = importlib.import_module(module)
-        importlib.reload(module_imp)
-    return getattr(importlib.import_module(module, package=None), cls)
-
-
-def instantiate_from_config(config):
-    if not "target" in config:
-        if config == '__is_first_stage__':
-            return None
-        elif config == "__is_unconditional__":
-            return None
-        raise KeyError("Expected key `target` to instantiate.")
-    return get_obj_from_str(config["target"])(**config.get("params", dict()))
-
-
-def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
-    if schedule == "linear":
-        betas = (
-                torch.linspace(linear_start ** 0.5, linear_end ** 0.5, n_timestep, dtype=torch.float64) ** 2
-        )
-
-    elif schedule == "cosine":
-        timesteps = (
-                torch.arange(n_timestep + 1, dtype=torch.float64) / n_timestep + cosine_s
-        )
-        alphas = timesteps / (1 + cosine_s) * np.pi / 2
-        alphas = torch.cos(alphas).pow(2)
-        alphas = alphas / alphas[0]
-        betas = 1 - alphas[1:] / alphas[:-1]
-        betas = np.clip(betas, a_min=0, a_max=0.999)
-
-    elif schedule == "sqrt_linear":
-        betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64)
-    elif schedule == "sqrt":
-        betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) ** 0.5
-    else:
-        raise ValueError(f"schedule '{schedule}' unknown.")
-    return betas.numpy()
-
-
-def make_ddim_timesteps(ddim_discr_method, num_ddim_timesteps, num_ddpm_timesteps, verbose=True):
-    if ddim_discr_method == 'uniform':
-        c = num_ddpm_timesteps // num_ddim_timesteps
-        ddim_timesteps = np.asarray(list(range(0, num_ddpm_timesteps, c)))
-    elif ddim_discr_method == 'quad':
-        ddim_timesteps = ((np.linspace(0, np.sqrt(num_ddpm_timesteps * .8), num_ddim_timesteps)) ** 2).astype(int)
-    else:
-        raise NotImplementedError(f'There is no ddim discretization method called "{ddim_discr_method}"')
-
-    # assert ddim_timesteps.shape[0] == num_ddim_timesteps
-    # add one to get the final alpha values right (the ones from first scale to data during sampling)
-    steps_out = ddim_timesteps + 1
-    if verbose:
-        print(f'Selected timesteps for ddim sampler: {steps_out}')
-    return steps_out
-
-
-def make_ddim_sampling_parameters(alphacums, ddim_timesteps, eta, verbose=True):
-    # select alphas for computing the variance schedule
-    alphas = alphacums[ddim_timesteps]
-    alphas_prev = np.asarray([alphacums[0]] + alphacums[ddim_timesteps[:-1]].tolist())
-
-    # according the the formula provided in https://arxiv.org/abs/2010.02502
-    sigmas = eta * np.sqrt((1 - alphas_prev) / (1 - alphas) * (1 - alphas / alphas_prev))
-    if verbose:
-        print(f'Selected alphas for ddim sampler: a_t: {alphas}; a_(t-1): {alphas_prev}')
-        print(f'For the chosen value of eta, which is {eta}, '
-              f'this results in the following sigma_t schedule for ddim sampler {sigmas}')
-    return sigmas, alphas, alphas_prev
-
-
-def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
-    """
-    Create a beta schedule that discretizes the given alpha_t_bar function,
-    which defines the cumulative product of (1-beta) over time from t = [0,1].
-    :param num_diffusion_timesteps: the number of betas to produce.
-    :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
-                      produces the cumulative product of (1-beta) up to that
-                      part of the diffusion process.
-    :param max_beta: the maximum beta to use; use values lower than 1 to
-                     prevent singularities.
-    """
-    betas = []
-    for i in range(num_diffusion_timesteps):
-        t1 = i / num_diffusion_timesteps
-        t2 = (i + 1) / num_diffusion_timesteps
-        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
-    return np.array(betas)
-
-
-def extract_into_tensor(a, t, x_shape):
-    b, *_ = t.shape
-    out = a.gather(-1, t)
-    return out.reshape(b, *((1,) * (len(x_shape) - 1)))
-
-
-def checkpoint(func, inputs, params, flag):
-    """
-    Evaluate a function without caching intermediate activations, allowing for
-    reduced memory at the expense of extra compute in the backward pass.
-    :param func: the function to evaluate.
-    :param inputs: the argument sequence to pass to `func`.
-    :param params: a sequence of parameters `func` depends on but does not
-                   explicitly take as arguments.
-    :param flag: if False, disable gradient checkpointing.
-    """
-    if flag:
-        args = tuple(inputs) + tuple(params)
-        return CheckpointFunction.apply(func, len(inputs), *args)
-    else:
-        return func(*inputs)
-
-
-class CheckpointFunction(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, run_function, length, *args):
-        ctx.run_function = run_function
-        ctx.input_tensors = list(args[:length])
-        ctx.input_params = list(args[length:])
-
-        with torch.no_grad():
-            output_tensors = ctx.run_function(*ctx.input_tensors)
-        return output_tensors
-
-    @staticmethod
-    def backward(ctx, *output_grads):
-        ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]
-        with torch.enable_grad():
-            # Fixes a bug where the first op in run_function modifies the
-            # Tensor storage in place, which is not allowed for detach()'d
-            # Tensors.
-            shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
-            output_tensors = ctx.run_function(*shallow_copies)
-        input_grads = torch.autograd.grad(
-            output_tensors,
-            ctx.input_tensors + ctx.input_params,
-            output_grads,
-            allow_unused=True,
-        )
-        del ctx.input_tensors
-        del ctx.input_params
-        del output_tensors
-        return (None, None) + input_grads
-
-
-def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
-    """
-    Create sinusoidal timestep embeddings.
-    :param timesteps: a 1-D Tensor of N indices, one per batch element.
-                      These may be fractional.
-    :param dim: the dimension of the output.
-    :param max_period: controls the minimum frequency of the embeddings.
-    :return: an [N x dim] Tensor of positional embeddings.
-    """
-    if not repeat_only:
-        half = dim // 2
-        freqs = torch.exp(
-            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
-        ).to(device=timesteps.device)
-        args = timesteps[:, None].float() * freqs[None]
-        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
-        if dim % 2:
-            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
-    else:
-        embedding = repeat(timesteps, 'b -> b d', d=dim)
-    return embedding
-
-
-def zero_module(module):
-    """
-    Zero out the parameters of a module and return it.
-    """
-    for p in module.parameters():
-        p.detach().zero_()
-    return module
-
-
-def scale_module(module, scale):
-    """
-    Scale the parameters of a module and return it.
-    """
-    for p in module.parameters():
-        p.detach().mul_(scale)
-    return module
-
-
-def mean_flat(tensor):
-    """
-    Take the mean over all non-batch dimensions.
-    """
-    return tensor.mean(dim=list(range(1, len(tensor.shape))))
-
-
-def normalization(channels):
-    """
-    Make a standard normalization layer.
-    :param channels: number of input channels.
-    :return: an nn.Module for normalization.
-    """
-    return GroupNorm32(32, channels)
-
-
-# PyTorch 1.7 has SiLU, but we support PyTorch 1.5.
-class SiLU(nn.Module):
-    def forward(self, x):
-        return x * torch.sigmoid(x)
-
-
-class GroupNorm32(nn.GroupNorm):
-    def forward(self, x):
-        return super().forward(x.float()).type(x.dtype)
-
-def conv_nd(dims, *args, **kwargs):
-    """
-    Create a 1D, 2D, or 3D convolution module.
-    """
-    if dims == 1:
-        return nn.Conv1d(*args, **kwargs)
-    elif dims == 2:
-        return nn.Conv2d(*args, **kwargs)
-    elif dims == 3:
-        return nn.Conv3d(*args, **kwargs)
-    raise ValueError(f"unsupported dimensions: {dims}")
-
-
-def linear(*args, **kwargs):
-    """
-    Create a linear module.
-    """
-    return nn.Linear(*args, **kwargs)
-
-
-def avg_pool_nd(dims, *args, **kwargs):
-    """
-    Create a 1D, 2D, or 3D average pooling module.
-    """
-    if dims == 1:
-        return nn.AvgPool1d(*args, **kwargs)
-    elif dims == 2:
-        return nn.AvgPool2d(*args, **kwargs)
-    elif dims == 3:
-        return nn.AvgPool3d(*args, **kwargs)
-    raise ValueError(f"unsupported dimensions: {dims}")
-
-
-class HybridConditioner(nn.Module):
-
-    def __init__(self, c_concat_config, c_crossattn_config):
-        super().__init__()
-        self.concat_conditioner = instantiate_from_config(c_concat_config)
-        self.crossattn_conditioner = instantiate_from_config(c_crossattn_config)
-
-    def forward(self, c_concat, c_crossattn):
-        c_concat = self.concat_conditioner(c_concat)
-        c_crossattn = self.crossattn_conditioner(c_crossattn)
-        return {'c_concat': [c_concat], 'c_crossattn': [c_crossattn]}
-
-
-def noise_like(shape, device, repeat=False):
-    repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
-    noise = lambda: torch.randn(shape, device=device)
-    return repeat_noise() if repeat else noise()
\ No newline at end of file
diff --git a/instructany2pix/llm/model/vae/modules/vggishish.py b/instructany2pix/llm/model/vae/modules/vggishish.py
deleted file mode 100644
index ff00c955d30d338d2575708121b403aebe8baa3d..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/vae/modules/vggishish.py
+++ /dev/null
@@ -1,77 +0,0 @@
-import torch
-import torch.nn as nn
-
-
-class VGGishish(nn.Module):
-
-    def __init__(self, conv_layers, use_bn, num_classes):
-        '''
-        Mostly from
-            https://pytorch.org/vision/0.8/_modules/torchvision/models/vgg.html
-        '''
-        super().__init__()
-        layers = []
-        in_channels = 1
-
-        # a list of channels with 'MP' (maxpool) from config
-        for v in conv_layers:
-            if v == 'MP':
-                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
-            else:
-                conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1, stride=1)
-                if use_bn:
-                    layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
-                else:
-                    layers += [conv2d, nn.ReLU(inplace=True)]
-                in_channels = v
-        self.features = nn.Sequential(*layers)
-
-        self.avgpool = nn.AdaptiveAvgPool2d((5, 10))
-
-        self.flatten = nn.Flatten()
-        self.classifier = nn.Sequential(
-            nn.Linear(512 * 5 * 10, 4096),
-            nn.ReLU(True),
-            nn.Linear(4096, 4096),
-            nn.ReLU(True),
-            nn.Linear(4096, num_classes)
-        )
-
-        # weight init
-        self.reset_parameters()
-
-    def forward(self, x):
-        # adding channel dim for conv2d (B, 1, F, T) <-
-        x = x.unsqueeze(1)
-        # backbone (B, 1, 5, 53) <- (B, 1, 80, 860)
-        x = self.features(x)
-        # adaptive avg pooling (B, 1, 5, 10) <- (B, 1, 5, 53) – if no MP is used as the end of VGG
-        x = self.avgpool(x)
-        # flatten
-        x = self.flatten(x)
-        # classify
-        x = self.classifier(x)
-        return x
-
-    def reset_parameters(self):
-        for m in self.modules():
-            if isinstance(m, nn.Conv2d):
-                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
-                if m.bias is not None:
-                    nn.init.constant_(m.bias, 0)
-            elif isinstance(m, nn.BatchNorm2d):
-                nn.init.constant_(m.weight, 1)
-                nn.init.constant_(m.bias, 0)
-            elif isinstance(m, nn.Linear):
-                nn.init.normal_(m.weight, 0, 0.01)
-                nn.init.constant_(m.bias, 0)
-
-
-if __name__ == '__main__':
-    num_classes = 309
-    inputs = torch.rand(3, 80, 848)
-    conv_layers = [64, 64, 'MP', 128, 128, 'MP', 256, 256, 256, 'MP', 512, 512, 512, 'MP', 512, 512, 512]
-    # conv_layers = [64, 'MP', 128, 'MP', 256, 256, 'MP', 512, 512, 'MP']
-    model = VGGishish(conv_layers, use_bn=False, num_classes=num_classes)
-    outputs = model(inputs)
-    print(outputs.shape)
\ No newline at end of file
diff --git a/instructany2pix/llm/model/vae/modules/vocoder.py b/instructany2pix/llm/model/vae/modules/vocoder.py
deleted file mode 100644
index 2a9195bd2126d5615d790360cfc055aeb4390b46..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/vae/modules/vocoder.py
+++ /dev/null
@@ -1,216 +0,0 @@
-import torch.nn as nn
-import torch.nn.functional as F
-import torch
-from librosa.filters import mel as librosa_mel_fn
-from torch.nn.utils import weight_norm
-import numpy as np
-import yaml
-from pathlib import Path
-
-def weights_init(m):
-    classname = m.__class__.__name__
-    if classname.find("Conv") != -1:
-        m.weight.data.normal_(0.0, 0.02)
-    elif classname.find("BatchNorm2d") != -1:
-        m.weight.data.normal_(1.0, 0.02)
-        m.bias.data.fill_(0)
-
-
-def WNConv1d(*args, **kwargs):
-    return weight_norm(nn.Conv1d(*args, **kwargs))
-
-
-def WNConvTranspose1d(*args, **kwargs):
-    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
-
-
-class Audio2Mel(nn.Module):
-    def __init__(
-        self,
-        n_fft=1024,
-        hop_length=256,
-        win_length=1024,
-        sampling_rate=22050,
-        n_mel_channels=80,
-        mel_fmin=0.0,
-        mel_fmax=None,
-    ):
-        super().__init__()
-        ##############################################
-        # FFT Parameters                              #
-        ##############################################
-        window = torch.hann_window(win_length).float()
-        mel_basis = librosa_mel_fn(
-            sampling_rate, n_fft, n_mel_channels, mel_fmin, mel_fmax
-        )
-        mel_basis = torch.from_numpy(mel_basis).float()
-        self.register_buffer("mel_basis", mel_basis)
-        self.register_buffer("window", window)
-        self.n_fft = n_fft
-        self.hop_length = hop_length
-        self.win_length = win_length
-        self.sampling_rate = sampling_rate
-        self.n_mel_channels = n_mel_channels
-
-    def forward(self, audio):
-        p = (self.n_fft - self.hop_length) // 2
-        audio = F.pad(audio, (p, p), "reflect").squeeze(1)
-        fft = torch.stft(
-            audio,
-            n_fft=self.n_fft,
-            hop_length=self.hop_length,
-            win_length=self.win_length,
-            window=self.window,
-            center=False,
-        )
-        real_part, imag_part = fft.unbind(-1)
-        magnitude = torch.sqrt(real_part ** 2 + imag_part ** 2)
-        mel_output = torch.matmul(self.mel_basis, magnitude)
-        log_mel_spec = torch.log10(torch.clamp(mel_output, min=1e-5))
-        return log_mel_spec
-
-
-class ResnetBlock(nn.Module):
-    def __init__(self, dim, dilation=1):
-        super().__init__()
-        self.block = nn.Sequential(
-            nn.LeakyReLU(0.2),
-            nn.ReflectionPad1d(dilation),
-            WNConv1d(dim, dim, kernel_size=3, dilation=dilation),
-            nn.LeakyReLU(0.2),
-            WNConv1d(dim, dim, kernel_size=1),
-        )
-        self.shortcut = WNConv1d(dim, dim, kernel_size=1)
-
-    def forward(self, x):
-        return self.shortcut(x) + self.block(x)
-
-
-class Generator(nn.Module):
-    def __init__(self, input_size, ngf, n_residual_layers):
-        super().__init__()
-        ratios = [8, 8, 2, 2]
-        self.hop_length = np.prod(ratios)
-        mult = int(2 ** len(ratios))
-
-        model = [
-            nn.ReflectionPad1d(3),
-            WNConv1d(input_size, mult * ngf, kernel_size=7, padding=0),
-        ]
-
-        # Upsample to raw audio scale
-        for i, r in enumerate(ratios):
-            model += [
-                nn.LeakyReLU(0.2),
-                WNConvTranspose1d(
-                    mult * ngf,
-                    mult * ngf // 2,
-                    kernel_size=r * 2,
-                    stride=r,
-                    padding=r // 2 + r % 2,
-                    output_padding=r % 2,
-                ),
-            ]
-
-            for j in range(n_residual_layers):
-                model += [ResnetBlock(mult * ngf // 2, dilation=3 ** j)]
-
-            mult //= 2
-
-        model += [
-            nn.LeakyReLU(0.2),
-            nn.ReflectionPad1d(3),
-            WNConv1d(ngf, 1, kernel_size=7, padding=0),
-            nn.Tanh(),
-        ]
-
-        self.model = nn.Sequential(*model)
-        self.apply(weights_init)
-
-    def forward(self, x):
-        return self.model(x)
-
-
-class NLayerDiscriminator(nn.Module):
-    def __init__(self, ndf, n_layers, downsampling_factor):
-        super().__init__()
-        model = nn.ModuleDict()
-
-        model["layer_0"] = nn.Sequential(
-            nn.ReflectionPad1d(7),
-            WNConv1d(1, ndf, kernel_size=15),
-            nn.LeakyReLU(0.2, True),
-        )
-
-        nf = ndf
-        stride = downsampling_factor
-        for n in range(1, n_layers + 1):
-            nf_prev = nf
-            nf = min(nf * stride, 1024)
-
-            model["layer_%d" % n] = nn.Sequential(
-                WNConv1d(
-                    nf_prev,
-                    nf,
-                    kernel_size=stride * 10 + 1,
-                    stride=stride,
-                    padding=stride * 5,
-                    groups=nf_prev // 4,
-                ),
-                nn.LeakyReLU(0.2, True),
-            )
-
-        nf = min(nf * 2, 1024)
-        model["layer_%d" % (n_layers + 1)] = nn.Sequential(
-            WNConv1d(nf_prev, nf, kernel_size=5, stride=1, padding=2),
-            nn.LeakyReLU(0.2, True),
-        )
-
-        model["layer_%d" % (n_layers + 2)] = WNConv1d(
-            nf, 1, kernel_size=3, stride=1, padding=1
-        )
-
-        self.model = model
-
-    def forward(self, x):
-        results = []
-        for key, layer in self.model.items():
-            x = layer(x)
-            results.append(x)
-        return results
-
-
-class Discriminator(nn.Module):
-    def __init__(self, num_D, ndf, n_layers, downsampling_factor):
-        super().__init__()
-        self.model = nn.ModuleDict()
-        for i in range(num_D):
-            self.model[f"disc_{i}"] = NLayerDiscriminator(
-                ndf, n_layers, downsampling_factor
-            )
-
-        self.downsample = nn.AvgPool1d(4, stride=2, padding=1, count_include_pad=False)
-        self.apply(weights_init)
-
-    def forward(self, x):
-        results = []
-        for key, disc in self.model.items():
-            results.append(disc(x))
-            x = self.downsample(x)
-        return results
-    
-
-def load_vocoder(ckpt_vocoder: str, eval_mode: bool):
-    ckpt_vocoder = Path(ckpt_vocoder)
-    vocoder_sd = torch.load(ckpt_vocoder / 'best_netG.pt', map_location='cpu')
-
-    with open(ckpt_vocoder / 'args.yml', 'r') as f:
-        args = yaml.load(f, Loader=yaml.UnsafeLoader)
-
-    vocoder = Generator(args.n_mel_channels, args.ngf, args.n_residual_layers)
-    vocoder.load_state_dict(vocoder_sd)
-
-    if eval_mode:
-        vocoder.eval()
-
-    return vocoder
\ No newline at end of file
diff --git a/instructany2pix/llm/model/vae/vq-npz.yaml b/instructany2pix/llm/model/vae/vq-npz.yaml
deleted file mode 100644
index 57900b188a6e40801f1dbb9136ff7c1d2b619e7e..0000000000000000000000000000000000000000
--- a/instructany2pix/llm/model/vae/vq-npz.yaml
+++ /dev/null
@@ -1,27 +0,0 @@
-use_clip: on
-model:
-  base_learning_rate: 4.5e-06
-  target: instructany2pix.llm.model.vae.clip.TensorLoader
-  params:
-    embed_dim: 1024
-    n_embed: 1024
-processor: npz
-
-data:
-  target: main.DataModuleFromConfig
-  
-  params:
-    batch_size: 8
-    num_workers: 16
-    wrap: true
-    train:
-      target: ldm.data.openimages.FullOpenImagesTrain
-      params:
-        crop_size: 256
-    validation:
-      target: ldm.data.openimages.FullOpenImagesValidation
-      params:
-        crop_size: 256
-
-ckpt: null
-image_size: 128
\ No newline at end of file
diff --git a/instructany2pix/pipeline.py b/instructany2pix/pipeline.py
deleted file mode 100644
index 944f73bdc7be98ce1dbbc71db6ef92d8ea87a15f..0000000000000000000000000000000000000000
--- a/instructany2pix/pipeline.py
+++ /dev/null
@@ -1,322 +0,0 @@
-from typing import Any
-import torch
-import os
-from .prior import InstructAny2PixPrior,prior_config
-
-from imagebind import data as image_bind_data
-from imagebind import imagebind_model
-from .llm.model import InstructAny2PixLMForCausalLM
-from .prior.model import MODALITY
-import transformers
-import torch
-from instructany2pix.llm.mm_utils import KeywordsStoppingCriteria
-from sdxl_img2img_pipeline import build_sdxl
-from instructany2pix.llm.conversation import conv_templates, SeparatorStyle
-import torchaudio
-from diffusers import StableDiffusionXLImg2ImgPipeline
-def build_lm(ckpt='ckpts/llm'):
-    any2pix_tokenizer = transformers.AutoTokenizer.from_pretrained(
-        ckpt,subfolder='tokenizer'
-    )
-
-    any2pix_lm = None
-    any2pix_lm = InstructAny2PixLMForCausalLM.from_pretrained(
-        ckpt,
-        load_in_4bit=True, bnb_4bit_compute_dtype=torch.float32,
-                            # config=config
-                                #cache_dir=training_args.cache_dir,
-                                #**bnb_model_from_pretrained_args
-                            )
-
-    any2pix_lm.initialize_vision_tokenizer(None,any2pix_tokenizer,True)
-    return any2pix_tokenizer,any2pix_lm
-from ddim.pnp_pipeline import SDXLDDIMPipeline
-from diffusers import DDIMScheduler
-from PIL import Image
-def resize_and_crop(img, size, crop_type='middle'):
-    """
-    Resize and crop an image to fit the specified size.
-    args:
-        img_path: path for the image to resize.
-        modified_path: path to store the modified image.
-        size: `(width, height)` tuple.
-        crop_type: can be 'top', 'middle' or 'bottom', depending on this
-            value, the image will cropped getting the 'top/left', 'midle' or
-            'bottom/rigth' of the image to fit the size.
-    raises:
-        Exception: if can not open the file in img_path of there is problems
-            to save the image.
-        ValueError: if an invalid `crop_type` is provided.
-    """
-    # If height is higher we resize vertically, if not we resize horizontally
-    # Get current and desired ratio for the images
-    img_ratio = img.size[0] / float(img.size[1])
-    ratio = size[0] / float(size[1])
-    #The image is scaled/cropped vertically or horizontally depending on the ratio
-    if ratio > img_ratio:
-        img = img.resize((size[0], int(size[0] * img.size[1] / img.size[0])),)
-        # Crop in the top, middle or bottom
-        if crop_type == 'top':
-            box = (0, 0, img.size[0], size[1])
-        elif crop_type == 'middle':
-            box = (0, (img.size[1] - size[1]) / 2, img.size[0], (img.size[1] + size[1]) / 2)
-        elif crop_type == 'bottom':
-            box = (0, img.size[1] - size[1], img.size[0], img.size[1])
-        else :
-            raise ValueError('ERROR: invalid value for crop_type')
-        img = img.crop(box)
-    elif ratio < img_ratio:
-        img = img.resize((int(size[1] * img.size[0] / img.size[1]), size[1]),
-                )
-        # Crop in the top, middle or bottom
-        if crop_type == 'top':
-            box = (0, 0, size[0], img.size[1])
-        elif crop_type == 'middle':
-            box = ((img.size[0] - size[0]) / 2, 0, (img.size[0] + size[0]) / 2, img.size[1])
-        elif crop_type == 'bottom':
-            box = (img.size[0] - size[0], 0, img.size[0], img.size[1])
-        else :
-            raise ValueError('ERROR: invalid value for crop_type')
-        img = img.crop(box)
-    else :
-        img = img.resize((int(size[0]), int(size[1])),)
-    return img
-class REPLACEMENT_TYPE:
-            INPUT = 0
-            BASE = 1
-            GEN = 2
-
-class InstructAny2PixPipeline:
-
-    def __init__(self,ckpt='ckpts') -> None:
-        model = InstructAny2PixPrior(**prior_config)
-        model = model.eval()
-        # model.load_state_dict(torch.load('../diffusion_prior.bin',map_location='cpu'))
-        pipe = build_sdxl(ckpt=os.path.join(ckpt,'sdxl'))
-
-        new_sch = DDIMScheduler.from_config(pipe.scheduler.config)
-        pipe_inversion = SDXLDDIMPipeline.from_pretrained(
-                "stabilityai/stable-diffusion-xl-base-1.0",
-                vae=pipe.vae,
-                text_encoder=pipe.text_encoder,
-                text_encoder_2=pipe.text_encoder_2,
-                unet=pipe.unet,
-                torch_dtype=torch.float16,
-                use_safetensors=True,
-                variant="fp16",
-        )
-        pipe_inversion.scheduler = new_sch
-        any2pix_tokenizer,any2pix_lm = build_lm(os.path.join(ckpt,'llm'))
-        model_imb = imagebind_model.imagebind_huge(pretrained=False)
-        model_imb.load_state_dict(torch.load(os.path.join(ckpt,'imagebind_huge.pth'),map_location='cpu'))
-        model.load_state_dict(torch.load(os.path.join(ckpt,'prior/model.bin'),map_location='cpu'))
-        #audioldm_dict = build_model_dict('cuda','/home/jacklishufan/AudioLDM2/audio_align/checkpoint-50000/pytorch_model.bin')
-        model_imb = model_imb.eval()
-        self.model = model
-        self.pipe_inversion = pipe_inversion
-        self.pipe = pipe.to('cuda').to(torch.float16)
-        self.model_imb = model_imb
-        #self.audioldm_dict = audioldm_dict
-        self.any2pix_tokenizer = any2pix_tokenizer
-        self.any2pix_lm = any2pix_lm
-        self.piperf = StableDiffusionXLImg2ImgPipeline.from_pretrained(
-            "stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
-        )
-        self.piperf = self.piperf.to("cuda").to(torch.float16)
-
-
-    def forward_llm(self,inst,mm_data=[]):
-        all_tensors = []
-        for r in mm_data:
-            dict_type = r['type']
-            fpath = r['fname']
-            inputs = {}
-            if dict_type == 'audio':
-                inputs[dict_type] = image_bind_data.load_and_transform_audio_data([fpath],'cpu')
-                res = self.model_imb(inputs)['audio']
-            elif dict_type == 'image':
-                inputs['vision'] = image_bind_data.load_and_transform_vision_data([fpath],'cpu')
-                res = self.model_imb(inputs)['vision']
-            
-            all_tensors.append(res)
-        aux_info = torch.cat(all_tensors)
-        aux_info = aux_info / (aux_info.norm(dim=-1,keepdim=True)+1e-9) * 20
-
-
-        extra_replacement = {
-            "data":aux_info,
-            "mask":torch.tensor([REPLACEMENT_TYPE.INPUT]*aux_info.shape[0],dtype=torch.long)
-        }
-
-        conv = conv_templates['vicuna_v1'].copy()
-        roles = {"human": conv.roles[0], "gpt": conv.roles[1]}
-        conv.append_message(conv.roles[0],inst)
-        #conv.append_message(conv.roles[0], inst)
-        conv.append_message(conv.roles[1], None)
-        prompt = conv.get_prompt()
-        input_ids = self.any2pix_tokenizer(prompt, return_tensors='pt').input_ids
-        stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
-        keywords = [stop_str]
-        stopping_criteria = KeywordsStoppingCriteria(keywords, self.any2pix_tokenizer, input_ids)
-
-        gen_seq = self.any2pix_tokenizer('<video>',add_special_tokens=False)
-        gen_seq = gen_seq.input_ids[0]
-        gen_seq
-        base_null = self.any2pix_tokenizer('<base_null>',add_special_tokens=False)
-        base_null = base_null.input_ids[0]
-        base_null
-        base_tkn = self.any2pix_tokenizer('<base>',add_special_tokens=False)
-        base_tkn = base_tkn.input_ids[0]
-        base_tkn
-
-        self.any2pix_lm.eval()
-
-        output_ids = self.any2pix_lm.generate(
-            input_ids.cuda(),
-            images=None,
-            do_sample=True,
-            temperature=0.3,
-            max_new_tokens=100,
-            output_hidden_states=True,
-            use_cache=False,
-            return_dict_in_generate=True,
-            extra_replacement=extra_replacement,
-            stopping_criteria=[stopping_criteria])
-
-        out_seq = output_ids.sequences[:,input_ids.shape[1]:]
-        assert len(output_ids.hidden_states) == out_seq.shape[1]
-        gen_idx = torch.where(out_seq.view(-1) == gen_seq)[0][-1].item()
-        with torch.no_grad():
-            image_embeds = self.any2pix_lm.get_model().vae_predictor_image(output_ids.hidden_states[-2][-1][:,-1:]).detach().cpu()
-            #image_embeds = image_embeds / image_embeds.norm() * 22
-
-        gen_idx = out_seq.view(-1).tolist().index(base_tkn)+1
-        with torch.no_grad():
-            base_embed = self.any2pix_lm.get_model().vae_predictor_image(output_ids.hidden_states[gen_idx][-1][:,-1:]).detach().cpu()
-        base_embed = base_embed[0]
-        base_idx = torch.einsum('ac,bc->ab',base_embed.float().detach().cpu() / base_embed.norm() * 20,aux_info.float() )[0].argmax().item()
-        b = mm_data[base_idx]['fname']
-        #assert a == 'image'
-        tp = self.any2pix_tokenizer.batch_decode(output_ids.sequences)
-        import re
-        output_caption = re.compile('\[([^\]]+)\]').findall(tp[0])[0]
-        return image_embeds,base_embed,output_caption,b
-    
-    def loas_base_img(self,base_img_path):
-        img_base = Image.open(base_img_path)
-        img_base = resize_and_crop(img_base,(1024,1024),crop_type='middle')
-        img_base = img_base.resize((1024,1024))
-        return img_base
-    
-    def polar_intrtpolate(self,x,y,alpha):
-        n0 = x.norm()
-        n1 = y.norm()
-        ll = x * alpha + y * (1-alpha)
-        n = n0 * alpha + n1 * (1-alpha)
-        return ll / ll.norm() * n
-    
-    def __call__(self, inst,mm_data,alpha = 0.7,h=[0.0,0.4,1.0],norm=20.0,refinement=0.5,llm_only=False,num_inference_steps=25) -> Any:
-        image_embeds,base_embed,output_caption,base_img_path = self.forward_llm(inst,mm_data)
-        if llm_only:
-            return None,None,output_caption
-        y = self.model.generate_diffusion(MODALITY.VIDEO,MODALITY.IMAGE,image_embeds / image_embeds.norm() * 100,device='cpu',
-                             no_diffusion=True,num_inference_steps=25,
-                             image_bind_overwrite=None,
-                             dtype=torch.float32,guidance_scale=10,
-                             force_guidence_t0=True,do_classifier_free_guidance=True,score=6.5)
-        # y = self.model.generate_diffusion(MODALITY.TEXT,MODALITY.IMAGE,[output_caption],device='cpu',
-        #                      no_diffusion=True,num_inference_steps=25,
-        #                      image_bind_overwrite=None,
-        #                      dtype=torch.float32,guidance_scale=10,
-        #                      force_guidence_t0=True,do_classifier_free_guidance=True,score=6.5)
-        
-        img_base = self.loas_base_img(base_img_path)
-
-        
-        latent_la = base_embed * h[0] + image_embeds * h[1] +y[0] / y[0].norm() *20.0 *h[2] #+y[0] #/y[0].norm()*20# + image_embeds/image_embeds.norm()*10
-        latent_la = latent_la.detach().clone()
-        latent_la = latent_la / latent_la.norm() * norm
-        null_prompt_embeds = torch.zeros(1,768)
-        extra_kwargs = {}
-
-        null_prompt_embeds = torch.zeros(1,768)
-        latent_inv = self.pipe_inversion.inverse(num_inference_steps=num_inference_steps,
-                                            prompt_embeds=torch.zeros(1,1,1024),pooled_prompt_embeds=null_prompt_embeds,image=img_base.resize((1024,1024)))
-        latent_inv = latent_inv.images.cpu()
-        #alpha = 0.7
-        latent_inv = self.polar_intrtpolate(
-            latent_inv,
-            torch.randn_like(latent_inv),
-            alpha,
-        )
-        extra_kwargs['latents']= latent_inv # latent_inv
-        # pipe = pipe.to('cuda').to(torch.float16)
-        images = self.pipe(original_size=(1024,1024),
-        height=1024,
-        width=1024,
-        prompt_embeds=latent_la.to(torch.bfloat16).cuda(),
-        pooled_prompt_embeds=null_prompt_embeds.to(torch.bfloat16).cuda(), generator=None,
-            num_inference_steps=num_inference_steps,guidance_scale=10,**extra_kwargs)
-        non_refined =  images[0][0]
-        if refinement > 0:
-            oo = self.piperf(image=images[0][0],prompt=output_caption,strength=refinement,).images[0]
-        else:
-            oo = images[0][0]
-        return non_refined,oo,output_caption
-        # z = 0
-        # while(os.path.exists(out_name+f'{z}.jpg')):
-        #     z += 1
-        # images[0][0].save(out_name+f'{z}.jpg')
-    
-def load_json(fp):
-    with open(fp) as f:
-        data = json.loads(f.read())
-
-    return data
-
-def preprocess_mm_data(z,img_folder,audio_folder):
-    new_data = []
-    for x in z:
-        payload = {k:v for k,v in x.items()}
-        if payload['type'] == 'image':
-            payload['fname'] = os.path.join(img_folder,payload['fname'])
-        elif payload['type'] == 'audio':
-            payload['fname'] = os.path.join(audio_folder,payload['fname'])
-        new_data.append(payload)
-    return new_data
-
-from pathlib import Path
-import os
-import shutil
-import json
-def dump_json(data,fp):
-    with open(fp,'w') as f:
-        f.write(json.dumps(data))
-    
-if __name__=='__main__':
-    img_folder = '/home/jacklishufan/public_eval/instruction_gen/out_img'
-    audio_folder = '/home/jacklishufan/public_eval/instruction_gen/out_audio'
-    pipe = InstructAny2PixPipeline()
-    data = load_json('/home/jacklishufan/public_eval/instruction_gen/all_data.json')
-    alpha = 1.0
-    h = [0.1,0.4,1.0]
-    refinement=0.3
-    norm=20.0
-    out_dir = f'test_r{refinement}_alpha{alpha}_h{h}_n{norm}'
-    if not os.path.exists(out_dir):
-        Path(out_dir).mkdir(exist_ok=True)
-    for idx,x in enumerate(data):
-        inst = x['inst']
-        mm_data = preprocess_mm_data(x['mm_data'],img_folder,audio_folder)
-        # breakpoint()
-        base = list([z for z in mm_data if z['is_base']])
-        if len(base)==0:
-            continue
-        base = base[0]
-        shutil.copy(base['fname'],os.path.join(out_dir,f'{idx}_0.jpg'))
-        res0,res,output_caption = pipe(inst,mm_data,alpha =alpha,h=h,norm=norm,refinement=refinement)
-        res0.save(os.path.join(out_dir,f'{idx}_1.jpg'))
-        res.save(os.path.join(out_dir,f'{idx}_2.jpg'))
-        x['output_caption']=output_caption
-        dump_json(x,os.path.join(out_dir,f'{idx}_0.json'))
-
diff --git a/instructany2pix/prior/__init__.py b/instructany2pix/prior/__init__.py
deleted file mode 100644
index ba09485993374b3f79908028cb003b5e12b36e77..0000000000000000000000000000000000000000
--- a/instructany2pix/prior/__init__.py
+++ /dev/null
@@ -1,41 +0,0 @@
-from .model import InstructAny2PixPrior
-prior_config ={
-        "always_output_audiomae_gt": False,
-        "learnable": True,
-        "device": "cuda",
-        "use_gt_mae_output": True,
-        "use_gt_mae_prob": 0.0,
-        "base_learning_rate": 0.0002,
-        "sequence_gen_length": 1,
-        "diffusion": True,
-        "use_warmup": True,
-        "sequence_input_key": [
-            "src_type",
-            "imagebind",
-            # "crossattn_flan_t5",
-            "crossattn_clip",
-            "score",
-            "noisy_inputs",
-            "noise_level"
-            "tgt_type",
-        ],
-        "sequence_input_embed_dim": [0,1024,1024,512,0,0,0],
-        "pretrained_name":'gpt2-medium',
-        "batchsize": 16,
-        "embed_dim":1024,
-        "output_dim":1024,
-        "cond_stage_config": {
-            # "crossattn_flan_t5": {
-            #     "cond_stage_key": "text",
-            #     "conditioning_key": "crossattn",
-            #     "target": "audioldm2.latent_diffusion.modules.encoders.modules.FlanT5HiddenState",
-            # },
-            "crossattn_clip": {
-                "cond_stage_key": "text",
-                "conditioning_key": "crossattn",
-                "target": "audioldm2.latent_diffusion.modules.encoders.modules.CLIPTextModelHiddenState",
-            },
-
-            
-        },
-}
\ No newline at end of file
diff --git a/instructany2pix/prior/model.py b/instructany2pix/prior/model.py
deleted file mode 100644
index 0cdc20704c5b25c88e957689ee2124aa4c77a2ae..0000000000000000000000000000000000000000
--- a/instructany2pix/prior/model.py
+++ /dev/null
@@ -1,794 +0,0 @@
-import torch
-import torch.nn as nn
-from ..utils import (
-    instantiate_from_config,
-)
-class MODALITY:
-    IMAGE = 0
-    AUDIO = 1 
-    TEXT = 2
-    VIDEO = 3
-# from latent_diffusion.modules.encoders.modules import CLAPAudioEmbeddingClassifierFreev2
-from transformers import GPT2Config, GPT2Model,AutoTokenizer,CLIPTextModel
-import torch.optim.lr_scheduler as lr_scheduler
-from torch.optim import AdamW
-from diffusers import DDPMScheduler
-from diffusers.models.embeddings import get_timestep_embedding
-import inspect
-from tqdm.cli import tqdm
-import logging
-class CLIPTextModelHiddenState(nn.Module):
-    """
-    llama = FlanT5HiddenState()
-    data = ["","this is not an empty sentence"]
-    encoder_hidden_states = llama(data)
-    import ipdb;ipdb.set_trace()
-    """
-
-    def __init__(
-        self, text_encoder_name='laion/CLIP-ViT-H-14-laion2B-s32B-b79K', freeze_text_encoder=True
-    ):
-        super().__init__()
-        self.freeze_text_encoder = freeze_text_encoder
-        self.tokenizer = AutoTokenizer.from_pretrained(text_encoder_name)
-        self.model = CLIPTextModel.from_pretrained(text_encoder_name)
-        if freeze_text_encoder:
-            self.model.eval()
-            for p in self.model.parameters():
-                p.requires_grad = False
-        else:
-            print("=> The text encoder is learnable")
-
-        self.empty_hidden_state_cfg = None
-        self.device = None
-
-    # Required
-    def get_unconditional_condition(self, batchsize):
-        param = next(self.model.parameters())
-        if self.freeze_text_encoder:
-            assert param.requires_grad == False
-
-        # device = param.device
-        if self.empty_hidden_state_cfg is None:
-            self.empty_hidden_state_cfg, _ = self([""])
-
-        hidden_state = torch.cat([self.empty_hidden_state_cfg] * batchsize).float()
-        attention_mask = (
-            torch.ones((batchsize, hidden_state.size(1)))
-            .to(hidden_state.device)
-            .float()
-        )
-        return [hidden_state, attention_mask]  # Need to return float type
-
-    def forward(self, batch):
-        param = next(self.model.parameters())
-        if self.freeze_text_encoder:
-            assert param.requires_grad == False
-
-        if self.device is None:
-            self.device = param.device
-
-        # print("Manually change text")
-        # for i in range(len(batch)):
-        #     batch[i] = "dog barking"
-        try:
-            return self.encode_text(batch)
-        except Exception as e:
-            print(e, batch)
-            logging.exception("An error occurred: %s", str(e))
-
-    def encode_text(self, prompt):
-        device = self.model.device
-        batch = self.tokenizer(
-            prompt,
-            max_length=77,  # self.tokenizer.model_max_length
-            padding=True,
-            truncation=True,
-            return_tensors="pt",
-        )
-        input_ids, attention_mask = batch.input_ids.to(device), batch.attention_mask.to(
-            device
-        )
-        # Get text encoding
-        if self.freeze_text_encoder:
-            with torch.no_grad():
-                encoder_hidden_states = self.model(
-                    input_ids=input_ids, attention_mask=attention_mask
-                )[0]
-        else:
-            encoder_hidden_states = self.model(
-                input_ids=input_ids, attention_mask=attention_mask
-            )[0]
-        return [
-            encoder_hidden_states.detach(),
-            attention_mask.float(),
-        ]  # Attention mask == 1 means usable token
-
-
-
-class InstructAny2PixPrior(nn.Module):
-    def __init__(
-        self,
-        base_learning_rate,
-        sequence_gen_length,
-        sequence_input_key,
-        sequence_input_embed_dim,
-        cond_stage_config,
-        optimizer_type="AdamW",
-        use_warmup=True,
-        use_ar_gen_loss=False,
-        use_audiomae_linear=False,
-        target_tokens_mask_ratio=0.0,
-        random_mask_ratio=False,
-        loss='L1',
-        pretrained_name='gpt2',
-        embed_dim=768,
-        output_dim=None,
-        weight_decay=0.,
-        diffusion=None,
-        **kwargs
-    ):
-        super().__init__()
-        self.diffusion = diffusion
-        if diffusion:
-            self.noise_scheduler = DDPMScheduler.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler")
-            #self.noise_scheduler. sconfig.prediction_type == "sample"
-        assert use_audiomae_linear == False
-        self.random_mask_ratio = random_mask_ratio
-        self.learning_rate = base_learning_rate
-        self.cond_stage_config = cond_stage_config
-        self.use_audiomae_linear = use_audiomae_linear
-        self.optimizer_type = optimizer_type
-        self.use_warmup = use_warmup
-        self.use_ar_gen_loss = use_ar_gen_loss
-        self.weight_decay = weight_decay
-        # Even though the LDM can be conditioned on mutliple pooling rate
-        # Our model always predict the higest pooling rate
-
-        # self.time_pool = max(self.cond_stage_config["crossattn_audiomae_pooled"]["params"]["time_pooling_factors"])
-        # self.freq_pool = max(self.cond_stage_config["crossattn_audiomae_pooled"]["params"]["freq_pooling_factors"])
-        # self.mae_token_num = int(512/(self.time_pool*self.freq_pool))
-
-        self.mae_token_num = sequence_gen_length
-        self.sequence_input_key = sequence_input_key
-        self.sequence_input_embed_dim = sequence_input_embed_dim
-        self.target_tokens_mask_ratio = target_tokens_mask_ratio
-
-        self.start_of_sequence_tokens = nn.Embedding(32, embed_dim)
-        self.end_of_sequence_tokens = nn.Embedding(32, embed_dim)
-
-        self.input_sequence_embed_linear = nn.ModuleList([])
-        self.initial_learning_rate = None
-
-        for dim in self.sequence_input_embed_dim:
-            if dim == 0:
-                self.input_sequence_embed_linear.append(nn.Identity())
-            else: 
-                self.input_sequence_embed_linear.append(nn.Linear(dim, embed_dim))
-
-        self.modality_embedding = nn.Embedding(10,embed_dim)
-
-        if output_dim is not None and output_dim != embed_dim:
-            self.output_proj = nn.Linear(embed_dim,output_dim)
-        else:
-            self.output_proj = nn.Identity()
-
-        self.cond_stage_models = nn.ModuleList([])
-        self.instantiate_cond_stage(cond_stage_config)
-        self.initialize_param_check_toolkit()
-
-        # configuration = GPT2Config(n_layer=1) # TODO
-        # self.model=GPT2Model(configuration)
-        ###################
-        # self.model=nn.Linear(768,768, bias=False) # TODO change the model
-        # with torch.no_grad():
-        #     self.model.weight.copy_(torch.eye(768))
-        ###################
-        self.model = GPT2Model(GPT2Config.from_pretrained(pretrained_name))
-        ###################
-        # self.model = nn.LSTM(input_size=768, hidden_size=768, num_layers=1,bias=False) # TODO
-
-        # self.loss_fn = nn.MSELoss()
-        if loss == 'L1':   
-            self.loss_fn = nn.L1Loss()
-        elif loss == 'mse':
-            self.loss_fn = nn.MSELoss()
-        else:
-            raise NotImplemented
-        
-
-        self.embed_dim = embed_dim
-
-        self.logger_save_dir = None
-        self.logger_exp_name = None
-        self.logger_exp_group_name = None
-        self.logger_version = None
-
-    
-    def get_eps(self,timestep,sample,model_output):
-        #print(sample.shape,model_output.shape)
-        t = timestep
-        prev_timestep = timestep - self.noise_scheduler.config.num_train_timesteps // self.noise_scheduler.num_inference_steps
-
-        if model_output.shape[1] == sample.shape[1] * 2 and self.variance_type in ["learned", "learned_range"]:
-            raise NotImplemented
-        else:
-            predicted_variance = None
-
-        prev_timestep = timestep - self.noise_scheduler.config.num_train_timesteps // self.noise_scheduler.num_inference_steps
-        # 1. compute alphas, betas
-        alpha_prod_t = self.noise_scheduler.alphas_cumprod[t]
-        alpha_prod_t_prev = self.noise_scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.noise_scheduler.one
-        beta_prod_t = 1 - alpha_prod_t
-        beta_prod_t_prev = 1 - alpha_prod_t_prev
-        # current_alpha_t = alpha_prod_t / alpha_prod_t_prev
-        # current_beta_t = 1 - current_alpha_t
-        # alpha_prod_t = self.noise_scheduler.alphas_cumprod[timestep]
-        # alpha_prod_t_prev = self.noise_scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.noise_scheduler.final_alpha_cumprod
-
-        #beta_prod_t = 1 - alpha_prod_t
-
-        # 2. compute predicted original sample from predicted noise also called
-        # "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
-        # if self.config.prediction_type == "epsilon":
-        #     pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
-        # solve for model_output
-        #      (sample - alpha_prod_t ** (0.5) * pred_original_sample ) / beta_prod_t ** (0.5) = model_output 
-        eps =  (sample - alpha_prod_t ** (0.5) * model_output ) / beta_prod_t ** (0.5) 
-        #print(eps.mean(),eps.std())
-        return eps
-
-
-    def set_log_dir(self, save_dir, exp_group_name, exp_name):
-        self.logger_save_dir = save_dir
-        self.logger_exp_group_name = exp_group_name
-        self.logger_exp_name = exp_name
-
-    def cfg_uncond(self, batch_size):
-        unconditional_conditioning = {}
-        for key in self.cond_stage_model_metadata:
-            model_idx = self.cond_stage_model_metadata[key]["model_idx"]
-            unconditional_conditioning[key] = self.cond_stage_models[
-                model_idx
-            ].get_unconditional_condition(batch_size)
-        assert (
-            "crossattn_audiomae_pooled" in unconditional_conditioning.keys()
-        ), "The module is not initialized with AudioMAE"
-        unconditional_conditioning[
-            "crossattn_clap_to_audiomae_feature"
-        ] = unconditional_conditioning["crossattn_audiomae_pooled"]
-        return unconditional_conditioning
-
-    def configure_optimizers(self,step_size=5):
-        lr = float(self.learning_rate)
-        # params = list(self.model.parameters()) + list(self.input_sequence_embed_linear.parameters())
-        params = list(self.parameters())
-
-        # opt = torch.optim.Adam(params, lr=lr, betas=(0.9, 0.98), eps=1e-9)
-        opt = eval(self.optimizer_type)(params, lr=lr,weight_decay=self.weight_decay)
-        scheduler = lr_scheduler.StepLR(opt, step_size=step_size, gamma=0.8)
-        return [opt], [scheduler]
-
-    def add_sos_eos_tokens(self, _id, sequence, attn_mask):
-        batchsize = sequence.size(0)
-
-        new_attn_mask_step = torch.ones((batchsize, 1)).to(sequence.device)
-        key_id = torch.tensor([_id]).to(sequence.device)
-
-        # Add two more steps to attn mask
-        new_attn_mask = torch.cat(
-            [new_attn_mask_step, attn_mask, new_attn_mask_step], dim=1
-        )
-
-        # Add two more tokens in the sequence
-        sos_token = self.start_of_sequence_tokens(key_id).expand(batchsize, 1, -1)
-        eos_token = self.end_of_sequence_tokens(key_id).expand(batchsize, 1, -1)
-        new_sequence = torch.cat([sos_token, sequence, eos_token], dim=1)
-        return new_sequence, new_attn_mask
-
-    def truncate_sequence_and_mask(self, sequence, mask, max_len=512):
-        if sequence.size(1) > max_len:
-            print(
-                "The input sequence length to GPT-2 model is too long:",
-                sequence.size(1),
-            )
-            return sequence[:, :max_len], mask[:, :max_len]
-        else:
-            return sequence, mask
-
-    def get_input_sequence_and_mask(self, cond_dict):
-        input_embeds = None
-        input_embeds_attn_mask = None
-        for _id, sequence_key in enumerate(self.sequence_input_key):
-            # assert sequence_key in cond_dict.keys(), (
-            #     "Invalid sequence key %s" % sequence_key
-            # )
-            if sequence_key not in cond_dict.keys():
-                continue
-            cond_embed = cond_dict[sequence_key]
-            if sequence_key in ['src_type','tgt_type']:
-                if len(cond_embed.shape) == 1:
-                    cond_embed = cond_embed[:,None]
-                item_input_embeds = self.modality_embedding(cond_embed)
-                item_attn_mask = torch.ones((cond_embed.size(0), cond_embed.size(1))).to(
-                    cond_embed.device
-                )
-                if input_embeds is None and input_embeds_attn_mask is None:
-                    input_embeds, input_embeds_attn_mask = (
-                        item_input_embeds,
-                        item_attn_mask,
-                    )
-                else:
-                    input_embeds, input_embeds_attn_mask = torch.cat(
-                        [input_embeds, item_input_embeds], dim=1
-                    ), torch.cat([input_embeds_attn_mask, item_attn_mask], dim=1)
-            elif isinstance(cond_embed, list):
-                assert (
-                    len(cond_embed) == 2
-                ), "The crossattn returned list should have length 2, including embed and attn_mask"
-                item_input_embeds, item_attn_mask = cond_embed
-
-                item_input_embeds = self.input_sequence_embed_linear[_id](
-                    item_input_embeds
-                )
-
-                item_input_embeds, item_attn_mask = self.add_sos_eos_tokens(
-                    _id, item_input_embeds, item_attn_mask
-                )
-
-                if input_embeds is None and input_embeds_attn_mask is None:
-                    input_embeds, input_embeds_attn_mask = (
-                        item_input_embeds,
-                        item_attn_mask,
-                    )
-                else:
-                    input_embeds = torch.cat(
-                        [input_embeds, item_input_embeds], dim=1
-                    )  # The 1-st dimension is time steps
-                    input_embeds_attn_mask = torch.cat(
-                        [input_embeds_attn_mask, item_attn_mask], dim=1
-                    )  # The 1-st dimension is time steps
-            else:
-                assert isinstance(cond_embed, torch.Tensor)
-                cond_embed = self.input_sequence_embed_linear[_id](cond_embed)
-                attn_mask = torch.ones((cond_embed.size(0), cond_embed.size(1))).to(
-                    cond_embed.device
-                )
-
-                item_input_embeds, item_attn_mask = self.add_sos_eos_tokens(
-                    _id, cond_embed, attn_mask
-                )
-
-                if input_embeds is None and input_embeds_attn_mask is None:
-                    input_embeds, input_embeds_attn_mask = (
-                        item_input_embeds,
-                        item_attn_mask,
-                    )
-                else:
-                    input_embeds, input_embeds_attn_mask = torch.cat(
-                        [input_embeds, item_input_embeds], dim=1
-                    ), torch.cat([input_embeds_attn_mask, item_attn_mask], dim=1)
-
-        assert input_embeds is not None and input_embeds_attn_mask is not None
-
-        input_embeds, input_embeds_attn_mask = self.truncate_sequence_and_mask(
-            input_embeds, input_embeds_attn_mask, int(1024 - self.mae_token_num)
-        )
-        cond_sequence_end_time_idx = input_embeds.size(
-            1
-        )  # The index that we start to collect the output embeds
-
-        return input_embeds, input_embeds_attn_mask, cond_sequence_end_time_idx
-
-    def warmup_step(self):
-        if self.initial_learning_rate is None:
-            self.initial_learning_rate = float(self.learning_rate)
-
-        # Only the first parameter group
-        if self.global_step <= 1000:
-            if self.global_step == 0:
-                print(
-                    "Warming up learning rate start with %s"
-                    % self.initial_learning_rate
-                )
-            self.trainer.optimizers[0].param_groups[0]["lr"] = (
-                self.global_step / 1000
-            ) * self.initial_learning_rate
-        else:
-            # TODO set learning rate here
-            self.trainer.optimizers[0].param_groups[0][
-                "lr"
-            ] = self.initial_learning_rate
-
-    def mask_target_sequence(self, target_embeds, target_embeds_attn_mask):
-        time_seq_mask = None
-        if self.target_tokens_mask_ratio > 1e-4:
-            batchsize, time_seq_len, embed_dim = target_embeds.size()
-            _, time_seq_len = target_embeds_attn_mask.size()
-            # Generate random mask
-            if self.random_mask_ratio:
-                mask_ratio = torch.rand(1).item() * self.target_tokens_mask_ratio
-            else:
-                mask_ratio = self.target_tokens_mask_ratio
-
-            time_seq_mask = (torch.rand((batchsize, time_seq_len)) > mask_ratio).to(
-                target_embeds.device
-            )
-            # Mask the target embedding
-            target_embeds = target_embeds * time_seq_mask.unsqueeze(-1)
-            target_embeds_attn_mask = target_embeds_attn_mask * time_seq_mask
-        return target_embeds, target_embeds_attn_mask, time_seq_mask
-
-    def generate_partial(self, batch, cond_dict=None, no_grad=False):
-        if cond_dict is None:
-            cond_dict = self.get_input(batch)
-
-        print("Generate partially prompted audio with in-context learning")
-        # self.model.train()
-        # assert self.model.training==True
-
-        target_embeds, target_embeds_attn_mask = (
-            cond_dict["crossattn_audiomae_pooled"][0],
-            cond_dict["crossattn_audiomae_pooled"][1],
-        )
-
-        target_time_steps = target_embeds.size(1)
-
-        (
-            input_embeds,
-            input_embeds_attn_mask,
-            cond_sequence_end_time_idx,
-        ) = self.get_input_sequence_and_mask(cond_dict)
-
-        model_input = torch.cat(
-            [input_embeds, target_embeds[:, : target_time_steps // 4, :]], dim=1
-        )
-        model_input_mask = torch.cat(
-            [
-                input_embeds_attn_mask,
-                target_embeds_attn_mask[:, : target_time_steps // 4],
-            ],
-            dim=1,
-        )
-
-        steps = self.mae_token_num
-
-        for _ in range(3 * steps // 4):
-            output = self.model(
-                inputs_embeds=model_input, attention_mask=model_input_mask
-            )["last_hidden_state"]
-            # Update the model input
-            model_input = torch.cat([model_input, output[:, -1:, :]], dim=1)
-            # Update the attention mask
-            attention_mask_new_step = torch.ones((model_input_mask.size(0), 1)).to(
-                model_input.device
-            )
-            model_input_mask = torch.cat(
-                [model_input_mask, attention_mask_new_step], dim=1
-            )
-
-        output = model_input[:, cond_sequence_end_time_idx:]
-
-        return output, cond_dict
-
-    def generate(self, batch, cond_dict=None, no_grad=False):
-        if cond_dict is None:
-            cond_dict = self.get_input(batch)
-        else:
-            z = self.get_input(cond_dict)
-            cond_dict.update(z)
-        # self.model.train()
-        # print("!!!!!!!!!!!!!train")
-
-        (
-            input_embeds,
-            input_embeds_attn_mask,
-            cond_sequence_end_time_idx,
-        ) = self.get_input_sequence_and_mask(cond_dict)
-        model_input = input_embeds
-        model_input_mask = input_embeds_attn_mask
-
-        steps = self.mae_token_num
-
-        for _ in range(steps):
-            output = self.model(
-                inputs_embeds=model_input, attention_mask=model_input_mask
-            )["last_hidden_state"]
-            # Update the model input
-            model_input = torch.cat([model_input, output[:, -1:, :]], dim=1)
-            # Update the attention mask
-            attention_mask_new_step = torch.ones((model_input_mask.size(0), 1)).to(
-                model_input.device
-            )
-            model_input_mask = torch.cat(
-                [model_input_mask, attention_mask_new_step], dim=1
-            )
-
-        return model_input[:, cond_sequence_end_time_idx:], cond_dict
-    
-    def prepare_extra_step_kwargs(self, generator, eta):
-        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
-        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
-        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
-        # and should be between [0, 1]
-
-        accepts_eta = "eta" in set(inspect.signature(self.noise_scheduler.step).parameters.keys())
-        extra_step_kwargs = {}
-        if accepts_eta:
-            extra_step_kwargs["eta"] = eta
-
-        # check if the scheduler accepts generator
-        accepts_generator = "generator" in set(inspect.signature(self.noise_scheduler.step).parameters.keys())
-        if accepts_generator:
-            extra_step_kwargs["generator"] = generator
-        return extra_step_kwargs
-    
-
-    @torch.no_grad()
-    def generate_diffusion(self,src_type,tgt_type,src, 
-                           
-                           no_grad=False,num_inference_steps=25,
-                           eta: float = 0.0,
-                           generator=None,
-                           image_bind_overwrite=None,
-                           guidance_scale=5,
-                           score=6.8,
-                           negative_score=2.0,
-                           do_classifier_free_guidance=True,
-                           device='cuda',
-                           dtype=torch.float16,
-                           no_diffusion=False,
-                           force_guidence_t0=False):
-
-                #         timesteps = torch.randint(
-                #         0, self.noise_scheduler.config.num_train_timesteps, (target.shape[0],), device=target.device
-                #     )
-                # noise = torch.randn_like(target)
-                # noisy_input = 10.0 * self.noise_scheduler.add_noise(target / 10.0,noise, timesteps)
-                # noise_level = get_timestep_embedding(
-                #     timesteps=timesteps, embedding_dim=noisy_input.shape[-1], flip_sin_to_cos=True, downscale_freq_shift=0)
-                # cond_dict['noisy_input'] = noisy_input
-                # cond_dict['noise_level'] = noise_level
-        if no_diffusion:
-            num_inference_steps = 1
-        bs = len(src)
-        raw_bs = bs
-        if src_type == MODALITY.TEXT:
-            src_key = 'text'
-        else:
-            src_key = 'imagebind'
-        if image_bind_overwrite is None:
-            image_bind_overwrite = torch.zeros(bs,1,1024).to(device).to(dtype)
-        cond_dict = dict(
-            src_type = torch.tensor(src_type).view(1,1).to(device).repeat(bs,1),
-            tgt_type=torch.tensor(tgt_type).view(1,1).to(device).repeat(bs,1),
-            score = get_timestep_embedding(
-                torch.tensor([score]).to(device).to(dtype),
-                512,flip_sin_to_cos=True,downscale_freq_shift=0
-            ).view(1,1,-1).repeat(bs,1,1),
-            text=[""],
-            imagebind = image_bind_overwrite.to(device).to(dtype),
-        )
-        if src_key == 'text':
-            cond_dict[src_key]= src
-        elif src_key == 'imagebind':
-            cond_dict[src_key] = src.view(bs,1,-1).to(device).to(dtype)
-        if do_classifier_free_guidance:
-            cond_dict['src_type'] = cond_dict['src_type'].repeat(2,1)
-            cond_dict['tgt_type'] = cond_dict['tgt_type'].repeat(2,1)
-            if 'text' in cond_dict:
-                cond_dict['text'] = cond_dict['text']+[''] * len(cond_dict['text'])
-            if 'imagebind' in cond_dict:
-                cond_dict['imagebind'] = torch.cat([cond_dict['imagebind'],cond_dict['imagebind']*0.],dim=0)
-            cond_dict['score'] = torch.cat([cond_dict['score'],cond_dict['score']*0.0+negative_score],dim=0)
-            bs = bs * 2
-        self.noise_scheduler.set_timesteps(num_inference_steps, device=device)
-        timesteps = self.noise_scheduler.timesteps
-        #print(cond_dict['imagebind'].shape)
-        z = self.get_input(cond_dict)
-        cond_dict.update(z)
-        # self.model.train()
-        
-        src_type = cond_dict['src_type']
-        if no_diffusion:
-            noisy_inputs_key = 'noisy_input' # hacl
-        else:
-            noisy_inputs_key = 'noisy_inputs'
-        cond_dict[noisy_inputs_key] = torch.randn(raw_bs,1,self.embed_dim).to(device).to(src_type)
-        if do_classifier_free_guidance:
-            cond_dict[noisy_inputs_key]  = cond_dict[noisy_inputs_key].repeat(2,1,1)
-        #cond_dict['noise_level'] = noise_level
-        # print("!!!!!!!!!!!!!train")
-
-        
-
-        steps = self.mae_token_num
-        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
-        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.noise_scheduler.order, 0)
-        with tqdm(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                timesteps = torch.ones(
-                    (cond_dict[noisy_inputs_key].shape[0],), device=cond_dict[noisy_inputs_key].device
-                )* t
-                noise_level = get_timestep_embedding(
-                     timesteps=timesteps, embedding_dim=cond_dict[noisy_inputs_key].shape[-1], flip_sin_to_cos=True, downscale_freq_shift=0)
-                cond_dict['noise_level'] = noise_level
-                (
-                input_embeds,
-                input_embeds_attn_mask,
-                cond_sequence_end_time_idx,
-                ) = self.get_input_sequence_and_mask(cond_dict)
-                model_input = input_embeds
-                model_input_mask = input_embeds_attn_mask
-                for _ in range(steps):
-                    output = self.model(
-                        inputs_embeds=model_input, attention_mask=model_input_mask
-                    )["last_hidden_state"]
-                    # Update the model input
-                    model_input = torch.cat([model_input, output[:, -1:, :]], dim=1)
-                    # Update the attention mask
-                    attention_mask_new_step = torch.ones((model_input_mask.size(0), 1)).to(
-                        model_input.device
-                    )
-                    model_input_mask = torch.cat(
-                        [model_input_mask, attention_mask_new_step], dim=1
-                    )
-                output = model_input[:, cond_sequence_end_time_idx:] #
-                #print(output.norm())
-                #prev_t = self.noise_scheduler.previous_timestep(t)
-                prev_t = self.noise_scheduler.config.num_train_timesteps // self.noise_scheduler.num_inference_steps
-                if prev_t >= 0 or force_guidence_t0:
-                    output = self.get_eps(t,cond_dict[noisy_inputs_key],output) # to eps transform
-                    if do_classifier_free_guidance:
-                        noise_pred_text,noise_pred_uncond  = output.chunk(2)
-                        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
-                    else:
-                        noise_pred = output
-
-                    latents = self.noise_scheduler.step(noise_pred, t, cond_dict[noisy_inputs_key][:raw_bs], **extra_step_kwargs, return_dict=False)[0]
-                    #print((latents== cond_dict[noisy_inputs_key][:raw_bs]).all())
-                    if do_classifier_free_guidance:
-                        latents = latents.repeat(2,1,1)
-                else:
-                    #print("HERE")
-                    latents = output[:raw_bs]
-                cond_dict[noisy_inputs_key] = latents
-                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.noise_scheduler.order == 0):
-                    progress_bar.update()
-        return cond_dict[noisy_inputs_key][:raw_bs], cond_dict
-
-    def forward(self,cond_dict,target=None,use_cache=None):
-        if self.diffusion:
-            if target is not None:
-                timesteps = torch.randint(
-                        0, self.noise_scheduler.config.num_train_timesteps, (target.shape[0],), device=target.device
-                    )
-                noise = torch.randn_like(target)
-                noisy_inputs = 10.0 * self.noise_scheduler.add_noise(target / 10.0,noise, timesteps)
-                noise_level = get_timestep_embedding(
-                    timesteps=timesteps, embedding_dim=noisy_inputs.shape[-1], flip_sin_to_cos=True, downscale_freq_shift=0)
-                cond_dict['noisy_inputs'] = noisy_inputs
-                cond_dict['noise_level'] = noise_level
-            else:
-                assert 'noisy_inputs' in cond_dict and 'noise_level' in cond_dict and 'noisy_inputs' in self.sequence_input_key and 'noise_level' in self.sequence_input_key
-
-        z = self.get_input(cond_dict)
-        cond_dict.update(z)
-        (
-            input_embeds,
-            input_embeds_attn_mask,
-            cond_sequence_end_time_idx,
-        ) = self.get_input_sequence_and_mask(cond_dict)
-        if target is not None:
-            input_embeds = torch.cat([input_embeds,target],dim=1)
-            input_embeds_attn_mask = torch.cat([input_embeds_attn_mask,torch.ones(input_embeds_attn_mask.shape[0],target.shape[1]).to(input_embeds_attn_mask)],dim=1)
-        output = self.model(
-                inputs_embeds=input_embeds, attention_mask=input_embeds_attn_mask,use_cache=use_cache
-        )
-        if self.training:
-            shifted_logits = output['last_hidden_state'][:,-1-target.shape[1]:-1]
-            loss = self.loss_fn(shifted_logits,target)
-            return loss,output
-        return output
-        
-    def get_input_item(self, batch, k):
-        ret = {}
-
-        if "log_mel_spec" in batch:
-            ret["fbank"] = (
-                batch["log_mel_spec"].unsqueeze(1).to(memory_format=torch.contiguous_format).float()
-            )
-        if 'stft' in batch:
-            ret["stft"] = batch["stft"].to(memory_format=torch.contiguous_format).float()
-        # ret["clip_label"] = clip_label.to(memory_format=torch.contiguous_format).float()
-        if 'waveform' in batch:
-            ret["waveform"] = batch["waveform"].to(memory_format=torch.contiguous_format).float()
-        if 'text' in batch:
-            ret["text"] = list(batch["text"])
-        if 'fname' in batch:
-            ret["fname"] = batch["fname"]
-
-        for key in batch.keys():
-            if key not in ret.keys():
-                ret[key] = batch[key]
-
-        return ret[k]
-
-    def get_input(self, batch):
-        cond_dict = {}
-        if len(self.cond_stage_model_metadata.keys()) > 0:
-            unconditional_cfg = False
-
-            for cond_model_key in self.cond_stage_model_metadata.keys():
-                cond_stage_key = self.cond_stage_model_metadata[cond_model_key][
-                    "cond_stage_key"
-                ]
-
-                # if(not self.training):
-                #     if(isinstance(self.cond_stage_models[self.cond_stage_model_metadata[cond_model_key]["model_idx"]], CLAPAudioEmbeddingClassifierFreev2)):
-                #         assert cond_stage_key == "text" # CLAP model should use text for evaluation
-
-                # The original data for conditioning
-                xc = self.get_input_item(batch, cond_stage_key)
-                if type(xc) == torch.Tensor:
-                    xc = xc.to(self.device)
-
-                c = self.get_learned_conditioning(
-                    xc, key=cond_model_key, unconditional_cfg=unconditional_cfg
-                )
-                cond_dict[cond_model_key] = c
-
-        return cond_dict
-
-    def instantiate_cond_stage(self, config):
-        self.cond_stage_model_metadata = {}
-
-        for i, cond_model_key in enumerate(config.keys()):
-            model = instantiate_from_config(config[cond_model_key])
-            self.cond_stage_models.append(model)
-            self.cond_stage_model_metadata[cond_model_key] = {
-                "model_idx": i,
-                "cond_stage_key": config[cond_model_key]["cond_stage_key"],
-                "conditioning_key": config[cond_model_key]["conditioning_key"],
-            }
-
-    def get_learned_conditioning(self, c, key, unconditional_cfg):
-        assert key in self.cond_stage_model_metadata.keys()
-
-        # Classifier-free guidance
-        if not unconditional_cfg:
-            c = self.cond_stage_models[
-                self.cond_stage_model_metadata[key]["model_idx"]
-            ](c)
-        else:
-            if isinstance(c, torch.Tensor):
-                batchsize = c.size(0)
-            elif isinstance(c, list):
-                batchsize = len(c)
-            else:
-                raise NotImplementedError()
-            c = self.cond_stage_models[
-                self.cond_stage_model_metadata[key]["model_idx"]
-            ].get_unconditional_condition(batchsize)
-
-        return c
-
-    def initialize_param_check_toolkit(self):
-        self.tracked_steps = 0
-        self.param_dict = {}
-
-    def statistic_require_grad_tensor_number(self, module, name=None):
-        requires_grad_num = 0
-        total_num = 0
-        require_grad_tensor = None
-        for p in module.parameters():
-            if p.requires_grad:
-                requires_grad_num += 1
-                if require_grad_tensor is None:
-                    require_grad_tensor = p
-            total_num += 1
-        print(
-            "Module: [%s] have %s trainable parameters out of %s total parameters (%.2f)"
-            % (name, requires_grad_num, total_num, requires_grad_num / total_num)
-        )
-        return require_grad_tensor
diff --git a/instructany2pix/utils/__init__.py b/instructany2pix/utils/__init__.py
deleted file mode 100644
index 3dbf4b6a7d5a51ccb5c1d4d654474ab25b998520..0000000000000000000000000000000000000000
--- a/instructany2pix/utils/__init__.py
+++ /dev/null
@@ -1,17 +0,0 @@
-
-import importlib
-def instantiate_from_config(config):
-    if not "target" in config:
-        if config == "__is_first_stage__":
-            return None
-        elif config == "__is_unconditional__":
-            return None
-        raise KeyError("Expected key `target` to instantiate.")
-    return get_obj_from_str(config["target"])(**config.get("params", dict()))
-
-def get_obj_from_str(string, reload=False):
-    module, cls = string.rsplit(".", 1)
-    if reload:
-        module_imp = importlib.import_module(module)
-        importlib.reload(module_imp)
-    return getattr(importlib.import_module(module, package=None), cls)
diff --git a/readme.MD b/readme.MD
new file mode 100644
index 0000000000000000000000000000000000000000..f2b9815ee43b2ff6ced1886c3adf55dd273ccf75
--- /dev/null
+++ b/readme.MD
@@ -0,0 +1,30 @@
+# InstructAny2Pix: Flexible Visual Editing via Multimodal Instruction Following
+
+PyTorch implementation of InstructAny2Pix: Flexible Visual Editing via Multimodal Instruction Following
+
+Shufan Li1*  Harkanwar Singh   Aditya Grover 
+
+University of California, Los Angeles
+
+
+![Img](assets/appendix1-3.png)
+## Installation
+
+```
+conda create --name instructany2pix python=3.10
+conda activate instructany2pix
+conda install ffmpeg
+pip3 install torch torchvision torchaudio
+pip3 install -r requirements.txt
+pip install git+https://github.com/facebookresearch/ImageBind.git --no-deps
+```
+
+## Running
+
+To serve gradio app, run
+```
+python serve.py
+```
+
+Alternatively, one can check `demo.ipynb` for a notebook demo
+
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..524febe337cd0da7755f9e0eaeadc124509ecfba
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,21 @@
+accelerate==0.21.0
+datasets==2.14.6
+diffusers==0.22.0
+transformers==4.34.1
+timm==0.6.13
+tokenizers==0.14.1
+torchlibrosa
+gradio
+pytorchvideo
+ftfy
+einops
+pytorch_lightning==2.1.2 
+sentencepiece==0.1.99
+protobuf==4.25.1
+bitsandbytes==0.41.2
+# torch
+# torchaudio
+# torchlibrosa
+# torchmetrics
+# torchvideo
+# torchvisio
\ No newline at end of file