Sombit commited on
Commit
36e08c0
·
verified ·
1 Parent(s): d8442f2

Upload config

Browse files
Files changed (3) hide show
  1. README.md +199 -0
  2. config.json +222 -0
  3. prismatic_config.py +309 -0
README.md ADDED
@@ -0,0 +1,199 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ library_name: transformers
3
+ tags: []
4
+ ---
5
+
6
+ # Model Card for Model ID
7
+
8
+ <!-- Provide a quick summary of what the model is/does. -->
9
+
10
+
11
+
12
+ ## Model Details
13
+
14
+ ### Model Description
15
+
16
+ <!-- Provide a longer summary of what this model is. -->
17
+
18
+ This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
19
+
20
+ - **Developed by:** [More Information Needed]
21
+ - **Funded by [optional]:** [More Information Needed]
22
+ - **Shared by [optional]:** [More Information Needed]
23
+ - **Model type:** [More Information Needed]
24
+ - **Language(s) (NLP):** [More Information Needed]
25
+ - **License:** [More Information Needed]
26
+ - **Finetuned from model [optional]:** [More Information Needed]
27
+
28
+ ### Model Sources [optional]
29
+
30
+ <!-- Provide the basic links for the model. -->
31
+
32
+ - **Repository:** [More Information Needed]
33
+ - **Paper [optional]:** [More Information Needed]
34
+ - **Demo [optional]:** [More Information Needed]
35
+
36
+ ## Uses
37
+
38
+ <!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
39
+
40
+ ### Direct Use
41
+
42
+ <!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
43
+
44
+ [More Information Needed]
45
+
46
+ ### Downstream Use [optional]
47
+
48
+ <!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
49
+
50
+ [More Information Needed]
51
+
52
+ ### Out-of-Scope Use
53
+
54
+ <!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
55
+
56
+ [More Information Needed]
57
+
58
+ ## Bias, Risks, and Limitations
59
+
60
+ <!-- This section is meant to convey both technical and sociotechnical limitations. -->
61
+
62
+ [More Information Needed]
63
+
64
+ ### Recommendations
65
+
66
+ <!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
67
+
68
+ Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
69
+
70
+ ## How to Get Started with the Model
71
+
72
+ Use the code below to get started with the model.
73
+
74
+ [More Information Needed]
75
+
76
+ ## Training Details
77
+
78
+ ### Training Data
79
+
80
+ <!-- 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. -->
81
+
82
+ [More Information Needed]
83
+
84
+ ### Training Procedure
85
+
86
+ <!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
87
+
88
+ #### Preprocessing [optional]
89
+
90
+ [More Information Needed]
91
+
92
+
93
+ #### Training Hyperparameters
94
+
95
+ - **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
96
+
97
+ #### Speeds, Sizes, Times [optional]
98
+
99
+ <!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
100
+
101
+ [More Information Needed]
102
+
103
+ ## Evaluation
104
+
105
+ <!-- This section describes the evaluation protocols and provides the results. -->
106
+
107
+ ### Testing Data, Factors & Metrics
108
+
109
+ #### Testing Data
110
+
111
+ <!-- This should link to a Dataset Card if possible. -->
112
+
113
+ [More Information Needed]
114
+
115
+ #### Factors
116
+
117
+ <!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
118
+
119
+ [More Information Needed]
120
+
121
+ #### Metrics
122
+
123
+ <!-- These are the evaluation metrics being used, ideally with a description of why. -->
124
+
125
+ [More Information Needed]
126
+
127
+ ### Results
128
+
129
+ [More Information Needed]
130
+
131
+ #### Summary
132
+
133
+
134
+
135
+ ## Model Examination [optional]
136
+
137
+ <!-- Relevant interpretability work for the model goes here -->
138
+
139
+ [More Information Needed]
140
+
141
+ ## Environmental Impact
142
+
143
+ <!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
144
+
145
+ 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).
146
+
147
+ - **Hardware Type:** [More Information Needed]
148
+ - **Hours used:** [More Information Needed]
149
+ - **Cloud Provider:** [More Information Needed]
150
+ - **Compute Region:** [More Information Needed]
151
+ - **Carbon Emitted:** [More Information Needed]
152
+
153
+ ## Technical Specifications [optional]
154
+
155
+ ### Model Architecture and Objective
156
+
157
+ [More Information Needed]
158
+
159
+ ### Compute Infrastructure
160
+
161
+ [More Information Needed]
162
+
163
+ #### Hardware
164
+
165
+ [More Information Needed]
166
+
167
+ #### Software
168
+
169
+ [More Information Needed]
170
+
171
+ ## Citation [optional]
172
+
173
+ <!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
174
+
175
+ **BibTeX:**
176
+
177
+ [More Information Needed]
178
+
179
+ **APA:**
180
+
181
+ [More Information Needed]
182
+
183
+ ## Glossary [optional]
184
+
185
+ <!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
186
+
187
+ [More Information Needed]
188
+
189
+ ## More Information [optional]
190
+
191
+ [More Information Needed]
192
+
193
+ ## Model Card Authors [optional]
194
+
195
+ [More Information Needed]
196
+
197
+ ## Model Card Contact
198
+
199
+ [More Information Needed]
config.json ADDED
@@ -0,0 +1,222 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "auto_map": {
3
+ "AutoConfig": "prismatic_config.TrajectoryVLAConfig"
4
+ },
5
+ "cheat": false,
6
+ "model_type": "trajectoryvla",
7
+ "num_timesteps": 6,
8
+ "prismatic_config": {
9
+ "_name_or_path": "",
10
+ "add_cross_attention": false,
11
+ "arch_specifier": "no-align+gelu-mlp",
12
+ "architectures": [
13
+ "TrajectoryVLA"
14
+ ],
15
+ "auto_map": {
16
+ "AutoModelForVision2Seq": "prismatic_model.TrajectoryVLA"
17
+ },
18
+ "bad_words_ids": null,
19
+ "begin_suppress_tokens": null,
20
+ "bos_token_id": null,
21
+ "chunk_size_feed_forward": 0,
22
+ "cross_attention_hidden_size": null,
23
+ "decoder_start_token_id": null,
24
+ "diversity_penalty": 0.0,
25
+ "do_sample": false,
26
+ "early_stopping": false,
27
+ "encoder_no_repeat_ngram_size": 0,
28
+ "eos_token_id": null,
29
+ "exponential_decay_length_penalty": null,
30
+ "finetuning_task": null,
31
+ "forced_bos_token_id": null,
32
+ "forced_eos_token_id": null,
33
+ "hf_llm_id": "meta-llama/Llama-2-7b-hf",
34
+ "id2label": {
35
+ "0": "LABEL_0",
36
+ "1": "LABEL_1"
37
+ },
38
+ "image_resize_strategy": "letterbox",
39
+ "image_sizes": [
40
+ 224,
41
+ 224
42
+ ],
43
+ "is_decoder": false,
44
+ "is_encoder_decoder": false,
45
+ "label2id": {
46
+ "LABEL_0": 0,
47
+ "LABEL_1": 1
48
+ },
49
+ "length_penalty": 1.0,
50
+ "llm_backbone_id": "llama2-7b-pure",
51
+ "llm_max_length": 2048,
52
+ "max_length": 20,
53
+ "min_length": 0,
54
+ "model_type": "prismatic",
55
+ "no_repeat_ngram_size": 0,
56
+ "num_beam_groups": 1,
57
+ "num_beams": 1,
58
+ "num_return_sequences": 1,
59
+ "output_attentions": false,
60
+ "output_hidden_states": false,
61
+ "output_projector_states": false,
62
+ "output_scores": false,
63
+ "pad_to_multiple_of": 64,
64
+ "pad_token_id": 32000,
65
+ "prefix": null,
66
+ "problem_type": null,
67
+ "pruned_heads": {},
68
+ "remove_invalid_values": false,
69
+ "repetition_penalty": 1.0,
70
+ "return_dict": false,
71
+ "return_dict_in_generate": false,
72
+ "sep_token_id": null,
73
+ "suppress_tokens": null,
74
+ "task_specific_params": null,
75
+ "temperature": 1.0,
76
+ "text_config": {
77
+ "_name_or_path": "",
78
+ "add_cross_attention": false,
79
+ "architectures": null,
80
+ "attention_bias": false,
81
+ "attention_dropout": 0.0,
82
+ "bad_words_ids": null,
83
+ "begin_suppress_tokens": null,
84
+ "bos_token_id": 1,
85
+ "chunk_size_feed_forward": 0,
86
+ "cross_attention_hidden_size": null,
87
+ "decoder_start_token_id": null,
88
+ "diversity_penalty": 0.0,
89
+ "do_sample": false,
90
+ "early_stopping": false,
91
+ "encoder_no_repeat_ngram_size": 0,
92
+ "eos_token_id": 2,
93
+ "exponential_decay_length_penalty": null,
94
+ "finetuning_task": null,
95
+ "forced_bos_token_id": null,
96
+ "forced_eos_token_id": null,
97
+ "hidden_act": "silu",
98
+ "hidden_size": 4096,
99
+ "id2label": {
100
+ "0": "LABEL_0",
101
+ "1": "LABEL_1"
102
+ },
103
+ "initializer_range": 0.02,
104
+ "intermediate_size": 11008,
105
+ "is_decoder": false,
106
+ "is_encoder_decoder": false,
107
+ "label2id": {
108
+ "LABEL_0": 0,
109
+ "LABEL_1": 1
110
+ },
111
+ "length_penalty": 1.0,
112
+ "max_length": 20,
113
+ "max_position_embeddings": 2048,
114
+ "min_length": 0,
115
+ "mlp_bias": false,
116
+ "model_type": "llama",
117
+ "no_repeat_ngram_size": 0,
118
+ "num_attention_heads": 32,
119
+ "num_beam_groups": 1,
120
+ "num_beams": 1,
121
+ "num_hidden_layers": 32,
122
+ "num_key_value_heads": 32,
123
+ "num_return_sequences": 1,
124
+ "output_attentions": false,
125
+ "output_hidden_states": false,
126
+ "output_scores": false,
127
+ "pad_token_id": null,
128
+ "prefix": null,
129
+ "pretraining_tp": 1,
130
+ "problem_type": null,
131
+ "pruned_heads": {},
132
+ "remove_invalid_values": false,
133
+ "repetition_penalty": 1.0,
134
+ "return_dict": true,
135
+ "return_dict_in_generate": false,
136
+ "rms_norm_eps": 1e-06,
137
+ "rope_scaling": null,
138
+ "rope_theta": 10000.0,
139
+ "sep_token_id": null,
140
+ "suppress_tokens": null,
141
+ "task_specific_params": null,
142
+ "temperature": 1.0,
143
+ "tf_legacy_loss": false,
144
+ "tie_encoder_decoder": false,
145
+ "tie_word_embeddings": false,
146
+ "tokenizer_class": null,
147
+ "top_k": 50,
148
+ "top_p": 1.0,
149
+ "torch_dtype": null,
150
+ "torchscript": false,
151
+ "typical_p": 1.0,
152
+ "use_bfloat16": false,
153
+ "use_cache": true,
154
+ "vocab_size": 32000
155
+ },
156
+ "tf_legacy_loss": false,
157
+ "tie_encoder_decoder": false,
158
+ "tie_word_embeddings": true,
159
+ "timm_model_ids": [
160
+ "vit_large_patch14_reg4_dinov2.lvd142m",
161
+ "vit_so400m_patch14_siglip_224"
162
+ ],
163
+ "timm_override_act_layers": [
164
+ null,
165
+ null
166
+ ],
167
+ "tokenizer_class": null,
168
+ "top_k": 50,
169
+ "top_p": 1.0,
170
+ "torch_dtype": "bfloat16",
171
+ "torchscript": false,
172
+ "typical_p": 1.0,
173
+ "use_bfloat16": false,
174
+ "use_fused_vision_backbone": true,
175
+ "vision_backbone_id": "dinosiglip-vit-so-224px"
176
+ },
177
+ "rotation_components": 9,
178
+ "seperate_control_proj": true,
179
+ "timestep_proj_config": {
180
+ "num_tokens": 3,
181
+ "pos_embed_scale": 8,
182
+ "proj_layers": [
183
+ 128,
184
+ 512,
185
+ 1024
186
+ ],
187
+ "time_delta_sec": 0.1
188
+ },
189
+ "token_proj_config": {
190
+ "control_tokens_layers": [
191
+ 4096,
192
+ 2048,
193
+ 1024
194
+ ],
195
+ "image_tokens_mode": "vit",
196
+ "llm_image_tokens_layers": [],
197
+ "vit_tokens_layers": [
198
+ 2176,
199
+ 1024
200
+ ]
201
+ },
202
+ "token_size": 1024,
203
+ "transformer_config": {
204
+ "decoder_block_config": {
205
+ "dropout": 0.0,
206
+ "feature_size": 1024,
207
+ "head_dim": 64,
208
+ "num_heads": 16
209
+ },
210
+ "encoder_block_config": {
211
+ "feature_size": 1024,
212
+ "head_dim": 64,
213
+ "num_heads": 16
214
+ },
215
+ "num_blocks": 2,
216
+ "pos_embed_config": {
217
+ "embedding_dim": 1024,
218
+ "num_embeddings": 300
219
+ }
220
+ },
221
+ "transformers_version": "4.44.2"
222
+ }
prismatic_config.py ADDED
@@ -0,0 +1,309 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ configuration_prismatic.py
3
+
4
+ HuggingFace-style configuration definition for Prismatic VLMs, inheriting from `transformers.PretrainedConfig`.
5
+ Default configuration specifies `siglip-224px+7b`.
6
+ """
7
+
8
+ from typing import Any, Dict, List, Optional
9
+ import transformers
10
+ from transformers import PretrainedConfig
11
+ from transformers.models.auto import CONFIG_MAPPING
12
+ import numpy as np
13
+
14
+ # === Utilities for Mapping Prismatic names to HF names ===
15
+ # fmt: off
16
+ VISION_BACKBONE_TO_RESOLUTION: Dict[str, List[int]] = {
17
+ "clip-vit-l": [224], "siglip-vit-so400m": [224], "dinov2-vit-l": [224], "in1k-vit-l": [224],
18
+
19
+ "clip-vit-l-336px": [336],
20
+ "siglip-vit-so400m-384px": [384],
21
+
22
+ "dinoclip-vit-l-336px": [336, 336],
23
+ "dinosiglip-vit-so-224px": [224, 224],
24
+ "dinosiglip-vit-so-384px": [384, 384],
25
+ }
26
+ VISION_BACKBONE_TO_TIMM_ID: Dict[str, List[str]] = {
27
+ "clip-vit-l": ["vit_large_patch14_clip_224.openai"],
28
+ "clip-vit-l-336px": ["vit_large_patch14_clip_336.openai"],
29
+
30
+ "dinov2-vit-l": ["vit_large_patch14_reg4_dinov2.lvd142m"],
31
+ "in1k-vit-l": ["vit_large_patch16_224.augreg_in21k_ft_in1k"],
32
+
33
+ "siglip-vit-so400m": ["vit_so400m_patch14_siglip_224"],
34
+ "siglip-vit-so400m-384px": ["vit_so400m_patch14_siglip_384"],
35
+
36
+ "dinoclip-vit-l-336px": ["vit_large_patch14_reg4_dinov2.lvd142m", "vit_large_patch14_clip_336.openai"],
37
+ "dinosiglip-vit-so-224px": ["vit_large_patch14_reg4_dinov2.lvd142m", "vit_so400m_patch14_siglip_224"],
38
+ "dinosiglip-vit-so-384px": ["vit_large_patch14_reg4_dinov2.lvd142m", "vit_so400m_patch14_siglip_384"],
39
+ }
40
+ TIMM_OVERRIDE_ACT_LAYER: Dict[str, List[Optional[str]]] = {
41
+ "clip-vit-l": ["quick_gelu"], "clip-vit-l-336px": ["quick_gelu"],
42
+ "dinov2-vit-l": [None], "in1k-vit-l": [None],
43
+ "siglip-vit-so400m": [None], "siglip-vit-so400m-384px": [None],
44
+ "dinoclip-vit-l-336px": [None, "quick_gelu"],
45
+ "dinosiglip-vit-so-224px": [None, None], "dinosiglip-vit-so-384px": [None, None]
46
+ }
47
+
48
+ LLM_BACKBONE_TO_HF_PATH = {
49
+ "llama2-7b-pure": "meta-llama/Llama-2-7b-hf", "llama2-13b-pure": "meta-llama/Llama-2-13b-hf",
50
+ "llama2-7b-chat": "meta-llama/Llama-2-7b-chat-hf", "llama2-13b-chat": "meta-llama/Llama-2-13b-chat-hf",
51
+
52
+ "vicuna-v15-7b": "lmsys/vicuna-7b-v1.5", "vicuna-v15-13b": "lmsys/vicuna-13b-v1.5",
53
+
54
+ "mistral-v0.1-7b-pure": "mistralai/Mistral-7B-v0.1",
55
+ "mistral-v0.1-7b-instruct": "mistralai/Mistral-7B-Instruct-v0.1",
56
+
57
+ "phi-2-3b": "microsoft/phi-2",
58
+ }
59
+ LLM_BACKBONE_TO_HF_METACLASS = {
60
+ "llama2-7b-pure": "llama", "llama2-13b-pure": "llama", "llama2-7b-chat": "llama", "llama2-13b-chat": "llama",
61
+ "vicuna-v15-7b": "llama", "vicuna-v15-13b": "llama",
62
+
63
+ "mistral-v0.1-7b-pure": "mistral", "mistral-v0.1-7b-instruct": "mistral",
64
+
65
+ "phi-2-3b": "phi",
66
+ }
67
+
68
+ VALID_VISION_BACKBONES = set(VISION_BACKBONE_TO_RESOLUTION.keys())
69
+ VALID_LLM_BACKBONES = set(LLM_BACKBONE_TO_HF_PATH)
70
+ # fmt: on
71
+
72
+ class WaypointTokenizer:
73
+ """
74
+ Wraps base LLM/VLM tokenizer and overloads least used token as a control token
75
+
76
+ NOTE: By default, assumes a BPE-style tokenizer akin to the LlamaTokenizer,
77
+ where *the least used tokens* appear at the end of the vocabulary!
78
+
79
+ TODO: Adding new token vs overloading? When I call `tokenizer.add_token()` vocab stays the same
80
+ """
81
+ model_type = "waypointer"
82
+ is_composition: bool = True
83
+ def __init__(self, tokenizer: transformers.PreTrainedTokenizerBase, num_tokens: int = 10) -> None:
84
+ self.tokenizer = tokenizer
85
+ self.num_tokens = num_tokens
86
+
87
+ def __call__(self, *_) -> str:
88
+ """Get the text token for control"""
89
+ return self.tokenizer.decode(self.control_token_ids)
90
+
91
+ @property
92
+ def control_token_ids(self) -> np.ndarray:
93
+ # Assumes we're overwriting the final tokens of the vocabulary (least used tokens)
94
+ return np.arange(self.num_tokens) + int(self.tokenizer.vocab_size - self.num_tokens)
95
+
96
+ @property
97
+ def num_control_tokens(self) -> int:
98
+ return self.num_tokens
99
+
100
+ class PrismaticConfig(PretrainedConfig):
101
+ model_type: str = "prismatic"
102
+ is_composition: bool = False
103
+
104
+ def __init__(
105
+ self,
106
+ vision_backbone_id: str = "dinosiglip-vit-so-224px",
107
+ llm_backbone_id: str = "llama2-7b-pure",
108
+ arch_specifier: str = "no-align+gelu-mlp", ## TODO: check
109
+ use_fused_vision_backbone: Optional[bool] = None, ## TODO: check
110
+ image_resize_strategy: str = "letterbox",
111
+ text_config: Optional[Dict[str, Any]] = None,
112
+ llm_max_length: int = 2048,
113
+ pad_token_id: int = 32000,
114
+ pad_to_multiple_of: int = 64,
115
+ output_projector_states: bool = False,
116
+ **kwargs: str,
117
+ ) -> None:
118
+ if vision_backbone_id not in VALID_VISION_BACKBONES:
119
+ raise ValueError(f"Vision backbone `{vision_backbone_id}` not in {VALID_VISION_BACKBONES = }")
120
+
121
+ if llm_backbone_id not in VALID_LLM_BACKBONES:
122
+ raise ValueError(f"LLM backbone `{llm_backbone_id}` not in {VALID_LLM_BACKBONES = }")
123
+
124
+ # Set Prismatic Configuration Fields
125
+ self.vision_backbone_id = vision_backbone_id
126
+ self.llm_backbone_id = llm_backbone_id
127
+ self.arch_specifier = arch_specifier
128
+ self.output_projector_states = output_projector_states
129
+
130
+ # [Contract] All vision backbone parameters are lists =>> supports fused backbones with different preprocessing
131
+ self.use_fused_vision_backbone = (
132
+ use_fused_vision_backbone
133
+ if use_fused_vision_backbone is not None
134
+ else any(self.vision_backbone_id.startswith(v) for v in ["dinoclip", "dinosiglip"])
135
+ )
136
+
137
+ self.timm_model_ids = VISION_BACKBONE_TO_TIMM_ID[self.vision_backbone_id]
138
+ self.timm_override_act_layers = TIMM_OVERRIDE_ACT_LAYER[self.vision_backbone_id]
139
+ self.image_sizes = VISION_BACKBONE_TO_RESOLUTION[self.vision_backbone_id]
140
+ self.image_resize_strategy = image_resize_strategy
141
+
142
+ self.hf_llm_id = LLM_BACKBONE_TO_HF_PATH[self.llm_backbone_id]
143
+ self.llm_max_length = llm_max_length
144
+ self.pad_token_id, self.pad_to_multiple_of = pad_token_id, pad_to_multiple_of
145
+
146
+ # [IMPORTANT] HF Utilities actually look for a `text_config` field... we need to use that specific naming!
147
+ self.text_config = (
148
+ CONFIG_MAPPING[LLM_BACKBONE_TO_HF_METACLASS[self.llm_backbone_id]](**text_config)
149
+ if text_config is not None
150
+ else CONFIG_MAPPING[LLM_BACKBONE_TO_HF_METACLASS[self.llm_backbone_id]]()
151
+ )
152
+
153
+ # Dispatch **kwargs to super() =>> note that `pad_token_id` collides, so we pass it in here as well...
154
+ super().__init__(pad_token_id=pad_token_id, **kwargs)
155
+
156
+ # Here we need trajectory_vla config, with
157
+ # prismatic_config fields and then the waypointer fields
158
+
159
+ class TrajectoryVLAConfig(PretrainedConfig):
160
+ model_type: str = "trajectoryvla"
161
+ is_composition: bool = True
162
+ def __init__(
163
+ self,
164
+ prismatic_config = {},
165
+ token_size: int = 1024, # Timestep token size
166
+ cheat: bool = False, # If True, cheat and use action tokens; Works only with OpenVLA checkpoint
167
+ num_timesteps: int = 20, # Number of prediction time steps
168
+ rotation_components: int = 9, # Number of rotation componens: euler -> 3, quaternion -> 4, rotmat -> 9
169
+ num_timestep_tokens : int = 3,
170
+ seperate_control_proj: bool = True, # If True, project control components separately
171
+ timestep_proj_config: Dict[str, Any] = {},
172
+ token_proj_config: Dict[str, Any] = {},
173
+ transformer_config: Dict[str, Any] = {},
174
+ # prismatic_config: PrismaticConfig,
175
+ # waypointer_config: Dict[str, Any],
176
+ # **kwargs: str,
177
+ ):
178
+
179
+ # super().__init__(**prismatic_config)
180
+ self.prismatic_config = PrismaticConfig(**prismatic_config)
181
+
182
+ self.token_size = token_size
183
+ self.cheat = cheat
184
+ self.num_timesteps = num_timesteps
185
+ self.rotation_components = rotation_components
186
+ self.seperate_control_proj = seperate_control_proj
187
+ self.timestep_proj_config = timestep_proj_config
188
+ self.token_proj_config = token_proj_config
189
+ self.transformer_config = transformer_config
190
+ # self.num_timestep_tokens = num_timestep_tokens
191
+
192
+ @property
193
+ def control_components(self) -> int:
194
+ # Number of control dimensions: 3 translation, N rotation, 1 gripper
195
+ return 3 + self.rotation_components + 1
196
+
197
+ @property
198
+ def num_timestep_tokens(self) -> int:
199
+ return self.timestep_proj_config['num_tokens']
200
+ # class WaypointerConfig(ConfigurableModuleConfig):
201
+ # token_size: int = 1024 # Timestep token size
202
+
203
+ # cheat: bool # If True, cheat and use action tokens; Works only with OpenVLA checkpoint
204
+
205
+ # timestep_proj_config: AutoConfig # Timestep tokens
206
+ # token_proj_config: TokenProjectorConfig # LLM output tokens projection and packing
207
+ # transformer_config: AutoConfig # Transformer config
208
+
209
+ # # Output configurations
210
+ # num_timesteps: int = 20 # Number of prediction time steps
211
+ # rotation_components: int = 3 # Number of rotation componens: euler -> 3, quaternion -> 4, rotmat -> 9
212
+ # separate_control_proj: bool = True # If True, project control components separately
213
+
214
+ # @property
215
+ # def control_components(self) -> int:
216
+ # # Number of control dimensions: 3 translation, N rotation, 1 gripper
217
+ # return 3 + self.rotation_components + 1
218
+
219
+ # @property
220
+ # def num_timestep_tokens(self) -> int:
221
+ # return self.timestep_proj_config.num_tokens
222
+
223
+
224
+ class OpenVLAConfig(PrismaticConfig):
225
+ model_type: str = "openvla"
226
+
227
+ def __init__(
228
+ self,
229
+ norm_stats: Optional[Dict[str, Dict[str, Dict[str, Dict[str, List[float]]]]]] = None,
230
+ n_action_bins: int = 256,
231
+ **kwargs: str,
232
+ ) -> None:
233
+ self.norm_stats, self.n_action_bins = norm_stats, n_action_bins
234
+
235
+ super().__init__(**kwargs)
236
+
237
+ if __name__ == "__main__" :
238
+ # yaml_file = 'barrel/pipes/vlams/configs/waypoints/waypointer_multistep_fractal.yaml'
239
+
240
+ prismatic_config = PrismaticConfig()
241
+ print(prismatic_config)
242
+
243
+ prismatic_config_dict = {
244
+ "vision_backbone_id":"dinosiglip-vit-so-224px",
245
+ # "llm_backbone_id":"llama2-7b-pure",meta-llama/Llama-2-7b-hf
246
+ "llm_backbone_id": "meta-llama/Llama-2-7b-hf",
247
+
248
+ "arch_specifier": "no-align+gelu-mlp", ## TODO: check
249
+ "use_fused_vision_backbone" :None, ## TODO: check
250
+ "image_resize_strategy" : "letterbox",
251
+ "text_config" : None,
252
+ "llm_max_length" : 2048,
253
+ "pad_token_id" :32000,
254
+ "pad_to_multiple_of" : 64,
255
+ "output_projector_states" : False,
256
+ }
257
+ token_proj_config = {
258
+ "vit_tokens_layers": [2176, 1024],
259
+ "control_tokens_layers": [4096, 2048, 1024],
260
+ "image_tokens_mode": 'vit',
261
+ }
262
+ timestep_proj_config = {
263
+ "pos_embed_scale": 1.0,
264
+ "proj_layers": [1024],
265
+ "time_delta_sec": 0.1,
266
+ "num_tokens":3
267
+ }
268
+
269
+ TrajectoryVlaConfig = {
270
+ "prismatic_config":prismatic_config_dict,
271
+ "token_size": 1024,
272
+ "cheat": False,
273
+ "num_timesteps": 20,
274
+ "rotation_components": 3,
275
+ "seperate_control_proj": True,
276
+ "timestep_proj_config": {},
277
+ "token_proj_config": {},
278
+ "transformer_config": {},
279
+ }
280
+
281
+ TrajectoryVLAConfig = TrajectoryVLAConfig( **TrajectoryVlaConfig)
282
+ print(TrajectoryVLAConfig)
283
+
284
+ class WaypointTokenizer:
285
+ """
286
+ Wraps base LLM/VLM tokenizer and overloads least used token as a control token
287
+
288
+ NOTE: By default, assumes a BPE-style tokenizer akin to the LlamaTokenizer,
289
+ where *the least used tokens* appear at the end of the vocabulary!
290
+
291
+ TODO: Adding new token vs overloading? When I call `tokenizer.add_token()` vocab stays the same
292
+ """
293
+
294
+ def __init__(self, tokenizer: transformers.PreTrainedTokenizerBase, num_tokens: int = 10) -> None:
295
+ self.tokenizer = tokenizer
296
+ self.num_tokens = num_tokens
297
+
298
+ def __call__(self, *_) -> str:
299
+ """Get the text token for control"""
300
+ return self.tokenizer.decode(self.control_token_ids)
301
+
302
+ @property
303
+ def control_token_ids(self) -> np.ndarray:
304
+ # Assumes we're overwriting the final tokens of the vocabulary (least used tokens)
305
+ return np.arange(self.num_tokens) + int(self.tokenizer.vocab_size - self.num_tokens)
306
+
307
+ @property
308
+ def num_control_tokens(self) -> int:
309
+ return self.num_tokens