Leacb4
/

gap-clip

@@ -1,840 +0,0 @@
-#!/usr/bin/env python3
-"""
-Main file for training the CLIP model with color and hierarchy alignment.
-This file centralizes all the logic for training the main model. It uses
-pre-trained color and hierarchy models to guide the main model's learning
-through contrastive and alignment loss functions. It handles data loading,
-training with validation, and checkpoint saving.
-"""
-import os
-# Set environment variable to disable tokenizers parallelism warnings
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-import pandas as pd
-import numpy as np
-import torch
-import torch.nn.functional as F
-from torch.utils.data import Dataset, DataLoader, random_split
-from torchvision import transforms
-from PIL import Image
-import matplotlib.pyplot as plt
-from transformers import CLIPProcessor, CLIPModel as CLIPModel_transformers
-import warnings
-from tqdm import tqdm
-import json
-import config
-# Suppress warnings
-warnings.filterwarnings("ignore", category=FutureWarning)
-warnings.filterwarnings("ignore", category=UserWarning)
-# -------------------------------
-# Loss Functions
-# -------------------------------
-def triple_contrastive_loss(text_features, image_features, attribute_features, temperature=0.07):
-    """
-    Calculate triple contrastive loss for text, image, and attribute features.
-    This loss combines text-image similarity with attribute-based similarities
-    (color and hierarchy) to learn aligned embeddings.
-    Args:
-        text_features: Text embeddings from main model [batch_size, embed_dim]
-        image_features: Image embeddings from main model [batch_size, embed_dim]
-        attribute_features: Concatenated color + hierarchy embeddings [batch_size, color_dim + hierarchy_dim]
-        temperature: Temperature scaling parameter for contrastive loss (default: 0.07)
-    Returns:
-        Contrastive loss value
-    """
-    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[:, config.color_emb_dim+config.hierarchy_emb_dim:] @ image_features[:, config.color_emb_dim+config.hierarchy_emb_dim:].T) / temperature
-    text_attr_logits = (text_features[:, :config.color_emb_dim+config.hierarchy_emb_dim] @ attribute_features.T) / temperature
-    image_attr_logits = (attribute_features @ image_features[:,:config.color_emb_dim+config.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, hierarchy_model, colors, hierarchies, temperature=0.07, alignment_weight=0.3):
-    """
-    Enhanced contrastive loss with direct alignment between color/hierarchy models and main model.
-    This loss combines the original triple contrastive loss with direct alignment losses
-    that force the main model's color and hierarchy dimensions to align with the
-    specialized color and hierarchy models.
-    Args:
-        text_features: Main model text embeddings [batch_size, embed_dim]
-        image_features: Main model image embeddings [batch_size, embed_dim]
-        attribute_features: Concatenated color + hierarchy features [batch_size, color_dim + hierarchy_dim]
-        color_model: Pre-trained color model for extracting color embeddings
-        hierarchy_model: Pre-trained hierarchy model for extracting hierarchy embeddings
-        colors: List of color strings for this batch [batch_size]
-        hierarchies: List of hierarchy strings for this batch [batch_size]
-        temperature: Temperature scaling parameter for contrastive loss (default: 0.07)
-        alignment_weight: Weight for the alignment loss component (default: 0.3)
-    Returns:
-        Tuple of (total_loss, metrics_dict) where metrics_dict contains detailed loss components
-    """
-    # 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[:, config.color_emb_dim+config.hierarchy_emb_dim:] @
-                        image_features_norm[:, config.color_emb_dim+config.hierarchy_emb_dim:].T) / temperature
-    text_attr_logits = (text_features_norm[:, :config.color_emb_dim+config.hierarchy_emb_dim] @
-                       attribute_features_norm.T) / temperature
-    image_attr_logits = (attribute_features_norm @
-                        image_features_norm[:,:config.color_emb_dim+config.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
-    # Direct alignment loss between color model and main model first 16 dims
-    with torch.no_grad():
-        color_embeddings = color_model.get_text_embeddings(colors)
-        hierarchy_embeddings = hierarchy_model.get_text_embeddings(hierarchies)
-    # Extract color dimensions from main model embeddings
-    main_color_text = text_features[:, :config.color_emb_dim]
-    main_color_image = image_features[:, :config.color_emb_dim]
-    # Extract hierarchy dimensions from main model embeddings
-    main_hierarchy_text = text_features[:, config.color_emb_dim:config.color_emb_dim+config.hierarchy_emb_dim]
-    main_hierarchy_image = image_features[:, config.color_emb_dim:config.color_emb_dim+config.hierarchy_emb_dim]
-    # 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)
-    hierarchy_embeddings_norm = F.normalize(hierarchy_embeddings, dim=-1)
-    main_hierarchy_text_norm = F.normalize(main_hierarchy_text, dim=-1)
-    main_hierarchy_image_norm = F.normalize(main_hierarchy_image, dim=-1)
-    # Color alignment loss using MSE and cosine similarity
-    color_text_alignment_loss = F.mse_loss(main_color_text_norm, color_embeddings_norm)
-    color_image_alignment_loss = F.mse_loss(main_color_image_norm, color_embeddings_norm)
-    color_text_cosine_loss = 1 - F.cosine_similarity(main_color_text_norm, color_embeddings_norm).mean()
-    color_image_cosine_loss = 1 - F.cosine_similarity(main_color_image_norm, color_embeddings_norm).mean()
-    # Color alignment loss
-    color_alignment_loss = (
-        color_text_alignment_loss + color_image_alignment_loss +
-        color_text_cosine_loss + color_image_cosine_loss
-    ) / 4
-    # Hierarchy alignment loss using MSE and cosine similarity
-    hierarchy_text_alignment_loss = F.mse_loss(main_hierarchy_text_norm, hierarchy_embeddings_norm)
-    hierarchy_image_alignment_loss = F.mse_loss(main_hierarchy_image_norm, hierarchy_embeddings_norm)
-    hierarchy_text_cosine_loss = 1 - F.cosine_similarity(main_hierarchy_text_norm, hierarchy_embeddings_norm).mean()
-    hierarchy_image_cosine_loss = 1 - F.cosine_similarity(main_hierarchy_image_norm, hierarchy_embeddings_norm).mean()
-    # Hierarchy alignment loss
-    hierarchy_alignment_loss = (
-        hierarchy_text_alignment_loss + hierarchy_image_alignment_loss +
-        hierarchy_text_cosine_loss + hierarchy_image_cosine_loss
-    ) / 4
-    # Combined alignment loss
-    alignment_loss = (color_alignment_loss + hierarchy_alignment_loss) / 2
-    # 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(),
-        'color_text_alignment': color_text_alignment_loss.item(),
-        'color_image_alignment': color_image_alignment_loss.item(),
-        'color_text_cosine': color_text_cosine_loss.item(),
-        'color_image_cosine': color_image_cosine_loss.item(),
-        'hierarchy_text_alignment': hierarchy_text_alignment_loss.item(),
-        'hierarchy_image_alignment': hierarchy_image_alignment_loss.item(),
-        'hierarchy_text_cosine': hierarchy_text_cosine_loss.item(),
-        'hierarchy_image_cosine': hierarchy_image_cosine_loss.item()
-    }
-# -------------------------------
-# Training Functions
-# -------------------------------
-def train_one_epoch(model, train_loader, optimizer, feature_models, device, clip_processor, temperature=0.07):
-    """
-    Train the model for one epoch using triple contrastive loss.
-    Args:
-        model: Main CLIP model to train
-        train_loader: DataLoader for training data
-        optimizer: Optimizer instance
-        feature_models: Dictionary containing color and hierarchy models
-        device: Device to train on
-        clip_processor: CLIP processor for text preprocessing
-        temperature: Temperature scaling parameter for contrastive loss (default: 0.07)
-    Returns:
-        Average training loss for the 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[config.color_column], 'get_color_name_embeddings'):
-            color_features = feature_models[config.color_column].get_color_name_embeddings(colors)
-        else:
-            color_features = feature_models[config.color_column].get_text_embeddings(colors)
-        hierarchy_features = feature_models[config.hierarchy_column].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 train_one_epoch_enhanced(model, train_loader, optimizer, feature_models, color_model, hierarchy_model,
-                           device, clip_processor, temperature=0.07, alignment_weight=0.3):
-    """
-    Enhanced training with direct color and hierarchy alignment loss.
-    This function trains the model using the enhanced contrastive loss that includes
-    direct alignment between the main model's color/hierarchy dimensions and the
-    specialized color/hierarchy models.
-    Args:
-        model: Main CLIP model to train
-        train_loader: DataLoader for training data
-        optimizer: Optimizer instance
-        feature_models: Dictionary containing color and hierarchy models
-        color_model: Pre-trained color model for alignment
-        hierarchy_model: Pre-trained hierarchy model for alignment
-        device: Device to train on
-        clip_processor: CLIP processor for text preprocessing
-        temperature: Temperature scaling parameter for contrastive loss (default: 0.07)
-        alignment_weight: Weight for the alignment loss component (default: 0.3)
-    Returns:
-        Tuple of (average_loss, metrics_dict) where metrics_dict contains detailed loss components
-    """
-    model.train()
-    total_loss = 0.0
-    total_metrics = {
-        'original_loss': 0.0,
-        'alignment_loss': 0.0,
-        'color_text_alignment': 0.0,
-        'color_image_alignment': 0.0,
-        'color_text_cosine': 0.0,
-        'color_image_cosine': 0.0,
-        'hierarchy_text_alignment': 0.0,
-        'hierarchy_image_alignment': 0.0,
-        'hierarchy_text_cosine': 0.0,
-        'hierarchy_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[config.color_column], 'get_color_name_embeddings'):
-            color_features = feature_models[config.color_column].get_color_name_embeddings(colors)
-        else:
-            color_features = feature_models[config.color_column].get_text_embeddings(colors)
-        hierarchy_features = feature_models[config.hierarchy_column].get_text_embeddings(hierarchy)
-        concat_features = torch.cat((color_features, hierarchy_features), dim=1)
-        # Calculate enhanced loss with hierarchy alignment
-        loss, metrics = enhanced_contrastive_loss(
-            text_features, image_features, concat_features,
-            color_model, hierarchy_model, colors, hierarchy, 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}',
-            'ColCos': f'{metrics["color_text_cosine"]:.3f}',
-            'HierCos': f'{metrics["hierarchy_text_cosine"]:.3f}'
-        })
-    avg_metrics = {key: value / num_batches for key, value in total_metrics.items()}
-    return total_loss / num_batches, avg_metrics
-def valid_one_epoch(model, val_loader, feature_models, device, clip_processor, temperature=0.07):
-    """
-    Validate the model for one epoch using triple contrastive loss.
-    Args:
-        model: Main CLIP model to validate
-        val_loader: DataLoader for validation data
-        feature_models: Dictionary containing color and hierarchy models
-        device: Device to validate on
-        clip_processor: CLIP processor for text preprocessing
-        temperature: Temperature scaling parameter for contrastive loss (default: 0.07)
-    Returns:
-        Average validation loss for the 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[config.color_column], 'get_color_name_embeddings'):
-                color_features = feature_models[config.color_column].get_color_name_embeddings(colors)
-            else:
-                color_features = feature_models[config.color_column].get_text_embeddings(colors)
-            hierarchy_features = feature_models[config.hierarchy_column].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
-# -------------------------------
-# Dataset
-# -------------------------------
-class CustomDataset(Dataset):
-    """
-    Custom dataset for main model training.
-    Handles loading images from local paths, extracting text descriptions,
-    and applying appropriate transformations for training and validation.
-    """
-    def __init__(self, dataframe, use_local_images=True, image_size=224):
-        """
-        Initialize the custom dataset.
-        Args:
-            dataframe: DataFrame with columns for image paths, text descriptions, colors, and hierarchy labels
-            use_local_images: Whether to use local images (default: True)
-            image_size: Size of images after resizing (default: 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):
-        """
-        Switch between training and validation transforms.
-        Args:
-            training: If True, use training transforms with augmentation; if False, use validation transforms
-        """
-        self.training_mode = training
-    def __len__(self):
-        """Return the number of samples in the dataset."""
-        return len(self.dataframe)
-    def __getitem__(self, idx):
-        """
-        Get a sample from the dataset.
-        Args:
-            idx: Index of the sample
-        Returns:
-            Tuple of (image_tensor, description_text, color_label, hierarchy_label)
-        """
-        row = self.dataframe.iloc[idx]
-        image_data = row[config.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[config.text_column]
-        color = row[config.color_column]
-        hierarchy = row[config.hierarchy_column]
-        return image, description, color, hierarchy
-# -------------------------------
-# Model Loading
-# -------------------------------
-def load_models():
-    """
-    Load color and hierarchy models from checkpoints.
-    This function loads the pre-trained color and hierarchy models along with
-    their tokenizers and extractors, and prepares them for use in main model training.
-    Returns:
-        Dictionary mapping model names to model instances:
-        - 'color': ColorCLIP model instance
-        - 'hierarchy': Hierarchy model instance
-    """
-    from color_model import ColorCLIP, Tokenizer
-    from hierarchy_model import Model, HierarchyExtractor
-    # Initialize tokenizer first
-    tokenizer = Tokenizer()
-    # Load vocabulary if available
-    if os.path.exists(config.tokeniser_path):
-        with open(config.tokeniser_path, 'r') as f:
-            vocab_dict = json.load(f)
-            tokenizer.load_vocab(vocab_dict)
-            print(f"Tokenizer vocabulary loaded from {config.tokeniser_path}")
-    else:
-        print(f"Warning: {config.tokeniser_path} not found. Using default tokenizer.")
-    # Load trained model first to get correct vocab size
-    checkpoint = torch.load(config.config.color_model_path, map_location=config.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=config.color_emb_dim).to(config.device)
-    color_model.tokenizer = tokenizer
-    # Load the checkpoint
-    color_model.load_state_dict(checkpoint)
-    print(f"Color model loaded from {config.color_model_path}")
-    color_model.eval()
-    color_model.name = config.color_column
-    # Load hierarchy model
-    hierarchy_checkpoint = torch.load(config.hierarchy_model_path, map_location=config.device)
-    hierarchy_classes = hierarchy_checkpoint.get('hierarchy_classes', [])
-    hierarchy_model = Model(
-        num_hierarchy_classes=len(hierarchy_classes),
-        embed_dim=config.hierarchy_emb_dim
-    ).to(config.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 = config.hierarchy_column
-    feature_models = {model.name: model for model in [color_model, hierarchy_model]}
-    return feature_models
-# -------------------------------
-# Main Training Function
-# -------------------------------
-def train_model(model, train_loader, val_loader, feature_models, device,
-                      num_epochs=20, learning_rate=1e-5, temperature=0.07,
-                      save_path=config.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.
-    This function handles the complete training process including:
-    - Training and validation loops
-    - Learning rate scheduling
-    - Early stopping
-    - Model checkpointing
-    - Training curve visualization
-    Args:
-        model: Main CLIP model to train
-        train_loader: DataLoader for training data
-        val_loader: DataLoader for validation data
-        feature_models: Dictionary containing color and hierarchy models
-        device: Device to train on
-        num_epochs: Number of training epochs (default: 20)
-        learning_rate: Learning rate for optimizer (default: 1e-5)
-        temperature: Temperature scaling parameter for contrastive loss (default: 0.07)
-        save_path: Path to save model checkpoints (default: main_model_path)
-        use_enhanced_loss: Whether to use enhanced contrastive loss with alignment (default: False)
-        alignment_weight: Weight for alignment loss component if using enhanced loss (default: 0.3)
-        color_alignment_model: Optional color model for alignment (default: None, uses feature_models)
-    Returns:
-        Tuple of (training_losses, validation_losses) lists
-    """
-    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[config.color_column]
-            hierarchy_model = feature_models[config.hierarchy_column]
-            train_loss, align_metrics = train_one_epoch_enhanced(
-                model, train_loader, optimizer, feature_models, color_alignment_model, hierarchy_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}",
-                'ColCos': f"{align_metrics['color_text_cosine']:.3f}",
-                'HierCos': f"{align_metrics['hierarchy_text_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.close()
-    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
-# -------------------------------
-# Main Function
-# -------------------------------
-def main():
-    print("="*80)
-    print("🚀 Training of the model with alignement color and hierarchy")
-    print("="*80)
-    # Configuration
-    num_epochs = 20
-    learning_rate = 1e-5
-    temperature = 0.07
-    alignment_weight = 0.5
-    batch_size = 32
-    subset_size = 10000
-    use_enhanced_loss = True
-    # Load the data
-    print(f"\n📂 Loading the data...")
-    df = pd.read_csv(config.local_dataset_path)
-    print(f"  Data downloaded: {len(df)} samples")
-    # filter the rows with NaN values
-    df_clean = df.dropna(subset=[config.column_local_image_path])
-    print(f"  After filtering NaN: {len(df_clean)} samples")
-    # Creation of datasets
-    dataset = CustomDataset(df_clean)
-    # Creation of a subset for a faster training
-    print(f"\n📊 Creation of a subset of {subset_size} samples...")
-    subset_size = min(subset_size, len(dataset))
-    train_size = int(0.8 * subset_size)
-    val_size = subset_size - train_size
-    # Creation of a subset with random indexes but reproductibles
-    np.random.seed(42)
-    subset_indices = np.random.choice(len(dataset), subset_size, replace=False)
-    subset_dataset = torch.utils.data.Subset(dataset, subset_indices)
-    train_dataset, val_dataset = random_split(
-        subset_dataset,
-        [train_size, val_size],
-        generator=torch.Generator().manual_seed(42)
-    )
-    # Creation of dataloaders
-    train_loader = DataLoader(
-        train_dataset,
-        batch_size=batch_size,
-        shuffle=True,
-        num_workers=2,
-        pin_memory=True if torch.cuda.is_available() else False
-    )
-    val_loader = DataLoader(
-        val_dataset,
-        batch_size=batch_size,
-        shuffle=False,
-        num_workers=2,
-        pin_memory=True if torch.cuda.is_available() else False
-    )
-    print(f"  Train: {len(train_dataset)} samples")
-    print(f"  Validation: {len(val_dataset)} samples")
-    # Loading models
-    print(f"\n🔧 Loading models...")
-    feature_models = load_models()
-    # Load or create the main model
-    print(f"\n📦 Loading main model...")
-    clip_model = CLIPModel_transformers.from_pretrained(
-        'laion/CLIP-ViT-B-32-laion2B-s34B-b79K'
-    )
-    # Load the model
-    if os.path.exists(config.main_model_path):
-        print(f"  Model found {config.main_model_path}")
-        print(f"  Loading checkpoint...")
-        checkpoint = torch.load(config.main_model_path, map_location=config.device)
-        if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
-            clip_model.load_state_dict(checkpoint['model_state_dict'])
-            print(f"  ✅ Checkpoint loaded from {checkpoint.get('epoch', '?')}")
-        else:
-            clip_model.load_state_dict(checkpoint)
-            print(f"  ✅ Checkpoint loaded")
-    else:
-        print(f"  New model, no checkpoint found")
-    # Move the model on the device
-    clip_model = clip_model.to(config.device)
-    # Training with enhanced loss
-    print(f"\n🎯 Beginning training...")
-    print(f"\n" + "="*80)
-    train_losses, val_losses = train_model(
-        model=clip_model,
-        train_loader=train_loader,
-        val_loader=val_loader,
-        feature_models=feature_models,
-        device=config.device,
-        num_epochs=num_epochs,
-        learning_rate=learning_rate,
-        temperature=temperature,
-        save_path=config.main_model_path,
-        use_enhanced_loss=use_enhanced_loss,
-        alignment_weight=alignment_weight,
-        color_alignment_model=feature_models[config.color_column]
-    )
-    print("\n" + "="*80)
-    print("✅ Traning finished!")
-    print(f"  Modèle sauvegardé: {config.main_model_path}")
-    print(f"  Training curves: training_curves.png")
-    print("\n📊 Final results:")
-    print(f"  Last train loss: {train_losses[-1]:.4f}")
-    print(f"  Last validation loss: {val_losses[-1]:.4f}")
-    print(f"  Best loss validation: {min(val_losses):.4f}")
-    print("="*80)
-if __name__ == "__main__":
-    main()