Upload 5 files

configs.py

@@ -0,0 +1,36 @@
+import torch
+import multiprocessing
+
+def get_config_phase1():
+    return {
+        "data_dir": "./data",
+        "clip_model_name": "openai/clip-vit-base-patch16",
+        "phi2_model_name": "microsoft/phi-2",
+        "train_batch_size": 2,
+        "val_batch_size": 1,
+        "device": torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+        "epochs": 2,
+        "max_tokens": 20,
+        "clip_embed": 768,
+        "phi_embed": 2560,
+        "num_workers": 32, 
+        "ckpts": "./ckpts"
+    }
+
+def get_config_phase2():
+    return {
+        "i150k_json": "./data/llava_instruct_150k.json",
+        "QA_datasetName": "OpenAssistant/oasst1",
+        "clip_model_name": "openai/clip-vit-base-patch16",
+        "phi2_model_name": "microsoft/phi-2",
+        "train_batch_size": 1,
+        "val_batch_size": 1,
+        "device": torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+        "epochs": 2,
+        "max_tokens": 20,
+        "clip_embed": 768,
+        "phi_embed": 2560,
+        "num_workers": 1, 
+        "ckpts": "./ckpts",
+        "vocab_size": 51200
+    }

dataset.py

@@ -0,0 +1,204 @@
+import os
+import json
+import torch
+from PIL import Image
+from torch.utils.data import Dataset
+from transformers import AutoProcessor
+from torch.utils.data import DataLoader
+import pickle
+import requests
+from datasets import Dataset, load_dataset
+import pandas as pd
+import numpy as np
+
+
+class ClipDataset(Dataset):
+  '''ClipDataset class for loading the CLIP dataset'''
+  def __init__(self, coco_data, model_name, tokenizer):
+
+    self.tokenizer  = tokenizer
+    self.processor  = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
+    self.caption_dataset = coco_data
+      
+  def __len__(self):
+    #Return the length of the dataset
+    return len(self.caption_dataset)
+
+  def __getitem__(self, idx):
+    #Get the image url and caption
+    img_url = self.caption_dataset[idx]["image_url"]
+    caption = self.caption_dataset[idx]["caption"]
+
+    #Get the image and caption embeddings
+    image = Image.open(requests.get(img_url,stream=True).raw)
+    width, height = image.size
+    new_width  = 224
+    new_height = new_width * height // width 
+    new_height = 224
+    new_width  = new_height * width // height
+    image = image.resize((new_width, new_height), Image.LANCZOS)
+    image_processed = self.processor(images=image, return_tensors="pt") ['pixel_values']
+    image_sqeezed = image_processed.squeeze(0)
+    tokenized_caption = self.tokenizer(caption, return_tensors="pt", return_attention_mask=False)
+    tokenized_caption_ids = tokenized_caption['input_ids'].squeeze(0)
+    return(image_sqeezed , tokenized_caption_ids)
+  
+
+def collate_fn_phase1(batch):
+    #Unzip the batch
+    image_embeddings, captions = zip(*batch)
+    #Stack the image embeddings
+    image_embeddings_stacked = torch.stack(image_embeddings, dim=0)
+    #Pad the captions, padded value is the <eos> token
+    captions_padded = torch.nn.utils.rnn.pad_sequence(captions, batch_first=True, padding_value=50256)
+    #Return the stacked image embeddings and padded captions
+    return (image_embeddings_stacked, captions_padded)
+
+
+def get_data_loaders_phase1(data_dir, clip_model_name, tokenizer, train_batch_size, val_batch_size, num_workers):
+    # Load the data
+    with open(os.path.join(data_dir, 'coco_train.pkl'), 'rb') as fp: 
+        train_pkl = pickle.load(fp)
+    with open(os.path.join(data_dir, "coco_val.pkl"), "rb") as fp:
+        val_pkl = pickle.load(fp)
+   # train data loaders
+    train_dataloader = DataLoader(ClipDataset(train_pkl, clip_model_name, tokenizer), collate_fn=collate_fn_phase1, batch_size=train_batch_size, num_workers = num_workers, shuffle=True, pin_memory=True)
+
+    # val data loaders
+    val_dataloader   = DataLoader(ClipDataset(val_pkl, clip_model_name, tokenizer), collate_fn=collate_fn_phase1, batch_size=val_batch_size, num_workers = num_workers, shuffle=False, pin_memory=True)
+    return train_dataloader, val_dataloader
+
+##################################### Phase 2 #########################################
+
+  
+class ClipDatasetPhase2(Dataset):
+  '''ClipDataset class for loading the CLIP dataset'''
+  def __init__(self, data_frame, model_name, tokenizer):
+
+    self.tokenizer  = tokenizer
+    self.processor  = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
+    self.df = data_frame
+      
+  def __len__(self):
+    #Return the length of the dataset
+    return len(self.df)
+
+  def __getitem__(self, idx):
+    #Get the image url and QAs
+    img_url = self.df.ImageUrl[idx[0]]
+    que = self.df.Question[idx[0]]
+    ans = self.df.Answer[idx[0]]
+
+    print("img_url", img_url)
+    print("que", que)
+    print("ans", ans)
+
+    #Get the image and caption embeddings
+    if img_url is None:
+        print("img_url is None")
+        image_sqeezed = None
+    else:
+        image = Image.open(requests.get(img_url,stream=True).raw)
+        width, height = image.size
+        new_width  = 224
+        new_height = new_width * height // width 
+        new_height = 224
+        new_width  = new_height * width // height
+        image = image.resize((new_width, new_height), Image.LANCZOS)
+        image_processed = self.processor(images=image, return_tensors="pt") ['pixel_values']
+        image_sqeezed = image_processed.squeeze(0)
+    que_ids = self.tokenizer(que, return_tensors="pt", return_attention_mask=False)['input_ids'].squeeze(0)
+    ans_ids = self.tokenizer(ans, return_tensors="pt", return_attention_mask=False)['input_ids'].squeeze(0)
+    return(image_sqeezed , que_ids, ans_ids)
+  
+
+def collate_fn_phase2(batch):
+    #Unzip the batch
+    image_embeddings, ques, ans = zip(*batch)
+    #Stack the image embeddings
+    if image_embeddings[0] is None:
+        image_embeddings_stacked = None
+    else:
+        image_embeddings_stacked = torch.stack(image_embeddings, dim=0)
+    #Pad the QAs, padded value is the <eos> token
+    ques_padded = torch.nn.utils.rnn.pad_sequence(ques, batch_first=True, padding_value=50256)
+    ans_padded = torch.nn.utils.rnn.pad_sequence(ans, batch_first=True, padding_value=50256)
+    #Return the stacked image embeddings and padded QAs
+    return (image_embeddings_stacked, ques_padded, ans_padded)
+
+
+def prep_data(df):
+    df_assistant = df[(df.role == "assistant") & (df["rank"] == 0.0)].copy()
+    df_prompter = df[(df.role == "prompter")].copy()
+    df_prompter = df_prompter.set_index("message_id")
+    df_assistant["Answer"] = df_assistant["text"].values
+
+    inputs = []
+    for _, row in df_assistant.iterrows():
+        input = df_prompter.loc[row.parent_id]
+        inputs.append(input.text)
+
+    df_assistant["Question"] = inputs
+    df_assistant["ImageUrl"] = None
+
+    df_assistant = df_assistant[df_assistant.lang == "en"]
+
+    df_assistant = df_assistant[
+        ["ImageUrl","Question", "Answer", "message_id"]
+    ].rename(columns={"message_id": "Ids"})
+
+    return df_assistant
+
+
+def get_i150_df(config):
+    with open(config.get("i150k_json"), "r") as fp: 
+        i150k_json_read = json.load(fp)
+    max_tokens = 100
+    image_urls = []
+    ques_list = []
+    ans_list = []
+    id_list = []
+    for idx, data in enumerate(i150k_json_read):
+        image = data['image']
+        image_url = 'http://images.cocodataset.org/train2017/' + image
+        id_ = data["id"]
+        iterator = iter(data['conversations'])
+        for i in iterator:
+            ques = i
+            ans = next(iterator)
+            if (len(ques["value"])>100 or len(ans["value"])>max_tokens):
+                continue
+            if ques["from"] == "human" and ans["from"] == "gpt":
+                image_urls.append(image_url)
+                ques_list.append(ques["value"].replace("<image>\n","").replace("<image>",""))
+                ans_list.append(ans["value"])
+                id_list.append(id_)
+    df_i150k = pd.DataFrame(list(zip(image_urls, ques_list, ans_list, id_list)),
+                  columns =["ImageUrl", "Question", "Answer", "Ids"])
+    msk = np.random.rand(len(df_i150k)) < 0.96
+
+    train_df = df_i150k[msk]
+    test_df = df_i150k[~msk]
+    return train_df, test_df
+
+
+def get_oas_df(config):
+    train_ds, val_ds = load_dataset(config.get("QA_datasetName"), split=["train", "validation"])
+    train_df = prep_data(train_ds.to_pandas())
+    test_df = prep_data(val_ds.to_pandas())
+    return train_df, test_df
+
+
+def get_data_loaders_phase2(tokenizer, config):
+
+    train_i150k, test_i150k = get_i150_df(config)
+    train_oas, test_oas = get_oas_df(config)
+
+    train_df = pd.concat([train_i150k, train_oas]).reset_index(drop=True)
+    val_df = pd.concat([test_i150k, test_oas]).reset_index(drop=True)
+   # train data loaders
+    train_dataloader = DataLoader(ClipDatasetPhase2(train_df, config.get("clip_model_name"), tokenizer), collate_fn=collate_fn_phase2, batch_size=config.get("train_batch_size"), num_workers = config.get("num_workers"), shuffle=True, pin_memory=True)
+
+    # val data loaders
+    val_dataloader   = DataLoader(ClipDatasetPhase2(val_df, config.get("clip_model_name"), tokenizer), collate_fn=collate_fn_phase2, batch_size=config.get("val_batch_size"), num_workers = config.get("num_workers"), shuffle=False, pin_memory=True)
+    return train_dataloader, val_dataloader

