Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

action_head--35000_checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b3cec4d149ab1e2105f4350ec2d63ed2c52d8b27ace1201d2b0fa0245fcb6090
+size 18984846

config.json

@@ -0,0 +1,308 @@
+{
+  "norm_stats": {
+    "libero_spatial_no_noops": {
+      "action": {
+        "mean": [
+          0.15312479436397552,
+          0.13707277178764343,
+          -0.15526802837848663,
+          -0.005176450591534376,
+          -0.01120874285697937,
+          -0.020194264128804207,
+          0.4578818082809448
+        ],
+        "std": [
+          0.41272708773612976,
+          0.34724321961402893,
+          0.50869220495224,
+          0.037266165018081665,
+          0.07244449853897095,
+          0.05762382969260216,
+          0.49827873706817627
+        ],
+        "max": [
+          0.9375,
+          0.9375,
+          0.9375,
+          0.1971428543329239,
+          0.33642858266830444,
+          0.375,
+          1.0
+        ],
+        "min": [
+          -0.9375,
+          -0.9375,
+          -0.9375,
+          -0.1875,
+          -0.3675000071525574,
+          -0.36000001430511475,
+          0.0
+        ],
+        "q01": [
+          -0.7454732114076613,
+          -0.6616071462631226,
+          -0.9375,
+          -0.1071428582072258,
+          -0.20678570866584778,
+          -0.1842857152223587,
+          0.0
+        ],
+        "q99": [
+          0.9375,
+          0.8758928775787354,
+          0.9321428537368774,
+          0.1039285734295845,
+          0.17678570747375488,
+          0.14571428298950195,
+          1.0
+        ],
+        "mask": [
+          true,
+          true,
+          true,
+          true,
+          true,
+          true,
+          false
+        ]
+      },
+      "proprio": {
+        "mean": [
+          -0.024462558329105377,
+          0.106529600918293,
+          1.0580483675003052,
+          3.0628468990325928,
+          -0.10464039444923401,
+          0.08307311683893204,
+          0.01995457336306572,
+          -0.020162804052233696
+        ],
+        "std": [
+          0.1101478561758995,
+          0.13784688711166382,
+          0.1044282391667366,
+          0.10451053828001022,
+          0.4112098217010498,
+          0.2176690548658371,
+          0.017260896041989326,
+          0.0171116404235363
+        ],
+        "max": [
+          0.1759040206670761,
+          0.3904820382595062,
+          1.3290715217590332,
+          3.4566118717193604,
+          1.2268599271774292,
+          1.0429412126541138,
+          0.041053611785173416,
+          0.000775813648942858
+        ],
+        "min": [
+          -0.3095473051071167,
+          -0.29250794649124146,
+          0.9095591306686401,
+          2.497488260269165,
+          -1.8006486892700195,
+          -0.7207611203193665,
+          -0.0004703797458205372,
+          -0.041536275297403336
+        ],
+        "q01": [
+          -0.2727657300233841,
+          -0.23721413239836692,
+          0.9160063165426254,
+          2.77949666261673,
+          -1.3187511622905732,
+          -0.41989982962608335,
+          0.001503719249740243,
+          -0.03989770736545324
+        ],
+        "q99": [
+          0.13529365032911292,
+          0.3629165390133857,
+          1.2862326657772063,
+          3.2829698753356933,
+          0.9332760351896285,
+          0.6325724506378171,
+          0.039933966137468815,
+          -0.001671919699292631
+        ]
+      },
+      "num_transitions": 52970,
+      "num_trajectories": 432
+    }
+  },
+  "n_action_bins": 256,
+  "vision_backbone_id": "dinosiglip-vit-so-224px",
+  "llm_backbone_id": "llama3.2-1b-pure",
+  "arch_specifier": "no-align+fused-gelu-mlp",
+  "output_projector_states": false,
+  "use_fused_vision_backbone": true,
+  "timm_model_ids": [
+    "vit_large_patch14_reg4_dinov2.lvd142m",
+    "vit_so400m_patch14_siglip_224"
+  ],
+  "timm_override_act_layers": [
+    null,
+    null
+  ],
+  "image_sizes": [
+    224,
+    224
+  ],
+  "image_resize_strategy": "resize-naive",
+  "hf_llm_id": "meta-llama/Llama-3.2-1B",
+  "llm_max_length": 2048,
+  "pad_token_id": 128256,
+  "pad_to_multiple_of": 64,
+  "text_config": {
+    "vocab_size": 128320,
+    "max_position_embeddings": 131072,
+    "hidden_size": 2048,
+    "intermediate_size": 8192,
+    "num_hidden_layers": 16,
+    "num_attention_heads": 32,
+    "num_key_value_heads": 8,
+    "hidden_act": "silu",
+    "initializer_range": 0.02,
+    "rms_norm_eps": 1e-06,
+    "pretraining_tp": 1,
+    "use_cache": true,
+    "rope_theta": 500000.0,
+    "rope_scaling": null,
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "mlp_bias": false,
+    "head_dim": 64,
+    "return_dict": true,
+    "output_hidden_states": false,
+    "output_attentions": false,
+    "torchscript": false,
+    "torch_dtype": "bfloat16",
+    "use_bfloat16": false,
+    "tf_legacy_loss": false,
+    "pruned_heads": {},
+    "tie_word_embeddings": false,
+    "chunk_size_feed_forward": 0,
+    "is_encoder_decoder": false,
+    "is_decoder": false,
+    "cross_attention_hidden_size": null,
+    "add_cross_attention": false,
+    "tie_encoder_decoder": false,
+    "max_length": 20,
+    "min_length": 0,
+    "do_sample": false,
+    "early_stopping": false,
+    "num_beams": 1,
+    "num_beam_groups": 1,
+    "diversity_penalty": 0.0,
+    "temperature": 1.0,
+    "top_k": 50,
+    "top_p": 1.0,
+    "typical_p": 1.0,
+    "repetition_penalty": 1.0,
+    "length_penalty": 1.0,
+    "no_repeat_ngram_size": 0,
+    "encoder_no_repeat_ngram_size": 0,
+    "bad_words_ids": null,
+    "num_return_sequences": 1,
+    "output_scores": false,
+    "return_dict_in_generate": false,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "remove_invalid_values": false,
+    "exponential_decay_length_penalty": null,
+    "suppress_tokens": null,
+    "begin_suppress_tokens": null,
+    "architectures": null,
+    "finetuning_task": null,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "tokenizer_class": null,
+    "prefix": null,
+    "bos_token_id": 1,
+    "pad_token_id": 128256,
+    "eos_token_id": 2,
+    "sep_token_id": null,
+    "decoder_start_token_id": null,
+    "task_specific_params": null,
+    "problem_type": null,
+    "_name_or_path": "",
+    "_attn_implementation_autoset": false,
+    "model_type": "llama"
+  },
+  "return_dict": true,
+  "output_hidden_states": false,
+  "output_attentions": false,
+  "torchscript": false,
+  "torch_dtype": "bfloat16",
+  "use_bfloat16": false,
+  "tf_legacy_loss": false,
+  "pruned_heads": {},
+  "tie_word_embeddings": true,
+  "chunk_size_feed_forward": 0,
+  "is_encoder_decoder": false,
+  "is_decoder": false,
+  "cross_attention_hidden_size": null,
+  "add_cross_attention": false,
+  "tie_encoder_decoder": false,
+  "max_length": 20,
+  "min_length": 0,
+  "do_sample": false,
+  "early_stopping": false,
+  "num_beams": 1,
+  "num_beam_groups": 1,
+  "diversity_penalty": 0.0,
+  "temperature": 1.0,
+  "top_k": 50,
+  "top_p": 1.0,
+  "typical_p": 1.0,
+  "repetition_penalty": 1.0,
+  "length_penalty": 1.0,
+  "no_repeat_ngram_size": 0,
+  "encoder_no_repeat_ngram_size": 0,
+  "bad_words_ids": null,
+  "num_return_sequences": 1,
+  "output_scores": false,
+  "return_dict_in_generate": false,
+  "forced_bos_token_id": null,
+  "forced_eos_token_id": null,
+  "remove_invalid_values": false,
+  "exponential_decay_length_penalty": null,
+  "suppress_tokens": null,
+  "begin_suppress_tokens": null,
+  "architectures": [
+    "PrismaticForConditionalGeneration"
+  ],
+  "finetuning_task": null,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "tokenizer_class": null,
+  "prefix": null,
+  "bos_token_id": null,
+  "eos_token_id": null,
+  "sep_token_id": null,
+  "decoder_start_token_id": null,
+  "task_specific_params": null,
+  "problem_type": null,
+  "_name_or_path": "/home/user1/workspace/juyi/lisaopenvla/hf-convert/prismatic-llama3.2-dinosiglip-224px-1b-vlm",
+  "_attn_implementation_autoset": true,
+  "transformers_version": "4.51.0",
+  "model_type": "openvla",
+  "auto_map": {
+    "AutoConfig": "configuration_prismatic.OpenVLAConfig",
+    "AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
+  }
+}

config.json.back.20250921_145013

@@ -0,0 +1,308 @@
+{
+  "norm_stats": {
+    "libero_spatial_no_noops": {
+      "action": {
+        "mean": [
+          0.15312479436397552,
+          0.13707277178764343,
+          -0.15526802837848663,
+          -0.005176450591534376,
+          -0.01120874285697937,
+          -0.020194264128804207,
+          0.4578818082809448
+        ],
+        "std": [
+          0.41272708773612976,
+          0.34724321961402893,
+          0.50869220495224,
+          0.037266165018081665,
+          0.07244449853897095,
+          0.05762382969260216,
+          0.49827873706817627
+        ],
+        "max": [
+          0.9375,
+          0.9375,
+          0.9375,
+          0.1971428543329239,
+          0.33642858266830444,
+          0.375,
+          1.0
+        ],
+        "min": [
+          -0.9375,
+          -0.9375,
+          -0.9375,
+          -0.1875,
+          -0.3675000071525574,
+          -0.36000001430511475,
+          0.0
+        ],
+        "q01": [
+          -0.7454732114076613,
+          -0.6616071462631226,
+          -0.9375,
+          -0.1071428582072258,
+          -0.20678570866584778,
+          -0.1842857152223587,
+          0.0
+        ],
+        "q99": [
+          0.9375,
+          0.8758928775787354,
+          0.9321428537368774,
+          0.1039285734295845,
+          0.17678570747375488,
+          0.14571428298950195,
+          1.0
+        ],
+        "mask": [
+          true,
+          true,
+          true,
+          true,
+          true,
+          true,
+          false
+        ]
+      },
+      "proprio": {
+        "mean": [
+          -0.024462558329105377,
+          0.106529600918293,
+          1.0580483675003052,
+          3.0628468990325928,
+          -0.10464039444923401,
+          0.08307311683893204,
+          0.01995457336306572,
+          -0.020162804052233696
+        ],
+        "std": [
+          0.1101478561758995,
+          0.13784688711166382,
+          0.1044282391667366,
+          0.10451053828001022,
+          0.4112098217010498,
+          0.2176690548658371,
+          0.017260896041989326,
+          0.0171116404235363
+        ],
+        "max": [
+          0.1759040206670761,
+          0.3904820382595062,
+          1.3290715217590332,
+          3.4566118717193604,
+          1.2268599271774292,
+          1.0429412126541138,
+          0.041053611785173416,
+          0.000775813648942858
+        ],
+        "min": [
+          -0.3095473051071167,
+          -0.29250794649124146,
+          0.9095591306686401,
+          2.497488260269165,
+          -1.8006486892700195,
+          -0.7207611203193665,
+          -0.0004703797458205372,
+          -0.041536275297403336
+        ],
+        "q01": [
+          -0.2727657300233841,
+          -0.23721413239836692,
+          0.9160063165426254,
+          2.77949666261673,
+          -1.3187511622905732,
+          -0.41989982962608335,
+          0.001503719249740243,
+          -0.03989770736545324
+        ],
+        "q99": [
+          0.13529365032911292,
+          0.3629165390133857,
+          1.2862326657772063,
+          3.2829698753356933,
+          0.9332760351896285,
+          0.6325724506378171,
+          0.039933966137468815,
+          -0.001671919699292631
+        ]
+      },
+      "num_transitions": 52970,
+      "num_trajectories": 432
+    }
+  },
+  "n_action_bins": 256,
+  "vision_backbone_id": "dinosiglip-vit-so-224px",
+  "llm_backbone_id": "llama3.2-1b-pure",
+  "arch_specifier": "no-align+fused-gelu-mlp",
+  "output_projector_states": false,
+  "use_fused_vision_backbone": true,
+  "timm_model_ids": [
+    "vit_large_patch14_reg4_dinov2.lvd142m",
+    "vit_so400m_patch14_siglip_224"
+  ],
+  "timm_override_act_layers": [
+    null,
+    null
+  ],
+  "image_sizes": [
+    224,
+    224
+  ],
+  "image_resize_strategy": "resize-naive",
+  "hf_llm_id": "meta-llama/Llama-3.2-1B",
+  "llm_max_length": 2048,
+  "pad_token_id": 128256,
+  "pad_to_multiple_of": 64,
+  "text_config": {
+    "vocab_size": 128320,
+    "max_position_embeddings": 131072,
+    "hidden_size": 2048,
+    "intermediate_size": 8192,
+    "num_hidden_layers": 16,
+    "num_attention_heads": 32,
+    "num_key_value_heads": 8,
+    "hidden_act": "silu",
+    "initializer_range": 0.02,
+    "rms_norm_eps": 1e-06,
+    "pretraining_tp": 1,
+    "use_cache": true,
+    "rope_theta": 500000.0,
+    "rope_scaling": null,
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "mlp_bias": false,
+    "head_dim": 64,
+    "return_dict": true,
+    "output_hidden_states": false,
+    "output_attentions": false,
+    "torchscript": false,
+    "torch_dtype": "bfloat16",
+    "use_bfloat16": false,
+    "tf_legacy_loss": false,
+    "pruned_heads": {},
+    "tie_word_embeddings": false,
+    "chunk_size_feed_forward": 0,
+    "is_encoder_decoder": false,
+    "is_decoder": false,
+    "cross_attention_hidden_size": null,
+    "add_cross_attention": false,
+    "tie_encoder_decoder": false,
+    "max_length": 20,
+    "min_length": 0,
+    "do_sample": false,
+    "early_stopping": false,
+    "num_beams": 1,
+    "num_beam_groups": 1,
+    "diversity_penalty": 0.0,
+    "temperature": 1.0,
+    "top_k": 50,
+    "top_p": 1.0,
+    "typical_p": 1.0,
+    "repetition_penalty": 1.0,
+    "length_penalty": 1.0,
+    "no_repeat_ngram_size": 0,
+    "encoder_no_repeat_ngram_size": 0,
+    "bad_words_ids": null,
+    "num_return_sequences": 1,
+    "output_scores": false,
+    "return_dict_in_generate": false,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "remove_invalid_values": false,
+    "exponential_decay_length_penalty": null,
+    "suppress_tokens": null,
+    "begin_suppress_tokens": null,
+    "architectures": null,
+    "finetuning_task": null,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "tokenizer_class": null,
+    "prefix": null,
+    "bos_token_id": 1,
+    "pad_token_id": 128256,
+    "eos_token_id": 2,
+    "sep_token_id": null,
+    "decoder_start_token_id": null,
+    "task_specific_params": null,
+    "problem_type": null,
+    "_name_or_path": "",
+    "_attn_implementation_autoset": false,
+    "model_type": "llama"
+  },
+  "return_dict": true,
+  "output_hidden_states": false,
+  "output_attentions": false,
+  "torchscript": false,
+  "torch_dtype": "bfloat16",
+  "use_bfloat16": false,
+  "tf_legacy_loss": false,
+  "pruned_heads": {},
+  "tie_word_embeddings": true,
+  "chunk_size_feed_forward": 0,
+  "is_encoder_decoder": false,
+  "is_decoder": false,
+  "cross_attention_hidden_size": null,
+  "add_cross_attention": false,
+  "tie_encoder_decoder": false,
+  "max_length": 20,
+  "min_length": 0,
+  "do_sample": false,
+  "early_stopping": false,
+  "num_beams": 1,
+  "num_beam_groups": 1,
+  "diversity_penalty": 0.0,
+  "temperature": 1.0,
+  "top_k": 50,
+  "top_p": 1.0,
+  "typical_p": 1.0,
+  "repetition_penalty": 1.0,
+  "length_penalty": 1.0,
+  "no_repeat_ngram_size": 0,
+  "encoder_no_repeat_ngram_size": 0,
+  "bad_words_ids": null,
+  "num_return_sequences": 1,
+  "output_scores": false,
+  "return_dict_in_generate": false,
+  "forced_bos_token_id": null,
+  "forced_eos_token_id": null,
+  "remove_invalid_values": false,
+  "exponential_decay_length_penalty": null,
+  "suppress_tokens": null,
+  "begin_suppress_tokens": null,
+  "architectures": [
+    "PrismaticForConditionalGeneration"
+  ],
+  "finetuning_task": null,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "tokenizer_class": null,
+  "prefix": null,
+  "bos_token_id": null,
+  "eos_token_id": null,
+  "sep_token_id": null,
+  "decoder_start_token_id": null,
+  "task_specific_params": null,
+  "problem_type": null,
+  "_name_or_path": "/home/user1/workspace/juyi/lisaopenvla/hf-convert/prismatic-llama3.2-dinosiglip-224px-1b-vlm",
+  "_attn_implementation_autoset": true,
+  "transformers_version": "4.51.0",
+  "model_type": "openvla",
+  "auto_map": {
+    "AutoConfig": "configuration_prismatic.OpenVLAConfig",
+    "AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
+  }
+}

config.json.back.20250921_182648

@@ -0,0 +1,308 @@
+{
+  "norm_stats": {
+    "libero_spatial_no_noops": {
+      "action": {
+        "mean": [
+          0.15312479436397552,
+          0.13707277178764343,
+          -0.15526802837848663,
+          -0.005176450591534376,
+          -0.01120874285697937,
+          -0.020194264128804207,
+          0.4578818082809448
+        ],
+        "std": [
+          0.41272708773612976,
+          0.34724321961402893,
+          0.50869220495224,
+          0.037266165018081665,
+          0.07244449853897095,
+          0.05762382969260216,
+          0.49827873706817627
+        ],
+        "max": [
+          0.9375,
+          0.9375,
+          0.9375,
+          0.1971428543329239,
+          0.33642858266830444,
+          0.375,
+          1.0
+        ],
+        "min": [
+          -0.9375,
+          -0.9375,
+          -0.9375,
+          -0.1875,
+          -0.3675000071525574,
+          -0.36000001430511475,
+          0.0
+        ],
+        "q01": [
+          -0.7454732114076613,
+          -0.6616071462631226,
+          -0.9375,
+          -0.1071428582072258,
+          -0.20678570866584778,
+          -0.1842857152223587,
+          0.0
+        ],
+        "q99": [
+          0.9375,
+          0.8758928775787354,
+          0.9321428537368774,
+          0.1039285734295845,
+          0.17678570747375488,
+          0.14571428298950195,
+          1.0
+        ],
+        "mask": [
+          true,
+          true,
+          true,
+          true,
+          true,
+          true,
+          false
+        ]
+      },
+      "proprio": {
+        "mean": [
+          -0.024462558329105377,
+          0.106529600918293,
+          1.0580483675003052,
+          3.0628468990325928,
+          -0.10464039444923401,
+          0.08307311683893204,
+          0.01995457336306572,
+          -0.020162804052233696
+        ],
+        "std": [
+          0.1101478561758995,
+          0.13784688711166382,
+          0.1044282391667366,
+          0.10451053828001022,
+          0.4112098217010498,
+          0.2176690548658371,
+          0.017260896041989326,
+          0.0171116404235363
+        ],
+        "max": [
+          0.1759040206670761,
+          0.3904820382595062,
+          1.3290715217590332,
+          3.4566118717193604,
+          1.2268599271774292,
+          1.0429412126541138,
+          0.041053611785173416,
+          0.000775813648942858
+        ],
+        "min": [
+          -0.3095473051071167,
+          -0.29250794649124146,
+          0.9095591306686401,
+          2.497488260269165,
+          -1.8006486892700195,
+          -0.7207611203193665,
+          -0.0004703797458205372,
+          -0.041536275297403336
+        ],
+        "q01": [
+          -0.2727657300233841,
+          -0.23721413239836692,
+          0.9160063165426254,
+          2.77949666261673,
+          -1.3187511622905732,
+          -0.41989982962608335,
+          0.001503719249740243,
+          -0.03989770736545324
+        ],
+        "q99": [
+          0.13529365032911292,
+          0.3629165390133857,
+          1.2862326657772063,
+          3.2829698753356933,
+          0.9332760351896285,
+          0.6325724506378171,
+          0.039933966137468815,
+          -0.001671919699292631
+        ]
+      },
+      "num_transitions": 52970,
+      "num_trajectories": 432
+    }
+  },
+  "n_action_bins": 256,
+  "vision_backbone_id": "dinosiglip-vit-so-224px",
+  "llm_backbone_id": "llama3.2-1b-pure",
+  "arch_specifier": "no-align+fused-gelu-mlp",
+  "output_projector_states": false,
+  "use_fused_vision_backbone": true,
+  "timm_model_ids": [
+    "vit_large_patch14_reg4_dinov2.lvd142m",
+    "vit_so400m_patch14_siglip_224"
+  ],
+  "timm_override_act_layers": [
+    null,
+    null
+  ],
+  "image_sizes": [
+    224,
+    224
+  ],
+  "image_resize_strategy": "resize-naive",
+  "hf_llm_id": "meta-llama/Llama-3.2-1B",
+  "llm_max_length": 2048,
+  "pad_token_id": 128256,
+  "pad_to_multiple_of": 64,
+  "text_config": {
+    "vocab_size": 128320,
+    "max_position_embeddings": 131072,
+    "hidden_size": 2048,
+    "intermediate_size": 8192,
+    "num_hidden_layers": 16,
+    "num_attention_heads": 32,
+    "num_key_value_heads": 8,
+    "hidden_act": "silu",
+    "initializer_range": 0.02,
+    "rms_norm_eps": 1e-06,
+    "pretraining_tp": 1,
+    "use_cache": true,
+    "rope_theta": 500000.0,
+    "rope_scaling": null,
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "mlp_bias": false,
+    "head_dim": 64,
+    "return_dict": true,
+    "output_hidden_states": false,
+    "output_attentions": false,
+    "torchscript": false,
+    "torch_dtype": "bfloat16",
+    "use_bfloat16": false,
+    "tf_legacy_loss": false,
+    "pruned_heads": {},
+    "tie_word_embeddings": false,
+    "chunk_size_feed_forward": 0,
+    "is_encoder_decoder": false,
+    "is_decoder": false,
+    "cross_attention_hidden_size": null,
+    "add_cross_attention": false,
+    "tie_encoder_decoder": false,
+    "max_length": 20,
+    "min_length": 0,
+    "do_sample": false,
+    "early_stopping": false,
+    "num_beams": 1,
+    "num_beam_groups": 1,
+    "diversity_penalty": 0.0,
+    "temperature": 1.0,
+    "top_k": 50,
+    "top_p": 1.0,
+    "typical_p": 1.0,
+    "repetition_penalty": 1.0,
+    "length_penalty": 1.0,
+    "no_repeat_ngram_size": 0,
+    "encoder_no_repeat_ngram_size": 0,
+    "bad_words_ids": null,
+    "num_return_sequences": 1,
+    "output_scores": false,
+    "return_dict_in_generate": false,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "remove_invalid_values": false,
+    "exponential_decay_length_penalty": null,
+    "suppress_tokens": null,
+    "begin_suppress_tokens": null,
+    "architectures": null,
+    "finetuning_task": null,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "tokenizer_class": null,
+    "prefix": null,
+    "bos_token_id": 1,
+    "pad_token_id": 128256,
+    "eos_token_id": 2,
+    "sep_token_id": null,
+    "decoder_start_token_id": null,
+    "task_specific_params": null,
+    "problem_type": null,
+    "_name_or_path": "",
+    "_attn_implementation_autoset": false,
+    "model_type": "llama"
+  },
+  "return_dict": true,
+  "output_hidden_states": false,
+  "output_attentions": false,
+  "torchscript": false,
+  "torch_dtype": "bfloat16",
+  "use_bfloat16": false,
+  "tf_legacy_loss": false,
+  "pruned_heads": {},
+  "tie_word_embeddings": true,
+  "chunk_size_feed_forward": 0,
+  "is_encoder_decoder": false,
+  "is_decoder": false,
+  "cross_attention_hidden_size": null,
+  "add_cross_attention": false,
+  "tie_encoder_decoder": false,
+  "max_length": 20,
+  "min_length": 0,
+  "do_sample": false,
+  "early_stopping": false,
+  "num_beams": 1,
+  "num_beam_groups": 1,
+  "diversity_penalty": 0.0,
+  "temperature": 1.0,
+  "top_k": 50,
+  "top_p": 1.0,
+  "typical_p": 1.0,
+  "repetition_penalty": 1.0,
+  "length_penalty": 1.0,
+  "no_repeat_ngram_size": 0,
+  "encoder_no_repeat_ngram_size": 0,
+  "bad_words_ids": null,
+  "num_return_sequences": 1,
+  "output_scores": false,
+  "return_dict_in_generate": false,
+  "forced_bos_token_id": null,
+  "forced_eos_token_id": null,
+  "remove_invalid_values": false,
+  "exponential_decay_length_penalty": null,
+  "suppress_tokens": null,
+  "begin_suppress_tokens": null,
+  "architectures": [
+    "PrismaticForConditionalGeneration"
+  ],
+  "finetuning_task": null,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "tokenizer_class": null,
+  "prefix": null,
+  "bos_token_id": null,
+  "eos_token_id": null,
+  "sep_token_id": null,
+  "decoder_start_token_id": null,
+  "task_specific_params": null,
+  "problem_type": null,
+  "_name_or_path": "/home/user1/workspace/juyi/lisaopenvla/hf-convert/prismatic-llama3.2-dinosiglip-224px-1b-vlm",
+  "_attn_implementation_autoset": true,
+  "transformers_version": "4.51.0",
+  "model_type": "openvla",
+  "auto_map": {
+    "AutoConfig": "configuration_prismatic.OpenVLAConfig",
+    "AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
+  }
+}

config.json.back.20250921_183706

@@ -0,0 +1,308 @@
+{
+  "norm_stats": {
+    "libero_spatial_no_noops": {
+      "action": {
+        "mean": [
+          0.15312479436397552,
+          0.13707277178764343,
+          -0.15526802837848663,
+          -0.005176450591534376,
+          -0.01120874285697937,
+          -0.020194264128804207,
+          0.4578818082809448
+        ],
+        "std": [
+          0.41272708773612976,
+          0.34724321961402893,
+          0.50869220495224,
+          0.037266165018081665,
+          0.07244449853897095,
+          0.05762382969260216,
+          0.49827873706817627
+        ],
+        "max": [
+          0.9375,
+          0.9375,
+          0.9375,
+          0.1971428543329239,
+          0.33642858266830444,
+          0.375,
+          1.0
+        ],
+        "min": [
+          -0.9375,
+          -0.9375,
+          -0.9375,
+          -0.1875,
+          -0.3675000071525574,
+          -0.36000001430511475,
+          0.0
+        ],
+        "q01": [
+          -0.7454732114076613,
+          -0.6616071462631226,
+          -0.9375,
+          -0.1071428582072258,
+          -0.20678570866584778,
+          -0.1842857152223587,
+          0.0
+        ],
+        "q99": [
+          0.9375,
+          0.8758928775787354,
+          0.9321428537368774,
+          0.1039285734295845,
+          0.17678570747375488,
+          0.14571428298950195,
+          1.0
+        ],
+        "mask": [
+          true,
+          true,
+          true,
+          true,
+          true,
+          true,
+          false
+        ]
+      },
+      "proprio": {
+        "mean": [
+          -0.024462558329105377,
+          0.106529600918293,
+          1.0580483675003052,
+          3.0628468990325928,
+          -0.10464039444923401,
+          0.08307311683893204,
+          0.01995457336306572,
+          -0.020162804052233696
+        ],
+        "std": [
+          0.1101478561758995,
+          0.13784688711166382,
+          0.1044282391667366,
+          0.10451053828001022,
+          0.4112098217010498,
+          0.2176690548658371,
+          0.017260896041989326,
+          0.0171116404235363
+        ],
+        "max": [
+          0.1759040206670761,
+          0.3904820382595062,
+          1.3290715217590332,
+          3.4566118717193604,
+          1.2268599271774292,
+          1.0429412126541138,
+          0.041053611785173416,
+          0.000775813648942858
+        ],
+        "min": [
+          -0.3095473051071167,
+          -0.29250794649124146,
+          0.9095591306686401,
+          2.497488260269165,
+          -1.8006486892700195,
+          -0.7207611203193665,
+          -0.0004703797458205372,
+          -0.041536275297403336
+        ],
+        "q01": [
+          -0.2727657300233841,
+          -0.23721413239836692,
+          0.9160063165426254,
+          2.77949666261673,
+          -1.3187511622905732,
+          -0.41989982962608335,
+          0.001503719249740243,
+          -0.03989770736545324
+        ],
+        "q99": [
+          0.13529365032911292,
+          0.3629165390133857,
+          1.2862326657772063,
+          3.2829698753356933,
+          0.9332760351896285,
+          0.6325724506378171,
+          0.039933966137468815,
+          -0.001671919699292631
+        ]
+      },
+      "num_transitions": 52970,
+      "num_trajectories": 432
+    }
+  },
+  "n_action_bins": 256,
+  "vision_backbone_id": "dinosiglip-vit-so-224px",
+  "llm_backbone_id": "llama3.2-1b-pure",
+  "arch_specifier": "no-align+fused-gelu-mlp",
+  "output_projector_states": false,
+  "use_fused_vision_backbone": true,
+  "timm_model_ids": [
+    "vit_large_patch14_reg4_dinov2.lvd142m",
+    "vit_so400m_patch14_siglip_224"
+  ],
+  "timm_override_act_layers": [
+    null,
+    null
+  ],
+  "image_sizes": [
+    224,
+    224
+  ],
+  "image_resize_strategy": "resize-naive",
+  "hf_llm_id": "meta-llama/Llama-3.2-1B",
+  "llm_max_length": 2048,
+  "pad_token_id": 128256,
+  "pad_to_multiple_of": 64,
+  "text_config": {
+    "vocab_size": 128320,
+    "max_position_embeddings": 131072,
+    "hidden_size": 2048,
+    "intermediate_size": 8192,
+    "num_hidden_layers": 16,
+    "num_attention_heads": 32,
+    "num_key_value_heads": 8,
+    "hidden_act": "silu",
+    "initializer_range": 0.02,
+    "rms_norm_eps": 1e-06,
+    "pretraining_tp": 1,
+    "use_cache": true,
+    "rope_theta": 500000.0,
+    "rope_scaling": null,
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "mlp_bias": false,
+    "head_dim": 64,
+    "return_dict": true,
+    "output_hidden_states": false,
+    "output_attentions": false,
+    "torchscript": false,
+    "torch_dtype": "bfloat16",
+    "use_bfloat16": false,
+    "tf_legacy_loss": false,
+    "pruned_heads": {},
+    "tie_word_embeddings": false,
+    "chunk_size_feed_forward": 0,
+    "is_encoder_decoder": false,
+    "is_decoder": false,
+    "cross_attention_hidden_size": null,
+    "add_cross_attention": false,
+    "tie_encoder_decoder": false,
+    "max_length": 20,
+    "min_length": 0,
+    "do_sample": false,
+    "early_stopping": false,
+    "num_beams": 1,
+    "num_beam_groups": 1,
+    "diversity_penalty": 0.0,
+    "temperature": 1.0,
+    "top_k": 50,
+    "top_p": 1.0,
+    "typical_p": 1.0,
+    "repetition_penalty": 1.0,
+    "length_penalty": 1.0,
+    "no_repeat_ngram_size": 0,
+    "encoder_no_repeat_ngram_size": 0,
+    "bad_words_ids": null,
+    "num_return_sequences": 1,
+    "output_scores": false,
+    "return_dict_in_generate": false,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "remove_invalid_values": false,
+    "exponential_decay_length_penalty": null,
+    "suppress_tokens": null,
+    "begin_suppress_tokens": null,
+    "architectures": null,
+    "finetuning_task": null,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "tokenizer_class": null,
+    "prefix": null,
+    "bos_token_id": 1,
+    "pad_token_id": 128256,
+    "eos_token_id": 2,
+    "sep_token_id": null,
+    "decoder_start_token_id": null,
+    "task_specific_params": null,
+    "problem_type": null,
+    "_name_or_path": "",
+    "_attn_implementation_autoset": false,
+    "model_type": "llama"
+  },
+  "return_dict": true,
+  "output_hidden_states": false,
+  "output_attentions": false,
+  "torchscript": false,
+  "torch_dtype": "bfloat16",
+  "use_bfloat16": false,
+  "tf_legacy_loss": false,
+  "pruned_heads": {},
+  "tie_word_embeddings": true,
+  "chunk_size_feed_forward": 0,
+  "is_encoder_decoder": false,
+  "is_decoder": false,
+  "cross_attention_hidden_size": null,
+  "add_cross_attention": false,
+  "tie_encoder_decoder": false,
+  "max_length": 20,
+  "min_length": 0,
+  "do_sample": false,
+  "early_stopping": false,
+  "num_beams": 1,
+  "num_beam_groups": 1,
+  "diversity_penalty": 0.0,
+  "temperature": 1.0,
+  "top_k": 50,
+  "top_p": 1.0,
+  "typical_p": 1.0,
+  "repetition_penalty": 1.0,
+  "length_penalty": 1.0,
+  "no_repeat_ngram_size": 0,
+  "encoder_no_repeat_ngram_size": 0,
+  "bad_words_ids": null,
+  "num_return_sequences": 1,
+  "output_scores": false,
+  "return_dict_in_generate": false,
+  "forced_bos_token_id": null,
+  "forced_eos_token_id": null,
+  "remove_invalid_values": false,
+  "exponential_decay_length_penalty": null,
+  "suppress_tokens": null,
+  "begin_suppress_tokens": null,
+  "architectures": [
+    "PrismaticForConditionalGeneration"
+  ],
+  "finetuning_task": null,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "tokenizer_class": null,
+  "prefix": null,
+  "bos_token_id": null,
+  "eos_token_id": null,
+  "sep_token_id": null,
+  "decoder_start_token_id": null,
+  "task_specific_params": null,
+  "problem_type": null,
+  "_name_or_path": "/home/user1/workspace/juyi/lisaopenvla/hf-convert/prismatic-llama3.2-dinosiglip-224px-1b-vlm",
+  "_attn_implementation_autoset": true,
+  "transformers_version": "4.51.0",
+  "model_type": "openvla",
+  "auto_map": {
+    "AutoConfig": "configuration_prismatic.OpenVLAConfig",
+    "AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
+  }
+}

