train_phi_with_siglip.py

import torch
import torch.nn as nn
from torch.optim import AdamW
from transformers import (
    SiglipVisionModel, 
    AutoTokenizer, 
    AutoImageProcessor, 
    AutoModelForCausalLM,
    BitsAndBytesConfig
)
from peft import prepare_model_for_kbit_training, LoraConfig, get_peft_model
from torchvision.datasets import CIFAR10
from torch.utils.data import DataLoader, Subset
import torchvision.transforms as transforms
from tqdm import tqdm
import os
from PIL import Image

class LinearProjection(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.linear = nn.Linear(input_dim, output_dim)
        
    def forward(self, x):
        return self.linear(x)

class ImageTextProjection(nn.Module):
    def __init__(self, image_dim, text_dim):
        super().__init__()
        self.image_projection = nn.Linear(image_dim, text_dim)
        
    def forward(self, x):
        return self.image_projection(x)

def get_image_embedding(image, siglip_model, siglip_processor, linear_proj, device):
    with torch.no_grad():
        # Process image through SigLIP
        inputs = siglip_processor(image, return_tensors="pt")
        # Move inputs to the same device as model
        inputs = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
        outputs = siglip_model(**inputs)
        image_features = outputs.pooler_output
        
        # Project through trained linear layer
        projected_features = linear_proj(image_features)
        
    return projected_features

def main(
    num_images=100,
    batch_size=4,  # Smaller batch size due to memory constraints
    num_epochs=100,
    learning_rate=2e-4,
    questions=None  # List of 5 questions to be provided
):
    if questions is None or len(questions) != 5:
        print("Please provide exactly 5 questions!")
        return

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # Load SigLIP model and processor
    siglip_model = SiglipVisionModel.from_pretrained("google/siglip-so400m-patch14-384").to(device)
    siglip_processor = AutoImageProcessor.from_pretrained("google/siglip-so400m-patch14-384")
    
    # Load trained linear projection
    dummy_image = Image.new('RGB', (384, 384), color='black')
    with torch.no_grad():
        siglip_inputs = siglip_processor(dummy_image, return_tensors="pt")
        # Move inputs to device
        siglip_inputs = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in siglip_inputs.items()}
        siglip_outputs = siglip_model(**siglip_inputs)
        siglip_output_dim = siglip_outputs.pooler_output.shape[-1]
    
    # First load the checkpoint to get the correct output dimension
    checkpoint = torch.load('linear_projection_final.pth', map_location=device)
    output_dim = checkpoint['linear.weight'].shape[0]  # Get the output dimension from saved weights
    print(f"Loading linear projection with output dimension: {output_dim}")
    
    # Initialize linear projection with correct dimensions
    linear_proj = LinearProjection(siglip_output_dim, output_dim).to(device)
    try:
        linear_proj.load_state_dict(checkpoint)
        print("Successfully loaded linear projection weights")
    except Exception as e:
        print(f"Error loading linear projection weights: {e}")
        return

    # Load Phi model with 4-bit quantization
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.float16,
        bnb_4bit_use_double_quant=False
    )

    phi_model = AutoModelForCausalLM.from_pretrained(
        "microsoft/Phi-3-mini-4k-instruct",
        quantization_config=bnb_config,
        device_map="auto"
    )
    phi_tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-4k-instruct")
    
    # Add padding token if not present
    if phi_tokenizer.pad_token is None:
        phi_tokenizer.pad_token = phi_tokenizer.eos_token
    
    # Get embedding dimension from phi model
    phi_embed_dim = phi_model.get_input_embeddings().weight.shape[1]
    
    # Create projection layer for image embeddings
    image_text_proj = ImageTextProjection(output_dim, phi_embed_dim).to(device)
    
    # Prepare model for k-bit training
    phi_model = prepare_model_for_kbit_training(phi_model)

    # Setup LoRA configuration
    lora_config = LoraConfig(
        r=16,
        lora_alpha=32,
        target_modules=["mlp.dense_h_to_4h", "mlp.dense_4h_to_h", "self_attn.qkv_proj", "self_attn.dense"],
        lora_dropout=0.05,
        bias="none",
        task_type="CAUSAL_LM"
    )

    # Get PEFT model
    phi_model = get_peft_model(phi_model, lora_config)
    
    # Freeze SigLIP and linear projection
    for param in siglip_model.parameters():
        param.requires_grad = False
    for param in linear_proj.parameters():
        param.requires_grad = False

    # Load CIFAR10 test dataset
    transform = transforms.Compose([
        transforms.Resize((384, 384)),
        transforms.ToTensor(),
    ])
    
    test_dataset = CIFAR10(root='./data', train=False, download=True, transform=transform)
    subset_indices = list(range(num_images))
    subset_dataset = Subset(test_dataset, subset_indices)
    dataloader = DataLoader(subset_dataset, batch_size=batch_size, shuffle=False)

    # Optimizer for both phi model and image projection
    optimizer = AdamW([
        {'params': phi_model.parameters()},
        {'params': image_text_proj.parameters()}
    ], lr=learning_rate)

    # Training loop
    for epoch in range(num_epochs):
        total_loss = 0
        phi_model.train()
        image_text_proj.train()
        
        progress_bar = tqdm(dataloader, desc=f'Epoch {epoch+1}/{num_epochs}')
        for batch_idx, (images, _) in enumerate(progress_bar):
            images = images.to(device)
            batch_size = images.size(0)

            # Get image embeddings
            image_embeddings = get_image_embedding(images, siglip_model, siglip_processor, linear_proj, device)

            # Process each question
            for q_idx, question in enumerate(questions):
                # Read corresponding answers
                answers = []
                for idx in range(batch_size):
                    global_idx = batch_idx * batch_size + idx
                    if global_idx < num_images:
                        file_path = f'qa_outputs/image_{global_idx}_extr.txt'
                        try:
                            with open(file_path, 'r') as f:
                                lines = f.readlines()
                                answer = lines[q_idx].strip() if q_idx < len(lines) else ""
                                answers.append(answer)
                        except:
                            answers.append("No answer available")

                # Tokenize questions and answers for the entire batch
                question_tokens = phi_tokenizer(
                    [question] * batch_size,
                    padding=True,
                    truncation=True,
                    max_length=512,
                    return_tensors="pt"
                ).to(device)

                target_tokens = phi_tokenizer(
                    answers,
                    padding=True,
                    truncation=True,
                    max_length=512,
                    return_tensors="pt"
                ).to(device)

                # Get question embeddings for the entire batch
                question_embeds = phi_model.get_input_embeddings()(question_tokens['input_ids'])  # [batch_size, seq_len, embed_dim]

                # Project and prepare image embeddings for the entire batch
                image_embeds = image_text_proj(image_embeddings)  # [batch_size, embed_dim]
                image_embeds = image_embeds.unsqueeze(1)  # [batch_size, 1, embed_dim]

                # Combine image embeddings with question embeddings
                combined_embedding = torch.cat([
                    image_embeds,  # [batch_size, 1, embed_dim]
                    question_embeds  # [batch_size, seq_len, embed_dim]
                ], dim=1)  # [batch_size, 1+seq_len, embed_dim]

                # Create attention mask for the combined sequence
                attention_mask = torch.ones(
                    (batch_size, combined_embedding.size(1)),
                    dtype=torch.long,
                    device=device
                )

                # Prepare labels by shifting them right
                labels = target_tokens['input_ids'].clone()
                labels = torch.cat([
                    torch.full((batch_size, combined_embedding.size(1) - 1), -100, device=device),
                    labels
                ], dim=1)[:, :combined_embedding.size(1)]

                # Forward pass
                outputs = phi_model(
                    inputs_embeds=combined_embedding,
                    attention_mask=attention_mask,
                    labels=labels
                )

                loss = outputs.loss
                total_loss += loss.item()

                # Backward pass
                loss.backward()
                optimizer.step()
                optimizer.zero_grad()

                progress_bar.set_postfix({'loss': loss.item()})

        avg_epoch_loss = total_loss / (len(dataloader) * len(questions) * batch_size)
        print(f'Epoch {epoch+1}/{num_epochs}, Average Loss: {avg_epoch_loss:.4f}')

    # Save the trained models
    phi_model.save_pretrained('phi_model_trained')
    torch.save(image_text_proj.state_dict(), 'image_text_proj.pth')
    print("Training completed. Models saved as 'phi_model_trained' and 'image_text_proj.pth'")

if __name__ == "__main__":
    # Example questions - replace with your actual questions
    questions = [
    "Give a description of the image?",
    "How does the main object in the image look like?",
    "How can the main object in the image be useful to humans?",
    "What is the color of the main object in the image?",
    "Describe the setting of the image?"
    ]
    
    main(questions=questions)