get_coco.py

@@ -0,0 +1,41 @@
+import os, shutil, json
+import pickle, argparse
+
+"""Unzip the data and and save it as a pickle file."""
+
+def make_pkl(data_dir, dataset_json, train_flag=False):
+    coco_data_list = []
+    for i, data in enumerate(dataset_json['annotations']):
+        image_id = data['image_id']
+        caption  = data['caption']
+        for img in dataset_json['images']:
+            if img['id'] == image_id:
+                image_url = img['coco_url']
+                file_name = img['file_name']
+                break
+        coco_data_list.append({'image_id': image_id,'image_url': image_url, 'file_name': file_name, 'caption': caption})
+    if train_flag:
+        with open(os.path.join(data_dir, f'coco_train.pkl'), 'wb') as f:
+            pickle.dump(coco_data_list, f)
+    else:
+        with open(os.path.join(data_dir, f'coco_val.pkl'), 'wb') as f:
+            pickle.dump(coco_data_list, f)
+
+
+def main(coco_path, data_dir):
+    coco_dir = os.path.dirname(coco_path)
+    # shutil.unpack_archive(coco_path, coco_dir)
+    with open(os.path.join(coco_dir, 'annotations/captions_train2017.json')) as f:
+        coco_train_dataset = json.load(f)
+    with open(os.path.join(coco_dir, 'annotations/captions_val2017.json')) as f:
+        coco_val_dataset = json.load(f)
+    make_pkl(data_dir, coco_train_dataset, train_flag=True)
+    # make_pkl(data_dir, coco_val_dataset)
+    
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--coco_path', type=str, default='coco.zip')
+    parser.add_argument('--data_dir', type=str, default='data')
+    args = parser.parse_args()
+    main(args.coco_path, args.data_dir)

