Upload main_clip_model.py with huggingface_hub

main_clip_model.py

@@ -0,0 +1,639 @@
+import os
+# Set environment variable to disable tokenizers parallelism warnings
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+import torch
+import pytorch_lightning as pl
+from torch.utils.data import DataLoader
+import torch.nn.functional as F
+from torch.utils.data import Dataset
+from torchvision import transforms
+from PIL import Image
+from config import local_dataset_path, column_local_image_path, color_emb_dim, hierarchy_emb_dim, color_model_path, hierarchy_model_path, device, main_model_path
+import matplotlib.pyplot as plt
+from transformers import CLIPProcessor, CLIPModel as CLIPModel_transformers
+import warnings
+from tqdm import tqdm
+import numpy as np
+
+# Suppress warnings
+warnings.filterwarnings("ignore", category=FutureWarning)
+warnings.filterwarnings("ignore", category=UserWarning)
+
+# -------------------------------
+# Step 1: Custom Training Functions
+# -------------------------------
+
+def train_one_epoch(model, train_loader, optimizer, feature_models, device, clip_processor, temperature=0.07):
+    """
+    Train the model for one epoch
+    """
+    model.train()
+    total_loss = 0.0
+    num_batches = 0
+    
+    # Create progress bar for training
+    pbar = tqdm(train_loader, desc="Training", leave=False)
+    
+    for batch_idx, (images, texts, colors, hierarchy) in enumerate(pbar):
+        # Move data to device
+        images = images.to(device)
+        images = images.expand(-1, 3, -1, -1)  # Ensure 3 channels
+        
+        # Process text inputs
+        text_inputs = clip_processor(text=texts, padding=True, return_tensors="pt")
+        text_inputs = {k: v.to(device) for k, v in text_inputs.items()}
+        
+        # Forward pass
+        optimizer.zero_grad()
+        outputs = model(**text_inputs, pixel_values=images)
+        
+        text_features = outputs.text_embeds
+        image_features = outputs.image_embeds
+        
+        # Get feature embeddings
+        # Use exact color-name embeddings if available (new color model)
+        if hasattr(feature_models['color'], 'get_color_name_embeddings'):
+            color_features = feature_models['color'].get_color_name_embeddings(colors)
+        else:
+            color_features = feature_models['color'].get_text_embeddings(colors)
+        hierarchy_features = feature_models['hierarchy'].get_text_embeddings(hierarchy)
+        concat_features = torch.cat((color_features, hierarchy_features), dim=1)
+        
+        # Calculate loss
+        loss = triple_contrastive_loss(text_features, image_features, concat_features, temperature)
+        
+        # Backward pass
+        loss.backward()
+        optimizer.step()
+        
+        total_loss += loss.item()
+        num_batches += 1
+        
+        # Update progress bar
+        pbar.set_postfix({
+            'Loss': f'{loss.item():.4f}',
+            'Avg Loss': f'{total_loss/num_batches:.4f}'
+        })
+    
+    return total_loss / num_batches
+
+def valid_one_epoch(model, val_loader, feature_models, device, clip_processor, temperature=0.07):
+    """
+    Validate the model for one epoch
+    """
+    model.eval()
+    total_loss = 0.0
+    num_batches = 0
+    
+    # Create progress bar for validation
+    pbar = tqdm(val_loader, desc="Validation", leave=False)
+    
+    with torch.no_grad():
+        for batch_idx, (images, texts, colors, hierarchy) in enumerate(pbar):
+            # Move data to device
+            images = images.to(device)
+            images = images.expand(-1, 3, -1, -1)  # Ensure 3 channels
+            
+            # Process text inputs
+            text_inputs = clip_processor(text=texts, padding=True, return_tensors="pt")
+            text_inputs = {k: v.to(device) for k, v in text_inputs.items()}
+            
+            # Forward pass
+            outputs = model(**text_inputs, pixel_values=images)
+            
+            text_features = outputs.text_embeds
+            image_features = outputs.image_embeds
+            
+            # Get feature embeddings
+            if hasattr(feature_models['color'], 'get_color_name_embeddings'):
+                color_features = feature_models['color'].get_color_name_embeddings(colors)
+            else:
+                color_features = feature_models['color'].get_text_embeddings(colors)
+            hierarchy_features = feature_models['hierarchy'].get_text_embeddings(hierarchy)
+            concat_features = torch.cat((color_features, hierarchy_features), dim=1)
+            
+            # Calculate loss
+            loss = triple_contrastive_loss(text_features, image_features, concat_features, temperature)
+            
+            total_loss += loss.item()
+            num_batches += 1
+            
+            # Update progress bar
+            pbar.set_postfix({
+                'Loss': f'{loss.item():.4f}',
+                'Avg Loss': f'{total_loss/num_batches:.4f}'
+            })
+    
+    return total_loss / num_batches
+
+def triple_contrastive_loss(text_features, image_features, attribute_features, temperature=0.07):
+    """
+    Calculate triple contrastive loss
+    """
+    text_features = F.normalize(text_features, dim=-1)
+    image_features = F.normalize(image_features, dim=-1)
+    attribute_features = F.normalize(attribute_features, dim=-1)
+
+    text_image_logits = (text_features[:, color_emb_dim+hierarchy_emb_dim:] @ image_features[:, color_emb_dim+hierarchy_emb_dim:].T) / temperature
+    text_attr_logits = (text_features[:, :color_emb_dim+hierarchy_emb_dim] @ attribute_features.T) / temperature
+    image_attr_logits = (attribute_features @ image_features[:,:color_emb_dim+hierarchy_emb_dim].T) / temperature
+
+    # Weight distribution
+    weight_text_image = 0.7
+    weight_attr_based = 0.15
+    
+    logits = (weight_text_image * text_image_logits + 
+             weight_attr_based * text_attr_logits + 
+             weight_attr_based * image_attr_logits)
+    
+    labels = torch.arange(len(text_features)).to(text_features.device)
+    loss = (F.cross_entropy(logits, labels) + F.cross_entropy(logits.T, labels)) / 2
+
+    return loss
+
+def enhanced_contrastive_loss(text_features, image_features, attribute_features, 
+                            color_model, colors, temperature=0.07, alignment_weight=0.3):
+    """
+    Enhanced contrastive loss with direct alignment between color model and main model
+    
+    Args:
+        text_features: Main model text embeddings
+        image_features: Main model image embeddings  
+        attribute_features: Concatenated color + hierarchy features
+        color_model: Pre-trained color model
+        colors: List of color strings for this batch
+        temperature: Temperature for contrastive loss
+        alignment_weight: Weight for the alignment loss
+    """
+    
+    # Original triple contrastive loss
+    text_features_norm = F.normalize(text_features, dim=-1)
+    image_features_norm = F.normalize(image_features, dim=-1)
+    attribute_features_norm = F.normalize(attribute_features, dim=-1)
+
+    text_image_logits = (text_features_norm[:, color_emb_dim+hierarchy_emb_dim:] @ 
+                        image_features_norm[:, color_emb_dim+hierarchy_emb_dim:].T) / temperature
+    text_attr_logits = (text_features_norm[:, :color_emb_dim+hierarchy_emb_dim] @ 
+                       attribute_features_norm.T) / temperature
+    image_attr_logits = (attribute_features_norm @ 
+                        image_features_norm[:,:color_emb_dim+hierarchy_emb_dim].T) / temperature
+
+    # Weight distribution for original loss
+    weight_text_image = 0.7
+    weight_attr_based = 0.15
+    
+    original_logits = (weight_text_image * text_image_logits + 
+                      weight_attr_based * text_attr_logits + 
+                      weight_attr_based * image_attr_logits)
+    
+    labels = torch.arange(len(text_features)).to(text_features.device)
+    original_loss = (F.cross_entropy(original_logits, labels) + 
+                    F.cross_entropy(original_logits.T, labels)) / 2
+
+    # NEW: Direct alignment loss between color model and main model first 16 logits
+    with torch.no_grad():
+        color_embeddings = color_model.get_text_embeddings(colors)  # [batch_size, 16]
+    
+    # Extract first 16 dimensions from main model text embeddings
+    main_color_text = text_features[:, :color_emb_dim]  # [batch_size, 16]
+    main_color_image = image_features[:, :color_emb_dim]  # [batch_size, 16]
+    
+    # Normalize for better correlation
+    color_embeddings_norm = F.normalize(color_embeddings, dim=-1)
+    main_color_text_norm = F.normalize(main_color_text, dim=-1)
+    main_color_image_norm = F.normalize(main_color_image, dim=-1)
+    
+    # Direct alignment loss using MSE and cosine similarity
+    text_alignment_loss = F.mse_loss(main_color_text_norm, color_embeddings_norm)
+    image_alignment_loss = F.mse_loss(main_color_image_norm, color_embeddings_norm)
+    
+    # Also encourage high cosine similarity
+    text_cosine_loss = 1 - F.cosine_similarity(main_color_text_norm, color_embeddings_norm).mean()
+    image_cosine_loss = 1 - F.cosine_similarity(main_color_image_norm, color_embeddings_norm).mean()
+    
+    alignment_loss = (text_alignment_loss + image_alignment_loss + 
+                     text_cosine_loss + image_cosine_loss) / 4
+    
+    # Combine losses
+    total_loss = (1 - alignment_weight) * original_loss + alignment_weight * alignment_loss
+    
+    return total_loss, {
+        'original_loss': original_loss.item(),
+        'alignment_loss': alignment_loss.item(),
+        'text_alignment': text_alignment_loss.item(),
+        'image_alignment': image_alignment_loss.item(),
+        'text_cosine': text_cosine_loss.item(),
+        'image_cosine': image_cosine_loss.item()
+    }
+
+def train_one_epoch_enhanced(model, train_loader, optimizer, feature_models, color_model, 
+                           device, clip_processor, temperature=0.07, alignment_weight=0.3):
+    """
+    Enhanced training with direct color alignment loss
+    """
+    model.train()
+    total_loss = 0.0
+    total_metrics = {
+        'original_loss': 0.0,
+        'alignment_loss': 0.0,
+        'text_alignment': 0.0,
+        'image_alignment': 0.0,
+        'text_cosine': 0.0,
+        'image_cosine': 0.0
+    }
+    num_batches = 0
+    
+    pbar = tqdm(train_loader, desc="Training Enhanced", leave=False)
+    
+    for batch_idx, (images, texts, colors, hierarchy) in enumerate(pbar):
+        # Move data to device
+        images = images.to(device)
+        images = images.expand(-1, 3, -1, -1)
+        
+        # Process text inputs
+        text_inputs = clip_processor(text=texts, padding=True, return_tensors="pt")
+        text_inputs = {k: v.to(device) for k, v in text_inputs.items()}
+        
+        # Forward pass
+        optimizer.zero_grad()
+        outputs = model(**text_inputs, pixel_values=images)
+        
+        text_features = outputs.text_embeds
+        image_features = outputs.image_embeds
+        
+        # Get feature embeddings
+        if hasattr(feature_models['color'], 'get_color_name_embeddings'):
+            color_features = feature_models['color'].get_color_name_embeddings(colors)
+        else:
+            color_features = feature_models['color'].get_text_embeddings(colors)
+        hierarchy_features = feature_models['hierarchy'].get_text_embeddings(hierarchy)
+        concat_features = torch.cat((color_features, hierarchy_features), dim=1)
+        
+        # Calculate enhanced loss
+        loss, metrics = enhanced_contrastive_loss(
+            text_features, image_features, concat_features, 
+            color_model, colors, temperature, alignment_weight
+        )
+        
+        # Backward pass
+        loss.backward()
+        optimizer.step()
+        
+        total_loss += loss.item()
+        for key, value in metrics.items():
+            total_metrics[key] += value
+        num_batches += 1
+        
+        # Update progress bar
+        pbar.set_postfix({
+            'Loss': f'{loss.item():.4f}',
+            'Align': f'{metrics["alignment_loss"]:.4f}',
+            'Text_Cos': f'{metrics["text_cosine"]:.4f}',
+            'Img_Cos': f'{metrics["image_cosine"]:.4f}'
+        })
+    
+    avg_metrics = {key: value / num_batches for key, value in total_metrics.items()}
+    return total_loss / num_batches, avg_metrics
+
+def validate_correlation(model, color_model, val_loader, clip_processor, device):
+    """
+    Validate the correlation between color model and main model embeddings
+    """
+    model.eval()
+    color_model.eval()
+    
+    all_color_embeddings = []
+    all_main_text_color = []
+    all_main_image_color = []
+    
+    with torch.no_grad():
+        for batch_idx, (images, texts, colors, hierarchy) in enumerate(tqdm(val_loader, desc="Validation Correlation", leave=False)):
+            if batch_idx >= 50:  # Limit validation samples
+                break
+                
+            images = images.to(device)
+            images = images.expand(-1, 3, -1, -1)
+            
+            text_inputs = clip_processor(text=texts, padding=True, return_tensors="pt")
+            text_inputs = {k: v.to(device) for k, v in text_inputs.items()}
+            
+            # Get embeddings
+            outputs = model(**text_inputs, pixel_values=images)
+            if hasattr(color_model, 'get_color_name_embeddings'):
+                color_emb = color_model.get_color_name_embeddings(colors)
+            else:
+                color_emb = color_model.get_text_embeddings(colors)
+            
+            # Extract first 16 dimensions
+            main_text_color = outputs.text_embeds[:, :color_emb_dim]
+            main_image_color = outputs.image_embeds[:, :color_emb_dim]
+            
+            all_color_embeddings.append(color_emb.cpu().numpy())
+            all_main_text_color.append(main_text_color.cpu().numpy())
+            all_main_image_color.append(main_image_color.cpu().numpy())
+    
+    # Compute correlations
+    color_emb = np.vstack(all_color_embeddings)
+    main_text = np.vstack(all_main_text_color)
+    main_image = np.vstack(all_main_image_color)
+    
+    # Flatten for correlation
+    color_flat = color_emb.flatten()
+    text_flat = main_text.flatten()
+    image_flat = main_image.flatten()
+    
+    text_correlation = np.corrcoef(color_flat, text_flat)[0, 1]
+    image_correlation = np.corrcoef(color_flat, image_flat)[0, 1]
+    
+    return {
+        'text_correlation': text_correlation,
+        'image_correlation': image_correlation
+    }
+
+# -------------------------------
+# Step 2: Define Dataset
+# -------------------------------
+
+class CustomDataset(Dataset):
+    def __init__(self, dataframe, use_local_images=True, image_size=224):
+        self.dataframe = dataframe
+        self.use_local_images = use_local_images
+        self.image_size = image_size
+        
+        # Transforms with augmentation for training
+        self.transform = transforms.Compose([
+            transforms.Resize((image_size, image_size)),
+            transforms.RandomHorizontalFlip(p=0.5),
+            transforms.RandomRotation(15),
+            transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.15),
+            transforms.RandomAffine(degrees=0, translate=(0.1, 0.1), scale=(0.9, 1.1)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        ])
+        
+        # Transforms for validation (no augmentation)
+        self.val_transform = transforms.Compose([
+            transforms.Resize((image_size, image_size)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        ])
+        
+        self.training_mode = True
+        
+    def set_training_mode(self, training=True):
+        self.training_mode = training
+
+    def __len__(self):
+        return len(self.dataframe)
+
+    def __getitem__(self, idx):
+        row = self.dataframe.iloc[idx]
+        
+        image_data = row[column_local_image_path]
+        image = Image.open(image_data).convert("RGB")
+        
+        # Apply appropriate transform
+        if self.training_mode:
+            image = self.transform(image)
+        else:
+            image = self.val_transform(image)
+
+        # Get text and labels
+        description = row['text']
+        color = row['color']
+        hierarchy = row['hierarchy']
+
+        return image, description, color, hierarchy
+
+def train_model(model, train_loader, val_loader, feature_models, device, 
+                      num_epochs=20, learning_rate=1e-5, temperature=0.07, 
+                      save_path=main_model_path, use_enhanced_loss=False, alignment_weight=0.3, color_alignment_model=None):
+    """
+    Custom training loop using train_one_epoch and valid_one_epoch functions
+    """
+    model = model.to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', patience=3, factor=0.5)
+    
+    train_losses = []
+    val_losses = []
+    best_val_loss = float('inf')
+    patience_counter = 0
+    patience = 5
+    
+    print(f"Starting training for {num_epochs} epochs...")
+    print(f"Learning rate: {learning_rate}")
+    print(f"Temperature: {temperature}")
+    print(f"Device: {device}")
+    print(f"Training samples: {len(train_loader.dataset)}")
+    print(f"Validation samples: {len(val_loader.dataset)}")
+    print(f"Batch size: {train_loader.batch_size}")
+    print(f"Estimated time per epoch: ~{len(train_loader) * 2 / 60:.1f} minutes")
+    
+    # Create processor once for efficiency
+    processor = CLIPProcessor.from_pretrained('laion/CLIP-ViT-B-32-laion2B-s34B-b79K')
+    
+    # Create progress bar for epochs
+    epoch_pbar = tqdm(range(num_epochs), desc="Training Progress", position=0)
+    
+    for epoch in epoch_pbar:
+        # Update epoch progress bar
+        epoch_pbar.set_description(f"Epoch {epoch+1}/{num_epochs}")
+        
+        # Training
+        if use_enhanced_loss:
+            if color_alignment_model is None:
+                color_alignment_model = feature_models['color']
+            train_loss, align_metrics = train_one_epoch_enhanced(
+                model, train_loader, optimizer, feature_models, color_alignment_model, device, processor, temperature, alignment_weight
+            )
+        else:
+            train_loss = train_one_epoch(model, train_loader, optimizer, feature_models, device, processor, temperature)
+            align_metrics = None
+        train_losses.append(train_loss)
+        
+        # Validation
+        val_loss = valid_one_epoch(model, val_loader, feature_models, device, processor, temperature)
+        val_losses.append(val_loss)
+        
+        # Learning rate scheduling
+        scheduler.step(val_loss)
+        
+        # Update epoch progress bar with metrics
+        postfix = {
+            'Train Loss': f'{train_loss:.4f}',
+            'Val Loss': f'{val_loss:.4f}',
+            'LR': f'{optimizer.param_groups[0]["lr"]:.2e}',
+            'Best Val': f'{best_val_loss:.4f}'
+        }
+        if align_metrics is not None:
+            postfix.update({'Align': f"{align_metrics['alignment_loss']:.3f}", 'TextCos': f"{align_metrics['text_cosine']:.3f}", 'ImgCos': f"{align_metrics['image_cosine']:.3f}"})
+        epoch_pbar.set_postfix(postfix)
+        
+        # Save best model
+        if val_loss < best_val_loss:
+            best_val_loss = val_loss
+            patience_counter = 0
+            
+            # Save checkpoint
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'train_loss': train_loss,
+                'val_loss': val_loss,
+                'best_val_loss': best_val_loss,
+            }, save_path)
+        else:
+            patience_counter += 1
+        
+        # Early stopping
+        if patience_counter >= patience:
+            print(f"\n🛑 Early stopping triggered after {patience_counter} epochs without improvement")
+            break
+    
+    # Plot training curves
+    plt.figure(figsize=(12, 4))
+    
+    plt.subplot(1, 2, 1)
+    plt.plot(train_losses, label='Train Loss', color='blue')
+    plt.plot(val_losses, label='Val Loss', color='red')
+    plt.title('Training and Validation Loss')
+    plt.xlabel('Epoch')
+    plt.ylabel('Loss')
+    plt.legend()
+    plt.grid(True, alpha=0.3)
+    
+    plt.subplot(1, 2, 2)
+    plt.plot(train_losses, label='Train Loss', color='blue')
+    plt.title('Training Loss')
+    plt.xlabel('Epoch')
+    plt.ylabel('Loss')
+    plt.legend()
+    plt.grid(True, alpha=0.3)
+    
+    plt.tight_layout()
+    plt.savefig('training_curves.png', dpi=300, bbox_inches='tight')
+    plt.show()
+    
+    print(f"\nTraining completed!")
+    print(f"Best validation loss: {best_val_loss:.4f}")
+    print(f"Final model saved to: {save_path}")
+    print(f"Training curves saved to: training_curves.png")
+    
+    return train_losses, val_losses
+
+def load_models():
+    # Load feature models
+    from color_model import ColorCLIP, SimpleTokenizer
+    from hierarchy_model import Model, HierarchyExtractor
+    import json
+
+        # Initialize tokenizer first
+    tokenizer = SimpleTokenizer()
+    
+    # Load vocabulary if available
+    vocab_path = 'tokenizer_vocab.json'
+    if os.path.exists(vocab_path):
+        with open(vocab_path, 'r') as f:
+            vocab_dict = json.load(f)
+            tokenizer.load_vocab(vocab_dict)
+            print(f"Tokenizer vocabulary loaded from {vocab_path}")
+    else:
+        print(f"Warning: {vocab_path} not found. Using default tokenizer.")
+    
+    # Load trained model first to get correct vocab size
+    checkpoint = torch.load(color_model_path, map_location=device)
+    
+    # Extract vocab size from the checkpoint's embedding layer
+    vocab_size_from_checkpoint = checkpoint['text_encoder.embedding.weight'].shape[0]
+    print(f"Vocab size from checkpoint: {vocab_size_from_checkpoint}")
+    print(f"Vocab size from tokenizer: {tokenizer.counter}")
+    
+    # Use the larger of the two to ensure compatibility
+    vocab_size = max(vocab_size_from_checkpoint, tokenizer.counter)
+    
+    # Initialize model with correct vocab size
+    color_model = ColorCLIP(vocab_size=vocab_size, embedding_dim=color_emb_dim).to(device)
+    color_model.tokenizer = tokenizer
+        
+    # Load the checkpoint
+    color_model.load_state_dict(checkpoint)
+    print(f"Model loaded from {color_model_path}")
+    
+    color_model.eval()
+    color_model.name = 'color'
+
+    # Load hierarchy model (embed_dim=64)
+    hierarchy_checkpoint = torch.load(hierarchy_model_path, map_location=device)
+    hierarchy_classes = hierarchy_checkpoint.get('hierarchy_classes', [])
+    hierarchy_model = Model(
+        num_hierarchy_classes=len(hierarchy_classes),
+        embed_dim=hierarchy_emb_dim
+    ).to(device)
+    hierarchy_model.load_state_dict(hierarchy_checkpoint['model_state'])
+
+    # Set up hierarchy extractor
+    hierarchy_extractor = HierarchyExtractor(hierarchy_classes, verbose=False)
+    hierarchy_model.set_hierarchy_extractor(hierarchy_extractor)
+    hierarchy_model.eval()
+    hierarchy_model.name = 'hierarchy'
+
+    feature_models = {model.name: model for model in [color_model, hierarchy_model]}
+
+    return feature_models
+
+if __name__ == "__main__":
+    # Load and prepare data
+    import pandas as pd
+    
+    print("Loading data...")
+    df = pd.read_csv(local_dataset_path)
+    print(f"Loaded {len(df)} samples")
+    
+    # Filter out rows with NaN values in image path
+    df_clean = df.dropna(subset=[column_local_image_path])
+    print(f"After filtering NaN image paths: {len(df_clean)} samples")
+    
+    # Create datasets
+    dataset = CustomDataset(df_clean)
+    
+    # Split for train/val - use only a subset for faster training
+    # Use 10% of data for faster training
+    subset_size = min(10000, len(dataset))  # Max 10k samples
+    train_size = int(0.8 * subset_size)
+    val_size = subset_size - train_size
+    
+    # Create subset with proper integer indices
+    subset_indices = np.random.choice(len(dataset), subset_size, replace=False)
+    subset_dataset = torch.utils.data.Subset(dataset, subset_indices)
+    
+    train_dataset, val_dataset = torch.utils.data.random_split(subset_dataset, [train_size, val_size])
+    
+    # Create dataloaders with optimized parameters
+    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=2, pin_memory=True)
+    val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=2, pin_memory=True)
+    
+    print(f"Train samples: {len(train_dataset)}")
+    print(f"Val samples: {len(val_dataset)}")
+
+    print("Loading models...")   
+    feature_models = load_models()
+    
+    # Create the main CLIP model
+    clip_model = CLIPModel_transformers.from_pretrained('laion/CLIP-ViT-B-32-laion2B-s34B-b79K')
+
+    print("Training...")
+    
+    # Train using custom training loop
+    train_losses, val_losses = train_model(
+        model=clip_model,
+        train_loader=train_loader,
+        val_loader=val_loader,
+        feature_models=feature_models,
+        device=device,
+        num_epochs=20,  # Reduced epochs for faster training
+        learning_rate=2e-5,  # Slightly higher learning rate
+        temperature=0.07
+    )