config.json.back.20250922_064615

@@ -0,0 +1,308 @@
+{
+  "norm_stats": {
+    "libero_spatial_no_noops": {
+      "action": {
+        "mean": [
+          0.15312479436397552,
+          0.13707277178764343,
+          -0.15526802837848663,
+          -0.005176450591534376,
+          -0.01120874285697937,
+          -0.020194264128804207,
+          0.4578818082809448
+        ],
+        "std": [
+          0.41272708773612976,
+          0.34724321961402893,
+          0.50869220495224,
+          0.037266165018081665,
+          0.07244449853897095,
+          0.05762382969260216,
+          0.49827873706817627
+        ],
+        "max": [
+          0.9375,
+          0.9375,
+          0.9375,
+          0.1971428543329239,
+          0.33642858266830444,
+          0.375,
+          1.0
+        ],
+        "min": [
+          -0.9375,
+          -0.9375,
+          -0.9375,
+          -0.1875,
+          -0.3675000071525574,
+          -0.36000001430511475,
+          0.0
+        ],
+        "q01": [
+          -0.7454732114076613,
+          -0.6616071462631226,
+          -0.9375,
+          -0.1071428582072258,
+          -0.20678570866584778,
+          -0.1842857152223587,
+          0.0
+        ],
+        "q99": [
+          0.9375,
+          0.8758928775787354,
+          0.9321428537368774,
+          0.1039285734295845,
+          0.17678570747375488,
+          0.14571428298950195,
+          1.0
+        ],
+        "mask": [
+          true,
+          true,
+          true,
+          true,
+          true,
+          true,
+          false
+        ]
+      },
+      "proprio": {
+        "mean": [
+          -0.024462558329105377,
+          0.106529600918293,
+          1.0580483675003052,
+          3.0628468990325928,
+          -0.10464039444923401,
+          0.08307311683893204,
+          0.01995457336306572,
+          -0.020162804052233696
+        ],
+        "std": [
+          0.1101478561758995,
+          0.13784688711166382,
+          0.1044282391667366,
+          0.10451053828001022,
+          0.4112098217010498,
+          0.2176690548658371,
+          0.017260896041989326,
+          0.0171116404235363
+        ],
+        "max": [
+          0.1759040206670761,
+          0.3904820382595062,
+          1.3290715217590332,
+          3.4566118717193604,
+          1.2268599271774292,
+          1.0429412126541138,
+          0.041053611785173416,
+          0.000775813648942858
+        ],
+        "min": [
+          -0.3095473051071167,
+          -0.29250794649124146,
+          0.9095591306686401,
+          2.497488260269165,
+          -1.8006486892700195,
+          -0.7207611203193665,
+          -0.0004703797458205372,
+          -0.041536275297403336
+        ],
+        "q01": [
+          -0.2727657300233841,
+          -0.23721413239836692,
+          0.9160063165426254,
+          2.77949666261673,
+          -1.3187511622905732,
+          -0.41989982962608335,
+          0.001503719249740243,
+          -0.03989770736545324
+        ],
+        "q99": [
+          0.13529365032911292,
+          0.3629165390133857,
+          1.2862326657772063,
+          3.2829698753356933,
+          0.9332760351896285,
+          0.6325724506378171,
+          0.039933966137468815,
+          -0.001671919699292631
+        ]
+      },
+      "num_transitions": 52970,
+      "num_trajectories": 432
+    }
+  },
+  "n_action_bins": 256,
+  "vision_backbone_id": "dinosiglip-vit-so-224px",
+  "llm_backbone_id": "llama3.2-1b-pure",
+  "arch_specifier": "no-align+fused-gelu-mlp",
+  "output_projector_states": false,
+  "use_fused_vision_backbone": true,
+  "timm_model_ids": [
+    "vit_large_patch14_reg4_dinov2.lvd142m",
+    "vit_so400m_patch14_siglip_224"
+  ],
+  "timm_override_act_layers": [
+    null,
+    null
+  ],
+  "image_sizes": [
+    224,
+    224
+  ],
+  "image_resize_strategy": "resize-naive",
+  "hf_llm_id": "meta-llama/Llama-3.2-1B",
+  "llm_max_length": 2048,
+  "pad_token_id": 128256,
+  "pad_to_multiple_of": 64,
+  "text_config": {
+    "vocab_size": 128320,
+    "max_position_embeddings": 131072,
+    "hidden_size": 2048,
+    "intermediate_size": 8192,
+    "num_hidden_layers": 16,
+    "num_attention_heads": 32,
+    "num_key_value_heads": 8,
+    "hidden_act": "silu",
+    "initializer_range": 0.02,
+    "rms_norm_eps": 1e-06,
+    "pretraining_tp": 1,
+    "use_cache": true,
+    "rope_theta": 500000.0,
+    "rope_scaling": null,
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "mlp_bias": false,
+    "head_dim": 64,
+    "return_dict": true,
+    "output_hidden_states": false,
+    "output_attentions": false,
+    "torchscript": false,
+    "torch_dtype": "bfloat16",
+    "use_bfloat16": false,
+    "tf_legacy_loss": false,
+    "pruned_heads": {},
+    "tie_word_embeddings": false,
+    "chunk_size_feed_forward": 0,
+    "is_encoder_decoder": false,
+    "is_decoder": false,
+    "cross_attention_hidden_size": null,
+    "add_cross_attention": false,
+    "tie_encoder_decoder": false,
+    "max_length": 20,
+    "min_length": 0,
+    "do_sample": false,
+    "early_stopping": false,
+    "num_beams": 1,
+    "num_beam_groups": 1,
+    "diversity_penalty": 0.0,
+    "temperature": 1.0,
+    "top_k": 50,
+    "top_p": 1.0,
+    "typical_p": 1.0,
+    "repetition_penalty": 1.0,
+    "length_penalty": 1.0,
+    "no_repeat_ngram_size": 0,
+    "encoder_no_repeat_ngram_size": 0,
+    "bad_words_ids": null,
+    "num_return_sequences": 1,
+    "output_scores": false,
+    "return_dict_in_generate": false,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "remove_invalid_values": false,
+    "exponential_decay_length_penalty": null,
+    "suppress_tokens": null,
+    "begin_suppress_tokens": null,
+    "architectures": null,
+    "finetuning_task": null,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "tokenizer_class": null,
+    "prefix": null,
+    "bos_token_id": 1,
+    "pad_token_id": 128256,
+    "eos_token_id": 2,
+    "sep_token_id": null,
+    "decoder_start_token_id": null,
+    "task_specific_params": null,
+    "problem_type": null,
+    "_name_or_path": "",
+    "_attn_implementation_autoset": false,
+    "model_type": "llama"
+  },
+  "return_dict": true,
+  "output_hidden_states": false,
+  "output_attentions": false,
+  "torchscript": false,
+  "torch_dtype": "bfloat16",
+  "use_bfloat16": false,
+  "tf_legacy_loss": false,
+  "pruned_heads": {},
+  "tie_word_embeddings": true,
+  "chunk_size_feed_forward": 0,
+  "is_encoder_decoder": false,
+  "is_decoder": false,
+  "cross_attention_hidden_size": null,
+  "add_cross_attention": false,
+  "tie_encoder_decoder": false,
+  "max_length": 20,
+  "min_length": 0,
+  "do_sample": false,
+  "early_stopping": false,
+  "num_beams": 1,
+  "num_beam_groups": 1,
+  "diversity_penalty": 0.0,
+  "temperature": 1.0,
+  "top_k": 50,
+  "top_p": 1.0,
+  "typical_p": 1.0,
+  "repetition_penalty": 1.0,
+  "length_penalty": 1.0,
+  "no_repeat_ngram_size": 0,
+  "encoder_no_repeat_ngram_size": 0,
+  "bad_words_ids": null,
+  "num_return_sequences": 1,
+  "output_scores": false,
+  "return_dict_in_generate": false,
+  "forced_bos_token_id": null,
+  "forced_eos_token_id": null,
+  "remove_invalid_values": false,
+  "exponential_decay_length_penalty": null,
+  "suppress_tokens": null,
+  "begin_suppress_tokens": null,
+  "architectures": [
+    "PrismaticForConditionalGeneration"
+  ],
+  "finetuning_task": null,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "tokenizer_class": null,
+  "prefix": null,
+  "bos_token_id": null,
+  "eos_token_id": null,
+  "sep_token_id": null,
+  "decoder_start_token_id": null,
+  "task_specific_params": null,
+  "problem_type": null,
+  "_name_or_path": "/home/user1/workspace/juyi/lisaopenvla/hf-convert/prismatic-llama3.2-dinosiglip-224px-1b-vlm",
+  "_attn_implementation_autoset": true,
+  "transformers_version": "4.51.0",
+  "model_type": "openvla",
+  "auto_map": {
+    "AutoConfig": "configuration_prismatic.OpenVLAConfig",
+    "AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
+  }
+}

config.json.back.20250922_073803

@@ -0,0 +1,308 @@
+{
+  "norm_stats": {
+    "libero_spatial_no_noops": {
+      "action": {
+        "mean": [
+          0.15312479436397552,
+          0.13707277178764343,
+          -0.15526802837848663,
+          -0.005176450591534376,
+          -0.01120874285697937,
+          -0.020194264128804207,
+          0.4578818082809448
+        ],
+        "std": [
+          0.41272708773612976,
+          0.34724321961402893,
+          0.50869220495224,
+          0.037266165018081665,
+          0.07244449853897095,
+          0.05762382969260216,
+          0.49827873706817627
+        ],
+        "max": [
+          0.9375,
+          0.9375,
+          0.9375,
+          0.1971428543329239,
+          0.33642858266830444,
+          0.375,
+          1.0
+        ],
+        "min": [
+          -0.9375,
+          -0.9375,
+          -0.9375,
+          -0.1875,
+          -0.3675000071525574,
+          -0.36000001430511475,
+          0.0
+        ],
+        "q01": [
+          -0.7454732114076613,
+          -0.6616071462631226,
+          -0.9375,
+          -0.1071428582072258,
+          -0.20678570866584778,
+          -0.1842857152223587,
+          0.0
+        ],
+        "q99": [
+          0.9375,
+          0.8758928775787354,
+          0.9321428537368774,
+          0.1039285734295845,
+          0.17678570747375488,
+          0.14571428298950195,
+          1.0
+        ],
+        "mask": [
+          true,
+          true,
+          true,
+          true,
+          true,
+          true,
+          false
+        ]
+      },
+      "proprio": {
+        "mean": [
+          -0.024462558329105377,
+          0.106529600918293,
+          1.0580483675003052,
+          3.0628468990325928,
+          -0.10464039444923401,
+          0.08307311683893204,
+          0.01995457336306572,
+          -0.020162804052233696
+        ],
+        "std": [
+          0.1101478561758995,
+          0.13784688711166382,
+          0.1044282391667366,
+          0.10451053828001022,
+          0.4112098217010498,
+          0.2176690548658371,
+          0.017260896041989326,
+          0.0171116404235363
+        ],
+        "max": [
+          0.1759040206670761,
+          0.3904820382595062,
+          1.3290715217590332,
+          3.4566118717193604,
+          1.2268599271774292,
+          1.0429412126541138,
+          0.041053611785173416,
+          0.000775813648942858
+        ],
+        "min": [
+          -0.3095473051071167,
+          -0.29250794649124146,
+          0.9095591306686401,
+          2.497488260269165,
+          -1.8006486892700195,
+          -0.7207611203193665,
+          -0.0004703797458205372,
+          -0.041536275297403336
+        ],
+        "q01": [
+          -0.2727657300233841,
+          -0.23721413239836692,
+          0.9160063165426254,
+          2.77949666261673,
+          -1.3187511622905732,
+          -0.41989982962608335,
+          0.001503719249740243,
+          -0.03989770736545324
+        ],
+        "q99": [
+          0.13529365032911292,
+          0.3629165390133857,
+          1.2862326657772063,
+          3.2829698753356933,
+          0.9332760351896285,
+          0.6325724506378171,
+          0.039933966137468815,
+          -0.001671919699292631
+        ]
+      },
+      "num_transitions": 52970,
+      "num_trajectories": 432
+    }
+  },
+  "n_action_bins": 256,
+  "vision_backbone_id": "dinosiglip-vit-so-224px",
+  "llm_backbone_id": "llama3.2-1b-pure",
+  "arch_specifier": "no-align+fused-gelu-mlp",
+  "output_projector_states": false,
+  "use_fused_vision_backbone": true,
+  "timm_model_ids": [
+    "vit_large_patch14_reg4_dinov2.lvd142m",
+    "vit_so400m_patch14_siglip_224"
+  ],
+  "timm_override_act_layers": [
+    null,
+    null
+  ],
+  "image_sizes": [
+    224,
+    224
+  ],
+  "image_resize_strategy": "resize-naive",
+  "hf_llm_id": "meta-llama/Llama-3.2-1B",
+  "llm_max_length": 2048,
+  "pad_token_id": 128256,
+  "pad_to_multiple_of": 64,
+  "text_config": {
+    "vocab_size": 128320,
+    "max_position_embeddings": 131072,
+    "hidden_size": 2048,
+    "intermediate_size": 8192,
+    "num_hidden_layers": 16,
+    "num_attention_heads": 32,
+    "num_key_value_heads": 8,
+    "hidden_act": "silu",
+    "initializer_range": 0.02,
+    "rms_norm_eps": 1e-06,
+    "pretraining_tp": 1,
+    "use_cache": true,
+    "rope_theta": 500000.0,
+    "rope_scaling": null,
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "mlp_bias": false,
+    "head_dim": 64,
+    "return_dict": true,
+    "output_hidden_states": false,
+    "output_attentions": false,
+    "torchscript": false,
+    "torch_dtype": "bfloat16",
+    "use_bfloat16": false,
+    "tf_legacy_loss": false,
+    "pruned_heads": {},
+    "tie_word_embeddings": false,
+    "chunk_size_feed_forward": 0,
+    "is_encoder_decoder": false,
+    "is_decoder": false,
+    "cross_attention_hidden_size": null,
+    "add_cross_attention": false,
+    "tie_encoder_decoder": false,
+    "max_length": 20,
+    "min_length": 0,
+    "do_sample": false,
+    "early_stopping": false,
+    "num_beams": 1,
+    "num_beam_groups": 1,
+    "diversity_penalty": 0.0,
+    "temperature": 1.0,
+    "top_k": 50,
+    "top_p": 1.0,
+    "typical_p": 1.0,
+    "repetition_penalty": 1.0,
+    "length_penalty": 1.0,
+    "no_repeat_ngram_size": 0,
+    "encoder_no_repeat_ngram_size": 0,
+    "bad_words_ids": null,
+    "num_return_sequences": 1,
+    "output_scores": false,
+    "return_dict_in_generate": false,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "remove_invalid_values": false,
+    "exponential_decay_length_penalty": null,
+    "suppress_tokens": null,
+    "begin_suppress_tokens": null,
+    "architectures": null,
+    "finetuning_task": null,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "tokenizer_class": null,
+    "prefix": null,
+    "bos_token_id": 1,
+    "pad_token_id": 128256,
+    "eos_token_id": 2,
+    "sep_token_id": null,
+    "decoder_start_token_id": null,
+    "task_specific_params": null,
+    "problem_type": null,
+    "_name_or_path": "",
+    "_attn_implementation_autoset": false,
+    "model_type": "llama"
+  },
+  "return_dict": true,
+  "output_hidden_states": false,
+  "output_attentions": false,
+  "torchscript": false,
+  "torch_dtype": "bfloat16",
+  "use_bfloat16": false,
+  "tf_legacy_loss": false,
+  "pruned_heads": {},
+  "tie_word_embeddings": true,
+  "chunk_size_feed_forward": 0,
+  "is_encoder_decoder": false,
+  "is_decoder": false,
+  "cross_attention_hidden_size": null,
+  "add_cross_attention": false,
+  "tie_encoder_decoder": false,
+  "max_length": 20,
+  "min_length": 0,
+  "do_sample": false,
+  "early_stopping": false,
+  "num_beams": 1,
+  "num_beam_groups": 1,
+  "diversity_penalty": 0.0,
+  "temperature": 1.0,
+  "top_k": 50,
+  "top_p": 1.0,
+  "typical_p": 1.0,
+  "repetition_penalty": 1.0,
+  "length_penalty": 1.0,
+  "no_repeat_ngram_size": 0,
+  "encoder_no_repeat_ngram_size": 0,
+  "bad_words_ids": null,
+  "num_return_sequences": 1,
+  "output_scores": false,
+  "return_dict_in_generate": false,
+  "forced_bos_token_id": null,
+  "forced_eos_token_id": null,
+  "remove_invalid_values": false,
+  "exponential_decay_length_penalty": null,
+  "suppress_tokens": null,
+  "begin_suppress_tokens": null,
+  "architectures": [
+    "PrismaticForConditionalGeneration"
+  ],
+  "finetuning_task": null,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "tokenizer_class": null,
+  "prefix": null,
+  "bos_token_id": null,
+  "eos_token_id": null,
+  "sep_token_id": null,
+  "decoder_start_token_id": null,
+  "task_specific_params": null,
+  "problem_type": null,
+  "_name_or_path": "/home/user1/workspace/juyi/lisaopenvla/hf-convert/prismatic-llama3.2-dinosiglip-224px-1b-vlm",
+  "_attn_implementation_autoset": true,
+  "transformers_version": "4.51.0",
+  "model_type": "openvla",
+  "auto_map": {
+    "AutoConfig": "configuration_prismatic.OpenVLAConfig",
+    "AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
+  }
+}

configuration_prismatic.py

@@ -0,0 +1,142 @@
+"""
+configuration_prismatic.py
+
+HuggingFace-style configuration definition for Prismatic VLMs, inheriting from `transformers.PretrainedConfig`.
+Default configuration specifies `siglip-224px+7b`.
+"""
+
+from typing import Any, Dict, List, Optional
+
+from transformers import PretrainedConfig
+from transformers.models.auto import CONFIG_MAPPING
+
+# === Utilities for Mapping Prismatic names to HF names ===
+# fmt: off
+VISION_BACKBONE_TO_RESOLUTION: Dict[str, List[int]] = {
+    "clip-vit-l": [224], "siglip-vit-so400m": [224], "dinov2-vit-l": [224], "in1k-vit-l": [224],
+
+    "clip-vit-l-336px": [336],
+    "siglip-vit-so400m-384px": [384],
+
+    "dinoclip-vit-l-336px": [336, 336],
+    "dinosiglip-vit-so-224px": [224, 224],
+    "dinosiglip-vit-so-384px": [384, 384],
+}
+VISION_BACKBONE_TO_TIMM_ID: Dict[str, List[str]] = {
+    "clip-vit-l": ["vit_large_patch14_clip_224.openai"],
+    "clip-vit-l-336px": ["vit_large_patch14_clip_336.openai"],
+
+    "dinov2-vit-l": ["vit_large_patch14_reg4_dinov2.lvd142m"],
+    "in1k-vit-l": ["vit_large_patch16_224.augreg_in21k_ft_in1k"],
+
+    "siglip-vit-so400m": ["vit_so400m_patch14_siglip_224"],
+    "siglip-vit-so400m-384px": ["vit_so400m_patch14_siglip_384"],
+
+    "dinoclip-vit-l-336px": ["vit_large_patch14_reg4_dinov2.lvd142m", "vit_large_patch14_clip_336.openai"],
+    "dinosiglip-vit-so-224px": ["vit_large_patch14_reg4_dinov2.lvd142m", "vit_so400m_patch14_siglip_224"],
+    "dinosiglip-vit-so-384px": ["vit_large_patch14_reg4_dinov2.lvd142m", "vit_so400m_patch14_siglip_384"],
+}
+TIMM_OVERRIDE_ACT_LAYER: Dict[str, List[Optional[str]]] = {
+    "clip-vit-l": ["quick_gelu"], "clip-vit-l-336px": ["quick_gelu"],
+    "dinov2-vit-l": [None], "in1k-vit-l": [None],
+    "siglip-vit-so400m": [None], "siglip-vit-so400m-384px": [None],
+    "dinoclip-vit-l-336px": [None, "quick_gelu"],
+    "dinosiglip-vit-so-224px": [None, None], "dinosiglip-vit-so-384px": [None, None]
+}
+
+LLM_BACKBONE_TO_HF_PATH = {
+    "llama2-7b-pure": "meta-llama/Llama-2-7b-hf", "llama2-13b-pure": "meta-llama/Llama-2-13b-hf",
+    "llama2-7b-chat": "meta-llama/Llama-2-7b-chat-hf", "llama2-13b-chat": "meta-llama/Llama-2-13b-chat-hf",
+
+    "vicuna-v15-7b": "lmsys/vicuna-7b-v1.5", "vicuna-v15-13b": "lmsys/vicuna-13b-v1.5",
+
+    "mistral-v0.1-7b-pure": "mistralai/Mistral-7B-v0.1",
+    "mistral-v0.1-7b-instruct": "mistralai/Mistral-7B-Instruct-v0.1",
+
+    "phi-2-3b": "microsoft/phi-2",
+    "llama3.2-1b-pure": "meta-llama/Llama-3.2-1B",
+}
+LLM_BACKBONE_TO_HF_METACLASS = {
+    "llama2-7b-pure": "llama", "llama2-13b-pure": "llama", "llama2-7b-chat": "llama", "llama2-13b-chat": "llama",
+    "vicuna-v15-7b": "llama", "vicuna-v15-13b": "llama",
+
+    "mistral-v0.1-7b-pure": "mistral", "mistral-v0.1-7b-instruct": "mistral",
+
+    "phi-2-3b": "phi",
+    "llama3.2-1b-pure": "llama",
+}
+
+VALID_VISION_BACKBONES = set(VISION_BACKBONE_TO_RESOLUTION.keys())
+VALID_LLM_BACKBONES = set(LLM_BACKBONE_TO_HF_PATH)
+# fmt: on
+
+
+class PrismaticConfig(PretrainedConfig):
+    model_type: str = "prismatic"
+    is_composition: bool = False
+
+    def __init__(
+        self,
+        vision_backbone_id: str = "siglip-vit-so400m",
+        llm_backbone_id: str = "vicuna-v15-7b",
+        arch_specifier: str = "no-align+gelu-mlp",
+        use_fused_vision_backbone: Optional[bool] = None,
+        image_resize_strategy: str = "letterbox",
+        text_config: Optional[Dict[str, Any]] = None,
+        llm_max_length: int = 2048,
+        pad_token_id: int = 32000,
+        pad_to_multiple_of: int = 64,
+        output_projector_states: bool = False,
+        **kwargs: str,
+    ) -> None:
+        if vision_backbone_id not in VALID_VISION_BACKBONES:
+            raise ValueError(f"Vision backbone `{vision_backbone_id}` not in {VALID_VISION_BACKBONES = }")
+
+        if llm_backbone_id not in VALID_LLM_BACKBONES:
+            raise ValueError(f"LLM backbone `{llm_backbone_id}` not in {VALID_LLM_BACKBONES = }")
+
+        # Set Prismatic Configuration Fields
+        self.vision_backbone_id = vision_backbone_id
+        self.llm_backbone_id = llm_backbone_id
+        self.arch_specifier = arch_specifier
+        self.output_projector_states = output_projector_states
+
+        # [Contract] All vision backbone parameters are lists =>> supports fused backbones with different preprocessing
+        self.use_fused_vision_backbone = (
+            use_fused_vision_backbone
+            if use_fused_vision_backbone is not None
+            else any(self.vision_backbone_id.startswith(v) for v in ["dinoclip", "dinosiglip"])
+        )
+
+        self.timm_model_ids = VISION_BACKBONE_TO_TIMM_ID[self.vision_backbone_id]
+        self.timm_override_act_layers = TIMM_OVERRIDE_ACT_LAYER[self.vision_backbone_id]
+        self.image_sizes = VISION_BACKBONE_TO_RESOLUTION[self.vision_backbone_id]
+        self.image_resize_strategy = image_resize_strategy
+
+        self.hf_llm_id = LLM_BACKBONE_TO_HF_PATH[self.llm_backbone_id]
+        self.llm_max_length = llm_max_length
+        self.pad_token_id, self.pad_to_multiple_of = pad_token_id, pad_to_multiple_of
+
+        # [IMPORTANT] HF Utilities actually look for a `text_config` field... we need to use that specific naming!
+        self.text_config = (
+            CONFIG_MAPPING[LLM_BACKBONE_TO_HF_METACLASS[self.llm_backbone_id]](**text_config)
+            if text_config is not None
+            else CONFIG_MAPPING[LLM_BACKBONE_TO_HF_METACLASS[self.llm_backbone_id]]()
+        )
+
+        # Dispatch **kwargs to super() =>> note that `pad_token_id` collides, so we pass it in here as well...
+        super().__init__(pad_token_id=pad_token_id, **kwargs)
+
+
+class OpenVLAConfig(PrismaticConfig):
+    model_type: str = "openvla"
+
+    def __init__(
+        self,
+        norm_stats: Optional[Dict[str, Dict[str, Dict[str, Dict[str, List[float]]]]]] = None,
+        n_action_bins: int = 256,
+        **kwargs: str,
+    ) -> None:
+        self.norm_stats, self.n_action_bins = norm_stats, n_action_bins
+
+        super().__init__(**kwargs)

dataset_statistics.json

@@ -0,0 +1,133 @@
+{
+  "libero_spatial_no_noops": {
+    "action": {
+      "mean": [
+        0.15312479436397552,
+        0.13707277178764343,
+        -0.15526802837848663,
+        -0.005176450591534376,
+        -0.01120874285697937,
+        -0.020194264128804207,
+        0.4578818082809448
+      ],
+      "std": [
+        0.41272708773612976,
+        0.34724321961402893,
+        0.50869220495224,
+        0.037266165018081665,
+        0.07244449853897095,
+        0.05762382969260216,
+        0.49827873706817627
+      ],
+      "max": [
+        0.9375,
+        0.9375,
+        0.9375,
+        0.1971428543329239,
+        0.33642858266830444,
+        0.375,
+        1.0
+      ],
+      "min": [
+        -0.9375,
+        -0.9375,
+        -0.9375,
+        -0.1875,
+        -0.3675000071525574,
+        -0.36000001430511475,
+        0.0
+      ],
+      "q01": [
+        -0.7454732114076613,
+        -0.6616071462631226,
+        -0.9375,
+        -0.1071428582072258,
+        -0.20678570866584778,
+        -0.1842857152223587,
+        0.0
+      ],
+      "q99": [
+        0.9375,
+        0.8758928775787354,
+        0.9321428537368774,
+        0.1039285734295845,
+        0.17678570747375488,
+        0.14571428298950195,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        -0.024462558329105377,
+        0.106529600918293,
+        1.0580483675003052,
+        3.0628468990325928,
+        -0.10464039444923401,
+        0.08307311683893204,
+        0.01995457336306572,
+        -0.020162804052233696
+      ],
+      "std": [
+        0.1101478561758995,
+        0.13784688711166382,
+        0.1044282391667366,
+        0.10451053828001022,
+        0.4112098217010498,
+        0.2176690548658371,
+        0.017260896041989326,
+        0.0171116404235363
+      ],
+      "max": [
+        0.1759040206670761,
+        0.3904820382595062,
+        1.3290715217590332,
+        3.4566118717193604,
+        1.2268599271774292,
+        1.0429412126541138,
+        0.041053611785173416,
+        0.000775813648942858
+      ],
+      "min": [
+        -0.3095473051071167,
+        -0.29250794649124146,
+        0.9095591306686401,
+        2.497488260269165,
+        -1.8006486892700195,
+        -0.7207611203193665,
+        -0.0004703797458205372,
+        -0.041536275297403336
+      ],
+      "q01": [
+        -0.2727657300233841,
+        -0.23721413239836692,
+        0.9160063165426254,
+        2.77949666261673,
+        -1.3187511622905732,
+        -0.41989982962608335,
+        0.001503719249740243,
+        -0.03989770736545324
+      ],
+      "q99": [
+        0.13529365032911292,
+        0.3629165390133857,
+        1.2862326657772063,
+        3.2829698753356933,
+        0.9332760351896285,
+        0.6325724506378171,
+        0.039933966137468815,
+        -0.001671919699292631
+      ]
+    },
+    "num_transitions": 52970,
+    "num_trajectories": 432
+  }
+}

lora_adapter/README.md

@@ -0,0 +1,202 @@
+---
+base_model: /home/user1/workspace/juyi/lisaopenvla/hf-convert/prismatic-llama3.2-dinosiglip-224px-1b-vlm
+library_name: peft
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]
+### Framework versions
+
+- PEFT 0.11.1

lora_adapter/adapter_config.json