main.py

@@ -0,0 +1,41 @@
+import torch
+from dataset import get_data_loaders_phase1, get_data_loaders_phase2
+from transformers import AutoTokenizer
+from model import CustomClipPhi2, MainQLoraModel, train_model_phase1, train_model_phase2
+from configs import get_config_phase1, get_config_phase2
+
+def phase_1():   
+    # get config
+    config = get_config_phase1() 
+    # tokenizer
+    tokenizer  = AutoTokenizer.from_pretrained(config.get("phi2_model_name"), trust_remote_code=True)
+
+    # data loaders
+    train_dataloader, val_dataloader = get_data_loaders_phase1(config.get("data_dir"), config.get("clip_model_name"), tokenizer, config.get("train_batch_size"), config.get("val_batch_size"), config.get("num_workers"))
+
+    llmModel = CustomClipPhi2(tokenizer, config.get("phi2_model_name"), config.get("clip_model_name"), clip_embed=768, phi_embed=2560).to(config.get("device"))
+    print(llmModel)
+    # optimizer
+    optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, llmModel.parameters()), lr=1e-3)
+    # train model
+    train_model_phase1(llmModel, train_dataloader, val_dataloader, optimizer, tokenizer, config)
+
+
+def phase_2():   
+    # get config
+    config = get_config_phase2() 
+    # tokenizer
+    tokenizer  = AutoTokenizer.from_pretrained(config.get("phi2_model_name"), trust_remote_code=True)
+
+    # data loaders
+    train_dataloader, val_dataloader = get_data_loaders_phase2(tokenizer, config)
+
+    llmModel = MainQLoraModel(tokenizer, config).to(config.get("device"))
+    print(llmModel)
+    # train model
+    train_model_phase2(llmModel, train_dataloader, val_dataloader, tokenizer, config)
+
+if __name__ == "__main__":
+    torch.set_float32_matmul_precision('medium')
+    phase_1()
+    # phase_2()

