├── LICENSE ├── README.md ├── assets ├── model_structure.png ├── radar.png └── visualization_vp.png ├── example.jpg ├── inference.py ├── requirements.txt └── visual_words ├── __init__.py ├── constants.py ├── conversation.py ├── mm_utils.py ├── model ├── __init__.py ├── builder.py ├── language_model │ ├── vw_llama.py │ ├── vw_mistral.py │ └── vw_pif_llama.py ├── multimodal_encoder │ ├── builder.py │ └── clip_encoder.py ├── multimodal_projector │ └── builder.py ├── vw_arch.py └── vw_pif_arch.py └── utils.py /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Multi-modal Auto-regressive Modeling via Visual Words 2 | 3 | [[`arXiv`](https://arxiv.org/abs/2403.07720)] [[`BibTeX`](#Citing)] 4 | 5 | This is the official repository for the multi-modal large language models: VW-LMM 6 | 7 |
8 | 9 |

10 | 11 | ## Introduction 12 | We propose VW-LMM, a large multi-modal model (LMM) that successfully performs multi-modal auto-regressive modeling with a unified objective for the first time. 13 | Specifically, we propose the concept of visual words, which maps the visual features to probability distributions over LLM's vocabulary, providing supervision information for visual modelling. 14 | We further explore the distribution of visual features in the semantic space within LMM and the possibility of using text embeddings to represent visual information. 15 | Experimental results and ablation studies on 5 VQA tasks and 4 benchmark toolkits validate the powerful performance of our proposed approach. 16 | For more technical details, please refer to our [paper](https://arxiv.org/abs/2403.07720). 17 | 18 |
19 | 20 |

21 | 22 | In order to verify whether the visual words learnt by VW-LMM can realistically reflect the image information, we take VW-LMM-Vicuna-7B as an example to explore. 23 | For each patch in the image, we select the token with the highest probability in its corresponding visual words, and compare the region of interest in the image with its visualisation result, visualization is as follows (Best viewed zoomed-in): 24 | 25 |
26 | 27 |

28 | 29 | ## Model Zoo 30 | 31 | | Version | Size | Support pseudo image features | Checkpoint | 32 | |-------------------|----------|-----------------------------------|---------------------------------------------------------------------------------------| 33 | | VW-LMM-Vicuna | 7B | False | [VW-LMM-Vicuna-7b](https://huggingface.co/MYTH-Lab/VW-LMM-Vicuna-7b) | 34 | | VW-LMM-Mistral | 7B | False | [VW-LMM-Mistral-7b](https://huggingface.co/MYTH-Lab/VW-LMM-Mistral-7b) | 35 | | VW-LMM-Vicuna-pif | 7B | True | [VW-LMM-Vicuna-pif-7b](https://huggingface.co/MYTH-Lab/VW-LMM-Vicuna-pif-7b) | 36 | 37 | 38 | VW-LMM, by constructing visual words to introduce visual supervisory information, achieves the best performance among models of the same scale of 7B, and obtains vision-language understanding capability competitive to or even surpassing that of 13B or even larger scale models. 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | 68 | 69 | 70 | 71 | 72 | 73 | 74 | 75 | 76 | 77 | 78 | 79 | 80 | 81 | 82 | 83 | 84 | 85 | 86 | 87 | 88 | 89 | 90 | 91 | 92 | 93 | 94 | 95 | 96 | 97 | 98 | 99 | 100 | 101 | 102 | 103 | 104 | 105 | 106 | 107 | 108 | 109 | 110 | 111 | 112 | 113 | 114 | 115 | 116 | 117 | 118 | 119 | 120 | 121 | 122 | 123 | 124 | 125 | 126 | 127 | 128 | 129 | 130 | 131 | 132 | 133 | 134 | 135 | 136 | 137 | 138 | 139 | 140 | 141 | 142 | 143 | 144 | 145 | 146 | 147 | 148 | 149 | 150 | 151 | 152 | 153 | 154 | 155 | 156 | 157 | 158 | 159 | 160 | 161 | 162 | 163 | 164 | 165 | 166 | 167 | 168 | 169 | 170 | 171 | 172 | 173 | 174 | 175 | 176 | 177 | 178 | 179 | 180 | 181 | 182 | 183 | 184 | 185 | 186 | 187 | 188 | 189 | 190 | 191 | 192 | 193 | 194 | 195 | 196 | 197 | 198 | 199 | 200 | 201 | 202 | 203 | 204 | 205 | 206 | 207 | 208 | 209 | 210 | 211 | 212 | 213 | 214 | 215 | 216 | 217 | 218 | 219 | 220 | 221 | 222 | 223 | 224 | 225 | 226 | 227 | 228 | 229 | 230 | 231 | 232 | 233 | 234 | 235 | 236 | 237 | 238 | 239 | 240 | 241 | 242 | 243 | 244 | 245 | 246 | 247 | 248 | 249 | 250 | 251 | 252 | 253 | 254 | 255 | 256 | 257 | 258 | 259 | 260 | 261 | 262 | 263 | 264 | 265 | 266 | 267 | 268 | 269 | 270 | 271 | 272 | 273 | 274 | 275 | 276 | 277 | 278 | 279 | 280 | 281 | 282 | 283 | 284 | 285 | 286 | 287 | 288 | 289 | 290 | 291 | 292 | 293 | 294 | 295 | 296 | 297 | 298 | 299 | 300 | 301 | 302 | 303 | 304 | 305 | 306 | 307 | 308 | 309 | 310 | 311 | 312 | 313 | 314 | 315 | 316 | 317 | 318 | 319 | 320 | 321 | 322 | 323 | 324 | 325 | 326 | 327 |
MethodsLLMRes.VQA^v2GQAVisWizSQA^IVQA^TPOPEMMBMMB^CNMM-Vet
*Language Modeling Method*
IDEFICS-80BLLaMA-65B22460.045.236.0--30.9--54.538.1--
InstructBLIPVicuna-13B224--49.533.463.150.778.9----25.6
BLIP-2Vicuna-13B22441.041.019.661.042.585.3----22.4
LLaVA-v1.5Vicuna-13B33680.063.353.671.661.385.967.763.635.4
InstructBLIPVicuna-7B224--49.234.560.550.1--3623.726.2
IDEFICS-9BLLaMA-7B22450.938.435.5--25.9--48.225.2--
Qwen-VLQwen-7B44878.859.335.267.163.8--38.27.4--
Qwen-VL-ChatQwen-7B44878.257.538.968.261.5--60.656.7--
LLaVA-v1.5Vicuna-7B33678.562.050.066.858.285.964.358.330.5
MoE-LLaVA-2.7B×4-Top2Phi-2-2.7B33677.661.443.968.551.486.365.2--34.3
*Multi-modal Modeling Method*
Emu2-ChatLLaMA-33B44884.965.154.965.566.6------48.5
Emu-ILLaMA-13B22462.046.038.3----------36.3
MM-Interleaved-SFTVicuna-13B22480.260.554.9--61.0--------
Unified-IO 2UIO-2-6.8B38479.4----86.2--87.771.5----
DreamLLMVicuna-7B22456.6--38.1--34.9--------
VL-GPT-ILLaMA-7B22467.251.538.9------------
LaVIT-v2LLaMA2-7B22468.347.941.0------------
VW-LMMVicuna-7B33678.962.748.368.157.685.965.959.831.3
VW-LMMMistral-7B33680.865.458.575.963.187.080.679.044.0
328 | 329 | ## Setup 330 | 331 | ### Requirements 332 | 333 | ```shell 334 | git clone https://github.com/pengts/VW-LMM.git 335 | cd VW-LMM 336 | pip install -r requirements.txt 337 | ``` 338 | 339 | ## Multi-modal Inference 340 | 341 | ### Model Configurations 342 | 343 | - VW-LMM-Vicuna 344 | ```python 345 | model_path="VW-LMM-Vicuna" 346 | conv_mode="vicuna_v1" 347 | model_base="llama" 348 | device = "cuda" 349 | ``` 350 | - VW-LMM-Mistral 351 | ```python 352 | model_path="VW-LMM-Mistral" 353 | conv_mode="mistral" 354 | model_base="mistral" 355 | device = "cuda" 356 | ``` 357 | 358 | VW-LMM-Vicuna-pif 359 | ```python 360 | model_path="VW-LMM-Vicuna-pif" 361 | conv_mode="vicuna_v1" 362 | model_base="llama" 363 | device = "cuda" 364 | ``` 365 | 366 | ### Model Initialization 367 | ```python 368 | disable_torch_init() 369 | model_path = os.path.expanduser(model_path) 370 | model_name = get_model_name_from_path(model_path) 371 | tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, model_name, model_base,device=device) 372 | ``` 373 | 374 | ### Input Processing 375 | ```python 376 | question="Write an exhaustive depiction of the given image." 377 | image_path="./example.jpg" 378 | qs = question 379 | qs = DEFAULT_IMAGE_TOKEN + '\n' + qs 380 | conv = conv_templates[conv_mode].copy() 381 | conv.append_message(conv.roles[0], qs) 382 | conv.append_message(conv.roles[1], None) 383 | prompt = conv.get_prompt() 384 | 385 | image = Image.open(image_path).convert('RGB') 386 | image_tensor = process_images([image], image_processor, model.config)[0].unsqueeze(0).to(device) 387 | input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(device) 388 | ``` 389 | 390 | ### Inference 391 | ```python 392 | with torch.inference_mode(): 393 | output_ids = model.generate( 394 | input_ids, 395 | images=image_tensor.to(dtype=torch.float16, device=device, non_blocking=True), 396 | do_sample= False, 397 | temperature=0, 398 | top_p=None, 399 | num_beams=1, 400 | max_new_tokens=128, 401 | use_cache=True) 402 | 403 | input_token_len = input_ids.shape[1] 404 | n_diff_input_output = (input_ids != output_ids[:, :input_token_len]).sum().item() 405 | if n_diff_input_output > 0: 406 | print(f'[Warning] {n_diff_input_output} output_ids are not the same as the input_ids') 407 | outputs = tokenizer.batch_decode(output_ids[:, input_token_len:], skip_special_tokens=True)[0] 408 | outputs = outputs.strip() 409 | print(outputs) 410 | ``` 411 | 412 | ## Acknowledgement 413 | We are grateful for the following awesome projects when implementing VW-LMM: 414 | * [LLaVA](https://github.com/haotian-liu/LLaVA/): Visual instruction tuning towards large language and vision models with GPT-4 level capabilities. 415 | * [LLaMA](https://github.com/facebookresearch/llama): Open and Efficient Foundation Language Models 416 | * [Vicuna](https://github.com/lm-sys/FastChat): Open-source LLM with amazing language capabilities! 417 | * [Mistral](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2): A 7B transformer model, fast-deployed and easily customisable. Small, yet very powerful for a variety of use cases. 418 | 419 | 420 | ## Citation 421 | Consider giving this repository a star and cite VW-LMM in your publications if it helps your research. 422 | 423 | ``` 424 | @misc{peng2024multimodal, 425 | title={Multi-modal Auto-regressive Modeling via Visual Words}, 426 | author={Tianshuo Peng and Zuchao Li and Lefei Zhang and Hai Zhao and Ping Wang and Bo Du}, 427 | year={2024}, 428 | eprint={2403.07720}, 429 | archivePrefix={arXiv}, 430 | primaryClass={cs.CV} 431 | } 432 | ``` 433 | -------------------------------------------------------------------------------- /assets/model_structure.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pengts/VW-LMM/787cf71cf9429cfb5845277ddc1a427170034325/assets/model_structure.png -------------------------------------------------------------------------------- /assets/radar.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pengts/VW-LMM/787cf71cf9429cfb5845277ddc1a427170034325/assets/radar.png -------------------------------------------------------------------------------- /assets/visualization_vp.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pengts/VW-LMM/787cf71cf9429cfb5845277ddc1a427170034325/assets/visualization_vp.png -------------------------------------------------------------------------------- /example.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pengts/VW-LMM/787cf71cf9429cfb5845277ddc1a427170034325/example.jpg -------------------------------------------------------------------------------- /inference.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import os 3 | import json 4 | from tqdm import tqdm 5 | from visual_words.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN 6 | from visual_words.conversation import conv_templates, SeparatorStyle 7 | from visual_words.model.builder import load_pretrained_model 8 | from visual_words.utils import disable_torch_init 9 | from visual_words.mm_utils import tokenizer_image_token, process_images, get_model_name_from_path 10 | from PIL import Image 11 | 12 | 13 | 14 | # Configurations 15 | 16 | # model_path="path_to_model_weights" 17 | # question="Write an exhaustive depiction of the given image." 18 | # image_path="./example.jpg" 19 | # conv_mode="vicuna_v1" 20 | # model_base="llama" 21 | # device = "cuda" 22 | 23 | # model_path="path_to_model_weights" 24 | # question="Write an exhaustive depiction of the given image." 25 | # image_path="./example.jpg" 26 | # conv_mode="mistral" 27 | # model_base="mistral" 28 | # device = "cuda" 29 | 30 | model_path="path_to_model_weights" 31 | question="Write an exhaustive depiction of the given image." 32 | image_path="./example.jpg" 33 | conv_mode="vicuna_v1" 34 | model_base="llama" 35 | device = "cuda" 36 | 37 | # Model 38 | disable_torch_init() 39 | model_path = os.path.expanduser(model_path) 40 | model_name = get_model_name_from_path(model_path) 41 | tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, model_name, model_base,device=device) 42 | 43 | # input 44 | qs = question 45 | qs = DEFAULT_IMAGE_TOKEN + '\n' + qs 46 | conv = conv_templates[conv_mode].copy() 47 | conv.append_message(conv.roles[0], qs) 48 | conv.append_message(conv.roles[1], None) 49 | prompt = conv.get_prompt() 50 | 51 | 52 | image = Image.open(image_path).convert('RGB') 53 | image_tensor = process_images([image], image_processor, model.config)[0].unsqueeze(0).to(device) 54 | input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(device) 55 | 56 | with torch.inference_mode(): 57 | output_ids = model.generate( 58 | input_ids, 59 | images=image_tensor.to(dtype=torch.float16, device=device, non_blocking=True), 60 | do_sample= False, 61 | temperature=0, 62 | top_p=None, 63 | num_beams=1, 64 | max_new_tokens=128, 65 | use_cache=True) 66 | 67 | input_token_len = input_ids.shape[1] 68 | n_diff_input_output = (input_ids != output_ids[:, :input_token_len]).sum().item() 69 | if n_diff_input_output > 0: 70 | print(f'[Warning] {n_diff_input_output} output_ids are not the same as the input_ids') 71 | outputs = tokenizer.batch_decode(output_ids[:, input_token_len:], skip_special_tokens=True)[0] 72 | outputs = outputs.strip() 73 | print(outputs) 74 | 75 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | einops 2 | fastapi 3 | gradio==3.35.2 4 | markdown2[all] 5 | numpy 6 | requests 7 | sentencepiece 8 | tokenizers>=0.12.1 9 | torch==2.0.1 10 | torchvision==0.15.2 11 | uvicorn 12 | wandb 13 | shortuuid 14 | httpx==0.24.0 15 | deepspeed==0.9.5 16 | peft==0.4.0 17 | transformers==4.35.0 18 | accelerate==0.21.0 19 | bitsandbytes==0.41.0 20 | scikit-learn==1.2.2 21 | sentencepiece==0.1.99 22 | einops==0.6.1 23 | einops-exts==0.0.4 24 | timm==0.6.13 25 | gradio_client==0.2.9 -------------------------------------------------------------------------------- /visual_words/__init__.py: -------------------------------------------------------------------------------- 1 | from .model import vwMistralForCausalLM,vwLlamaForCausalLM,vwPifLlamaForCausalLM 2 | -------------------------------------------------------------------------------- /visual_words/constants.py: -------------------------------------------------------------------------------- 1 | CONTROLLER_HEART_BEAT_EXPIRATION = 30 2 | WORKER_HEART_BEAT_INTERVAL = 15 3 | 4 | LOGDIR = "." 5 | 6 | # Model Constants 7 | IGNORE_INDEX = -100 8 | IMAGE_TOKEN_INDEX = -200 9 | DEFAULT_IMAGE_TOKEN = "" 10 | DEFAULT_IMAGE_PATCH_TOKEN = "" 11 | DEFAULT_IM_START_TOKEN = "" 12 | DEFAULT_IM_END_TOKEN = "" 13 | -------------------------------------------------------------------------------- /visual_words/conversation.py: -------------------------------------------------------------------------------- 1 | import dataclasses 2 | from enum import auto, Enum 3 | from typing import List, Tuple 4 | 5 | 6 | class SeparatorStyle(Enum): 7 | """Different separator style.""" 8 | SINGLE = auto() 9 | TWO = auto() 10 | MPT = auto() 11 | PLAIN = auto() 12 | LLAMA_2 = auto() 13 | MISTRAL = auto() 14 | 15 | 16 | @dataclasses.dataclass 17 | class Conversation: 18 | """A class that keeps all conversation history.""" 19 | system: str 20 | roles: List[str] 21 | messages: List[List[str]] 22 | offset: int 23 | sep_style: SeparatorStyle = SeparatorStyle.SINGLE 24 | sep: str = "###" 25 | sep2: str = None 26 | version: str = "Unknown" 27 | 28 | skip_next: bool = False 29 | 30 | def get_prompt(self): 31 | messages = self.messages 32 | if len(messages) > 0 and type(messages[0][1]) is tuple: 33 | messages = self.messages.copy() 34 | init_role, init_msg = messages[0].copy() 35 | init_msg = init_msg[0].replace("", "").strip() 36 | if 'mmtag' in self.version: 37 | messages[0] = (init_role, init_msg) 38 | messages.insert(0, (self.roles[0], "")) 39 | messages.insert(1, (self.roles[1], "Received.")) 40 | else: 41 | messages[0] = (init_role, "\n" + init_msg) 42 | 43 | if self.sep_style == SeparatorStyle.SINGLE: 44 | ret = self.system + self.sep 45 | for role, message in messages: 46 | if message: 47 | if type(message) is tuple: 48 | message, _, _ = message 49 | ret += role + ": " + message + self.sep 50 | else: 51 | ret += role + ":" 52 | elif self.sep_style == SeparatorStyle.TWO: 53 | seps = [self.sep, self.sep2] 54 | ret = self.system + seps[0] 55 | for i, (role, message) in enumerate(messages): 56 | if message: 57 | if type(message) is tuple: 58 | message, _, _ = message 59 | ret += role + ": " + message + seps[i % 2] 60 | else: 61 | ret += role + ":" 62 | elif self.sep_style == SeparatorStyle.MPT: 63 | ret = self.system + self.sep 64 | for role, message in messages: 65 | if message: 66 | if type(message) is tuple: 67 | message, _, _ = message 68 | ret += role + message + self.sep 69 | else: 70 | ret += role 71 | elif self.sep_style == SeparatorStyle.LLAMA_2: 72 | wrap_sys = lambda msg: f"<>\n{msg}\n<>\n\n" 73 | wrap_inst = lambda msg: f"[INST] {msg} [/INST]" 74 | ret = "" 75 | 76 | for i, (role, message) in enumerate(messages): 77 | if i == 0: 78 | assert message, "first message should not be none" 79 | assert role == self.roles[0], "first message should come from user" 80 | if message: 81 | if type(message) is tuple: 82 | message, _, _ = message 83 | if i == 0: message = wrap_sys(self.system) + message 84 | if i % 2 == 0: 85 | message = wrap_inst(message) 86 | ret += self.sep + message 87 | else: 88 | ret += " " + message + " " + self.sep2 89 | else: 90 | ret += "" 91 | ret = ret.lstrip(self.sep) 92 | 93 | elif self.sep_style == SeparatorStyle.MISTRAL: 94 | wrap_inst = lambda msg: f"[INST] {msg} [/INST]" 95 | ret = "" 96 | for i, (role, message) in enumerate(messages): 97 | if i == 0: 98 | assert message, "first message should not be none" 99 | assert role == self.roles[0], "first message should come from user" 100 | if message: 101 | if type(message) is tuple: 102 | message, _, _ = message 103 | # if i == 0: message = wrap_sys(self.system) + message 104 | if i % 2 == 0: 105 | message = wrap_inst(message) 106 | ret += self.sep + message 107 | else: 108 | ret += message + self.sep2 109 | else: 110 | ret += "" 111 | ret = ret.lstrip(self.sep) 112 | 113 | elif self.sep_style == SeparatorStyle.PLAIN: 114 | seps = [self.sep, self.sep2] 115 | ret = self.system 116 | for i, (role, message) in enumerate(messages): 117 | if message: 118 | if type(message) is tuple: 119 | message, _, _ = message 120 | ret += message + seps[i % 2] 121 | else: 122 | ret += "" 123 | else: 124 | raise ValueError(f"Invalid style: {self.sep_style}") 125 | 126 | return ret 127 | 128 | def append_message(self, role, message): 129 | self.messages.append([role, message]) 130 | 131 | def get_images(self, return_pil=False): 132 | images = [] 133 | for i, (role, msg) in enumerate(self.messages[self.offset:]): 134 | if i % 2 == 0: 135 | if type(msg) is tuple: 136 | import base64 137 | from io import BytesIO 138 | from PIL import Image 139 | msg, image, image_process_mode = msg 140 | if image_process_mode == "Pad": 141 | def expand2square(pil_img, background_color=(122, 116, 104)): 142 | width, height = pil_img.size 143 | if width == height: 144 | return pil_img 145 | elif width > height: 146 | result = Image.new(pil_img.mode, (width, width), background_color) 147 | result.paste(pil_img, (0, (width - height) // 2)) 148 | return result 149 | else: 150 | result = Image.new(pil_img.mode, (height, height), background_color) 151 | result.paste(pil_img, ((height - width) // 2, 0)) 152 | return result 153 | image = expand2square(image) 154 | elif image_process_mode in ["Default", "Crop"]: 155 | pass 156 | elif image_process_mode == "Resize": 157 | image = image.resize((336, 336)) 158 | else: 159 | raise ValueError(f"Invalid image_process_mode: {image_process_mode}") 160 | max_hw, min_hw = max(image.size), min(image.size) 161 | aspect_ratio = max_hw / min_hw 162 | max_len, min_len = 800, 400 163 | shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw)) 164 | longest_edge = int(shortest_edge * aspect_ratio) 165 | W, H = image.size 166 | if longest_edge != max(image.size): 167 | if H > W: 168 | H, W = longest_edge, shortest_edge 169 | else: 170 | H, W = shortest_edge, longest_edge 171 | image = image.resize((W, H)) 172 | if return_pil: 173 | images.append(image) 174 | else: 175 | buffered = BytesIO() 176 | image.save(buffered, format="PNG") 177 | img_b64_str = base64.b64encode(buffered.getvalue()).decode() 178 | images.append(img_b64_str) 179 | return images 180 | 181 | def to_gradio_chatbot(self): 182 | ret = [] 183 | for i, (role, msg) in enumerate(self.messages[self.offset:]): 184 | if i % 2 == 0: 185 | if type(msg) is tuple: 186 | import base64 187 | from io import BytesIO 188 | msg, image, image_process_mode = msg 189 | max_hw, min_hw = max(image.size), min(image.size) 190 | aspect_ratio = max_hw / min_hw 191 | max_len, min_len = 800, 400 192 | shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw)) 193 | longest_edge = int(shortest_edge * aspect_ratio) 194 | W, H = image.size 195 | if H > W: 196 | H, W = longest_edge, shortest_edge 197 | else: 198 | H, W = shortest_edge, longest_edge 199 | image = image.resize((W, H)) 200 | buffered = BytesIO() 201 | image.save(buffered, format="JPEG") 202 | img_b64_str = base64.b64encode(buffered.getvalue()).decode() 203 | img_str = f'user upload image' 204 | msg = img_str + msg.replace('', '').strip() 205 | ret.append([msg, None]) 206 | else: 207 | ret.append([msg, None]) 208 | else: 209 | ret[-1][-1] = msg 210 | return ret 211 | 212 | def copy(self): 213 | return Conversation( 214 | system=self.system, 215 | roles=self.roles, 216 | messages=[[x, y] for x, y in self.messages], 217 | offset=self.offset, 218 | sep_style=self.sep_style, 219 | sep=self.sep, 220 | sep2=self.sep2, 221 | version=self.version) 222 | 223 | def dict(self): 224 | if len(self.get_images()) > 0: 225 | return { 226 | "system": self.system, 227 | "roles": self.roles, 228 | "messages": [[x, y[0] if type(y) is tuple else y] for x, y in self.messages], 229 | "offset": self.offset, 230 | "sep": self.sep, 231 | "sep2": self.sep2, 232 | } 233 | return { 234 | "system": self.system, 235 | "roles": self.roles, 236 | "messages": self.messages, 237 | "offset": self.offset, 238 | "sep": self.sep, 239 | "sep2": self.sep2, 240 | } 241 | 242 | 243 | conv_vicuna_v0 = Conversation( 244 | system="A chat between a curious human and an artificial intelligence assistant. " 245 | "The assistant gives helpful, detailed, and polite answers to the human's questions.", 246 | roles=("Human", "Assistant"), 247 | messages=( 248 | ("Human", "What are the key differences between renewable and non-renewable energy sources?"), 249 | ("Assistant", 250 | "Renewable energy sources are those that can be replenished naturally in a relatively " 251 | "short amount of time, such as solar, wind, hydro, geothermal, and biomass. " 252 | "Non-renewable energy sources, on the other hand, are finite and will eventually be " 253 | "depleted, such as coal, oil, and natural gas. Here are some key differences between " 254 | "renewable and non-renewable energy sources:\n" 255 | "1. Availability: Renewable energy sources are virtually inexhaustible, while non-renewable " 256 | "energy sources are finite and will eventually run out.\n" 257 | "2. Environmental impact: Renewable energy sources have a much lower environmental impact " 258 | "than non-renewable sources, which can lead to air and water pollution, greenhouse gas emissions, " 259 | "and other negative effects.\n" 260 | "3. Cost: Renewable energy sources can be more expensive to initially set up, but they typically " 261 | "have lower operational costs than non-renewable sources.\n" 262 | "4. Reliability: Renewable energy sources are often more reliable and can be used in more remote " 263 | "locations than non-renewable sources.\n" 264 | "5. Flexibility: Renewable energy sources are often more flexible and can be adapted to different " 265 | "situations and needs, while non-renewable sources are more rigid and inflexible.\n" 266 | "6. Sustainability: Renewable energy sources are more sustainable over the long term, while " 267 | "non-renewable sources are not, and their depletion can lead to economic and social instability.\n") 268 | ), 269 | offset=2, 270 | sep_style=SeparatorStyle.SINGLE, 271 | sep="###", 272 | ) 273 | 274 | conv_vicuna_v1 = Conversation( 275 | system="A chat between a curious user and an artificial intelligence assistant. " 276 | "The assistant gives helpful, detailed, and polite answers to the user's questions.", 277 | roles=("USER", "ASSISTANT"), 278 | version="v1", 279 | messages=(), 280 | offset=0, 281 | sep_style=SeparatorStyle.TWO, 282 | sep=" ", 283 | sep2="", 284 | ) 285 | 286 | conv_llama_2 = Conversation( 287 | system="""You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature. 288 | 289 | If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.""", 290 | roles=("USER", "ASSISTANT"), 291 | version="llama_v2", 292 | messages=(), 293 | offset=0, 294 | sep_style=SeparatorStyle.LLAMA_2, 295 | sep="", 296 | sep2="", 297 | ) 298 | 299 | conv_llava_llama_2 = Conversation( 300 | system="You are a helpful language and vision assistant. " 301 | "You are able to understand the visual content that the user provides, " 302 | "and assist the user with a variety of tasks using natural language.", 303 | roles=("USER", "ASSISTANT"), 304 | version="llama_v2", 305 | messages=(), 306 | offset=0, 307 | sep_style=SeparatorStyle.LLAMA_2, 308 | sep="", 309 | sep2="", 310 | ) 311 | 312 | conv_mpt = Conversation( 313 | system="""<|im_start|>system 314 | A conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.""", 315 | roles=("<|im_start|>user\n", "<|im_start|>assistant\n"), 316 | version="mpt", 317 | messages=(), 318 | offset=0, 319 | sep_style=SeparatorStyle.MPT, 320 | sep="<|im_end|>", 321 | ) 322 | 323 | conv_llava_plain = Conversation( 324 | system="", 325 | roles=("", ""), 326 | messages=( 327 | ), 328 | offset=0, 329 | sep_style=SeparatorStyle.PLAIN, 330 | sep="\n", 331 | ) 332 | 333 | conv_llava_v0 = Conversation( 334 | system="A chat between a curious human and an artificial intelligence assistant. " 335 | "The assistant gives helpful, detailed, and polite answers to the human's questions.", 336 | roles=("Human", "Assistant"), 337 | messages=( 338 | ), 339 | offset=0, 340 | sep_style=SeparatorStyle.SINGLE, 341 | sep="###", 342 | ) 343 | 344 | conv_llava_v0_mmtag = Conversation( 345 | system="A chat between a curious user and an artificial intelligence assistant. " 346 | "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language." 347 | "The visual content will be provided with the following format: visual content.", 348 | roles=("Human", "Assistant"), 349 | messages=( 350 | ), 351 | offset=0, 352 | sep_style=SeparatorStyle.SINGLE, 353 | sep="###", 354 | version="v0_mmtag", 355 | ) 356 | 357 | conv_llava_v1 = Conversation( 358 | system="A chat between a curious human and an artificial intelligence assistant. " 359 | "The assistant gives helpful, detailed, and polite answers to the human's questions.", 360 | roles=("USER", "ASSISTANT"), 361 | version="v1", 362 | messages=(), 363 | offset=0, 364 | sep_style=SeparatorStyle.TWO, 365 | sep=" ", 366 | sep2="", 367 | ) 368 | 369 | conv_llava_v1_mmtag = Conversation( 370 | system="A chat between a curious user and an artificial intelligence assistant. " 371 | "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language." 372 | "The visual content will be provided with the following format: visual content.", 373 | roles=("USER", "ASSISTANT"), 374 | messages=(), 375 | offset=0, 376 | sep_style=SeparatorStyle.TWO, 377 | sep=" ", 378 | sep2="", 379 | version="v1_mmtag", 380 | ) 381 | 382 | mistral = Conversation( 383 | system="", 384 | roles=("USER", "ASSISTANT"), 385 | version="mistral", 386 | messages=(), 387 | offset=0, 388 | sep_style=SeparatorStyle.MISTRAL, 389 | sep="", 390 | sep2="", 391 | ) 392 | 393 | 394 | default_conversation = conv_vicuna_v0 395 | conv_templates = { 396 | "default": conv_vicuna_v0, 397 | "v0": conv_vicuna_v0, 398 | "v1": conv_vicuna_v1, 399 | "vicuna_v1": conv_vicuna_v1, 400 | "llama_2": conv_llama_2, 401 | 402 | "plain": conv_llava_plain, 403 | "v0_plain": conv_llava_plain, 404 | "llava_v0": conv_llava_v0, 405 | "v0_mmtag": conv_llava_v0_mmtag, 406 | "llava_v1": conv_llava_v1, 407 | "v1_mmtag": conv_llava_v1_mmtag, 408 | "llava_llama_2": conv_llava_llama_2, 409 | 410 | "mpt": conv_mpt, 411 | "mistral": mistral 412 | } 413 | 414 | 415 | if __name__ == "__main__": 416 | print(default_conversation.get_prompt()) 417 | -------------------------------------------------------------------------------- /visual_words/mm_utils.py: -------------------------------------------------------------------------------- 1 | from PIL import Image 2 | from io import BytesIO 3 | import base64 4 | 5 | import torch 6 | from transformers import StoppingCriteria 7 | from visual_words.constants import IMAGE_TOKEN_INDEX 8 | 9 | 10 | def load_image_from_base64(image): 11 | return Image.open(BytesIO(base64.b64decode(image))) 12 | 13 | 14 | def expand2square(pil_img, background_color): 15 | width, height = pil_img.size 16 | if width == height: 17 | return pil_img 18 | elif width > height: 19 | result = Image.new(pil_img.mode, (width, width), background_color) 20 | result.paste(pil_img, (0, (width - height) // 2)) 21 | return result 22 | else: 23 | result = Image.new(pil_img.mode, (height, height), background_color) 24 | result.paste(pil_img, ((height - width) // 2, 0)) 25 | return result 26 | 27 | 28 | def process_images(images, image_processor, model_cfg): 29 | image_aspect_ratio = getattr(model_cfg, "image_aspect_ratio", None) 30 | new_images = [] 31 | if image_aspect_ratio == 'pad': 32 | for image in images: 33 | image = expand2square(image, tuple(int(x*255) for x in image_processor.image_mean)) 34 | image = image_processor.preprocess(image, return_tensors='pt')['pixel_values'][0] 35 | new_images.append(image) 36 | else: 37 | return image_processor(images, return_tensors='pt')['pixel_values'] 38 | if all(x.shape == new_images[0].shape for x in new_images): 39 | new_images = torch.stack(new_images, dim=0) 40 | return new_images 41 | 42 | 43 | def tokenizer_image_token(prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors=None): 44 | prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split('')] 45 | 46 | def insert_separator(X, sep): 47 | return [ele for sublist in zip(X, [sep]*len(X)) for ele in sublist][:-1] 48 | 49 | input_ids = [] 50 | offset = 0 51 | if len(prompt_chunks) > 0 and len(prompt_chunks[0]) > 0 and prompt_chunks[0][0] == tokenizer.bos_token_id: 52 | offset = 1 53 | input_ids.append(prompt_chunks[0][0]) 54 | 55 | for x in insert_separator(prompt_chunks, [image_token_index] * (offset + 1)): 56 | input_ids.extend(x[offset:]) 57 | 58 | if return_tensors is not None: 59 | if return_tensors == 'pt': 60 | return torch.tensor(input_ids, dtype=torch.long) 61 | raise ValueError(f'Unsupported tensor type: {return_tensors}') 62 | return input_ids 63 | 64 | 65 | def get_model_name_from_path(model_path): 66 | model_path = model_path.strip("/") 67 | model_paths = model_path.split("/") 68 | if model_paths[-1].startswith('checkpoint-'): 69 | return model_paths[-2] + "_" + model_paths[-1] 70 | else: 71 | return model_paths[-1] 72 | 73 | 74 | 75 | 76 | class KeywordsStoppingCriteria(StoppingCriteria): 77 | def __init__(self, keywords, tokenizer, input_ids): 78 | self.keywords = keywords 79 | self.keyword_ids = [] 80 | self.max_keyword_len = 0 81 | for keyword in keywords: 82 | cur_keyword_ids = tokenizer(keyword).input_ids 83 | if len(cur_keyword_ids) > 1 and cur_keyword_ids[0] == tokenizer.bos_token_id: 84 | cur_keyword_ids = cur_keyword_ids[1:] 85 | if len(cur_keyword_ids) > self.max_keyword_len: 86 | self.max_keyword_len = len(cur_keyword_ids) 87 | self.keyword_ids.append(torch.tensor(cur_keyword_ids)) 88 | self.tokenizer = tokenizer 89 | self.start_len = input_ids.shape[1] 90 | 91 | def __call__(self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool: 92 | assert output_ids.shape[0] == 1, "Only support batch size 1 (yet)" # TODO 93 | offset = min(output_ids.shape[1] - self.start_len, self.max_keyword_len) 94 | self.keyword_ids = [keyword_id.to(output_ids.device) for keyword_id in self.keyword_ids] 95 | for keyword_id in self.keyword_ids: 96 | if (output_ids[0, -keyword_id.shape[0]:] == keyword_id).all(): 97 | return True 98 | outputs = self.tokenizer.batch_decode(output_ids[:, -offset:], skip_special_tokens=True)[0] 99 | for keyword in self.keywords: 100 | if keyword in outputs: 101 | return True 102 | return False -------------------------------------------------------------------------------- /visual_words/model/__init__.py: -------------------------------------------------------------------------------- 1 | from .language_model.vw_mistral import vwMistralForCausalLM 2 | from .language_model.vw_llama import vwLlamaForCausalLM 3 | from .language_model.vw_pif_llama import vwPifLlamaForCausalLM -------------------------------------------------------------------------------- /visual_words/model/builder.py: -------------------------------------------------------------------------------- 1 | import os 2 | import warnings 3 | import shutil 4 | from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, BitsAndBytesConfig 5 | import torch 6 | from visual_words.model import * 7 | from visual_words.constants import DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN 8 | 9 | 10 | def load_pretrained_model(model_path, model_name, model_base="llama", load_8bit=False, load_4bit=False, device_map="auto", device="cuda"): 11 | kwargs = {"device_map": device_map} 12 | 13 | if load_8bit: 14 | kwargs['load_in_8bit'] = True 15 | elif load_4bit: 16 | kwargs['load_in_4bit'] = True 17 | kwargs['quantization_config'] = BitsAndBytesConfig( 18 | load_in_4bit=True, 19 | bnb_4bit_compute_dtype=torch.float16, 20 | bnb_4bit_use_double_quant=True, 21 | bnb_4bit_quant_type='nf4' 22 | ) 23 | else: 24 | kwargs['torch_dtype'] = torch.float16 25 | 26 | tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False) 27 | 28 | if 'llama' in model_base.lower() or 'vicuna' in model_base.lower(): 29 | if "pif" in model_name: 30 | model = vwPifLlamaForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs) 31 | else: 32 | model = vwLlamaForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs) 33 | 34 | elif "mistral" in model_base.lower(): 35 | model = vwMistralForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs) 36 | 37 | else: 38 | raise Exception("undefined model") 39 | 40 | tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False) 41 | 42 | image_processor = None 43 | mm_use_im_start_end = getattr(model.config, "mm_use_im_start_end", False) 44 | mm_use_im_patch_token = getattr(model.config, "mm_use_im_patch_token", True) 45 | if mm_use_im_patch_token: 46 | tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True) 47 | if mm_use_im_start_end: 48 | tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True) 49 | model.resize_token_embeddings(len(tokenizer)) 50 | 51 | vision_tower = model.get_vision_tower() 52 | if not vision_tower.is_loaded: 53 | vision_tower.load_model() 54 | vision_tower.to(device=device, dtype=torch.float16) 55 | image_processor = vision_tower.image_processor 56 | 57 | if hasattr(model.config, "max_sequence_length"): 58 | context_len = model.config.max_sequence_length 59 | else: 60 | context_len = 2048 61 | 62 | return tokenizer, model, image_processor, context_len 63 | -------------------------------------------------------------------------------- /visual_words/model/language_model/vw_llama.py: -------------------------------------------------------------------------------- 1 | from typing import List, Optional, Tuple, Union 2 | import torch 3 | import torch.nn as nn 4 | from torch.nn import CrossEntropyLoss, KLDivLoss 5 | from transformers import AutoConfig, AutoModelForCausalLM, \ 6 | LlamaConfig, LlamaModel, LlamaForCausalLM 7 | from transformers.modeling_outputs import CausalLMOutputWithPast 8 | from ..vw_arch import vwMetaModel, vwMetaForCausalLM 9 | 10 | class LlavaConfig(LlamaConfig): 11 | model_type = "llava" 12 | 13 | class vwLlamaModel(vwMetaModel, LlamaModel): 14 | config_class = LlavaConfig 15 | 16 | def __init__(self, config: LlamaConfig): 17 | super(vwLlamaModel, self).__init__(config) 18 | 19 | 20 | class vwLlamaForCausalLM(LlamaForCausalLM, vwMetaForCausalLM): 21 | config_class = LlavaConfig 22 | 23 | def __init__(self, config): 24 | super(LlamaForCausalLM, self).__init__(config) 25 | self.model = vwLlamaModel(config) 26 | 27 | self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) 28 | 29 | self.vm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) 30 | # Initialize weights and apply final processing 31 | self.post_init() 32 | 33 | def get_model(self): 34 | return self.model 35 | 36 | def forward( 37 | self, 38 | input_ids: torch.LongTensor = None, 39 | attention_mask: Optional[torch.Tensor] = None, 40 | past_key_values: Optional[List[torch.FloatTensor]] = None, 41 | inputs_embeds: Optional[torch.FloatTensor] = None, 42 | labels: Optional[torch.LongTensor] = None, 43 | use_cache: Optional[bool] = None, 44 | output_attentions: Optional[bool] = None, 45 | output_hidden_states: Optional[bool] = None, 46 | images: Optional[torch.FloatTensor] = None, 47 | return_dict: Optional[bool] = None, 48 | ) -> Union[Tuple, CausalLMOutputWithPast]: 49 | output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions 50 | output_hidden_states = ( 51 | output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states 52 | ) 53 | return_dict = return_dict if return_dict is not None else self.config.use_return_dict 54 | input_ids, attention_mask, past_key_values, inputs_embeds, labels, image_position_labels = self.prepare_inputs_labels_for_multimodal(input_ids, attention_mask, past_key_values, labels, images) 55 | ### 56 | if image_position_labels is not None: 57 | image_position_labels=image_position_labels.bool() 58 | visual_labels=self.vm_head(inputs_embeds) 59 | else: 60 | visual_labels=None 61 | ### 62 | 63 | outputs = self.model( 64 | input_ids=input_ids, 65 | attention_mask=attention_mask, 66 | past_key_values=past_key_values, 67 | inputs_embeds=inputs_embeds, 68 | use_cache=use_cache, 69 | output_attentions=output_attentions, 70 | output_hidden_states=output_hidden_states, 71 | return_dict=return_dict 72 | ) 73 | 74 | hidden_states = outputs[0] 75 | logits = self.lm_head(hidden_states) 76 | 77 | loss = None 78 | loss_lm = None 79 | loss_vm = None 80 | if labels is not None: 81 | # Shift so that tokens < n predict n 82 | shift_logits = logits[..., :-1, :].contiguous() 83 | shift_labels = labels[..., 1:].contiguous() 84 | # Flatten the tokens 85 | loss_fct_lm = CrossEntropyLoss() 86 | shift_logits = shift_logits.view(-1, self.config.vocab_size) 87 | shift_labels = shift_labels.view(-1) 88 | # Enable model/pipeline parallelism 89 | shift_labels = shift_labels.to(shift_logits.device) 90 | loss_lm = loss_fct_lm(shift_logits, shift_labels) 91 | 92 | ### 93 | if visual_labels is not None: 94 | shift_logits = logits[..., :-1, :].contiguous() 95 | shift_visual_labels = visual_labels[..., 1:, :].contiguous() 96 | shift_image_position_labels = image_position_labels[..., 1:].contiguous() 97 | 98 | loss_fct_vm = KLDivLoss(reduction="batchmean") 99 | # loss_fct_vm = CrossEntropyLoss() 100 | shift_visual_logits = shift_logits[shift_image_position_labels] 101 | shift_visual_labels = shift_visual_labels[shift_image_position_labels] 102 | # shift_visual_logits = shift_logits 103 | assert shift_visual_logits.shape == shift_visual_labels.shape 104 | shift_visual_labels = shift_visual_labels.to(shift_visual_logits.device) 105 | 106 | # stage3 107 | # shift_visual_logits = nn.functional.softmax(shift_visual_logits,dim=-1) 108 | # shift_visual_labels = nn.functional.log_softmax(shift_visual_labels, dim=-1) 109 | # loss_vm = loss_fct_vm(shift_visual_labels, shift_visual_logits) 110 | 111 | # stage4 112 | shift_visual_logits = nn.functional.log_softmax(shift_visual_logits,dim=-1) 113 | shift_visual_labels = nn.functional.softmax(shift_visual_labels, dim=-1) 114 | loss_vm = loss_fct_vm(shift_visual_logits, shift_visual_labels) 115 | 116 | 117 | if loss_lm is not None and loss_vm is not None: 118 | loss = loss_lm + loss_vm 119 | elif loss_lm is None: 120 | loss = loss_vm 121 | elif loss_vm is None: 122 | loss = loss_lm 123 | else: 124 | loss = None 125 | ### 126 | 127 | if not return_dict: 128 | output = (logits,) + outputs[1:] 129 | return (loss,) + output if loss is not None else output 130 | 131 | return CausalLMOutputWithPast( 132 | loss=loss, 133 | logits=logits, 134 | past_key_values=outputs.past_key_values, 135 | hidden_states=outputs.hidden_states, 136 | attentions=outputs.attentions, 137 | ) 138 | 139 | def prepare_inputs_for_generation( 140 | self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs 141 | ): 142 | if past_key_values: 143 | input_ids = input_ids[:, -1:] 144 | 145 | # if `inputs_embeds` are passed, we only want to use them in the 1st generation step 146 | if inputs_embeds is not None and past_key_values is None: 147 | model_inputs = {"inputs_embeds": inputs_embeds} 148 | else: 149 | model_inputs = {"input_ids": input_ids} 150 | 151 | model_inputs.update( 152 | { 153 | "past_key_values": past_key_values, 154 | "use_cache": kwargs.get("use_cache"), 155 | "attention_mask": attention_mask, 156 | "images": kwargs.get("images", None), 157 | } 158 | ) 159 | return model_inputs 160 | 161 | -------------------------------------------------------------------------------- /visual_words/model/language_model/vw_mistral.py: -------------------------------------------------------------------------------- 1 | from typing import List, Optional, Tuple, Union 2 | import torch 3 | import torch.nn as nn 4 | from torch.nn import CrossEntropyLoss, KLDivLoss 5 | from transformers import AutoConfig, AutoModelForCausalLM, \ 6 | MistralConfig, MistralModel, MistralForCausalLM 7 | from transformers.modeling_outputs import CausalLMOutputWithPast 8 | from ..vw_arch import vwMetaModel, vwMetaForCausalLM 9 | 10 | class LlavaConfig(MistralConfig): 11 | model_type = "llava" 12 | 13 | class vwMistralModel(vwMetaModel, MistralModel): 14 | config_class = LlavaConfig 15 | 16 | def __init__(self, config: MistralConfig): 17 | super(vwMistralModel, self).__init__(config) 18 | 19 | 20 | class vwMistralForCausalLM(MistralForCausalLM, vwMetaForCausalLM): 21 | config_class = LlavaConfig 22 | 23 | def __init__(self, config): 24 | super(MistralForCausalLM, self).__init__(config) 25 | self.model = vwMistralModel(config) 26 | 27 | self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) 28 | 29 | self.vm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) 30 | # Initialize weights and apply final processing 31 | self.post_init() 32 | 33 | def get_model(self): 34 | return self.model 35 | 36 | def forward( 37 | self, 38 | input_ids: torch.LongTensor = None, 39 | attention_mask: Optional[torch.Tensor] = None, 40 | past_key_values: Optional[List[torch.FloatTensor]] = None, 41 | inputs_embeds: Optional[torch.FloatTensor] = None, 42 | labels: Optional[torch.LongTensor] = None, 43 | use_cache: Optional[bool] = None, 44 | output_attentions: Optional[bool] = None, 45 | output_hidden_states: Optional[bool] = None, 46 | images: Optional[torch.FloatTensor] = None, 47 | return_dict: Optional[bool] = None, 48 | ) -> Union[Tuple, CausalLMOutputWithPast]: 49 | output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions 50 | output_hidden_states = ( 51 | output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states 52 | ) 53 | return_dict = return_dict if return_dict is not None else self.config.use_return_dict 54 | input_ids, attention_mask, past_key_values, inputs_embeds, labels, image_position_labels = self.prepare_inputs_labels_for_multimodal(input_ids, attention_mask, past_key_values, labels, images) 55 | ### 56 | if image_position_labels is not None: 57 | image_position_labels=image_position_labels.bool() 58 | visual_labels=self.vm_head(inputs_embeds) 59 | else: 60 | visual_labels=None 61 | ### 62 | 63 | outputs = self.model( 64 | input_ids=input_ids, 65 | attention_mask=attention_mask, 66 | past_key_values=past_key_values, 67 | inputs_embeds=inputs_embeds, 68 | use_cache=use_cache, 69 | output_attentions=output_attentions, 70 | output_hidden_states=output_hidden_states, 71 | return_dict=return_dict 72 | ) 73 | 74 | hidden_states = outputs[0] 75 | logits = self.lm_head(hidden_states) 76 | 77 | loss = None 78 | loss_lm = None 79 | loss_vm = None 80 | if labels is not None: 81 | # Shift so that tokens < n predict n 82 | shift_logits = logits[..., :-1, :].contiguous() 83 | shift_labels = labels[..., 1:].contiguous() 84 | # Flatten the tokens 85 | loss_fct_lm = CrossEntropyLoss() 86 | shift_logits = shift_logits.view(-1, self.config.vocab_size) 87 | shift_labels = shift_labels.view(-1) 88 | # Enable model/pipeline parallelism 89 | shift_labels = shift_labels.to(shift_logits.device) 90 | loss_lm = loss_fct_lm(shift_logits, shift_labels) 91 | 92 | ### 93 | if visual_labels is not None: 94 | shift_logits = logits[..., :-1, :].contiguous() 95 | shift_visual_labels = visual_labels[..., 1:, :].contiguous() 96 | shift_image_position_labels = image_position_labels[..., 1:].contiguous() 97 | 98 | loss_fct_vm = KLDivLoss(reduction="batchmean") 99 | # loss_fct_vm = CrossEntropyLoss() 100 | shift_visual_logits = shift_logits[shift_image_position_labels] 101 | shift_visual_labels = shift_visual_labels[shift_image_position_labels] 102 | # shift_visual_logits = shift_logits 103 | assert shift_visual_logits.shape == shift_visual_labels.shape 104 | shift_visual_labels = shift_visual_labels.to(shift_visual_logits.device) 105 | 106 | # stage3 107 | # shift_visual_logits = nn.functional.softmax(shift_visual_logits,dim=-1) 108 | # shift_visual_labels = nn.functional.log_softmax(shift_visual_labels, dim=-1) 109 | # loss_vm = loss_fct_vm(shift_visual_labels, shift_visual_logits) 110 | 111 | # stage4 112 | shift_visual_logits = nn.functional.log_softmax(shift_visual_logits,dim=-1) 113 | shift_visual_labels = nn.functional.softmax(shift_visual_labels, dim=-1) 114 | loss_vm = loss_fct_vm(shift_visual_logits, shift_visual_labels) 115 | 116 | 117 | if loss_lm is not None and loss_vm is not None: 118 | loss = loss_lm + loss_vm 119 | elif loss_lm is None: 120 | loss = loss_vm 121 | elif loss_vm is None: 122 | loss = loss_lm 123 | else: 124 | loss = None 125 | ### 126 | 127 | if not return_dict: 128 | output = (logits,) + outputs[1:] 129 | return (loss,) + output if loss is not None else output 130 | 131 | return CausalLMOutputWithPast( 132 | loss=loss, 133 | logits=logits, 134 | past_key_values=outputs.past_key_values, 135 | hidden_states=outputs.hidden_states, 136 | attentions=outputs.attentions, 137 | ) 138 | 139 | def prepare_inputs_for_generation( 140 | self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs 141 | ): 142 | if past_key_values: 143 | input_ids = input_ids[:, -1:] 144 | 145 | # if `inputs_embeds` are passed, we only want to use them in the 1st generation step 146 | if inputs_embeds is not None and past_key_values is None: 147 | model_inputs = {"inputs_embeds": inputs_embeds} 148 | else: 149 | model_inputs = {"input_ids": input_ids} 150 | 151 | model_inputs.update( 152 | { 153 | "past_key_values": past_key_values, 154 | "use_cache": kwargs.get("use_cache"), 155 | "attention_mask": attention_mask, 156 | "images": kwargs.get("images", None), 157 | } 158 | ) 159 | return model_inputs 160 | 161 | -------------------------------------------------------------------------------- /visual_words/model/language_model/vw_pif_llama.py: -------------------------------------------------------------------------------- 1 | from typing import List, Optional, Tuple, Union 2 | import torch 3 | import torch.nn as nn 4 | from torch.nn import CrossEntropyLoss, KLDivLoss 5 | from transformers import AutoConfig, AutoModelForCausalLM, \ 6 | LlamaConfig, LlamaModel, LlamaForCausalLM 7 | from transformers.modeling_outputs import CausalLMOutputWithPast 8 | from ..vw_pif_arch import vwMetaModel, vwPifMetaForCausalLM 9 | 10 | class LlavaConfig(LlamaConfig): 11 | model_type = "llava" 12 | 13 | class vwLlamaModel(vwMetaModel, LlamaModel): 14 | config_class = LlavaConfig 15 | 16 | def __init__(self, config: LlamaConfig): 17 | super(vwLlamaModel, self).__init__(config) 18 | 19 | 20 | class vwPifLlamaForCausalLM(LlamaForCausalLM, vwPifMetaForCausalLM): 21 | config_class = LlavaConfig 22 | 23 | def __init__(self, config): 24 | super(LlamaForCausalLM, self).__init__(config) 25 | self.model = vwLlamaModel(config) 26 | 27 | self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) 28 | 29 | self.vm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) 30 | # Initialize weights and apply final processing 31 | self.post_init() 32 | 33 | def get_model(self): 34 | return self.model 35 | 36 | def forward( 37 | self, 38 | input_ids: torch.LongTensor = None, 39 | attention_mask: Optional[torch.Tensor] = None, 40 | past_key_values: Optional[List[torch.FloatTensor]] = None, 41 | inputs_embeds: Optional[torch.FloatTensor] = None, 42 | labels: Optional[torch.LongTensor] = None, 43 | use_cache: Optional[bool] = None, 44 | output_attentions: Optional[bool] = None, 45 | output_hidden_states: Optional[bool] = None, 46 | images: Optional[torch.FloatTensor] = None, 47 | return_dict: Optional[bool] = None, 48 | ) -> Union[Tuple, CausalLMOutputWithPast]: 49 | output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions 50 | output_hidden_states = ( 51 | output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states 52 | ) 53 | return_dict = return_dict if return_dict is not None else self.config.use_return_dict 54 | input_ids, attention_mask, past_key_values, inputs_embeds, labels, image_position_labels, textual_image_labels = self.prepare_inputs_labels_for_multimodal(input_ids, attention_mask, past_key_values, labels, images) 55 | ### 56 | if image_position_labels is not None: 57 | image_position_labels=image_position_labels.bool() 58 | visual_labels=textual_image_labels 59 | else: 60 | visual_labels=None 61 | ### 62 | 63 | # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) 64 | outputs = self.model( 65 | input_ids=input_ids, 66 | attention_mask=attention_mask, 67 | past_key_values=past_key_values, 68 | inputs_embeds=inputs_embeds, 69 | use_cache=use_cache, 70 | output_attentions=output_attentions, 71 | output_hidden_states=output_hidden_states, 72 | return_dict=return_dict 73 | ) 74 | 75 | hidden_states = outputs[0] 76 | logits = self.lm_head(hidden_states) 77 | 78 | loss = None 79 | loss_lm = None 80 | loss_vm = None 81 | if labels is not None: 82 | # Shift so that tokens < n predict n 83 | shift_logits = logits[..., :-1, :].contiguous() 84 | shift_labels = labels[..., 1:].contiguous() 85 | # Flatten the tokens 86 | loss_fct_lm = CrossEntropyLoss() 87 | shift_logits = shift_logits.view(-1, self.config.vocab_size) 88 | shift_labels = shift_labels.view(-1) 89 | # Enable model/pipeline parallelism 90 | shift_labels = shift_labels.to(shift_logits.device) 91 | loss_lm = loss_fct_lm(shift_logits, shift_labels) 92 | 93 | ### 94 | if visual_labels is not None: 95 | shift_logits = logits[..., :-1, :].contiguous() 96 | shift_visual_labels = visual_labels[..., 1:, :].contiguous() 97 | shift_image_position_labels = image_position_labels[..., 1:].contiguous() 98 | loss_fct_vm = KLDivLoss(reduction="batchmean") 99 | # loss_fct_vm = CrossEntropyLoss() 100 | shift_visual_logits = shift_logits[shift_image_position_labels] 101 | shift_visual_labels = shift_visual_labels[shift_image_position_labels] 102 | # shift_visual_logits = shift_logits 103 | assert shift_visual_logits.shape == shift_visual_labels.shape 104 | assert (shift_visual_labels>= 0).all() 105 | shift_visual_labels = shift_visual_labels.to(shift_visual_logits.device) 106 | 107 | # stage3 108 | # shift_visual_logits = nn.functional.softmax(shift_visual_logits,dim=-1) 109 | # shift_visual_labels = nn.functional.log_softmax(shift_visual_labels, dim=-1) 110 | # loss_vm = loss_fct_vm(shift_visual_labels, shift_visual_logits) 111 | 112 | # stage4 113 | shift_visual_logits = nn.functional.log_softmax(shift_visual_logits,dim=-1) 114 | # Since the labels used here are already softmaxed in prepare_inputs_labels_for_multimodal(), they are commented out here. 115 | # shift_visual_labels = nn.functional.softmax(shift_visual_labels, dim=-1) 116 | shift_visual_labels=shift_visual_labels.to(shift_visual_logits.dtype) 117 | loss_vm = loss_fct_vm(shift_visual_logits, shift_visual_labels) 118 | 119 | 120 | if loss_lm is not None and loss_vm is not None: 121 | loss = loss_lm + loss_vm 122 | elif loss_lm is None: 123 | loss = loss_vm 124 | elif loss_vm is None: 125 | loss = loss_lm 126 | else: 127 | loss = None 128 | ### 129 | 130 | if not return_dict: 131 | output = (logits,) + outputs[1:] 132 | return (loss,) + output if loss is not None else output 133 | 134 | return CausalLMOutputWithPast( 135 | loss=loss, 136 | logits=logits, 137 | past_key_values=outputs.past_key_values, 138 | hidden_states=outputs.hidden_states, 139 | attentions=outputs.attentions, 140 | ) 141 | 142 | def prepare_inputs_for_generation( 143 | self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs 144 | ): 145 | if past_key_values: 146 | input_ids = input_ids[:, -1:] 147 | 148 | # if `inputs_embeds` are passed, we only want to use them in the 1st generation step 149 | if inputs_embeds is not None and past_key_values is None: 150 | model_inputs = {"inputs_embeds": inputs_embeds} 151 | else: 152 | model_inputs = {"input_ids": input_ids} 153 | 154 | model_inputs.update( 155 | { 156 | "past_key_values": past_key_values, 157 | "use_cache": kwargs.get("use_cache"), 158 | "attention_mask": attention_mask, 159 | "images": kwargs.get("images", None), 160 | } 161 | ) 162 | return model_inputs 163 | 164 | -------------------------------------------------------------------------------- /visual_words/model/multimodal_encoder/builder.py: -------------------------------------------------------------------------------- 1 | import os 2 | from .clip_encoder import CLIPVisionTower 3 | 4 | 5 | def build_vision_tower(vision_tower_cfg, **kwargs): 6 | vision_tower = getattr(vision_tower_cfg, 'mm_vision_tower', getattr(vision_tower_cfg, 'vision_tower', None)) 7 | is_absolute_path_exists = os.path.exists(vision_tower) 8 | if is_absolute_path_exists or vision_tower.startswith("openai") or vision_tower.startswith("laion"): 9 | return CLIPVisionTower(vision_tower, args=vision_tower_cfg, **kwargs) 10 | 11 | raise ValueError(f'Unknown vision tower: {vision_tower}') 12 | -------------------------------------------------------------------------------- /visual_words/model/multimodal_encoder/clip_encoder.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from transformers import CLIPVisionModel, CLIPImageProcessor, CLIPVisionConfig 5 | 6 | 7 | class CLIPVisionTower(nn.Module): 8 | def __init__(self, vision_tower, args, delay_load=False): 9 | super().__init__() 10 | 11 | self.is_loaded = False 12 | 13 | self.vision_tower_name = vision_tower 14 | self.select_layer = args.mm_vision_select_layer 15 | self.select_feature = getattr(args, 'mm_vision_select_feature', 'patch') 16 | 17 | if not delay_load: 18 | self.load_model() 19 | else: 20 | self.cfg_only = CLIPVisionConfig.from_pretrained(self.vision_tower_name) 21 | 22 | def load_model(self): 23 | self.image_processor = CLIPImageProcessor.from_pretrained(self.vision_tower_name) 24 | self.vision_tower = CLIPVisionModel.from_pretrained(self.vision_tower_name) 25 | self.vision_tower.requires_grad_(False) 26 | 27 | self.is_loaded = True 28 | 29 | def feature_select(self, image_forward_outs): 30 | image_features = image_forward_outs.hidden_states[self.select_layer] 31 | if self.select_feature == 'patch': 32 | image_features = image_features[:, 1:] 33 | elif self.select_feature == 'cls_patch': 34 | image_features = image_features 35 | else: 36 | raise ValueError(f'Unexpected select feature: {self.select_feature}') 37 | return image_features 38 | 39 | @torch.no_grad() 40 | def forward(self, images): 41 | if type(images) is list: 42 | image_features = [] 43 | for image in images: 44 | image_forward_out = self.vision_tower(image.to(device=self.device, dtype=self.dtype).unsqueeze(0), output_hidden_states=True) 45 | image_feature = self.feature_select(image_forward_out).to(image.dtype) 46 | image_features.append(image_feature) 47 | else: 48 | image_forward_outs = self.vision_tower(images.to(device=self.device, dtype=self.dtype), output_hidden_states=True) 49 | image_features = self.feature_select(image_forward_outs).to(images.dtype) 50 | 51 | return image_features 52 | 53 | @property 54 | def dummy_feature(self): 55 | return torch.zeros(1, self.hidden_size, device=self.device, dtype=self.dtype) 56 | 57 | @property 58 | def dtype(self): 59 | return self.vision_tower.dtype 60 | 61 | @property 62 | def device(self): 63 | return self.vision_tower.device 64 | 65 | @property 66 | def config(self): 67 | if self.is_loaded: 68 | return self.vision_tower.config 69 | else: 70 | return self.cfg_only 71 | 72 | @property 73 | def hidden_size(self): 74 | return self.config.hidden_size 75 | 76 | @property 77 | def num_patches(self): 78 | return (self.config.image_size // self.config.patch_size) ** 2 79 | -------------------------------------------------------------------------------- /visual_words/model/multimodal_projector/builder.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import re 4 | 5 | 6 | class IdentityMap(nn.Module): 7 | def __init__(self): 8 | super().__init__() 9 | 10 | def forward(self, x, *args, **kwargs): 11 | return x 12 | 13 | @property 14 | def config(self): 15 | return {"mm_projector_type": 'identity'} 16 | 17 | 18 | class SimpleResBlock(nn.Module): 19 | def __init__(self, channels): 20 | super().__init__() 21 | self.pre_norm = nn.LayerNorm(channels) 22 | 23 | self.proj = nn.Sequential( 24 | nn.Linear(channels, channels), 25 | nn.GELU(), 26 | nn.Linear(channels, channels) 27 | ) 28 | def forward(self, x): 29 | x = self.pre_norm(x) 30 | return x + self.proj(x) 31 | 32 | 33 | def build_vision_projector(config, delay_load=False, **kwargs): 34 | projector_type = getattr(config, 'mm_projector_type', 'linear') 35 | 36 | if projector_type == 'linear': 37 | return nn.Linear(config.mm_hidden_size, config.hidden_size) 38 | 39 | mlp_gelu_match = re.match(r'^mlp(\d+)x_gelu$', projector_type) 40 | if mlp_gelu_match: 41 | mlp_depth = int(mlp_gelu_match.group(1)) 42 | modules = [nn.Linear(config.mm_hidden_size, config.hidden_size)] 43 | for _ in range(1, mlp_depth): 44 | modules.append(nn.GELU()) 45 | modules.append(nn.Linear(config.hidden_size, config.hidden_size)) 46 | return nn.Sequential(*modules) 47 | 48 | if projector_type == 'identity': 49 | return IdentityMap() 50 | 51 | raise ValueError(f'Unknown projector type: {projector_type}') 52 | -------------------------------------------------------------------------------- /visual_words/model/vw_arch.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | import torch 3 | import torch.nn as nn 4 | from .multimodal_encoder.builder import build_vision_tower 5 | from .multimodal_projector.builder import build_vision_projector 6 | from visual_words.constants import IGNORE_INDEX, IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN 7 | 8 | class vwMetaModel: 9 | 10 | def __init__(self, config): 11 | super(vwMetaModel, self).__init__(config) 12 | 13 | if hasattr(config, "mm_vision_tower"): 14 | self.vision_tower = build_vision_tower(config, delay_load=True) 15 | self.mm_projector = build_vision_projector(config) 16 | 17 | def get_vision_tower(self): 18 | vision_tower = getattr(self, 'vision_tower', None) 19 | if type(vision_tower) is list: 20 | vision_tower = vision_tower[0] 21 | return vision_tower 22 | 23 | def initialize_vision_modules(self, model_args, fsdp=None): 24 | vision_tower = model_args.vision_tower 25 | mm_vision_select_layer = model_args.mm_vision_select_layer 26 | mm_vision_select_feature = model_args.mm_vision_select_feature 27 | pretrain_mm_mlp_adapter = model_args.pretrain_mm_mlp_adapter 28 | 29 | self.config.mm_vision_tower = vision_tower 30 | 31 | if self.get_vision_tower() is None: 32 | vision_tower = build_vision_tower(model_args) 33 | 34 | if fsdp is not None and len(fsdp) > 0: 35 | self.vision_tower = [vision_tower] 36 | else: 37 | self.vision_tower = vision_tower 38 | else: 39 | if fsdp is not None and len(fsdp) > 0: 40 | vision_tower = self.vision_tower[0] 41 | else: 42 | vision_tower = self.vision_tower 43 | vision_tower.load_model() 44 | 45 | self.config.use_mm_proj = True 46 | self.config.mm_projector_type = getattr(model_args, 'mm_projector_type', 'linear') 47 | self.config.mm_hidden_size = vision_tower.hidden_size 48 | self.config.mm_vision_select_layer = mm_vision_select_layer 49 | self.config.mm_vision_select_feature = mm_vision_select_feature 50 | 51 | if getattr(self, 'mm_projector', None) is None: 52 | self.mm_projector = build_vision_projector(self.config) 53 | 54 | if pretrain_mm_mlp_adapter is not None: 55 | mm_projector_weights = torch.load(pretrain_mm_mlp_adapter, map_location='cpu') 56 | def get_w(weights, keyword): 57 | return {k.split(keyword + '.')[1]: v for k, v in weights.items() if keyword in k} 58 | 59 | self.mm_projector.load_state_dict(get_w(mm_projector_weights, 'mm_projector')) 60 | 61 | 62 | class vwMetaForCausalLM(ABC): 63 | 64 | @abstractmethod 65 | def get_model(self): 66 | pass 67 | 68 | def get_vision_tower(self): 69 | return self.get_model().get_vision_tower() 70 | 71 | def encode_images(self, images): 72 | image_features = self.get_model().get_vision_tower()(images) 73 | image_features = self.get_model().mm_projector(image_features) 74 | return image_features 75 | 76 | def prepare_inputs_labels_for_multimodal( 77 | self, input_ids, attention_mask, past_key_values, labels, images 78 | ): 79 | vision_tower = self.get_vision_tower() 80 | if vision_tower is None or images is None or input_ids.shape[1] == 1: 81 | if past_key_values is not None and vision_tower is not None and images is not None and input_ids.shape[1] == 1: 82 | attention_mask = torch.ones((attention_mask.shape[0], past_key_values[-1][-1].shape[-2] + 1), dtype=attention_mask.dtype, device=attention_mask.device) 83 | return input_ids, attention_mask, past_key_values, None, labels, None 84 | 85 | if type(images) is list or images.ndim == 5: 86 | concat_images = torch.cat([image for image in images], dim=0) 87 | image_features = self.encode_images(concat_images) 88 | split_sizes = [image.shape[0] for image in images] 89 | image_features = torch.split(image_features, split_sizes, dim=0) 90 | image_features = [x.flatten(0, 1) for x in image_features] 91 | else: 92 | image_features = self.encode_images(images) 93 | 94 | new_input_embeds = [] 95 | new_labels = [] if labels is not None else None 96 | new_image_position_labels = [] if labels is not None else None 97 | cur_image_idx = 0 98 | for batch_idx, cur_input_ids in enumerate(input_ids): 99 | if (cur_input_ids == IMAGE_TOKEN_INDEX).sum() == 0: 100 | # multimodal LLM, but the current sample is not multimodal 101 | half_len = cur_input_ids.shape[0] // 2 102 | cur_image_features = image_features[cur_image_idx] 103 | cur_input_embeds_1 = self.get_model().embed_tokens(cur_input_ids[:half_len]) 104 | cur_input_embeds_2 = self.get_model().embed_tokens(cur_input_ids[half_len:]) 105 | cur_input_embeds = torch.cat([cur_input_embeds_1, cur_image_features[0:0], cur_input_embeds_2], dim=0) 106 | new_input_embeds.append(cur_input_embeds) 107 | if labels is not None: 108 | new_labels.append(labels[batch_idx]) 109 | new_image_position_labels.append(torch.zeros_like(labels[batch_idx],dtype=labels.dtype, device=labels.device)) 110 | cur_image_idx += 1 111 | continue 112 | image_token_indices = torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0] 113 | cur_new_input_embeds = [] 114 | if labels is not None: 115 | cur_labels = labels[batch_idx] 116 | cur_new_labels = [] 117 | cur_image_position_labels=torch.zeros_like(labels[batch_idx],dtype=labels.dtype, device=labels.device) 118 | cur_new_image_position_labels=[] 119 | assert cur_labels.shape == cur_input_ids.shape 120 | while image_token_indices.numel() > 0: 121 | cur_image_features = image_features[cur_image_idx] 122 | image_token_start = image_token_indices[0] 123 | 124 | cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids[:image_token_start])) 125 | cur_new_input_embeds.append(cur_image_features) 126 | if labels is not None: 127 | cur_new_labels.append(cur_labels[:image_token_start]) 128 | cur_new_labels.append(torch.full((cur_image_features.shape[0],), IGNORE_INDEX, device=labels.device, dtype=labels.dtype)) 129 | cur_labels = cur_labels[image_token_start+1:] 130 | cur_new_image_position_labels.append(cur_image_position_labels[:image_token_start]) 131 | cur_new_image_position_labels.append(torch.full((cur_image_features.shape[0],), 1.0, device=labels.device, dtype=labels.dtype)) 132 | cur_image_position_labels=cur_image_position_labels[image_token_start+1:] 133 | 134 | cur_image_idx += 1 135 | if getattr(self.config, 'tune_mm_mlp_adapter', False) and getattr(self.config, 'mm_use_im_start_end', False): 136 | cur_input_ids = cur_input_ids[image_token_start+2:] 137 | else: 138 | cur_input_ids = cur_input_ids[image_token_start+1:] 139 | image_token_indices = torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0] 140 | if cur_input_ids.numel() > 0: 141 | if getattr(self.config, 'tune_mm_mlp_adapter', False) and getattr(self.config, 'mm_use_im_start_end', False): 142 | cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids).detach()) 143 | else: 144 | cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids)) 145 | if labels is not None: 146 | cur_new_labels.append(cur_labels) 147 | cur_new_image_position_labels.append(cur_image_position_labels) 148 | cur_new_input_embeds = [x.to(device=self.device) for x in cur_new_input_embeds] 149 | cur_new_input_embeds = torch.cat(cur_new_input_embeds, dim=0) 150 | new_input_embeds.append(cur_new_input_embeds) 151 | if labels is not None: 152 | cur_new_labels = torch.cat(cur_new_labels, dim=0) 153 | new_labels.append(cur_new_labels) 154 | cur_new_image_position_labels = torch.cat(cur_new_image_position_labels, dim=0) 155 | new_image_position_labels.append(cur_new_image_position_labels) 156 | 157 | if any(x.shape != new_input_embeds[0].shape for x in new_input_embeds): 158 | max_len = max(x.shape[0] for x in new_input_embeds) 159 | 160 | new_input_embeds_align = [] 161 | for cur_new_embed in new_input_embeds: 162 | cur_new_embed = torch.cat((cur_new_embed, torch.zeros((max_len - cur_new_embed.shape[0], cur_new_embed.shape[1]), dtype=cur_new_embed.dtype, device=cur_new_embed.device)), dim=0) 163 | new_input_embeds_align.append(cur_new_embed) 164 | new_input_embeds = torch.stack(new_input_embeds_align, dim=0) 165 | 166 | if labels is not None: 167 | new_labels_align = [] 168 | _new_labels = new_labels 169 | for cur_new_label in new_labels: 170 | cur_new_label = torch.cat((cur_new_label, torch.full((max_len - cur_new_label.shape[0],), IGNORE_INDEX, dtype=cur_new_label.dtype, device=cur_new_label.device)), dim=0) 171 | new_labels_align.append(cur_new_label) 172 | new_labels = torch.stack(new_labels_align, dim=0) 173 | new_image_position_labels_align = [] 174 | for cur_new_image_position_label in new_image_position_labels: 175 | cur_new_image_position_label = torch.cat((cur_new_image_position_label, torch.full((max_len - cur_new_image_position_label.shape[0],), 0.0, dtype=cur_new_image_position_label.dtype, device=cur_new_image_position_label.device)), dim=0) 176 | new_image_position_labels_align.append(cur_new_image_position_label) 177 | new_image_position_labels = torch.stack(new_image_position_labels_align, dim=0) 178 | 179 | if attention_mask is not None: 180 | new_attention_mask = [] 181 | for cur_attention_mask, cur_new_labels, cur_new_labels_align in zip(attention_mask, _new_labels, new_labels): 182 | new_attn_mask_pad_left = torch.full((cur_new_labels.shape[0] - labels.shape[1],), True, dtype=attention_mask.dtype, device=attention_mask.device) 183 | new_attn_mask_pad_right = torch.full((cur_new_labels_align.shape[0] - cur_new_labels.shape[0],), False, dtype=attention_mask.dtype, device=attention_mask.device) 184 | cur_new_attention_mask = torch.cat((new_attn_mask_pad_left, cur_attention_mask, new_attn_mask_pad_right), dim=0) 185 | new_attention_mask.append(cur_new_attention_mask) 186 | attention_mask = torch.stack(new_attention_mask, dim=0) 187 | assert attention_mask.shape == new_labels.shape 188 | else: 189 | new_input_embeds = torch.stack(new_input_embeds, dim=0) 190 | if labels is not None: 191 | new_labels = torch.stack(new_labels, dim=0) 192 | new_image_position_labels=torch.stack(new_image_position_labels) 193 | 194 | if attention_mask is not None: 195 | new_attn_mask_pad_left = torch.full((attention_mask.shape[0], new_input_embeds.shape[1] - input_ids.shape[1]), True, dtype=attention_mask.dtype, device=attention_mask.device) 196 | attention_mask = torch.cat((new_attn_mask_pad_left, attention_mask), dim=1) 197 | assert attention_mask.shape == new_input_embeds.shape[:2] 198 | if new_image_position_labels is not None and new_labels is not None: 199 | assert new_image_position_labels.shape==new_labels.shape 200 | return None, attention_mask, past_key_values, new_input_embeds, new_labels, new_image_position_labels 201 | -------------------------------------------------------------------------------- /visual_words/model/vw_pif_arch.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | import torch 3 | import torch.nn as nn 4 | from .multimodal_encoder.builder import build_vision_tower 5 | from .multimodal_projector.builder import build_vision_projector 6 | from visual_words.constants import IGNORE_INDEX, IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN 7 | 8 | class vwMetaModel: 9 | 10 | def __init__(self, config): 11 | super(vwMetaModel, self).__init__(config) 12 | 13 | if hasattr(config, "mm_vision_tower"): 14 | self.vision_tower = build_vision_tower(config, delay_load=True) 15 | self.mm_projector = build_vision_projector(config) 16 | 17 | def get_vision_tower(self): 18 | vision_tower = getattr(self, 'vision_tower', None) 19 | if type(vision_tower) is list: 20 | vision_tower = vision_tower[0] 21 | return vision_tower 22 | 23 | def initialize_vision_modules(self, model_args, fsdp=None): 24 | vision_tower = model_args.vision_tower 25 | mm_vision_select_layer = model_args.mm_vision_select_layer 26 | mm_vision_select_feature = model_args.mm_vision_select_feature 27 | pretrain_mm_mlp_adapter = model_args.pretrain_mm_mlp_adapter 28 | 29 | self.config.mm_vision_tower = vision_tower 30 | 31 | if self.get_vision_tower() is None: 32 | vision_tower = build_vision_tower(model_args) 33 | 34 | if fsdp is not None and len(fsdp) > 0: 35 | self.vision_tower = [vision_tower] 36 | else: 37 | self.vision_tower = vision_tower 38 | else: 39 | if fsdp is not None and len(fsdp) > 0: 40 | vision_tower = self.vision_tower[0] 41 | else: 42 | vision_tower = self.vision_tower 43 | vision_tower.load_model() 44 | 45 | self.config.use_mm_proj = True 46 | self.config.mm_projector_type = getattr(model_args, 'mm_projector_type', 'linear') 47 | self.config.mm_hidden_size = vision_tower.hidden_size 48 | self.config.mm_vision_select_layer = mm_vision_select_layer 49 | self.config.mm_vision_select_feature = mm_vision_select_feature 50 | 51 | if getattr(self, 'mm_projector', None) is None: 52 | self.mm_projector = build_vision_projector(self.config) 53 | 54 | if pretrain_mm_mlp_adapter is not None: 55 | mm_projector_weights = torch.load(pretrain_mm_mlp_adapter, map_location='cpu') 56 | def get_w(weights, keyword): 57 | return {k.split(keyword + '.')[1]: v for k, v in weights.items() if keyword in k} 58 | 59 | self.mm_projector.load_state_dict(get_w(mm_projector_weights, 'mm_projector')) 60 | 61 | 62 | class vwPifMetaForCausalLM(ABC): 63 | 64 | @abstractmethod 65 | def get_model(self): 66 | pass 67 | 68 | def get_vision_tower(self): 69 | return self.get_model().get_vision_tower() 70 | 71 | def encode_images(self, images): 72 | image_features = self.get_model().get_vision_tower()(images) 73 | image_features = self.get_model().mm_projector(image_features) 74 | return image_features 75 | 76 | def prepare_inputs_labels_for_multimodal( 77 | self, input_ids, attention_mask, past_key_values, labels, images 78 | ): 79 | vision_tower = self.get_vision_tower() 80 | if vision_tower is None or images is None or input_ids.shape[1] == 1: 81 | if past_key_values is not None and vision_tower is not None and images is not None and input_ids.shape[1] == 1: 82 | attention_mask = torch.ones((attention_mask.shape[0], past_key_values[-1][-1].shape[-2] + 1), dtype=attention_mask.dtype, device=attention_mask.device) 83 | ### 84 | return input_ids, attention_mask, past_key_values, None, labels, None, None 85 | ### 86 | 87 | if type(images) is list or images.ndim == 5: 88 | concat_images = torch.cat([image for image in images], dim=0) 89 | image_features = self.encode_images(concat_images) 90 | ### 91 | image_features_shape=image_features.shape 92 | textual_image_score = self.vm_head(image_features) 93 | textual_image_score = nn.functional.softmax(textual_image_score) 94 | embed_tokens = self.get_model().embed_tokens.weight 95 | image_features=torch.einsum("id,li->ld",[embed_tokens,textual_image_score.to(embed_tokens.device)]) 96 | assert image_features_shape==image_features.shape 97 | ### 98 | split_sizes = [image.shape[0] for image in images] 99 | image_features = torch.split(image_features, split_sizes, dim=0) 100 | image_features = [x.flatten(0, 1) for x in image_features] 101 | else: 102 | image_features = self.encode_images(images) 103 | ### 104 | image_features_shape=image_features.shape 105 | textual_image_score = self.vm_head(image_features) 106 | textual_image_score = nn.functional.softmax(textual_image_score,dim=-1) 107 | embed_tokens = self.get_model().embed_tokens.weight 108 | image_features=torch.einsum("id,bli->bld",[embed_tokens,textual_image_score.to(embed_tokens.device)]) 109 | assert image_features_shape==image_features.shape 110 | ### 111 | new_input_embeds = [] 112 | new_labels = [] if labels is not None else None 113 | ### 114 | new_image_position_labels = [] if labels is not None else None 115 | new_textual_image_labels = [] if labels is not None else None 116 | ### 117 | cur_image_idx = 0 118 | for batch_idx, cur_input_ids in enumerate(input_ids): 119 | if (cur_input_ids == IMAGE_TOKEN_INDEX).sum() == 0: 120 | # multimodal LLM, but the current sample is not multimodal 121 | half_len = cur_input_ids.shape[0] // 2 122 | cur_image_features = image_features[cur_image_idx] 123 | cur_input_embeds_1 = self.get_model().embed_tokens(cur_input_ids[:half_len]) 124 | cur_input_embeds_2 = self.get_model().embed_tokens(cur_input_ids[half_len:]) 125 | cur_input_embeds = torch.cat([cur_input_embeds_1, cur_image_features[0:0], cur_input_embeds_2], dim=0) 126 | new_input_embeds.append(cur_input_embeds) 127 | if labels is not None: 128 | new_labels.append(labels[batch_idx]) 129 | ### 130 | image_label_l=labels[batch_idx].size()[0] 131 | image_label_d=self.get_model().vocab_size 132 | new_textual_image_labels.append( 133 | torch.full((image_label_l,image_label_d),-100,dtype=labels.dtype, device=labels.device) 134 | ) 135 | new_image_position_labels.append(torch.zeros_like(labels[batch_idx],dtype=labels.dtype, device=labels.device)) 136 | ### 137 | cur_image_idx += 1 138 | continue 139 | image_token_indices = torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0] 140 | cur_new_input_embeds = [] 141 | if labels is not None: 142 | cur_labels = labels[batch_idx] 143 | cur_new_labels = [] 144 | ### 145 | image_label_l=labels[batch_idx].size()[0] 146 | image_label_d=self.get_model().vocab_size 147 | cur_textual_image_labels=torch.full((image_label_l,image_label_d),-100,dtype=labels.dtype, device=labels.device) 148 | cur_new_textual_image_labels=[] 149 | cur_image_position_labels=torch.zeros_like(labels[batch_idx],dtype=labels.dtype, device=labels.device) 150 | cur_new_image_position_labels=[] 151 | ### 152 | assert cur_labels.shape == cur_input_ids.shape 153 | while image_token_indices.numel() > 0: 154 | cur_image_features = image_features[cur_image_idx] 155 | image_token_start = image_token_indices[0] 156 | cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids[:image_token_start])) 157 | cur_new_input_embeds.append(cur_image_features) 158 | if labels is not None: 159 | cur_new_labels.append(cur_labels[:image_token_start]) 160 | cur_new_labels.append(torch.full((cur_image_features.shape[0],), IGNORE_INDEX, device=labels.device, dtype=labels.dtype)) 161 | cur_labels = cur_labels[image_token_start+1:] 162 | ### 163 | cur_new_textual_image_labels.append(cur_textual_image_labels[:image_token_start]) 164 | cur_new_textual_image_labels.append(textual_image_score[cur_image_idx]) 165 | cur_textual_image_labels=cur_textual_image_labels[image_token_start+1:] 166 | cur_new_image_position_labels.append(cur_image_position_labels[:image_token_start]) 167 | cur_new_image_position_labels.append(torch.full((cur_image_features.shape[0],), 1.0, device=labels.device, dtype=labels.dtype)) 168 | cur_image_position_labels=cur_image_position_labels[image_token_start+1:] 169 | ### 170 | cur_image_idx += 1 171 | if getattr(self.config, 'tune_mm_mlp_adapter', False) and getattr(self.config, 'mm_use_im_start_end', False): 172 | cur_input_ids = cur_input_ids[image_token_start+2:] 173 | else: 174 | cur_input_ids = cur_input_ids[image_token_start+1:] 175 | image_token_indices = torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0] 176 | if cur_input_ids.numel() > 0: 177 | if getattr(self.config, 'tune_mm_mlp_adapter', False) and getattr(self.config, 'mm_use_im_start_end', False): 178 | cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids).detach()) 179 | else: 180 | cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids)) 181 | if labels is not None: 182 | cur_new_labels.append(cur_labels) 183 | ### 184 | cur_new_textual_image_labels.append(cur_textual_image_labels) 185 | cur_new_image_position_labels.append(cur_image_position_labels) 186 | ### 187 | cur_new_input_embeds = [x.to(device=self.device) for x in cur_new_input_embeds] 188 | cur_new_input_embeds = torch.cat(cur_new_input_embeds, dim=0) 189 | new_input_embeds.append(cur_new_input_embeds) 190 | if labels is not None: 191 | cur_new_labels = torch.cat(cur_new_labels, dim=0) 192 | new_labels.append(cur_new_labels) 193 | ### 194 | cur_new_textual_image_labels = torch.cat(cur_new_textual_image_labels, dim=0) 195 | new_textual_image_labels.append(cur_new_textual_image_labels) 196 | cur_new_image_position_labels = torch.cat(cur_new_image_position_labels, dim=0) 197 | new_image_position_labels.append(cur_new_image_position_labels) 198 | ### 199 | 200 | if any(x.shape != new_input_embeds[0].shape for x in new_input_embeds): 201 | max_len = max(x.shape[0] for x in new_input_embeds) 202 | 203 | new_input_embeds_align = [] 204 | for cur_new_embed in new_input_embeds: 205 | cur_new_embed = torch.cat((cur_new_embed, torch.zeros((max_len - cur_new_embed.shape[0], cur_new_embed.shape[1]), dtype=cur_new_embed.dtype, device=cur_new_embed.device)), dim=0) 206 | new_input_embeds_align.append(cur_new_embed) 207 | new_input_embeds = torch.stack(new_input_embeds_align, dim=0) 208 | 209 | if labels is not None: 210 | new_labels_align = [] 211 | _new_labels = new_labels 212 | for cur_new_label in new_labels: 213 | cur_new_label = torch.cat((cur_new_label, torch.full((max_len - cur_new_label.shape[0],), IGNORE_INDEX, dtype=cur_new_label.dtype, device=cur_new_label.device)), dim=0) 214 | new_labels_align.append(cur_new_label) 215 | new_labels = torch.stack(new_labels_align, dim=0) 216 | ### 217 | new_textual_image_labels_align = [] 218 | for cur_new_textual_image_label in new_textual_image_labels: 219 | cur_new_textual_image_label = torch.cat((cur_new_textual_image_label, torch.full((max_len - cur_new_textual_image_label.shape[0],cur_new_textual_image_label.shape[1]), -100, dtype=cur_new_textual_image_label.dtype, device=cur_new_textual_image_label.device)), dim=0) 220 | new_textual_image_labels_align.append(cur_new_textual_image_label) 221 | new_textual_image_labels = torch.stack(new_textual_image_labels_align, dim=0) 222 | 223 | new_image_position_labels_align = [] 224 | for cur_new_image_position_label in new_image_position_labels: 225 | cur_new_image_position_label = torch.cat((cur_new_image_position_label, torch.full((max_len - cur_new_image_position_label.shape[0],), 0.0, dtype=cur_new_image_position_label.dtype, device=cur_new_image_position_label.device)), dim=0) 226 | new_image_position_labels_align.append(cur_new_image_position_label) 227 | new_image_position_labels = torch.stack(new_image_position_labels_align, dim=0) 228 | ### 229 | 230 | if attention_mask is not None: 231 | new_attention_mask = [] 232 | for cur_attention_mask, cur_new_labels, cur_new_labels_align in zip(attention_mask, _new_labels, new_labels): 233 | new_attn_mask_pad_left = torch.full((cur_new_labels.shape[0] - labels.shape[1],), True, dtype=attention_mask.dtype, device=attention_mask.device) 234 | new_attn_mask_pad_right = torch.full((cur_new_labels_align.shape[0] - cur_new_labels.shape[0],), False, dtype=attention_mask.dtype, device=attention_mask.device) 235 | cur_new_attention_mask = torch.cat((new_attn_mask_pad_left, cur_attention_mask, new_attn_mask_pad_right), dim=0) 236 | new_attention_mask.append(cur_new_attention_mask) 237 | attention_mask = torch.stack(new_attention_mask, dim=0) 238 | assert attention_mask.shape == new_labels.shape 239 | else: 240 | new_input_embeds = torch.stack(new_input_embeds, dim=0) 241 | if labels is not None: 242 | new_labels = torch.stack(new_labels, dim=0) 243 | ### 244 | new_textual_image_labels=torch.stack(new_textual_image_labels) 245 | new_image_position_labels=torch.stack(new_image_position_labels) 246 | ### 247 | 248 | if attention_mask is not None: 249 | new_attn_mask_pad_left = torch.full((attention_mask.shape[0], new_input_embeds.shape[1] - input_ids.shape[1]), True, dtype=attention_mask.dtype, device=attention_mask.device) 250 | attention_mask = torch.cat((new_attn_mask_pad_left, attention_mask), dim=1) 251 | assert attention_mask.shape == new_input_embeds.shape[:2] 252 | ### 253 | if new_textual_image_labels is not None and new_labels is not None: 254 | assert new_textual_image_labels.shape[:-1]==new_labels.shape 255 | if new_image_position_labels is not None and new_labels is not None: 256 | assert new_image_position_labels.shape==new_labels.shape 257 | return None, attention_mask, past_key_values, new_input_embeds, new_labels, new_image_position_labels, new_textual_image_labels 258 | -------------------------------------------------------------------------------- /visual_words/utils.py: -------------------------------------------------------------------------------- 1 | import datetime 2 | import logging 3 | import logging.handlers 4 | import os 5 | import sys 6 | 7 | import requests 8 | 9 | from visual_words.constants import LOGDIR 10 | 11 | server_error_msg = "**NETWORK ERROR DUE TO HIGH TRAFFIC. PLEASE REGENERATE OR REFRESH THIS PAGE.**" 12 | moderation_msg = "YOUR INPUT VIOLATES OUR CONTENT MODERATION GUIDELINES. PLEASE TRY AGAIN." 13 | 14 | handler = None 15 | 16 | 17 | def build_logger(logger_name, logger_filename): 18 | global handler 19 | 20 | formatter = logging.Formatter( 21 | fmt="%(asctime)s | %(levelname)s | %(name)s | %(message)s", 22 | datefmt="%Y-%m-%d %H:%M:%S", 23 | ) 24 | 25 | # Set the format of root handlers 26 | if not logging.getLogger().handlers: 27 | logging.basicConfig(level=logging.INFO) 28 | logging.getLogger().handlers[0].setFormatter(formatter) 29 | 30 | # Redirect stdout and stderr to loggers 31 | stdout_logger = logging.getLogger("stdout") 32 | stdout_logger.setLevel(logging.INFO) 33 | sl = StreamToLogger(stdout_logger, logging.INFO) 34 | sys.stdout = sl 35 | 36 | stderr_logger = logging.getLogger("stderr") 37 | stderr_logger.setLevel(logging.ERROR) 38 | sl = StreamToLogger(stderr_logger, logging.ERROR) 39 | sys.stderr = sl 40 | 41 | # Get logger 42 | logger = logging.getLogger(logger_name) 43 | logger.setLevel(logging.INFO) 44 | 45 | # Add a file handler for all loggers 46 | if handler is None: 47 | os.makedirs(LOGDIR, exist_ok=True) 48 | filename = os.path.join(LOGDIR, logger_filename) 49 | handler = logging.handlers.TimedRotatingFileHandler( 50 | filename, when='D', utc=True) 51 | handler.setFormatter(formatter) 52 | 53 | for name, item in logging.root.manager.loggerDict.items(): 54 | if isinstance(item, logging.Logger): 55 | item.addHandler(handler) 56 | 57 | return logger 58 | 59 | 60 | class StreamToLogger(object): 61 | """ 62 | Fake file-like stream object that redirects writes to a logger instance. 63 | """ 64 | def __init__(self, logger, log_level=logging.INFO): 65 | self.terminal = sys.stdout 66 | self.logger = logger 67 | self.log_level = log_level 68 | self.linebuf = '' 69 | 70 | def __getattr__(self, attr): 71 | return getattr(self.terminal, attr) 72 | 73 | def write(self, buf): 74 | temp_linebuf = self.linebuf + buf 75 | self.linebuf = '' 76 | for line in temp_linebuf.splitlines(True): 77 | # From the io.TextIOWrapper docs: 78 | # On output, if newline is None, any '\n' characters written 79 | # are translated to the system default line separator. 80 | # By default sys.stdout.write() expects '\n' newlines and then 81 | # translates them so this is still cross platform. 82 | if line[-1] == '\n': 83 | self.logger.log(self.log_level, line.rstrip()) 84 | else: 85 | self.linebuf += line 86 | 87 | def flush(self): 88 | if self.linebuf != '': 89 | self.logger.log(self.log_level, self.linebuf.rstrip()) 90 | self.linebuf = '' 91 | 92 | 93 | def disable_torch_init(): 94 | """ 95 | Disable the redundant torch default initialization to accelerate model creation. 96 | """ 97 | import torch 98 | setattr(torch.nn.Linear, "reset_parameters", lambda self: None) 99 | setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None) 100 | 101 | 102 | def violates_moderation(text): 103 | """ 104 | Check whether the text violates OpenAI moderation API. 105 | """ 106 | url = "https://api.openai.com/v1/moderations" 107 | headers = {"Content-Type": "application/json", 108 | "Authorization": "Bearer " + os.environ["OPENAI_API_KEY"]} 109 | text = text.replace("\n", "") 110 | data = "{" + '"input": ' + f'"{text}"' + "}" 111 | data = data.encode("utf-8") 112 | try: 113 | ret = requests.post(url, headers=headers, data=data, timeout=5) 114 | flagged = ret.json()["results"][0]["flagged"] 115 | except requests.exceptions.RequestException as e: 116 | flagged = False 117 | except KeyError as e: 118 | flagged = False 119 | 120 | return flagged 121 | 122 | 123 | def pretty_print_semaphore(semaphore): 124 | if semaphore is None: 125 | return "None" 126 | return f"Semaphore(value={semaphore._value}, locked={semaphore.locked()})" 127 | --------------------------------------------------------------------------------