@@ -0,0 +1,45 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": {
+    "base_model_class": "OpenVLAForActionPrediction",
+    "parent_library": "transformers_modules.prismatic-llama3.2-dinosiglip-224px-1b-vlm.modeling_prismatic"
+  },
+  "base_model_name_or_path": "/home/user1/workspace/juyi/lisaopenvla/hf-convert/prismatic-llama3.2-dinosiglip-224px-1b-vlm",
+  "bias": "none",
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": "gaussian",
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 16,
+  "lora_dropout": 0.0,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "r": 32,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "k_proj",
+    "v_proj",
+    "q",
+    "lm_head",
+    "fc1",
+    "qkv",
+    "kv",
+    "fc3",
+    "fc2",
+    "o_proj",
+    "up_proj",
+    "q_proj",
+    "gate_proj",
+    "proj",
+    "down_proj"
+  ],
+  "task_type": null,
+  "use_dora": false,
+  "use_rslora": false
+}

lora_adapter/adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:914179b95993ea021e03d3d3f62cc0528f2adc3ded5086e5df998cb80759c68a
+size 638172008

modeling_prismatic.py

@@ -0,0 +1,1556 @@
+"""
+modeling_prismatic.py
+
+Core HuggingFace-style PrismaticPreTrainedModel and PrismaticForConditionalGeneration class definitions.
+Inherits from the default `transformers.PretrainedModel`. Meant to be standalone and self-contained,
+but exactly replicate the logic in `prismatic.models.vlms.prismatic.py`.
+"""
+
+import logging
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import timm
+import tokenizers
+import torch
+import torch.nn as nn
+import transformers
+from timm.models.vision_transformer import LayerScale
+from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+from prismatic.models.action_heads import  L1RegressionActionHead, DiTActionHead, FlowMatchingActionHead
+from prismatic.training.train_utils import (
+    get_current_action_mask,
+    get_next_actions_mask,
+)
+from prismatic.vla.constants import (
+    ACTION_DIM,
+    ACTION_PROPRIO_NORMALIZATION_TYPE,
+    IGNORE_INDEX,
+    NUM_ACTIONS_CHUNK,
+    ACTION_TOKEN_IDX,
+    NormalizationType,
+)
+
+from .configuration_prismatic import OpenVLAConfig, PrismaticConfig
+
+# Set up logger
+logger = logging.getLogger(__name__)
+
+
+# === Utility Functions for Monkey-Patching ===
+def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]:
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        result = fn(*args, **kwargs)
+        return result[0] if isinstance(result, tuple) else result
+
+    return wrapper
+
+
+# HF Transformers overwrites parameters with names containing `gamma`; we're going to patch VisionBackbone.LayerScale.
+#   =>> TIMM :: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L109
+#   =>> Transformers :: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L3960
+def _ls_new_forward(self, x: torch.Tensor) -> torch.Tensor:
+    return x.mul_(self.scale_factor) if self.inplace else x * self.scale_factor
+
+
+def ls_apply_patch(ls_module: LayerScale):
+    ls_module.scale_factor = nn.Parameter(ls_module.gamma.clone())
+    ls_module.forward = _ls_new_forward.__get__(ls_module, LayerScale)
+    del ls_module.gamma
+
+
+# === Prismatic Vision Backbone (nn.Module) Definitions (w/ Fused Backbone Support) ===
+class PrismaticVisionBackbone(nn.Module):
+    """
+    Vision backbone for Prismatic models that handles image feature extraction.
+
+    Supports both single backbone (e.g., SigLIP) and fused backbone (e.g., SigLIP + DINOv2) configurations.
+    For fused backbones, features from both models are concatenated along the feature dimension.
+    """
+
+    def __init__(
+        self,
+        use_fused_vision_backbone: bool,
+        image_sizes: List[int],
+        timm_model_ids: List[str],
+        timm_override_act_layers: List[Optional[str]],
+    ) -> None:
+        """
+        Initialize the vision backbone.
+
+        Args:
+            use_fused_vision_backbone: Whether to use two backbones and fuse their features
+            image_sizes: List of image sizes for each backbone
+            timm_model_ids: List of TIMM model IDs to use for each backbone
+            timm_override_act_layers: List of activation layer overrides for each backbone
+        """
+        super().__init__()
+        self.use_fused_vision_backbone = use_fused_vision_backbone
+        self.num_images_in_input = 1  # Default value, can be overridden later
+
+        # Validate number of (fused) vision backbones
+        if len(timm_model_ids) > 2:
+            raise ValueError("Prismatic models only support up to 2 (fused) vision backbones!")
+
+        # Create primary featurizer
+        self.featurizer = self._create_featurizer(
+            model_id=timm_model_ids[0], img_size=image_sizes[0], act_layer=timm_override_act_layers[0]
+        )
+        self.embed_dim = self.featurizer.embed_dim
+
+        # Create secondary featurizer if using fused backbone
+        if self.use_fused_vision_backbone:
+            self.fused_featurizer = self._create_featurizer(
+                model_id=timm_model_ids[1], img_size=image_sizes[1], act_layer=timm_override_act_layers[1]
+            )
+            self.embed_dim += self.fused_featurizer.embed_dim
+
+        # Patch LayerScale modules for HF compatibility
+        self._patch_layer_scales()
+
+    def _create_featurizer(self, model_id: str, img_size: int, act_layer: Optional[str]) -> nn.Module:
+        """
+        Create a TIMM-based featurizer model with appropriate configurations.
+
+        Args:
+            model_id: The TIMM model ID to load
+            img_size: Input image size for the model
+            act_layer: Override for the activation layer type
+
+        Returns:
+            A configured featurizer model
+        """
+        featurizer = timm.create_model(
+            model_id,
+            pretrained=False,
+            num_classes=0,
+            img_size=img_size,
+            act_layer=act_layer,
+        )
+
+        # Monkey-patch the forward function to extract the second-to-last layer features
+        num_blocks = len(featurizer.blocks)
+        featurizer.forward = unpack_tuple(partial(featurizer.get_intermediate_layers, n={num_blocks - 2}))
+
+        return featurizer
+
+    def _patch_layer_scales(self) -> None:
+        """
+        Patch all LayerScale modules to be compatible with HF's parameter naming.
+
+        HF Transformers overwrites parameters with names containing 'gamma',
+        so we need to rename and modify the forward method.
+        """
+        # Patch primary featurizer
+        for module in self.featurizer.modules():
+            if isinstance(module, LayerScale):
+                ls_apply_patch(module)
+
+        # Patch secondary featurizer if it exists
+        if self.use_fused_vision_backbone:
+            for module in self.fused_featurizer.modules():
+                if isinstance(module, LayerScale):
+                    ls_apply_patch(module)
+
+    def get_num_patches(self) -> int:
+        """
+        Returns the number of vision patches output by the vision backbone.
+
+        Returns:
+            Number of patches per image
+        """
+        return self.featurizer.patch_embed.num_patches
+
+    def get_num_images_in_input(self) -> int:
+        """
+        Returns the number of input images for the vision backbone.
+
+        Returns:
+            Number of images expected in the input
+        """
+        return self.num_images_in_input
+
+    def set_num_images_in_input(self, num_images_in_input: int) -> None:
+        """
+        Sets the number of input images for the vision backbone.
+
+        Args:
+            num_images_in_input: Number of images to expect in the input
+        """
+        self.num_images_in_input = num_images_in_input
+
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        """
+        Implements the forward pass for the vision backbone.
+
+        If `self.use_fused_vision_backbone == True`, uses both SigLIP and DINOv2 transformers to extract visual features
+        (otherwise uses SigLIP only). Allows multi-image inputs (but only for fused vision backbone).
+
+        Args:
+            pixel_values (torch.Tensor): Pixels for input image(s), (B, C, H, W).
+        """
+        if self.num_images_in_input == 1:
+            if not self.use_fused_vision_backbone:
+                return self.featurizer(pixel_values)
+
+            # Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack
+            img, img_fused = torch.split(pixel_values, [3, 3], dim=1)
+            patches, patches_fused = self.featurizer(img), self.fused_featurizer(img_fused)
+
+            return torch.cat([patches, patches_fused], dim=2)
+
+        else:
+            assert self.use_fused_vision_backbone, "Multi-image inputs require using fused backbone!"
+
+            # Split `pixel_values` into individual images (each with 6 channels: 3 for SigLIP + 3 for DINOv2)
+            images = torch.split(pixel_values, [6] * self.num_images_in_input, dim=1)
+
+            # Process each image and collect patches
+            all_patches = []
+            for img in images:
+                # Split each image further into two stacks of channels (each with 3 channels)
+                img_regular, img_fused = torch.split(img, [3, 3], dim=1)
+
+                # Get patches from both SigLIP and DINOv2 vision transformers
+                patches = self.featurizer(img_regular)
+                patches_fused = self.fused_featurizer(img_fused)
+
+                # Concatenate SigLIP and DINOv2 patches along the hidden dimension
+                combined_patches = torch.cat([patches, patches_fused], dim=2)
+                all_patches.append(combined_patches)
+
+            # Concatenate all patches along the patch dimension
+            return torch.cat(all_patches, dim=1)
+
+
+# === Prismatic Projector (nn.Module) Definitions ===
+class PrismaticProjector(nn.Module):
+    def __init__(self, use_fused_vision_backbone: bool, vision_dim: int, llm_dim: int) -> None:
+        super().__init__()
+        self.use_fused_vision_backbone = use_fused_vision_backbone
+        self.vision_dim, self.llm_dim = vision_dim, llm_dim
+
+        # Switch on `use_fused_vision_backbone` =>> use slightly different MLPs and projection factors!
+        if not self.use_fused_vision_backbone:
+            self.fc1 = nn.Linear(self.vision_dim, self.llm_dim, bias=True)
+            self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
+            self.act_fn1 = nn.GELU()
+        else:
+            initial_projection_dim = 4 * vision_dim
+            self.fc1 = nn.Linear(self.vision_dim, initial_projection_dim, bias=True)
+            self.fc2 = nn.Linear(initial_projection_dim, self.llm_dim, bias=True)
+            self.fc3 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
+            self.act_fn1 = nn.GELU()
+            self.act_fn2 = nn.GELU()
+
+    def forward(self, img_patches: torch.Tensor) -> torch.Tensor:
+        if not self.use_fused_vision_backbone:
+            projected_features = self.fc1(img_patches)
+            projected_features = self.act_fn1(projected_features)
+            projected_features = self.fc2(projected_features)
+        else:
+            projected_features = self.fc1(img_patches)
+            projected_features = self.act_fn1(projected_features)
+            projected_features = self.fc2(projected_features)
+            projected_features = self.act_fn2(projected_features)
+            projected_features = self.fc3(projected_features)
+
+        return projected_features
+
+
+# === Main HF Class Definitions ===
+@dataclass
+class PrismaticCausalLMOutputWithPast(ModelOutput):
+    """Base class for Prismatic casual (visually-conditioned) language model outputs; also exposes visual features."""
+
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+    # Additions for VLMs
+    projector_features: Optional[torch.FloatTensor] = None
+
+
+class PrismaticPreTrainedModel(PreTrainedModel):
+    config_class: PretrainedConfig = PrismaticConfig
+    base_model_prefix: str = "model"
+    supports_gradient_checkpointing: bool = True
+
+    _no_split_modules: ClassVar[List[str]] = ["PrismaticProjector"]
+    _skip_keys_device_placement: str = "past_key_values"
+    _supports_flash_attn_2: bool = True
+
+    def _init_weights(self, module: nn.Module) -> None:
+        # Important :: this HF ported version is *not* meant for training from scratch; only inference and fine-tuning!
+        #   => As such, this init_weights code is not correct; if training VLMs from scratch, use the main codebase at
+        #      https://github.com/TRI-ML/prismatic-vlms
+        std = (
+            self.config.initializer_range
+            if hasattr(self.config, "initializer_range")
+            else self.config.text_config.initializer_range
+        )
+
+        if hasattr(module, "class_embedding"):
+            module.class_embedding.data.normal_(mean=0.0, std=std)
+
+        if isinstance(module, (nn.Linear, nn.Conv2d)):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+    @property
+    def _supports_sdpa(self) -> bool:
+        """Check LLM supports SDPA Attention"""
+        return self.language_model._supports_sdpa
+
+
+class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
+    def __init__(self, config: PrismaticConfig) -> None:
+        super().__init__(config)
+
+        # [Validation] Lightweight Validate on `config` Fields + Dependency Versions
+        if config.use_fused_vision_backbone is None:
+            raise ValueError("Missing config field `use_fused_vision_backbone`")
+
+        if timm.__version__ not in {"0.9.10", "0.9.11", "0.9.12", "0.9.16"}:
+            raise NotImplementedError(
+                "TIMM Version must be >= 0.9.10 and < 1.0.0 (breaking); please raise a GitHub Issue "
+                "if you urgently need support for latest TIMM versions."
+            )
+
+        if (transformers.__version__ != "4.40.1") or (tokenizers.__version__ != "0.19.1"):
+            logger.warning(
+                f"Expected `transformers==4.40.1` and `tokenizers==0.19.1` but got "
+                f"`transformers=={transformers.__version__}` and `tokenizers=={tokenizers.__version__}`; "
+                f"there might be inference-time regressions due to dependency changes. If in doubt, please"
+                f"use the above versions."
+            )
+
+        # Instantiate PrismaticVisionBackbone (w/ Potential Fused Backbone)
+        self.vision_backbone = PrismaticVisionBackbone(
+            config.use_fused_vision_backbone, config.image_sizes, config.timm_model_ids, config.timm_override_act_layers
+        )
+
+        # Create Multimodal Projector
+        self.projector = PrismaticProjector(
+            config.use_fused_vision_backbone,
+            vision_dim=self.vision_backbone.embed_dim,
+            llm_dim=config.text_config.hidden_size,
+        )
+
+        # Instantiate LLM Backbone
+        self.language_model = AutoModelForCausalLM.from_config(
+            config.text_config, attn_implementation=config._attn_implementation
+        )
+        self.vocab_size = config.text_config.vocab_size
+        self.pad_token_id = config.pad_token_id
+        self.llm_dim = config.text_config.hidden_size
+
+        # HF Boilerplate =>> initializes weights via `_init_weights()` and sets gradient checkpointing
+        self.post_init()
+
+    # === `PreTrainedModel` Boilerplate ===
+    def get_input_embeddings(self) -> nn.Module:
+        return self.language_model.get_input_embeddings()
+
+    def set_input_embeddings(self, value: nn.Module) -> None:
+        self.language_model.set_input_embeddings(value)
+
+    def get_output_embeddings(self) -> nn.Module:
+        return self.language_model.get_output_embeddings()
+
+    def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
+        self.language_model.set_output_embeddings(new_embeddings)
+
+    def get_decoder(self) -> nn.Module:
+        return self.language_model.get_decoder()
+
+    def set_decoder(self, decoder: nn.Module) -> None:
+        self.language_model.set_decoder(decoder)
+
+    def tie_weights(self) -> None:
+        self.language_model.tie_weights()  # Note: `Llama-2` and `Mistral` don't tie weights (no-op)
+
+    def resize_token_embeddings(
+        self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None
+    ) -> nn.Embedding:
+        updated_embeddings = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
+
+        # Update config/instance variables
+        self.config.text_config.vocab_size = updated_embeddings.num_embeddings
+        self.vocab_size = updated_embeddings.num_embeddings
+
+        return updated_embeddings
+
+    def _replace_input_embeddings(self, input_embeddings, all_actions_mask, noisy_action_features):
+        """
+        Replace embeddings in input_embeddings at positions where all_actions_mask is True
+        with embeddings from noisy_action_features, using vectorized operations.
+
+        Args:
+            input_embeddings: Tensor of shape (B, S, D)
+            all_actions_mask: Boolean tensor of shape (B, S)
+            noisy_action_features: Tensor of shape (B, K, D) where K is the number of True values in mask per sample
+
+        Returns:
+            Modified input_embeddings tensor
+        """
+        # Clone input to avoid modifying the original tensor
+        new_input_embeddings = input_embeddings.clone()
+
+        # Create a tensor with the same shape of input_embeddings to hold the noisy action features
+        repositioned_noisy_action_features = torch.zeros_like(input_embeddings)
+
+        # Create batch indices for splicing
+        batch_indices = torch.arange(input_embeddings.shape[0], device=input_embeddings.device)
+        batch_indices = batch_indices.unsqueeze(1).expand(-1, noisy_action_features.shape[1])
+
+        # Get indices where mask is True for each sample
+        masked_indices = torch.stack([torch.where(mask)[0] for mask in all_actions_mask])
+
+        # Move the noisy action features into their correct positions
+        repositioned_noisy_action_features[batch_indices, masked_indices] = noisy_action_features
+
+        # Combine original input embeddings and noisy action embeddings using the mask
+        new_input_embeddings = torch.where(
+            all_actions_mask.unsqueeze(-1), repositioned_noisy_action_features, new_input_embeddings
+        )
+
+        return new_input_embeddings
+
+    def _process_action_masks(self, labels):
+        """Helper to get action masks from labels"""
+        current_action_mask = get_current_action_mask(labels) # (B, seq_len)
+        next_actions_mask = get_next_actions_mask(labels) # (B, seq_len)
+        all_actions_mask = current_action_mask | next_actions_mask  # (B, seq_len)
+        return all_actions_mask
+
+    def _process_vision_features(self, pixel_values):
+        """Process vision features with optional FiLM conditioning"""
+        patch_features = self.vision_backbone(pixel_values)  # (bsz, 256 * num_images, D)
+
+        # Project patch embeddings into language embedding space
+        return self.projector(patch_features)
+
+    def _process_proprio_features(self, projected_patch_embeddings, proprio, proprio_projector):
+        """Process proprioceptive features and append to vision features"""
+        if proprio_projector is not None and proprio is not None:
+            # projected_patch_embeddings: (bsz, num_patches * num_images, llm_dim)
+            # proprio: (bsz, proprio_dim) or (propro_dim,)
+            proprio = proprio.reshape(projected_patch_embeddings.shape[0], -1)  # (bsz, proprio_dim)
+            proprio_features = proprio_projector(proprio)  # (bsz, llm_dim)
+            proprio_features = proprio_features.unsqueeze(dim=1)  # (bsz, 1, llm_dim)
+            # For simplicity, just append proprio token to the end of projected vision patch tokens
+            return torch.cat((projected_patch_embeddings, proprio_features), dim=1)
+        return projected_patch_embeddings
+
+    def _build_multimodal_attention(self, input_embeddings, projected_patch_embeddings, attention_mask):
+        """Build multimodal embeddings and attention mask"""
+        # juyi: Update attention mask 是不是要改成下三角? 不用, 因为generate会自动屏蔽
+        projected_patch_attention_mask = None
+        if attention_mask is not None:
+            projected_patch_attention_mask = torch.full(
+                (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
+                fill_value=True,
+                dtype=attention_mask.dtype,
+                device=attention_mask.device,
+            )
+
+        # Build multimodal embeddings & attention mask; insert embeddings after <BOS> token (1:)
+        multimodal_embeddings = torch.cat(
+            [input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1
+        )
+
+        multimodal_attention_mask = None
+        if attention_mask is not None:
+            multimodal_attention_mask = torch.cat(
+                [attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1
+            )
+
+        return multimodal_embeddings, multimodal_attention_mask
+
+    def _build_multimodal_labels(self, labels, projected_patch_embeddings):
+        """Build multimodal labels with IGNORE_INDEX for patch embeddings"""
+        if labels is not None:
+            projected_patch_labels = torch.full(
+                (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
+                fill_value=IGNORE_INDEX, # 这些位置不需要计算损失。
+                dtype=labels.dtype,
+                device=labels.device,
+            )
+            return torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1) # 第一个token是<BOS>
+        return None
+
+    # === Core Prismatic VLM `forward()` Logic ===
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_projector_features: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        proprio=None,
+        proprio_projector=None,
+        noisy_actions=None,
+        noisy_action_projector=None,
+        diffusion_timestep_embeddings=None,
+        use_film: bool = False,
+    ) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
+        """Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
+        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
+        )
+        output_projector_features = output_projector_features if output_projector_features is not None else False
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
+        use_cache = use_cache and not self.training
+
+        # Instantiate Placeholder for Projector Features
+        projected_patch_embeddings = None
+
+        # === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
+        if input_ids.shape[1] == 1:
+            assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
+            assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
+            assert labels is None, "Unexpected key `labels` provided during cached generation!"
+
+            language_model_output = self.language_model(
+                input_ids=input_ids,
+                attention_mask=None,
+                position_ids=None,
+                past_key_values=past_key_values,
+                inputs_embeds=None,
+                labels=None,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+        # === Handle Unimodal Forward ===
+        elif pixel_values is None:
+            assert (input_ids is not None) and (inputs_embeds is None), "Missing `input_ids` in language-only forward!"
+            assert past_key_values is None, "Unexpected key `past_key_values` provided during language-only forward!"
+
+            language_model_output = self.language_model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                position_ids=None,
+                past_key_values=None,
+                inputs_embeds=None,
+                labels=labels,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+        # === Handle Multimodal Forward ===
+        elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
+            assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
+
+            # Get input embeddings (from language model embeddings)
+            input_embeddings = self.get_input_embeddings()(input_ids)  # (B, seq_len, D)
+
+            # Extract action masks
+            all_actions_mask = self._process_action_masks(labels)
+
+            # Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
+            language_embeddings = input_embeddings[~all_actions_mask].reshape(
+                input_embeddings.shape[0], -1, input_embeddings.shape[2]
+            )  # (B, lang_seq_len, llm_dim)
+
+            # Get visual features
+            projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
+            # bug: TypeError: PrismaticForConditionalGeneration._process_vision_features() takes 2 positional arguments but 4 were given
+
+            # Add proprioceptive state if provided
+            projected_patch_embeddings = self._process_proprio_features(
+                projected_patch_embeddings, proprio, proprio_projector
+            )
+
+            # [Diffusion] Add diffusion timestep embedding if provided
+            if diffusion_timestep_embeddings is not None:
+                # For simplicity, just append diffusion timestep embedding to the end of projected vision patch tokens
+                projected_patch_embeddings = torch.cat(
+                    (projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
+                )
+                
+            # Process action embeddings
+            if noisy_actions is not None:
+                # Get mask corresponding to all action tokens
+                all_actions_mask = self._process_action_masks(labels)
+
+                # Reshape noisy actions into individual action tokens
+                # noisy_actions: (B, chunk_len, action_dim) -> (B, chunk_len * action_dim, 1)
+                B = noisy_actions.shape[0]
+                noisy_actions = noisy_actions.reshape(B, -1).unsqueeze(-1)
+
+                # Project noisy action tokens into language model embedding space
+                noisy_action_features = noisy_action_projector(noisy_actions)  # (B, chunk_len * action_dim, llm_dim)
+
+                # Replace embeddings of the action tokens with noisy action embeddings
+                input_embeddings = self._replace_input_embeddings(
+                    input_embeddings, all_actions_mask, noisy_action_features
+                )
+            else:
+                # Replace the embeddings of the action tokens with zeros
+                # (Later on, the positional embeddings will be added to them)
+                all_actions_mask = all_actions_mask.unsqueeze(-1)  # (B, seq_len, 1)
+                input_embeddings = input_embeddings * ~all_actions_mask
+
+            # Build multimodal embeddings & attention mask
+            multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+                input_embeddings, projected_patch_embeddings, attention_mask
+            )
+
+            # Build labels for multimodal sequence if needed
+            multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
+
+            # Dispatch to language model
+            language_model_output = self.language_model(
+                input_ids=None,
+                attention_mask=multimodal_attention_mask,
+                position_ids=None,
+                past_key_values=None,
+                inputs_embeds=multimodal_embeddings,
+                labels=multimodal_labels,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+        # === Otherwise =>> Assume Invalid! ===
+        elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
+            raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
+
+        else:
+            raise ValueError(
+                "Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
+                f"=> `input_ids` = {input_ids is not None}\n"
+                f"=> `attention_mask` = {attention_mask is not None}\n"
+                f"=> `pixel_values` = {pixel_values is not None}\n"
+                f"=> `labels` = {labels is not None}\n"
+                f"=> `input_embeds` = {inputs_embeds is not None}\n"
+                f"=> `past_key_values` = {past_key_values is not None}\n"
+                f"=> `use_cache` = {use_cache}"
+            )
+
+        # Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
+        if not return_dict:
+            if output_projector_features and (projected_patch_embeddings is not None):
+                return *language_model_output, projected_patch_embeddings
+
+            return language_model_output
+
+        return PrismaticCausalLMOutputWithPast(
+            loss=language_model_output.loss,
+            logits=language_model_output.logits,
+            past_key_values=language_model_output.past_key_values,
+            hidden_states=language_model_output.hidden_states,
+            attentions=language_model_output.attentions,
+            projector_features=projected_patch_embeddings,
+        )
+
+    # === GenerationMixin Methods ===
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs: str,
+    ) -> Dict[str, torch.Tensor]:
+        """Borrowed from `LlamaForCausalLM` and simplified for batch size = 1; mirrors original PrismaticVLM logic."""
+        if ((input_ids is not None) and (input_ids.shape[0] > 1)) or (
+            (inputs_embeds is not None) and (inputs_embeds.shape[0] > 1)
+        ):
+            raise ValueError("Generation with batch size > 1 is not currently supported!")
+
+        # Handle `past_key_values` (cache) =>> assume `input_ids` just has unprocessed tokens
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1:]
+
+        # If `input_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 = {"input_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        # Make sure `pixel_values` are preserved in `model_inputs`
+        model_inputs.update(
+            {
+                "attention_mask": attention_mask,
+                "pixel_values": pixel_values,
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+            }
+        )
+
+        return model_inputs
+
+    # Defer to Language Model (all handle this differently, with different return types)
+    def _reorder_cache(self, *args, **kwargs) -> Any:
+        return self.language_model._reorder_cache(*args, **kwargs)
+
+
+class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
+    config_class: PretrainedConfig = OpenVLAConfig
+
+    def __init__(self, config: OpenVLAConfig) -> None:
+        super().__init__(config)
+        self.norm_stats = config.norm_stats
+
+        # Compute action bins
+        self.bins = np.linspace(-1, 1, config.n_action_bins)
+        self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0
+
+        # Compute vocab size for de-tokenization -- revert added "multiple of"
+        self.vocab_size = self.config.text_config.vocab_size - self.config.pad_to_multiple_of
+
+    def _prepare_input_for_action_prediction(self, input_ids, attention_mask):
+        # eval 会用到这里
+        """Prepares input for action prediction by adding necessary tokens"""
+        # Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens
+        placeholder_action_token_ids = (
+            torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
+        ) 
+        input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1) # torch.Size([1,  35 + 56= 91])
+
+        # Extend the attention mask to fit the new shape of input
+        # Note: Only batch size == 1 supported right now
+        mask_extension = (
+            torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
+            .to(attention_mask.device)
+            .to(attention_mask.dtype)
+        )
+        attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
+
+        return input_ids, attention_mask
+
+    def _prepare_labels_for_action_prediction(self, labels, input_ids):
+        """Creates labels tensor for action prediction if not provided"""
+        # eval 会用到这里 , 
+        # Extends label tensors with fake action labels
+        # Adds stop tokens at the end of sequences
+        # Handles label preparation for action prediction tasks
+        # 他为啥可以随便一个?  xuan说 你自定义一个值 ,然后一直指定这个 , PAD token可以吗? 
+        #TODO: 这里是否要改?  感觉不需要改. 随便写就行了因为labels不重要只是要一个mask. 为什么需要这个函数? 确保 action 预测任务的标签（labels）符合模型的输入长度，并正确地处理序列终止
+        # Extend labels tensor with fake action labels
+        ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_IDX # = 为了mask正确生成, action_tokens_only_mask = (labels == ACTION_TOKEN_IDX ), 所以这里也填上ACTION_TOKEN_IDX
+        labels_extension = (
+            torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
+            * ARBITRARY_ACTION_TOKEN_IDX
+        ) #torch.Size([1, 57]),全是 ARBITRARY_ACTION_TOKEN_IDX
+        labels = torch.cat([labels, labels_extension], dim=-1)
+
+        return labels
+
+    def _unnormalize_actions(self, normalized_actions, unnorm_key=None):
+        """Unnormalize actions using dataset statistics"""
+        action_norm_stats = self.get_action_stats(unnorm_key)
+
+        if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:
+            mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))
+            action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])
+        elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:
+            mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
+            action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
+        else:
+            raise ValueError("Unsupported action/proprio normalization type detected!")
+
+        actions = np.where(
+            mask,
+            0.5 * (normalized_actions + 1) * (action_high - action_low + 1e-8) + action_low,
+            normalized_actions,
+        )
+
+        return actions
+
+
+    def _normalize_actions(self, actions, norm_key=None):
+        """Normalize actions to [-1, 1] using dataset statistics"""
+        action_norm_stats = self.get_action_stats(norm_key)
+
+        if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:
+            mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))
+            action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])
+        elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:
+            mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
+            action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
+        else:
+            raise ValueError("Unsupported action/proprio normalization type detected!")
+
+        normalized = np.where(
+            mask,
+            2 * (actions - action_low) / (action_high - action_low + 1e-8) - 1,
+            actions,
+        )
+
+        return normalized
+
+    def _run_diffusion_prediction(
+        self,
+        input_embeddings,
+        all_actions_mask,
+        noise,
+        action_head,
+        projected_patch_embeddings,
+        labels,
+        attention_mask,
+        NUM_PATCHES,
+        NUM_PROMPT_TOKENS,
+        noisy_action_projector,
+    ):
+        """Run diffusion-based action prediction"""
+        # Set diffusion timestep values
+        action_head.noise_scheduler.set_timesteps(action_head.num_diffusion_steps)
+        # Clone embedding for reuse in each timestep
+        orig_projected_patch_embeddings = projected_patch_embeddings.clone()
+        curr_noisy_actions = noise
+
+        # Reverse diffusion: Iteratively denoise to generate action prediction
+        for t in action_head.noise_scheduler.timesteps:
+            # Get diffusion model's noise prediction (conditioned on VLA latent embedding, current noisy action
+            # embedding, and diffusion timestep embedding)
+            timesteps = torch.Tensor([t]).to(labels.device)
+            diffusion_timestep_embeddings = (
+                action_head.time_encoder(timesteps).to(curr_noisy_actions.dtype).to(curr_noisy_actions.device)
+            )  # (B, llm_dim)
+            diffusion_timestep_embeddings = diffusion_timestep_embeddings.unsqueeze(1)  # (B, 1, llm_dim)
+
+            # [Diffusion] Replace the embeddings of the action tokens with noisy actions
+            # (Later on, the positional embeddings will be added to them)
+
+            # For simplicity, append diffusion timestep embedding to the end of projected vision tokens
+            projected_patch_embeddings = torch.cat(
+                (orig_projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
+            )
+
+            # Reshape and project noisy actions into language embedding space
+            B = curr_noisy_actions.shape[0]
+            orig_curr_noisy_actions_shape = curr_noisy_actions.shape
+            curr_noisy_actions = curr_noisy_actions.reshape(B, -1).unsqueeze(-1)
+            noisy_action_features = noisy_action_projector(curr_noisy_actions)
+            curr_noisy_actions = curr_noisy_actions.reshape(orig_curr_noisy_actions_shape)
+
+            # Replace action token embeddings with noisy action embeddings
+            input_embeddings = self._replace_input_embeddings(
+                input_embeddings.clone(), all_actions_mask, noisy_action_features
+            )
+
+            # Build multimodal embeddings and attention mask
+            multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+                input_embeddings, projected_patch_embeddings, attention_mask
+            )
+
+            # Forward pass through language model
+            language_model_output = self.language_model(
+                input_ids=None,
+                attention_mask=multimodal_attention_mask,
+                position_ids=None,
+                past_key_values=None,
+                inputs_embeds=multimodal_embeddings,
+                labels=None,
+                use_cache=None,
+                output_attentions=False,
+                output_hidden_states=True,
+                return_dict=True,
+            )
+
+            # Extract hidden states for action portion of response
+            last_hidden_states = language_model_output.hidden_states[-1]  # (B, seq_len, D)
+            actions_hidden_states = last_hidden_states[
+                :,
+                NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
+                :,
+            ]  # (B, act_chunk_len, D)
+
+            # Predict noise and update noisy actions: x_t -> x_{t-1}
+            noise_pred = action_head.predict_noise(actions_hidden_states)
+            curr_noisy_actions = action_head.noise_scheduler.step(noise_pred, t, curr_noisy_actions).prev_sample
+
+        curr_noisy_actions = curr_noisy_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
+
+        # Return final actions
+        return curr_noisy_actions.float().cpu().detach().numpy(), actions_hidden_states
+
+    def _regression_or_discrete_prediction(
+        self,
+        input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
+        all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action . 
+        projected_patch_embeddings: torch.FloatTensor, 
+        attention_mask: torch.BoolTensor,
+        labels: torch.LongTensor,
+        NUM_PATCHES: int,
+        NUM_PROMPT_TOKENS: int,
+        action_head: L1RegressionActionHead,
+        **kwargs,
+    ):
+        """Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
+        # Extract hidden states for action tokens
+        # last_hidden_states = language_model_output.hidden_states[-1]  # (B, seq_len, D)
+
+        # actions_hidden_states = last_hidden_states[:, NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + NUM_ACTIONS_CHUNK * tokennum, :]# (B, act_chunk_len, D)
+        # 都不需要取了, 直接就给 token对应的hidden state了 ,太方便了. 
+        # 为什么第一个是torch.Size([1, 535, 4096])? 我应该选哪个?  https://discuss.huggingface.co/t/get-each-generated-token-last-layer-hidden-state/145921
+        # language_model_output.sequences tensor([[29871, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 32001,     2]], device='cuda:0')
+        cfg = kwargs.pop("cfg", None)
+        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+            input_embeddings, projected_patch_embeddings, attention_mask
+        ) 
+        # multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'  
+        language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
+        #  is tuple (1 token,  33 layers, torch.Size([1, 314, 4096]))
+        hidden_states = language_model_output.hidden_states[0][-1]
+        actions_hidden_states = hidden_states[:, -NUM_ACTIONS_CHUNK:]
+
+        normalized_actions = action_head.predict_action(actions_hidden_states)
+        normalized_actions = normalized_actions.reshape(cfg.num_actions_chunk, ACTION_DIM)
+        normalized_actions = normalized_actions.float().cpu().detach().numpy()
+        if cfg.mode == "mul": 
+            if cfg.num_actions_chunk//cfg.num_actions_per_token > 1:
+                token_num = cfg.num_actions_chunk//cfg.num_actions_per_token
+                language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,min_new_tokens=token_num, max_new_tokens=token_num,output_hidden_states=True,return_dict_in_generate=True)
+                actions_hidden_states0 = language_model_output.hidden_states[0][-1][:, -1]  # 第一个token的hidden state的last layer
+                
+                actions_hidden_states_list = [actions_hidden_states0]
+                for i in range(1, token_num):
+                    token_hidden_state = language_model_output.hidden_states[i][-1][0]  # i+1个token的hidden state的last layer
+                    actions_hidden_states_list.append(token_hidden_state)
+                # 将所有hidden states拼接起来
+                combined_hidden_states = torch.stack(actions_hidden_states_list, dim=0)  # (16, 4096)
+                actions_hidden_states = combined_hidden_states.reshape(multimodal_embeddings.shape[0], token_num, multimodal_embeddings.shape[-1])
+            elif cfg.num_actions_chunk//cfg.num_actions_per_token == 1:
+                language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
+                # assert language_model_output.sequences == torch.tensor([[32001]], device=multimodal_embeddings.device)
+                actions_hidden_states = language_model_output.hidden_states[0][-1] 
+                actions_hidden_states = actions_hidden_states[:, -1]
+            else:
+                raise NotImplementedError
+        else:
+            raise NotImplementedError    
+        return normalized_actions, actions_hidden_states
+
+    def hist_predict_action(
+        self,
+        input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
+        all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action . 
+        projected_patch_embeddings: torch.FloatTensor, 
+        attention_mask: torch.BoolTensor,
+        action_head: L1RegressionActionHead,
+        **kwargs,
+    ):
+        cfg = kwargs.get("cfg", None) 
+        action_history = kwargs.get("action_history", None)
+        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+            input_embeddings, projected_patch_embeddings, attention_mask
+        ) 
+        # multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'  
+        # first language_model_output.hidden_states  , is tuple (1 token,  33 layers, torch.Size([1, 314, 4096]))
+        # the following is  (num of tokens,)
+        if cfg.mode == "mul": 
+            if cfg.num_actions_chunk//cfg.num_actions_per_token > 1:
+                raise NotImplementedError
+                # token_num = cfg.num_actions_chunk//cfg.num_actions_per_token
+                # language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,min_new_tokens=token_num, max_new_tokens=token_num,output_hidden_states=True,return_dict_in_generate=True)
+                # actions_hidden_states0 = language_model_output.hidden_states[0][-1][:, -1]  # 第一个token的hidden state的last layer
+                # actions_hidden_states_list = [actions_hidden_states0]
+                # for i in range(1, token_num):
+                #     token_hidden_state = language_model_output.hidden_states[i][-1][0]  # i+1个token的hidden state的last layer
+                #     actions_hidden_states_list.append(token_hidden_state)
+                # combined_hidden_states = torch.stack(actions_hidden_states_list, dim=0)  # (16, 4096)
+                # actions_hidden_states = combined_hidden_states.reshape(multimodal_embeddings.shape[0], token_num, multimodal_embeddings.shape[-1])
+            elif cfg.num_actions_chunk//cfg.num_actions_per_token == 1:
+                language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
+                # assert language_model_output.sequences == torch.tensor([[32001]], device=multimodal_embeddings.device)
+                actions_hidden_states = language_model_output.hidden_states[0][-1] 
+                actions_hidden_states = actions_hidden_states[:, -1]
+                # 在中间加一个 1 维度
+                actions_hidden_states = actions_hidden_states.unsqueeze(1) # for match 3 dim 
+            else:
+                raise NotImplementedError
+        else:
+            raise NotImplementedError    
+
+        normalized_actions = action_head.predict_action(actions_hidden_states, action_history)
+        normalized_actions = normalized_actions.reshape(cfg.num_actions_chunk, ACTION_DIM)
+        normalized_actions = normalized_actions.float().cpu().detach().numpy()
+
+        return normalized_actions, actions_hidden_states
+
+    def mul_regression_or_discrete_prediction(
+        self,
+        input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
+        all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action . 
+        projected_patch_embeddings: torch.FloatTensor, 
+        attention_mask: torch.BoolTensor,
+        labels: torch.LongTensor,
+        NUM_PATCHES: int,
+        NUM_PROMPT_TOKENS: int,
+        action_head: L1RegressionActionHead,
+        **kwargs,
+    ):
+        cfg = kwargs.get("cfg", None) 
+        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+            input_embeddings, projected_patch_embeddings, attention_mask
+        ) 
+        # multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'  
+        # first language_model_output.hidden_states  , is tuple (1 token,  33 layers, torch.Size([1, 314, 4096]))
+        # the following is  (num of tokens,)
+        if cfg.mode == "mul": 
+            if cfg.num_actions_chunk//cfg.num_actions_per_token > 1:
+                token_num = cfg.num_actions_chunk//cfg.num_actions_per_token
+                language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,min_new_tokens=token_num, max_new_tokens=token_num,output_hidden_states=True,return_dict_in_generate=True)
+                actions_hidden_states0 = language_model_output.hidden_states[0][-1][:, -1]  # 第一个token的hidden state的last layer
+                
+                actions_hidden_states_list = [actions_hidden_states0]
+                for i in range(1, token_num):
+                    token_hidden_state = language_model_output.hidden_states[i][-1][0]  # i+1个token的hidden state的last layer
+                    actions_hidden_states_list.append(token_hidden_state)
+                combined_hidden_states = torch.stack(actions_hidden_states_list, dim=0)  # (16, 4096)
+                actions_hidden_states = combined_hidden_states.reshape(multimodal_embeddings.shape[0], token_num, multimodal_embeddings.shape[-1])
+            elif cfg.num_actions_chunk//cfg.num_actions_per_token == 1:
+                language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
+                actions_hidden_states = language_model_output.hidden_states[0][-1] 
+                actions_hidden_states = actions_hidden_states[:, -1]
+            else:
+                raise NotImplementedError
+        else:
+            raise NotImplementedError    
+
+        normalized_actions = action_head.predict_action(actions_hidden_states)
+        normalized_actions = normalized_actions.reshape(cfg.num_actions_chunk, ACTION_DIM)
+        # print(f"*** normalized_actions[]: {normalized_actions} ***")
+        if cfg.action_head_name == "medusa":
+            normalized_actions[:, 6] = torch.sigmoid(normalized_actions[:,  6]) # without bs dim.
+        # print(f"*** normalized_actions[]: {normalized_actions} ***")
+        normalized_actions = normalized_actions.float().cpu().detach().numpy()
+
+        return normalized_actions, actions_hidden_states
+
+    def predict_action(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        unnorm_key: Optional[str] = None,
+        proprio=None,
+        proprio_projector=None,
+        action_head=None,
+        noisy_action_projector=None,
+        use_film: bool = False,
+        **kwargs: str,
+    ) -> np.ndarray:
+        """Predict actions from input sequence, with options for different prediction methods.
+
+        Args:
+            input_ids: Input token ids
+            unnorm_key: Key for unnormalization statistics
+            proprio: Proprioceptive features
+            proprio_projector: Projector for proprioceptive features
+            action_head: Optional head for L1 regression or diffusion-based prediction
+            noisy_action_projector: Projector for noisy actions in diffusion-based prediction
+            use_film: Whether to use FiLM conditioning
+            **kwargs: Additional arguments including pixel_values and attention_mask
+
+        Returns:
+            Tuple of (unnormalized_actions, action_hidden_states)
+        """
+        # If the special empty token ('') does not already appear after the colon (':') token in the prompt
+        # (after "OUT:" or "ASSISTANT:"), insert it to match the inputs seen at training time
+        if not torch.all(input_ids[:, -1] == 29871):
+            input_ids = torch.cat(
+                (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
+            )
+
+        pixel_values = kwargs["pixel_values"]
+        attention_mask = kwargs["attention_mask"]
+
+        # Create fake labels tensor (needed for action mask)
+        labels = input_ids.clone()
+        labels[:] = IGNORE_INDEX # 输入都ignore IGNORE_INDEX = -100
+
+        # Get number of tokens in prompt (excluding the start token)
+        NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1  # Subtract action tokens and stop token
+
+        # Prepare inputs by adding necessary tokens
+        input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)
+
+        # Update labels tensor for action mask computation later
+        labels = self._prepare_labels_for_action_prediction(labels, input_ids)
+
+        # Get input embeddings and action masks
+        input_embeddings = self.get_input_embeddings()(input_ids)
+        all_actions_mask = self._process_action_masks(labels)
+
+        # Extract language embeddings
+        language_embeddings = input_embeddings[~all_actions_mask].reshape(
+            input_embeddings.shape[0], -1, input_embeddings.shape[2]
+        )
+
+        # Process vision features
+        projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
+
+        # Add proprioceptive features if provided
+        use_proprio = proprio_projector is not None and proprio is not None
+        if use_proprio:
+            proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
+            projected_patch_embeddings = self._process_proprio_features(
+                projected_patch_embeddings, proprio, proprio_projector
+            )
+
+        # Use diffusion if provided, otherwise use regression or discrete prediction
+        use_diffusion = noisy_action_projector is not None and hasattr(action_head, "noise_scheduler")
+
+        # Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
+        NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
+        if use_proprio:
+            NUM_PATCHES += 1
+
+        normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
+            input_embeddings,
+            all_actions_mask,
+            projected_patch_embeddings,
+            attention_mask,
+            labels,
+            NUM_PATCHES,
+            NUM_PROMPT_TOKENS,
+            action_head,
+        )
+
+        # Unnormalize predicted actions
+        actions = self._unnormalize_actions(normalized_actions, unnorm_key)
+
+        return actions, actions_hidden_states
+
+    @staticmethod
+    def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_key: Optional[str]) -> str:
+        """Validate and resolve the unnormalization key for action statistics"""
+        if unnorm_key is None:
+            assert len(norm_stats) == 1, (
+                f"Your model was trained on more than one dataset, "
+                f"please pass a `unnorm_key` from the following options to choose the statistics "
+                f"used for un-normalizing actions: {norm_stats.keys()}"
+            )
+            unnorm_key = next(iter(norm_stats.keys()))
+        # norm states没有加载libero, 为什么? 
+        assert unnorm_key in norm_stats, (
+            f"The `unnorm_key` you chose is not in the set of available dataset statistics, "
+            f"please choose from: {norm_stats.keys()}"
+        )
+        return unnorm_key
+
+    def get_action_dim(self, unnorm_key: Optional[str] = None) -> int:
+        """Get the dimensionality of the policy's action space."""
+        unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
+        return len(self.norm_stats[unnorm_key]["action"]["min"])
+
+    def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[str, Any]:
+        """Get all the logged statistics for the given dataset."""
+        unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
+        return self.norm_stats[unnorm_key]["action"]
+
+
+    def lisa_forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_projector_features: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        proprio=None,
+        proprio_projector=None,
+        **kwargs
+    ) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
+        """Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
+        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
+        )
+        output_projector_features = output_projector_features if output_projector_features is not None else False
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
+        use_cache = use_cache and not self.training
+
+        # Instantiate Placeholder for Projector Features
+        projected_patch_embeddings = None
+
+        # === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
+        if input_ids.shape[1] == 1:
+            assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
+            assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
+            assert labels is None, "Unexpected key `labels` provided during cached generation!"
+
+            language_model_output = self.language_model(
+                input_ids=input_ids,
+                attention_mask=None,
+                position_ids=None,
+                past_key_values=past_key_values,
+                inputs_embeds=None,
+                labels=None,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+        # === Handle Unimodal Forward ===
+        elif pixel_values is None:
+            raise NotImplementedError
+
+        # === Handle Multimodal Forward ===
+        elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
+            assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
+
+            # Get input embeddings (from language model embeddings)
+            input_embeddings = self.get_input_embeddings()(input_ids)  # (B, seq_len, D)
+            # Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
+            # language_embeddings = input_embeddings[~all_actions_mask].reshape(
+            #     input_embeddings.shape[0], -1, input_embeddings.shape[2]
+            # )  # (B, lang_seq_len, llm_dim) 这里就会把结尾的 stop index和padding 也算进去. 没问题吗? 没问题因为ignore了 我直接删了因为不用film
+            # Get visual features
+            projected_patch_embeddings = self._process_vision_features(pixel_values)
+
+            # Add proprioceptive state if provided
+            projected_patch_embeddings = self._process_proprio_features(
+                projected_patch_embeddings, proprio, proprio_projector
+            )
+
+            all_actions_mask = (labels == ACTION_TOKEN_IDX) #和run forward pass不一样, run forward pass要手动算token number来找偏移. 
+            all_actions_mask = all_actions_mask.unsqueeze(-1)  # (B, seq_len, 1)
+            input_embeddings = input_embeddings * ~all_actions_mask
+
+            # Build multimodal embeddings & attention mask
+            multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+                input_embeddings, projected_patch_embeddings, attention_mask
+            )
+
+            # Build labels for multimodal sequence if needed
+            multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
+
+            # Dispatch to language model
+            language_model_output = self.language_model(
+                input_ids=None,
+                attention_mask=multimodal_attention_mask,
+                position_ids=None,
+                past_key_values=None,
+                inputs_embeds=multimodal_embeddings,
+                labels=multimodal_labels,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+
+        # === Otherwise =>> Assume Invalid! ===
+        elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
+            raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
+
+        else:
+            raise ValueError(
+                "Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
+                f"=> `input_ids` = {input_ids is not None}\n"
+                f"=> `attention_mask` = {attention_mask is not None}\n"
+                f"=> `pixel_values` = {pixel_values is not None}\n"
+                f"=> `labels` = {labels is not None}\n"
+                f"=> `input_embeds` = {inputs_embeds is not None}\n"
+                f"=> `past_key_values` = {past_key_values is not None}\n"
+                f"=> `use_cache` = {use_cache}"
+            )
+
+        # Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
+        if not return_dict:
+            if output_projector_features and (projected_patch_embeddings is not None):
+                return *language_model_output, projected_patch_embeddings
+
+            return language_model_output
+
+        return PrismaticCausalLMOutputWithPast(
+            loss=language_model_output.loss,
+            logits=language_model_output.logits,
+            past_key_values=language_model_output.past_key_values,
+            hidden_states=language_model_output.hidden_states,
+            attentions=language_model_output.attentions,
+            projector_features=projected_patch_embeddings,
+        )
+
+
+
+    def mul_predict_action(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,  #就是 language instruction.
+        unnorm_key: Optional[str] = None,
+        proprio=None,
+        proprio_projector=None,
+        action_head:L1RegressionActionHead=None,
+        noisy_action_projector=None,
+        use_film: bool = False,
+        **kwargs: str,
+    ) -> np.ndarray:
+    # only use in evaluation.
+        cfg = kwargs.get("cfg", None)  # Extract cfg from kwargs
+        action_history = kwargs.get("action_history", None)
+
+        emptytoken = 220 # for llama3.2
+        # emptytoken = 29871 # for openvla oft
+
+        if not torch.all(input_ids[:, -1] == emptytoken):
+            input_ids = torch.cat(
+                (input_ids, torch.unsqueeze(torch.Tensor([emptytoken]).long(), dim=0).to(input_ids.device)), dim=1
+            )
+
+        
+        pixel_values = kwargs["pixel_values"]
+        attention_mask = kwargs["attention_mask"]
+
+        #  input id  '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
+        # import ipdb; ipdb.set_trace()
+        
+        input_embeddings = self.get_input_embeddings()(input_ids) 
+        
+        projected_patch_embeddings = self._process_vision_features(pixel_values)
+
+        use_proprio = proprio_projector is not None and proprio is not None
+        if use_proprio:
+            proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
+            projected_patch_embeddings = self._process_proprio_features(
+                projected_patch_embeddings, proprio, proprio_projector
+            )
+        if cfg.action_head_name == "hist":
+            normalized_actions, actions_hidden_states = self.hist_predict_action(
+                input_embeddings,
+                None, 
+                projected_patch_embeddings,
+                attention_mask,
+                action_head,
+                cfg=cfg,
+                action_history=action_history,
+            )
+        else:
+            normalized_actions, actions_hidden_states = self.mul_regression_or_discrete_prediction(
+                input_embeddings,
+                None, 
+                projected_patch_embeddings,
+                attention_mask,
+                None, #推理不需要labels
+                None, #推理不需要NUM_PATCHES
+                None, #推理不需要NUM_PROMPT_TOKENS
+                action_head,
+                cfg=cfg,
+            )
+        
+
+        actions = self._unnormalize_actions(normalized_actions, unnorm_key) #在这里 unorm, 所以出来的已经是unorm的了.  所以我 history 也要记录 norm 的. 
+
+        return actions, normalized_actions
+
+
+    def flow_matching_predict_action(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        unnorm_key: Optional[str] = None,
+        proprio=None,
+        proprio_projector=None,
+        action_head: FlowMatchingActionHead = None,
+        noisy_action_projector=None,
+        use_film: bool = False,
+        **kwargs: str,
+    ) -> np.ndarray:
+        """Predict actions using Flow Matching"""
+        cfg = kwargs.get("cfg", None)
+
+        if not torch.all(input_ids[:, -1] == 29871):
+            input_ids = torch.cat(
+                (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
+            )
+
+        pixel_values = kwargs["pixel_values"]
+        attention_mask = kwargs["attention_mask"]
+
+        input_embeddings = self.get_input_embeddings()(input_ids)
+        projected_patch_embeddings = self._process_vision_features(pixel_values)
+
+        use_proprio = proprio_projector is not None and proprio is not None
+        if use_proprio:
+            proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
+            projected_patch_embeddings = self._process_proprio_features(
+                projected_patch_embeddings, proprio, proprio_projector
+            )
+
+        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+            input_embeddings, projected_patch_embeddings, attention_mask
+        )
+
+        if cfg.mode == "flow_matching":
+            if cfg.num_actions_chunk//cfg.num_actions_per_token == 1:
+                language_model_output = self.language_model.generate(
+                    inputs_embeds=multimodal_embeddings,
+                    max_new_tokens=1,
+                    output_hidden_states=True,
+                    return_dict_in_generate=True
+                )
+                
+                actions_hidden_states = language_model_output.hidden_states[0][-1]
+                cognition_features = actions_hidden_states[:, -1]
+                assert (cognition_features.shape[0], cognition_features.shape[1]) == (1, 4096), "Batch size must be 1 for action prediction"
+
+                model_dtype = next(action_head.net.parameters()).dtype
+                cognition_features = cognition_features.unsqueeze(1).to(model_dtype)
+
+                # Sample actions using flow matching
+                normalized_actions = action_head.sample_actions(
+                    cognition_features,
+                    num_steps=getattr(cfg, 'num_flow_steps', 20)
+                )
+                normalized_actions = normalized_actions[0].cpu().numpy()
+            else:
+                raise NotImplementedError("Multi-token flow matching not yet implemented")
+        else:
+            raise NotImplementedError
+
+        actions = self._unnormalize_actions(normalized_actions, unnorm_key)
+        return actions, actions_hidden_states
+
+    def diffusion_predict_action(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,  #就是 language instruction.
+        unnorm_key: Optional[str] = None,
+        proprio=None,
+        proprio_projector=None,
+        action_head:DiTActionHead=None,
+        noisy_action_projector=None,
+        use_film: bool = False,
+        **kwargs: str,
+    ) -> np.ndarray:
+        cfg = kwargs.get("cfg", None)  # Extract cfg from kwargs
+
+        if not torch.all(input_ids[:, -1] == 29871):
+            input_ids = torch.cat(
+                (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
+            )
+
+        
+        pixel_values = kwargs["pixel_values"]
+        attention_mask = kwargs["attention_mask"]
+
+        #  input id  '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:' 
+        
+        input_embeddings = self.get_input_embeddings()(input_ids) 
+        
+        projected_patch_embeddings = self._process_vision_features(pixel_values)
+
+        use_proprio = proprio_projector is not None and proprio is not None
+        if use_proprio:
+            proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
+            projected_patch_embeddings = self._process_proprio_features(
+                projected_patch_embeddings, proprio, proprio_projector
+            )
+        
+        # normalized_actions, actions_hidden_states = self.mul_regression_or_discrete_prediction(
+        #     input_embeddings,
+        #     None, 
+        #     projected_patch_embeddings,
+        #     attention_mask,
+        #     None, #推理不需要labels
+        #     None, #推理不需要NUM_PATCHES
+        #     None, #推理不需要NUM_PROMPT_TOKENS
+        #     action_head,
+        #     cfg=cfg,
+        # )
+
+        # cfg = kwargs.get("cfg", None) 
+        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+            input_embeddings, projected_patch_embeddings, attention_mask
+        ) 
+        # multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'  
+        # first language_model_output.hidden_states  , is tuple (1 token,  33 layers, torch.Size([1, 314, 4096]))
+        # the following is  (num of tokens,)
+
+        if cfg.mode == "dit": 
+            if cfg.num_actions_chunk//cfg.num_actions_per_token > 1:
+                # token_num = cfg.num_actions_chunk//cfg.num_actions_per_token
+                # language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,min_new_tokens=token_num, max_new_tokens=token_num,output_hidden_states=True,return_dict_in_generate=True)
+                # actions_hidden_states0 = language_model_output.hidden_states[0][-1][:, -1]  # 第一个token的hidden state的last layer
+                
+                # actions_hidden_states_list = [actions_hidden_states0]
+                # for i in range(1, token_num):
+                #     token_hidden_state = language_model_output.hidden_states[i][-1][0]  # i+1个token的hidden state的last layer
+                #     actions_hidden_states_list.append(token_hidden_state)
+                # combined_hidden_states = torch.stack(actions_hidden_states_list, dim=0)  # (16, 4096)
+                # actions_hidden_states = combined_hidden_states.reshape(multimodal_embeddings.shape[0], token_num, multimodal_embeddings.shape[-1])
+                raise NotImplementedError
+            elif cfg.num_actions_chunk//cfg.num_actions_per_token == 1:
+                language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
+                # assert language_model_output.sequences == torch.tensor([[32001]], device=multimodal_embeddings.device)
+                actions_hidden_states = language_model_output.hidden_states[0][-1] 
+                cognition_features = actions_hidden_states[:, -1]
+                assert (cognition_features.shape[0], cognition_features.shape[1]) == (1,4096), "Batch size must be 1 for action prediction"
+                using_cfg = cfg.cfg_scale > 1.0
+
+                model_dtype = next(action_head.net.parameters()).dtype
+                B = cognition_features.shape[0]
+
+                cognition_features = cognition_features.unsqueeze(1).to(model_dtype) 
+                
+                noise = torch.randn(B, cfg.num_actions_chunk, action_head.net.in_channels, device=cognition_features.device).to(model_dtype) 
+
+                # TODO: Setup classifier-free guidance: now use cfg   
+                noise = torch.cat([noise, noise], 0) #  noise.shape torch.Size([2, 16, 7])
+                uncondition = action_head.net.z_embedder.uncondition # torch.Size([1, 4096])
+                uncondition = uncondition.unsqueeze(0)  #[1, D] # torch.Size([1, 1, 4096])
+                uncondition = uncondition.expand(B, 1, -1) #[B, 1, D]
+                z = torch.cat([cognition_features, uncondition], 0) # z  shape torch.Size([2, 1, 4096])
+                model_kwargs = dict(z=z, cfg_scale=cfg.cfg_scale)
+                sample_fn = action_head.net.forward_with_cfg
+                # default use ddim
+                if action_head.ddim_diffusion is None:
+                    action_head.create_ddim(ddim_step=cfg.num_ddim_steps)
+                samples = action_head.ddim_diffusion.ddim_sample_loop(sample_fn, 
+                                                                    noise.shape, 
+                                                                    noise, 
+                                                                    clip_denoised=False,
+                                                                    model_kwargs=model_kwargs,
+                                                                    progress=False,
+                                                                    device=cognition_features.device,
+                                                                    eta=0.0
+                                                                    )
+                if using_cfg:
+                    samples, _ = samples.chunk(2, dim=0)  # Remove null class samples
+                normalized_actions = samples[0].cpu().numpy()
+            else:
+                raise NotImplementedError
+        else:
+            raise NotImplementedError    
+
+
+
+        # normalized_actions = normalized_actions.float().cpu().detach().numpy()
+        actions = self._unnormalize_actions(normalized_actions, unnorm_key)
+
+        return actions, actions_hidden_states

modeling_prismatic.py.back.20250921_182648

@@ -0,0 +1,1552 @@
+"""
+modeling_prismatic.py
+
+Core HuggingFace-style PrismaticPreTrainedModel and PrismaticForConditionalGeneration class definitions.
+Inherits from the default `transformers.PretrainedModel`. Meant to be standalone and self-contained,
+but exactly replicate the logic in `prismatic.models.vlms.prismatic.py`.
+"""
+
+import logging
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import timm
+import tokenizers
+import torch
+import torch.nn as nn
+import transformers
+from timm.models.vision_transformer import LayerScale
+from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+from prismatic.models.action_heads import  L1RegressionActionHead, DiTActionHead, FlowMatchingActionHead
+from prismatic.training.train_utils import (
+    get_current_action_mask,
+    get_next_actions_mask,
+)
+from prismatic.vla.constants import (
+    ACTION_DIM,
+    ACTION_PROPRIO_NORMALIZATION_TYPE,
+    IGNORE_INDEX,
+    NUM_ACTIONS_CHUNK,
+    ACTION_TOKEN_IDX,
+    NormalizationType,
+)
+
+from .configuration_prismatic import OpenVLAConfig, PrismaticConfig
+
+# Set up logger
+logger = logging.getLogger(__name__)
+
+
+# === Utility Functions for Monkey-Patching ===
+def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]:
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        result = fn(*args, **kwargs)
+        return result[0] if isinstance(result, tuple) else result
+
+    return wrapper
+
+
+# HF Transformers overwrites parameters with names containing `gamma`; we're going to patch VisionBackbone.LayerScale.
+#   =>> TIMM :: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L109
+#   =>> Transformers :: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L3960
+def _ls_new_forward(self, x: torch.Tensor) -> torch.Tensor:
+    return x.mul_(self.scale_factor) if self.inplace else x * self.scale_factor
+
+
+def ls_apply_patch(ls_module: LayerScale):
+    ls_module.scale_factor = nn.Parameter(ls_module.gamma.clone())
+    ls_module.forward = _ls_new_forward.__get__(ls_module, LayerScale)
+    del ls_module.gamma
+
+
+# === Prismatic Vision Backbone (nn.Module) Definitions (w/ Fused Backbone Support) ===
+class PrismaticVisionBackbone(nn.Module):
+    """
+    Vision backbone for Prismatic models that handles image feature extraction.
+
+    Supports both single backbone (e.g., SigLIP) and fused backbone (e.g., SigLIP + DINOv2) configurations.
+    For fused backbones, features from both models are concatenated along the feature dimension.
+    """
+
+    def __init__(
+        self,
+        use_fused_vision_backbone: bool,
+        image_sizes: List[int],
+        timm_model_ids: List[str],
+        timm_override_act_layers: List[Optional[str]],
+    ) -> None:
+        """
+        Initialize the vision backbone.
+
+        Args:
+            use_fused_vision_backbone: Whether to use two backbones and fuse their features
+            image_sizes: List of image sizes for each backbone
+            timm_model_ids: List of TIMM model IDs to use for each backbone
+            timm_override_act_layers: List of activation layer overrides for each backbone
+        """
+        super().__init__()
+        self.use_fused_vision_backbone = use_fused_vision_backbone
+        self.num_images_in_input = 1  # Default value, can be overridden later
+
+        # Validate number of (fused) vision backbones
+        if len(timm_model_ids) > 2:
+            raise ValueError("Prismatic models only support up to 2 (fused) vision backbones!")
+
+        # Create primary featurizer
+        self.featurizer = self._create_featurizer(
+            model_id=timm_model_ids[0], img_size=image_sizes[0], act_layer=timm_override_act_layers[0]
+        )
+        self.embed_dim = self.featurizer.embed_dim
+
+        # Create secondary featurizer if using fused backbone
+        if self.use_fused_vision_backbone:
+            self.fused_featurizer = self._create_featurizer(
+                model_id=timm_model_ids[1], img_size=image_sizes[1], act_layer=timm_override_act_layers[1]
+            )
+            self.embed_dim += self.fused_featurizer.embed_dim
+
+        # Patch LayerScale modules for HF compatibility
+        self._patch_layer_scales()
+
+    def _create_featurizer(self, model_id: str, img_size: int, act_layer: Optional[str]) -> nn.Module:
+        """
+        Create a TIMM-based featurizer model with appropriate configurations.
+
+        Args:
+            model_id: The TIMM model ID to load
+            img_size: Input image size for the model
+            act_layer: Override for the activation layer type
+
+        Returns:
+            A configured featurizer model
+        """
+        featurizer = timm.create_model(
+            model_id,
+            pretrained=False,
+            num_classes=0,
+            img_size=img_size,
+            act_layer=act_layer,
+        )
+
+        # Monkey-patch the forward function to extract the second-to-last layer features
+        num_blocks = len(featurizer.blocks)
+        featurizer.forward = unpack_tuple(partial(featurizer.get_intermediate_layers, n={num_blocks - 2}))
+
+        return featurizer
+
+    def _patch_layer_scales(self) -> None:
+        """
+        Patch all LayerScale modules to be compatible with HF's parameter naming.
+
+        HF Transformers overwrites parameters with names containing 'gamma',
+        so we need to rename and modify the forward method.
+        """
+        # Patch primary featurizer
+        for module in self.featurizer.modules():
+            if isinstance(module, LayerScale):
+                ls_apply_patch(module)
+
+        # Patch secondary featurizer if it exists
+        if self.use_fused_vision_backbone:
+            for module in self.fused_featurizer.modules():
+                if isinstance(module, LayerScale):
+                    ls_apply_patch(module)
+
+    def get_num_patches(self) -> int:
+        """
+        Returns the number of vision patches output by the vision backbone.
+
+        Returns:
+            Number of patches per image
+        """
+        return self.featurizer.patch_embed.num_patches
+
+    def get_num_images_in_input(self) -> int:
+        """
+        Returns the number of input images for the vision backbone.
+
+        Returns:
+            Number of images expected in the input
+        """
+        return self.num_images_in_input
+
+    def set_num_images_in_input(self, num_images_in_input: int) -> None:
+        """
+        Sets the number of input images for the vision backbone.
+
+        Args:
+            num_images_in_input: Number of images to expect in the input
+        """
+        self.num_images_in_input = num_images_in_input
+
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        """
+        Implements the forward pass for the vision backbone.
+
+        If `self.use_fused_vision_backbone == True`, uses both SigLIP and DINOv2 transformers to extract visual features
+        (otherwise uses SigLIP only). Allows multi-image inputs (but only for fused vision backbone).
+
+        Args:
+            pixel_values (torch.Tensor): Pixels for input image(s), (B, C, H, W).
+        """
+        if self.num_images_in_input == 1:
+            if not self.use_fused_vision_backbone:
+                return self.featurizer(pixel_values)
+
+            # Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack
+            img, img_fused = torch.split(pixel_values, [3, 3], dim=1)
+            patches, patches_fused = self.featurizer(img), self.fused_featurizer(img_fused)
+
+            return torch.cat([patches, patches_fused], dim=2)
+
+        else:
+            assert self.use_fused_vision_backbone, "Multi-image inputs require using fused backbone!"
+
+            # Split `pixel_values` into individual images (each with 6 channels: 3 for SigLIP + 3 for DINOv2)
+            images = torch.split(pixel_values, [6] * self.num_images_in_input, dim=1)
+
+            # Process each image and collect patches
+            all_patches = []
+            for img in images:
+                # Split each image further into two stacks of channels (each with 3 channels)
+                img_regular, img_fused = torch.split(img, [3, 3], dim=1)
+
+                # Get patches from both SigLIP and DINOv2 vision transformers
+                patches = self.featurizer(img_regular)
+                patches_fused = self.fused_featurizer(img_fused)
+
+                # Concatenate SigLIP and DINOv2 patches along the hidden dimension
+                combined_patches = torch.cat([patches, patches_fused], dim=2)
+                all_patches.append(combined_patches)
+
+            # Concatenate all patches along the patch dimension
+            return torch.cat(all_patches, dim=1)
+
+
+# === Prismatic Projector (nn.Module) Definitions ===
+class PrismaticProjector(nn.Module):
+    def __init__(self, use_fused_vision_backbone: bool, vision_dim: int, llm_dim: int) -> None:
+        super().__init__()
+        self.use_fused_vision_backbone = use_fused_vision_backbone
+        self.vision_dim, self.llm_dim = vision_dim, llm_dim
+
+        # Switch on `use_fused_vision_backbone` =>> use slightly different MLPs and projection factors!
+        if not self.use_fused_vision_backbone:
+            self.fc1 = nn.Linear(self.vision_dim, self.llm_dim, bias=True)
+            self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
+            self.act_fn1 = nn.GELU()
+        else:
+            initial_projection_dim = 4 * vision_dim
+            self.fc1 = nn.Linear(self.vision_dim, initial_projection_dim, bias=True)
+            self.fc2 = nn.Linear(initial_projection_dim, self.llm_dim, bias=True)
+            self.fc3 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
+            self.act_fn1 = nn.GELU()
+            self.act_fn2 = nn.GELU()
+
+    def forward(self, img_patches: torch.Tensor) -> torch.Tensor:
+        if not self.use_fused_vision_backbone:
+            projected_features = self.fc1(img_patches)
+            projected_features = self.act_fn1(projected_features)
+            projected_features = self.fc2(projected_features)
+        else:
+            projected_features = self.fc1(img_patches)
+            projected_features = self.act_fn1(projected_features)
+            projected_features = self.fc2(projected_features)
+            projected_features = self.act_fn2(projected_features)
+            projected_features = self.fc3(projected_features)
+
+        return projected_features
+
+
+# === Main HF Class Definitions ===
+@dataclass
+class PrismaticCausalLMOutputWithPast(ModelOutput):
+    """Base class for Prismatic casual (visually-conditioned) language model outputs; also exposes visual features."""
+
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+    # Additions for VLMs
+    projector_features: Optional[torch.FloatTensor] = None
+
+
+class PrismaticPreTrainedModel(PreTrainedModel):
+    config_class: PretrainedConfig = PrismaticConfig
+    base_model_prefix: str = "model"
+    supports_gradient_checkpointing: bool = True
+
+    _no_split_modules: ClassVar[List[str]] = ["PrismaticProjector"]
+    _skip_keys_device_placement: str = "past_key_values"
+    _supports_flash_attn_2: bool = True
+
+    def _init_weights(self, module: nn.Module) -> None:
+        # Important :: this HF ported version is *not* meant for training from scratch; only inference and fine-tuning!
+        #   => As such, this init_weights code is not correct; if training VLMs from scratch, use the main codebase at
+        #      https://github.com/TRI-ML/prismatic-vlms
+        std = (
+            self.config.initializer_range
+            if hasattr(self.config, "initializer_range")
+            else self.config.text_config.initializer_range
+        )
+
+        if hasattr(module, "class_embedding"):
+            module.class_embedding.data.normal_(mean=0.0, std=std)
+
+        if isinstance(module, (nn.Linear, nn.Conv2d)):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+    @property
+    def _supports_sdpa(self) -> bool:
+        """Check LLM supports SDPA Attention"""
+        return self.language_model._supports_sdpa
+
+
+class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
+    def __init__(self, config: PrismaticConfig) -> None:
+        super().__init__(config)
+
+        # [Validation] Lightweight Validate on `config` Fields + Dependency Versions
+        if config.use_fused_vision_backbone is None:
+            raise ValueError("Missing config field `use_fused_vision_backbone`")
+
+        if timm.__version__ not in {"0.9.10", "0.9.11", "0.9.12", "0.9.16"}:
+            raise NotImplementedError(
+                "TIMM Version must be >= 0.9.10 and < 1.0.0 (breaking); please raise a GitHub Issue "
+                "if you urgently need support for latest TIMM versions."
+            )
+
+        if (transformers.__version__ != "4.40.1") or (tokenizers.__version__ != "0.19.1"):
+            logger.warning(
+                f"Expected `transformers==4.40.1` and `tokenizers==0.19.1` but got "
+                f"`transformers=={transformers.__version__}` and `tokenizers=={tokenizers.__version__}`; "
+                f"there might be inference-time regressions due to dependency changes. If in doubt, please"
+                f"use the above versions."
+            )
+
+        # Instantiate PrismaticVisionBackbone (w/ Potential Fused Backbone)
+        self.vision_backbone = PrismaticVisionBackbone(
+            config.use_fused_vision_backbone, config.image_sizes, config.timm_model_ids, config.timm_override_act_layers
+        )
+
+        # Create Multimodal Projector
+        self.projector = PrismaticProjector(
+            config.use_fused_vision_backbone,
+            vision_dim=self.vision_backbone.embed_dim,
+            llm_dim=config.text_config.hidden_size,
+        )
+
+        # Instantiate LLM Backbone
+        self.language_model = AutoModelForCausalLM.from_config(
+            config.text_config, attn_implementation=config._attn_implementation
+        )
+        self.vocab_size = config.text_config.vocab_size
+        self.pad_token_id = config.pad_token_id
+        self.llm_dim = config.text_config.hidden_size
+
+        # HF Boilerplate =>> initializes weights via `_init_weights()` and sets gradient checkpointing
+        self.post_init()
+
+    # === `PreTrainedModel` Boilerplate ===
+    def get_input_embeddings(self) -> nn.Module:
+        return self.language_model.get_input_embeddings()
+
+    def set_input_embeddings(self, value: nn.Module) -> None:
+        self.language_model.set_input_embeddings(value)
+
+    def get_output_embeddings(self) -> nn.Module:
+        return self.language_model.get_output_embeddings()
+
+    def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
+        self.language_model.set_output_embeddings(new_embeddings)
+
+    def get_decoder(self) -> nn.Module:
+        return self.language_model.get_decoder()
+
+    def set_decoder(self, decoder: nn.Module) -> None:
+        self.language_model.set_decoder(decoder)
+
+    def tie_weights(self) -> None:
+        self.language_model.tie_weights()  # Note: `Llama-2` and `Mistral` don't tie weights (no-op)
+
+    def resize_token_embeddings(
+        self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None
+    ) -> nn.Embedding:
+        updated_embeddings = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
+
+        # Update config/instance variables
+        self.config.text_config.vocab_size = updated_embeddings.num_embeddings
+        self.vocab_size = updated_embeddings.num_embeddings
+
+        return updated_embeddings
+
+    def _replace_input_embeddings(self, input_embeddings, all_actions_mask, noisy_action_features):
+        """
+        Replace embeddings in input_embeddings at positions where all_actions_mask is True
+        with embeddings from noisy_action_features, using vectorized operations.
+
+        Args:
+            input_embeddings: Tensor of shape (B, S, D)
+            all_actions_mask: Boolean tensor of shape (B, S)
+            noisy_action_features: Tensor of shape (B, K, D) where K is the number of True values in mask per sample
+
+        Returns:
+            Modified input_embeddings tensor
+        """
+        # Clone input to avoid modifying the original tensor
+        new_input_embeddings = input_embeddings.clone()
+
+        # Create a tensor with the same shape of input_embeddings to hold the noisy action features
+        repositioned_noisy_action_features = torch.zeros_like(input_embeddings)
+
+        # Create batch indices for splicing
+        batch_indices = torch.arange(input_embeddings.shape[0], device=input_embeddings.device)
+        batch_indices = batch_indices.unsqueeze(1).expand(-1, noisy_action_features.shape[1])
+
+        # Get indices where mask is True for each sample
+        masked_indices = torch.stack([torch.where(mask)[0] for mask in all_actions_mask])
+
+        # Move the noisy action features into their correct positions
+        repositioned_noisy_action_features[batch_indices, masked_indices] = noisy_action_features
+
+        # Combine original input embeddings and noisy action embeddings using the mask
+        new_input_embeddings = torch.where(
+            all_actions_mask.unsqueeze(-1), repositioned_noisy_action_features, new_input_embeddings
+        )
+
+        return new_input_embeddings
+
+    def _process_action_masks(self, labels):
+        """Helper to get action masks from labels"""
+        current_action_mask = get_current_action_mask(labels) # (B, seq_len)
+        next_actions_mask = get_next_actions_mask(labels) # (B, seq_len)
+        all_actions_mask = current_action_mask | next_actions_mask  # (B, seq_len)
+        return all_actions_mask
+
+    def _process_vision_features(self, pixel_values):
+        """Process vision features with optional FiLM conditioning"""
+        patch_features = self.vision_backbone(pixel_values)  # (bsz, 256 * num_images, D)
+
+        # Project patch embeddings into language embedding space
+        return self.projector(patch_features)
+
+    def _process_proprio_features(self, projected_patch_embeddings, proprio, proprio_projector):
+        """Process proprioceptive features and append to vision features"""
+        if proprio_projector is not None and proprio is not None:
+            # projected_patch_embeddings: (bsz, num_patches * num_images, llm_dim)
+            # proprio: (bsz, proprio_dim) or (propro_dim,)
+            proprio = proprio.reshape(projected_patch_embeddings.shape[0], -1)  # (bsz, proprio_dim)
+            proprio_features = proprio_projector(proprio)  # (bsz, llm_dim)
+            proprio_features = proprio_features.unsqueeze(dim=1)  # (bsz, 1, llm_dim)
+            # For simplicity, just append proprio token to the end of projected vision patch tokens
+            return torch.cat((projected_patch_embeddings, proprio_features), dim=1)
+        return projected_patch_embeddings
+
+    def _build_multimodal_attention(self, input_embeddings, projected_patch_embeddings, attention_mask):
+        """Build multimodal embeddings and attention mask"""
+        # juyi: Update attention mask 是不是要改成下三角? 不用, 因为generate会自动屏蔽
+        projected_patch_attention_mask = None
+        if attention_mask is not None:
+            projected_patch_attention_mask = torch.full(
+                (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
+                fill_value=True,
+                dtype=attention_mask.dtype,
+                device=attention_mask.device,
+            )
+
+        # Build multimodal embeddings & attention mask; insert embeddings after <BOS> token (1:)
+        multimodal_embeddings = torch.cat(
+            [input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1
+        )
+
+        multimodal_attention_mask = None
+        if attention_mask is not None:
+            multimodal_attention_mask = torch.cat(
+                [attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1
+            )
+
+        return multimodal_embeddings, multimodal_attention_mask
+
+    def _build_multimodal_labels(self, labels, projected_patch_embeddings):
+        """Build multimodal labels with IGNORE_INDEX for patch embeddings"""
+        if labels is not None:
+            projected_patch_labels = torch.full(
+                (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
+                fill_value=IGNORE_INDEX, # 这些位置不需要计算损失。
+                dtype=labels.dtype,
+                device=labels.device,
+            )
+            return torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1) # 第一个token是<BOS>
+        return None
+
+    # === Core Prismatic VLM `forward()` Logic ===
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_projector_features: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        proprio=None,
+        proprio_projector=None,
+        noisy_actions=None,
+        noisy_action_projector=None,
+        diffusion_timestep_embeddings=None,
+        use_film: bool = False,
+    ) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
+        """Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
+        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
+        )
+        output_projector_features = output_projector_features if output_projector_features is not None else False
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
+        use_cache = use_cache and not self.training
+
+        # Instantiate Placeholder for Projector Features
+        projected_patch_embeddings = None
+
+        # === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
+        if input_ids.shape[1] == 1:
+            assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
+            assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
+            assert labels is None, "Unexpected key `labels` provided during cached generation!"
+
+            language_model_output = self.language_model(
+                input_ids=input_ids,
+                attention_mask=None,
+                position_ids=None,
+                past_key_values=past_key_values,
+                inputs_embeds=None,
+                labels=None,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+        # === Handle Unimodal Forward ===
+        elif pixel_values is None:
+            assert (input_ids is not None) and (inputs_embeds is None), "Missing `input_ids` in language-only forward!"
+            assert past_key_values is None, "Unexpected key `past_key_values` provided during language-only forward!"
+
+            language_model_output = self.language_model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                position_ids=None,
+                past_key_values=None,
+                inputs_embeds=None,
+                labels=labels,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+        # === Handle Multimodal Forward ===
+        elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
+            assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
+
+            # Get input embeddings (from language model embeddings)
+            input_embeddings = self.get_input_embeddings()(input_ids)  # (B, seq_len, D)
+
+            # Extract action masks
+            all_actions_mask = self._process_action_masks(labels)
+
+            # Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
+            language_embeddings = input_embeddings[~all_actions_mask].reshape(
+                input_embeddings.shape[0], -1, input_embeddings.shape[2]
+            )  # (B, lang_seq_len, llm_dim)
+
+            # Get visual features
+            projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
+            # bug: TypeError: PrismaticForConditionalGeneration._process_vision_features() takes 2 positional arguments but 4 were given
+
+            # Add proprioceptive state if provided
+            projected_patch_embeddings = self._process_proprio_features(
+                projected_patch_embeddings, proprio, proprio_projector
+            )
+
+            # [Diffusion] Add diffusion timestep embedding if provided
+            if diffusion_timestep_embeddings is not None:
+                # For simplicity, just append diffusion timestep embedding to the end of projected vision patch tokens
+                projected_patch_embeddings = torch.cat(
+                    (projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
+                )
+                
+            # Process action embeddings
+            if noisy_actions is not None:
+                # Get mask corresponding to all action tokens
+                all_actions_mask = self._process_action_masks(labels)
+
+                # Reshape noisy actions into individual action tokens
+                # noisy_actions: (B, chunk_len, action_dim) -> (B, chunk_len * action_dim, 1)
+                B = noisy_actions.shape[0]
+                noisy_actions = noisy_actions.reshape(B, -1).unsqueeze(-1)
+
+                # Project noisy action tokens into language model embedding space
+                noisy_action_features = noisy_action_projector(noisy_actions)  # (B, chunk_len * action_dim, llm_dim)
+
+                # Replace embeddings of the action tokens with noisy action embeddings
+                input_embeddings = self._replace_input_embeddings(
+                    input_embeddings, all_actions_mask, noisy_action_features
+                )
+            else:
+                # Replace the embeddings of the action tokens with zeros
+                # (Later on, the positional embeddings will be added to them)
+                all_actions_mask = all_actions_mask.unsqueeze(-1)  # (B, seq_len, 1)
+                input_embeddings = input_embeddings * ~all_actions_mask
+
+            # Build multimodal embeddings & attention mask
+            multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+                input_embeddings, projected_patch_embeddings, attention_mask
+            )
+
+            # Build labels for multimodal sequence if needed
+            multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
+
+            # Dispatch to language model
+            language_model_output = self.language_model(
+                input_ids=None,
+                attention_mask=multimodal_attention_mask,
+                position_ids=None,
+                past_key_values=None,
+                inputs_embeds=multimodal_embeddings,
+                labels=multimodal_labels,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+        # === Otherwise =>> Assume Invalid! ===
+        elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
+            raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
+
+        else:
+            raise ValueError(
+                "Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
+                f"=> `input_ids` = {input_ids is not None}\n"
+                f"=> `attention_mask` = {attention_mask is not None}\n"
+                f"=> `pixel_values` = {pixel_values is not None}\n"
+                f"=> `labels` = {labels is not None}\n"
+                f"=> `input_embeds` = {inputs_embeds is not None}\n"
+                f"=> `past_key_values` = {past_key_values is not None}\n"
+                f"=> `use_cache` = {use_cache}"
+            )
+
+        # Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
+        if not return_dict:
+            if output_projector_features and (projected_patch_embeddings is not None):
+                return *language_model_output, projected_patch_embeddings
+
+            return language_model_output
+
+        return PrismaticCausalLMOutputWithPast(
+            loss=language_model_output.loss,
+            logits=language_model_output.logits,
+            past_key_values=language_model_output.past_key_values,
+            hidden_states=language_model_output.hidden_states,
+            attentions=language_model_output.attentions,
+            projector_features=projected_patch_embeddings,
+        )
+
+    # === GenerationMixin Methods ===
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs: str,
+    ) -> Dict[str, torch.Tensor]:
+        """Borrowed from `LlamaForCausalLM` and simplified for batch size = 1; mirrors original PrismaticVLM logic."""
+        if ((input_ids is not None) and (input_ids.shape[0] > 1)) or (
+            (inputs_embeds is not None) and (inputs_embeds.shape[0] > 1)
+        ):
+            raise ValueError("Generation with batch size > 1 is not currently supported!")
+
+        # Handle `past_key_values` (cache) =>> assume `input_ids` just has unprocessed tokens
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1:]
+
+        # If `input_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 = {"input_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        # Make sure `pixel_values` are preserved in `model_inputs`
+        model_inputs.update(
+            {
+                "attention_mask": attention_mask,
+                "pixel_values": pixel_values,
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+            }
+        )
+
+        return model_inputs
+
+    # Defer to Language Model (all handle this differently, with different return types)
+    def _reorder_cache(self, *args, **kwargs) -> Any:
+        return self.language_model._reorder_cache(*args, **kwargs)
+
+
+class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
+    config_class: PretrainedConfig = OpenVLAConfig
+
+    def __init__(self, config: OpenVLAConfig) -> None:
+        super().__init__(config)
+        self.norm_stats = config.norm_stats
+
+        # Compute action bins
+        self.bins = np.linspace(-1, 1, config.n_action_bins)
+        self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0
+
+        # Compute vocab size for de-tokenization -- revert added "multiple of"
+        self.vocab_size = self.config.text_config.vocab_size - self.config.pad_to_multiple_of
+
+    def _prepare_input_for_action_prediction(self, input_ids, attention_mask):
+        # eval 会用到这里
+        """Prepares input for action prediction by adding necessary tokens"""
+        # Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens
+        placeholder_action_token_ids = (
+            torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
+        ) 
+        input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1) # torch.Size([1,  35 + 56= 91])
+
+        # Extend the attention mask to fit the new shape of input
+        # Note: Only batch size == 1 supported right now
+        mask_extension = (
+            torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
+            .to(attention_mask.device)
+            .to(attention_mask.dtype)
+        )
+        attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
+
+        return input_ids, attention_mask
+
+    def _prepare_labels_for_action_prediction(self, labels, input_ids):
+        """Creates labels tensor for action prediction if not provided"""
+        # eval 会用到这里 , 
+        # Extends label tensors with fake action labels
+        # Adds stop tokens at the end of sequences
+        # Handles label preparation for action prediction tasks
+        # 他为啥可以随便一个?  xuan说 你自定义一个值 ,然后一直指定这个 , PAD token可以吗? 
+        #TODO: 这里是否要改?  感觉不需要改. 随便写就行了因为labels不重要只是要一个mask. 为什么需要这个函数? 确保 action 预测任务的标签（labels）符合模型的输入长度，并正确地处理序列终止
+        # Extend labels tensor with fake action labels
+        ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_IDX # = 为了mask正确生成, action_tokens_only_mask = (labels == ACTION_TOKEN_IDX ), 所以这里也填上ACTION_TOKEN_IDX
+        labels_extension = (
+            torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
+            * ARBITRARY_ACTION_TOKEN_IDX
+        ) #torch.Size([1, 57]),全是 ARBITRARY_ACTION_TOKEN_IDX
+        labels = torch.cat([labels, labels_extension], dim=-1)
+
+        return labels
+
+    def _unnormalize_actions(self, normalized_actions, unnorm_key=None):
+        """Unnormalize actions using dataset statistics"""
+        action_norm_stats = self.get_action_stats(unnorm_key)
+
+        if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:
+            mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))
+            action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])
+        elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:
+            mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
+            action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
+        else:
+            raise ValueError("Unsupported action/proprio normalization type detected!")
+
+        actions = np.where(
+            mask,
+            0.5 * (normalized_actions + 1) * (action_high - action_low + 1e-8) + action_low,
+            normalized_actions,
+        )
+
+        return actions
+
+
+    def _normalize_actions(self, actions, norm_key=None):
+        """Normalize actions to [-1, 1] using dataset statistics"""
+        action_norm_stats = self.get_action_stats(norm_key)
+
+        if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:
+            mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))
+            action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])
+        elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:
+            mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
+            action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
+        else:
+            raise ValueError("Unsupported action/proprio normalization type detected!")
+
+        normalized = np.where(
+            mask,
+            2 * (actions - action_low) / (action_high - action_low + 1e-8) - 1,
+            actions,
+        )
+
+        return normalized
+
+    def _run_diffusion_prediction(
+        self,
+        input_embeddings,
+        all_actions_mask,
+        noise,
+        action_head,
+        projected_patch_embeddings,
+        labels,
+        attention_mask,
+        NUM_PATCHES,
+        NUM_PROMPT_TOKENS,
+        noisy_action_projector,
+    ):
+        """Run diffusion-based action prediction"""
+        # Set diffusion timestep values
+        action_head.noise_scheduler.set_timesteps(action_head.num_diffusion_steps)
+        # Clone embedding for reuse in each timestep
+        orig_projected_patch_embeddings = projected_patch_embeddings.clone()
+        curr_noisy_actions = noise
+
+        # Reverse diffusion: Iteratively denoise to generate action prediction
+        for t in action_head.noise_scheduler.timesteps:
+            # Get diffusion model's noise prediction (conditioned on VLA latent embedding, current noisy action
+            # embedding, and diffusion timestep embedding)
+            timesteps = torch.Tensor([t]).to(labels.device)
+            diffusion_timestep_embeddings = (
+                action_head.time_encoder(timesteps).to(curr_noisy_actions.dtype).to(curr_noisy_actions.device)
+            )  # (B, llm_dim)
+            diffusion_timestep_embeddings = diffusion_timestep_embeddings.unsqueeze(1)  # (B, 1, llm_dim)
+
+            # [Diffusion] Replace the embeddings of the action tokens with noisy actions
+            # (Later on, the positional embeddings will be added to them)
+
+            # For simplicity, append diffusion timestep embedding to the end of projected vision tokens
+            projected_patch_embeddings = torch.cat(
+                (orig_projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
+            )
+
+            # Reshape and project noisy actions into language embedding space
+            B = curr_noisy_actions.shape[0]
+            orig_curr_noisy_actions_shape = curr_noisy_actions.shape
+            curr_noisy_actions = curr_noisy_actions.reshape(B, -1).unsqueeze(-1)
+            noisy_action_features = noisy_action_projector(curr_noisy_actions)
+            curr_noisy_actions = curr_noisy_actions.reshape(orig_curr_noisy_actions_shape)
+
+            # Replace action token embeddings with noisy action embeddings
+            input_embeddings = self._replace_input_embeddings(
+                input_embeddings.clone(), all_actions_mask, noisy_action_features
+            )
+
+            # Build multimodal embeddings and attention mask
+            multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+                input_embeddings, projected_patch_embeddings, attention_mask
+            )
+
+            # Forward pass through language model
+            language_model_output = self.language_model(
+                input_ids=None,
+                attention_mask=multimodal_attention_mask,
+                position_ids=None,
+                past_key_values=None,
+                inputs_embeds=multimodal_embeddings,
+                labels=None,
+                use_cache=None,
+                output_attentions=False,
+                output_hidden_states=True,
+                return_dict=True,
+            )
+
+            # Extract hidden states for action portion of response
+            last_hidden_states = language_model_output.hidden_states[-1]  # (B, seq_len, D)
+            actions_hidden_states = last_hidden_states[
+                :,
+                NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
+                :,
+            ]  # (B, act_chunk_len, D)
+
+            # Predict noise and update noisy actions: x_t -> x_{t-1}
+            noise_pred = action_head.predict_noise(actions_hidden_states)
+            curr_noisy_actions = action_head.noise_scheduler.step(noise_pred, t, curr_noisy_actions).prev_sample
+
+        curr_noisy_actions = curr_noisy_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
+
+        # Return final actions
+        return curr_noisy_actions.float().cpu().detach().numpy(), actions_hidden_states
+
+    def _regression_or_discrete_prediction(
+        self,
+        input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
+        all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action . 
+        projected_patch_embeddings: torch.FloatTensor, 
+        attention_mask: torch.BoolTensor,
+        labels: torch.LongTensor,
+        NUM_PATCHES: int,
+        NUM_PROMPT_TOKENS: int,
+        action_head: L1RegressionActionHead,
+        **kwargs,
+    ):
+        """Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
+        # Extract hidden states for action tokens
+        # last_hidden_states = language_model_output.hidden_states[-1]  # (B, seq_len, D)
+
+        # actions_hidden_states = last_hidden_states[:, NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + NUM_ACTIONS_CHUNK * tokennum, :]# (B, act_chunk_len, D)
+        # 都不需要取了, 直接就给 token对应的hidden state了 ,太方便了. 
+        # 为什么第一个是torch.Size([1, 535, 4096])? 我应该选哪个?  https://discuss.huggingface.co/t/get-each-generated-token-last-layer-hidden-state/145921
+        # language_model_output.sequences tensor([[29871, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 32001,     2]], device='cuda:0')
+        cfg = kwargs.pop("cfg", None)
+        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+            input_embeddings, projected_patch_embeddings, attention_mask
+        ) 
+        # multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'  
+        language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
+        #  is tuple (1 token,  33 layers, torch.Size([1, 314, 4096]))
+        hidden_states = language_model_output.hidden_states[0][-1]
+        actions_hidden_states = hidden_states[:, -NUM_ACTIONS_CHUNK:]
+
+        normalized_actions = action_head.predict_action(actions_hidden_states)
+        normalized_actions = normalized_actions.reshape(cfg.num_actions_chunk, ACTION_DIM)
+        normalized_actions = normalized_actions.float().cpu().detach().numpy()
+        if cfg.mode == "mul": 
+            if cfg.num_actions_chunk//cfg.num_actions_per_token > 1:
+                token_num = cfg.num_actions_chunk//cfg.num_actions_per_token
+                language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,min_new_tokens=token_num, max_new_tokens=token_num,output_hidden_states=True,return_dict_in_generate=True)
+                actions_hidden_states0 = language_model_output.hidden_states[0][-1][:, -1]  # 第一个token的hidden state的last layer
+                
+                actions_hidden_states_list = [actions_hidden_states0]
+                for i in range(1, token_num):
+                    token_hidden_state = language_model_output.hidden_states[i][-1][0]  # i+1个token的hidden state的last layer
+                    actions_hidden_states_list.append(token_hidden_state)
+                # 将所有hidden states拼接起来
+                combined_hidden_states = torch.stack(actions_hidden_states_list, dim=0)  # (16, 4096)
+                actions_hidden_states = combined_hidden_states.reshape(multimodal_embeddings.shape[0], token_num, multimodal_embeddings.shape[-1])
+            elif cfg.num_actions_chunk//cfg.num_actions_per_token == 1:
+                language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
+                # assert language_model_output.sequences == torch.tensor([[32001]], device=multimodal_embeddings.device)
+                actions_hidden_states = language_model_output.hidden_states[0][-1] 
+                actions_hidden_states = actions_hidden_states[:, -1]
+            else:
+                raise NotImplementedError
+        else:
+            raise NotImplementedError    
+        return normalized_actions, actions_hidden_states
+
+    def hist_predict_action(
+        self,
+        input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
+        all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action . 
+        projected_patch_embeddings: torch.FloatTensor, 
+        attention_mask: torch.BoolTensor,
+        action_head: L1RegressionActionHead,
+        **kwargs,
+    ):
+        cfg = kwargs.get("cfg", None) 
+        action_history = kwargs.get("action_history", None)
+        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+            input_embeddings, projected_patch_embeddings, attention_mask
+        ) 
+        # multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'  
+        # first language_model_output.hidden_states  , is tuple (1 token,  33 layers, torch.Size([1, 314, 4096]))
+        # the following is  (num of tokens,)
+        if cfg.mode == "mul": 
+            if cfg.num_actions_chunk//cfg.num_actions_per_token > 1:
+                raise NotImplementedError
+                # token_num = cfg.num_actions_chunk//cfg.num_actions_per_token
+                # language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,min_new_tokens=token_num, max_new_tokens=token_num,output_hidden_states=True,return_dict_in_generate=True)
+                # actions_hidden_states0 = language_model_output.hidden_states[0][-1][:, -1]  # 第一个token的hidden state的last layer
+                # actions_hidden_states_list = [actions_hidden_states0]
+                # for i in range(1, token_num):
+                #     token_hidden_state = language_model_output.hidden_states[i][-1][0]  # i+1个token的hidden state的last layer
+                #     actions_hidden_states_list.append(token_hidden_state)
+                # combined_hidden_states = torch.stack(actions_hidden_states_list, dim=0)  # (16, 4096)
+                # actions_hidden_states = combined_hidden_states.reshape(multimodal_embeddings.shape[0], token_num, multimodal_embeddings.shape[-1])
+            elif cfg.num_actions_chunk//cfg.num_actions_per_token == 1:
+                language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
+                # assert language_model_output.sequences == torch.tensor([[32001]], device=multimodal_embeddings.device)
+                actions_hidden_states = language_model_output.hidden_states[0][-1] 
+                actions_hidden_states = actions_hidden_states[:, -1]
+                # 在中间加一个 1 维度
+                actions_hidden_states = actions_hidden_states.unsqueeze(1) # for match 3 dim 
+            else:
+                raise NotImplementedError
+        else:
+            raise NotImplementedError    
+
+        normalized_actions = action_head.predict_action(actions_hidden_states, action_history)
+        normalized_actions = normalized_actions.reshape(cfg.num_actions_chunk, ACTION_DIM)
+        normalized_actions = normalized_actions.float().cpu().detach().numpy()
+
+        return normalized_actions, actions_hidden_states
+
+    def mul_regression_or_discrete_prediction(
+        self,
+        input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
+        all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action . 
+        projected_patch_embeddings: torch.FloatTensor, 
+        attention_mask: torch.BoolTensor,
+        labels: torch.LongTensor,
+        NUM_PATCHES: int,
+        NUM_PROMPT_TOKENS: int,
+        action_head: L1RegressionActionHead,
+        **kwargs,
+    ):
+        cfg = kwargs.get("cfg", None) 
+        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+            input_embeddings, projected_patch_embeddings, attention_mask
+        ) 
+        # multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'  
+        # first language_model_output.hidden_states  , is tuple (1 token,  33 layers, torch.Size([1, 314, 4096]))
+        # the following is  (num of tokens,)
+        if cfg.mode == "mul": 
+            if cfg.num_actions_chunk//cfg.num_actions_per_token > 1:
+                token_num = cfg.num_actions_chunk//cfg.num_actions_per_token
+                language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,min_new_tokens=token_num, max_new_tokens=token_num,output_hidden_states=True,return_dict_in_generate=True)
+                actions_hidden_states0 = language_model_output.hidden_states[0][-1][:, -1]  # 第一个token的hidden state的last layer
+                
+                actions_hidden_states_list = [actions_hidden_states0]
+                for i in range(1, token_num):
+                    token_hidden_state = language_model_output.hidden_states[i][-1][0]  # i+1个token的hidden state的last layer
+                    actions_hidden_states_list.append(token_hidden_state)
+                combined_hidden_states = torch.stack(actions_hidden_states_list, dim=0)  # (16, 4096)
+                actions_hidden_states = combined_hidden_states.reshape(multimodal_embeddings.shape[0], token_num, multimodal_embeddings.shape[-1])
+            elif cfg.num_actions_chunk//cfg.num_actions_per_token == 1:
+                language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
+                actions_hidden_states = language_model_output.hidden_states[0][-1] 
+                actions_hidden_states = actions_hidden_states[:, -1]
+            else:
+                raise NotImplementedError
+        else:
+            raise NotImplementedError    
+
+        normalized_actions = action_head.predict_action(actions_hidden_states)
+        normalized_actions = normalized_actions.reshape(cfg.num_actions_chunk, ACTION_DIM)
+        # print(f"*** normalized_actions[]: {normalized_actions} ***")
+        if cfg.action_head_name == "medusa":
+            normalized_actions[:, 6] = torch.sigmoid(normalized_actions[:,  6]) # without bs dim.
+        # print(f"*** normalized_actions[]: {normalized_actions} ***")
+        normalized_actions = normalized_actions.float().cpu().detach().numpy()
+
+        return normalized_actions, actions_hidden_states
+
+    def predict_action(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        unnorm_key: Optional[str] = None,
+        proprio=None,
+        proprio_projector=None,
+        action_head=None,
+        noisy_action_projector=None,
+        use_film: bool = False,
+        **kwargs: str,
+    ) -> np.ndarray:
+        """Predict actions from input sequence, with options for different prediction methods.
+
+        Args:
+            input_ids: Input token ids
+            unnorm_key: Key for unnormalization statistics
+            proprio: Proprioceptive features
+            proprio_projector: Projector for proprioceptive features
+            action_head: Optional head for L1 regression or diffusion-based prediction
+            noisy_action_projector: Projector for noisy actions in diffusion-based prediction
+            use_film: Whether to use FiLM conditioning
+            **kwargs: Additional arguments including pixel_values and attention_mask
+
+        Returns:
+            Tuple of (unnormalized_actions, action_hidden_states)
+        """
+        # If the special empty token ('') does not already appear after the colon (':') token in the prompt
+        # (after "OUT:" or "ASSISTANT:"), insert it to match the inputs seen at training time
+        if not torch.all(input_ids[:, -1] == 29871):
+            input_ids = torch.cat(
+                (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
+            )
+
+        pixel_values = kwargs["pixel_values"]
+        attention_mask = kwargs["attention_mask"]
+
+        # Create fake labels tensor (needed for action mask)
+        labels = input_ids.clone()
+        labels[:] = IGNORE_INDEX # 输入都ignore IGNORE_INDEX = -100
+
+        # Get number of tokens in prompt (excluding the start token)
+        NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1  # Subtract action tokens and stop token
+
+        # Prepare inputs by adding necessary tokens
+        input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)
+
+        # Update labels tensor for action mask computation later
+        labels = self._prepare_labels_for_action_prediction(labels, input_ids)
+
+        # Get input embeddings and action masks
+        input_embeddings = self.get_input_embeddings()(input_ids)
+        all_actions_mask = self._process_action_masks(labels)
+
+        # Extract language embeddings
+        language_embeddings = input_embeddings[~all_actions_mask].reshape(
+            input_embeddings.shape[0], -1, input_embeddings.shape[2]
+        )
+
+        # Process vision features
+        projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
+
+        # Add proprioceptive features if provided
+        use_proprio = proprio_projector is not None and proprio is not None
+        if use_proprio:
+            proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
+            projected_patch_embeddings = self._process_proprio_features(
+                projected_patch_embeddings, proprio, proprio_projector
+            )
+
+        # Use diffusion if provided, otherwise use regression or discrete prediction
+        use_diffusion = noisy_action_projector is not None and hasattr(action_head, "noise_scheduler")
+
+        # Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
+        NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
+        if use_proprio:
+            NUM_PATCHES += 1
+
+        normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
+            input_embeddings,
+            all_actions_mask,
+            projected_patch_embeddings,
+            attention_mask,
+            labels,
+            NUM_PATCHES,
+            NUM_PROMPT_TOKENS,
+            action_head,
+        )
+
+        # Unnormalize predicted actions
+        actions = self._unnormalize_actions(normalized_actions, unnorm_key)
+
+        return actions, actions_hidden_states
+
+    @staticmethod
+    def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_key: Optional[str]) -> str:
+        """Validate and resolve the unnormalization key for action statistics"""
+        if unnorm_key is None:
+            assert len(norm_stats) == 1, (
+                f"Your model was trained on more than one dataset, "
+                f"please pass a `unnorm_key` from the following options to choose the statistics "
+                f"used for un-normalizing actions: {norm_stats.keys()}"
+            )
+            unnorm_key = next(iter(norm_stats.keys()))
+        # norm states没有加载libero, 为什么? 
+        assert unnorm_key in norm_stats, (
+            f"The `unnorm_key` you chose is not in the set of available dataset statistics, "
+            f"please choose from: {norm_stats.keys()}"
+        )
+        return unnorm_key
+
+    def get_action_dim(self, unnorm_key: Optional[str] = None) -> int:
+        """Get the dimensionality of the policy's action space."""
+        unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
+        return len(self.norm_stats[unnorm_key]["action"]["min"])
+
+    def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[str, Any]:
+        """Get all the logged statistics for the given dataset."""
+        unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
+        return self.norm_stats[unnorm_key]["action"]
+
+
+    def lisa_forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_projector_features: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        proprio=None,
+        proprio_projector=None,
+        **kwargs
+    ) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
+        """Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
+        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
+        )
+        output_projector_features = output_projector_features if output_projector_features is not None else False
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
+        use_cache = use_cache and not self.training
+
+        # Instantiate Placeholder for Projector Features
+        projected_patch_embeddings = None
+
+        # === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
+        if input_ids.shape[1] == 1:
+            assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
+            assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
+            assert labels is None, "Unexpected key `labels` provided during cached generation!"
+
+            language_model_output = self.language_model(
+                input_ids=input_ids,
+                attention_mask=None,
+                position_ids=None,
+                past_key_values=past_key_values,
+                inputs_embeds=None,
+                labels=None,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+        # === Handle Unimodal Forward ===
+        elif pixel_values is None:
+            raise NotImplementedError
+
+        # === Handle Multimodal Forward ===
+        elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
+            assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
+
+            # Get input embeddings (from language model embeddings)
+            input_embeddings = self.get_input_embeddings()(input_ids)  # (B, seq_len, D)
+            # Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
+            # language_embeddings = input_embeddings[~all_actions_mask].reshape(
+            #     input_embeddings.shape[0], -1, input_embeddings.shape[2]
+            # )  # (B, lang_seq_len, llm_dim) 这里就会把结尾的 stop index和padding 也算进去. 没问题吗? 没问题因为ignore了 我直接删了因为不用film
+            # Get visual features
+            projected_patch_embeddings = self._process_vision_features(pixel_values)
+
+            # Add proprioceptive state if provided
+            projected_patch_embeddings = self._process_proprio_features(
+                projected_patch_embeddings, proprio, proprio_projector
+            )
+
+            all_actions_mask = (labels == ACTION_TOKEN_IDX) #和run forward pass不一样, run forward pass要手动算token number来找偏移. 
+            all_actions_mask = all_actions_mask.unsqueeze(-1)  # (B, seq_len, 1)
+            input_embeddings = input_embeddings * ~all_actions_mask
+
+            # Build multimodal embeddings & attention mask
+            multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+                input_embeddings, projected_patch_embeddings, attention_mask
+            )
+
+            # Build labels for multimodal sequence if needed
+            multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
+
+            # Dispatch to language model
+            language_model_output = self.language_model(
+                input_ids=None,
+                attention_mask=multimodal_attention_mask,
+                position_ids=None,
+                past_key_values=None,
+                inputs_embeds=multimodal_embeddings,
+                labels=multimodal_labels,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+
+        # === Otherwise =>> Assume Invalid! ===
+        elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
+            raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
+
+        else:
+            raise ValueError(
+                "Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
+                f"=> `input_ids` = {input_ids is not None}\n"
+                f"=> `attention_mask` = {attention_mask is not None}\n"
+                f"=> `pixel_values` = {pixel_values is not None}\n"
+                f"=> `labels` = {labels is not None}\n"
+                f"=> `input_embeds` = {inputs_embeds is not None}\n"
+                f"=> `past_key_values` = {past_key_values is not None}\n"
+                f"=> `use_cache` = {use_cache}"
+            )
+
+        # Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
+        if not return_dict:
+            if output_projector_features and (projected_patch_embeddings is not None):
+                return *language_model_output, projected_patch_embeddings
+
+            return language_model_output
+
+        return PrismaticCausalLMOutputWithPast(
+            loss=language_model_output.loss,
+            logits=language_model_output.logits,
+            past_key_values=language_model_output.past_key_values,
+            hidden_states=language_model_output.hidden_states,
+            attentions=language_model_output.attentions,
+            projector_features=projected_patch_embeddings,
+        )
+
+
+
+    def mul_predict_action(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,  #就是 language instruction.
+        unnorm_key: Optional[str] = None,
+        proprio=None,
+        proprio_projector=None,
+        action_head:L1RegressionActionHead=None,
+        noisy_action_projector=None,
+        use_film: bool = False,
+        **kwargs: str,
+    ) -> np.ndarray:
+    # only use in evaluation.
+        cfg = kwargs.get("cfg", None)  # Extract cfg from kwargs
+        action_history = kwargs.get("action_history", None)
+
+        if not torch.all(input_ids[:, -1] == 29871):
+            input_ids = torch.cat(
+                (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
+            )
+
+        
+        pixel_values = kwargs["pixel_values"]
+        attention_mask = kwargs["attention_mask"]
+
+        #  input id  '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:' 
+        
+        input_embeddings = self.get_input_embeddings()(input_ids) 
+        
+        projected_patch_embeddings = self._process_vision_features(pixel_values)
+
+        use_proprio = proprio_projector is not None and proprio is not None
+        if use_proprio:
+            proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
+            projected_patch_embeddings = self._process_proprio_features(
+                projected_patch_embeddings, proprio, proprio_projector
+            )
+        if cfg.action_head_name == "hist":
+            normalized_actions, actions_hidden_states = self.hist_predict_action(
+                input_embeddings,
+                None, 
+                projected_patch_embeddings,
+                attention_mask,
+                action_head,
+                cfg=cfg,
+                action_history=action_history,
+            )
+        else:
+            normalized_actions, actions_hidden_states = self.mul_regression_or_discrete_prediction(
+                input_embeddings,
+                None, 
+                projected_patch_embeddings,
+                attention_mask,
+                None, #推理不需要labels
+                None, #推理不需要NUM_PATCHES
+                None, #推理不需要NUM_PROMPT_TOKENS
+                action_head,
+                cfg=cfg,
+            )
+        
+
+        actions = self._unnormalize_actions(normalized_actions, unnorm_key) #在这里 unorm, 所以出来的已经是unorm的了.  所以我 history 也要记录 norm 的. 
+
+        return actions, normalized_actions
+
+
+    def flow_matching_predict_action(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        unnorm_key: Optional[str] = None,
+        proprio=None,
+        proprio_projector=None,
+        action_head: FlowMatchingActionHead = None,
+        noisy_action_projector=None,
+        use_film: bool = False,
+        **kwargs: str,
+    ) -> np.ndarray:
+        """Predict actions using Flow Matching"""
+        cfg = kwargs.get("cfg", None)
+
+        if not torch.all(input_ids[:, -1] == 29871):
+            input_ids = torch.cat(
+                (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
+            )
+
+        pixel_values = kwargs["pixel_values"]
+        attention_mask = kwargs["attention_mask"]
+
+        input_embeddings = self.get_input_embeddings()(input_ids)
+        projected_patch_embeddings = self._process_vision_features(pixel_values)
+
+        use_proprio = proprio_projector is not None and proprio is not None
+        if use_proprio:
+            proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
+            projected_patch_embeddings = self._process_proprio_features(
+                projected_patch_embeddings, proprio, proprio_projector
+            )
+
+        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+            input_embeddings, projected_patch_embeddings, attention_mask
+        )
+
+        if cfg.mode == "flow_matching":
+            if cfg.num_actions_chunk//cfg.num_actions_per_token == 1:
+                language_model_output = self.language_model.generate(
+                    inputs_embeds=multimodal_embeddings,
+                    max_new_tokens=1,
+                    output_hidden_states=True,
+                    return_dict_in_generate=True
+                )
+                
+                actions_hidden_states = language_model_output.hidden_states[0][-1]
+                cognition_features = actions_hidden_states[:, -1]
+                assert (cognition_features.shape[0], cognition_features.shape[1]) == (1, 4096), "Batch size must be 1 for action prediction"
+
+                model_dtype = next(action_head.net.parameters()).dtype
+                cognition_features = cognition_features.unsqueeze(1).to(model_dtype)
+
+                # Sample actions using flow matching
+                normalized_actions = action_head.sample_actions(
+                    cognition_features,
+                    num_steps=getattr(cfg, 'num_flow_steps', 20)
+                )
+                normalized_actions = normalized_actions[0].cpu().numpy()
+            else:
+                raise NotImplementedError("Multi-token flow matching not yet implemented")
+        else:
+            raise NotImplementedError
+
+        actions = self._unnormalize_actions(normalized_actions, unnorm_key)
+        return actions, actions_hidden_states
+
+    def diffusion_predict_action(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,  #就是 language instruction.
+        unnorm_key: Optional[str] = None,
+        proprio=None,
+        proprio_projector=None,
+        action_head:DiTActionHead=None,
+        noisy_action_projector=None,
+        use_film: bool = False,
+        **kwargs: str,
+    ) -> np.ndarray:
+        cfg = kwargs.get("cfg", None)  # Extract cfg from kwargs
+
+        if not torch.all(input_ids[:, -1] == 29871):
+            input_ids = torch.cat(
+                (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
+            )
+
+        
+        pixel_values = kwargs["pixel_values"]
+        attention_mask = kwargs["attention_mask"]
+
+        #  input id  '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:' 
+        
+        input_embeddings = self.get_input_embeddings()(input_ids) 
+        
+        projected_patch_embeddings = self._process_vision_features(pixel_values)
+
+        use_proprio = proprio_projector is not None and proprio is not None
+        if use_proprio:
+            proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
+            projected_patch_embeddings = self._process_proprio_features(
+                projected_patch_embeddings, proprio, proprio_projector
+            )
+        
+        # normalized_actions, actions_hidden_states = self.mul_regression_or_discrete_prediction(
+        #     input_embeddings,
+        #     None, 
+        #     projected_patch_embeddings,
+        #     attention_mask,
+        #     None, #推理不需要labels
+        #     None, #推理不需要NUM_PATCHES
+        #     None, #推理不需要NUM_PROMPT_TOKENS
+        #     action_head,
+        #     cfg=cfg,
+        # )
+
+        # cfg = kwargs.get("cfg", None) 
+        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+            input_embeddings, projected_patch_embeddings, attention_mask
+        ) 
+        # multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'  
+        # first language_model_output.hidden_states  , is tuple (1 token,  33 layers, torch.Size([1, 314, 4096]))
+        # the following is  (num of tokens,)
+
+        if cfg.mode == "dit": 
+            if cfg.num_actions_chunk//cfg.num_actions_per_token > 1:
+                # token_num = cfg.num_actions_chunk//cfg.num_actions_per_token
+                # language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,min_new_tokens=token_num, max_new_tokens=token_num,output_hidden_states=True,return_dict_in_generate=True)
+                # actions_hidden_states0 = language_model_output.hidden_states[0][-1][:, -1]  # 第一个token的hidden state的last layer
+                
+                # actions_hidden_states_list = [actions_hidden_states0]
+                # for i in range(1, token_num):
+                #     token_hidden_state = language_model_output.hidden_states[i][-1][0]  # i+1个token的hidden state的last layer
+                #     actions_hidden_states_list.append(token_hidden_state)
+                # combined_hidden_states = torch.stack(actions_hidden_states_list, dim=0)  # (16, 4096)
+                # actions_hidden_states = combined_hidden_states.reshape(multimodal_embeddings.shape[0], token_num, multimodal_embeddings.shape[-1])
+                raise NotImplementedError
+            elif cfg.num_actions_chunk//cfg.num_actions_per_token == 1:
+                language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
+                # assert language_model_output.sequences == torch.tensor([[32001]], device=multimodal_embeddings.device)
+                actions_hidden_states = language_model_output.hidden_states[0][-1] 
+                cognition_features = actions_hidden_states[:, -1]
+                assert (cognition_features.shape[0], cognition_features.shape[1]) == (1,4096), "Batch size must be 1 for action prediction"
+                using_cfg = cfg.cfg_scale > 1.0
+
+                model_dtype = next(action_head.net.parameters()).dtype
+                B = cognition_features.shape[0]
+
+                cognition_features = cognition_features.unsqueeze(1).to(model_dtype) 
+                
+                noise = torch.randn(B, cfg.num_actions_chunk, action_head.net.in_channels, device=cognition_features.device).to(model_dtype) 
+
+                # TODO: Setup classifier-free guidance: now use cfg   
+                noise = torch.cat([noise, noise], 0) #  noise.shape torch.Size([2, 16, 7])
+                uncondition = action_head.net.z_embedder.uncondition # torch.Size([1, 4096])
+                uncondition = uncondition.unsqueeze(0)  #[1, D] # torch.Size([1, 1, 4096])
+                uncondition = uncondition.expand(B, 1, -1) #[B, 1, D]
+                z = torch.cat([cognition_features, uncondition], 0) # z  shape torch.Size([2, 1, 4096])
+                model_kwargs = dict(z=z, cfg_scale=cfg.cfg_scale)
+                sample_fn = action_head.net.forward_with_cfg
+                # default use ddim
+                if action_head.ddim_diffusion is None:
+                    action_head.create_ddim(ddim_step=cfg.num_ddim_steps)
+                samples = action_head.ddim_diffusion.ddim_sample_loop(sample_fn, 
+                                                                    noise.shape, 
+                                                                    noise, 
+                                                                    clip_denoised=False,
+                                                                    model_kwargs=model_kwargs,
+                                                                    progress=False,
+                                                                    device=cognition_features.device,
+                                                                    eta=0.0
+                                                                    )
+                if using_cfg:
+                    samples, _ = samples.chunk(2, dim=0)  # Remove null class samples
+                normalized_actions = samples[0].cpu().numpy()
+            else:
+                raise NotImplementedError
+        else:
+            raise NotImplementedError    
+
+
+
+        # normalized_actions = normalized_actions.float().cpu().detach().numpy()
+        actions = self._unnormalize_actions(normalized_actions, unnorm_key)
+
+        return actions, actions_hidden_states

modeling_prismatic.py.back.20250921_183706

@@ -0,0 +1,1553 @@
+"""
+modeling_prismatic.py
+
+Core HuggingFace-style PrismaticPreTrainedModel and PrismaticForConditionalGeneration class definitions.
+Inherits from the default `transformers.PretrainedModel`. Meant to be standalone and self-contained,
+but exactly replicate the logic in `prismatic.models.vlms.prismatic.py`.
+"""
+
+import logging
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import timm
+import tokenizers
+import torch
+import torch.nn as nn
+import transformers
+from timm.models.vision_transformer import LayerScale
+from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+from prismatic.models.action_heads import  L1RegressionActionHead, DiTActionHead, FlowMatchingActionHead
+from prismatic.training.train_utils import (
+    get_current_action_mask,
+    get_next_actions_mask,
+)
+from prismatic.vla.constants import (
+    ACTION_DIM,
+    ACTION_PROPRIO_NORMALIZATION_TYPE,
+    IGNORE_INDEX,
+    NUM_ACTIONS_CHUNK,
+    ACTION_TOKEN_IDX,
+    NormalizationType,
+)
+
+from .configuration_prismatic import OpenVLAConfig, PrismaticConfig
+
+# Set up logger
+logger = logging.getLogger(__name__)
+
+
+# === Utility Functions for Monkey-Patching ===
+def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]:
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        result = fn(*args, **kwargs)
+        return result[0] if isinstance(result, tuple) else result
+
+    return wrapper
+
+
+# HF Transformers overwrites parameters with names containing `gamma`; we're going to patch VisionBackbone.LayerScale.
+#   =>> TIMM :: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L109
+#   =>> Transformers :: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L3960
+def _ls_new_forward(self, x: torch.Tensor) -> torch.Tensor:
+    return x.mul_(self.scale_factor) if self.inplace else x * self.scale_factor
+
+
+def ls_apply_patch(ls_module: LayerScale):
+    ls_module.scale_factor = nn.Parameter(ls_module.gamma.clone())
+    ls_module.forward = _ls_new_forward.__get__(ls_module, LayerScale)
+    del ls_module.gamma
+
+
+# === Prismatic Vision Backbone (nn.Module) Definitions (w/ Fused Backbone Support) ===
+class PrismaticVisionBackbone(nn.Module):
+    """
+    Vision backbone for Prismatic models that handles image feature extraction.
+
+    Supports both single backbone (e.g., SigLIP) and fused backbone (e.g., SigLIP + DINOv2) configurations.
+    For fused backbones, features from both models are concatenated along the feature dimension.
+    """
+
+    def __init__(
+        self,
+        use_fused_vision_backbone: bool,
+        image_sizes: List[int],
+        timm_model_ids: List[str],
+        timm_override_act_layers: List[Optional[str]],
+    ) -> None:
+        """
+        Initialize the vision backbone.
+
+        Args:
+            use_fused_vision_backbone: Whether to use two backbones and fuse their features
+            image_sizes: List of image sizes for each backbone
+            timm_model_ids: List of TIMM model IDs to use for each backbone
+            timm_override_act_layers: List of activation layer overrides for each backbone
+        """
+        super().__init__()
+        self.use_fused_vision_backbone = use_fused_vision_backbone
+        self.num_images_in_input = 1  # Default value, can be overridden later
+
+        # Validate number of (fused) vision backbones
+        if len(timm_model_ids) > 2:
+            raise ValueError("Prismatic models only support up to 2 (fused) vision backbones!")
+
+        # Create primary featurizer
+        self.featurizer = self._create_featurizer(
+            model_id=timm_model_ids[0], img_size=image_sizes[0], act_layer=timm_override_act_layers[0]
+        )
+        self.embed_dim = self.featurizer.embed_dim
+
+        # Create secondary featurizer if using fused backbone
+        if self.use_fused_vision_backbone:
+            self.fused_featurizer = self._create_featurizer(
+                model_id=timm_model_ids[1], img_size=image_sizes[1], act_layer=timm_override_act_layers[1]
+            )
+            self.embed_dim += self.fused_featurizer.embed_dim
+
+        # Patch LayerScale modules for HF compatibility
+        self._patch_layer_scales()
+
+    def _create_featurizer(self, model_id: str, img_size: int, act_layer: Optional[str]) -> nn.Module:
+        """
+        Create a TIMM-based featurizer model with appropriate configurations.
+
+        Args:
+            model_id: The TIMM model ID to load
+            img_size: Input image size for the model
+            act_layer: Override for the activation layer type
+
+        Returns:
+            A configured featurizer model
+        """
+        featurizer = timm.create_model(
+            model_id,
+            pretrained=False,
+            num_classes=0,
+            img_size=img_size,
+            act_layer=act_layer,
+        )
+
+        # Monkey-patch the forward function to extract the second-to-last layer features
+        num_blocks = len(featurizer.blocks)
+        featurizer.forward = unpack_tuple(partial(featurizer.get_intermediate_layers, n={num_blocks - 2}))
+
+        return featurizer
+
+    def _patch_layer_scales(self) -> None:
+        """
+        Patch all LayerScale modules to be compatible with HF's parameter naming.
+
+        HF Transformers overwrites parameters with names containing 'gamma',
+        so we need to rename and modify the forward method.
+        """
+        # Patch primary featurizer
+        for module in self.featurizer.modules():
+            if isinstance(module, LayerScale):
+                ls_apply_patch(module)
+
+        # Patch secondary featurizer if it exists
+        if self.use_fused_vision_backbone:
+            for module in self.fused_featurizer.modules():
+                if isinstance(module, LayerScale):
+                    ls_apply_patch(module)
+
+    def get_num_patches(self) -> int:
+        """
+        Returns the number of vision patches output by the vision backbone.
+
+        Returns:
+            Number of patches per image
+        """
+        return self.featurizer.patch_embed.num_patches
+
+    def get_num_images_in_input(self) -> int:
+        """
+        Returns the number of input images for the vision backbone.
+
+        Returns:
+            Number of images expected in the input
+        """
+        return self.num_images_in_input
+
+    def set_num_images_in_input(self, num_images_in_input: int) -> None:
+        """
+        Sets the number of input images for the vision backbone.
+
+        Args:
+            num_images_in_input: Number of images to expect in the input
+        """
+        self.num_images_in_input = num_images_in_input
+
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        """
+        Implements the forward pass for the vision backbone.
+
+        If `self.use_fused_vision_backbone == True`, uses both SigLIP and DINOv2 transformers to extract visual features
+        (otherwise uses SigLIP only). Allows multi-image inputs (but only for fused vision backbone).
+
+        Args:
+            pixel_values (torch.Tensor): Pixels for input image(s), (B, C, H, W).
+        """
+        if self.num_images_in_input == 1:
+            if not self.use_fused_vision_backbone:
+                return self.featurizer(pixel_values)
+
+            # Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack
+            img, img_fused = torch.split(pixel_values, [3, 3], dim=1)
+            patches, patches_fused = self.featurizer(img), self.fused_featurizer(img_fused)
+
+            return torch.cat([patches, patches_fused], dim=2)
+
+        else:
+            assert self.use_fused_vision_backbone, "Multi-image inputs require using fused backbone!"
+
+            # Split `pixel_values` into individual images (each with 6 channels: 3 for SigLIP + 3 for DINOv2)
+            images = torch.split(pixel_values, [6] * self.num_images_in_input, dim=1)
+
+            # Process each image and collect patches
+            all_patches = []
+            for img in images:
+                # Split each image further into two stacks of channels (each with 3 channels)
+                img_regular, img_fused = torch.split(img, [3, 3], dim=1)
+
+                # Get patches from both SigLIP and DINOv2 vision transformers
+                patches = self.featurizer(img_regular)
+                patches_fused = self.fused_featurizer(img_fused)
+
+                # Concatenate SigLIP and DINOv2 patches along the hidden dimension
+                combined_patches = torch.cat([patches, patches_fused], dim=2)
+                all_patches.append(combined_patches)
+
+            # Concatenate all patches along the patch dimension
+            return torch.cat(all_patches, dim=1)
+
+
+# === Prismatic Projector (nn.Module) Definitions ===
+class PrismaticProjector(nn.Module):
+    def __init__(self, use_fused_vision_backbone: bool, vision_dim: int, llm_dim: int) -> None:
+        super().__init__()
+        self.use_fused_vision_backbone = use_fused_vision_backbone
+        self.vision_dim, self.llm_dim = vision_dim, llm_dim
+
+        # Switch on `use_fused_vision_backbone` =>> use slightly different MLPs and projection factors!
+        if not self.use_fused_vision_backbone:
+            self.fc1 = nn.Linear(self.vision_dim, self.llm_dim, bias=True)
+            self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
+            self.act_fn1 = nn.GELU()
+        else:
+            initial_projection_dim = 4 * vision_dim
+            self.fc1 = nn.Linear(self.vision_dim, initial_projection_dim, bias=True)
+            self.fc2 = nn.Linear(initial_projection_dim, self.llm_dim, bias=True)
+            self.fc3 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
+            self.act_fn1 = nn.GELU()
+            self.act_fn2 = nn.GELU()
+
+    def forward(self, img_patches: torch.Tensor) -> torch.Tensor:
+        if not self.use_fused_vision_backbone:
+            projected_features = self.fc1(img_patches)
+            projected_features = self.act_fn1(projected_features)
+            projected_features = self.fc2(projected_features)
+        else:
+            projected_features = self.fc1(img_patches)
+            projected_features = self.act_fn1(projected_features)
+            projected_features = self.fc2(projected_features)
+            projected_features = self.act_fn2(projected_features)
+            projected_features = self.fc3(projected_features)
+
+        return projected_features
+
+
+# === Main HF Class Definitions ===
+@dataclass
+class PrismaticCausalLMOutputWithPast(ModelOutput):
+    """Base class for Prismatic casual (visually-conditioned) language model outputs; also exposes visual features."""
+
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+    # Additions for VLMs
+    projector_features: Optional[torch.FloatTensor] = None
+
+
+class PrismaticPreTrainedModel(PreTrainedModel):
+    config_class: PretrainedConfig = PrismaticConfig
+    base_model_prefix: str = "model"
+    supports_gradient_checkpointing: bool = True
+
+    _no_split_modules: ClassVar[List[str]] = ["PrismaticProjector"]
+    _skip_keys_device_placement: str = "past_key_values"
+    _supports_flash_attn_2: bool = True
+
+    def _init_weights(self, module: nn.Module) -> None:
+        # Important :: this HF ported version is *not* meant for training from scratch; only inference and fine-tuning!
+        #   => As such, this init_weights code is not correct; if training VLMs from scratch, use the main codebase at
+        #      https://github.com/TRI-ML/prismatic-vlms
+        std = (
+            self.config.initializer_range
+            if hasattr(self.config, "initializer_range")
+            else self.config.text_config.initializer_range
+        )
+
+        if hasattr(module, "class_embedding"):
+            module.class_embedding.data.normal_(mean=0.0, std=std)
+
+        if isinstance(module, (nn.Linear, nn.Conv2d)):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+    @property
+    def _supports_sdpa(self) -> bool:
+        """Check LLM supports SDPA Attention"""
+        return self.language_model._supports_sdpa
+
+
+class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
+    def __init__(self, config: PrismaticConfig) -> None:
+        super().__init__(config)
+
+        # [Validation] Lightweight Validate on `config` Fields + Dependency Versions
+        if config.use_fused_vision_backbone is None:
+            raise ValueError("Missing config field `use_fused_vision_backbone`")
+
+        if timm.__version__ not in {"0.9.10", "0.9.11", "0.9.12", "0.9.16"}:
+            raise NotImplementedError(
+                "TIMM Version must be >= 0.9.10 and < 1.0.0 (breaking); please raise a GitHub Issue "
+                "if you urgently need support for latest TIMM versions."
+            )
+
+        if (transformers.__version__ != "4.40.1") or (tokenizers.__version__ != "0.19.1"):
+            logger.warning(
+                f"Expected `transformers==4.40.1` and `tokenizers==0.19.1` but got "
+                f"`transformers=={transformers.__version__}` and `tokenizers=={tokenizers.__version__}`; "
+                f"there might be inference-time regressions due to dependency changes. If in doubt, please"
+                f"use the above versions."
+            )
+
+        # Instantiate PrismaticVisionBackbone (w/ Potential Fused Backbone)
+        self.vision_backbone = PrismaticVisionBackbone(
+            config.use_fused_vision_backbone, config.image_sizes, config.timm_model_ids, config.timm_override_act_layers
+        )
+
+        # Create Multimodal Projector
+        self.projector = PrismaticProjector(
+            config.use_fused_vision_backbone,
+            vision_dim=self.vision_backbone.embed_dim,
+            llm_dim=config.text_config.hidden_size,
+        )
+
+        # Instantiate LLM Backbone
+        self.language_model = AutoModelForCausalLM.from_config(
+            config.text_config, attn_implementation=config._attn_implementation
+        )
+        self.vocab_size = config.text_config.vocab_size
+        self.pad_token_id = config.pad_token_id
+        self.llm_dim = config.text_config.hidden_size
+
+        # HF Boilerplate =>> initializes weights via `_init_weights()` and sets gradient checkpointing
+        self.post_init()
+
+    # === `PreTrainedModel` Boilerplate ===
+    def get_input_embeddings(self) -> nn.Module:
+        return self.language_model.get_input_embeddings()
+
+    def set_input_embeddings(self, value: nn.Module) -> None:
+        self.language_model.set_input_embeddings(value)
+
+    def get_output_embeddings(self) -> nn.Module:
+        return self.language_model.get_output_embeddings()
+
+    def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
+        self.language_model.set_output_embeddings(new_embeddings)
+
+    def get_decoder(self) -> nn.Module:
+        return self.language_model.get_decoder()
+
+    def set_decoder(self, decoder: nn.Module) -> None:
+        self.language_model.set_decoder(decoder)
+
+    def tie_weights(self) -> None:
+        self.language_model.tie_weights()  # Note: `Llama-2` and `Mistral` don't tie weights (no-op)
+
+    def resize_token_embeddings(
+        self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None
+    ) -> nn.Embedding:
+        updated_embeddings = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
+
+        # Update config/instance variables
+        self.config.text_config.vocab_size = updated_embeddings.num_embeddings
+        self.vocab_size = updated_embeddings.num_embeddings
+
+        return updated_embeddings
+
+    def _replace_input_embeddings(self, input_embeddings, all_actions_mask, noisy_action_features):
+        """
+        Replace embeddings in input_embeddings at positions where all_actions_mask is True
+        with embeddings from noisy_action_features, using vectorized operations.
+
+        Args:
+            input_embeddings: Tensor of shape (B, S, D)
+            all_actions_mask: Boolean tensor of shape (B, S)
+            noisy_action_features: Tensor of shape (B, K, D) where K is the number of True values in mask per sample
+
+        Returns:
+            Modified input_embeddings tensor
+        """
+        # Clone input to avoid modifying the original tensor
+        new_input_embeddings = input_embeddings.clone()
+
+        # Create a tensor with the same shape of input_embeddings to hold the noisy action features
+        repositioned_noisy_action_features = torch.zeros_like(input_embeddings)
+
+        # Create batch indices for splicing
+        batch_indices = torch.arange(input_embeddings.shape[0], device=input_embeddings.device)
+        batch_indices = batch_indices.unsqueeze(1).expand(-1, noisy_action_features.shape[1])
+
+        # Get indices where mask is True for each sample
+        masked_indices = torch.stack([torch.where(mask)[0] for mask in all_actions_mask])
+
+        # Move the noisy action features into their correct positions
+        repositioned_noisy_action_features[batch_indices, masked_indices] = noisy_action_features
+
+        # Combine original input embeddings and noisy action embeddings using the mask
+        new_input_embeddings = torch.where(
+            all_actions_mask.unsqueeze(-1), repositioned_noisy_action_features, new_input_embeddings
+        )
+
+        return new_input_embeddings
+
+    def _process_action_masks(self, labels):
+        """Helper to get action masks from labels"""
+        current_action_mask = get_current_action_mask(labels) # (B, seq_len)
+        next_actions_mask = get_next_actions_mask(labels) # (B, seq_len)
+        all_actions_mask = current_action_mask | next_actions_mask  # (B, seq_len)
+        return all_actions_mask
+
+    def _process_vision_features(self, pixel_values):
+        """Process vision features with optional FiLM conditioning"""
+        patch_features = self.vision_backbone(pixel_values)  # (bsz, 256 * num_images, D)
+
+        # Project patch embeddings into language embedding space
+        return self.projector(patch_features)
+
+    def _process_proprio_features(self, projected_patch_embeddings, proprio, proprio_projector):
+        """Process proprioceptive features and append to vision features"""
+        if proprio_projector is not None and proprio is not None:
+            # projected_patch_embeddings: (bsz, num_patches * num_images, llm_dim)
+            # proprio: (bsz, proprio_dim) or (propro_dim,)
+            proprio = proprio.reshape(projected_patch_embeddings.shape[0], -1)  # (bsz, proprio_dim)
+            proprio_features = proprio_projector(proprio)  # (bsz, llm_dim)
+            proprio_features = proprio_features.unsqueeze(dim=1)  # (bsz, 1, llm_dim)
+            # For simplicity, just append proprio token to the end of projected vision patch tokens
+            return torch.cat((projected_patch_embeddings, proprio_features), dim=1)
+        return projected_patch_embeddings
+
+    def _build_multimodal_attention(self, input_embeddings, projected_patch_embeddings, attention_mask):
+        """Build multimodal embeddings and attention mask"""
+        # juyi: Update attention mask 是不是要改成下三角? 不用, 因为generate会自动屏蔽
+        projected_patch_attention_mask = None
+        if attention_mask is not None:
+            projected_patch_attention_mask = torch.full(
+                (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
+                fill_value=True,
+                dtype=attention_mask.dtype,
+                device=attention_mask.device,
+            )
+
+        # Build multimodal embeddings & attention mask; insert embeddings after <BOS> token (1:)
+        multimodal_embeddings = torch.cat(
+            [input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1
+        )
+
+        multimodal_attention_mask = None
+        if attention_mask is not None:
+            multimodal_attention_mask = torch.cat(
+                [attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1
+            )
+
+        return multimodal_embeddings, multimodal_attention_mask
+
+    def _build_multimodal_labels(self, labels, projected_patch_embeddings):
+        """Build multimodal labels with IGNORE_INDEX for patch embeddings"""
+        if labels is not None:
+            projected_patch_labels = torch.full(
+                (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
+                fill_value=IGNORE_INDEX, # 这些位置不需要计算损失。
+                dtype=labels.dtype,
+                device=labels.device,
+            )
+            return torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1) # 第一个token是<BOS>
+        return None
+
+    # === Core Prismatic VLM `forward()` Logic ===
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_projector_features: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        proprio=None,
+        proprio_projector=None,
+        noisy_actions=None,
+        noisy_action_projector=None,
+        diffusion_timestep_embeddings=None,
+        use_film: bool = False,
+    ) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
+        """Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
+        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
+        )
+        output_projector_features = output_projector_features if output_projector_features is not None else False
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
+        use_cache = use_cache and not self.training
+
+        # Instantiate Placeholder for Projector Features
+        projected_patch_embeddings = None
+
+        # === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
+        if input_ids.shape[1] == 1:
+            assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
+            assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
+            assert labels is None, "Unexpected key `labels` provided during cached generation!"
+
+            language_model_output = self.language_model(
+                input_ids=input_ids,
+                attention_mask=None,
+                position_ids=None,
+                past_key_values=past_key_values,
+                inputs_embeds=None,
+                labels=None,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+        # === Handle Unimodal Forward ===
+        elif pixel_values is None:
+            assert (input_ids is not None) and (inputs_embeds is None), "Missing `input_ids` in language-only forward!"
+            assert past_key_values is None, "Unexpected key `past_key_values` provided during language-only forward!"
+
+            language_model_output = self.language_model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                position_ids=None,
+                past_key_values=None,
+                inputs_embeds=None,
+                labels=labels,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+        # === Handle Multimodal Forward ===
+        elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
+            assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
+
+            # Get input embeddings (from language model embeddings)
+            input_embeddings = self.get_input_embeddings()(input_ids)  # (B, seq_len, D)
+
+            # Extract action masks
+            all_actions_mask = self._process_action_masks(labels)
+
+            # Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
+            language_embeddings = input_embeddings[~all_actions_mask].reshape(
+                input_embeddings.shape[0], -1, input_embeddings.shape[2]
+            )  # (B, lang_seq_len, llm_dim)
+
+            # Get visual features
+            projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
+            # bug: TypeError: PrismaticForConditionalGeneration._process_vision_features() takes 2 positional arguments but 4 were given
+
+            # Add proprioceptive state if provided
+            projected_patch_embeddings = self._process_proprio_features(
+                projected_patch_embeddings, proprio, proprio_projector
+            )
+
+            # [Diffusion] Add diffusion timestep embedding if provided
+            if diffusion_timestep_embeddings is not None:
+                # For simplicity, just append diffusion timestep embedding to the end of projected vision patch tokens
+                projected_patch_embeddings = torch.cat(
+                    (projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
+                )
+                
+            # Process action embeddings
+            if noisy_actions is not None:
+                # Get mask corresponding to all action tokens
+                all_actions_mask = self._process_action_masks(labels)
+
+                # Reshape noisy actions into individual action tokens
+                # noisy_actions: (B, chunk_len, action_dim) -> (B, chunk_len * action_dim, 1)
+                B = noisy_actions.shape[0]
+                noisy_actions = noisy_actions.reshape(B, -1).unsqueeze(-1)
+
+                # Project noisy action tokens into language model embedding space
+                noisy_action_features = noisy_action_projector(noisy_actions)  # (B, chunk_len * action_dim, llm_dim)
+
+                # Replace embeddings of the action tokens with noisy action embeddings
+                input_embeddings = self._replace_input_embeddings(
+                    input_embeddings, all_actions_mask, noisy_action_features
+                )
+            else:
+                # Replace the embeddings of the action tokens with zeros
+                # (Later on, the positional embeddings will be added to them)
+                all_actions_mask = all_actions_mask.unsqueeze(-1)  # (B, seq_len, 1)
+                input_embeddings = input_embeddings * ~all_actions_mask
+
+            # Build multimodal embeddings & attention mask
+            multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+                input_embeddings, projected_patch_embeddings, attention_mask
+            )
+
+            # Build labels for multimodal sequence if needed
+            multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
+
+            # Dispatch to language model
+            language_model_output = self.language_model(
+                input_ids=None,
+                attention_mask=multimodal_attention_mask,
+                position_ids=None,
+                past_key_values=None,
+                inputs_embeds=multimodal_embeddings,
+                labels=multimodal_labels,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+        # === Otherwise =>> Assume Invalid! ===
+        elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
+            raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
+
+        else:
+            raise ValueError(
+                "Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
+                f"=> `input_ids` = {input_ids is not None}\n"
+                f"=> `attention_mask` = {attention_mask is not None}\n"
+                f"=> `pixel_values` = {pixel_values is not None}\n"
+                f"=> `labels` = {labels is not None}\n"
+                f"=> `input_embeds` = {inputs_embeds is not None}\n"
+                f"=> `past_key_values` = {past_key_values is not None}\n"
+                f"=> `use_cache` = {use_cache}"
+            )
+
+        # Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
+        if not return_dict:
+            if output_projector_features and (projected_patch_embeddings is not None):
+                return *language_model_output, projected_patch_embeddings
+
+            return language_model_output
+
+        return PrismaticCausalLMOutputWithPast(
+            loss=language_model_output.loss,
+            logits=language_model_output.logits,
+            past_key_values=language_model_output.past_key_values,
+            hidden_states=language_model_output.hidden_states,
+            attentions=language_model_output.attentions,
+            projector_features=projected_patch_embeddings,
+        )
+
+    # === GenerationMixin Methods ===
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs: str,
+    ) -> Dict[str, torch.Tensor]:
+        """Borrowed from `LlamaForCausalLM` and simplified for batch size = 1; mirrors original PrismaticVLM logic."""
+        if ((input_ids is not None) and (input_ids.shape[0] > 1)) or (
+            (inputs_embeds is not None) and (inputs_embeds.shape[0] > 1)
+        ):
+            raise ValueError("Generation with batch size > 1 is not currently supported!")
+
+        # Handle `past_key_values` (cache) =>> assume `input_ids` just has unprocessed tokens
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1:]
+
+        # If `input_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 = {"input_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        # Make sure `pixel_values` are preserved in `model_inputs`
+        model_inputs.update(
+            {
+                "attention_mask": attention_mask,
+                "pixel_values": pixel_values,
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+            }
+        )
+
+        return model_inputs
+
+    # Defer to Language Model (all handle this differently, with different return types)
+    def _reorder_cache(self, *args, **kwargs) -> Any:
+        return self.language_model._reorder_cache(*args, **kwargs)
+
+
+class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
+    config_class: PretrainedConfig = OpenVLAConfig
+
+    def __init__(self, config: OpenVLAConfig) -> None:
+        super().__init__(config)
+        self.norm_stats = config.norm_stats
+
+        # Compute action bins
+        self.bins = np.linspace(-1, 1, config.n_action_bins)
+        self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0
+
+        # Compute vocab size for de-tokenization -- revert added "multiple of"
+        self.vocab_size = self.config.text_config.vocab_size - self.config.pad_to_multiple_of
+
+    def _prepare_input_for_action_prediction(self, input_ids, attention_mask):
+        # eval 会用到这里
+        """Prepares input for action prediction by adding necessary tokens"""
+        # Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens
+        placeholder_action_token_ids = (
+            torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
+        ) 
+        input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1) # torch.Size([1,  35 + 56= 91])
+
+        # Extend the attention mask to fit the new shape of input
+        # Note: Only batch size == 1 supported right now
+        mask_extension = (
+            torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
+            .to(attention_mask.device)
+            .to(attention_mask.dtype)
+        )
+        attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
+
+        return input_ids, attention_mask
+
+    def _prepare_labels_for_action_prediction(self, labels, input_ids):
+        """Creates labels tensor for action prediction if not provided"""
+        # eval 会用到这里 , 
+        # Extends label tensors with fake action labels
+        # Adds stop tokens at the end of sequences
+        # Handles label preparation for action prediction tasks
+        # 他为啥可以随便一个?  xuan说 你自定义一个值 ,然后一直指定这个 , PAD token可以吗? 
+        #TODO: 这里是否要改?  感觉不需要改. 随便写就行了因为labels不重要只是要一个mask. 为什么需要这个函数? 确保 action 预测任务的标签（labels）符合模型的输入长度，并正确地处理序列终止
+        # Extend labels tensor with fake action labels
+        ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_IDX # = 为了mask正确生成, action_tokens_only_mask = (labels == ACTION_TOKEN_IDX ), 所以这里也填上ACTION_TOKEN_IDX
+        labels_extension = (
+            torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
+            * ARBITRARY_ACTION_TOKEN_IDX
+        ) #torch.Size([1, 57]),全是 ARBITRARY_ACTION_TOKEN_IDX
+        labels = torch.cat([labels, labels_extension], dim=-1)
+
+        return labels
+
+    def _unnormalize_actions(self, normalized_actions, unnorm_key=None):
+        """Unnormalize actions using dataset statistics"""
+        action_norm_stats = self.get_action_stats(unnorm_key)
+
+        if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:
+            mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))
+            action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])
+        elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:
+            mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
+            action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
+        else:
+            raise ValueError("Unsupported action/proprio normalization type detected!")
+
+        actions = np.where(
+            mask,
+            0.5 * (normalized_actions + 1) * (action_high - action_low + 1e-8) + action_low,
+            normalized_actions,
+        )
+
+        return actions
+
+
+    def _normalize_actions(self, actions, norm_key=None):
+        """Normalize actions to [-1, 1] using dataset statistics"""
+        action_norm_stats = self.get_action_stats(norm_key)
+
+        if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:
+            mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))
+            action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])
+        elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:
+            mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
+            action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
+        else:
+            raise ValueError("Unsupported action/proprio normalization type detected!")
+
+        normalized = np.where(
+            mask,
+            2 * (actions - action_low) / (action_high - action_low + 1e-8) - 1,
+            actions,
+        )
+
+        return normalized
+
+    def _run_diffusion_prediction(
+        self,
+        input_embeddings,
+        all_actions_mask,
+        noise,
+        action_head,
+        projected_patch_embeddings,
+        labels,
+        attention_mask,
+        NUM_PATCHES,
+        NUM_PROMPT_TOKENS,
+        noisy_action_projector,
+    ):
+        """Run diffusion-based action prediction"""
+        # Set diffusion timestep values
+        action_head.noise_scheduler.set_timesteps(action_head.num_diffusion_steps)
+        # Clone embedding for reuse in each timestep
+        orig_projected_patch_embeddings = projected_patch_embeddings.clone()
+        curr_noisy_actions = noise
+
+        # Reverse diffusion: Iteratively denoise to generate action prediction
+        for t in action_head.noise_scheduler.timesteps:
+            # Get diffusion model's noise prediction (conditioned on VLA latent embedding, current noisy action
+            # embedding, and diffusion timestep embedding)
+            timesteps = torch.Tensor([t]).to(labels.device)
+            diffusion_timestep_embeddings = (
+                action_head.time_encoder(timesteps).to(curr_noisy_actions.dtype).to(curr_noisy_actions.device)
+            )  # (B, llm_dim)
+            diffusion_timestep_embeddings = diffusion_timestep_embeddings.unsqueeze(1)  # (B, 1, llm_dim)
+
+            # [Diffusion] Replace the embeddings of the action tokens with noisy actions
+            # (Later on, the positional embeddings will be added to them)
+
+            # For simplicity, append diffusion timestep embedding to the end of projected vision tokens
+            projected_patch_embeddings = torch.cat(
+                (orig_projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
+            )
+
+            # Reshape and project noisy actions into language embedding space
+            B = curr_noisy_actions.shape[0]
+            orig_curr_noisy_actions_shape = curr_noisy_actions.shape
+            curr_noisy_actions = curr_noisy_actions.reshape(B, -1).unsqueeze(-1)
+            noisy_action_features = noisy_action_projector(curr_noisy_actions)
+            curr_noisy_actions = curr_noisy_actions.reshape(orig_curr_noisy_actions_shape)
+
+            # Replace action token embeddings with noisy action embeddings
+            input_embeddings = self._replace_input_embeddings(
+                input_embeddings.clone(), all_actions_mask, noisy_action_features
+            )
+
+            # Build multimodal embeddings and attention mask
+            multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+                input_embeddings, projected_patch_embeddings, attention_mask
+            )
+
+            # Forward pass through language model
+            language_model_output = self.language_model(
+                input_ids=None,
+                attention_mask=multimodal_attention_mask,
+                position_ids=None,
+                past_key_values=None,
+                inputs_embeds=multimodal_embeddings,
+                labels=None,
+                use_cache=None,
+                output_attentions=False,
+                output_hidden_states=True,
+                return_dict=True,
+            )
+
+            # Extract hidden states for action portion of response
+            last_hidden_states = language_model_output.hidden_states[-1]  # (B, seq_len, D)
+            actions_hidden_states = last_hidden_states[
+                :,
+                NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
+                :,
+            ]  # (B, act_chunk_len, D)
+
+            # Predict noise and update noisy actions: x_t -> x_{t-1}
+            noise_pred = action_head.predict_noise(actions_hidden_states)
+            curr_noisy_actions = action_head.noise_scheduler.step(noise_pred, t, curr_noisy_actions).prev_sample
+
+        curr_noisy_actions = curr_noisy_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
+
+        # Return final actions
+        return curr_noisy_actions.float().cpu().detach().numpy(), actions_hidden_states
+
+    def _regression_or_discrete_prediction(
+        self,
+        input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
+        all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action . 
+        projected_patch_embeddings: torch.FloatTensor, 
+        attention_mask: torch.BoolTensor,
+        labels: torch.LongTensor,
+        NUM_PATCHES: int,
+        NUM_PROMPT_TOKENS: int,
+        action_head: L1RegressionActionHead,
+        **kwargs,
+    ):
+        """Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
+        # Extract hidden states for action tokens
+        # last_hidden_states = language_model_output.hidden_states[-1]  # (B, seq_len, D)
+
+        # actions_hidden_states = last_hidden_states[:, NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + NUM_ACTIONS_CHUNK * tokennum, :]# (B, act_chunk_len, D)
+        # 都不需要取了, 直接就给 token对应的hidden state了 ,太方便了. 
+        # 为什么第一个是torch.Size([1, 535, 4096])? 我应该选哪个?  https://discuss.huggingface.co/t/get-each-generated-token-last-layer-hidden-state/145921
+        # language_model_output.sequences tensor([[29871, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 32001,     2]], device='cuda:0')
+        cfg = kwargs.pop("cfg", None)
+        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+            input_embeddings, projected_patch_embeddings, attention_mask
+        ) 
+        # multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'  
+        language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
+        #  is tuple (1 token,  33 layers, torch.Size([1, 314, 4096]))
+        hidden_states = language_model_output.hidden_states[0][-1]
+        actions_hidden_states = hidden_states[:, -NUM_ACTIONS_CHUNK:]
+
+        normalized_actions = action_head.predict_action(actions_hidden_states)
+        normalized_actions = normalized_actions.reshape(cfg.num_actions_chunk, ACTION_DIM)
+        normalized_actions = normalized_actions.float().cpu().detach().numpy()
+        if cfg.mode == "mul": 
+            if cfg.num_actions_chunk//cfg.num_actions_per_token > 1:
+                token_num = cfg.num_actions_chunk//cfg.num_actions_per_token
+                language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,min_new_tokens=token_num, max_new_tokens=token_num,output_hidden_states=True,return_dict_in_generate=True)
+                actions_hidden_states0 = language_model_output.hidden_states[0][-1][:, -1]  # 第一个token的hidden state的last layer
+                
+                actions_hidden_states_list = [actions_hidden_states0]
+                for i in range(1, token_num):
+                    token_hidden_state = language_model_output.hidden_states[i][-1][0]  # i+1个token的hidden state的last layer
+                    actions_hidden_states_list.append(token_hidden_state)
+                # 将所有hidden states拼接起来
+                combined_hidden_states = torch.stack(actions_hidden_states_list, dim=0)  # (16, 4096)
+                actions_hidden_states = combined_hidden_states.reshape(multimodal_embeddings.shape[0], token_num, multimodal_embeddings.shape[-1])
+            elif cfg.num_actions_chunk//cfg.num_actions_per_token == 1:
+                language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
+                # assert language_model_output.sequences == torch.tensor([[32001]], device=multimodal_embeddings.device)
+                actions_hidden_states = language_model_output.hidden_states[0][-1] 
+                actions_hidden_states = actions_hidden_states[:, -1]
+            else:
+                raise NotImplementedError
+        else:
+            raise NotImplementedError    
+        return normalized_actions, actions_hidden_states
+
+    def hist_predict_action(
+        self,
+        input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
+        all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action . 
+        projected_patch_embeddings: torch.FloatTensor, 
+        attention_mask: torch.BoolTensor,
+        action_head: L1RegressionActionHead,
+        **kwargs,
+    ):
+        cfg = kwargs.get("cfg", None) 
+        action_history = kwargs.get("action_history", None)
+        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+            input_embeddings, projected_patch_embeddings, attention_mask
+        ) 
+        # multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'  
+        # first language_model_output.hidden_states  , is tuple (1 token,  33 layers, torch.Size([1, 314, 4096]))
+        # the following is  (num of tokens,)
+        if cfg.mode == "mul": 
+            if cfg.num_actions_chunk//cfg.num_actions_per_token > 1:
+                raise NotImplementedError
+                # token_num = cfg.num_actions_chunk//cfg.num_actions_per_token
+                # language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,min_new_tokens=token_num, max_new_tokens=token_num,output_hidden_states=True,return_dict_in_generate=True)
+                # actions_hidden_states0 = language_model_output.hidden_states[0][-1][:, -1]  # 第一个token的hidden state的last layer
+                # actions_hidden_states_list = [actions_hidden_states0]
+                # for i in range(1, token_num):
+                #     token_hidden_state = language_model_output.hidden_states[i][-1][0]  # i+1个token的hidden state的last layer
+                #     actions_hidden_states_list.append(token_hidden_state)
+                # combined_hidden_states = torch.stack(actions_hidden_states_list, dim=0)  # (16, 4096)
+                # actions_hidden_states = combined_hidden_states.reshape(multimodal_embeddings.shape[0], token_num, multimodal_embeddings.shape[-1])
+            elif cfg.num_actions_chunk//cfg.num_actions_per_token == 1:
+                language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
+                # assert language_model_output.sequences == torch.tensor([[32001]], device=multimodal_embeddings.device)
+                actions_hidden_states = language_model_output.hidden_states[0][-1] 
+                actions_hidden_states = actions_hidden_states[:, -1]
+                # 在中间加一个 1 维度
+                actions_hidden_states = actions_hidden_states.unsqueeze(1) # for match 3 dim 
+            else:
+                raise NotImplementedError
+        else:
+            raise NotImplementedError    
+
+        normalized_actions = action_head.predict_action(actions_hidden_states, action_history)
+        normalized_actions = normalized_actions.reshape(cfg.num_actions_chunk, ACTION_DIM)
+        normalized_actions = normalized_actions.float().cpu().detach().numpy()
+
+        return normalized_actions, actions_hidden_states
+
+    def mul_regression_or_discrete_prediction(
+        self,
+        input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
+        all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action . 
+        projected_patch_embeddings: torch.FloatTensor, 
+        attention_mask: torch.BoolTensor,
+        labels: torch.LongTensor,
+        NUM_PATCHES: int,
+        NUM_PROMPT_TOKENS: int,
+        action_head: L1RegressionActionHead,
+        **kwargs,
+    ):
+        cfg = kwargs.get("cfg", None) 
+        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+            input_embeddings, projected_patch_embeddings, attention_mask
+        ) 
+        # multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'  
+        # first language_model_output.hidden_states  , is tuple (1 token,  33 layers, torch.Size([1, 314, 4096]))
+        # the following is  (num of tokens,)
+        if cfg.mode == "mul": 
+            if cfg.num_actions_chunk//cfg.num_actions_per_token > 1:
+                token_num = cfg.num_actions_chunk//cfg.num_actions_per_token
+                language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,min_new_tokens=token_num, max_new_tokens=token_num,output_hidden_states=True,return_dict_in_generate=True)
+                actions_hidden_states0 = language_model_output.hidden_states[0][-1][:, -1]  # 第一个token的hidden state的last layer
+                
+                actions_hidden_states_list = [actions_hidden_states0]
+                for i in range(1, token_num):
+                    token_hidden_state = language_model_output.hidden_states[i][-1][0]  # i+1个token的hidden state的last layer
+                    actions_hidden_states_list.append(token_hidden_state)
+                combined_hidden_states = torch.stack(actions_hidden_states_list, dim=0)  # (16, 4096)
+                actions_hidden_states = combined_hidden_states.reshape(multimodal_embeddings.shape[0], token_num, multimodal_embeddings.shape[-1])
+            elif cfg.num_actions_chunk//cfg.num_actions_per_token == 1:
+                language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
+                actions_hidden_states = language_model_output.hidden_states[0][-1] 
+                actions_hidden_states = actions_hidden_states[:, -1]
+            else:
+                raise NotImplementedError
+        else:
+            raise NotImplementedError    
+
+        normalized_actions = action_head.predict_action(actions_hidden_states)
+        normalized_actions = normalized_actions.reshape(cfg.num_actions_chunk, ACTION_DIM)
+        # print(f"*** normalized_actions[]: {normalized_actions} ***")
+        if cfg.action_head_name == "medusa":
+            normalized_actions[:, 6] = torch.sigmoid(normalized_actions[:,  6]) # without bs dim.
+        # print(f"*** normalized_actions[]: {normalized_actions} ***")
+        normalized_actions = normalized_actions.float().cpu().detach().numpy()
+
+        return normalized_actions, actions_hidden_states
+
+    def predict_action(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        unnorm_key: Optional[str] = None,
+        proprio=None,
+        proprio_projector=None,
+        action_head=None,
+        noisy_action_projector=None,
+        use_film: bool = False,
+        **kwargs: str,
+    ) -> np.ndarray:
+        """Predict actions from input sequence, with options for different prediction methods.
+
+        Args:
+            input_ids: Input token ids
+            unnorm_key: Key for unnormalization statistics
+            proprio: Proprioceptive features
+            proprio_projector: Projector for proprioceptive features
+            action_head: Optional head for L1 regression or diffusion-based prediction
+            noisy_action_projector: Projector for noisy actions in diffusion-based prediction
+            use_film: Whether to use FiLM conditioning
+            **kwargs: Additional arguments including pixel_values and attention_mask
+
+        Returns:
+            Tuple of (unnormalized_actions, action_hidden_states)
+        """
+        # If the special empty token ('') does not already appear after the colon (':') token in the prompt
+        # (after "OUT:" or "ASSISTANT:"), insert it to match the inputs seen at training time
+        if not torch.all(input_ids[:, -1] == 29871):
+            input_ids = torch.cat(
+                (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
+            )
+
+        pixel_values = kwargs["pixel_values"]
+        attention_mask = kwargs["attention_mask"]
+
+        # Create fake labels tensor (needed for action mask)
+        labels = input_ids.clone()
+        labels[:] = IGNORE_INDEX # 输入都ignore IGNORE_INDEX = -100
+
+        # Get number of tokens in prompt (excluding the start token)
+        NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1  # Subtract action tokens and stop token
+
+        # Prepare inputs by adding necessary tokens
+        input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)
+
+        # Update labels tensor for action mask computation later
+        labels = self._prepare_labels_for_action_prediction(labels, input_ids)
+
+        # Get input embeddings and action masks
+        input_embeddings = self.get_input_embeddings()(input_ids)
+        all_actions_mask = self._process_action_masks(labels)
+
+        # Extract language embeddings
+        language_embeddings = input_embeddings[~all_actions_mask].reshape(
+            input_embeddings.shape[0], -1, input_embeddings.shape[2]
+        )
+
+        # Process vision features
+        projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
+
+        # Add proprioceptive features if provided
+        use_proprio = proprio_projector is not None and proprio is not None
+        if use_proprio:
+            proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
+            projected_patch_embeddings = self._process_proprio_features(
+                projected_patch_embeddings, proprio, proprio_projector
+            )
+
+        # Use diffusion if provided, otherwise use regression or discrete prediction
+        use_diffusion = noisy_action_projector is not None and hasattr(action_head, "noise_scheduler")
+
+        # Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
+        NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
+        if use_proprio:
+            NUM_PATCHES += 1
+
+        normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
+            input_embeddings,
+            all_actions_mask,
+            projected_patch_embeddings,
+            attention_mask,
+            labels,
+            NUM_PATCHES,
+            NUM_PROMPT_TOKENS,
+            action_head,
+        )
+
+        # Unnormalize predicted actions
+        actions = self._unnormalize_actions(normalized_actions, unnorm_key)
+
+        return actions, actions_hidden_states
+
+    @staticmethod
+    def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_key: Optional[str]) -> str:
+        """Validate and resolve the unnormalization key for action statistics"""
+        if unnorm_key is None:
+            assert len(norm_stats) == 1, (
+                f"Your model was trained on more than one dataset, "
+                f"please pass a `unnorm_key` from the following options to choose the statistics "
+                f"used for un-normalizing actions: {norm_stats.keys()}"
+            )
+            unnorm_key = next(iter(norm_stats.keys()))
+        # norm states没有加载libero, 为什么? 
+        assert unnorm_key in norm_stats, (
+            f"The `unnorm_key` you chose is not in the set of available dataset statistics, "
+            f"please choose from: {norm_stats.keys()}"
+        )
+        return unnorm_key
+
+    def get_action_dim(self, unnorm_key: Optional[str] = None) -> int:
+        """Get the dimensionality of the policy's action space."""
+        unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
+        return len(self.norm_stats[unnorm_key]["action"]["min"])
+
+    def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[str, Any]:
+        """Get all the logged statistics for the given dataset."""
+        unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
+        return self.norm_stats[unnorm_key]["action"]
+
+
+    def lisa_forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_projector_features: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        proprio=None,
+        proprio_projector=None,
+        **kwargs
+    ) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
+        """Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
+        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
+        )
+        output_projector_features = output_projector_features if output_projector_features is not None else False
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
+        use_cache = use_cache and not self.training
+
+        # Instantiate Placeholder for Projector Features
+        projected_patch_embeddings = None
+
+        # === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
+        if input_ids.shape[1] == 1:
+            assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
+            assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
+            assert labels is None, "Unexpected key `labels` provided during cached generation!"
+
+            language_model_output = self.language_model(
+                input_ids=input_ids,
+                attention_mask=None,
+                position_ids=None,
+                past_key_values=past_key_values,
+                inputs_embeds=None,
+                labels=None,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+        # === Handle Unimodal Forward ===
+        elif pixel_values is None:
+            raise NotImplementedError
+
+        # === Handle Multimodal Forward ===
+        elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
+            assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
+
+            # Get input embeddings (from language model embeddings)
+            input_embeddings = self.get_input_embeddings()(input_ids)  # (B, seq_len, D)
+            # Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
+            # language_embeddings = input_embeddings[~all_actions_mask].reshape(
+            #     input_embeddings.shape[0], -1, input_embeddings.shape[2]
+            # )  # (B, lang_seq_len, llm_dim) 这里就会把结尾的 stop index和padding 也算进去. 没问题吗? 没问题因为ignore了 我直接删了因为不用film
+            # Get visual features
+            projected_patch_embeddings = self._process_vision_features(pixel_values)
+
+            # Add proprioceptive state if provided
+            projected_patch_embeddings = self._process_proprio_features(
+                projected_patch_embeddings, proprio, proprio_projector
+            )
+
+            all_actions_mask = (labels == ACTION_TOKEN_IDX) #和run forward pass不一样, run forward pass要手动算token number来找偏移. 
+            all_actions_mask = all_actions_mask.unsqueeze(-1)  # (B, seq_len, 1)
+            input_embeddings = input_embeddings * ~all_actions_mask
+
+            # Build multimodal embeddings & attention mask
+            multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+                input_embeddings, projected_patch_embeddings, attention_mask
+            )
+
+            # Build labels for multimodal sequence if needed
+            multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
+
+            # Dispatch to language model
+            language_model_output = self.language_model(
+                input_ids=None,
+                attention_mask=multimodal_attention_mask,
+                position_ids=None,
+                past_key_values=None,
+                inputs_embeds=multimodal_embeddings,
+                labels=multimodal_labels,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+
+        # === Otherwise =>> Assume Invalid! ===
+        elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
+            raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
+
+        else:
+            raise ValueError(
+                "Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
+                f"=> `input_ids` = {input_ids is not None}\n"
+                f"=> `attention_mask` = {attention_mask is not None}\n"
+                f"=> `pixel_values` = {pixel_values is not None}\n"
+                f"=> `labels` = {labels is not None}\n"
+                f"=> `input_embeds` = {inputs_embeds is not None}\n"
+                f"=> `past_key_values` = {past_key_values is not None}\n"
+                f"=> `use_cache` = {use_cache}"
+            )
+
+        # Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
+        if not return_dict:
+            if output_projector_features and (projected_patch_embeddings is not None):
+                return *language_model_output, projected_patch_embeddings
+
+            return language_model_output
+
+        return PrismaticCausalLMOutputWithPast(
+            loss=language_model_output.loss,
+            logits=language_model_output.logits,
+            past_key_values=language_model_output.past_key_values,
+            hidden_states=language_model_output.hidden_states,
+            attentions=language_model_output.attentions,
+            projector_features=projected_patch_embeddings,
+        )
+
+
+
+    def mul_predict_action(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,  #就是 language instruction.
+        unnorm_key: Optional[str] = None,
+        proprio=None,
+        proprio_projector=None,
+        action_head:L1RegressionActionHead=None,
+        noisy_action_projector=None,
+        use_film: bool = False,
+        **kwargs: str,
+    ) -> np.ndarray:
+    # only use in evaluation.
+        cfg = kwargs.get("cfg", None)  # Extract cfg from kwargs
+        action_history = kwargs.get("action_history", None)
+
+        if not torch.all(input_ids[:, -1] == 29871):
+            input_ids = torch.cat(
+                (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
+            )
+
+        
+        pixel_values = kwargs["pixel_values"]
+        attention_mask = kwargs["attention_mask"]
+
+        #  input id  '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
+        import ipdb; ipdb.set_trace()
+        
+        input_embeddings = self.get_input_embeddings()(input_ids) 
+        
+        projected_patch_embeddings = self._process_vision_features(pixel_values)
+
+        use_proprio = proprio_projector is not None and proprio is not None
+        if use_proprio:
+            proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
+            projected_patch_embeddings = self._process_proprio_features(
+                projected_patch_embeddings, proprio, proprio_projector
+            )
+        if cfg.action_head_name == "hist":
+            normalized_actions, actions_hidden_states = self.hist_predict_action(
+                input_embeddings,
+                None, 
+                projected_patch_embeddings,
+                attention_mask,
+                action_head,
+                cfg=cfg,
+                action_history=action_history,
+            )
+        else:
+            normalized_actions, actions_hidden_states = self.mul_regression_or_discrete_prediction(
+                input_embeddings,
+                None, 
+                projected_patch_embeddings,
+                attention_mask,
+                None, #推理不需要labels
+                None, #推理不需要NUM_PATCHES
+                None, #推理不需要NUM_PROMPT_TOKENS
+                action_head,
+                cfg=cfg,
+            )
+        
+
+        actions = self._unnormalize_actions(normalized_actions, unnorm_key) #在这里 unorm, 所以出来的已经是unorm的了.  所以我 history 也要记录 norm 的. 
+
+        return actions, normalized_actions
+
+
+    def flow_matching_predict_action(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        unnorm_key: Optional[str] = None,
+        proprio=None,
+        proprio_projector=None,
+        action_head: FlowMatchingActionHead = None,
+        noisy_action_projector=None,
+        use_film: bool = False,
+        **kwargs: str,
+    ) -> np.ndarray:
+        """Predict actions using Flow Matching"""
+        cfg = kwargs.get("cfg", None)
+
+        if not torch.all(input_ids[:, -1] == 29871):
+            input_ids = torch.cat(
+                (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
+            )
+
+        pixel_values = kwargs["pixel_values"]
+        attention_mask = kwargs["attention_mask"]
+
+        input_embeddings = self.get_input_embeddings()(input_ids)
+        projected_patch_embeddings = self._process_vision_features(pixel_values)
+
+        use_proprio = proprio_projector is not None and proprio is not None
+        if use_proprio:
+            proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
+            projected_patch_embeddings = self._process_proprio_features(
+                projected_patch_embeddings, proprio, proprio_projector
+            )
+
+        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+            input_embeddings, projected_patch_embeddings, attention_mask
+        )
+
+        if cfg.mode == "flow_matching":
+            if cfg.num_actions_chunk//cfg.num_actions_per_token == 1:
+                language_model_output = self.language_model.generate(
+                    inputs_embeds=multimodal_embeddings,
+                    max_new_tokens=1,
+                    output_hidden_states=True,
+                    return_dict_in_generate=True
+                )
+                
+                actions_hidden_states = language_model_output.hidden_states[0][-1]
+                cognition_features = actions_hidden_states[:, -1]
+                assert (cognition_features.shape[0], cognition_features.shape[1]) == (1, 4096), "Batch size must be 1 for action prediction"
+
+                model_dtype = next(action_head.net.parameters()).dtype
+                cognition_features = cognition_features.unsqueeze(1).to(model_dtype)
+
+                # Sample actions using flow matching
+                normalized_actions = action_head.sample_actions(
+                    cognition_features,
+                    num_steps=getattr(cfg, 'num_flow_steps', 20)
+                )
+                normalized_actions = normalized_actions[0].cpu().numpy()
+            else:
+                raise NotImplementedError("Multi-token flow matching not yet implemented")
+        else:
+            raise NotImplementedError
+
+        actions = self._unnormalize_actions(normalized_actions, unnorm_key)
+        return actions, actions_hidden_states
+
+    def diffusion_predict_action(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,  #就是 language instruction.
+        unnorm_key: Optional[str] = None,
+        proprio=None,
+        proprio_projector=None,
+        action_head:DiTActionHead=None,
+        noisy_action_projector=None,
+        use_film: bool = False,
+        **kwargs: str,
+    ) -> np.ndarray:
+        cfg = kwargs.get("cfg", None)  # Extract cfg from kwargs
+
+        if not torch.all(input_ids[:, -1] == 29871):
+            input_ids = torch.cat(
+                (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
+            )
+
+        
+        pixel_values = kwargs["pixel_values"]
+        attention_mask = kwargs["attention_mask"]
+
+        #  input id  '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:' 
+        
+        input_embeddings = self.get_input_embeddings()(input_ids) 
+        
+        projected_patch_embeddings = self._process_vision_features(pixel_values)
+
+        use_proprio = proprio_projector is not None and proprio is not None
+        if use_proprio:
+            proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
+            projected_patch_embeddings = self._process_proprio_features(
+                projected_patch_embeddings, proprio, proprio_projector
+            )
+        
+        # normalized_actions, actions_hidden_states = self.mul_regression_or_discrete_prediction(
+        #     input_embeddings,
+        #     None, 
+        #     projected_patch_embeddings,
+        #     attention_mask,
+        #     None, #推理不需要labels
+        #     None, #推理不需要NUM_PATCHES
+        #     None, #推理不需要NUM_PROMPT_TOKENS
+        #     action_head,
+        #     cfg=cfg,
+        # )
+
+        # cfg = kwargs.get("cfg", None) 
+        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+            input_embeddings, projected_patch_embeddings, attention_mask
+        ) 
+        # multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'  
+        # first language_model_output.hidden_states  , is tuple (1 token,  33 layers, torch.Size([1, 314, 4096]))
+        # the following is  (num of tokens,)
+
+        if cfg.mode == "dit": 
+            if cfg.num_actions_chunk//cfg.num_actions_per_token > 1:
+                # token_num = cfg.num_actions_chunk//cfg.num_actions_per_token
+                # language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,min_new_tokens=token_num, max_new_tokens=token_num,output_hidden_states=True,return_dict_in_generate=True)
+                # actions_hidden_states0 = language_model_output.hidden_states[0][-1][:, -1]  # 第一个token的hidden state的last layer
+                
+                # actions_hidden_states_list = [actions_hidden_states0]
+                # for i in range(1, token_num):
+                #     token_hidden_state = language_model_output.hidden_states[i][-1][0]  # i+1个token的hidden state的last layer
+                #     actions_hidden_states_list.append(token_hidden_state)
+                # combined_hidden_states = torch.stack(actions_hidden_states_list, dim=0)  # (16, 4096)
+                # actions_hidden_states = combined_hidden_states.reshape(multimodal_embeddings.shape[0], token_num, multimodal_embeddings.shape[-1])
+                raise NotImplementedError
+            elif cfg.num_actions_chunk//cfg.num_actions_per_token == 1:
+                language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
+                # assert language_model_output.sequences == torch.tensor([[32001]], device=multimodal_embeddings.device)
+                actions_hidden_states = language_model_output.hidden_states[0][-1] 
+                cognition_features = actions_hidden_states[:, -1]
+                assert (cognition_features.shape[0], cognition_features.shape[1]) == (1,4096), "Batch size must be 1 for action prediction"
+                using_cfg = cfg.cfg_scale > 1.0
+
+                model_dtype = next(action_head.net.parameters()).dtype
+                B = cognition_features.shape[0]
+
+                cognition_features = cognition_features.unsqueeze(1).to(model_dtype) 
+                
+                noise = torch.randn(B, cfg.num_actions_chunk, action_head.net.in_channels, device=cognition_features.device).to(model_dtype) 
+
+                # TODO: Setup classifier-free guidance: now use cfg   
+                noise = torch.cat([noise, noise], 0) #  noise.shape torch.Size([2, 16, 7])
+                uncondition = action_head.net.z_embedder.uncondition # torch.Size([1, 4096])
+                uncondition = uncondition.unsqueeze(0)  #[1, D] # torch.Size([1, 1, 4096])
+                uncondition = uncondition.expand(B, 1, -1) #[B, 1, D]
+                z = torch.cat([cognition_features, uncondition], 0) # z  shape torch.Size([2, 1, 4096])
+                model_kwargs = dict(z=z, cfg_scale=cfg.cfg_scale)
+                sample_fn = action_head.net.forward_with_cfg
+                # default use ddim
+                if action_head.ddim_diffusion is None:
+                    action_head.create_ddim(ddim_step=cfg.num_ddim_steps)
+                samples = action_head.ddim_diffusion.ddim_sample_loop(sample_fn, 
+                                                                    noise.shape, 
+                                                                    noise, 
+                                                                    clip_denoised=False,
+                                                                    model_kwargs=model_kwargs,
+                                                                    progress=False,
+                                                                    device=cognition_features.device,
+                                                                    eta=0.0
+                                                                    )
+                if using_cfg:
+                    samples, _ = samples.chunk(2, dim=0)  # Remove null class samples
+                normalized_actions = samples[0].cpu().numpy()
+            else:
+                raise NotImplementedError
+        else:
+            raise NotImplementedError    
+
+
+
+        # normalized_actions = normalized_actions.float().cpu().detach().numpy()
+        actions = self._unnormalize_actions(normalized_actions, unnorm_key)
+
+        return actions, actions_hidden_states

preprocessor_config.json

@@ -0,0 +1,110 @@
+{
+  "image_processor_type": "PrismaticImageProcessor",
+  "image_resize_strategy": "resize-naive",
+  "input_sizes": [
+    [
+      3,
+      224,
+      224
+    ],
+    [
+      3,
+      224,
+      224
+    ]
+  ],
+  "interpolations": [
+    "bicubic",
+    "bicubic"
+  ],
+  "means": [
+    [
+      0.485,
+      0.456,
+      0.406
+    ],
+    [
+      0.5,
+      0.5,
+      0.5
+    ]
+  ],
+  "processor_class": "PrismaticProcessor",
+  "stds": [
+    [
+      0.229,
+      0.224,
+      0.225
+    ],
+    [
+      0.5,
+      0.5,
+      0.5
+    ]
+  ],
+  "tvf_crop_params": [
+    {
+      "output_size": [
+        224,
+        224
+      ]
+    },
+    {
+      "output_size": [
+        224,
+        224
+      ]
+    }
+  ],
+  "tvf_do_letterbox": false,
+  "tvf_letterbox_fill": null,
+  "tvf_normalize_params": [
+    {
+      "inplace": false,
+      "mean": [
+        0.484375,
+        0.455078125,
+        0.40625
+      ],
+      "std": [
+        0.228515625,
+        0.2236328125,
+        0.224609375
+      ]
+    },
+    {
+      "inplace": false,
+      "mean": [
+        0.5,
+        0.5,
+        0.5
+      ],
+      "std": [
+        0.5,
+        0.5,
+        0.5
+      ]
+    }
+  ],
+  "tvf_resize_params": [
+    {
+      "antialias": true,
+      "interpolation": 3,
+      "max_size": null,
+      "size": [
+        224,
+        224
+      ]
+    },
+    {
+      "antialias": true,
+      "interpolation": 3,
+      "max_size": null,
+      "size": [
+        224,
+        224
+      ]
+    }
+  ],
+  "use_fused_vision_backbone": true
+}

special_tokens_map.json

@@ -0,0 +1,32 @@
+{
+  "additional_special_tokens": [
+    {
+      "content": "<ACT>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    }
+  ],
+  "bos_token": {
+    "content": "<|begin_of_text|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|end_of_text|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<PAD>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:181432f1f6f7a4a71b3b17f3da5d631873d77f153394c8807bfff2a7473ab217
+size 17210284

tokenizer_config.json

@@ -0,0 +1,2084 @@
+{
+  "added_tokens_decoder": {
+    "128000": {
+      "content": "<|begin_of_text|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128001": {
+      "content": "<|end_of_text|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128002": {
+      "content": "<|reserved_special_token_0|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128003": {
+      "content": "<|reserved_special_token_1|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128004": {
+      "content": "<|finetune_right_pad_id|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128005": {
+      "content": "<|reserved_special_token_2|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128006": {
+      "content": "<|start_header_id|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128007": {
+      "content": "<|end_header_id|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128008": {
+      "content": "<|eom_id|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128009": {
+      "content": "<|eot_id|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128010": {
+      "content": "<|python_tag|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128011": {
+      "content": "<|reserved_special_token_3|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128012": {
+      "content": "<|reserved_special_token_4|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128013": {
+      "content": "<|reserved_special_token_5|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128014": {
+      "content": "<|reserved_special_token_6|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128015": {
+      "content": "<|reserved_special_token_7|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128016": {
+      "content": "<|reserved_special_token_8|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128017": {
+      "content": "<|reserved_special_token_9|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128018": {
+      "content": "<|reserved_special_token_10|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128019": {
+      "content": "<|reserved_special_token_11|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128020": {
+      "content": "<|reserved_special_token_12|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128021": {
+      "content": "<|reserved_special_token_13|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128022": {
+      "content": "<|reserved_special_token_14|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128023": {
+      "content": "<|reserved_special_token_15|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128024": {
+      "content": "<|reserved_special_token_16|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128025": {
+      "content": "<|reserved_special_token_17|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128026": {
+      "content": "<|reserved_special_token_18|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128027": {
+      "content": "<|reserved_special_token_19|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128028": {
+      "content": "<|reserved_special_token_20|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128029": {
+      "content": "<|reserved_special_token_21|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128030": {
+      "content": "<|reserved_special_token_22|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128031": {
+      "content": "<|reserved_special_token_23|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128032": {
+      "content": "<|reserved_special_token_24|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128033": {
+      "content": "<|reserved_special_token_25|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128034": {
+      "content": "<|reserved_special_token_26|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128035": {
+      "content": "<|reserved_special_token_27|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128036": {
+      "content": "<|reserved_special_token_28|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128037": {
+      "content": "<|reserved_special_token_29|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128038": {
+      "content": "<|reserved_special_token_30|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128039": {
+      "content": "<|reserved_special_token_31|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128040": {
+      "content": "<|reserved_special_token_32|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128041": {
+      "content": "<|reserved_special_token_33|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128042": {
+      "content": "<|reserved_special_token_34|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128043": {
+      "content": "<|reserved_special_token_35|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128044": {
+      "content": "<|reserved_special_token_36|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128045": {
+      "content": "<|reserved_special_token_37|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128046": {
+      "content": "<|reserved_special_token_38|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128047": {
+      "content": "<|reserved_special_token_39|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128048": {
+      "content": "<|reserved_special_token_40|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128049": {
+      "content": "<|reserved_special_token_41|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128050": {
+      "content": "<|reserved_special_token_42|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128051": {
+      "content": "<|reserved_special_token_43|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128052": {
+      "content": "<|reserved_special_token_44|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128053": {
+      "content": "<|reserved_special_token_45|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128054": {
+      "content": "<|reserved_special_token_46|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128055": {
+      "content": "<|reserved_special_token_47|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128056": {
+      "content": "<|reserved_special_token_48|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128057": {
+      "content": "<|reserved_special_token_49|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128058": {
+      "content": "<|reserved_special_token_50|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128059": {
+      "content": "<|reserved_special_token_51|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128060": {
+      "content": "<|reserved_special_token_52|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128061": {
+      "content": "<|reserved_special_token_53|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128062": {
+      "content": "<|reserved_special_token_54|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128063": {
+      "content": "<|reserved_special_token_55|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128064": {
+      "content": "<|reserved_special_token_56|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128065": {
+      "content": "<|reserved_special_token_57|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128066": {
+      "content": "<|reserved_special_token_58|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128067": {
+      "content": "<|reserved_special_token_59|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128068": {
+      "content": "<|reserved_special_token_60|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128069": {
+      "content": "<|reserved_special_token_61|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128070": {
+      "content": "<|reserved_special_token_62|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128071": {
+      "content": "<|reserved_special_token_63|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128072": {
+      "content": "<|reserved_special_token_64|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128073": {
+      "content": "<|reserved_special_token_65|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128074": {
+      "content": "<|reserved_special_token_66|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128075": {
+      "content": "<|reserved_special_token_67|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128076": {
+      "content": "<|reserved_special_token_68|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128077": {
+      "content": "<|reserved_special_token_69|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128078": {
+      "content": "<|reserved_special_token_70|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128079": {
+      "content": "<|reserved_special_token_71|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128080": {
+      "content": "<|reserved_special_token_72|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128081": {
+      "content": "<|reserved_special_token_73|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128082": {
+      "content": "<|reserved_special_token_74|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128083": {
+      "content": "<|reserved_special_token_75|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128084": {
+      "content": "<|reserved_special_token_76|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128085": {
+      "content": "<|reserved_special_token_77|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128086": {
+      "content": "<|reserved_special_token_78|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128087": {
+      "content": "<|reserved_special_token_79|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128088": {
+      "content": "<|reserved_special_token_80|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128089": {
+      "content": "<|reserved_special_token_81|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128090": {
+      "content": "<|reserved_special_token_82|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128091": {
+      "content": "<|reserved_special_token_83|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128092": {
+      "content": "<|reserved_special_token_84|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128093": {
+      "content": "<|reserved_special_token_85|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128094": {
+      "content": "<|reserved_special_token_86|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128095": {
+      "content": "<|reserved_special_token_87|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128096": {
+      "content": "<|reserved_special_token_88|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128097": {
+      "content": "<|reserved_special_token_89|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128098": {
+      "content": "<|reserved_special_token_90|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128099": {
+      "content": "<|reserved_special_token_91|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128100": {
+      "content": "<|reserved_special_token_92|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128101": {
+      "content": "<|reserved_special_token_93|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128102": {
+      "content": "<|reserved_special_token_94|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128103": {
+      "content": "<|reserved_special_token_95|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128104": {
+      "content": "<|reserved_special_token_96|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128105": {
+      "content": "<|reserved_special_token_97|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128106": {
+      "content": "<|reserved_special_token_98|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128107": {
+      "content": "<|reserved_special_token_99|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128108": {
+      "content": "<|reserved_special_token_100|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128109": {
+      "content": "<|reserved_special_token_101|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128110": {
+      "content": "<|reserved_special_token_102|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128111": {
+      "content": "<|reserved_special_token_103|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128112": {
+      "content": "<|reserved_special_token_104|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128113": {
+      "content": "<|reserved_special_token_105|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128114": {
+      "content": "<|reserved_special_token_106|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128115": {
+      "content": "<|reserved_special_token_107|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128116": {
+      "content": "<|reserved_special_token_108|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128117": {
+      "content": "<|reserved_special_token_109|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128118": {
+      "content": "<|reserved_special_token_110|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128119": {
+      "content": "<|reserved_special_token_111|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128120": {
+      "content": "<|reserved_special_token_112|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128121": {
+      "content": "<|reserved_special_token_113|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128122": {
+      "content": "<|reserved_special_token_114|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128123": {
+      "content": "<|reserved_special_token_115|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128124": {
+      "content": "<|reserved_special_token_116|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128125": {
+      "content": "<|reserved_special_token_117|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128126": {
+      "content": "<|reserved_special_token_118|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128127": {
+      "content": "<|reserved_special_token_119|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128128": {
+      "content": "<|reserved_special_token_120|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128129": {
+      "content": "<|reserved_special_token_121|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128130": {
+      "content": "<|reserved_special_token_122|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128131": {
+      "content": "<|reserved_special_token_123|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128132": {
+      "content": "<|reserved_special_token_124|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128133": {
+      "content": "<|reserved_special_token_125|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128134": {
+      "content": "<|reserved_special_token_126|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128135": {
+      "content": "<|reserved_special_token_127|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128136": {
+      "content": "<|reserved_special_token_128|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128137": {
+      "content": "<|reserved_special_token_129|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128138": {
+      "content": "<|reserved_special_token_130|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128139": {
+      "content": "<|reserved_special_token_131|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128140": {
+      "content": "<|reserved_special_token_132|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128141": {
+      "content": "<|reserved_special_token_133|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128142": {
+      "content": "<|reserved_special_token_134|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128143": {
+      "content": "<|reserved_special_token_135|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128144": {
+      "content": "<|reserved_special_token_136|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128145": {
+      "content": "<|reserved_special_token_137|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128146": {
+      "content": "<|reserved_special_token_138|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128147": {
+      "content": "<|reserved_special_token_139|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128148": {
+      "content": "<|reserved_special_token_140|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128149": {
+      "content": "<|reserved_special_token_141|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128150": {
+      "content": "<|reserved_special_token_142|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128151": {
+      "content": "<|reserved_special_token_143|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128152": {
+      "content": "<|reserved_special_token_144|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128153": {
+      "content": "<|reserved_special_token_145|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128154": {
+      "content": "<|reserved_special_token_146|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128155": {
+      "content": "<|reserved_special_token_147|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128156": {
+      "content": "<|reserved_special_token_148|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128157": {
+      "content": "<|reserved_special_token_149|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128158": {
+      "content": "<|reserved_special_token_150|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128159": {
+      "content": "<|reserved_special_token_151|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128160": {
+      "content": "<|reserved_special_token_152|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128161": {
+      "content": "<|reserved_special_token_153|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128162": {
+      "content": "<|reserved_special_token_154|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128163": {
+      "content": "<|reserved_special_token_155|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128164": {
+      "content": "<|reserved_special_token_156|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128165": {
+      "content": "<|reserved_special_token_157|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128166": {
+      "content": "<|reserved_special_token_158|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128167": {
+      "content": "<|reserved_special_token_159|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128168": {
+      "content": "<|reserved_special_token_160|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128169": {
+      "content": "<|reserved_special_token_161|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128170": {
+      "content": "<|reserved_special_token_162|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128171": {
+      "content": "<|reserved_special_token_163|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128172": {
+      "content": "<|reserved_special_token_164|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128173": {
+      "content": "<|reserved_special_token_165|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128174": {
+      "content": "<|reserved_special_token_166|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128175": {
+      "content": "<|reserved_special_token_167|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128176": {
+      "content": "<|reserved_special_token_168|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128177": {
+      "content": "<|reserved_special_token_169|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128178": {
+      "content": "<|reserved_special_token_170|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128179": {
+      "content": "<|reserved_special_token_171|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128180": {
+      "content": "<|reserved_special_token_172|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128181": {
+      "content": "<|reserved_special_token_173|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128182": {
+      "content": "<|reserved_special_token_174|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128183": {
+      "content": "<|reserved_special_token_175|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128184": {
+      "content": "<|reserved_special_token_176|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128185": {
+      "content": "<|reserved_special_token_177|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128186": {
+      "content": "<|reserved_special_token_178|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128187": {
+      "content": "<|reserved_special_token_179|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128188": {
+      "content": "<|reserved_special_token_180|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128189": {
+      "content": "<|reserved_special_token_181|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128190": {
+      "content": "<|reserved_special_token_182|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128191": {
+      "content": "<|reserved_special_token_183|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128192": {
+      "content": "<|reserved_special_token_184|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128193": {
+      "content": "<|reserved_special_token_185|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128194": {
+      "content": "<|reserved_special_token_186|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128195": {
+      "content": "<|reserved_special_token_187|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128196": {
+      "content": "<|reserved_special_token_188|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128197": {
+      "content": "<|reserved_special_token_189|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128198": {
+      "content": "<|reserved_special_token_190|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128199": {
+      "content": "<|reserved_special_token_191|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128200": {
+      "content": "<|reserved_special_token_192|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128201": {
+      "content": "<|reserved_special_token_193|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128202": {
+      "content": "<|reserved_special_token_194|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128203": {
+      "content": "<|reserved_special_token_195|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128204": {
+      "content": "<|reserved_special_token_196|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128205": {
+      "content": "<|reserved_special_token_197|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128206": {
+      "content": "<|reserved_special_token_198|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128207": {
+      "content": "<|reserved_special_token_199|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128208": {
+      "content": "<|reserved_special_token_200|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128209": {
+      "content": "<|reserved_special_token_201|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128210": {
+      "content": "<|reserved_special_token_202|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128211": {
+      "content": "<|reserved_special_token_203|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128212": {
+      "content": "<|reserved_special_token_204|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128213": {
+      "content": "<|reserved_special_token_205|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128214": {
+      "content": "<|reserved_special_token_206|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128215": {
+      "content": "<|reserved_special_token_207|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128216": {
+      "content": "<|reserved_special_token_208|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128217": {
+      "content": "<|reserved_special_token_209|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128218": {
+      "content": "<|reserved_special_token_210|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128219": {
+      "content": "<|reserved_special_token_211|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128220": {
+      "content": "<|reserved_special_token_212|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128221": {
+      "content": "<|reserved_special_token_213|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128222": {
+      "content": "<|reserved_special_token_214|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128223": {
+      "content": "<|reserved_special_token_215|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128224": {
+      "content": "<|reserved_special_token_216|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128225": {
+      "content": "<|reserved_special_token_217|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128226": {
+      "content": "<|reserved_special_token_218|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128227": {
+      "content": "<|reserved_special_token_219|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128228": {
+      "content": "<|reserved_special_token_220|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128229": {
+      "content": "<|reserved_special_token_221|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128230": {
+      "content": "<|reserved_special_token_222|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128231": {
+      "content": "<|reserved_special_token_223|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128232": {
+      "content": "<|reserved_special_token_224|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128233": {
+      "content": "<|reserved_special_token_225|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128234": {
+      "content": "<|reserved_special_token_226|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128235": {
+      "content": "<|reserved_special_token_227|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128236": {
+      "content": "<|reserved_special_token_228|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128237": {
+      "content": "<|reserved_special_token_229|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128238": {
+      "content": "<|reserved_special_token_230|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128239": {
+      "content": "<|reserved_special_token_231|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128240": {
+      "content": "<|reserved_special_token_232|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128241": {
+      "content": "<|reserved_special_token_233|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128242": {
+      "content": "<|reserved_special_token_234|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128243": {
+      "content": "<|reserved_special_token_235|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128244": {
+      "content": "<|reserved_special_token_236|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128245": {
+      "content": "<|reserved_special_token_237|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128246": {
+      "content": "<|reserved_special_token_238|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128247": {
+      "content": "<|reserved_special_token_239|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128248": {
+      "content": "<|reserved_special_token_240|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128249": {
+      "content": "<|reserved_special_token_241|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128250": {
+      "content": "<|reserved_special_token_242|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128251": {
+      "content": "<|reserved_special_token_243|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128252": {
+      "content": "<|reserved_special_token_244|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128253": {
+      "content": "<|reserved_special_token_245|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128254": {
+      "content": "<|reserved_special_token_246|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128255": {
+      "content": "<|reserved_special_token_247|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128256": {
+      "content": "<PAD>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128257": {
+      "content": "<ACT>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<ACT>"
+  ],
+  "bos_token": "<|begin_of_text|>",
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "<|end_of_text|>",
+  "extra_special_tokens": {},
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 2048,
+  "pad_token": "<PAD>",
+  "padding_side": "right",
+  "processor_class": "PrismaticProcessor",
+  "tokenizer_class": "PreTrainedTokenizer"
+}