new files

config.py

@@ -0,0 +1,173 @@
+import torch
+from transformers import PretrainedConfig, BitsAndBytesConfig
+import math
+from typing import Optional
+
+class VisionProjectorConfig(PretrainedConfig):
+    def __init__(
+            self,
+            input_dim=768,
+            hidden_dim=256,
+            num_tokens=1,
+            output_dim=2560,
+            **kwargs
+    ):
+        #super.__init__(**kwargs)
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.output_dim = output_dim
+        self.num_tokens = num_tokens
+        self.kwargs = kwargs
+
+
+class CustomPhiConfig(PretrainedConfig):
+    model_type = "phi-msft"
+    attribute_map = {
+        "max_position_embeddings": "n_positions",
+        "hidden_size": "n_embd",
+        "num_attention_heads": "n_head",
+        "num_hidden_layers": "n_layer",
+    }
+
+    def __init__(
+            self,
+            vocab_size: int = 51200,
+            n_positions: int = 2048,
+            n_embd: int = 2560,
+            n_layer: int = 32,
+            n_inner: Optional[int] = None,
+            n_head: int = 32,
+            n_head_kv: Optional[int] = None,
+            rotary_dim: Optional[int] = 32,
+            activation_function: Optional[str] = "gelu_new",
+            flash_attn: bool = False,
+            flash_rotary: bool = False,
+            fused_dense: bool = False,
+            attn_pdrop: float = 0.0,
+            embd_pdrop: float = 0.0,
+            resid_pdrop: float = 0.1,
+            layer_norm_epsilon: float = 1e-05,
+            initializer_range: float = 0.02,
+            tie_word_embeddings: bool = False,
+            pad_vocab_size_multiple: int = 64,
+            **kwargs
+    ) -> None:
+        self.vocab_size = int(math.ceil(vocab_size / pad_vocab_size_multiple) * pad_vocab_size_multiple)
+        self.n_positions = n_positions
+        self.n_embd = n_embd
+        self.n_layer = n_layer
+        self.n_inner = n_inner
+        self.n_head = n_head
+        self.n_head_kv = n_head_kv
+        self.rotary_dim = min(rotary_dim, n_embd // n_head)
+        self.activation_function = activation_function
+        self.flash_attn = flash_attn
+        self.flash_rotary = flash_rotary
+        self.fused_dense = fused_dense
+        self.attn_pdrop = attn_pdrop
+        self.embd_pdrop = embd_pdrop
+        self.resid_pdrop = resid_pdrop
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.initializer_range = initializer_range
+
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+
+
+class CLIPVisionToPhiConfig(PretrainedConfig):
+    def __init__(self,
+                 vision_projector_config: VisionProjectorConfig,
+                 phi_config: CustomPhiConfig,
+                 **kwargs
+                 ):
+
+        #super().__init__(**kwargs)
+
+        self.vision_projector_config = vision_projector_config
+        self.phi_config = phi_config
+        self.tokenizer = kwargs.get('tokenizer')
+        self.freeze_phi_model = True
+
+
+'''
+phi_config_obj = CustomPhiConfig(
+    **{
+      "_name_or_path": "microsoft/phi-2",
+      "architectures": [
+        "PhiForCausalLM"
+      ],
+      "auto_map": {
+        "AutoConfig": "configuration_phi.PhiConfig",
+        "AutoModelForCausalLM": "modeling_phi.PhiForCausalLM"
+      },
+      "img_processor": None,
+      "model_type": "phi-msft",
+      "torch_dtype": "float16",
+      "transformers_version": "4.35.2"
+    }
+
+)
+
+'''
+from peft import LoraConfig
+
+bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.float16
+)
+
+peft_config = LoraConfig(
+    lora_alpha=16,
+    lora_dropout=0.1,
+    r=64,
+    bias="none",
+    task_type="CAUSAL_LM",
+    target_modules=[
+        "q_proj",
+        "k_proj",
+        "v_proj",
+        "dense",
+        "fc1",
+        "fc2"
+    ]
+)
+
+class MultiInstructModelConfig(PretrainedConfig):
+    def __init__(self,
+                 vision_projector_config: Optional[VisionProjectorConfig] = None,
+                 **kwargs
+                 ):
+
+        self.vision_projector_config = vision_projector_config
+        self.quantization_config = bnb_config
+
+        self.peft_config = peft_config
+
+        self.tokenizer = kwargs.get('tokenizer')
+        self.freeze_vision_projector = True
+
+
+extra = dict(
+    num_epochs=1,
+    resume=False,
+    data_dir='../data',
+    checkpoint_dir='../checkpoints',
+    max_seqlen=80,
+    batch_size=2,
+    live_image_processing=True,
+    vision_projector_file='/Users/piyushgrover/Downloads/old_vt_proj/vp_ckpt_0.pth',
+    validation_phase=False
+)
+
+qlora_config = dict(
+    num_steps=1000,
+    max_seqlen=512,
+    max_caption_len=100,
+    batch_size=8,
+    micro_batch_size=2,
+    data_dir='../data',
+    output_dir="./results",
+    vision_model=True,
+    vision_projector_file='models/vision_projector/vp_ckpt_0.pth',
+    resume=False
+)