model.py

@@ -0,0 +1,336 @@
+import torch
+import torch.nn as nn
+from torch.nn.functional import cross_entropy
+from transformers import CLIPVisionModel, AutoModelForCausalLM, BitsAndBytesConfig
+from peft import LoraConfig
+from tqdm import tqdm
+import os, peft
+
+
+class CustomClipPhi2(nn.Module):
+    def __init__(self,tokenizer, phi2_model_name, clip_model_name, clip_embed=768, phi_embed=2560):
+        super().__init__()
+
+        self.tokenizer = tokenizer
+        # These two models are not finetuned
+        # pretrained Microsoft phi2 model
+        self.phi2_model = AutoModelForCausalLM.from_pretrained(phi2_model_name,torch_dtype=torch.float32, trust_remote_code=True)
+        # pretrained OpenAI clip model
+        self.clip_model = CLIPVisionModel.from_pretrained(clip_model_name)
+
+        self.EOS_TOKEN_ID    = self.tokenizer.eos_token_id # 50256
+        self.IMAGE_TOKEN_ID  = 23903 # token for Comments
+        self.clip_embed      = clip_embed
+        self.phi_embed       = phi_embed        
+
+        # projection layers
+        # Trainable projection layer
+        self.projection_layer = torch.nn.Linear(clip_embed, phi_embed)
+
+        # Freeze Weights
+        for models in [self.phi2_model, self.clip_model]:
+            for param in models.parameters():
+                param.requires_grad_(False)
+
+        # load checkpoint weights
+        if os.path.exists('./ckpts/model_phase1.pth'):
+            self.projection_layer.load_state_dict(torch.load('./ckpts/model_phase1.pth', map_location='cpu'))
+            print("Loaded checkpoint weights for projection layer")
+        else:
+            print("No checkpoint weights for projection layer")
+            print("Initializing projection layer with random weights")
+            self.projection_layer.weight.data.normal_(mean=0.0, std=0.02)
+            self.projection_layer.bias.data.zero_()
+
+
+    def generate(self, images, tokenizer, config):
+        clip_outputs = self.clip_model(**images)
+        # remove cls token
+        images = clip_outputs.last_hidden_state[:, 1:, :]
+        image_embeddings = self.projection_layer(images).to(torch.float16)
+
+        batch_size = images.size()[0]
+        predicted_caption = torch.full((batch_size, config.get("max_tokens")), self.EOS_TOKEN_ID, dtype=torch.long, device=config.get('device'))
+        img_token_tensor = torch.tensor(self.IMAGE_TOKEN_ID).repeat(batch_size, 1)
+        img_token_embeds = self.phi2_model.model.embed_tokens(img_token_tensor.to(image_embeddings.device))
+        combined_embeds  = torch.cat([image_embeddings, img_token_embeds], dim=1)
+
+        for pos in range(config.get("max_tokens") - 1):
+            model_output_logits = self.phi2_model.forward(inputs_embeds = combined_embeds)['logits']
+            predicted_word_token_logits = model_output_logits[:, -1, :].unsqueeze(1)
+            predicted_word_token = torch.argmax(predicted_word_token_logits, dim = -1)
+            predicted_caption[:, pos] = predicted_word_token.view(1,-1).to('cpu')
+            next_token_embeds = self.phi2_model.model.embed_tokens(predicted_word_token)
+            combined_embeds   = torch.cat([combined_embeds, next_token_embeds], dim=1)
+        return predicted_caption
+
+
+    def forward(self, images, target_captions):
+
+        batch_size    = target_captions.size()[0]
+        target_length = target_captions.size()[1]
+
+        # clip model output for image
+        clip_outputs = self.clip_model(**images) # See this for loading https://huggingface.co/openai/clip-vit-base-patch36
+        images = clip_outputs.last_hidden_state[:, 1:, :] # remove CLS token
+
+        # projection layer
+        image_embeddings = self.projection_layer(images).to(torch.float16)
+
+        # add comment token from phi2
+        img_token_tensor = torch.tensor(self.IMAGE_TOKEN_ID).repeat(batch_size, 1)
+        img_token_embeds = self.phi2_model.model.embed_tokens(img_token_tensor.to(image_embeddings.device))
+        combined_embeds  = torch.cat([image_embeddings, img_token_embeds], dim=1) # 4,49,2560
+        del clip_outputs
+        del image_embeddings
+
+        # for loss
+        loss = 0
+        for pos in range(target_length - 1):
+           
+            model_output_logits = self.phi2_model.forward(inputs_embeds = combined_embeds)['logits']
+            predicted_word_token_logits = model_output_logits[:, -1, :].unsqueeze(1)
+            pos_loss = cross_entropy(predicted_word_token_logits.view(-1,predicted_word_token_logits.size(-1)), target_captions[:, pos].contiguous().view(-1), ignore_index=self.EOS_TOKEN_ID,label_smoothing=0.1)
+            loss += pos_loss
+
+            predicted_word_token = torch.argmax(predicted_word_token_logits, dim=-1)
+            next_token_embeds = self.phi2_model.model.embed_tokens(predicted_word_token) 
+            combined_embeds   = torch.cat([combined_embeds, next_token_embeds], dim=1)
+        loss = loss / target_length
+
+        # Delete variables to free up memory
+        del combined_embeds
+        del model_output_logits
+        torch.cuda.empty_cache()
+
+        return loss
+        
+
+def show_results_for_samples_phase1(model, val_dataloader, tokenizer, config, num_samples = 2):
+    model.eval()
+    with torch.no_grad():
+        for i in range(num_samples):
+            for images, target_captions in val_dataloader:
+                images = {'pixel_values': images.to(config.get('device'))}
+                target_captions = target_captions.to(config.get('device'))
+                target_captions_decoded = tokenizer.batch_decode(target_captions, ignore_index = tokenizer.eos_token_id)
+                predicted_captions = model.generate(images,  tokenizer, config)
+                predicted_captions_decoded = tokenizer.batch_decode(predicted_captions,ignore_index = tokenizer.eos_token_id)
+
+                for idx, pc in enumerate(predicted_captions_decoded):
+                    print(f"{idx} - Target captions: {target_captions_decoded[idx]} \n {'---------------------'*10} \n Predicted_captions:{pc} ")
+                break
+
+
+def validate_model_phase1(model, val_dataloader, tokenizer, config):
+    model.eval()
+    total_loss = 0
+    with torch.no_grad():
+        try:
+            for images, target_captions in tqdm(val_dataloader):
+                images = {'pixel_values': images.to(config.get('device'))}
+                target_captions = target_captions.to(config.get('device'))
+                loss = model(images, target_captions)
+                total_loss+=loss.item()
+            print(f"Validation Loss: {total_loss/len(val_dataloader)}")
+        except Exception as e:
+            pass
+    model.train()
+
+    
+def train_model_phase1(model, train_loader, val_dataloader, optimizer, tokenizer, config):
+    model.train()
+
+    pbar = tqdm(train_loader)
+    for epoch in range(1, config.get("epochs")):
+        print(f"Epoch: {epoch}")
+        torch.cuda.empty_cache()
+        step = 1
+        try:
+            for idx, (images, target_captions) in enumerate(pbar):
+                try:
+                    if target_captions.shape[1] >= config.get("max_tokens"):
+                        # print(f"Skipping batch {idx} due to long caption")
+                        continue 
+        
+                    images = {'pixel_values': images.to(config.get('device'))}
+                    target_captions = target_captions.to(config.get('device'))
+        
+                    optimizer.zero_grad()
+                    loss = model(images, target_captions)
+                    loss.backward()
+                    optimizer.step()
+                    pbar.set_description(f"Epoch: {epoch}: Training Loss = {loss.item()}")
+                    torch.cuda.empty_cache()
+                    step+=1
+                    if (step%1000==0):
+                        torch.save(model.projection_layer.state_dict(), './ckpts/model_phase1.pth')
+                except Exception as e:
+                    continue
+                 
+            # # save model
+            # if ((epoch % 2) == 0):
+                # Only save last checkpoint
+            validate_model_phase1(model, val_dataloader, tokenizer, config)
+            show_results_for_samples_phase1(model, val_dataloader, tokenizer, config)
+            torch.save(model.projection_layer.state_dict(), './ckpts/model_phase1.pth')
+
+        except Exception as e:
+            continue
+
+
+
+
+######################################## Phase 2 #########################################
+
+class MainQLoraModel(nn.Module):
+    def __init__(self, tokenizer, config):
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.config = config
+        self.clip_model = CLIPVisionModel.from_pretrained(config.get("clip_model_name"))
+
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.float16,
+        )
+
+        phi2_model = AutoModelForCausalLM.from_pretrained(
+            config.get("phi2_model_name"),
+            quantization_config=bnb_config,
+            trust_remote_code=True
+        )
+        phi2_model.config.use_cache = False
+
+        ## 4 - LORA config
+
+        lora_alpha = 16
+        lora_dropout = 0.1
+        lora_r = 64
+
+        peft_config = LoraConfig(
+            lora_alpha = lora_alpha,
+            lora_dropout = lora_dropout,
+            r = lora_r,
+            bias="none",
+            task_type="CAUSAL_LM",
+            target_modules=[
+                "q_proj",
+                "k_proj",
+                "v_proj",
+                "dense",
+                "fc1",
+                "fc2"
+            ]
+        )
+        self.phi2_model = peft.get_peft_model(phi2_model, peft_config).to(config.get("device"))
+
+        self.EOS_TOKEN_ID    = self.tokenizer.eos_token_id
+        self.IMAGE_TOKEN_ID  = 23903 # token for Comments
+        self.clip_embed      = config.get("clip_embed")
+        self.phi_embed       = config.get("phi_embed")        
+
+        # projection layers
+        # Trainable projection layer
+        self.projection_layer = torch.nn.Linear(self.clip_embed, self.phi_embed)
+
+        # Freeze Weights
+        for models in [self.clip_model]:
+            for param in models.parameters():
+                param.requires_grad_(False)
+
+        # load checkpoint weights
+        if os.path.exists('./ckpts/model_phase2.pth'):
+            self.projection_layer.load_state_dict(torch.load('./ckpts/model_phase2.pth', map_location=config.get("device")))
+            self.phi2_model.from_pretrained(self.phi2_model,'./ckpts/Qlora_adaptor')
+            print("Loaded checkpoint weights for projection layer")
+        else:
+            # Load weights from phase 1
+            self.projection_layer.load_state_dict(torch.load('./ckpts/model_phase1.pth', map_location=config.get("device")))
+
+
+    def forward(self, images, ques, ans):
+
+        batch_size = ques.size()[0]
+        questions  = ques.to(self.config.get("device"))
+        answers    = ans.to(self.config.get("device"))
+
+        questions_embed  = peft_model.model.model.embed_tokens(questions)
+        if images is None:
+            combined_embeds = questions_embed
+        else:
+            images = {'pixel_values': images.to(self.config.get("device"))}
+            clip_outputs  = clip_model(**images)
+            images_embeds = clip_outputs.last_hidden_state[:,1:,:] # remove cls token
+            
+            # projection
+            image_embeds  = projection(images_embeds).to(torch.float16)
+            img_token_tensor = torch.tensor(self.IMAGE_TOKEN_ID).repeat(batch_size, 1).to(self.config.get("device"))
+            img_token_embeds = peft_model.model.model.embed_tokens(img_token_tensor)
+            combined_embeds = torch.cat([image_embeds, img_token_embeds, questions_embed], dim=1) 
+        
+        phi_output_logits = peft_model(inputs_embeds=combined_embeds)['logits']
+
+        if images is not None:
+            # remove image and image token embeddings
+            phi_output_logits = phi_output_logits[:,images_embeds.shape[1] + 2 : ,:]
+        
+        phi_output_logits = phi_output_logits.reshape(-1, self.config.get("vocab_size"))
+        
+        loss = cross_entropy(phi_output_logits, answers.contiguous().view(-1), ignore_index=self.EOS_TOKEN_ID, label_smoothing=0.1)
+
+        return loss
+
+def validate_model_phase2(model, val_dataloader, tokenizer, config):
+    model.eval()
+    total_loss = 0
+    with torch.no_grad():
+        try:
+            for images, ques, ans in tqdm(val_dataloader):
+                loss = model(images, ques, ans)
+                total_loss+=loss.item()
+            print(f"Validation Loss: {total_loss/len(val_dataloader)}")
+        except Exception as e:
+            pass
+    model.train()
+
+
+def train_model_phase2(model, train_loader, val_dataloader, tokenizer, config):
+    phi2_optim = torch.optim.Adam(filter(lambda p: p.requires_grad, model.phi2_model.parameters()), lr=1e-5)
+    proj_optim = torch.optim.Adam(filter(lambda p: p.requires_grad, model.projection_layer.parameters()), lr=1e-5)
+    model.phi2_model.train()
+    model.projection_layer.train()
+
+    pbar = tqdm(train_loader)
+    for epoch in range(1, config.get("epochs")):
+        print(f"Epoch: {epoch}")
+        torch.cuda.empty_cache()
+        step = 1
+        try:
+            for idx, (images, ques, ans) in enumerate(pbar):
+                try:
+                    print("hi")
+                    phi2_optim.zero_grad()
+                    proj_optim.zero_grad()
+                    loss = model(images, ques, ans)
+                    loss.backward()
+                    phi2_optim.step()
+                    proj_optim.step()
+                    pbar.set_description(f"Epoch: {epoch}: Training Loss = {loss.item()}")
+                    torch.cuda.empty_cache()
+                    step+=1
+                    if (step%1000==0):
+                        torch.save(model.projection_layer.state_dict(), './ckpts/model_phase2.pth')
+                        model.phi2_model.save_pretrained('./ckpts/Qlora_adaptor/', save_adapter=True, save_config=True)
+                except Exception as e:
+                    print(e)
+                    continue
+                 
+            validate_model_phase2(model, val_dataloader, tokenizer, config)
+            torch.save(model.projection_layer.state_dict(), './ckpts/model_phase2.pth')
+            model.phi2_model.save_pretrained('./ckpts/Qlora_adaptor/', save_adapter=True, save_config=True)
+
+        except Exception as e:
+            print(e)
+            continue