Rename modal.py to model.py

modal.py

@@ -1,300 +0,0 @@
-import os
-import numpy as np
-import torch
-import torch.nn as nn
-from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, WhisperProcessor, WhisperForConditionalGeneration, PreTrainedModel,BitsAndBytesConfig
-from peft import prepare_model_for_kbit_training, LoraConfig, get_peft_model, TaskType
-from tqdm import tqdm
-import json
-import librosa
-from dataclasses import dataclass
-from typing import Any, Dict, List, Union
-import gc
-import torch.nn.functional as F
-
-# # Define multimodal projector class
-# class ProjectionBlock(nn.Module):
-#     def __init__(self, input_dim, output_dim):
-#         super().__init__()
-#         self.pre_norm = nn.LayerNorm(input_dim)
-#         self.proj = nn.Sequential(nn.Linear(input_dim, output_dim), nn.GELU(), nn.Linear(output_dim, output_dim))
-
-#     def forward(self, x):
-#         return self.proj(self.pre_norm(x))
-
-import torch
-import torch.nn as nn
-
-class CrossAttentionBlock(nn.Module):
-    def __init__(self, embed_dim, num_heads=8, dropout=0.1):
-        super().__init__()
-        self.attention = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout)
-        self.norm1 = nn.LayerNorm(embed_dim)
-        self.norm2 = nn.LayerNorm(embed_dim)
-        self.ffn = nn.Sequential(
-            nn.Linear(embed_dim, embed_dim * 4),
-            nn.GELU(),
-            nn.Linear(embed_dim * 4, embed_dim),
-            nn.Dropout(dropout)
-        )
-        
-    def forward(self, x, context):
-        # Self attention
-        attended, _ = self.attention(
-            query=self.norm1(x),
-            key=self.norm1(context),
-            value=self.norm1(context)
-        )
-        x = x + attended
-        
-        # FFN
-        x = x + self.ffn(self.norm2(x))
-        return x
-
-## Updated on 23rd November
-class ProjectionBlock(nn.Module):
-    def __init__(self, input_dim, output_dim):
-        super().__init__()
-        self.pre_norm = nn.LayerNorm(input_dim)
-        self.proj = nn.Sequential(
-            nn.Linear(input_dim, output_dim * 2),  # Increase intermediate dimension
-            nn.GELU(),
-            nn.Linear(output_dim * 2, output_dim)  # Project to final dimension
-        )
-
-    def forward(self, x):
-        # Add shape validation
-        if len(x.shape) == 2:  # If input is [batch_size, features]
-            return self.proj(self.pre_norm(x))
-        elif len(x.shape) == 3:  # If input is [batch_size, seq_len, features]
-            return self.proj(self.pre_norm(x.mean(dim=1)))  # Pool sequence dimension
-        else:
-            raise ValueError(f"Unexpected input shape: {x.shape}")
-
-##Updated on 23rd November
-# class EnhancedMultimodalProjector(nn.Module):
-    # def __init__(self, image_input_dim, audio_input_dim, output_dim, num_heads=8):
-    #     super().__init__()
-        
-    #     # Adjust projectors to match Phi-3's hidden size (1024)
-    #     self.image_proj = ProjectionBlock(image_input_dim, output_dim)
-    #     self.audio_proj = ProjectionBlock(audio_input_dim, output_dim)
-        
-    #     # Cross-attention blocks
-    #     self.image_audio_cross_attn = CrossAttentionBlock(output_dim, num_heads)
-    #     self.audio_image_cross_attn = CrossAttentionBlock(output_dim, num_heads)
-        
-    #     # Final fusion layer
-    #     self.fusion_layer = nn.Sequential(
-    #         nn.LayerNorm(output_dim * 2),
-    #         nn.Linear(output_dim * 2, output_dim),
-    #         nn.GELU(),
-    #         nn.Linear(output_dim, output_dim)
-    #     )
-class EnhancedMultimodalProjector(nn.Module):
-    def __init__(self, image_input_dim, audio_input_dim=1024, output_dim=1024, num_heads=8):
-        super().__init__()
-        self.image_proj = ProjectionBlock(image_input_dim, output_dim)
-        self.audio_proj = ProjectionBlock(audio_input_dim, output_dim)
-        self.image_audio_cross_attn = CrossAttentionBlock(output_dim, num_heads)
-        self.audio_image_cross_attn = CrossAttentionBlock(output_dim, num_heads)
-        self.fusion_layer = nn.Sequential(
-            nn.LayerNorm(output_dim * 2),
-            nn.Linear(output_dim * 2, output_dim),
-            nn.GELU(),
-            nn.Linear(output_dim, output_dim)
-        )
-
-    def forward(self, image_embedding=None, audio_embedding=None):
-        # Add shape validation and adjustment
-        if image_embedding is not None and image_embedding.dim() < 2:
-          raise ValueError("Expected `image_embedding` to have at least 2 dimensions.")
-        if audio_embedding is not None and audio_embedding.dim() < 2:
-          raise ValueError("Expected `audio_embedding` to have at least 2 dimensions.")
-        if image_embedding is not None and len(image_embedding.shape) == 2:
-            image_embedding = image_embedding.unsqueeze(1)  # Add sequence dimension
-        if audio_embedding is not None and len(audio_embedding.shape) == 2:
-            audio_embedding = audio_embedding.unsqueeze(1)  # Add sequence dimension
-
-        # Initial projections
-        projected_image = self.image_proj(image_embedding) if image_embedding is not None else None
-        projected_audio = self.audio_proj(audio_embedding) if audio_embedding is not None else None
-
-        if projected_image is not None and projected_audio is not None:
-            # Ensure correct shapes for cross-attention
-            attended_image = self.image_audio_cross_attn(projected_image, projected_audio)
-            attended_audio = self.audio_image_cross_attn(projected_audio, projected_image)
-            
-            # Combine the attended features
-            fused_features = torch.cat([attended_image, attended_audio], dim=-1)
-            final_output = self.fusion_layer(fused_features)
-            
-            return final_output, final_output
-            
-        elif projected_image is not None:
-            return projected_image, None
-        elif projected_audio is not None:
-            return None, projected_audio
-        else:
-            return None, None
-
-# Update the Phi3WithProjector to use the enhanced projector
-class Phi3WithProjector(PreTrainedModel):
-    def __init__(self, config, phi3_model, projector):
-        super().__init__(config)
-        self.phi3_model = phi3_model
-        self.projector = projector
-        self.supports_gradient_checkpointing = True
-
-    def forward(self, input_ids=None, attention_mask=None, inputs_embeds=None, 
-                image_embeddings=None, audio_embeddings=None, labels=None, **kwargs):
-        if inputs_embeds is None:
-            inputs_embeds = self.phi3_model.get_input_embeddings()(input_ids)
-
-        # Get fused embeddings from enhanced projector
-        projected_features, _ = self.projector(image_embeddings, audio_embeddings)
-
-        # Concatenate embeddings if we have projected features
-        if projected_features is not None:
-            combined_embeddings = torch.cat([inputs_embeds, projected_features.unsqueeze(1)], dim=1)
-            # Extend attention mask
-            extended_attention_mask = torch.cat([
-                attention_mask,
-                torch.ones((attention_mask.shape[0], 1), device=attention_mask.device)
-            ], dim=1)
-        else:
-            combined_embeddings = inputs_embeds
-            extended_attention_mask = attention_mask
-
-        # Adjust labels if needed
-        if labels is not None and projected_features is not None:
-            labels = torch.cat([
-                labels,
-                torch.full((labels.shape[0], 1), -100, dtype=labels.dtype, device=labels.device)
-            ], dim=1)
-
-        return self.phi3_model(
-            inputs_embeds=combined_embeddings,
-            attention_mask=extended_attention_mask,
-            labels=labels,
-            **kwargs
-        )
-
-
-class MultimodalProjector(nn.Module):
-    def __init__(self, image_input_dim, audio_input_dim, output_dim):
-        super().__init__()
-        self.image_proj = ProjectionBlock(image_input_dim, output_dim)
-        self.audio_proj = ProjectionBlock(audio_input_dim, output_dim)
-
-    def forward(self, image_embedding=None, audio_embedding=None):
-        projected_image = self.image_proj(image_embedding) if image_embedding is not None else None
-        projected_audio = self.audio_proj(audio_embedding) if audio_embedding is not None else None
-        return projected_image, projected_audio
-
-
-
-class Phi3WithProjector(PreTrainedModel):
-    def __init__(self, config, phi3_model, projector):
-        super().__init__(config)
-        self.phi3_model = phi3_model
-        self.projector = projector
-        self.supports_gradient_checkpointing = True 
-
-
-    def forward(self, input_ids=None, attention_mask=None, inputs_embeds=None, image_embeddings=None, audio_embeddings=None, labels=None, **kwargs):
-        # Use get_input_embeddings() to retrieve the embeddings layer
-        if inputs_embeds is None:
-            inputs_embeds = self.phi3_model.get_input_embeddings()(input_ids)
-
-        # Project both image and audio embeddings to the appropriate dimension
-        projected_image, projected_audio = self.projector(image_embeddings, audio_embeddings)
-
-        # Concatenate the embeddings
-        embeddings_to_concat = [inputs_embeds]
-        if projected_image is not None:
-            embeddings_to_concat.append(projected_image.unsqueeze(1))
-        if projected_audio is not None:
-            embeddings_to_concat.append(projected_audio.unsqueeze(1))
-
-        combined_embeddings = torch.cat(embeddings_to_concat, dim=1)
-
-        # Modify how the attention mask is extended
-        extended_attention_mask = attention_mask.clone()  # Start with a copy
-
-        # Extend for image and audio, if present
-        if projected_image is not None:
-            extended_attention_mask = torch.cat([extended_attention_mask, torch.ones_like(extended_attention_mask[:, :1])], dim=1)
-        if projected_audio is not None:
-            extended_attention_mask = torch.cat([extended_attention_mask, torch.ones_like(extended_attention_mask[:, :1])], dim=1)
-
-        # Adjust labels to match the extended input sequence length
-        if labels is not None:
-            # Pad labels with -100 to ignore the added tokens in the loss calculation
-            num_added_tokens = sum(1 for emb in [projected_image, projected_audio] if emb is not None)
-            labels = torch.cat([labels, torch.full((labels.shape[0], num_added_tokens), -100, dtype=labels.dtype, device=labels.device)], dim=1)
-        outputs = self.phi3_model(
-            inputs_embeds=combined_embeddings,
-            attention_mask=extended_attention_mask,
-            labels=labels,
-            **kwargs
-        )
-        
-        # Add auxiliary losses for multimodal alignment
-        if image_embeddings is not None or audio_embeddings is not None:
-            loss = outputs.loss
-            
-            # Add contrastive loss for multimodal alignment
-            if image_embeddings is not None and audio_embeddings is not None:
-                img_proj, audio_proj = self.projector(image_embeddings, audio_embeddings)
-                contrastive_loss = self.compute_contrastive_loss(img_proj, audio_proj)
-                loss = loss + 0.1 * contrastive_loss  # Weight the auxiliary loss
-                
-            outputs.loss = loss
-
-
-
-
-        return outputs
-
-    def get_input_embeddings(self):
-        """Returns the model's input embeddings."""
-        return self.phi3_model.get_input_embeddings()
-
-    def set_input_embeddings(self, value):
-        """Sets the model's input embeddings."""
-        self.phi3_model.set_input_embeddings(value)
-    
-
-    # Instead, use the built-in gradient checkpointing
-    def enable_gradient_checkpointing(self):
-        """Enable gradient checkpointing for the model."""
-        if hasattr(self.phi3_model, "gradient_checkpointing_enable"):
-            self.phi3_model.gradient_checkpointing_enable()
-        else:
-            self.phi3_model.config.use_cache = False
-            self.phi3_model.train()  # Ensure model is in training mode
-            
-    def disable_gradient_checkpointing(self):
-        """Disable gradient checkpointing for the model."""
-        if hasattr(self.phi3_model, "gradient_checkpointing_disable"):
-            self.phi3_model.gradient_checkpointing_disable()
-        else:
-            self.phi3_model.config.use_cache = True
-
-    def compute_contrastive_loss(self, img_features, audio_features):
-        # Normalize features
-        img_features = F.normalize(img_features, dim=-1)
-        audio_features = F.normalize(audio_features, dim=-1)
-        
-        # Compute similarity matrix
-        similarity = torch.matmul(img_features, audio_features.transpose(0, 1))
-        
-        # Temperature-scaled cross entropy loss
-        temperature = 0.07
-        labels = torch.arange(similarity.size(0)).to(similarity.device)
-        loss = F.cross_entropy(similarity / temperature, labels)
-        
-        return loss
-    
-  