constants.py

@@ -0,0 +1,2 @@
+IGNORE_INDEX = -100
+IMAGE_TOKEN_INDEX = -200

models/vision_projector_model.py

@@ -0,0 +1,44 @@
+import torch
+import torch.nn as nn
+
+from config import VisionProjectorConfig
+
+'''
+class VisionProjector(nn.Module):
+
+    def __init__(self, config: VisionProjectorConfig):
+        super().__init__()
+        self.config = config
+        self.input_dim = config.input_dim
+        self.hidden_dim = config.hidden_dim
+        self.output_dim = config.output_dim
+        self.num_tokens = config.num_tokens
+
+        self.pre_norm = nn.LayerNorm(self.input_dim)
+
+        self.proj = nn.Sequential(
+            nn.GELU(),
+            nn.Linear(self.input_dim, self.num_tokens * self.output_dim)
+        )
+
+    def forward(self, x):
+        x = self.pre_norm(x)
+        x = self.proj(x)
+        x = x.reshape( (-1, self.num_tokens, self.output_dim) )
+        return x
+
+'''
+
+class VisionProjector(nn.Module):
+
+    def __init__(self, config: VisionProjectorConfig):
+        super().__init__()
+        self.config = config
+        self.input_dim = config.input_dim
+        self.output_dim = config.output_dim
+
+        self.proj = nn.Linear(self.input_dim, self.output_dim)
+
+    def forward(self, x):
+        x = self.proj(x)
+        return x

utils.py

