Upload 19 files

Experiments/clip_expt.ipynb

@@ -0,0 +1,840 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "9fe51ce7-4c87-4186-9fd3-0fb18ac43e56",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from PIL import Image\n",
+    "import requests\n",
+    "from transformers import AutoProcessor, CLIPVisionModel"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "0f4c21dd-4258-461d-8511-5be089d068a8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = CLIPVisionModel.from_pretrained(\"openai/clip-vit-base-patch32\", device_map=\"cuda:0\")\n",
+    "processor = AutoProcessor.from_pretrained(\"openai/clip-vit-base-patch32\", device_map=\"cuda:0\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "98b9f906-ffaa-4be4-8671-4ecf65f12c49",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# url = \"http://images.cocodataset.org/val2017/000000039769.jpg\"\n",
+    "# image = Image.open(requests.get(url, stream=True).raw)\n",
+    "image = Image.open(\"002579.jpg\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "54b2e4ce-b77b-4314-87f6-ca2a1970fc79",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "cdd65c58-007f-450b-8deb-f8b4f372a823",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# image = None"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "e9066c2e-c78b-49d1-979b-10d0f4f09441",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "inputs = processor(images=image, return_tensors=\"pt\", device_map=\"cuda:0\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "e98b211d-29d9-4662-be0b-e011e89b0101",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# inputs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "b030bd3d-4282-4074-98fe-97e658bd0f50",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "torch.Size([1, 3, 224, 224])"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "inputs[\"pixel_values\"].shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "0ce68f11-1c88-4dd7-8b17-0d1de5811fe6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "outputs = model(inputs[\"pixel_values\"].to(\"cuda:0\"))\n",
+    "last_hidden_state = outputs.last_hidden_state\n",
+    "pooled_output = outputs.pooler_output  # pooled CLS states"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "id": "30cb0918-a30e-4246-b540-6b8e0d876807",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "torch.Size([1, 768])"
+      ]
+     },
+     "execution_count": 23,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "pooled_output.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "id": "6399543a-f23f-426d-8289-3bb52d293ece",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "torch.Size([1, 50, 768])"
+      ]
+     },
+     "execution_count": 24,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "last_hidden_state.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "id": "19a70443-5942-4937-b3ea-6a52d76e2b08",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "torch.Size([1, 768])"
+      ]
+     },
+     "execution_count": 25,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "outputs[1].shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "fa13903f-a94a-4839-ae5a-8df4f55c68b6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from torch import nn\n",
+    "from transformers import CLIPVisionConfig,CLIPPreTrainedModel"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "b2bd9198-42f0-40c3-80e1-d167c0b038fb",
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "NameError",
+     "evalue": "name 'Optional' is not defined",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[9], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mclass\u001b[39;00m \u001b[38;5;21;01mCLIPVisionModelWithProjection\u001b[39;00m(CLIPPreTrainedModel):\n\u001b[1;32m      2\u001b[0m     config_class \u001b[38;5;241m=\u001b[39m CLIPVisionConfig\n\u001b[1;32m      3\u001b[0m     main_input_name \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpixel_values\u001b[39m\u001b[38;5;124m\"\u001b[39m\n",
+      "Cell \u001b[0;32mIn[9], line 20\u001b[0m, in \u001b[0;36mCLIPVisionModelWithProjection\u001b[0;34m()\u001b[0m\n\u001b[1;32m     15\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mget_input_embeddings\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m nn\u001b[38;5;241m.\u001b[39mModule:\n\u001b[1;32m     16\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvision_model\u001b[38;5;241m.\u001b[39membeddings\u001b[38;5;241m.\u001b[39mpatch_embedding\n\u001b[1;32m     18\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\n\u001b[1;32m     19\u001b[0m     \u001b[38;5;28mself\u001b[39m,\n\u001b[0;32m---> 20\u001b[0m     pixel_values: \u001b[43mOptional\u001b[49m[torch\u001b[38;5;241m.\u001b[39mFloatTensor] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     21\u001b[0m     output_attentions: Optional[\u001b[38;5;28mbool\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     22\u001b[0m     output_hidden_states: Optional[\u001b[38;5;28mbool\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     23\u001b[0m     return_dict: Optional[\u001b[38;5;28mbool\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     24\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Union[Tuple, CLIPVisionModelOutput]:\n\u001b[1;32m     25\u001b[0m     return_dict \u001b[38;5;241m=\u001b[39m return_dict \u001b[38;5;28;01mif\u001b[39;00m return_dict \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mconfig\u001b[38;5;241m.\u001b[39muse_return_dict\n\u001b[1;32m     27\u001b[0m     vision_outputs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvision_model(\n\u001b[1;32m     28\u001b[0m         pixel_values\u001b[38;5;241m=\u001b[39mpixel_values,\n\u001b[1;32m     29\u001b[0m         output_attentions\u001b[38;5;241m=\u001b[39moutput_attentions,\n\u001b[1;32m     30\u001b[0m         output_hidden_states\u001b[38;5;241m=\u001b[39moutput_hidden_states,\n\u001b[1;32m     31\u001b[0m         return_dict\u001b[38;5;241m=\u001b[39mreturn_dict,\n\u001b[1;32m     32\u001b[0m     )\n",
+      "\u001b[0;31mNameError\u001b[0m: name 'Optional' is not defined"
+     ]
+    }
+   ],
+   "source": [
+    "class CLIPVisionModelWithProjection(CLIPPreTrainedModel):\n",
+    "    config_class = CLIPVisionConfig\n",
+    "    main_input_name = \"pixel_values\"\n",
+    "\n",
+    "    def __init__(self, config: CLIPVisionConfig):\n",
+    "        super().__init__(config)\n",
+    "\n",
+    "        self.vision_model = CLIPVisionTransformer(config)\n",
+    "\n",
+    "        self.visual_projection = nn.Linear(config.hidden_size, config.projection_dim, bias=False)\n",
+    "\n",
+    "        # Initialize weights and apply final processing\n",
+    "        self.post_init()\n",
+    "\n",
+    "    def get_input_embeddings(self) -> nn.Module:\n",
+    "        return self.vision_model.embeddings.patch_embedding\n",
+    "\n",
+    "    def forward(\n",
+    "        self,\n",
+    "        pixel_values: Optional[torch.FloatTensor] = None,\n",
+    "        output_attentions: Optional[bool] = None,\n",
+    "        output_hidden_states: Optional[bool] = None,\n",
+    "        return_dict: Optional[bool] = None,\n",
+    "    ) -> Union[Tuple, CLIPVisionModelOutput]:\n",
+    "        return_dict = return_dict if return_dict is not None else self.config.use_return_dict\n",
+    "\n",
+    "        vision_outputs = self.vision_model(\n",
+    "            pixel_values=pixel_values,\n",
+    "            output_attentions=output_attentions,\n",
+    "            output_hidden_states=output_hidden_states,\n",
+    "            return_dict=return_dict,\n",
+    "        )\n",
+    "\n",
+    "        pooled_output = vision_outputs[1]  # pooled_output\n",
+    "\n",
+    "        image_embeds = self.visual_projection(pooled_output)\n",
+    "\n",
+    "        if not return_dict:\n",
+    "            outputs = (image_embeds, vision_outputs[0]) + vision_outputs[2:]\n",
+    "            return tuple(output for output in outputs if output is not None)\n",
+    "\n",
+    "        return CLIPVisionModelOutput(\n",
+    "            image_embeds=image_embeds,\n",
+    "            last_hidden_state=vision_outputs.last_hidden_state,\n",
+    "            hidden_states=vision_outputs.hidden_states,\n",
+    "            attentions=vision_outputs.attentions,\n",
+    "        )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "id": "68a9ee4a-d977-4725-842d-e64e0dd2f61d",
+   "metadata": {
+    "collapsed": true,
+    "jupyter": {
+     "outputs_hidden": true
+    }
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "loading configuration file config.json from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/config.json\n",
+      "`text_config` is `None`. Initializing the `CLIPTextConfig` with default values.\n",
+      "`vision_config` is `None`. initializing the `CLIPVisionConfig` with default values.\n",
+      "Model config CLIPConfig {\n",
+      "  \"_name_or_path\": \"openai/clip-vit-base-patch32\",\n",
+      "  \"architectures\": [\n",
+      "    \"CLIPModel\"\n",
+      "  ],\n",
+      "  \"initializer_factor\": 1.0,\n",
+      "  \"logit_scale_init_value\": 2.6592,\n",
+      "  \"model_type\": \"clip\",\n",
+      "  \"projection_dim\": 512,\n",
+      "  \"text_config\": {\n",
+      "    \"bos_token_id\": 0,\n",
+      "    \"dropout\": 0.0,\n",
+      "    \"eos_token_id\": 2,\n",
+      "    \"model_type\": \"clip_text_model\"\n",
+      "  },\n",
+      "  \"transformers_version\": \"4.36.2\",\n",
+      "  \"vision_config\": {\n",
+      "    \"dropout\": 0.0,\n",
+      "    \"model_type\": \"clip_vision_model\"\n",
+      "  }\n",
+      "}\n",
+      "\n",
+      "loading weights file pytorch_model.bin from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/pytorch_model.bin\n",
+      "All model checkpoint weights were used when initializing CLIPModel.\n",
+      "\n",
+      "All the weights of CLIPModel were initialized from the model checkpoint at openai/clip-vit-base-patch32.\n",
+      "If your task is similar to the task the model of the checkpoint was trained on, you can already use CLIPModel for predictions without further training.\n",
+      "loading configuration file preprocessor_config.json from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/preprocessor_config.json\n",
+      "loading configuration file preprocessor_config.json from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/preprocessor_config.json\n",
+      "loading configuration file config.json from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/config.json\n",
+      "`text_config` is `None`. Initializing the `CLIPTextConfig` with default values.\n",
+      "`vision_config` is `None`. initializing the `CLIPVisionConfig` with default values.\n",
+      "Model config CLIPConfig {\n",
+      "  \"_name_or_path\": \"openai/clip-vit-base-patch32\",\n",
+      "  \"architectures\": [\n",
+      "    \"CLIPModel\"\n",
+      "  ],\n",
+      "  \"initializer_factor\": 1.0,\n",
+      "  \"logit_scale_init_value\": 2.6592,\n",
+      "  \"model_type\": \"clip\",\n",
+      "  \"projection_dim\": 512,\n",
+      "  \"text_config\": {\n",
+      "    \"bos_token_id\": 0,\n",
+      "    \"dropout\": 0.0,\n",
+      "    \"eos_token_id\": 2,\n",
+      "    \"model_type\": \"clip_text_model\"\n",
+      "  },\n",
+      "  \"transformers_version\": \"4.36.2\",\n",
+      "  \"vision_config\": {\n",
+      "    \"dropout\": 0.0,\n",
+      "    \"model_type\": \"clip_vision_model\"\n",
+      "  }\n",
+      "}\n",
+      "\n",
+      "loading configuration file preprocessor_config.json from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/preprocessor_config.json\n",
+      "size should be a dictionary on of the following set of keys: ({'width', 'height'}, {'shortest_edge'}, {'longest_edge', 'shortest_edge'}, {'longest_edge'}), got 224. Converted to {'shortest_edge': 224}.\n",
+      "crop_size should be a dictionary on of the following set of keys: ({'width', 'height'}, {'shortest_edge'}, {'longest_edge', 'shortest_edge'}, {'longest_edge'}), got 224. Converted to {'height': 224, 'width': 224}.\n",
+      "Image processor CLIPImageProcessor {\n",
+      "  \"crop_size\": {\n",
+      "    \"height\": 224,\n",
+      "    \"width\": 224\n",
+      "  },\n",
+      "  \"do_center_crop\": true,\n",
+      "  \"do_convert_rgb\": true,\n",
+      "  \"do_normalize\": true,\n",
+      "  \"do_rescale\": true,\n",
+      "  \"do_resize\": true,\n",
+      "  \"feature_extractor_type\": \"CLIPFeatureExtractor\",\n",
+      "  \"image_mean\": [\n",
+      "    0.48145466,\n",
+      "    0.4578275,\n",
+      "    0.40821073\n",
+      "  ],\n",
+      "  \"image_processor_type\": \"CLIPImageProcessor\",\n",
+      "  \"image_std\": [\n",
+      "    0.26862954,\n",
+      "    0.26130258,\n",
+      "    0.27577711\n",
+      "  ],\n",
+      "  \"resample\": 3,\n",
+      "  \"rescale_factor\": 0.00392156862745098,\n",
+      "  \"size\": {\n",
+      "    \"shortest_edge\": 224\n",
+      "  }\n",
+      "}\n",
+      "\n",
+      "loading file vocab.json from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/vocab.json\n",
+      "loading file merges.txt from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/merges.txt\n",
+      "loading file tokenizer.json from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/tokenizer.json\n",
+      "loading file added_tokens.json from cache at None\n",
+      "loading file special_tokens_map.json from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/special_tokens_map.json\n",
+      "loading file tokenizer_config.json from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/tokenizer_config.json\n",
+      "loading configuration file config.json from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/config.json\n",
+      "`text_config` is `None`. Initializing the `CLIPTextConfig` with default values.\n",
+      "`vision_config` is `None`. initializing the `CLIPVisionConfig` with default values.\n",
+      "Model config CLIPConfig {\n",
+      "  \"_name_or_path\": \"openai/clip-vit-base-patch32\",\n",
+      "  \"architectures\": [\n",
+      "    \"CLIPModel\"\n",
+      "  ],\n",
+      "  \"initializer_factor\": 1.0,\n",
+      "  \"logit_scale_init_value\": 2.6592,\n",
+      "  \"model_type\": \"clip\",\n",
+      "  \"projection_dim\": 512,\n",
+      "  \"text_config\": {\n",
+      "    \"bos_token_id\": 0,\n",
+      "    \"dropout\": 0.0,\n",
+      "    \"eos_token_id\": 2,\n",
+      "    \"model_type\": \"clip_text_model\"\n",
+      "  },\n",
+      "  \"transformers_version\": \"4.36.2\",\n",
+      "  \"vision_config\": {\n",
+      "    \"dropout\": 0.0,\n",
+      "    \"model_type\": \"clip_vision_model\"\n",
+      "  }\n",
+      "}\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from PIL import Image\n",
+    "import requests\n",
+    "from transformers import AutoProcessor, CLIPModel\n",
+    "\n",
+    "model = CLIPModel.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
+    "processor = AutoProcessor.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
+    "\n",
+    "url = \"http://images.cocodataset.org/val2017/000000039769.jpg\"\n",
+    "image = Image.open(requests.get(url, stream=True).raw)\n",
+    "\n",
+    "inputs = processor(images=image, return_tensors=\"pt\")\n",
+    "\n",
+    "image_features = model.get_image_features(**inputs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "id": "9ff63766-b706-452b-b735-bf9000fb9c20",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "torch.Size([1, 512])"
+      ]
+     },
+     "execution_count": 29,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "image_features.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 30,
+   "id": "82566e7b-3c91-421a-94c5-f1e2b3e91c8c",
+   "metadata": {
+    "collapsed": true,
+    "jupyter": {
+     "outputs_hidden": true
+    }
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "loading configuration file config.json from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/config.json\n",
+      "Model config CLIPVisionConfig {\n",
+      "  \"attention_dropout\": 0.0,\n",
+      "  \"dropout\": 0.0,\n",
+      "  \"hidden_act\": \"quick_gelu\",\n",
+      "  \"hidden_size\": 768,\n",
+      "  \"image_size\": 224,\n",
+      "  \"initializer_factor\": 1.0,\n",
+      "  \"initializer_range\": 0.02,\n",
+      "  \"intermediate_size\": 3072,\n",
+      "  \"layer_norm_eps\": 1e-05,\n",
+      "  \"model_type\": \"clip_vision_model\",\n",
+      "  \"num_attention_heads\": 12,\n",
+      "  \"num_channels\": 3,\n",
+      "  \"num_hidden_layers\": 12,\n",
+      "  \"patch_size\": 32,\n",
+      "  \"projection_dim\": 512,\n",
+      "  \"transformers_version\": \"4.36.2\"\n",
+      "}\n",
+      "\n",
+      "loading weights file pytorch_model.bin from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/pytorch_model.bin\n",
+      "Some weights of the model checkpoint at openai/clip-vit-base-patch32 were not used when initializing CLIPVisionModel: ['text_model.encoder.layers.8.mlp.fc2.weight', 'text_model.encoder.layers.4.self_attn.v_proj.weight', 'text_model.encoder.layers.2.mlp.fc2.bias', 'text_model.encoder.layers.2.self_attn.q_proj.bias', 'text_model.encoder.layers.4.self_attn.v_proj.bias', 'text_model.encoder.layers.6.self_attn.k_proj.bias', 'text_model.encoder.layers.9.self_attn.k_proj.weight', 'text_model.encoder.layers.11.self_attn.q_proj.bias', 'text_model.encoder.layers.3.self_attn.out_proj.bias', 'text_model.encoder.layers.11.self_attn.k_proj.bias', 'text_model.encoder.layers.7.self_attn.k_proj.bias', 'text_model.encoder.layers.1.self_attn.q_proj.weight', 'text_model.encoder.layers.11.layer_norm1.bias', 'text_model.encoder.layers.11.mlp.fc2.bias', 'text_model.encoder.layers.10.layer_norm1.bias', 'text_model.encoder.layers.7.self_attn.q_proj.bias', 'text_model.encoder.layers.11.self_attn.k_proj.weight', 'text_model.encoder.layers.5.self_attn.q_proj.bias', 'text_model.encoder.layers.0.self_attn.v_proj.bias', 'logit_scale', 'text_model.encoder.layers.8.mlp.fc1.bias', 'text_model.encoder.layers.6.layer_norm1.bias', 'text_model.encoder.layers.5.self_attn.out_proj.weight', 'text_model.encoder.layers.7.self_attn.out_proj.bias', 'text_model.embeddings.token_embedding.weight', 'text_model.encoder.layers.8.layer_norm2.bias', 'text_model.encoder.layers.9.self_attn.v_proj.weight', 'text_model.encoder.layers.1.mlp.fc1.weight', 'text_model.encoder.layers.0.layer_norm1.weight', 'text_model.encoder.layers.6.self_attn.k_proj.weight', 'text_model.encoder.layers.3.self_attn.q_proj.weight', 'text_model.encoder.layers.2.layer_norm1.weight', 'text_model.encoder.layers.0.self_attn.v_proj.weight', 'text_model.encoder.layers.7.self_attn.q_proj.weight', 'text_model.encoder.layers.1.layer_norm2.weight', 'text_model.encoder.layers.2.self_attn.out_proj.weight', 'text_model.encoder.layers.3.self_attn.k_proj.weight', 'text_model.encoder.layers.7.mlp.fc2.bias', 'text_model.encoder.layers.10.self_attn.out_proj.weight', 'text_model.encoder.layers.2.self_attn.q_proj.weight', 'text_model.encoder.layers.1.self_attn.k_proj.weight', 'text_model.encoder.layers.4.layer_norm1.weight', 'text_model.encoder.layers.0.self_attn.q_proj.weight', 'text_model.encoder.layers.0.self_attn.out_proj.bias', 'text_model.encoder.layers.4.self_attn.out_proj.bias', 'text_model.encoder.layers.5.self_attn.k_proj.weight', 'visual_projection.weight', 'text_model.encoder.layers.6.layer_norm2.bias', 'text_model.encoder.layers.6.layer_norm1.weight', 'text_model.encoder.layers.4.self_attn.out_proj.weight', 'text_model.encoder.layers.10.mlp.fc2.bias', 'text_model.encoder.layers.10.mlp.fc1.weight', 'text_model.encoder.layers.6.self_attn.out_proj.weight', 'text_model.encoder.layers.9.layer_norm1.weight', 'text_model.encoder.layers.11.layer_norm2.weight', 'text_model.encoder.layers.6.self_attn.q_proj.bias', 'text_model.encoder.layers.5.mlp.fc1.weight', 'text_model.encoder.layers.2.mlp.fc1.weight', 'text_model.encoder.layers.11.self_attn.out_proj.weight', 'text_model.encoder.layers.0.self_attn.out_proj.weight', 'text_model.encoder.layers.11.mlp.fc2.weight', 'text_model.encoder.layers.7.layer_norm2.weight', 'text_model.encoder.layers.10.self_attn.v_proj.bias', 'text_model.encoder.layers.9.mlp.fc1.bias', 'text_model.encoder.layers.8.self_attn.v_proj.weight', 'text_model.encoder.layers.3.layer_norm1.bias', 'text_model.encoder.layers.6.self_attn.v_proj.bias', 'text_model.encoder.layers.1.self_attn.v_proj.bias', 'text_model.encoder.layers.9.self_attn.q_proj.weight', 'text_model.encoder.layers.4.self_attn.k_proj.weight', 'text_model.encoder.layers.7.layer_norm1.weight', 'text_model.encoder.layers.10.self_attn.k_proj.weight', 'text_model.encoder.layers.7.self_attn.v_proj.bias', 'text_model.encoder.layers.7.mlp.fc1.bias', 'text_model.encoder.layers.11.mlp.fc1.weight', 'text_model.encoder.layers.2.mlp.fc1.bias', 'text_model.encoder.layers.3.mlp.fc2.bias', 'text_model.encoder.layers.8.self_attn.q_proj.weight', 'text_model.encoder.layers.0.mlp.fc1.weight', 'text_model.encoder.layers.11.self_attn.out_proj.bias', 'text_model.encoder.layers.1.self_attn.v_proj.weight', 'text_model.encoder.layers.0.self_attn.k_proj.weight', 'text_model.encoder.layers.9.layer_norm1.bias', 'text_model.final_layer_norm.weight', 'text_model.encoder.layers.3.layer_norm1.weight', 'text_model.encoder.layers.4.mlp.fc1.bias', 'text_model.encoder.layers.1.layer_norm1.weight', 'text_model.encoder.layers.10.layer_norm2.bias', 'text_model.encoder.layers.9.self_attn.v_proj.bias', 'text_model.encoder.layers.10.self_attn.k_proj.bias', 'text_model.encoder.layers.8.mlp.fc2.bias', 'text_model.encoder.layers.5.mlp.fc2.bias', 'text_model.encoder.layers.6.self_attn.q_proj.weight', 'text_model.encoder.layers.5.self_attn.out_proj.bias', 'text_model.encoder.layers.9.mlp.fc2.bias', 'text_model.encoder.layers.5.layer_norm2.weight', 'text_model.encoder.layers.2.mlp.fc2.weight', 'text_model.encoder.layers.3.self_attn.out_proj.weight', 'text_model.encoder.layers.6.mlp.fc2.weight', 'text_model.encoder.layers.1.self_attn.out_proj.weight', 'text_model.encoder.layers.1.mlp.fc2.bias', 'text_model.encoder.layers.7.mlp.fc2.weight', 'text_model.encoder.layers.10.self_attn.v_proj.weight', 'text_model.encoder.layers.11.self_attn.v_proj.bias', 'text_model.encoder.layers.4.layer_norm1.bias', 'text_model.encoder.layers.4.layer_norm2.bias', 'text_model.encoder.layers.8.self_attn.q_proj.bias', 'text_model.embeddings.position_ids', 'text_model.encoder.layers.10.layer_norm2.weight', 'text_model.encoder.layers.1.self_attn.out_proj.bias', 'text_model.encoder.layers.2.layer_norm2.weight', 'text_model.encoder.layers.10.self_attn.q_proj.weight', 'text_model.encoder.layers.4.mlp.fc1.weight', 'text_model.encoder.layers.8.layer_norm1.bias', 'text_model.encoder.layers.2.self_attn.k_proj.weight', 'text_model.encoder.layers.5.mlp.fc1.bias', 'text_model.encoder.layers.9.self_attn.out_proj.bias', 'text_model.encoder.layers.7.self_attn.v_proj.weight', 'text_model.encoder.layers.2.self_attn.k_proj.bias', 'text_model.encoder.layers.5.self_attn.k_proj.bias', 'text_model.encoder.layers.8.self_attn.out_proj.bias', 'text_model.encoder.layers.7.self_attn.k_proj.weight', 'text_model.encoder.layers.6.mlp.fc1.weight', 'text_model.encoder.layers.6.mlp.fc1.bias', 'text_model.encoder.layers.3.self_attn.v_proj.weight', 'text_model.encoder.layers.3.self_attn.q_proj.bias', 'text_model.encoder.layers.9.self_attn.out_proj.weight', 'text_model.encoder.layers.3.mlp.fc1.bias', 'text_model.encoder.layers.0.self_attn.q_proj.bias', 'text_model.encoder.layers.1.layer_norm2.bias', 'text_model.encoder.layers.8.layer_norm2.weight', 'text_model.encoder.layers.5.self_attn.q_proj.weight', 'text_model.encoder.layers.4.layer_norm2.weight', 'text_model.encoder.layers.4.mlp.fc2.bias', 'text_model.encoder.layers.9.mlp.fc2.weight', 'text_model.encoder.layers.8.self_attn.k_proj.weight', 'text_model.encoder.layers.10.layer_norm1.weight', 'text_model.encoder.layers.0.self_attn.k_proj.bias', 'text_model.encoder.layers.8.self_attn.k_proj.bias', 'text_model.encoder.layers.9.layer_norm2.weight', 'text_model.encoder.layers.4.self_attn.k_proj.bias', 'text_model.encoder.layers.6.layer_norm2.weight', 'text_model.encoder.layers.0.layer_norm2.weight', 'text_model.encoder.layers.5.self_attn.v_proj.bias', 'text_model.encoder.layers.3.layer_norm2.bias', 'text_model.encoder.layers.8.mlp.fc1.weight', 'text_model.encoder.layers.4.self_attn.q_proj.bias', 'text_model.encoder.layers.8.layer_norm1.weight', 'text_model.encoder.layers.2.self_attn.v_proj.weight', 'text_model.encoder.layers.3.self_attn.v_proj.bias', 'text_model.encoder.layers.11.mlp.fc1.bias', 'text_model.encoder.layers.6.mlp.fc2.bias', 'text_model.encoder.layers.1.mlp.fc1.bias', 'text_model.encoder.layers.2.self_attn.v_proj.bias', 'text_model.encoder.layers.5.mlp.fc2.weight', 'text_model.encoder.layers.8.self_attn.v_proj.bias', 'text_model.encoder.layers.10.self_attn.out_proj.bias', 'text_model.encoder.layers.5.layer_norm1.bias', 'text_model.encoder.layers.5.self_attn.v_proj.weight', 'text_model.encoder.layers.10.self_attn.q_proj.bias', 'text_model.encoder.layers.2.layer_norm2.bias', 'text_model.encoder.layers.7.layer_norm1.bias', 'text_model.encoder.layers.4.mlp.fc2.weight', 'text_model.encoder.layers.10.mlp.fc2.weight', 'text_model.encoder.layers.3.mlp.fc1.weight', 'text_model.encoder.layers.5.layer_norm2.bias', 'text_model.encoder.layers.9.self_attn.q_proj.bias', 'text_model.encoder.layers.1.self_attn.k_proj.bias', 'text_model.encoder.layers.7.self_attn.out_proj.weight', 'text_model.encoder.layers.0.mlp.fc2.weight', 'text_model.encoder.layers.11.self_attn.v_proj.weight', 'text_model.encoder.layers.1.layer_norm1.bias', 'text_model.encoder.layers.1.mlp.fc2.weight', 'text_model.encoder.layers.9.layer_norm2.bias', 'text_model.encoder.layers.9.self_attn.k_proj.bias', 'text_model.encoder.layers.11.layer_norm1.weight', 'text_model.encoder.layers.8.self_attn.out_proj.weight', 'text_model.encoder.layers.0.layer_norm1.bias', 'text_model.encoder.layers.7.mlp.fc1.weight', 'text_model.encoder.layers.0.mlp.fc1.bias', 'text_model.encoder.layers.0.layer_norm2.bias', 'text_model.encoder.layers.3.self_attn.k_proj.bias', 'text_model.encoder.layers.5.layer_norm1.weight', 'text_model.encoder.layers.3.layer_norm2.weight', 'text_model.encoder.layers.1.self_attn.q_proj.bias', 'text_model.encoder.layers.2.self_attn.out_proj.bias', 'text_model.encoder.layers.3.mlp.fc2.weight', 'text_model.encoder.layers.11.self_attn.q_proj.weight', 'text_model.final_layer_norm.bias', 'text_model.encoder.layers.6.self_attn.v_proj.weight', 'text_model.encoder.layers.0.mlp.fc2.bias', 'text_model.encoder.layers.7.layer_norm2.bias', 'text_model.encoder.layers.10.mlp.fc1.bias', 'text_model.embeddings.position_embedding.weight', 'text_model.encoder.layers.6.self_attn.out_proj.bias', 'text_model.encoder.layers.2.layer_norm1.bias', 'text_model.encoder.layers.9.mlp.fc1.weight', 'text_projection.weight', 'text_model.encoder.layers.11.layer_norm2.bias', 'text_model.encoder.layers.4.self_attn.q_proj.weight']\n",
+      "- This IS expected if you are initializing CLIPVisionModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
+      "- This IS NOT expected if you are initializing CLIPVisionModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
+      "All the weights of CLIPVisionModel were initialized from the model checkpoint at openai/clip-vit-base-patch32.\n",
+      "If your task is similar to the task the model of the checkpoint was trained on, you can already use CLIPVisionModel for predictions without further training.\n",
+      "loading configuration file preprocessor_config.json from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/preprocessor_config.json\n",
+      "loading configuration file preprocessor_config.json from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/preprocessor_config.json\n",
+      "loading configuration file config.json from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/config.json\n",
+      "`text_config` is `None`. Initializing the `CLIPTextConfig` with default values.\n",
+      "`vision_config` is `None`. initializing the `CLIPVisionConfig` with default values.\n",
+      "Model config CLIPConfig {\n",
+      "  \"_name_or_path\": \"openai/clip-vit-base-patch32\",\n",
+      "  \"architectures\": [\n",
+      "    \"CLIPModel\"\n",
+      "  ],\n",
+      "  \"initializer_factor\": 1.0,\n",
+      "  \"logit_scale_init_value\": 2.6592,\n",
+      "  \"model_type\": \"clip\",\n",
+      "  \"projection_dim\": 512,\n",
+      "  \"text_config\": {\n",
+      "    \"bos_token_id\": 0,\n",
+      "    \"dropout\": 0.0,\n",
+      "    \"eos_token_id\": 2,\n",
+      "    \"model_type\": \"clip_text_model\"\n",
+      "  },\n",
+      "  \"transformers_version\": \"4.36.2\",\n",
+      "  \"vision_config\": {\n",
+      "    \"dropout\": 0.0,\n",
+      "    \"model_type\": \"clip_vision_model\"\n",
+      "  }\n",
+      "}\n",
+      "\n",
+      "loading configuration file preprocessor_config.json from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/preprocessor_config.json\n",
+      "size should be a dictionary on of the following set of keys: ({'width', 'height'}, {'shortest_edge'}, {'longest_edge', 'shortest_edge'}, {'longest_edge'}), got 224. Converted to {'shortest_edge': 224}.\n",
+      "crop_size should be a dictionary on of the following set of keys: ({'width', 'height'}, {'shortest_edge'}, {'longest_edge', 'shortest_edge'}, {'longest_edge'}), got 224. Converted to {'height': 224, 'width': 224}.\n",
+      "Image processor CLIPImageProcessor {\n",
+      "  \"crop_size\": {\n",
+      "    \"height\": 224,\n",
+      "    \"width\": 224\n",
+      "  },\n",
+      "  \"do_center_crop\": true,\n",
+      "  \"do_convert_rgb\": true,\n",
+      "  \"do_normalize\": true,\n",
+      "  \"do_rescale\": true,\n",
+      "  \"do_resize\": true,\n",
+      "  \"feature_extractor_type\": \"CLIPFeatureExtractor\",\n",
+      "  \"image_mean\": [\n",
+      "    0.48145466,\n",
+      "    0.4578275,\n",
+      "    0.40821073\n",
+      "  ],\n",
+      "  \"image_processor_type\": \"CLIPImageProcessor\",\n",
+      "  \"image_std\": [\n",
+      "    0.26862954,\n",
+      "    0.26130258,\n",
+      "    0.27577711\n",
+      "  ],\n",
+      "  \"resample\": 3,\n",
+      "  \"rescale_factor\": 0.00392156862745098,\n",
+      "  \"size\": {\n",
+      "    \"shortest_edge\": 224\n",
+      "  }\n",
+      "}\n",
+      "\n",
+      "loading file vocab.json from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/vocab.json\n",
+      "loading file merges.txt from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/merges.txt\n",
+      "loading file tokenizer.json from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/tokenizer.json\n",
+      "loading file added_tokens.json from cache at None\n",
+      "loading file special_tokens_map.json from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/special_tokens_map.json\n",
+      "loading file tokenizer_config.json from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/tokenizer_config.json\n",
+      "loading configuration file config.json from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/config.json\n",
+      "`text_config` is `None`. Initializing the `CLIPTextConfig` with default values.\n",
+      "`vision_config` is `None`. initializing the `CLIPVisionConfig` with default values.\n",
+      "Model config CLIPConfig {\n",
+      "  \"_name_or_path\": \"openai/clip-vit-base-patch32\",\n",
+      "  \"architectures\": [\n",
+      "    \"CLIPModel\"\n",
+      "  ],\n",
+      "  \"initializer_factor\": 1.0,\n",
+      "  \"logit_scale_init_value\": 2.6592,\n",
+      "  \"model_type\": \"clip\",\n",
+      "  \"projection_dim\": 512,\n",
+      "  \"text_config\": {\n",
+      "    \"bos_token_id\": 0,\n",
+      "    \"dropout\": 0.0,\n",
+      "    \"eos_token_id\": 2,\n",
+      "    \"model_type\": \"clip_text_model\"\n",
+      "  },\n",
+      "  \"transformers_version\": \"4.36.2\",\n",
+      "  \"vision_config\": {\n",
+      "    \"dropout\": 0.0,\n",
+      "    \"model_type\": \"clip_vision_model\"\n",
+      "  }\n",
+      "}\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from PIL import Image\n",
+    "import requests\n",
+    "from transformers import AutoProcessor, CLIPVisionModel\n",
+    "\n",
+    "model = CLIPVisionModel.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
+    "processor = AutoProcessor.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
+    "\n",
+    "url = \"http://images.cocodataset.org/val2017/000000039769.jpg\"\n",
+    "image = Image.open(requests.get(url, stream=True).raw)\n",
+    "\n",
+    "inputs = processor(images=image, return_tensors=\"pt\")\n",
+    "\n",
+    "outputs = model(**inputs)\n",
+    "last_hidden_state = outputs.last_hidden_state\n",
+    "pooled_output = outputs.pooler_output  # pooled CLS states"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "id": "bcf0a7b3-6cbb-492e-bc2c-42e3edbe6a0c",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "torch.Size([1, 768])"
+      ]
+     },
+     "execution_count": 31,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "pooled_output.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "67240294-c7a0-4e94-a8c1-86bfe1b21977",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from transformers import CLIPPreTrainedModel\n",
+    "from transformers.models.clip.modeling_clip import CLIPVisionModelOutput, CLIPVisionTransformer\n",
+    "from typing import Optional, Union, Tuple"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 54,
+   "id": "cc9b20db-7f84-44c3-9c78-e84164ccc192",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class VisionLanguageConnector(nn.Module):\n",
+    "    def __init__(self, hidden_size, projection_dim):\n",
+    "        super().__init__()\n",
+    "        self.mlp = nn.Sequential(\n",
+    "            nn.Linear(hidden_size, hidden_size, bias=False),\n",
+    "            nn.GELU(),\n",
+    "            nn.Linear(hidden_size, projection_dim, bias=False)\n",
+    "        )\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        return self.mlp(x)\n",
+    "        \n",
+    "class ClipWithProjection(CLIPPreTrainedModel):\n",
+    "    config_class = CLIPVisionConfig\n",
+    "    main_input_name = \"pixel_values\"\n",
+    "\n",
+    "    def __init__(self, config: CLIPVisionConfig):\n",
+    "        super().__init__(config)\n",
+    "\n",
+    "        self.vision_model = CLIPVisionTransformer(config)\n",
+    "        self.vision_model.\n",
+    "        self.vision_language_connector = VisionLanguageConnector(config.hidden_size, config.projection_dim)\n",
+    "\n",
+    "        # Initialize weights and apply final processing\n",
+    "        self.post_init()\n",
+    "\n",
+    "    def forward(\n",
+    "        self,\n",
+    "        pixel_values: Optional[torch.FloatTensor] = None,\n",
+    "        output_attentions: Optional[bool] = None,\n",
+    "        output_hidden_states: Optional[bool] = None,\n",
+    "        return_dict: Optional[bool] = None,\n",
+    "    ) -> Union[Tuple, CLIPVisionModelOutput]:\n",
+    "        return_dict = return_dict if return_dict is not None else self.config.use_return_dict\n",
+    "\n",
+    "        vision_outputs = self.vision_model(\n",
+    "            pixel_values=pixel_values,\n",
+    "            output_attentions=output_attentions,\n",
+    "            output_hidden_states=output_hidden_states,\n",
+    "            return_dict=return_dict,\n",
+    "        )\n",
+    "\n",
+    "        pooled_output = vision_outputs[1]  # pooled_output\n",
+    "\n",
+    "        image_embeds = self.vision_language_connector(pooled_output)\n",
+    "\n",
+    "        if not return_dict:\n",
+    "            outputs = (image_embeds, vision_outputs[0]) + vision_outputs[2:]\n",
+    "            return tuple(output for output in outputs if output is not None)\n",
+    "\n",
+    "        return CLIPVisionModelOutput(\n",
+    "            image_embeds=image_embeds,\n",
+    "            last_hidden_state=vision_outputs.last_hidden_state,\n",
+    "            hidden_states=vision_outputs.hidden_states,\n",
+    "            attentions=vision_outputs.attentions,\n",
+    "        )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 55,
+   "id": "a4892ab8-39d2-41c9-ad2a-04711c22b95f",
+   "metadata": {
+    "collapsed": true,
+    "jupyter": {
+     "outputs_hidden": true
+    }
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "loading configuration file config.json from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/config.json\n",
+      "Model config CLIPVisionConfig {\n",
+      "  \"attention_dropout\": 0.0,\n",
+      "  \"dropout\": 0.0,\n",
+      "  \"hidden_act\": \"quick_gelu\",\n",
+      "  \"hidden_size\": 768,\n",
+      "  \"image_size\": 224,\n",
+      "  \"initializer_factor\": 1.0,\n",
+      "  \"initializer_range\": 0.02,\n",
+      "  \"intermediate_size\": 3072,\n",
+      "  \"layer_norm_eps\": 1e-05,\n",
+      "  \"model_type\": \"clip_vision_model\",\n",
+      "  \"num_attention_heads\": 12,\n",
+      "  \"num_channels\": 3,\n",
+      "  \"num_hidden_layers\": 12,\n",
+      "  \"patch_size\": 32,\n",
+      "  \"projection_dim\": 512,\n",
+      "  \"transformers_version\": \"4.36.2\"\n",
+      "}\n",
+      "\n",
+      "loading weights file pytorch_model.bin from cache at /home/ravi.naik/.cache/huggingface/hub/models--openai--clip-vit-base-patch32/snapshots/e6a30b603a447e251fdaca1c3056b2a16cdfebeb/pytorch_model.bin\n",
+      "Some weights of the model checkpoint at openai/clip-vit-base-patch32 were not used when initializing ClipWithProjection: ['text_model.encoder.layers.8.mlp.fc2.weight', 'text_model.encoder.layers.4.self_attn.v_proj.weight', 'text_model.encoder.layers.2.mlp.fc2.bias', 'text_model.encoder.layers.2.self_attn.q_proj.bias', 'text_model.encoder.layers.4.self_attn.v_proj.bias', 'text_model.encoder.layers.6.self_attn.k_proj.bias', 'text_model.encoder.layers.9.self_attn.k_proj.weight', 'text_model.encoder.layers.11.self_attn.q_proj.bias', 'text_model.encoder.layers.3.self_attn.out_proj.bias', 'text_model.encoder.layers.11.self_attn.k_proj.bias', 'text_model.encoder.layers.7.self_attn.k_proj.bias', 'text_model.encoder.layers.1.self_attn.q_proj.weight', 'text_model.encoder.layers.11.layer_norm1.bias', 'text_model.encoder.layers.11.mlp.fc2.bias', 'text_model.encoder.layers.10.layer_norm1.bias', 'text_model.encoder.layers.7.self_attn.q_proj.bias', 'text_model.encoder.layers.11.self_attn.k_proj.weight', 'text_model.encoder.layers.5.self_attn.q_proj.bias', 'text_model.encoder.layers.0.self_attn.v_proj.bias', 'logit_scale', 'text_model.encoder.layers.8.mlp.fc1.bias', 'text_model.encoder.layers.6.layer_norm1.bias', 'text_model.encoder.layers.5.self_attn.out_proj.weight', 'text_model.encoder.layers.7.self_attn.out_proj.bias', 'text_model.embeddings.token_embedding.weight', 'text_model.encoder.layers.8.layer_norm2.bias', 'text_model.encoder.layers.9.self_attn.v_proj.weight', 'text_model.encoder.layers.1.mlp.fc1.weight', 'text_model.encoder.layers.0.layer_norm1.weight', 'text_model.encoder.layers.6.self_attn.k_proj.weight', 'text_model.encoder.layers.3.self_attn.q_proj.weight', 'text_model.encoder.layers.2.layer_norm1.weight', 'text_model.encoder.layers.0.self_attn.v_proj.weight', 'text_model.encoder.layers.7.self_attn.q_proj.weight', 'text_model.encoder.layers.1.layer_norm2.weight', 'text_model.encoder.layers.2.self_attn.out_proj.weight', 'text_model.encoder.layers.3.self_attn.k_proj.weight', 'text_model.encoder.layers.7.mlp.fc2.bias', 'text_model.encoder.layers.10.self_attn.out_proj.weight', 'text_model.encoder.layers.2.self_attn.q_proj.weight', 'text_model.encoder.layers.1.self_attn.k_proj.weight', 'text_model.encoder.layers.4.layer_norm1.weight', 'text_model.encoder.layers.0.self_attn.q_proj.weight', 'text_model.encoder.layers.0.self_attn.out_proj.bias', 'text_model.encoder.layers.4.self_attn.out_proj.bias', 'text_model.encoder.layers.5.self_attn.k_proj.weight', 'visual_projection.weight', 'text_model.encoder.layers.6.layer_norm2.bias', 'text_model.encoder.layers.6.layer_norm1.weight', 'text_model.encoder.layers.4.self_attn.out_proj.weight', 'text_model.encoder.layers.10.mlp.fc2.bias', 'text_model.encoder.layers.10.mlp.fc1.weight', 'text_model.encoder.layers.6.self_attn.out_proj.weight', 'text_model.encoder.layers.9.layer_norm1.weight', 'text_model.encoder.layers.11.layer_norm2.weight', 'text_model.encoder.layers.6.self_attn.q_proj.bias', 'text_model.encoder.layers.5.mlp.fc1.weight', 'text_model.encoder.layers.2.mlp.fc1.weight', 'text_model.encoder.layers.11.self_attn.out_proj.weight', 'text_model.encoder.layers.0.self_attn.out_proj.weight', 'text_model.encoder.layers.11.mlp.fc2.weight', 'text_model.encoder.layers.7.layer_norm2.weight', 'text_model.encoder.layers.10.self_attn.v_proj.bias', 'text_model.encoder.layers.9.mlp.fc1.bias', 'text_model.encoder.layers.8.self_attn.v_proj.weight', 'text_model.encoder.layers.3.layer_norm1.bias', 'text_model.encoder.layers.6.self_attn.v_proj.bias', 'text_model.encoder.layers.1.self_attn.v_proj.bias', 'text_model.encoder.layers.9.self_attn.q_proj.weight', 'text_model.encoder.layers.4.self_attn.k_proj.weight', 'text_model.encoder.layers.7.layer_norm1.weight', 'text_model.encoder.layers.10.self_attn.k_proj.weight', 'text_model.encoder.layers.7.self_attn.v_proj.bias', 'text_model.encoder.layers.7.mlp.fc1.bias', 'text_model.encoder.layers.11.mlp.fc1.weight', 'text_model.encoder.layers.2.mlp.fc1.bias', 'text_model.encoder.layers.3.mlp.fc2.bias', 'text_model.encoder.layers.8.self_attn.q_proj.weight', 'text_model.encoder.layers.0.mlp.fc1.weight', 'text_model.encoder.layers.11.self_attn.out_proj.bias', 'text_model.encoder.layers.1.self_attn.v_proj.weight', 'text_model.encoder.layers.0.self_attn.k_proj.weight', 'text_model.encoder.layers.9.layer_norm1.bias', 'text_model.final_layer_norm.weight', 'text_model.encoder.layers.3.layer_norm1.weight', 'text_model.encoder.layers.4.mlp.fc1.bias', 'text_model.encoder.layers.1.layer_norm1.weight', 'text_model.encoder.layers.10.layer_norm2.bias', 'text_model.encoder.layers.9.self_attn.v_proj.bias', 'text_model.encoder.layers.10.self_attn.k_proj.bias', 'text_model.encoder.layers.8.mlp.fc2.bias', 'text_model.encoder.layers.5.mlp.fc2.bias', 'text_model.encoder.layers.6.self_attn.q_proj.weight', 'text_model.encoder.layers.5.self_attn.out_proj.bias', 'text_model.encoder.layers.9.mlp.fc2.bias', 'text_model.encoder.layers.5.layer_norm2.weight', 'text_model.encoder.layers.2.mlp.fc2.weight', 'text_model.encoder.layers.3.self_attn.out_proj.weight', 'text_model.encoder.layers.6.mlp.fc2.weight', 'text_model.encoder.layers.1.self_attn.out_proj.weight', 'text_model.encoder.layers.1.mlp.fc2.bias', 'text_model.encoder.layers.7.mlp.fc2.weight', 'text_model.encoder.layers.10.self_attn.v_proj.weight', 'text_model.encoder.layers.11.self_attn.v_proj.bias', 'text_model.encoder.layers.4.layer_norm1.bias', 'text_model.encoder.layers.4.layer_norm2.bias', 'text_model.encoder.layers.8.self_attn.q_proj.bias', 'text_model.embeddings.position_ids', 'text_model.encoder.layers.10.layer_norm2.weight', 'text_model.encoder.layers.1.self_attn.out_proj.bias', 'text_model.encoder.layers.2.layer_norm2.weight', 'text_model.encoder.layers.10.self_attn.q_proj.weight', 'text_model.encoder.layers.4.mlp.fc1.weight', 'text_model.encoder.layers.8.layer_norm1.bias', 'text_model.encoder.layers.2.self_attn.k_proj.weight', 'text_model.encoder.layers.5.mlp.fc1.bias', 'text_model.encoder.layers.9.self_attn.out_proj.bias', 'text_model.encoder.layers.7.self_attn.v_proj.weight', 'text_model.encoder.layers.2.self_attn.k_proj.bias', 'text_model.encoder.layers.5.self_attn.k_proj.bias', 'text_model.encoder.layers.8.self_attn.out_proj.bias', 'text_model.encoder.layers.7.self_attn.k_proj.weight', 'text_model.encoder.layers.6.mlp.fc1.weight', 'text_model.encoder.layers.6.mlp.fc1.bias', 'text_model.encoder.layers.3.self_attn.v_proj.weight', 'text_model.encoder.layers.3.self_attn.q_proj.bias', 'text_model.encoder.layers.9.self_attn.out_proj.weight', 'text_model.encoder.layers.3.mlp.fc1.bias', 'text_model.encoder.layers.0.self_attn.q_proj.bias', 'text_model.encoder.layers.1.layer_norm2.bias', 'text_model.encoder.layers.8.layer_norm2.weight', 'text_model.encoder.layers.5.self_attn.q_proj.weight', 'text_model.encoder.layers.4.layer_norm2.weight', 'text_model.encoder.layers.4.mlp.fc2.bias', 'text_model.encoder.layers.9.mlp.fc2.weight', 'text_model.encoder.layers.8.self_attn.k_proj.weight', 'text_model.encoder.layers.10.layer_norm1.weight', 'text_model.encoder.layers.0.self_attn.k_proj.bias', 'text_model.encoder.layers.8.self_attn.k_proj.bias', 'text_model.encoder.layers.9.layer_norm2.weight', 'text_model.encoder.layers.4.self_attn.k_proj.bias', 'text_model.encoder.layers.6.layer_norm2.weight', 'text_model.encoder.layers.0.layer_norm2.weight', 'text_model.encoder.layers.5.self_attn.v_proj.bias', 'text_model.encoder.layers.3.layer_norm2.bias', 'text_model.encoder.layers.8.mlp.fc1.weight', 'text_model.encoder.layers.4.self_attn.q_proj.bias', 'text_model.encoder.layers.8.layer_norm1.weight', 'text_model.encoder.layers.2.self_attn.v_proj.weight', 'text_model.encoder.layers.3.self_attn.v_proj.bias', 'text_model.encoder.layers.11.mlp.fc1.bias', 'text_model.encoder.layers.6.mlp.fc2.bias', 'text_model.encoder.layers.1.mlp.fc1.bias', 'text_model.encoder.layers.2.self_attn.v_proj.bias', 'text_model.encoder.layers.5.mlp.fc2.weight', 'text_model.encoder.layers.8.self_attn.v_proj.bias', 'text_model.encoder.layers.10.self_attn.out_proj.bias', 'text_model.encoder.layers.5.layer_norm1.bias', 'text_model.encoder.layers.5.self_attn.v_proj.weight', 'text_model.encoder.layers.10.self_attn.q_proj.bias', 'text_model.encoder.layers.2.layer_norm2.bias', 'text_model.encoder.layers.7.layer_norm1.bias', 'text_model.encoder.layers.4.mlp.fc2.weight', 'text_model.encoder.layers.10.mlp.fc2.weight', 'text_model.encoder.layers.3.mlp.fc1.weight', 'text_model.encoder.layers.5.layer_norm2.bias', 'text_model.encoder.layers.9.self_attn.q_proj.bias', 'text_model.encoder.layers.1.self_attn.k_proj.bias', 'text_model.encoder.layers.7.self_attn.out_proj.weight', 'text_model.encoder.layers.0.mlp.fc2.weight', 'text_model.encoder.layers.11.self_attn.v_proj.weight', 'text_model.encoder.layers.1.layer_norm1.bias', 'text_model.encoder.layers.1.mlp.fc2.weight', 'text_model.encoder.layers.9.layer_norm2.bias', 'text_model.encoder.layers.9.self_attn.k_proj.bias', 'text_model.encoder.layers.11.layer_norm1.weight', 'text_model.encoder.layers.8.self_attn.out_proj.weight', 'text_model.encoder.layers.0.layer_norm1.bias', 'text_model.encoder.layers.7.mlp.fc1.weight', 'text_model.encoder.layers.0.mlp.fc1.bias', 'text_model.encoder.layers.0.layer_norm2.bias', 'text_model.encoder.layers.3.self_attn.k_proj.bias', 'text_model.encoder.layers.5.layer_norm1.weight', 'text_model.encoder.layers.3.layer_norm2.weight', 'text_model.encoder.layers.1.self_attn.q_proj.bias', 'text_model.encoder.layers.2.self_attn.out_proj.bias', 'text_model.encoder.layers.3.mlp.fc2.weight', 'text_model.encoder.layers.11.self_attn.q_proj.weight', 'text_model.final_layer_norm.bias', 'text_model.encoder.layers.6.self_attn.v_proj.weight', 'text_model.encoder.layers.0.mlp.fc2.bias', 'text_model.encoder.layers.7.layer_norm2.bias', 'text_model.encoder.layers.10.mlp.fc1.bias', 'text_model.embeddings.position_embedding.weight', 'text_model.encoder.layers.6.self_attn.out_proj.bias', 'text_model.encoder.layers.2.layer_norm1.bias', 'text_model.encoder.layers.9.mlp.fc1.weight', 'text_projection.weight', 'text_model.encoder.layers.11.layer_norm2.bias', 'text_model.encoder.layers.4.self_attn.q_proj.weight']\n",
+      "- This IS expected if you are initializing ClipWithProjection from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
+      "- This IS NOT expected if you are initializing ClipWithProjection from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
+      "Some weights of ClipWithProjection were not initialized from the model checkpoint at openai/clip-vit-base-patch32 and are newly initialized: ['vision_language_connector.mlp.2.weight', 'vision_language_connector.mlp.0.weight']\n",
+      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
+     ]
+    }
+   ],
+   "source": [
+    "model = ClipWithProjection.from_pretrained(\"openai/clip-vit-base-patch32\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 56,
+   "id": "588ef914-5be9-49e1-b68d-b899e0e74edd",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "768"
+      ]
+     },
+     "execution_count": 56,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model.config.hidden_size"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 57,
+   "id": "05d95b9e-9831-4415-860e-94793e29d210",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "outputs = model(**inputs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 61,
+   "id": "185b1bff-6ffe-4cce-9255-ee7629feba54",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "torch.Size([1, 512])"
+      ]
+     },
+     "execution_count": 61,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "outputs[0].shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "04414a35-c7b3-4986-a79e-1d363916caa4",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "485dbbcb-06df-4926-b257-dfd1a4081d44",
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "NameError",
+     "evalue": "name 'outputs' is not defined",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[1], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43moutputs\u001b[49m[\u001b[38;5;241m0\u001b[39m]\n",
+      "\u001b[0;31mNameError\u001b[0m: name 'outputs' is not defined"
+     ]
+    }
+   ],
+   "source": [
+    "outputs[0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f983313c-8e0f-4805-af14-25bb69afd04c",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

Experiments/eval.ipynb

@@ -0,0 +1,782 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "215cfd2f-62b0-4a86-a407-777a1d32597f",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[2024-01-24 15:18:49,948] [INFO] [real_accelerator.py:110:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n"
+     ]
+    }
+   ],
+   "source": [
+    "from PIL import Image\n",
+    "import requests\n",
+    "\n",
+    "import torch\n",
+    "from torch import nn\n",
+    "from transformers import AutoProcessor, CLIPVisionModel, CLIPVisionConfig, CLIPPreTrainedModel\n",
+    "from transformers.models.clip.modeling_clip import CLIPVisionModelOutput, CLIPVisionTransformer\n",
+    "from transformers import WhisperProcessor, WhisperForConditionalGeneration\n",
+    "from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, AutoTokenizer"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "2244e8f3-fcc7-4309-9d4d-fea557f89f79",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from llava_phi import LlavaPhiForCausalLM"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "587883e1-3419-4b14-b16b-38fabbc8bfaa",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# model = LlavaPhiForCausalLM.from_pretrained(\"./llava-phi/checkpoints/llavaPhi-v0-3b-finetune/checkpoint-4000\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "0e27a7db-e2ab-4d65-b21d-497222e318ad",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# processor = AutoProcessor.from_pretrained(\"./llava-phi/checkpoints/llavaPhi-v0-3b-finetune/checkpoint-4000\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "663efdd8-ea21-4231-a2ae-bcc0fb47b46a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# prompt = \"<image>\\nUSER: What's the content of the image?\\nASSISTANT:\"\n",
+    "# url = \"https://www.ilankelman.org/stopsigns/australia.jpg\"\n",
+    "# image = Image.open(requests.get(url, stream=True).raw)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "f622609f-f6a7-4ec1-ac35-c1d33d9436ca",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# # Generate\n",
+    "# generate_ids = model.generate(**inputs, max_length=30)\n",
+    "# processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "45f5ba72-2e41-4ccc-84c1-97d542ebee63",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from llava_phi.model.builder import load_pretrained_model\n",
+    "from llava_phi.mm_utils import tokenizer_image_token, get_model_name_from_path\n",
+    "from llava_phi.utils import disable_torch_init\n",
+    "from llava_phi.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN\n",
+    "from llava_phi.conversation import conv_templates, SeparatorStyle"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "b98ac5d3-5503-4430-81d1-19a4f8d6bd75",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_path = \"checkpoints/llavaPhi-v0-3b-finetune/checkpoint-4000\"\n",
+    "model_name = get_model_name_from_path(model_path)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "42fd5721-75a7-475b-bd30-5ee23aeaac64",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'llavaPhi-v0-3b-finetune_checkpoint-4000'"
+      ]
+     },
+     "execution_count": 12,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model_name"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "8c2076b5-3bfc-48fd-917b-5dfd06fc532f",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "load llaVA-Phi MLLM!!!\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
+      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "20b86f2c01744081b537620c8780f12e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'device_map': 'cuda'}\n"
+     ]
+    }
+   ],
+   "source": [
+    "tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "4e46221e-0907-453e-8126-76199828493e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "qs = \"What's the content of the image?\"\n",
+    "qs = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\\n' + qs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "07355444-0eb8-4d4d-ad50-48b91c969664",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "conv = conv_templates[\"default\"].copy()\n",
+    "conv.append_message(conv.roles[0], qs)\n",
+    "conv.append_message(conv.roles[1], None)\n",
+    "prompt = conv.get_prompt()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "ccb5674f-aff8-456e-b61b-1d167864f1a6",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "\"A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: <im_start><image><im_end>\\nWhat's the content of the image? ASSISTANT:\""
+      ]
+     },
+     "execution_count": 16,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "prompt"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "a89cc181-2214-4844-b966-164a41744e54",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "url = \"https://www.ilankelman.org/stopsigns/australia.jpg\"\n",
+    "image = Image.open(requests.get(url, stream=True).raw)\n",
+    "image_tensor = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].cuda()\n",
+    "\n",
+    "input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda()\n",
+    "\n",
+    "stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "id": "0d519851-64d4-4cf5-b2eb-19474f9aa260",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "torch.Size([1, 55])"
+      ]
+     },
+     "execution_count": 25,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "input_ids.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "id": "1694ff36-f214-4ed3-b2f3-d3dbd0a1a25b",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
+     ]
+    }
+   ],
+   "source": [
+    "from datasets import load_dataset\n",
+    "audio_ds = load_dataset(\"hf-internal-testing/librispeech_asr_dummy\", \"clean\", split=\"validation\")\n",
+    "audio = audio_ds[0][\"audio\"]\n",
+    "\n",
+    "whisper_w_proj = WhisperWithProjection(projection_dim=512)\n",
+    "audio_embed = whisper_w_proj(audio)[\"input_ids\"]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 28,
+   "id": "9c4a9fae-d6ed-4fc2-ba02-97df64cddd93",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(torch.Size([1, 33]), device(type='cpu'))"
+      ]
+     },
+     "execution_count": 28,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "audio_embed.shape, audio_embed.device"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "id": "c3fffe29-98fb-4f4b-ac51-4bdda9e46752",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "input_ids = torch.concat([input_ids, audio_embed.to(\"cuda:0\")], dim=1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 30,
+   "id": "5dee1ec8-2db2-4f65-99e8-d34bd2735c9c",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "torch.Size([1, 88])"
+      ]
+     },
+     "execution_count": 30,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "input_ids.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "id": "96033b43-4f57-4f0c-bcf7-37b57ca02e47",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with torch.inference_mode():\n",
+    "        output_ids = model.generate(\n",
+    "            input_ids,\n",
+    "            images=image_tensor,\n",
+    "            do_sample=True,\n",
+    "            temperature=0.2,\n",
+    "            max_new_tokens=1024,\n",
+    "            eos_token_id=tokenizer.eos_token_id,  # End of sequence token\n",
+    "            pad_token_id=tokenizer.eos_token_id,  # Pad token\n",
+    "            use_cache=True,\n",
+    "        )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 32,
+   "id": "741e8da5-0d18-4c11-b559-76054ce4ca3a",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "is a Japanese character from the story of Jesus, who is a Chinese monk who is also known for his teachings. The story is based on the story of the story of Jesus Christ, and it is a representation of the story of Jesus and the story of Jesus Christ.\n"
+     ]
+    }
+   ],
+   "source": [
+    "input_token_len = input_ids.shape[1]\n",
+    "n_diff_input_output = (input_ids != output_ids[:, :input_token_len]).sum().item()\n",
+    "if n_diff_input_output > 0:\n",
+    "    print(f'[Warning] {n_diff_input_output} output_ids are not the same as the input_ids')\n",
+    "outputs = tokenizer.batch_decode(output_ids[:, input_token_len:], skip_special_tokens=True)[0]\n",
+    "outputs = outputs.strip()\n",
+    "if outputs.endswith(stop_str):\n",
+    "    outputs = outputs[:-len(stop_str)]\n",
+    "outputs = outputs.strip()\n",
+    "print(outputs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "69d494d4-d768-4645-b4d6-5c455791b50d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8a340856-a13f-4b18-9911-126a4ba37816",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3c56fdea-c7a1-4e67-9832-e2ed077d8704",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 52,
+   "id": "89e84d39-8ed8-45db-ae82-27c156ee6dd1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class AudioLanguageConnector:\n",
+    "    def __init__(self, projection_dim):\n",
+    "        model_name = \"microsoft/phi-2\"\n",
+    "        self.phi2_tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n",
+    "        self.phi2_tokenizer.pad_token = self.phi2_tokenizer.eos_token\n",
+    "        self.phi2_tokenizer.max_length = projection_dim\n",
+    "\n",
+    "    def __call__(self, text):\n",
+    "        text = f\"<audio_start> {text} <audio_end>\"\n",
+    "        tokens = self.phi2_tokenizer(text, return_tensors=\"pt\", return_attention_mask=False)\n",
+    "        return tokens\n",
+    "        \n",
+    "\n",
+    "class WhisperWithProjection:\n",
+    "    def __init__(self, projection_dim, device):\n",
+    "        self.device = device\n",
+    "        self.processor = WhisperProcessor.from_pretrained(\"openai/whisper-tiny\", device_map=device)\n",
+    "        self.model = WhisperForConditionalGeneration.from_pretrained(\"openai/whisper-tiny\", device_map=device)\n",
+    "        self.model.config.forced_decoder_ids = None\n",
+    "        self.audio_language_connector = AudioLanguageConnector(projection_dim)\n",
+    "        \n",
+    "    def __call__(self, audio):\n",
+    "        input_features = self.processor(audio[\"array\"],\n",
+    "                                   sampling_rate=audio[\"sampling_rate\"],\n",
+    "                                   return_tensors=\"pt\").input_features\n",
+    "        # generate token ids\n",
+    "        predicted_ids = self.model.generate(input_features.to(self.device))\n",
+    "        # decode token ids to text        \n",
+    "        transcription = self.processor.batch_decode(predicted_ids, skip_special_tokens=True)\n",
+    "\n",
+    "        audio_embeddings = self.audio_language_connector(transcription)\n",
+    "        return audio_embeddings.to(self.device)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 53,
+   "id": "75e24be0-b236-4047-83ef-5c344e262476",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class MultiModalPhi2:\n",
+    "    def __init__(self, model_path=\"checkpoints/llavaPhi-v0-3b-finetune/checkpoint-4000\",\n",
+    "                temperature=0.2,\n",
+    "                max_new_tokens=1024,\n",
+    "                device=\"cuda\"):\n",
+    "        self.temperature = temperature\n",
+    "        self.max_new_tokens = max_new_tokens\n",
+    "        self.device = device\n",
+    "        model_name = get_model_name_from_path(model_path)\n",
+    "        self.tokenizer, self.model, self.image_processor, self.context_len = load_pretrained_model(model_path, None, model_name, device_map=device)\n",
+    "        self.whisper_w_proj = WhisperWithProjection(projection_dim=512, device=device)\n",
+    "        \n",
+    "        \n",
+    "    def __call__(self, text, audio, image):\n",
+    "        qs = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\\n' + text\n",
+    "        conv = conv_templates[\"default\"].copy()\n",
+    "        conv.append_message(conv.roles[0], qs)\n",
+    "        conv.append_message(conv.roles[1], None)\n",
+    "        prompt = conv.get_prompt()\n",
+    "\n",
+    "        image_tensor = self.image_processor.preprocess(image, return_tensors='pt')['pixel_values'].cuda()\n",
+    "        \n",
+    "        input_ids = tokenizer_image_token(prompt, self.tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda()\n",
+    "\n",
+    "        audio_embed = self.whisper_w_proj(audio)[\"input_ids\"]\n",
+    "        \n",
+    "        stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2\n",
+    "\n",
+    "        input_ids = torch.concat([input_ids, audio_embed], dim=1)\n",
+    "\n",
+    "        with torch.inference_mode():\n",
+    "            output_ids = self.model.generate(\n",
+    "                input_ids,\n",
+    "                images=image_tensor,\n",
+    "                do_sample=True,\n",
+    "                temperature=self.temperature,\n",
+    "                max_new_tokens=self.max_new_tokens,\n",
+    "                eos_token_id=tokenizer.eos_token_id,  # End of sequence token\n",
+    "                pad_token_id=tokenizer.eos_token_id,  # Pad token\n",
+    "                use_cache=True,\n",
+    "            )\n",
+    "\n",
+    "        input_token_len = input_ids.shape[1]\n",
+    "        n_diff_input_output = (input_ids != output_ids[:, :input_token_len]).sum().item()\n",
+    "        if n_diff_input_output > 0:\n",
+    "            print(f'[Warning] {n_diff_input_output} output_ids are not the same as the input_ids')\n",
+    "        outputs = self.tokenizer.batch_decode(output_ids[:, input_token_len:], skip_special_tokens=True)[0]\n",
+    "        outputs = outputs.strip()\n",
+    "        if outputs.endswith(stop_str):\n",
+    "            outputs = outputs[:-len(stop_str)]\n",
+    "        outputs = outputs.strip()\n",
+    "        return outputs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 54,
+   "id": "4efdbad4-d88a-4477-a3a0-f5591cd0b172",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.\n",
+      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
+      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "load llaVA-Phi MLLM!!!\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "492c17cf54f34d4d9e4f288fc9e72e79",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'device_map': 'cuda'}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
+     ]
+    }
+   ],
+   "source": [
+    "multimodal_phi2 = MultiModalPhi2()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 57,
+   "id": "9a6de0b0-a231-4d50-88e8-e40c6f7216c3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "text = \"tell me about the audio\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 58,
+   "id": "b4919948-6a75-4d19-ba95-9ba233a7d3d9",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'is a popular Japanese drama series featuring a man in a red and white costume, who is dressed as Santa Claus, is walking down the street. The scene takes place in a busy city environment, with people walking and standing on the sidewalk, likely enjoying the festive atmosphere and the festive atmosphere.'"
+      ]
+     },
+     "execution_count": 58,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "multimodal_phi2(text, audio, image)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "590f2d64-62ed-4e6f-b7c8-b0cf68aecaab",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 64,
+   "id": "c921eb63-feb5-4fa9-993b-2faeb6dfe1db",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, BitsAndBytesConfig, CLIPImageProcessor"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 65,
+   "id": "b470a2c4-806a-435d-9fc2-f17448dbe5fc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from llava_phi.model import LlavaPhiConfig"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 66,
+   "id": "4f7bc91a-0a41-45e5-92a4-daa1e3eea0da",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.\n",
+      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
+      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "993bc3a38cb84de4a2e3a79a3448c4d6",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "device_map = \"cuda:0\"\n",
+    "load_8bit = False\n",
+    "load_4bit = False\n",
+    "kwargs = {\"device_map\": device_map}\n",
+    "if load_8bit:\n",
+    "    kwargs['load_in_8bit'] = True\n",
+    "elif load_4bit:\n",
+    "    kwargs['load_in_4bit'] = True\n",
+    "    kwargs['quantization_config'] = BitsAndBytesConfig(\n",
+    "        load_in_4bit=True,\n",
+    "        bnb_4bit_compute_dtype=torch.float16,\n",
+    "        bnb_4bit_use_double_quant=True,\n",
+    "        bnb_4bit_quant_type='nf4'\n",
+    "    )\n",
+    "config = LlavaPhiConfig.from_pretrained(model_path, trust_remote_code=True)\n",
+    "tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)\n",
+    "model = LlavaPhiForCausalLM.from_pretrained(\n",
+    "    model_path, \n",
+    "    config=config, \n",
+    "    use_safetensors=True, \n",
+    "    **kwargs).to(\"cuda\")\n",
+    "image_processor = CLIPImageProcessor.from_pretrained(model_path)\n",
+    "mm_use_im_start_end = getattr(model.config, \"mm_use_im_start_end\", False)\n",
+    "mm_use_im_patch_token = getattr(model.config, \"mm_use_im_patch_token\", True)\n",
+    "\n",
+    "# TODO: the tokenizer length of phi-2 is 50295, but the output class of lm_head is 51200\n",
+    "if mm_use_im_patch_token:\n",
+    "    tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)\n",
+    "if mm_use_im_start_end:\n",
+    "    tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True)\n",
+    "    \n",
+    "if hasattr(model.config, \"max_sequence_length\"):\n",
+    "        context_len = model.config.max_sequence_length\n",
+    "else:\n",
+    "    context_len = 2048"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 70,
+   "id": "99355837-a297-4a25-aeb3-1670af7e9251",
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "KeyboardInterrupt",
+     "evalue": "",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[70], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msave_pretrained\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mLlava-Phi-Checkpoint\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m~/miniconda3/envs/torchenv/lib/python3.10/site-packages/transformers/modeling_utils.py:2376\u001b[0m, in \u001b[0;36mPreTrainedModel.save_pretrained\u001b[0;34m(self, save_directory, is_main_process, state_dict, save_function, push_to_hub, max_shard_size, safe_serialization, variant, token, save_peft_format, **kwargs)\u001b[0m\n\u001b[1;32m   2372\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m shard_file, shard \u001b[38;5;129;01min\u001b[39;00m shards\u001b[38;5;241m.\u001b[39mitems():\n\u001b[1;32m   2373\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m safe_serialization:\n\u001b[1;32m   2374\u001b[0m         \u001b[38;5;66;03m# At some point we will need to deal better with save_function (used for TPU and other distributed\u001b[39;00m\n\u001b[1;32m   2375\u001b[0m         \u001b[38;5;66;03m# joyfulness), but for now this enough.\u001b[39;00m\n\u001b[0;32m-> 2376\u001b[0m         \u001b[43msafe_save_file\u001b[49m\u001b[43m(\u001b[49m\u001b[43mshard\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mos\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpath\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mjoin\u001b[49m\u001b[43m(\u001b[49m\u001b[43msave_directory\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mshard_file\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmetadata\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m{\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mformat\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mpt\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m}\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2377\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   2378\u001b[0m         save_function(shard, os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mjoin(save_directory, shard_file))\n",
+      "File \u001b[0;32m~/miniconda3/envs/torchenv/lib/python3.10/site-packages/safetensors/torch.py:281\u001b[0m, in \u001b[0;36msave_file\u001b[0;34m(tensors, filename, metadata)\u001b[0m\n\u001b[1;32m    250\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21msave_file\u001b[39m(\n\u001b[1;32m    251\u001b[0m     tensors: Dict[\u001b[38;5;28mstr\u001b[39m, torch\u001b[38;5;241m.\u001b[39mTensor],\n\u001b[1;32m    252\u001b[0m     filename: Union[\u001b[38;5;28mstr\u001b[39m, os\u001b[38;5;241m.\u001b[39mPathLike],\n\u001b[1;32m    253\u001b[0m     metadata: Optional[Dict[\u001b[38;5;28mstr\u001b[39m, \u001b[38;5;28mstr\u001b[39m]] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    254\u001b[0m ):\n\u001b[1;32m    255\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    256\u001b[0m \u001b[38;5;124;03m    Saves a dictionary of tensors into raw bytes in safetensors format.\u001b[39;00m\n\u001b[1;32m    257\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    279\u001b[0m \u001b[38;5;124;03m    ```\u001b[39;00m\n\u001b[1;32m    280\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 281\u001b[0m     \u001b[43mserialize_file\u001b[49m\u001b[43m(\u001b[49m\u001b[43m_flatten\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtensors\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfilename\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmetadata\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmetadata\u001b[49m\u001b[43m)\u001b[49m\n",
+      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
+     ]
+    }
+   ],
+   "source": [
+    "model.save_pretrained(\"Llava-Phi-Checkpoint\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fa0bec34-a148-4340-a30c-6f09dd5e71ca",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.push_to_hub(\"RaviNaik/Llava-Phi2\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 73,
+   "id": "382f74b0-2967-408a-badc-a90918810d74",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "CommitInfo(commit_url='https://huggingface.co/RaviNaik/Llava-Phi2/commit/fa8f7240058241243f6bdc3d6ab44bb691f76e39', commit_message='Upload tokenizer', commit_description='', oid='fa8f7240058241243f6bdc3d6ab44bb691f76e39', pr_url=None, pr_revision=None, pr_num=None)"
+      ]
+     },
+     "execution_count": 73,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "tokenizer.push_to_hub(\"RaviNaik/Llava-Phi2\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b851459b-d3ac-4fb8-99b6-17a648adc41f",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

Experiments/instruct_data.py

@@ -0,0 +1,39 @@
+from datasets import Dataset, IterableDataset
+from PIL import Image
+
+# ChatML format
+templates = {
+    "assistant": "<|im_start|>assistant\n{msg}<|im_end|>",      # message by assistant
+    "user": "<|im_start|>user\n{msg}<|im_end|>"       # message by user
+}
+
+ds = Dataset.from_json("llava_instruct_150k.json", split="train")
+ds_stream = ds.to_iterable_dataset()
+
+
+def get_image(image_path):
+    image_path = f"train2014/COCO_train2014_{image_path}"
+    img = Image.open(image_path)
+    return img
+    
+def get_chatml_text(conversations):
+    chatml_text = ""
+    for conversation in conversations:
+        role = conversation["from"]
+        role = "user" if role == "human" else "assistant"
+        content = conversation["value"]
+
+        formatted_text = templates[role].format(msg=content)
+        chatml_text += formatted_text + "\n"
+    return chatml_text
+
+def instruct_data_generator():
+    for sample in ds_stream:
+        image_path = sample["image"]
+        conversations = sample["conversations"]
+        
+        image = get_image(image_path)
+        text = get_chatml_text(conversations)
+        yield {"text": text, "image": image}
+
+instruct_ds = IterableDataset.from_generator(generator=instruct_data_generator)

Experiments/llava_exp.ipynb

@@ -0,0 +1,145 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "99576983-f881-47c8-8b5e-c6f561a93e71",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import transformers"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "58ba19f2-4b91-4f90-a33d-4c1ed17e202a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from transformers import LlavaForConditionalGeneration, LlavaConfig, CLIPVisionConfig, PhiConfig\n",
+    "\n",
+    "# Initializing a CLIP-vision config\n",
+    "vision_config = CLIPVisionConfig()\n",
+    "\n",
+    "# Initializing a Llama config\n",
+    "text_config = PhiConfig()\n",
+    "\n",
+    "# Initializing a Llava llava-1.5-7b style configuration\n",
+    "configuration = LlavaConfig(vision_config, text_config)\n",
+    "\n",
+    "# Initializing a model from the llava-1.5-7b style configuration\n",
+    "model = LlavaForConditionalGeneration(configuration)\n",
+    "\n",
+    "# Accessing the model configuration\n",
+    "configuration = model.config"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "a806a07a-fe72-45a3-8ceb-8e942c6c845d",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "LlavaConfig {\n",
+       "  \"ignore_index\": -100,\n",
+       "  \"image_token_index\": 32000,\n",
+       "  \"model_type\": \"llava\",\n",
+       "  \"projector_hidden_act\": \"gelu\",\n",
+       "  \"text_config\": {\n",
+       "    \"embd_pdrop\": 0.0,\n",
+       "    \"hidden_act\": \"gelu_new\",\n",
+       "    \"hidden_size\": 2048,\n",
+       "    \"intermediate_size\": 8192,\n",
+       "    \"layer_norm_eps\": 1e-05,\n",
+       "    \"model_type\": \"phi\",\n",
+       "    \"num_hidden_layers\": 24,\n",
+       "    \"partial_rotary_factor\": 0.5,\n",
+       "    \"qk_layernorm\": false,\n",
+       "    \"resid_pdrop\": 0.0,\n",
+       "    \"vocab_size\": 51200\n",
+       "  },\n",
+       "  \"transformers_version\": \"4.36.2\",\n",
+       "  \"vision_config\": {\n",
+       "    \"hidden_size\": 768,\n",
+       "    \"image_size\": 224,\n",
+       "    \"intermediate_size\": 3072,\n",
+       "    \"model_type\": \"clip_vision_model\",\n",
+       "    \"num_attention_heads\": 12,\n",
+       "    \"num_hidden_layers\": 12,\n",
+       "    \"patch_size\": 32,\n",
+       "    \"projection_dim\": 512\n",
+       "  },\n",
+       "  \"vision_feature_layer\": -2,\n",
+       "  \"vision_feature_select_strategy\": \"default\",\n",
+       "  \"vocab_size\": 32000\n",
+       "}"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model.config"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "79efbc6b-f005-4a5c-82a1-112fa37f1904",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Cloning into 'llava-phi'...\n",
+      "remote: Enumerating objects: 151, done.\u001b[K\n",
+      "remote: Counting objects: 100% (151/151), done.\u001b[K\n",
+      "remote: Compressing objects: 100% (116/116), done.\u001b[K\n",
+      "remote: Total 151 (delta 36), reused 133 (delta 25), pack-reused 0\u001b[K\n",
+      "Receiving objects: 100% (151/151), 333.89 KiB | 112.00 KiB/s, done.\n",
+      "Resolving deltas: 100% (36/36), done.\n"
+     ]
+    }
+   ],
+   "source": [
+    "!git clone https://github.com/zhuyiche/llava-phi.git"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cf827184-f334-4d86-ace1-fe9c92f84d66",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

Experiments/multimodal_exp.ipynb

@@ -0,0 +1,362 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "id": "d4bed9ef-4bff-4d61-a4f9-a585f377f136",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from PIL import Image\n",
+    "import requests\n",
+    "\n",
+    "import torch\n",
+    "from torch import nn\n",
+    "from transformers import AutoProcessor, CLIPVisionModel, CLIPVisionConfig, CLIPPreTrainedModel\n",
+    "from transformers.models.clip.modeling_clip import CLIPVisionModelOutput, CLIPVisionTransformer\n",
+    "from transformers import WhisperProcessor, WhisperForConditionalGeneration\n",
+    "from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, AutoTokenizer\n",
+    "from typing import Optional, Union, Tuple"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 43,
+   "id": "952314f0-ee9d-45e7-85b8-1e3e44c1a2fd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class VisionLanguageConnector(nn.Module):\n",
+    "    def __init__(self, hidden_size, projection_dim):\n",
+    "        super().__init__()\n",
+    "        self.mlp = nn.Sequential(\n",
+    "            nn.Linear(hidden_size, hidden_size, bias=False),\n",
+    "            nn.GELU(),\n",
+    "            nn.Linear(hidden_size, projection_dim, bias=False)\n",
+    "        )\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        return self.mlp(x)\n",
+    "        \n",
+    "class ClipWithProjection():\n",
+    "    config_class = CLIPVisionConfig\n",
+    "    main_input_name = \"pixel_values\"\n",
+    "\n",
+    "    def __init__(self, hidden_size, projection_dim):\n",
+    "        super().__init__()\n",
+    "        \n",
+    "        self.processor = AutoProcessor.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
+    "        self.vision_model = CLIPVisionModel.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
+    "        self.vision_language_connector = VisionLanguageConnector(hidden_size, projection_dim)\n",
+    "\n",
+    "    def forward(\n",
+    "        self,\n",
+    "        image = None,\n",
+    "        output_attentions: Optional[bool] = None,\n",
+    "        output_hidden_states: Optional[bool] = None,\n",
+    "        return_dict: Optional[bool] = None,\n",
+    "    ) -> Union[Tuple, CLIPVisionModelOutput]:\n",
+    "        \n",
+    "        pixel_values = self.processor(images=image, return_tensors=\"pt\")[\"pixel_values\"]\n",
+    "        vision_outputs = self.vision_model(\n",
+    "            pixel_values=pixel_values,\n",
+    "            output_attentions=output_attentions,\n",
+    "            output_hidden_states=output_hidden_states,\n",
+    "            return_dict=return_dict,\n",
+    "        )\n",
+    "\n",
+    "        pooled_output = vision_outputs[1]  # pooled_output\n",
+    "\n",
+    "        image_embeds = self.vision_language_connector(pooled_output)\n",
+    "\n",
+    "        return CLIPVisionModelOutput(\n",
+    "            image_embeds=image_embeds,\n",
+    "            last_hidden_state=vision_outputs.last_hidden_state,\n",
+    "            hidden_states=vision_outputs.hidden_states,\n",
+    "            attentions=vision_outputs.attentions,\n",
+    "        )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 44,
+   "id": "bd2889fe-be85-44a3-afe8-65b47f7a93c3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "url = \"http://images.cocodataset.org/val2017/000000039769.jpg\"\n",
+    "image = Image.open(requests.get(url, stream=True).raw)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 46,
+   "id": "17c72699-fe98-4b96-b63c-5c8ab7c1a65f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# model = ClipWithProjection(768, 512)\n",
+    "# model.forward(image)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 47,
+   "id": "70806156-38a9-45a2-bf9f-e72047a0173f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class AudioLanguageConnector:\n",
+    "    def __init__(self, projection_dim):\n",
+    "        model_name = \"microsoft/phi-2\"\n",
+    "        self.phi2_tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n",
+    "        self.phi2_tokenizer.pad_token = self.phi2_tokenizer.eos_token\n",
+    "        self.phi2_tokenizer.max_length = projection_dim\n",
+    "\n",
+    "    def __call__(self, text):\n",
+    "        text = f\"<audio_start> {text} <audio_end>\"\n",
+    "        tokens = self.phi2_tokenizer(text, return_tensors=\"pt\", return_attention_mask=False)\n",
+    "        return tokens\n",
+    "        \n",
+    "\n",
+    "class WhisperWithProjection:\n",
+    "    def __init__(self, projection_dim):\n",
+    "        self.processor = WhisperProcessor.from_pretrained(\"openai/whisper-tiny\")\n",
+    "        self.model = WhisperForConditionalGeneration.from_pretrained(\"openai/whisper-tiny\")\n",
+    "        self.model.config.forced_decoder_ids = None\n",
+    "        self.audio_language_connector = AudioLanguageConnector(projection_dim)\n",
+    "        \n",
+    "    def forward(self, audio):\n",
+    "        input_features = self.processor(audio[\"array\"],\n",
+    "                                   sampling_rate=audio[\"sampling_rate\"],\n",
+    "                                   return_tensors=\"pt\").input_features\n",
+    "        # generate token ids\n",
+    "        predicted_ids = self.model.generate(input_features)\n",
+    "        # decode token ids to text        \n",
+    "        transcription = self.processor.batch_decode(predicted_ids, skip_special_tokens=True)\n",
+    "\n",
+    "        audio_embeddings = self.audio_language_connector(transcription)\n",
+    "        return audio_embeddings"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 48,
+   "id": "79cc4d98-498b-4042-bd71-143b2477733d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class TextModality:\n",
+    "    def __init__(self, projection_dim):\n",
+    "        model_name = \"microsoft/phi-2\"\n",
+    "        self.phi2_tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n",
+    "        self.phi2_tokenizer.pad_token = self.phi2_tokenizer.eos_token\n",
+    "        self.phi2_tokenizer.max_length = projection_dim\n",
+    "\n",
+    "\n",
+    "    def __call__(self, text):\n",
+    "        tokens = self.phi2_tokenizer(text, return_tensors=\"pt\", return_attention_mask=False)\n",
+    "        return tokens"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 77,
+   "id": "ba4c4772-923f-48e8-a4af-b7d9c192dd4b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class MultiModalPhi2:\n",
+    "    def __init__(self):\n",
+    "        self.text_modality = TextModality(projection_dim=768)\n",
+    "        self.whisper_w_proj = WhisperWithProjection(projection_dim=512)\n",
+    "        self.clip_w_proj = ClipWithProjection(hidden_size=768, projection_dim=768)\n",
+    "        self.llm = self.load_llm()\n",
+    "\n",
+    "    def load_llm(self):\n",
+    "        model_name = \"microsoft/phi-2\"\n",
+    "        \n",
+    "        bnb_config = BitsAndBytesConfig(\n",
+    "        load_in_4bit=True,\n",
+    "        bnb_4bit_quant_type=\"nf4\",\n",
+    "        bnb_4bit_compute_dtype=torch.float16)\n",
+    "    \n",
+    "        model = AutoModelForCausalLM.from_pretrained(\n",
+    "            model_name,\n",
+    "            quantization_config=bnb_config,\n",
+    "            trust_remote_code=True,\n",
+    "            device_map=\"cuda:0\"\n",
+    "        )\n",
+    "        model.config.use_cache = False\n",
+    "        return model\n",
+    "\n",
+    "    def forward(self, audio, image, text):\n",
+    "        if text is not None:\n",
+    "            text_embed = self.text_modality(text)[\"input_ids\"]\n",
+    "        if audio is not None:\n",
+    "            audio_embed = self.whisper_w_proj.forward(audio)[\"input_ids\"]\n",
+    "        if image is not None:\n",
+    "            image_embed = self.clip_w_proj.forward(image)[0]\n",
+    "        print(text_embed.shape, text_embed.dtype)\n",
+    "        print(audio_embed.shape, audio_embed.dtype)\n",
+    "        print(image_embed.shape, image_embed.dtype)\n",
+    "        \n",
+    "        inputs = torch.concat([text_embed, audio_embed, image_embed], dim=1)\n",
+    "        print(inputs.shape, inputs.dtype)\n",
+    "        outputs = self.llm(inputs)\n",
+    "\n",
+    "        return outputs \n",
+    "        \n",
+    "\n",
+    "    def generate(self, audio, text):\n",
+    "        text_embeddings = self.text_modality(text)\n",
+    "        audio_embeddings = self.whisper_w_proj.forward(audio)\n",
+    "        inputs = torch.concat([text_embed[\"input_ids\"], audio_embed[\"input_ids\"]], dim=1)\n",
+    "        \n",
+    "        outputs = self.llm.generate(inputs, max_length=200)\n",
+    "        text = self.text_modality.phi2_tokenizer.batch_decode(outputs)[0]\n",
+    "        print(text)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 74,
+   "id": "7ca694eb-8009-4eb9-9a4c-eac406ab9584",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datasets import load_dataset\n",
+    "audio_ds = load_dataset(\"hf-internal-testing/librispeech_asr_dummy\", \"clean\", split=\"validation\")\n",
+    "audio = audio_ds[0][\"audio\"]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 58,
+   "id": "37be28c5-4cc3-4471-b394-032c7602accc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "text = \"explain about the audio\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 59,
+   "id": "c0705114-1670-4937-bc3e-3660e5a5d2c5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 78,
+   "id": "0d7e5b49-b4bd-477c-87b8-91ef70857677",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
+      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "733dc7b2208b4853a89aea49bff9a55c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "model = MultiModalPhi2()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 79,
+   "id": "0b6471c4-4553-47f3-b38f-46057dcf80f2",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([1, 5]) torch.int64\n",
+      "torch.Size([1, 33]) torch.int64\n",
+      "torch.Size([1, 768]) torch.float32\n",
+      "torch.Size([1, 806]) torch.float32\n"
+     ]
+    },
+    {
+     "ename": "RuntimeError",
+     "evalue": "Expected tensor for argument #1 'indices' to have one of the following scalar types: Long, Int; but got torch.cuda.FloatTensor instead (while checking arguments for embedding)",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[79], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mforward\u001b[49m\u001b[43m(\u001b[49m\u001b[43maudio\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mimage\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtext\u001b[49m\u001b[43m)\u001b[49m\n",
+      "Cell \u001b[0;32mIn[77], line 38\u001b[0m, in \u001b[0;36mMultiModalPhi2.forward\u001b[0;34m(self, audio, image, text)\u001b[0m\n\u001b[1;32m     36\u001b[0m inputs \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mconcat([text_embed, audio_embed, image_embed], dim\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m)\n\u001b[1;32m     37\u001b[0m \u001b[38;5;28mprint\u001b[39m(inputs\u001b[38;5;241m.\u001b[39mshape, inputs\u001b[38;5;241m.\u001b[39mdtype)\n\u001b[0;32m---> 38\u001b[0m outputs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mllm\u001b[49m\u001b[43m(\u001b[49m\u001b[43minputs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     40\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m outputs\n",
+      "File \u001b[0;32m~/miniconda3/envs/torchenv/lib/python3.10/site-packages/torch/nn/modules/module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1516\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1517\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1518\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m~/miniconda3/envs/torchenv/lib/python3.10/site-packages/torch/nn/modules/module.py:1527\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1522\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1523\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1524\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1525\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1526\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1529\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1530\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
+      "File \u001b[0;32m~/miniconda3/envs/torchenv/lib/python3.10/site-packages/accelerate/hooks.py:165\u001b[0m, in \u001b[0;36madd_hook_to_module.<locals>.new_forward\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    163\u001b[0m         output \u001b[38;5;241m=\u001b[39m old_forward(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m    164\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 165\u001b[0m     output \u001b[38;5;241m=\u001b[39m \u001b[43mold_forward\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    166\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m module\u001b[38;5;241m.\u001b[39m_hf_hook\u001b[38;5;241m.\u001b[39mpost_forward(module, output)\n",
+      "File \u001b[0;32m~/.cache/huggingface/modules/transformers_modules/microsoft/phi-2/85d00b03fee509307549d823fdd095473ba5197c/modeling_phi.py:1049\u001b[0m, in \u001b[0;36mPhiForCausalLM.forward\u001b[0;34m(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict)\u001b[0m\n\u001b[1;32m   1046\u001b[0m return_dict \u001b[38;5;241m=\u001b[39m return_dict \u001b[38;5;28;01mif\u001b[39;00m return_dict \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mconfig\u001b[38;5;241m.\u001b[39muse_return_dict\n\u001b[1;32m   1048\u001b[0m \u001b[38;5;66;03m# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)\u001b[39;00m\n\u001b[0;32m-> 1049\u001b[0m outputs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmodel\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1050\u001b[0m \u001b[43m    \u001b[49m\u001b[43minput_ids\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minput_ids\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1051\u001b[0m \u001b[43m    \u001b[49m\u001b[43mattention_mask\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mattention_mask\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1052\u001b[0m \u001b[43m    \u001b[49m\u001b[43mposition_ids\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mposition_ids\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1053\u001b[0m \u001b[43m    \u001b[49m\u001b[43mpast_key_values\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mpast_key_values\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1054\u001b[0m \u001b[43m    \u001b[49m\u001b[43minputs_embeds\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minputs_embeds\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1055\u001b[0m \u001b[43m    \u001b[49m\u001b[43muse_cache\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43muse_cache\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1056\u001b[0m \u001b[43m    \u001b[49m\u001b[43moutput_attentions\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43moutput_attentions\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1057\u001b[0m \u001b[43m    \u001b[49m\u001b[43moutput_hidden_states\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43moutput_hidden_states\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1058\u001b[0m \u001b[43m    \u001b[49m\u001b[43mreturn_dict\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreturn_dict\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1059\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1061\u001b[0m hidden_states \u001b[38;5;241m=\u001b[39m outputs[\u001b[38;5;241m0\u001b[39m]\n\u001b[1;32m   1062\u001b[0m logits \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlm_head(hidden_states)\n",
+      "File \u001b[0;32m~/miniconda3/envs/torchenv/lib/python3.10/site-packages/torch/nn/modules/module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1516\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1517\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1518\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m~/miniconda3/envs/torchenv/lib/python3.10/site-packages/torch/nn/modules/module.py:1527\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1522\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1523\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1524\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1525\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1526\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1529\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1530\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
+      "File \u001b[0;32m~/miniconda3/envs/torchenv/lib/python3.10/site-packages/accelerate/hooks.py:165\u001b[0m, in \u001b[0;36madd_hook_to_module.<locals>.new_forward\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    163\u001b[0m         output \u001b[38;5;241m=\u001b[39m old_forward(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m    164\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 165\u001b[0m     output \u001b[38;5;241m=\u001b[39m \u001b[43mold_forward\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    166\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m module\u001b[38;5;241m.\u001b[39m_hf_hook\u001b[38;5;241m.\u001b[39mpost_forward(module, output)\n",
+      "File \u001b[0;32m~/.cache/huggingface/modules/transformers_modules/microsoft/phi-2/85d00b03fee509307549d823fdd095473ba5197c/modeling_phi.py:893\u001b[0m, in \u001b[0;36mPhiModel.forward\u001b[0;34m(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict)\u001b[0m\n\u001b[1;32m    890\u001b[0m     position_ids \u001b[38;5;241m=\u001b[39m position_ids\u001b[38;5;241m.\u001b[39munsqueeze(\u001b[38;5;241m0\u001b[39m)\n\u001b[1;32m    892\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m inputs_embeds \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m--> 893\u001b[0m     inputs_embeds \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43membed_tokens\u001b[49m\u001b[43m(\u001b[49m\u001b[43minput_ids\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    895\u001b[0m inputs_embeds \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39membed_dropout(inputs_embeds)\n\u001b[1;32m    897\u001b[0m \u001b[38;5;66;03m# Attention mask.\u001b[39;00m\n",
+      "File \u001b[0;32m~/miniconda3/envs/torchenv/lib/python3.10/site-packages/torch/nn/modules/module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1516\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1517\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1518\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m~/miniconda3/envs/torchenv/lib/python3.10/site-packages/torch/nn/modules/module.py:1527\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1522\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1523\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1524\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1525\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1526\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1529\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1530\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
+      "File \u001b[0;32m~/miniconda3/envs/torchenv/lib/python3.10/site-packages/accelerate/hooks.py:165\u001b[0m, in \u001b[0;36madd_hook_to_module.<locals>.new_forward\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    163\u001b[0m         output \u001b[38;5;241m=\u001b[39m old_forward(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m    164\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 165\u001b[0m     output \u001b[38;5;241m=\u001b[39m \u001b[43mold_forward\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    166\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m module\u001b[38;5;241m.\u001b[39m_hf_hook\u001b[38;5;241m.\u001b[39mpost_forward(module, output)\n",
+      "File \u001b[0;32m~/miniconda3/envs/torchenv/lib/python3.10/site-packages/torch/nn/modules/sparse.py:162\u001b[0m, in \u001b[0;36mEmbedding.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    161\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m: Tensor) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tensor:\n\u001b[0;32m--> 162\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mF\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43membedding\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    163\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mweight\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpadding_idx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmax_norm\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    164\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnorm_type\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mscale_grad_by_freq\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msparse\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m~/miniconda3/envs/torchenv/lib/python3.10/site-packages/torch/nn/functional.py:2233\u001b[0m, in \u001b[0;36membedding\u001b[0;34m(input, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse)\u001b[0m\n\u001b[1;32m   2227\u001b[0m     \u001b[38;5;66;03m# Note [embedding_renorm set_grad_enabled]\u001b[39;00m\n\u001b[1;32m   2228\u001b[0m     \u001b[38;5;66;03m# XXX: equivalent to\u001b[39;00m\n\u001b[1;32m   2229\u001b[0m     \u001b[38;5;66;03m# with torch.no_grad():\u001b[39;00m\n\u001b[1;32m   2230\u001b[0m     \u001b[38;5;66;03m#   torch.embedding_renorm_\u001b[39;00m\n\u001b[1;32m   2231\u001b[0m     \u001b[38;5;66;03m# remove once script supports set_grad_enabled\u001b[39;00m\n\u001b[1;32m   2232\u001b[0m     _no_grad_embedding_renorm_(weight, \u001b[38;5;28minput\u001b[39m, max_norm, norm_type)\n\u001b[0;32m-> 2233\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43membedding\u001b[49m\u001b[43m(\u001b[49m\u001b[43mweight\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mpadding_idx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mscale_grad_by_freq\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msparse\u001b[49m\u001b[43m)\u001b[49m\n",
+      "\u001b[0;31mRuntimeError\u001b[0m: Expected tensor for argument #1 'indices' to have one of the following scalar types: Long, Int; but got torch.cuda.FloatTensor instead (while checking arguments for embedding)"
+     ]
+    }
+   ],
+   "source": [
+    "model.forward(audio, image, text)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4ca96caf-82e2-4f07-87b3-8654dfdc89aa",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

Experiments/pretrain_data_check.ipynb

@@ -0,0 +1,304 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "61c272f2-edbe-4b7d-8fec-3ab431400cd3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "e9dfd7d7-1685-4fc7-bbb9-3905c32d8ba1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with open(\"metadata.json\", \"rb\") as f:\n",
+    "    metadata = json.load(f)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "70bdba48-db01-42ac-8d89-edc69d7d7672",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "595375"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "len(metadata)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "59e193cc-0dd8-4f7e-959a-fbad0133d76c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with open(\"blip_laion_cc_sbu_558k.jsonblip_laion_cc_sbu_558k.json\", \"rb\") as f:\n",
+    "    data = json.load(f)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "f3157f41-269b-4f7a-b3ba-9be711babe02",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'id': '004539375',\n",
+       " 'image': '00453/004539375.jpg',\n",
+       " 'conversations': [{'from': 'human',\n",
+       "   'value': 'Render a clear and concise summary of the photo.\\n<image>'},\n",
+       "  {'from': 'gpt',\n",
+       "   'value': 'select luxury furniture 3 - inch gel memory foam mattress topper'}]}"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "data[0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "50d8a051-1526-47dd-ad71-d3c66f7bd34e",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'id': '004374662',\n",
+       " 'image': '00437/004374662.jpg',\n",
+       " 'conversations': [{'from': 'human',\n",
+       "   'value': 'Give a brief description of the image.\\n<image>'},\n",
+       "  {'from': 'gpt', 'value': 'the north face duffel bag camo large'}]}"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "data[234]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "2e6d5664-4583-49a6-93cc-079ee2d1ff6c",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "558128"
+      ]
+     },
+     "execution_count": 17,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "len(data)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "11ed106d-6bef-482c-a456-5eaaf2025534",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'id': 'GCC_train_001749371',\n",
+       " 'image': 'GCC_train_001749371.jpg',\n",
+       " 'caption': 'if you are dreaming of simpler or off - the - grid living , a yurt is a fantastic option',\n",
+       " 'blip_caption': 'a white and tan yurt sitting on a dirt road',\n",
+       " 'url': 'https://i.pinimg.com/736x/14/7b/64/147b64467ee966d9a578097bb70475ad--yurt-kits-small-space-living.jpg'}"
+      ]
+     },
+     "execution_count": 10,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "metadata[67]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "ce8adcec-2499-4be3-be1d-7313fe54e96a",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'id': '000466761',\n",
+       " 'image': '00046/000466761.jpg',\n",
+       " 'conversations': [{'from': 'human',\n",
+       "   'value': '<image>\\nProvide a brief description of the given image.'},\n",
+       "  {'from': 'gpt',\n",
+       "   'value': 'a clipboard and a pen with the words public health emergency next to it on a white table'}]}"
+      ]
+     },
+     "execution_count": 15,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "data[67]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "068313b6-6379-4ca2-892c-682634d3581e",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "list"
+      ]
+     },
+     "execution_count": 16,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "type(data)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "id": "9ec33b51-4a0b-4a1e-81f7-2fda7cddb25f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sample_data = data[:200000]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "id": "095685e5-40f1-4d84-8280-ef74fa56c5a2",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "200000"
+      ]
+     },
+     "execution_count": 25,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "len(sample_data)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "id": "ffbad552-23fd-475f-8e9a-7118bcc4f51e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with open(\"llava-phi/pretrain_data/blip_sample.json\", \"w\") as f:\n",
+    "    json.dump(sample_data, f)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "id": "69a05d25-6f3b-40c0-a3b5-e185ff526471",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with open(\"llava-phi/pretrain_data/blip_sample.json\", \"rb\") as f:\n",
+    "    sample = json.load(f)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 28,
+   "id": "200eea06-dfd6-4b3a-bb91-82af7d363951",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "200000"
+      ]
+     },
+     "execution_count": 28,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "len(sample)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f86caa1e-edea-4a9c-934f-5420ede80d0d",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

Experiments/whispher_exp.ipynb

@@ -0,0 +1,500 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "bb4dd66b-0c17-48d4-9d34-f48cece2feb5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# !pip install soundfile\n",
+    "# !pip install librosa"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "6e9386ea-4862-4f5b-a02f-d656e1a5ab9e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from transformers import WhisperProcessor, WhisperForConditionalGeneration\n",
+    "from datasets import load_dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "914ab2b4-389d-4c48-8d1d-1250356646ac",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# load model and processor\n",
+    "processor = WhisperProcessor.from_pretrained(\"openai/whisper-tiny\")\n",
+    "model = WhisperForConditionalGeneration.from_pretrained(\"openai/whisper-tiny\")\n",
+    "model.config.forced_decoder_ids = None\n",
+    "\n",
+    "# load dummy dataset and read audio files\n",
+    "ds = load_dataset(\"hf-internal-testing/librispeech_asr_dummy\", \"clean\", split=\"validation\")\n",
+    "sample = ds[0][\"audio\"]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "2b299bab-1228-48d9-a8a5-3d5b6c52162d",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'path': '/home/ravi.naik/.cache/huggingface/datasets/downloads/extracted/431c2c946d216530b2666a0e7ffa5ac3f5b3da89dd28858a9de6c78fae7caa4a/dev_clean/1272/128104/1272-128104-0000.flac',\n",
+       " 'array': array([0.00238037, 0.0020752 , 0.00198364, ..., 0.00042725, 0.00057983,\n",
+       "        0.0010376 ]),\n",
+       " 'sampling_rate': 16000}"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "sample"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "b7e570a1-cf5c-450c-a7b6-49b45a10d2df",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "input_features = processor(sample[\"array\"], sampling_rate=sample[\"sampling_rate\"], return_tensors=\"pt\").input_features "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "584e920b-a7fd-402d-95dd-3b9128cd34bb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# generate token ids\n",
+    "predicted_ids = model.generate(input_features)\n",
+    "# decode token ids to text\n",
+    "transcription = processor.batch_decode(predicted_ids, skip_special_tokens=False)\n",
+    "\n",
+    "transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "b27ab660-861b-49d1-81f9-f51cb7f9d8d8",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[' Mr. Quilter is the apostle of the middle classes and we are glad to welcome his gospel.']"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "transcription"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "eca553b8-68f6-493d-b567-3d526b49ae1b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from torch import nn"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "c619a4cf-9068-4e4d-8139-e16d15345f4f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, AutoTokenizer"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "47d5b1ff-ab0f-4d11-af64-d2fa2be39286",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
+     ]
+    }
+   ],
+   "source": [
+    "model_name = \"microsoft/phi-2\"\n",
+    "phi2_tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n",
+    "phi2_tokenizer.pad_token = phi2_tokenizer.eos_token"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "0b36b3f0-db5b-4029-9072-0a53bcab315a",
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "NameError",
+     "evalue": "name 'transcription' is not defined",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[6], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m tokens \u001b[38;5;241m=\u001b[39m phi2_tokenizer(\u001b[38;5;241m*\u001b[39m\u001b[43mtranscription\u001b[49m, return_tensors\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpt\u001b[39m\u001b[38;5;124m\"\u001b[39m, return_attention_mask\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m)\n",
+      "\u001b[0;31mNameError\u001b[0m: name 'transcription' is not defined"
+     ]
+    }
+   ],
+   "source": [
+    "tokens = phi2_tokenizer(*transcription, return_tensors=\"pt\", return_attention_mask=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "91f6d3d3-bb00-434f-a91e-6952375890d0",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'input_ids': tensor([[ 1770,    13,  2264,   346,   353,   318,   262, 46329,   286,   262,\n",
+       "          3504,  6097,   290,   356,   389,  9675,   284,  7062,   465, 21443,\n",
+       "            13]])}"
+      ]
+     },
+     "execution_count": 22,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "tokens"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "533191d9-4b3b-417a-918d-6fe854f24b50",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "A new version of the following files was downloaded from https://huggingface.co/microsoft/phi-2:\n",
+      "- configuration_phi.py\n",
+      ". Make sure to double-check they do not contain any added malicious code. To avoid downloading new versions of the code file, you can pin a revision.\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2a65a119388b4cb4b123b532176e786e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "modeling_phi.py:   0%|          | 0.00/62.7k [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "A new version of the following files was downloaded from https://huggingface.co/microsoft/phi-2:\n",
+      "- modeling_phi.py\n",
+      ". Make sure to double-check they do not contain any added malicious code. To avoid downloading new versions of the code file, you can pin a revision.\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "7183811844304c16b72d53fe11098a74",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Downloading shards:   0%|          | 0/2 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "3e78fe144e8f42139a4d7a1830dbf192",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "bnb_config = BitsAndBytesConfig(\n",
+    "    load_in_4bit=True,\n",
+    "    bnb_4bit_quant_type=\"nf4\",\n",
+    "    bnb_4bit_compute_dtype=torch.float16,\n",
+    ")\n",
+    "\n",
+    "model = AutoModelForCausalLM.from_pretrained(\n",
+    "    model_name,\n",
+    "    quantization_config=bnb_config,\n",
+    "    trust_remote_code=True,\n",
+    "    device_map=\"cuda:0\"\n",
+    ")\n",
+    "model.config.use_cache = False"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "155c054a-a00f-4ed5-bfff-1ad64889e7f1",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[' Mr. Quilter is the apostle of the middle classes and we are glad to welcome his gospel.\\n']"
+      ]
+     },
+     "execution_count": 19,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "phi2_tokenizer.batch_decode(model.generate(**tokens))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "04f940c9-586d-4937-ae31-cc0f96d33e92",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class AudioLanguageConnector:\n",
+    "    def __init__(self):\n",
+    "        model_name = \"microsoft/phi-2\"\n",
+    "        self.phi2_tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n",
+    "        self.phi2_tokenizer.pad_token = self.phi2_tokenizer.eos_token\n",
+    "\n",
+    "    def __call__(self, text):\n",
+    "        text = f\"<audio_start> {text} <audio_end>\"\n",
+    "        tokens = self.phi2_tokenizer(text, return_tensors=\"pt\", return_attention_mask=False)\n",
+    "        return tokens\n",
+    "        \n",
+    "\n",
+    "class WhisperWithProjection:\n",
+    "    def __init__(self):\n",
+    "        self.processor = WhisperProcessor.from_pretrained(\"openai/whisper-tiny\")\n",
+    "        self.model = WhisperForConditionalGeneration.from_pretrained(\"openai/whisper-tiny\")\n",
+    "        self.model.config.forced_decoder_ids = None\n",
+    "        self.audio_language_connector = AudioLanguageConnector()\n",
+    "        \n",
+    "    def forward(self, audio):\n",
+    "        input_features = self.processor(audio[\"array\"],\n",
+    "                                   sampling_rate=audio[\"sampling_rate\"],\n",
+    "                                   return_tensors=\"pt\").input_features\n",
+    "        # generate token ids\n",
+    "        predicted_ids = self.model.generate(input_features)\n",
+    "        # decode token ids to text        \n",
+    "        transcription = self.processor.batch_decode(predicted_ids, skip_special_tokens=True)\n",
+    "\n",
+    "        audio_embeddings = self.audio_language_connector(transcription)\n",
+    "        return audio_embeddings"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "2b1f8f44-bfe6-413c-9e32-c38fa5517981",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class TextModality:\n",
+    "    def __init__(self):\n",
+    "        model_name = \"microsoft/phi-2\"\n",
+    "        self.phi2_tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n",
+    "        self.phi2_tokenizer.pad_token = self.phi2_tokenizer.eos_token\n",
+    "\n",
+    "    def __call__(self, text):\n",
+    "        tokens = self.phi2_tokenizer(text, return_tensors=\"pt\", return_attention_mask=False)\n",
+    "        return tokens"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "21c51648-abb6-4bbd-b4c1-509967a69337",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class MultiModalPhi2:\n",
+    "    def __init__(self):\n",
+    "        self.text_modality = TextModality()\n",
+    "        self.whisper_w_proj = WhisperWithProjection()\n",
+    "        self.llm = self.load_llm()\n",
+    "\n",
+    "    def load_llm(self):\n",
+    "        bnb_config = BitsAndBytesConfig(\n",
+    "        load_in_4bit=True,\n",
+    "        bnb_4bit_quant_type=\"nf4\",\n",
+    "        bnb_4bit_compute_dtype=torch.float16)\n",
+    "    \n",
+    "        model = AutoModelForCausalLM.from_pretrained(\n",
+    "            model_name,\n",
+    "            quantization_config=bnb_config,\n",
+    "            trust_remote_code=True,\n",
+    "            device_map=\"cuda:0\"\n",
+    "        )\n",
+    "        model.config.use_cache = False\n",
+    "        return model\n",
+    "\n",
+    "    def generate(self, audio, text):\n",
+    "        text_embeddings = self.text_modality(text)\n",
+    "        audio_embeddings = self.whisper_w_proj.forward(audio)\n",
+    "        inputs = torch.concat([text_embeddings[\"input_ids\"], audio_embeddings[\"input_ids\"]], dim=1)\n",
+    "        \n",
+    "        # outputs = self.llm.generate(inputs, max_length=200)\n",
+    "        outputs = self.llm(inputs)\n",
+    "        return outputs\n",
+    "        \n",
+    "        # text = self.text_modality.phi2_tokenizer.batch_decode(outputs)[0]\n",
+    "        # print(text)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "472a00cb-bae9-4c09-a0ef-bc57881b5e2c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
+      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2236e6b1e26d444fa3d48181ba1a6cf9",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "multi_modal_phi = MultiModalPhi2()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "c350f2d3-0929-4c46-b63d-ff92dea437f3",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "CausalLMOutputWithPast(loss={'logits': tensor([[[ 6.9531,  9.9375,  7.0234,  ...,  2.0020,  2.0020,  2.0000],\n",
+       "         [ 8.9062, 12.1172,  7.5977,  ..., -1.2012, -1.2012, -1.2012],\n",
+       "         [ 7.0273,  5.3477,  3.6328,  ..., -4.2070, -4.2070, -4.2070],\n",
+       "         ...,\n",
+       "         [ 7.0234,  7.4414,  9.1016,  ...,  1.0117,  1.0127,  1.0117],\n",
+       "         [ 9.4531, 10.0391,  9.7578,  ...,  0.0776,  0.0775,  0.0764],\n",
+       "         [ 8.0703,  6.6445,  5.5156,  ..., -1.9268, -1.9268, -1.9277]]],\n",
+       "       grad_fn=<ToCopyBackward0>)}, logits=tensor([[[ 6.9531,  9.9375,  7.0234,  ...,  2.0020,  2.0020,  2.0000],\n",
+       "         [ 8.9062, 12.1172,  7.5977,  ..., -1.2012, -1.2012, -1.2012],\n",
+       "         [ 7.0273,  5.3477,  3.6328,  ..., -4.2070, -4.2070, -4.2070],\n",
+       "         ...,\n",
+       "         [ 7.0234,  7.4414,  9.1016,  ...,  1.0117,  1.0127,  1.0117],\n",
+       "         [ 9.4531, 10.0391,  9.7578,  ...,  0.0776,  0.0775,  0.0764],\n",
+       "         [ 8.0703,  6.6445,  5.5156,  ..., -1.9268, -1.9268, -1.9277]]],\n",
+       "       grad_fn=<ToCopyBackward0>), past_key_values=None, hidden_states=None, attentions=None)"
+      ]
+     },
+     "execution_count": 17,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "audio = sample\n",
+    "text = \"explain about the audio\"\n",
+    "multi_modal_phi.generate(audio, text)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "46aa9c66-a5bb-4760-8895-92673f49345f",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

inference/conversation.py

@@ -0,0 +1,224 @@
+import dataclasses
+from enum import auto, Enum
+from typing import List, Tuple
+
+
+class SeparatorStyle(Enum):
+    """Different separator style."""
+    SINGLE = auto()
+    TWO = auto()
+    MPT = auto()
+    PLAIN = auto()
+    LLAMA_2 = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that keeps all conversation history."""
+    system: str
+    roles: List[str]
+    messages: List[List[str]]
+    offset: int
+    sep_style: SeparatorStyle = SeparatorStyle.SINGLE
+    sep: str = "###"
+    sep2: str = None
+    version: str = "Unknown"
+
+    skip_next: bool = False
+
+    def get_prompt(self):
+        messages = self.messages
+        if len(messages) > 0 and type(messages[0][1]) is tuple:
+            messages = self.messages.copy()
+            init_role, init_msg = messages[0].copy()
+            init_msg = init_msg[0].replace("<image>", "").strip()
+            if 'mmtag' in self.version:
+                messages[0] = (init_role, init_msg)
+                messages.insert(0, (self.roles[0], "<Image><image></Image>"))
+                messages.insert(1, (self.roles[1], "Received."))
+            else:
+                messages[0] = (init_role, "<image>\n" + init_msg)
+
+        if self.sep_style == SeparatorStyle.SINGLE:
+            ret = self.system + self.sep
+            for role, message in messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + ": " + message + self.sep
+                else:
+                    ret += role + ":"
+        elif self.sep_style == SeparatorStyle.TWO:
+            seps = [self.sep, self.sep2]
+            ret = self.system + seps[0]
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + ": " + message + seps[i % 2]
+                else:
+                    ret += role + ":"
+        elif self.sep_style == SeparatorStyle.PLAIN:
+            seps = [self.sep, self.sep2]
+            ret = self.system
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += message + seps[i % 2]
+                else:
+                    ret += ""
+        else:
+            raise ValueError(f"Invalid style: {self.sep_style}")
+
+        return ret
+
+    def append_message(self, role, message):
+        self.messages.append([role, message])
+
+    def get_images(self, return_pil=False):
+        images = []
+        for i, (role, msg) in enumerate(self.messages[self.offset:]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    import base64
+                    from io import BytesIO
+                    from PIL import Image
+                    msg, image, image_process_mode = msg
+                    if image_process_mode == "Pad":
+                        def expand2square(pil_img, background_color=(122, 116, 104)):
+                            width, height = pil_img.size
+                            if width == height:
+                                return pil_img
+                            elif width > height:
+                                result = Image.new(pil_img.mode, (width, width), background_color)
+                                result.paste(pil_img, (0, (width - height) // 2))
+                                return result
+                            else:
+                                result = Image.new(pil_img.mode, (height, height), background_color)
+                                result.paste(pil_img, ((height - width) // 2, 0))
+                                return result
+                        image = expand2square(image)
+                    elif image_process_mode in ["Default", "Crop"]:
+                        pass
+                    elif image_process_mode == "Resize":
+                        image = image.resize((336, 336))
+                    else:
+                        raise ValueError(f"Invalid image_process_mode: {image_process_mode}")
+                    max_hw, min_hw = max(image.size), min(image.size)
+                    aspect_ratio = max_hw / min_hw
+                    max_len, min_len = 800, 400
+                    shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))
+                    longest_edge = int(shortest_edge * aspect_ratio)
+                    W, H = image.size
+                    if longest_edge != max(image.size):
+                        if H > W:
+                            H, W = longest_edge, shortest_edge
+                        else:
+                            H, W = shortest_edge, longest_edge
+                        image = image.resize((W, H))
+                    if return_pil:
+                        images.append(image)
+                    else:
+                        buffered = BytesIO()
+                        image.save(buffered, format="PNG")
+                        img_b64_str = base64.b64encode(buffered.getvalue()).decode()
+                        images.append(img_b64_str)
+        return images
+
+    def to_gradio_chatbot(self):
+        ret = []
+        for i, (role, msg) in enumerate(self.messages[self.offset:]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    import base64
+                    from io import BytesIO
+                    msg, image, image_process_mode = msg
+                    max_hw, min_hw = max(image.size), min(image.size)
+                    aspect_ratio = max_hw / min_hw
+                    max_len, min_len = 800, 400
+                    shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))
+                    longest_edge = int(shortest_edge * aspect_ratio)
+                    W, H = image.size
+                    if H > W:
+                        H, W = longest_edge, shortest_edge
+                    else:
+                        H, W = shortest_edge, longest_edge
+                    image = image.resize((W, H))
+                    buffered = BytesIO()
+                    image.save(buffered, format="JPEG")
+                    img_b64_str = base64.b64encode(buffered.getvalue()).decode()
+                    img_str = f'<img src="data:image/png;base64,{img_b64_str}" alt="user upload image" />'
+                    msg = img_str + msg.replace('<image>', '').strip()
+                    ret.append([msg, None])
+                else:
+                    ret.append([msg, None])
+            else:
+                ret[-1][-1] = msg
+        return ret
+
+    def copy(self):
+        return Conversation(
+            system=self.system,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            version=self.version)
+
+    def dict(self):
+        if len(self.get_images()) > 0:
+            return {
+                "system": self.system,
+                "roles": self.roles,
+                "messages": [[x, y[0] if type(y) is tuple else y] for x, y in self.messages],
+                "offset": self.offset,
+                "sep": self.sep,
+                "sep2": self.sep2,
+            }
+        return {
+            "system": self.system,
+            "roles": self.roles,
+            "messages": self.messages,
+            "offset": self.offset,
+            "sep": self.sep,
+            "sep2": self.sep2,
+        }
+
+
+conv_phi_v0 = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v0",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="<|endoftext|>",
+)
+
+conv_llava_plain = Conversation(
+    system="",
+    roles=("", ""),
+    messages=(
+    ),
+    offset=0,
+    sep_style=SeparatorStyle.PLAIN,
+    sep="\n",
+)
+
+default_conversation = conv_phi_v0
+conv_templates = {
+    "default": conv_phi_v0,
+    "v0": conv_phi_v0,
+    "phi-2_v0": conv_phi_v0,
+
+    "plain": conv_llava_plain,
+}
+
+
+if __name__ == "__main__":
+    print(default_conversation.get_prompt())

inference/inference.ipynb

@@ -0,0 +1,369 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "cdad6b21-030a-40d3-9b31-a229e5b6196d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, AutoTokenizer, AutoConfig, CLIPImageProcessor"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "1f832710-0e8c-42ec-b581-1b15fd2a6acc",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[2024-01-25 14:31:58,511] [INFO] [real_accelerator.py:110:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n"
+     ]
+    }
+   ],
+   "source": [
+    "from model import LlavaPhiForCausalLM"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "9e68f1d4-1ae3-4d45-b818-4600218d2215",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "e5e13e666e3a43d4ad26cc70904abee8",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "model_name = \"RaviNaik/Llava-Phi2\"\n",
+    "model = LlavaPhiForCausalLM.from_pretrained(model_name)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "49edfa0d-e08a-4d3c-a1d6-34068b122419",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
+     ]
+    }
+   ],
+   "source": [
+    "tokenizer = AutoTokenizer.from_pretrained(model_name)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "dcec20cd-d946-42d7-8e10-c198cd49b486",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image_processor = CLIPImageProcessor.from_pretrained(model_name)\n",
+    "mm_use_im_start_end = getattr(model.config, \"mm_use_im_start_end\", False)\n",
+    "mm_use_im_patch_token = getattr(model.config, \"mm_use_im_patch_token\", True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "443c13c4-b7e6-4bc5-b6c7-c577bd4708f6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if mm_use_im_patch_token:\n",
+    "    tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)\n",
+    "if mm_use_im_start_end:\n",
+    "    tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True)\n",
+    "    \n",
+    "if hasattr(model.config, \"max_sequence_length\"):\n",
+    "        context_len = model.config.max_sequence_length\n",
+    "else:\n",
+    "    context_len = 2048"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "d8caee43-0d2a-46d4-bdbc-2cfc7dec9e52",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from transformers import WhisperProcessor, WhisperForConditionalGeneration"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "3acea526-d8ae-4eb6-8dfc-4ea72651b547",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class AudioLanguageConnector:\n",
+    "    def __init__(self, projection_dim):\n",
+    "        model_name = \"microsoft/phi-2\"\n",
+    "        self.phi2_tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n",
+    "        self.phi2_tokenizer.pad_token = self.phi2_tokenizer.eos_token\n",
+    "        self.phi2_tokenizer.max_length = projection_dim\n",
+    "\n",
+    "    def __call__(self, text):\n",
+    "        text = f\"<audio_start> {text} <audio_end>\"\n",
+    "        tokens = self.phi2_tokenizer(text, return_tensors=\"pt\", return_attention_mask=False)\n",
+    "        return tokens\n",
+    "        \n",
+    "\n",
+    "class WhisperWithProjection:\n",
+    "    def __init__(self, projection_dim, device):\n",
+    "        self.device = device\n",
+    "        self.processor = WhisperProcessor.from_pretrained(\"openai/whisper-tiny\", device_map=device)\n",
+    "        self.model = WhisperForConditionalGeneration.from_pretrained(\"openai/whisper-tiny\", device_map=device)\n",
+    "        self.model.config.forced_decoder_ids = None\n",
+    "        # self.audio_language_connector = AudioLanguageConnector(projection_dim)\n",
+    "        \n",
+    "    def __call__(self, audio):\n",
+    "        input_features = self.processor(audio[\"array\"],\n",
+    "                                   sampling_rate=audio[\"sampling_rate\"],\n",
+    "                                   return_tensors=\"pt\").input_features\n",
+    "        # generate token ids\n",
+    "        predicted_ids = self.model.generate(input_features.to(self.device))\n",
+    "        # decode token ids to text        \n",
+    "        transcription = self.processor.batch_decode(predicted_ids, skip_special_tokens=True)\n",
+    "\n",
+    "        # audio_embeddings = self.audio_language_connector(transcription)\n",
+    "        return transcription"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "a2757c91-2ec1-4fe7-9216-03740bf80061",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "IGNORE_INDEX = -100\n",
+    "IMAGE_TOKEN_INDEX = -200\n",
+    "DEFAULT_IMAGE_TOKEN = \"<image>\"\n",
+    "DEFAULT_IMAGE_PATCH_TOKEN = \"<im_patch>\"\n",
+    "DEFAULT_IM_START_TOKEN = \"<im_start>\"\n",
+    "DEFAULT_IM_END_TOKEN = \"<im_end>\"\n",
+    "\n",
+    "from conversation import conv_templates, SeparatorStyle\n",
+    "\n",
+    "class MultiModalPhi2:\n",
+    "    def __init__(self, modelname_or_path=\"RaviNaik/Llava-Phi2\",\n",
+    "                temperature=0.2,\n",
+    "                max_new_tokens=1024,\n",
+    "                device=\"cuda:0\"):\n",
+    "        self.model_name = modelname_or_path\n",
+    "        self.temperature = temperature\n",
+    "        self.max_new_tokens = max_new_tokens\n",
+    "        self.device = device\n",
+    "        self.disable_torch_init()\n",
+    "        self.whisper_w_proj = WhisperWithProjection(projection_dim=512, device=device)\n",
+    "        self.load_pretrained_model()\n",
+    "        \n",
+    "    def disable_torch_init(self):\n",
+    "        \"\"\"\n",
+    "        Disable the redundant torch default initialization to accelerate model creation.\n",
+    "        \"\"\"\n",
+    "        setattr(torch.nn.Linear, \"reset_parameters\", lambda self: None)\n",
+    "        setattr(torch.nn.LayerNorm, \"reset_parameters\", lambda self: None)\n",
+    "        \n",
+    "    def load_pretrained_model(self):\n",
+    "        self.model = LlavaPhiForCausalLM.from_pretrained(self.model_name, device_map=self.device)\n",
+    "        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)\n",
+    "        self.image_processor = CLIPImageProcessor.from_pretrained(self.model_name)\n",
+    "        mm_use_im_start_end = getattr(self.model.config, \"mm_use_im_start_end\", False)\n",
+    "        mm_use_im_patch_token = getattr(self.model.config, \"mm_use_im_patch_token\", True)\n",
+    "        if mm_use_im_patch_token:\n",
+    "            self.tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)\n",
+    "        if mm_use_im_start_end:\n",
+    "            self.tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True)\n",
+    "    \n",
+    "    def tokenizer_image_token(self, prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors=None):\n",
+    "        prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split('<image>')]\n",
+    "    \n",
+    "        def insert_separator(X, sep):\n",
+    "            return [ele for sublist in zip(X, [sep]*len(X)) for ele in sublist][:-1]\n",
+    "    \n",
+    "        input_ids = []\n",
+    "        offset = 0\n",
+    "        if len(prompt_chunks) > 0 and len(prompt_chunks[0]) > 0 and prompt_chunks[0][0] == tokenizer.bos_token_id:\n",
+    "            offset = 1\n",
+    "            input_ids.append(prompt_chunks[0][0])\n",
+    "        for x in insert_separator(prompt_chunks, [image_token_index] * (offset + 1)):\n",
+    "            input_ids.extend(x[offset:])\n",
+    "    \n",
+    "        if return_tensors is not None:\n",
+    "            if return_tensors == 'pt':\n",
+    "                return torch.tensor(input_ids, dtype=torch.long)\n",
+    "            raise ValueError(f'Unsupported tensor type: {return_tensors}')\n",
+    "        return input_ids\n",
+    "        \n",
+    "    def __call__(self, text, audio, image):\n",
+    "        qs = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\\n' + text\n",
+    "        conv = conv_templates[\"phi-2_v0\"].copy()\n",
+    "        conv.append_message(conv.roles[0], qs)\n",
+    "        conv.append_message(conv.roles[1], None)\n",
+    "        prompt = conv.get_prompt()\n",
+    "\n",
+    "        image_tensor = self.image_processor.preprocess(image, return_tensors='pt')['pixel_values'].to(self.device)\n",
+    "        \n",
+    "        input_ids = self.tokenizer_image_token(prompt, self.tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0)\n",
+    "        if audio is not None:\n",
+    "            audio_transcript = self.whisper_w_proj(audio)\n",
+    "            audio_embed = self.tokenizer(audio_transcript, return_tensors='pt')[\"input_ids\"]\n",
+    "            input_ids = torch.concat([input_ids, audio_embed], dim=1)\n",
+    "        input_ids = input_ids.to(self.device)\n",
+    "            \n",
+    "        stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2\n",
+    "\n",
+    "        with torch.inference_mode():\n",
+    "            output_ids = self.model.generate(\n",
+    "                input_ids,\n",
+    "                images=image_tensor,\n",
+    "                do_sample=True,\n",
+    "                temperature=self.temperature,\n",
+    "                max_new_tokens=self.max_new_tokens,\n",
+    "                eos_token_id=self.tokenizer.eos_token_id,  # End of sequence token\n",
+    "                pad_token_id=self.tokenizer.eos_token_id,  # Pad token\n",
+    "                use_cache=True,\n",
+    "            )\n",
+    "\n",
+    "        input_token_len = input_ids.shape[1]\n",
+    "        n_diff_input_output = (input_ids != output_ids[:, :input_token_len]).sum().item()\n",
+    "        if n_diff_input_output > 0:\n",
+    "            print(f'[Warning] {n_diff_input_output} output_ids are not the same as the input_ids')\n",
+    "        outputs = self.tokenizer.batch_decode(output_ids[:, input_token_len:], skip_special_tokens=True)[0]\n",
+    "        outputs = outputs.strip()\n",
+    "        if outputs.endswith(stop_str):\n",
+    "            outputs = outputs[:-len(stop_str)]\n",
+    "        outputs = outputs.strip()\n",
+    "        return outputs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "cc47e6a0-3544-4a60-930f-ccae87ef945a",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "9ef56077307d4cef907e25b092061611",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
+     ]
+    }
+   ],
+   "source": [
+    "multimodal_phi2 = MultiModalPhi2()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "cb8aca1b-7d75-45e7-b5a4-71d151f792e1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from PIL import Image\n",
+    "import requests\n",
+    "\n",
+    "url = \"https://www.ilankelman.org/stopsigns/australia.jpg\"\n",
+    "image = Image.open(requests.get(url, stream=True).raw)\n",
+    "\n",
+    "from datasets import load_dataset\n",
+    "audio_ds = load_dataset(\"hf-internal-testing/librispeech_asr_dummy\", \"clean\", split=\"validation\")\n",
+    "audio = audio_ds[0][\"audio\"]\n",
+    "\n",
+    "text = \"tell me about the image\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "6767efc6-be4f-44d3-84ff-34db57d9f940",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'In the image, there is a Chinese writing on a pole in a foreign language. This suggests that the image was taken in a foreign country, possibly in a foreign country. The sign is in a foreign language, which might be in a foreign language. The sign is written in Japanese, which is a common language in Japan. The sign is also written in two different languages, which suggests that it is written in a language that is not in the native language.'"
+      ]
+     },
+     "execution_count": 14,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "multimodal_phi2(text, None, image)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0bdd0b8a-709b-4c82-ac1d-dc746d3a0748",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

inference/main.py

@@ -0,0 +1,226 @@
+import soundfile as sf
+import librosa
+import torch
+from transformers import (
+    AutoTokenizer,
+    CLIPImageProcessor,
+    WhisperProcessor,
+    WhisperForConditionalGeneration,
+)
+
+from .model import LlavaPhiForCausalLM
+from .conversation import conv_templates, SeparatorStyle
+
+IGNORE_INDEX = -100
+IMAGE_TOKEN_INDEX = -200
+DEFAULT_IMAGE_TOKEN = "<image>"
+DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>"
+DEFAULT_IM_START_TOKEN = "<im_start>"
+DEFAULT_IM_END_TOKEN = "<im_end>"
+
+
+class AudioLanguageConnector:
+    def __init__(self, projection_dim):
+        model_name = "microsoft/phi-2"
+        self.phi2_tokenizer = AutoTokenizer.from_pretrained(
+            model_name, trust_remote_code=True
+        )
+        self.phi2_tokenizer.pad_token = self.phi2_tokenizer.eos_token
+        self.phi2_tokenizer.max_length = projection_dim
+
+    def __call__(self, text):
+        text = f"<audio_start> {text} <audio_end>"
+        tokens = self.phi2_tokenizer(
+            text, return_tensors="pt", return_attention_mask=False
+        )
+        return tokens
+
+
+class WhisperWithProjection:
+    def __init__(self, projection_dim, device):
+        self.device = device
+        self.processor = WhisperProcessor.from_pretrained(
+            "openai/whisper-tiny", device_map=device
+        )
+        self.model = WhisperForConditionalGeneration.from_pretrained(
+            "openai/whisper-tiny", device_map=device
+        )
+        self.model.config.forced_decoder_ids = None
+        # self.audio_language_connector = AudioLanguageConnector(projection_dim)
+
+    def __call__(self, audio):
+        array, sampling_rate = sf.read(audio)
+        resampled_array = librosa.resample(
+            array,
+            orig_sr=sampling_rate,
+            target_sr=16000,
+        )
+        input_features = self.processor(
+            resampled_array, sampling_rate=16000, return_tensors="pt"
+        ).input_features
+        # generate token ids
+        predicted_ids = self.model.generate(input_features.to(self.device))
+        # decode token ids to text
+        transcription = self.processor.batch_decode(
+            predicted_ids, skip_special_tokens=True
+        )
+
+        # audio_embeddings = self.audio_language_connector(transcription)
+        return transcription
+
+
+class MultiModalPhi2:
+    def __init__(
+        self,
+        modelname_or_path="RaviNaik/Llava-Phi2",
+        temperature=0.2,
+        max_new_tokens=1024,
+        device="cuda:0",
+    ):
+        self.model_name = modelname_or_path
+        self.temperature = temperature
+        self.max_new_tokens = max_new_tokens
+        self.device = device
+        self.disable_torch_init()
+        self.whisper_w_proj = WhisperWithProjection(projection_dim=512, device=device)
+        self.load_pretrained_model()
+
+    def disable_torch_init(self):
+        """
+        Disable the redundant torch default initialization to accelerate model creation.
+        """
+        setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
+        setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
+
+    def load_pretrained_model(self):
+        self.model = LlavaPhiForCausalLM.from_pretrained(
+            self.model_name, device_map=self.device
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        self.image_processor = CLIPImageProcessor.from_pretrained(self.model_name)
+        mm_use_im_start_end = getattr(self.model.config, "mm_use_im_start_end", False)
+        mm_use_im_patch_token = getattr(
+            self.model.config, "mm_use_im_patch_token", True
+        )
+        if mm_use_im_patch_token:
+            self.tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
+        if mm_use_im_start_end:
+            self.tokenizer.add_tokens(
+                [DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True
+            )
+
+    def tokenizer_image_token(
+        self,
+        prompt,
+        tokenizer,
+        image_token_index=IMAGE_TOKEN_INDEX,
+        return_tensors=None,
+    ):
+        prompt_chunks = [
+            tokenizer(chunk).input_ids for chunk in prompt.split("<image>")
+        ]
+
+        def insert_separator(X, sep):
+            return [ele for sublist in zip(X, [sep] * len(X)) for ele in sublist][:-1]
+
+        input_ids = []
+        offset = 0
+        if (
+            len(prompt_chunks) > 0
+            and len(prompt_chunks[0]) > 0
+            and prompt_chunks[0][0] == tokenizer.bos_token_id
+        ):
+            offset = 1
+            input_ids.append(prompt_chunks[0][0])
+        for x in insert_separator(prompt_chunks, [image_token_index] * (offset + 1)):
+            input_ids.extend(x[offset:])
+
+        if return_tensors is not None:
+            if return_tensors == "pt":
+                return torch.tensor(input_ids, dtype=torch.long)
+            raise ValueError(f"Unsupported tensor type: {return_tensors}")
+        return input_ids
+
+    def __call__(self, text, audio, image):
+        if text is None:
+            text = ""
+        if image is not None:
+            qs = (
+                DEFAULT_IM_START_TOKEN
+                + DEFAULT_IMAGE_TOKEN
+                + DEFAULT_IM_END_TOKEN
+                + "\n"
+                + text
+            )
+            conv = conv_templates["phi-2_v0"].copy()
+            conv.append_message(conv.roles[0], qs)
+            conv.append_message(conv.roles[1], None)
+            prompt = conv.get_prompt()
+
+            input_ids = self.tokenizer_image_token(
+                prompt, self.tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt"
+            ).unsqueeze(0)
+
+            image_tensor = self.image_processor.preprocess(image, return_tensors="pt")[
+                "pixel_values"
+            ].to(self.device)
+        else:
+            qs = text
+            conv = conv_templates["phi-2_v0"].copy()
+            conv.append_message(conv.roles[0], qs)
+            conv.append_message(conv.roles[1], None)
+            prompt = conv.get_prompt()
+
+            input_ids = self.tokenizer(prompt, return_tensors="pt")["input_ids"]
+
+            image_tensor = None
+
+        if audio is not None:
+            audio_transcript = self.whisper_w_proj(audio)
+            audio_embed = self.tokenizer(audio_transcript, return_tensors="pt")[
+                "input_ids"
+            ]
+            input_ids = torch.concat([input_ids, audio_embed], dim=1)
+        input_ids = input_ids.to(self.device)
+
+        stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
+
+        with torch.inference_mode():
+            if image is not None:
+                output_ids = self.model.generate(
+                    input_ids,
+                    images=image_tensor,
+                    do_sample=True,
+                    temperature=self.temperature,
+                    max_new_tokens=self.max_new_tokens,
+                    eos_token_id=self.tokenizer.eos_token_id,  # End of sequence token
+                    pad_token_id=self.tokenizer.eos_token_id,  # Pad token
+                    use_cache=True,
+                )
+            else:
+                output_ids = self.model.generate(
+                    input_ids,
+                    do_sample=True,
+                    temperature=self.temperature,
+                    max_new_tokens=self.max_new_tokens,
+                    eos_token_id=self.tokenizer.eos_token_id,  # End of sequence token
+                    pad_token_id=self.tokenizer.eos_token_id,  # Pad token
+                    use_cache=True,
+                )
+
+        input_token_len = input_ids.shape[1]
+        n_diff_input_output = (
+            (input_ids != output_ids[:, :input_token_len]).sum().item()
+        )
+        if n_diff_input_output > 0:
+            print(
+                f"[Warning] {n_diff_input_output} output_ids are not the same as the input_ids"
+            )
+        outputs = self.tokenizer.batch_decode(
+            output_ids[:, input_token_len:], skip_special_tokens=True
+        )[0]
+        outputs = outputs.strip()
+        if outputs.endswith(stop_str):
+            outputs = outputs[: -len(stop_str)]
+        outputs = outputs.strip()
+        return outputs

inference/model/__init__.py

@@ -0,0 +1,2 @@
+from .language_model.llava_phi import LlavaPhiForCausalLM
+from .language_model.configuration_llava_phi import LlavaPhiConfig, LlavaPhiVisionConfig, ProjectorConfig

inference/model/builder.py

@@ -0,0 +1,180 @@
+import os
+import warnings
+
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    AutoConfig,
+    BitsAndBytesConfig,
+    CLIPImageProcessor,
+)
+import torch
+from .language_model.llava_phi import LlavaPhiForCausalLM
+from .language_model.configuration_llava_phi import LlavaPhiConfig
+
+IGNORE_INDEX = -100
+IMAGE_TOKEN_INDEX = -200
+DEFAULT_IMAGE_TOKEN = "<image>"
+DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>"
+DEFAULT_IM_START_TOKEN = "<im_start>"
+DEFAULT_IM_END_TOKEN = "<im_end>"
+
+
+def load_pretrained_model(
+    model_path,
+    model_base,
+    model_name,
+    load_8bit=False,
+    load_4bit=False,
+    device_map="cuda",
+    device="cuda",
+):
+    kwargs = {"device_map": device_map}
+    if load_8bit:
+        kwargs["load_in_8bit"] = True
+    elif load_4bit:
+        kwargs["load_in_4bit"] = True
+        kwargs["quantization_config"] = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+        )
+    # else:  # TODO: after fine-tuning LLava-Phi, load the model weights with fp16 will pose nan
+    #     kwargs['torch_dtype'] = torch.float16
+
+    if "phi" in model_name.lower():
+        # Load LLaVA-Phi model
+        if "lora" in model_name.lower() and model_base is None:
+            warnings.warn(
+                "There is `lora` in model name but no `model_base` is provided. If you are loading a LoRA model, please provide the `model_base` argument."
+            )
+        if "lora" in model_name.lower() and model_base is not None:
+            lora_cfg_pretrained = AutoConfig.from_pretrained(model_path)
+            tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
+            print("Loading LLaVA-Phi from base model...")
+            model = LlavaPhiForCausalLM.from_pretrained(
+                model_base, low_cpu_mem_usage=True, config=lora_cfg_pretrained, **kwargs
+            )
+            token_num, tokem_dim = model.lm_head.out_features, model.lm_head.in_features
+            if model.lm_head.weight.shape[0] != token_num:
+                model.lm_head.weight = torch.nn.Parameter(
+                    torch.empty(
+                        token_num, tokem_dim, device=model.device, dtype=model.dtype
+                    )
+                )
+                model.model.embed_tokens.weight = torch.nn.Parameter(
+                    torch.empty(
+                        token_num, tokem_dim, device=model.device, dtype=model.dtype
+                    )
+                )
+
+            print("Loading additional LLaVA-Phi weights...")
+            if os.path.exists(os.path.join(model_path, "non_lora_trainables.bin")):
+                non_lora_trainables = torch.load(
+                    os.path.join(model_path, "non_lora_trainables.bin"),
+                    map_location="cpu",
+                )
+            else:
+                # this is probably from HF Hub
+                from huggingface_hub import hf_hub_download
+
+                def load_from_hf(repo_id, filename, subfolder=None):
+                    cache_file = hf_hub_download(
+                        repo_id=repo_id, filename=filename, subfolder=subfolder
+                    )
+                    return torch.load(cache_file, map_location="cpu")
+
+                non_lora_trainables = load_from_hf(
+                    model_path, "non_lora_trainables.bin"
+                )
+            non_lora_trainables = {
+                (k[11:] if k.startswith("base_model.") else k): v
+                for k, v in non_lora_trainables.items()
+            }
+            if any(k.startswith("model.model.") for k in non_lora_trainables):
+                non_lora_trainables = {
+                    (k[6:] if k.startswith("model.") else k): v
+                    for k, v in non_lora_trainables.items()
+                }
+            model.load_state_dict(non_lora_trainables, strict=False)
+
+            from peft import PeftModel
+
+            print("Loading LoRA weights...")
+            model = PeftModel.from_pretrained(model, model_path)
+            print("Merging LoRA weights...")
+            model = model.merge_and_unload()
+            print("Model is loaded...")
+        elif model_base is not None:
+            # this may be mm projector only
+            print("Loading LLaVA-Phi from base model...")
+            tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
+            cfg_pretrained = AutoConfig.from_pretrained(model_path)
+            model = LlavaPhiForCausalLM.from_pretrained(
+                model_base, low_cpu_mem_usage=True, config=cfg_pretrained, **kwargs
+            )
+
+            mm_projector_weights = torch.load(
+                os.path.join(model_path, "mm_projector.bin"), map_location="cpu"
+            )
+            mm_projector_weights = {
+                k: v.to(torch.float16) for k, v in mm_projector_weights.items()
+            }
+            model.load_state_dict(mm_projector_weights, strict=False)
+        else:
+            print("load llaVA-Phi MLLM!!!")
+            config = LlavaPhiConfig.from_pretrained(model_path, trust_remote_code=True)
+            tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)
+            model = LlavaPhiForCausalLM.from_pretrained(
+                model_path, config=config, use_safetensors=True, **kwargs
+            ).to("cuda")
+    else:
+        # Load language model
+        if model_base is not None:
+            # PEFT model
+            from peft import PeftModel
+
+            tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
+            model = AutoModelForCausalLM.from_pretrained(
+                model_base,
+                torch_dtype=torch.float16,
+                low_cpu_mem_usage=True,
+                device_map="auto",
+            )
+            print(f"Loading LoRA weights from {model_path}")
+            model = PeftModel.from_pretrained(model, model_path)
+            print(f"Merging weights")
+            model = model.merge_and_unload()
+            print("Convert to FP16...")
+            model.to(torch.float16)
+        else:
+            tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+            model = AutoModelForCausalLM.from_pretrained(
+                model_path, low_cpu_mem_usage=True, **kwargs
+            )
+
+    image_processor = CLIPImageProcessor.from_pretrained(model_path)
+
+    if "phi" in model_name.lower():
+        mm_use_im_start_end = getattr(model.config, "mm_use_im_start_end", False)
+        mm_use_im_patch_token = getattr(model.config, "mm_use_im_patch_token", True)
+
+        # TODO: the tokenizer length of phi-2 is 50295, but the output class of lm_head is 51200
+        if mm_use_im_patch_token:
+            tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
+        if mm_use_im_start_end:
+            tokenizer.add_tokens(
+                [DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True
+            )
+            # model.resize_token_embeddings(len(tokenizer))
+    else:
+        raise ValueError(f"Unsupported model name: {model_name}")
+
+    if hasattr(model.config, "max_sequence_length"):
+        context_len = model.config.max_sequence_length
+    else:
+        context_len = 2048
+    model.to(device="cuda")
+    print(kwargs)
+    return tokenizer, model, image_processor, context_len

inference/model/language_model/configuration_llava_phi.py

@@ -0,0 +1,191 @@
+import os
+from typing import Union
+from transformers import PretrainedConfig, PhiConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class LlavaPhiVisionConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`CLIPVisionModel`]. It is used to instantiate a
+    CLIP vision encoder according to the specified arguments, defining the model architecture. Instantiating a
+    configuration with the defaults will yield a similar configuration to that of the vision encoder of the CLIP
+    [openai/clip-vit-base-patch32](https://huggingface.co/openai/clip-vit-base-patch32) architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        hidden_size (`int`, *optional*, defaults to 768):
+            Dimensionality of the encoder layers and the pooler layer.
+        intermediate_size (`int`, *optional*, defaults to 3072):
+            Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
+        projection_dim (`int`, *optional*, defaults to 512):
+            Dimentionality of text and vision projection layers.
+        num_hidden_layers (`int`, *optional*, defaults to 12):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 12):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        num_channels (`int`, *optional*, defaults to 3):
+            The number of input channels.
+        image_size (`int`, *optional*, defaults to 224):
+            The size (resolution) of each image.
+        patch_size (`int`, *optional*, defaults to 32):
+            The size (resolution) of each patch.
+        hidden_act (`str` or `function`, *optional*, defaults to `"quick_gelu"`):
+            The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
+            `"relu"`, `"selu"` and `"gelu_new"` ``"quick_gelu"` are supported.
+        layer_norm_eps (`float`, *optional*, defaults to 1e-05):
+            The epsilon used by the layer normalization layers.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        initializer_factor (`float`, *optional*, defaults to 1.0):
+            A factor for initializing all weight matrices (should be kept to 1, used internally for initialization
+            testing).
+        mm_vision_select_feature (`str`, *optional*, defaults to `"patch"`):
+            The feature to select from the vision encoder output. Can be one of `"patch"` or `"cls_patch"`.
+        mm_vision_select_layer (`int`, *optional*, defaults to `-2`):
+            The layer to select from the vision encoder output.
+
+    Example:
+
+    ```python
+    >>> from transformers import CLIPVisionConfig, CLIPVisionModel
+
+    >>> # Initializing a CLIPVisionConfig with openai/clip-vit-base-patch32 style configuration
+    >>> configuration = CLIPVisionConfig()
+
+    >>> # Initializing a CLIPVisionModel (with random weights) from the openai/clip-vit-base-patch32 style configuration
+    >>> model = CLIPVisionModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "llava_phi_clip_vision_model"
+
+    def __init__(
+        self,
+        hidden_size=768,
+        intermediate_size=3072,
+        projection_dim=512,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        num_channels=3,
+        image_size=224,
+        patch_size=32,
+        hidden_act="quick_gelu",
+        layer_norm_eps=1e-5,
+        attention_dropout=0.0,
+        initializer_range=0.02,
+        initializer_factor=1.0,
+        mm_vision_select_feature="patch",
+        mm_vision_select_layer=-2,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.projection_dim = projection_dim
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_channels = num_channels
+        self.patch_size = patch_size
+        self.image_size = image_size
+        self.initializer_range = initializer_range
+        self.initializer_factor = initializer_factor
+        self.attention_dropout = attention_dropout
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_act = hidden_act
+        self.mm_vision_select_feature = mm_vision_select_feature
+        self.mm_vision_select_layer = mm_vision_select_layer
+
+    @classmethod
+    def from_pretrained(
+        cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs
+    ) -> "PretrainedConfig":
+        cls._set_token_in_kwargs(kwargs)
+
+        config_dict, kwargs = cls.get_config_dict(
+            pretrained_model_name_or_path, **kwargs
+        )
+
+        # get the vision config dict if we are loading from CLIPConfig
+        if config_dict.get("model_type") == "llava_phi-phi":
+            config_dict = config_dict["vision_config"]
+
+        if (
+            "model_type" in config_dict
+            and hasattr(cls, "model_type")
+            and config_dict["model_type"] != cls.model_type
+        ):
+            logger.warning(
+                f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
+                f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
+            )
+
+        return cls.from_dict(config_dict, **kwargs)
+
+
+class ProjectorConfig(PretrainedConfig):
+    model_type = "llava_phi_projector"
+
+    def __init__(
+        self, mm_projector_type="linear", mm_hidden_size=768, hidden_size=2560, **kwargs
+    ):
+        self.mm_projector_type = mm_projector_type
+        self.mm_hidden_size = mm_hidden_size
+        self.hidden_size = hidden_size
+        super().__init__(**kwargs)
+
+    @classmethod
+    def from_pretrained(
+        cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs
+    ) -> "PretrainedConfig":
+        cls._set_token_in_kwargs(kwargs)
+
+        config_dict, kwargs = cls.get_config_dict(
+            pretrained_model_name_or_path, **kwargs
+        )
+
+        # get the vision config dict if we are loading from CLIPConfig
+        if config_dict.get("model_type") == "llava_phi-phi":
+            config_dict = config_dict["projector_config"]
+
+        if (
+            "model_type" in config_dict
+            and hasattr(cls, "model_type")
+            and config_dict["model_type"] != cls.model_type
+        ):
+            logger.warning(
+                f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
+                f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
+            )
+
+        return cls.from_dict(config_dict, **kwargs)
+
+
+DEFAULT_VISUAL_CONFIG = {
+    "vision_tower": LlavaPhiVisionConfig().to_dict(),
+    "mm_projector": ProjectorConfig().to_dict(),
+}
+
+
+class LlavaPhiConfig(PhiConfig):
+    model_type = "llava_phi"
+
+    def __init__(self, vision_config=None, **kwargs):
+        if vision_config is None:
+            self.vision_config = DEFAULT_VISUAL_CONFIG
+        else:
+            self.vision_config = vision_config
+
+        super().__init__(**kwargs)
+
+
+if __name__ == "__main__":
+    print(LlavaPhiVisionConfig())

inference/model/language_model/llava_phi.py

@@ -0,0 +1,126 @@
+import os
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from torch.nn import CrossEntropyLoss
+
+from transformers import AutoConfig, AutoModelForCausalLM, \
+    PhiModel, PhiPreTrainedModel
+
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from ..llava_arch import LlavaMetaModel, LlavaMetaForCausalLM
+from transformers.utils import logging
+from .configuration_llava_phi import LlavaPhiConfig
+
+logger = logging.get_logger(__name__)
+
+
+class LLavaPhiModel(LlavaMetaModel, PhiModel):
+    config_class = LlavaPhiConfig
+
+    def __init__(self, config):
+        super(LLavaPhiModel, self).__init__(config)
+
+
+class LlavaPhiForCausalLM(PhiPreTrainedModel, LlavaMetaForCausalLM):
+    config_class = LlavaPhiConfig
+
+    def __init__(self, config):
+        super(PhiPreTrainedModel, self).__init__(config)
+        self.model = LLavaPhiModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=True)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_model(self):
+        return self.model
+
+    def forward(
+            self,
+            input_ids: torch.LongTensor = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            past_key_values: Optional[List[torch.FloatTensor]] = None,
+            inputs_embeds: Optional[torch.FloatTensor] = None,
+            labels: Optional[torch.LongTensor] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            images: Optional[torch.FloatTensor] = None,
+            return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        input_ids, attention_mask, past_key_values, inputs_embeds, labels = self.prepare_inputs_labels_for_multimodal(
+            input_ids, attention_mask, past_key_values, labels, images)
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict
+        )
+
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model/pipeline parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+        
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+            self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+    ):
+        if past_key_values:
+            input_ids = input_ids[:, -1:]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+                "images": kwargs.get("images", None),
+            }
+        )
+        return model_inputs
+
+
+AutoConfig.register("llava_phi", LlavaPhiConfig)
+AutoModelForCausalLM.register(LlavaPhiConfig, LlavaPhiForCausalLM)

inference/model/llava_arch.py

@@ -0,0 +1,330 @@
+#    Copyright 2023 Haotian Liu
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+
+from abc import ABC, abstractmethod
+
+import torch
+
+from .multimodal_encoder.clip_encoder import CLIPVisionTower
+from .multimodal_projector.builder import build_vision_projector
+from .language_model.configuration_llava_phi import (
+    LlavaPhiConfig,
+    LlavaPhiVisionConfig,
+    ProjectorConfig,
+)
+
+# from llava_phi.constants import IGNORE_INDEX, IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
+IGNORE_INDEX = -100
+IMAGE_TOKEN_INDEX = -200
+DEFAULT_IMAGE_TOKEN = "<image>"
+DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>"
+DEFAULT_IM_START_TOKEN = "<im_start>"
+DEFAULT_IM_END_TOKEN = "<im_end>"
+
+
+class LlavaMetaModel:
+    def __init__(self, config):
+        super(LlavaMetaModel, self).__init__(config)
+        self.vision_tower = CLIPVisionTower(
+            LlavaPhiVisionConfig(**config.vision_config["vision_tower"])
+        )
+        self.mm_projector = build_vision_projector(
+            ProjectorConfig(**config.vision_config["mm_projector"])
+        )
+
+    def get_vision_tower(self):
+        vision_tower = getattr(self, "vision_tower", None)
+        if type(vision_tower) is list:
+            vision_tower = vision_tower[0]
+        return vision_tower
+
+
+class LlavaMetaForCausalLM(ABC):
+    @abstractmethod
+    def get_model(self):
+        pass
+
+    def get_vision_tower(self):
+        return self.get_model().get_vision_tower()
+
+    def encode_images(self, images):
+        image_features = self.get_model().get_vision_tower()(images)
+        image_features = self.get_model().mm_projector(image_features)
+        return image_features
+
+    def prepare_inputs_labels_for_multimodal(
+        self, input_ids, attention_mask, past_key_values, labels, images
+    ):
+        vision_tower = self.get_vision_tower()
+        if vision_tower is None or images is None or input_ids.shape[1] == 1:
+            if (
+                past_key_values is not None
+                and vision_tower is not None
+                and images is not None
+                and input_ids.shape[1] == 1
+            ):
+                attention_mask = torch.ones(
+                    (attention_mask.shape[0], past_key_values[-1][-1].shape[-2] + 1),
+                    dtype=attention_mask.dtype,
+                    device=attention_mask.device,
+                )
+            return input_ids, attention_mask, past_key_values, None, labels
+
+        if type(images) is list or images.ndim == 5:
+            concat_images = torch.cat([image for image in images], dim=0)
+            image_features = self.encode_images(concat_images)
+            split_sizes = [image.shape[0] for image in images]
+            image_features = torch.split(image_features, split_sizes, dim=0)
+            image_features = [x.flatten(0, 1) for x in image_features]
+        else:
+            image_features = self.encode_images(images)
+
+        new_input_embeds = []
+        new_labels = [] if labels is not None else None
+        cur_image_idx = 0
+        for batch_idx, cur_input_ids in enumerate(input_ids):
+            if (cur_input_ids == IMAGE_TOKEN_INDEX).sum() == 0:
+                # multimodal LLM, but the current sample is not multimodal
+                # FIXME: this is a hacky fix, for deepspeed zero3 to work
+                half_len = cur_input_ids.shape[0] // 2
+                cur_image_features = image_features[cur_image_idx]
+                cur_input_embeds_1 = self.get_model().embed_tokens(
+                    cur_input_ids[:half_len]
+                )
+                cur_input_embeds_2 = self.get_model().embed_tokens(
+                    cur_input_ids[half_len:]
+                )
+                cur_input_embeds = torch.cat(
+                    [cur_input_embeds_1, cur_image_features[0:0], cur_input_embeds_2],
+                    dim=0,
+                )
+                new_input_embeds.append(cur_input_embeds)
+                if labels is not None:
+                    new_labels.append(labels[batch_idx])
+                cur_image_idx += 1
+                continue
+            image_token_indices = torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0]
+            cur_new_input_embeds = []
+            if labels is not None:
+                cur_labels = labels[batch_idx]
+                cur_new_labels = []
+                assert cur_labels.shape == cur_input_ids.shape
+            while image_token_indices.numel() > 0:
+                cur_image_features = image_features[cur_image_idx]
+                image_token_start = image_token_indices[0]
+                if getattr(self.config, "tune_mm_mlp_adapter", False) and getattr(
+                    self.config, "mm_use_im_start_end", False
+                ):
+                    cur_new_input_embeds.append(
+                        self.get_model()
+                        .embed_tokens(cur_input_ids[: image_token_start - 1])
+                        .detach()
+                    )
+                    cur_new_input_embeds.append(
+                        self.get_model().embed_tokens(
+                            cur_input_ids[image_token_start - 1 : image_token_start]
+                        )
+                    )
+                    cur_new_input_embeds.append(cur_image_features)
+                    cur_new_input_embeds.append(
+                        self.get_model().embed_tokens(
+                            cur_input_ids[image_token_start + 1 : image_token_start + 2]
+                        )
+                    )
+                    if labels is not None:
+                        cur_new_labels.append(cur_labels[:image_token_start])
+                        cur_new_labels.append(
+                            torch.full(
+                                (cur_image_features.shape[0],),
+                                IGNORE_INDEX,
+                                device=labels.device,
+                                dtype=labels.dtype,
+                            )
+                        )
+                        cur_new_labels.append(
+                            cur_labels[image_token_start : image_token_start + 1]
+                        )
+                        cur_labels = cur_labels[image_token_start + 2 :]
+                else:
+                    cur_new_input_embeds.append(
+                        self.get_model().embed_tokens(cur_input_ids[:image_token_start])
+                    )
+                    cur_new_input_embeds.append(cur_image_features)
+                    if labels is not None:
+                        cur_new_labels.append(cur_labels[:image_token_start])
+                        cur_new_labels.append(
+                            torch.full(
+                                (cur_image_features.shape[0],),
+                                IGNORE_INDEX,
+                                device=labels.device,
+                                dtype=labels.dtype,
+                            )
+                        )
+                        cur_labels = cur_labels[image_token_start + 1 :]
+                cur_image_idx += 1
+                if getattr(self.config, "tune_mm_mlp_adapter", False) and getattr(
+                    self.config, "mm_use_im_start_end", False
+                ):
+                    cur_input_ids = cur_input_ids[image_token_start + 2 :]
+                else:
+                    cur_input_ids = cur_input_ids[image_token_start + 1 :]
+                image_token_indices = torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0]
+            if cur_input_ids.numel() > 0:
+                if getattr(self.config, "tune_mm_mlp_adapter", False) and getattr(
+                    self.config, "mm_use_im_start_end", False
+                ):
+                    cur_new_input_embeds.append(
+                        self.get_model().embed_tokens(cur_input_ids).detach()
+                    )
+                else:
+                    cur_new_input_embeds.append(
+                        self.get_model().embed_tokens(cur_input_ids)
+                    )
+                if labels is not None:
+                    cur_new_labels.append(cur_labels)
+            cur_new_input_embeds = [
+                x.to(device=self.device) for x in cur_new_input_embeds
+            ]
+            cur_new_input_embeds = torch.cat(cur_new_input_embeds, dim=0)
+            new_input_embeds.append(cur_new_input_embeds)
+            if labels is not None:
+                cur_new_labels = torch.cat(cur_new_labels, dim=0)
+                new_labels.append(cur_new_labels)
+
+        if any(x.shape != new_input_embeds[0].shape for x in new_input_embeds):
+            max_len = max(x.shape[0] for x in new_input_embeds)
+
+            new_input_embeds_align = []
+            for cur_new_embed in new_input_embeds:
+                cur_new_embed = torch.cat(
+                    (
+                        cur_new_embed,
+                        torch.zeros(
+                            (max_len - cur_new_embed.shape[0], cur_new_embed.shape[1]),
+                            dtype=cur_new_embed.dtype,
+                            device=cur_new_embed.device,
+                        ),
+                    ),
+                    dim=0,
+                )
+                new_input_embeds_align.append(cur_new_embed)
+            new_input_embeds = torch.stack(new_input_embeds_align, dim=0)
+
+            if labels is not None:
+                new_labels_align = []
+                _new_labels = new_labels
+                for cur_new_label in new_labels:
+                    cur_new_label = torch.cat(
+                        (
+                            cur_new_label,
+                            torch.full(
+                                (max_len - cur_new_label.shape[0],),
+                                IGNORE_INDEX,
+                                dtype=cur_new_label.dtype,
+                                device=cur_new_label.device,
+                            ),
+                        ),
+                        dim=0,
+                    )
+                    new_labels_align.append(cur_new_label)
+                new_labels = torch.stack(new_labels_align, dim=0)
+
+            if attention_mask is not None:
+                new_attention_mask = []
+                for cur_attention_mask, cur_new_labels, cur_new_labels_align in zip(
+                    attention_mask, _new_labels, new_labels
+                ):
+                    new_attn_mask_pad_left = torch.full(
+                        (cur_new_labels.shape[0] - labels.shape[1],),
+                        True,
+                        dtype=attention_mask.dtype,
+                        device=attention_mask.device,
+                    )
+                    new_attn_mask_pad_right = torch.full(
+                        (cur_new_labels_align.shape[0] - cur_new_labels.shape[0],),
+                        False,
+                        dtype=attention_mask.dtype,
+                        device=attention_mask.device,
+                    )
+                    cur_new_attention_mask = torch.cat(
+                        (
+                            new_attn_mask_pad_left,
+                            cur_attention_mask,
+                            new_attn_mask_pad_right,
+                        ),
+                        dim=0,
+                    )
+                    new_attention_mask.append(cur_new_attention_mask)
+                attention_mask = torch.stack(new_attention_mask, dim=0)
+                assert attention_mask.shape == new_labels.shape
+        else:
+            new_input_embeds = torch.stack(new_input_embeds, dim=0)
+            if labels is not None:
+                new_labels = torch.stack(new_labels, dim=0)
+
+            if attention_mask is not None:
+                new_attn_mask_pad_left = torch.full(
+                    (
+                        attention_mask.shape[0],
+                        new_input_embeds.shape[1] - input_ids.shape[1],
+                    ),
+                    True,
+                    dtype=attention_mask.dtype,
+                    device=attention_mask.device,
+                )
+                attention_mask = torch.cat(
+                    (new_attn_mask_pad_left, attention_mask), dim=1
+                )
+                assert attention_mask.shape == new_input_embeds.shape[:2]
+
+        return None, attention_mask, past_key_values, new_input_embeds, new_labels
+
+    def initialize_vision_tokenizer(self, model_args, tokenizer):
+        if model_args.mm_use_im_patch_token:
+            tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
+            self.resize_token_embeddings(len(tokenizer))
+
+        if model_args.mm_use_im_start_end:
+            num_new_tokens = tokenizer.add_tokens(
+                [DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True
+            )
+            self.resize_token_embeddings(len(tokenizer))
+
+            if num_new_tokens > 0:
+                input_embeddings = self.get_input_embeddings().weight.data
+                output_embeddings = self.get_output_embeddings().weight.data
+
+                input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(
+                    dim=0, keepdim=True
+                )
+                output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(
+                    dim=0, keepdim=True
+                )
+
+                input_embeddings[-num_new_tokens:] = input_embeddings_avg
+                output_embeddings[-num_new_tokens:] = output_embeddings_avg
+
+            if model_args.tune_mm_mlp_adapter:
+                for p in self.get_input_embeddings().parameters():
+                    p.requires_grad = True
+                for p in self.get_output_embeddings().parameters():
+                    p.requires_grad = False
+
+        elif model_args.mm_use_im_patch_token:
+            if model_args.tune_mm_mlp_adapter:
+                for p in self.get_input_embeddings().parameters():
+                    p.requires_grad = False
+                for p in self.get_output_embeddings().parameters():
+                    p.requires_grad = False

inference/model/multimodal_encoder/clip_encoder.py

@@ -0,0 +1,89 @@
+from abc import ABC
+
+import torch
+import torch.nn as nn
+
+from transformers import CLIPPreTrainedModel, CLIPVisionConfig
+from transformers.models.clip.modeling_clip import CLIPVisionTransformer
+from inference.model.language_model.configuration_llava_phi import LlavaPhiVisionConfig
+
+
+class CLIPVisionTower(CLIPPreTrainedModel):
+    config_class = LlavaPhiVisionConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.vision_model = CLIPVisionTransformer(config)
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.vision_model.embeddings.patch_embedding
+
+    def feature_select(self, image_forward_outs):
+        image_features = image_forward_outs.hidden_states[
+            self.config.mm_vision_select_layer
+        ]
+        if self.config.mm_vision_select_feature == "patch":
+            image_features = image_features[:, 1:]
+        elif self.config.mm_vision_select_feature == "cls_patch":
+            image_features = image_features
+        else:
+            raise ValueError(
+                f"Unexpected select feature: {self.config.mm_vision_select_feature}"
+            )
+        return image_features
+
+    def forward(self, images):
+        if type(images) is list:
+            image_features = []
+            for image in images:
+                image_forward_out = self.vision_model(
+                    image.to(device=self.device, dtype=self.dtype).unsqueeze(0),
+                    output_hidden_states=True,
+                )
+                image_feature = self.feature_select(image_forward_out).to(image.dtype)
+                image_features.append(image_feature)
+        else:
+            image_forward_outs = self.vision_model(
+                images.to(device=self.device, dtype=self.dtype),
+                output_hidden_states=True,
+            )
+            image_features = self.feature_select(image_forward_outs).to(images.dtype)
+
+        return image_features
+
+    @property
+    def dummy_feature(self):
+        return torch.zeros(1, self.hidden_size, device=self.device, dtype=self.dtype)
+
+    @property
+    def dtype(self):
+        return list(self.vision_model.parameters())[0].dtype
+
+    @property
+    def device(self):
+        return list(self.vision_model.parameters())[0].device
+
+    @property
+    def hidden_size(self):
+        return self.config.hidden_size
+
+    @property
+    def num_patches(self):
+        return (self.config.image_size // self.config.patch_size) ** 2
+
+
+if __name__ == "__main__":
+    clip_config = CLIPVisionConfig.from_pretrained(
+        "/data/private/zhumj/GPTcode/mm-phi/openai/clip-vit-large-patch14-336"
+    )
+    print("################ clip_config ##############")
+    print(clip_config)
+    phi_vis_config = LlavaPhiVisionConfig(**clip_config.to_dict())
+    print("################ phi_vis_config ##############")
+    print(phi_vis_config)
+
+    model = CLIPVisionTower(clip_config)
+    # print(list(model.vision_model.parameters())[0].dtype)

inference/model/multimodal_projector/builder.py

@@ -0,0 +1,50 @@
+import torch
+import torch.nn as nn
+import re
+
+
+class IdentityMap(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, *args, **kwargs):
+        return x
+
+    @property
+    def config(self):
+        return {"mm_projector_type": "identity"}
+
+
+class SimpleResBlock(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.pre_norm = nn.LayerNorm(channels)
+
+        self.proj = nn.Sequential(
+            nn.Linear(channels, channels), nn.GELU(), nn.Linear(channels, channels)
+        )
+
+    def forward(self, x):
+        x = self.pre_norm(x)
+        return x + self.proj(x)
+
+
+def build_vision_projector(config):
+    projector_type = getattr(config, "mm_projector_type", "linear")
+
+    if projector_type == "linear":
+        return nn.Linear(config.mm_hidden_size, config.hidden_size)
+
+    mlp_gelu_match = re.match(r"^mlp(\d+)x_gelu$", projector_type)
+    if mlp_gelu_match:
+        mlp_depth = int(mlp_gelu_match.group(1))
+        modules = [nn.Linear(config.mm_hidden_size, config.hidden_size)]
+        for _ in range(1, mlp_depth):
+            modules.append(nn.GELU())
+            modules.append(nn.Linear(config.hidden_size, config.hidden_size))
+        return nn.Sequential(*modules)
+
+    if projector_type == "identity":
+        return IdentityMap()
+
+    raise ValueError(f"Unknown projector type: {projector_type}")