model.py

@@ -0,0 +1,106 @@
+# Step 2: Import necessary libraries
+import gradio as gr
+from PIL import Image
+from transformers import CLIPProcessor, CLIPModel, AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from peft import PeftConfig, PeftModel
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers.cache_utils import DynamicCache, StaticCache
+
+
+# Step 3: Set device and default dtype
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+torch.set_default_dtype(torch.float16)
+
+# Step 4: Load CLIP model and processor
+clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32", torch_dtype=torch.float16).to(DEVICE)
+clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32", use_fast=True)
+
+# Step 5: Define the MultiModalModel class
+class MultiModalModel(nn.Module):
+    def __init__(self, phi_model_name="microsoft/phi-3-mini-4k-instruct", clip_model_name="openai/clip-vit-base-patch32"):
+        super().__init__()
+        self.phi = None  # Will be set after loading the PEFT model
+        self.tokenizer = AutoTokenizer.from_pretrained(phi_model_name, trust_remote_code=True)
+        self.tokenizer.add_special_tokens({"additional_special_tokens": ["[IMG]"], "pad_token": "<pad>"})
+        self.clip = CLIPModel.from_pretrained(clip_model_name, torch_dtype=torch.float16).eval().to(DEVICE)
+        image_embedding_dim = self.clip.config.projection_dim
+        phi_hidden_size = 3072  # Hardcoded for Phi-3 mini
+        self.image_projection = nn.Sequential(
+            nn.Linear(image_embedding_dim, phi_hidden_size, dtype=torch.float16),
+            nn.LayerNorm(phi_hidden_size, dtype=torch.float16),
+            nn.Dropout(0.1)
+        ).to(DEVICE)
+        nn.init.xavier_uniform_(self.image_projection[0].weight, gain=1.0)
+        nn.init.zeros_(self.image_projection[0].bias)
+
+    def forward(self, text_input_ids, attention_mask=None, image_embedding=None):
+        image_embedding = torch.clamp(image_embedding, min=-1e4, max=1e4)
+        image_embedding = F.normalize(image_embedding, dim=-1, eps=1e-5).to(torch.float16)
+        with torch.no_grad():
+            self.image_projection[0].weight.clamp_(-1.0, 1.0)
+            self.image_projection[0].bias.clamp_(-1.0, 1.0)
+        projected_image = 1.0 * self.image_projection(image_embedding)
+        projected_image = torch.clamp(projected_image, min=-1e4, max=1e4)
+        if torch.isnan(projected_image).any() or torch.isinf(projected_image).any():
+            print("Warning: Projected image contains NaN or Inf values after clamping, replacing with zeros")
+            projected_image = torch.where(
+                torch.logical_or(torch.isnan(projected_image), torch.isinf(projected_image)),
+                torch.zeros_like(projected_image),
+                projected_image
+            )
+        if projected_image.dim() == 2:
+            projected_image = projected_image.unsqueeze(1)
+        text_embeddings = self.phi.get_input_embeddings()(text_input_ids)
+        fused_embeddings = text_embeddings.clone()
+        img_token_id = self.tokenizer.convert_tokens_to_ids("[IMG]")
+        img_token_mask = (text_input_ids == img_token_id)
+        for i in range(fused_embeddings.shape[0]):
+            img_positions = img_token_mask[i].nonzero(as_tuple=True)[0]
+            if img_positions.numel() > 0:
+                fused_embeddings[i, img_positions[0], :] = projected_image[i, 0, :]
+        if torch.isnan(fused_embeddings).any() or torch.isinf(fused_embeddings).any():
+            print("Warning: Fused embeddings contain NaN or Inf values, replacing with zeros")
+            fused_embeddings = torch.where(
+                torch.logical_or(torch.isnan(fused_embeddings), torch.isinf(fused_embeddings)),
+                torch.zeros_like(fused_embeddings),
+                fused_embeddings
+            )
+        return fused_embeddings
+
+# Step 6: Load the fine-tuned model weights from Epoch_0
+def load_model():
+    
+    # 1. Load PEFT Config
+    peft_model_id = "/content/drive/MyDrive/V6_Checkpoints/Epoch_peft_1" # Path to the saved PEFT directory
+    config = PeftConfig.from_pretrained(peft_model_id) #Use the config to determine the base model
+
+    attn_implementation = "eager"
+    cache = DynamicCache()
+   
+    quantization_config = BitsAndBytesConfig(load_in_8bit=True)
+    base_model = AutoModelForCausalLM.from_pretrained(
+        config.base_model_name_or_path,
+        return_dict=True,
+        quantization_config=quantization_config,
+        device_map="auto",
+        trust_remote_code=False,
+        torch_dtype=torch.float16,
+        attn_implementation="eager"
+    )
+    base_model.gradient_checkpointing_enable()
+    peft_model = PeftModel.from_pretrained(base_model, peft_model_id)
+    tokenizer = AutoTokenizer.from_pretrained(peft_model_id)
+    special_tokens = {"additional_special_tokens": ["[IMG]"], "pad_token": "<pad>"}
+    tokenizer.add_special_tokens(special_tokens)
+    peft_model.resize_token_embeddings(len(tokenizer))
+    tokenizer.pad_token = tokenizer.eos_token
+
+    model = MultiModalModel(phi_model_name=config.base_model_name_or_path)
+    #model.load_state_dict(torch.load("/content/drive/MyDrive/V7_Checkpoints_2ndMarch2025/Epoch_0/model_state_dict.pth"), strict=False)
+    model.phi = peft_model
+    model.to(DEVICE)
+    model.eval()
+    return model, tokenizer
+