@@ -0,0 +1,151 @@
+from constants import *
+import torch
+import torch.nn.functional as F
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+def tokenizer_image_token(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 top_k_top_p_filtering(logits, top_k=0, top_p=0.0, filter_value=-float('Inf')):
+    """ Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
+        Args:
+            logits: logits distribution shape (batch size x vocabulary size)
+            top_k > 0: keep only top k tokens with highest probability (top-k filtering).
+            top_p > 0.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
+    """
+    top_k = min(top_k, logits.size(-1))  # Safety check
+    if top_k > 0:
+        # Remove all tokens with a probability less than the last token of the top-k
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits[indices_to_remove] = filter_value
+
+    if top_p > 0.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+        # Remove tokens with cumulative probability above the threshold
+        sorted_indices_to_remove = cumulative_probs > top_p
+        # Shift the indices to the right to keep also the first token above the threshold
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = 0
+
+        indices_to_remove = sorted_indices[sorted_indices_to_remove]
+        logits[indices_to_remove] = filter_value
+    return logits
+
+'''
+def get_image_feature_for_vision_projector(image_url):
+    image_url = 'http://images.cocodataset.org/%s/%s' % (self.directory, self.image_indices_json[image_index])
+
+    image = Image.open(requests.get(image_url, stream=True).raw)
+    inputs = self.processor(images=image, return_tensors="pt")
+    x = self.model(**inputs, output_hidden_states=True)
+    image_feature = x.hidden_states[-2][:, 1:].squeeze(0).cpu().detach()
+'''
+
+
+def generate_output(model, tokenizer, length, input_ids=None, image_features=None, inputs_embeds=None, labels=None,
+                    temperature=1, top_k=0, top_p=0.0):
+    if inputs_embeds is None and (image_features is None or input_ids is None):
+        print("image_features or input_ids missing.. returning")
+        return
+
+    ie_size = inputs_embeds.size(1) - 1
+    inputs = inputs_embeds
+    predicted_tokens = [] #torch.tensor([[]]).to(device)
+
+    label_size = labels.size(1)
+    out = {}
+    if labels is None:
+        with torch.no_grad():
+            for idx in range(length):
+                outputs = model.phi_model(inputs_embeds=inputs)
+                logits = outputs['logits']
+                next_token_logits = logits[:, -1, :] / temperature  # Apply temperature
+
+                filtered_logits = top_k_top_p_filtering(next_token_logits, top_k=top_k,
+                                                        top_p=top_p)  # Apply top-k and/or top-p
+                next_token = torch.multinomial(F.softmax(filtered_logits, dim=-1), num_samples=1)  # Sample
+
+                predicted_tokens.append(next_token)
+                next_token_embed = model.text_embedding(next_token)
+                inputs = torch.cat((inputs, next_token_embed), dim=1)
+
+            predicted_tokens = torch.cat([x.unsqueeze(1) for x in predicted_tokens], dim=1)
+            out['pred'] = predicted_tokens
+            out['logits'] = logits[:, ie_size:, :]
+
+            return out
+    else:
+            # traverse_len = labels.size(1) - inputs_embeds.size(1)
+        for idx in range(length):
+            outputs = model.phi_model(inputs_embeds=inputs)
+            logits = outputs['logits']
+            next_token_logits = logits[:, -1, :] / temperature  # Apply temperature
+
+            filtered_logits = top_k_top_p_filtering(next_token_logits, top_k=top_k,
+                                                    top_p=top_p)  # Apply top-k and/or top-p
+            next_token = torch.multinomial(F.softmax(filtered_logits, dim=-1), num_samples=1)  # Sample
+
+            predicted_tokens.append(next_token)
+
+            tf_token = labels[:, idx : idx+1 ].to(device)
+            tf_token_embed = model.text_embedding(tf_token)
+
+            inputs = torch.cat((inputs, tf_token_embed), dim=1)  # Add the token to the generated text
+
+        predicted_tokens = torch.cat([x.unsqueeze(1) for x in predicted_tokens], dim=1).to(device)
+        #predicted_token_logits = torch.cat([x.unsqueeze(1) for x in predicted_token_logits], dim=1).to(device)
+
+        out = dict(pred=predicted_tokens,
+                   logits=logits)
+
+        labels = labels.contiguous().type(torch.LongTensor).to(device)
+
+        logits = logits[:, ie_size:ie_size+label_size, :].contiguous()
+
+        loss = model.loss(logits.view(-1, logits.size(-1)), labels.view(-1))
+
+        out['loss'] = loss
+
+        #model.train()
+
+        return out
+
+
+def generate_with_logits(logits, temperature=1, top_k=0, top_p=0.0):
+    predicted_tokens = []
+
+    for idx in range(logits.size(1)):
+        next_token_logits = logits[:, idx, :] / temperature  # Apply temperature
+
+        filtered_logits = top_k_top_p_filtering(next_token_logits, top_k=top_k,
+                                                top_p=top_p)  # Apply top-k and/or top-p
+        next_token = torch.multinomial(F.softmax(filtered_logits, dim=-1), num_samples=1)  # Sample
+
+        predicted_tokens.append(next_token)
+
+    predicted_tokens = torch.cat([x.unsqueeze(1) for x in predicted_tokens], dim=1).to(device)
+
+    out = dict(pred=predicted_tokens,
+               logits=logits)
+    return out