Upload evaluation/hierarchy_evaluation_with_clip_baseline.py with huggingface_hub

evaluation/hierarchy_evaluation_with_clip_baseline.py

@@ -0,0 +1,808 @@
+"""
+Hierarchy embedding evaluation with CLIP baseline comparison.
+This file evaluates the quality of hierarchy embeddings from the custom model and compares them
+with a CLIP baseline model (OpenAI CLIP). It calculates similarity metrics, classification accuracies,
+and generates confusion matrices for both models to measure relative performance. It also supports
+evaluation on Fashion-MNIST and kagl Marqo datasets.
+"""
+
+import torch
+import pandas as pd
+import numpy as np
+import matplotlib.pyplot as plt
+import seaborn as sns
+from sklearn.metrics.pairwise import cosine_similarity
+from sklearn.metrics import confusion_matrix, classification_report, accuracy_score, f1_score
+from collections import defaultdict
+import os
+import requests
+from tqdm import tqdm
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
+from sklearn.model_selection import train_test_split
+from io import BytesIO
+from PIL import Image
+import warnings
+warnings.filterwarnings('ignore')
+
+# Import transformers CLIP
+from transformers import CLIPProcessor, CLIPModel as TransformersCLIPModel
+
+# Import your custom model
+from hierarchy_model import Model, HierarchyExtractor, HierarchyDataset, collate_fn
+import config
+
+def convert_fashion_mnist_to_image(pixel_values):
+    """Convert Fashion-MNIST pixel values to PIL image"""
+    # Reshape to 28x28 and convert to PIL Image
+    image_array = np.array(pixel_values).reshape(28, 28).astype(np.uint8)
+    # Convert to RGB by duplicating the grayscale channel
+    image_array = np.stack([image_array] * 3, axis=-1)
+    image = Image.fromarray(image_array)
+    return image
+
+def get_fashion_mnist_labels():
+    """Get Fashion-MNIST class labels"""
+    return {
+        0: "T-shirt/top", 
+        1: "Trouser", 
+        2: "Pullover", 
+        3: "Dress", 
+        4: "Coat", 
+        5: "Sandal", 
+        6: "Shirt", 
+        7: "Sneaker", 
+        8: "Bag", 
+        9: "Ankle boot"
+    }
+
+class FashionMNISTDataset(Dataset):
+    def __init__(self, dataframe, image_size=224):
+        self.dataframe = dataframe
+        self.image_size = image_size
+        self.labels_map = get_fashion_mnist_labels()
+        
+        # Simple transforms for validation/inference
+        self.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])
+        ])
+        
+    def __len__(self):
+        return len(self.dataframe)
+
+    def __getitem__(self, idx):
+        row = self.dataframe.iloc[idx]
+        
+        # Get pixel values (columns 1-784)
+        pixel_cols = [f'pixel{i}' for i in range(1, 785)]
+        pixel_values = row[pixel_cols].values
+        
+        # Convert to image
+        image = convert_fashion_mnist_to_image(pixel_values)
+        image = self.transform(image)
+
+        # Get text description
+        text = row['text']
+        
+        # Get hierarchy label
+        hierarchy = row['hierarchy']
+
+        return image, text, hierarchy
+
+class CLIPBaselineEvaluator:
+    def __init__(self, device='mps'):
+        self.device = torch.device(device)
+        
+        # Load CLIP model and processor
+        print("🤗 Loading CLIP baseline model from transformers...")
+        self.clip_model = TransformersCLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(self.device)
+        self.clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+        self.clip_model.eval()
+        print("✅ CLIP model loaded successfully")
+        
+    def extract_clip_embeddings(self, images, texts):
+        """Extract CLIP embeddings for images and texts"""
+        all_image_embeddings = []
+        all_text_embeddings = []
+        
+        with torch.no_grad():
+            for i in tqdm(range(len(images)), desc="Extracting CLIP embeddings"):
+                # Process image
+                if isinstance(images[i], torch.Tensor):
+                    # Convert tensor back to PIL Image
+                    image_tensor = images[i]
+                    if image_tensor.dim() == 3:
+                        # Denormalize
+                        mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
+                        std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
+                        image_tensor = image_tensor * std + mean
+                        image_tensor = torch.clamp(image_tensor, 0, 1)
+                        
+                        # Convert to PIL
+                        image_pil = transforms.ToPILImage()(image_tensor)
+                elif isinstance(images[i], Image.Image):
+                    image_pil = images[i]
+                else:
+                    raise ValueError(f"Unsupported image type: {type(images[i])}")
+                
+                # Process with CLIP
+                inputs = self.clip_processor(
+                    text=texts[i], 
+                    images=image_pil, 
+                    return_tensors="pt", 
+                    padding=True
+                ).to(self.device)
+                
+                outputs = self.clip_model(**inputs)
+                
+                # Get normalized embeddings
+                image_emb = outputs.image_embeds / outputs.image_embeds.norm(p=2, dim=-1, keepdim=True)
+                text_emb = outputs.text_embeds / outputs.text_embeds.norm(p=2, dim=-1, keepdim=True)
+                
+                all_image_embeddings.append(image_emb.cpu().numpy())
+                all_text_embeddings.append(text_emb.cpu().numpy())
+        
+        return np.vstack(all_image_embeddings), np.vstack(all_text_embeddings)
+
+
+class EmbeddingEvaluator:
+    def __init__(self, model_path, directory):
+        self.device = config.device
+        self.directory = directory
+
+        # 1. Load the dataset
+        CSV = config.local_dataset_path  
+        print(f"📁 Using dataset with local images: {CSV}")
+        df = pd.read_csv(CSV)
+        
+        print(f"📁 Loaded {len(df)} samples")
+        
+        # 2. Get unique hierarchy classes from the dataset
+        hierarchy_classes = sorted(df[config.hierarchy_column].unique().tolist())
+        print(f"📋 Found {len(hierarchy_classes)} hierarchy classes")
+
+        _, self.val_df = train_test_split(df, test_size=0.2, random_state=42, stratify=df['hierarchy'])
+        
+        # 3. Load the model
+        if os.path.exists(model_path):
+            checkpoint = torch.load(model_path, map_location=self.device)
+            
+            config = checkpoint.get('config', {})
+            saved_hierarchy_classes = checkpoint['hierarchy_classes']
+            
+            # Use the saved hierarchy classes
+            self.hierarchy_classes = saved_hierarchy_classes
+            
+            # Create the hierarchy extractor
+            self.vocab = HierarchyExtractor(saved_hierarchy_classes)
+            
+            # Create the model with the saved configuration
+            self.model = Model(
+                num_hierarchy_classes=len(saved_hierarchy_classes),
+                embed_dim=config['embed_dim'],
+                dropout=config['dropout']
+            ).to(self.device)   
+                
+            self.model.load_state_dict(checkpoint['model_state'])
+            
+            print(f"✅ Custom model loaded with:")
+            print(f"📋 Hierarchy classes: {len(saved_hierarchy_classes)}")
+            print(f"🎯 Embed dim: {config['embed_dim']}")
+            print(f"💧 Dropout: {config['dropout']}")
+            print(f"📅 Epoch: {checkpoint.get('epoch', 'unknown')}")
+                
+        else:
+            raise FileNotFoundError(f"Model file {model_path} not found")
+        
+        self.model.eval()
+        
+        # Initialize CLIP baseline
+        self.clip_evaluator = CLIPBaselineEvaluator(device)
+    
+    def create_dataloader(self, dataframe, batch_size=16):
+        """Create a dataloader for custom model"""
+        # Check if this is Fashion-MNIST data (has pixel1 column)
+        if 'pixel1' in dataframe.columns:
+            print("🔍 Detected Fashion-MNIST data, using FashionMNISTDataset")
+            dataset = FashionMNISTDataset(dataframe, image_size=224)
+        else:
+            dataset = HierarchyDataset(dataframe, image_size=224)
+        
+        dataloader = DataLoader(
+            dataset, 
+            batch_size=batch_size, 
+            shuffle=False, 
+            collate_fn=lambda batch: collate_fn(batch, self.vocab), 
+            num_workers=0
+        )
+        
+        return dataloader
+    
+    def create_clip_dataloader(self, dataframe, batch_size=16):
+        """Create a dataloader for CLIP baseline"""
+        # Check if this is Fashion-MNIST data (has pixel1 column)
+        if 'pixel1' in dataframe.columns:
+            print("🔍 Detected Fashion-MNIST data for CLIP, using FashionMNISTDataset")
+            dataset = FashionMNISTDataset(dataframe, image_size=224)
+        else:
+            dataset = CLIPDataset(dataframe)
+        
+        dataloader = DataLoader(
+            dataset, 
+            batch_size=batch_size, 
+            shuffle=False, 
+            num_workers=0
+        )
+        
+        return dataloader
+    
+    def extract_custom_embeddings(self, dataloader, embedding_type='text'):
+        """Extract embeddings from custom model"""
+        all_embeddings = []
+        all_labels = []
+        all_texts = []
+        
+        with torch.no_grad():
+            for batch in tqdm(dataloader, desc=f"Extracting custom {embedding_type} embeddings"):
+                images = batch['image'].to(self.device)
+                hierarchy_indices = batch['hierarchy_indices'].to(self.device)
+                hierarchy_labels = batch['hierarchy']
+                
+                # Forward pass
+                out = self.model(image=images, hierarchy_indices=hierarchy_indices)
+                embeddings = out['z_txt'] if embedding_type == 'text' else out['z_img'] if embedding_type == 'image' else out['z_txt']
+                
+                all_embeddings.append(embeddings.cpu().numpy())
+                all_labels.extend(hierarchy_labels)
+                all_texts.extend(hierarchy_labels)
+        
+        return np.vstack(all_embeddings), all_labels, all_texts
+
+    def compute_similarity_metrics(self, embeddings, labels):
+        """Compute intra-class and inter-class similarities"""
+        similarities = cosine_similarity(embeddings)
+        
+        # Group embeddings by hierarchy
+        hierarchy_groups = defaultdict(list)
+        for i, hierarchy in enumerate(labels):
+            hierarchy_groups[hierarchy].append(i)
+        
+        # Calculate intra-class similarities (same hierarchy)
+        intra_class_similarities = []
+        for hierarchy, indices in hierarchy_groups.items():
+            if len(indices) > 1:
+                for i in range(len(indices)):
+                    for j in range(i+1, len(indices)):
+                        sim = similarities[indices[i], indices[j]]
+                        intra_class_similarities.append(sim)
+        
+        # Calculate inter-class similarities (different hierarchies)
+        inter_class_similarities = []
+        hierarchies = list(hierarchy_groups.keys())
+        for i in range(len(hierarchies)):
+            for j in range(i+1, len(hierarchies)):
+                hierarchy1_indices = hierarchy_groups[hierarchies[i]]
+                hierarchy2_indices = hierarchy_groups[hierarchies[j]]
+                
+                for idx1 in hierarchy1_indices:
+                    for idx2 in hierarchy2_indices:
+                        sim = similarities[idx1, idx2]
+                        inter_class_similarities.append(sim)
+        
+        # Calculate classification accuracy using nearest neighbor in embedding space
+        nn_accuracy = self.compute_embedding_accuracy(embeddings, labels, similarities)
+        
+        # Calculate classification accuracy using centroids
+        centroid_accuracy = self.compute_centroid_accuracy(embeddings, labels)
+        
+        return {
+            'intra_class_similarities': intra_class_similarities,
+            'inter_class_similarities': inter_class_similarities,
+            'intra_class_mean': np.mean(intra_class_similarities) if intra_class_similarities else 0,
+            'inter_class_mean': np.mean(inter_class_similarities) if inter_class_similarities else 0,
+            'separation_score': np.mean(intra_class_similarities) - np.mean(inter_class_similarities) if intra_class_similarities and inter_class_similarities else 0,
+            'accuracy': nn_accuracy,
+            'centroid_accuracy': centroid_accuracy
+        }
+    
+    def compute_embedding_accuracy(self, embeddings, labels, similarities):
+        """Compute classification accuracy using nearest neighbor in embedding space"""
+        correct_predictions = 0
+        total_predictions = len(labels)
+        
+        for i in range(len(embeddings)):
+            true_label = labels[i]
+            
+            # Find the most similar embedding (excluding itself)
+            similarities_row = similarities[i].copy()
+            similarities_row[i] = -1  # Exclude self-similarity
+            nearest_neighbor_idx = np.argmax(similarities_row)
+            predicted_label = labels[nearest_neighbor_idx]
+            
+            if predicted_label == true_label:
+                correct_predictions += 1
+        
+        return correct_predictions / total_predictions if total_predictions > 0 else 0
+    
+    def compute_centroid_accuracy(self, embeddings, labels):
+        """Compute classification accuracy using hierarchy centroids"""
+        # Create centroids for each hierarchy
+        unique_hierarchies = list(set(labels))
+        centroids = {}
+        
+        for hierarchy in unique_hierarchies:
+            hierarchy_indices = [i for i, label in enumerate(labels) if label == hierarchy]
+            hierarchy_embeddings = embeddings[hierarchy_indices]
+            centroids[hierarchy] = np.mean(hierarchy_embeddings, axis=0)
+        
+        # Classify each embedding to nearest centroid
+        correct_predictions = 0
+        total_predictions = len(labels)
+        
+        for i, embedding in enumerate(embeddings):
+            true_label = labels[i]
+            
+            # Find closest centroid
+            best_similarity = -1
+            predicted_label = None
+            
+            for hierarchy, centroid in centroids.items():
+                similarity = cosine_similarity([embedding], [centroid])[0][0]
+                if similarity > best_similarity:
+                    best_similarity = similarity
+                    predicted_label = hierarchy
+            
+            if predicted_label == true_label:
+                correct_predictions += 1
+        
+        return correct_predictions / total_predictions if total_predictions > 0 else 0
+
+    def create_confusion_matrix(self, true_labels, predicted_labels, title="Confusion Matrix"):
+        """Create and plot confusion matrix"""
+        # Get unique labels
+        unique_labels = sorted(list(set(true_labels + predicted_labels)))
+        
+        # Create confusion matrix
+        cm = confusion_matrix(true_labels, predicted_labels, labels=unique_labels)
+        
+        # Calculate accuracy
+        accuracy = accuracy_score(true_labels, predicted_labels)
+        
+        # Plot confusion matrix
+        plt.figure(figsize=(12, 10))
+        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
+                   xticklabels=unique_labels, yticklabels=unique_labels)
+        plt.title(f'{title}\nAccuracy: {accuracy:.3f} ({accuracy*100:.1f}%)')
+        plt.ylabel('True Hierarchy')
+        plt.xlabel('Predicted Hierarchy')
+        plt.xticks(rotation=45)
+        plt.yticks(rotation=0)
+        plt.tight_layout()
+        
+        return plt.gcf(), accuracy, cm
+
+    def predict_hierarchy_from_embeddings(self, embeddings, labels):
+        """Predict hierarchy from embeddings using centroid-based classification"""
+        # Create hierarchy centroids from training data
+        unique_hierarchies = list(set(labels))
+        centroids = {}
+        
+        for hierarchy in unique_hierarchies:
+            hierarchy_indices = [i for i, label in enumerate(labels) if label == hierarchy]
+            hierarchy_embeddings = embeddings[hierarchy_indices]
+            centroids[hierarchy] = np.mean(hierarchy_embeddings, axis=0)
+        
+        # Predict hierarchy for all embeddings
+        predictions = []
+        
+        for i, embedding in enumerate(embeddings):
+            # Find closest centroid
+            best_similarity = -1
+            predicted_hierarchy = None
+            
+            for hierarchy, centroid in centroids.items():
+                similarity = cosine_similarity([embedding], [centroid])[0][0]
+                if similarity > best_similarity:
+                    best_similarity = similarity
+                    predicted_hierarchy = hierarchy
+            
+            predictions.append(predicted_hierarchy)
+        
+        return predictions
+
+    def evaluate_classification_performance(self, embeddings, labels, embedding_type="Embeddings"):
+        """Evaluate classification performance and create confusion matrix"""
+        # Predict hierarchy
+        predictions = self.predict_hierarchy_from_embeddings(embeddings, labels)
+        
+        # Calculate accuracy
+        accuracy = accuracy_score(labels, predictions)
+        
+        # Calculate F1 scores
+        unique_labels = sorted(list(set(labels)))
+        f1_macro = f1_score(labels, predictions, labels=unique_labels, average='macro', zero_division=0)
+        f1_weighted = f1_score(labels, predictions, labels=unique_labels, average='weighted', zero_division=0)
+        f1_per_class = f1_score(labels, predictions, labels=unique_labels, average=None, zero_division=0)
+        
+        # Create confusion matrix
+        fig, acc, cm = self.create_confusion_matrix(labels, predictions, 
+                                                   f"{embedding_type} - Hierarchy Classification")
+        
+        # Generate classification report
+        report = classification_report(labels, predictions, labels=unique_labels, 
+                                     target_names=unique_labels, output_dict=True)
+        
+        return {
+            'accuracy': accuracy,
+            'f1_macro': f1_macro,
+            'f1_weighted': f1_weighted,
+            'f1_per_class': f1_per_class,
+            'predictions': predictions,
+            'confusion_matrix': cm,
+            'classification_report': report,
+            'figure': fig
+        }
+
+    def evaluate_dataset_with_baselines(self, dataframe, dataset_name="Dataset"):
+        """Evaluate embeddings on a given dataset with both custom model and CLIP baseline"""
+        print(f"\n{'='*60}")
+        print(f"Evaluating {dataset_name}")
+        print(f"{'='*60}")
+        
+        results = {}
+        
+        # ===== CUSTOM MODEL EVALUATION =====
+        print(f"\n🔧 Evaluating Custom Model on {dataset_name}")
+        print("-" * 40)
+        
+        # Create dataloader for custom model
+        custom_dataloader = self.create_dataloader(dataframe, batch_size=16)
+        
+        # Evaluate text embeddings
+        text_embeddings, text_labels, texts = self.extract_custom_embeddings(custom_dataloader, 'text')
+        text_metrics = self.compute_similarity_metrics(text_embeddings, text_labels)
+        text_classification = self.evaluate_classification_performance(text_embeddings, text_labels, "Custom Text Embeddings")
+        text_metrics.update(text_classification)
+        results['custom_text'] = text_metrics
+        
+        # Evaluate image embeddings
+        image_embeddings, image_labels, _ = self.extract_custom_embeddings(custom_dataloader, 'image')
+        image_metrics = self.compute_similarity_metrics(image_embeddings, image_labels)
+        image_classification = self.evaluate_classification_performance(image_embeddings, image_labels, "Custom Image Embeddings")
+        image_metrics.update(image_classification)
+        results['custom_image'] = image_metrics
+        
+        # ===== CLIP BASELINE EVALUATION =====
+        print(f"\n🤗 Evaluating CLIP Baseline on {dataset_name}")
+        print("-" * 40)
+        
+        # Create dataloader for CLIP
+        clip_dataloader = self.create_clip_dataloader(dataframe, batch_size=8)  # Smaller batch for CLIP
+        
+        # Extract data for CLIP
+        all_images = []
+        all_texts = []
+        all_labels = []
+        
+        for batch in tqdm(clip_dataloader, desc="Preparing data for CLIP"):
+            images, texts, labels = batch
+            all_images.extend(images)
+            all_texts.extend(texts)
+            all_labels.extend(labels)
+        
+        # Get CLIP embeddings
+        clip_image_embeddings, clip_text_embeddings = self.clip_evaluator.extract_clip_embeddings(all_images, all_texts)
+        
+        # Evaluate CLIP text embeddings
+        clip_text_metrics = self.compute_similarity_metrics(clip_text_embeddings, all_labels)
+        clip_text_classification = self.evaluate_classification_performance(clip_text_embeddings, all_labels, "CLIP Text Embeddings")
+        clip_text_metrics.update(clip_text_classification)
+        results['clip_text'] = clip_text_metrics
+        
+        # Evaluate CLIP image embeddings
+        clip_image_metrics = self.compute_similarity_metrics(clip_image_embeddings, all_labels)
+        clip_image_classification = self.evaluate_classification_performance(clip_image_embeddings, all_labels, "CLIP Image Embeddings")
+        clip_image_metrics.update(clip_image_classification)
+        results['clip_image'] = clip_image_metrics
+        
+        # ===== PRINT COMPARISON RESULTS =====
+        print(f"\n{dataset_name} Results Comparison:")
+        print(f"Dataset size: {len(dataframe)} samples")
+        print("=" * 80)
+        print(f"{'Model':<20} {'Embedding':<10} {'Sep Score':<10} {'NN Acc':<8} {'Centroid Acc':<12} {'F1 Macro':<10}")
+        print("-" * 80)
+        
+        for model_type in ['custom', 'clip']:
+            for emb_type in ['text', 'image']:
+                key = f"{model_type}_{emb_type}"
+                if key in results:
+                    metrics = results[key]
+                    model_name = "Custom Model" if model_type == 'custom' else "CLIP Baseline"
+                    print(f"{model_name:<20} {emb_type.capitalize():<10} {metrics['separation_score']:<10.4f} {metrics['accuracy']*100:<8.1f}% {metrics['centroid_accuracy']*100:<12.1f}% {metrics['f1_macro']*100:<10.1f}%")
+        
+        # ===== SAVE VISUALIZATIONS =====
+        os.makedirs(f'{self.directory}', exist_ok=True)
+        
+        # Save confusion matrices
+        for key, metrics in results.items():
+            if 'figure' in metrics:
+                metrics['figure'].savefig(f'{self.directory}/{dataset_name.lower()}_{key}_confusion_matrix.png', dpi=300, bbox_inches='tight')
+                plt.close(metrics['figure'])
+        
+        return results
+
+
+class CLIPDataset(Dataset):
+    def __init__(self, dataframe):
+        self.dataframe = dataframe
+        # Use VALIDATION transforms (no augmentation)
+        self.transform = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        ])
+
+    def __len__(self):
+        return len(self.dataframe)
+
+    def __getitem__(self, idx):
+        row = self.dataframe.iloc[idx]
+        
+        # Handle image loading (same as HierarchyDataset)
+        if config.column_local_image_path in row.index and pd.notna(row[config.column_local_image_path]):
+            local_path = row[config.column_local_image_path]
+            try:
+                if os.path.exists(local_path):
+                    image = Image.open(local_path).convert("RGB")
+                else:
+                    print(f"⚠️ Local image not found: {local_path}")
+                    image = Image.new('RGB', (224, 224), color='gray')
+            except Exception as e:
+                print(f"⚠️ Failed to load local image {idx}: {e}")
+                image = Image.new('RGB', (224, 224), color='gray')
+        elif isinstance(row[config.column_url_image], dict):
+            image = Image.open(BytesIO(row[config.column_url_image]['bytes'])).convert('RGB')
+        elif isinstance(row['image_url'], (list, np.ndarray)):
+            pixels = np.array(row[config.column_url_image]).reshape(28, 28)
+            image = Image.fromarray(pixels.astype(np.uint8)).convert("RGB")
+        elif isinstance(row[config.column_url_image], Image.Image):
+            # Handle PIL Image objects directly (for Fashion-MNIST)
+            image = row[config.column_url_image].convert("RGB")
+        else:
+            try:
+                response = requests.get(row[config.column_url_image], timeout=10)
+                response.raise_for_status()
+                image = Image.open(BytesIO(response.content)).convert("RGB")
+            except Exception as e:
+                print(f"⚠️ Failed to load image {idx}: {e}")
+                image = Image.new('RGB', (224, 224), color='gray')
+        
+        # Apply transforms
+        image_tensor = self.transform(image)
+
+        description = row[config.text_column]  
+        hierarchy = row[config.hierarchy_column]
+
+        return image_tensor, description, hierarchy
+
+
+def load_fashion_mnist_dataset(evaluator):
+    """Load and prepare Fashion-MNIST test dataset"""
+    print("Loading Fashion-MNIST test dataset...")
+    
+    # Load the dataset
+    df = pd.read_csv(config.fashion_mnist_test_path)
+    print(f"✅ Fashion-MNIST dataset loaded")
+    print(f"📊 Total samples: {len(df)}")
+    
+    # Fashion-MNIST class labels mapping
+    fashion_mnist_labels = get_fashion_mnist_labels()
+    
+    # Map labels to hierarchy classes
+    hierarchy_mapping = {
+        'T-shirt/top': 'top',
+        'Trouser': 'bottom', 
+        'Pullover': 'top',
+        'Dress': 'dress',
+        'Coat': 'top',
+        'Sandal': 'shoes',
+        'Shirt': 'top',
+        'Sneaker': 'shoes',
+        'Bag': 'bag',
+        'Ankle boot': 'shoes'
+    }
+    
+    # Apply label mapping
+    df['hierarchy'] = df['label'].map(fashion_mnist_labels).map(hierarchy_mapping)
+    
+    # Filter to only include hierarchies that exist in our model
+    valid_hierarchies = df['hierarchy'].dropna().unique()
+    print(f"🎯 Valid hierarchies found: {sorted(valid_hierarchies)}")
+    print(f"🎯 Model hierarchies: {sorted(evaluator.hierarchy_classes)}")
+    
+    # Filter to only include hierarchies that exist in our model
+    df = df[df['hierarchy'].isin(evaluator.hierarchy_classes)]
+    print(f"📊 After filtering to model hierarchies: {len(df)} samples")
+    
+    if len(df) == 0:
+        print("❌ No samples left after hierarchy filtering.")
+        return pd.DataFrame()
+    
+    # Keep pixel columns as they are (FashionMNISTDataset will handle them)
+    
+    # Create text descriptions based on hierarchy
+    text_descriptions = {
+        'top': 'A top clothing item',
+        'bottom': 'A bottom clothing item', 
+        'dress': 'A dress',
+        'shoes': 'A pair of shoes',
+        'bag': 'A bag'
+    }
+    
+    df['text'] = df['hierarchy'].map(text_descriptions)
+    
+    # Show sample of data
+    print(f"📝 Sample data:")
+    for i, (hierarchy, text) in enumerate(zip(df['hierarchy'].head(3), df['text'].head(3))):
+        print(f"  {i+1}. [{hierarchy}] {text}")
+    
+    df_test = df.copy()
+    
+    print(f"📊 After sampling: {len(df_test)} samples")
+    print(f"📊 Samples per hierarchy:")
+    for hierarchy in sorted(df_test['hierarchy'].unique()):
+        count = len(df_test[df_test['hierarchy'] == hierarchy])
+        print(f"  {hierarchy}: {count} samples")
+    
+    # Create formatted dataset with proper column names
+    # Keep all pixel columns for FashionMNISTDataset
+    pixel_cols = [f'pixel{i}' for i in range(1, 785)]
+    fashion_mnist_formatted = df_test[['label'] + pixel_cols + ['text', 'hierarchy']].copy()
+    
+    print(f"📊 Final dataset size: {len(fashion_mnist_formatted)} samples")
+    return fashion_mnist_formatted
+
+
+def load_kagl_marqo_dataset(evaluator):
+    """Load and prepare kagl dataset"""
+    from datasets import load_dataset
+    print("Loading kagl dataset...")
+
+    # Load the dataset
+    dataset = load_dataset("Marqo/KAGL")
+    df = dataset["data"].to_pandas()
+    print(f"✅ Dataset kagl loaded")
+    print(f"📊 Before filtering: {len(df)} samples")
+    print(f"📋 Available columns: {list(df.columns)}")
+
+    # Check available categories and map them to our hierarchy
+    print(f"🎨 Available categories: {sorted(df['category2'].unique())}")
+    # Apply mapping
+    df['hierarchy'] = df['category2'].str.lower()
+    df['hierarchy'] = df['hierarchy'].replace('bags', 'bag').replace('topwear', 'top').replace('flip flops', 'shoes').replace('sandal', 'shoes')
+
+    # Filter to only include valid hierarchies that exist in our model
+    valid_hierarchies = df['hierarchy'].dropna().unique()
+    print(f"🎯 Valid hierarchies found: {sorted(valid_hierarchies)}")
+    print(f"🎯 Model hierarchies: {sorted(evaluator.hierarchy_classes)}")
+    
+    # Filter to only include hierarchies that exist in our model
+    df = df[df['hierarchy'].isin(evaluator.hierarchy_classes)]
+    print(f"📊 After filtering to model hierarchies: {len(df)} samples")
+    
+    if len(df) == 0:
+        print("❌ No samples left after hierarchy filtering.")
+        return pd.DataFrame()
+    
+    # Ensure we have text and image data
+    df = df.dropna(subset=['text', 'image'])
+    print(f"📊 After removing missing text/image: {len(df)} samples")
+    
+    # Show sample of text data to verify quality
+    print(f"📝 Sample texts:")
+    for i, (text, hierarchy) in enumerate(zip(df['text'].head(3), df['hierarchy'].head(3))):
+        print(f"  {i+1}. [{hierarchy}] {text[:100]}...")
+    
+    print(f"📊 After sampling: {len(df_test)} samples")
+    print(f"📊 Samples per hierarchy:")
+    for hierarchy in sorted(df_test['hierarchy'].unique()):
+        count = len(df_test[df_test['hierarchy'] == hierarchy])
+        print(f"  {hierarchy}: {count} samples")
+    
+    # Create formatted dataset with proper column names
+    kagl_formatted = pd.DataFrame({
+        'image_url': df_test['image'],
+        'text': df_test['text'],
+        'hierarchy': df_test['hierarchy']
+    })
+    
+    print(f"📊 Final dataset size: {len(kagl_formatted)} samples")
+    return kagl_formatted
+
+
+if __name__ == "__main__":
+    device = config.device
+    directory = config.evaluation_directory
+
+    print(f"🚀 Starting evaluation with custom model: {config.hierarchy_model_path}")
+    print(f"🤗 Including CLIP baseline comparison")
+    
+    evaluator = EmbeddingEvaluator(config.hierarchy_model_path, directory, device=device)
+    
+    print(f"📊 Final hierarchy classes after initialization: {len(evaluator.vocab.hierarchy_classes)} classes")
+    
+    # Evaluate on validation dataset (same subset as during training)
+    print("\n" + "="*60)
+    print("EVALUATING VALIDATION DATASET - CUSTOM MODEL vs CLIP BASELINE")
+    print("="*60)
+    val_results = evaluator.evaluate_dataset_with_baselines(evaluator.val_df, "Validation Dataset")
+    
+    print("\n" + "="*60)
+    print("EVALUATING FASHION-MNIST TEST DATASET - CUSTOM MODEL vs CLIP BASELINE") 
+    print("="*60)
+    df_fashion_mnist = load_fashion_mnist_dataset(evaluator)
+    if len(df_fashion_mnist) > 0:
+        fashion_mnist_results = evaluator.evaluate_dataset_with_baselines(df_fashion_mnist, "Fashion-MNIST Test Dataset")
+    else:
+        fashion_mnist_results = {}
+    
+    print("\n" + "="*60)
+    print("EVALUATING kagl MARQO DATASET - CUSTOM MODEL vs CLIP BASELINE") 
+    print("="*60)
+    df_kagl_marqo = load_kagl_marqo_dataset(evaluator)
+    if len(df_kagl_marqo) > 0:
+        kagl_results = evaluator.evaluate_dataset_with_baselines(df_kagl_marqo, "kagl Marqo Dataset")
+    else:
+        kagl_results = {}
+    
+    # Compare results
+    print(f"\n{'='*80}")
+    print("FINAL EVALUATION SUMMARY - CUSTOM MODEL vs CLIP BASELINE")
+    print(f"{'='*80}")
+    
+    print("\n🔍 VALIDATION DATASET RESULTS:")
+    print(f"Dataset size: {len(evaluator.val_df)} samples")
+    print(f"{'Model':<20} {'Embedding':<10} {'Sep Score':<12} {'NN Acc':<10} {'Centroid Acc':<12} {'F1 Macro':<10}")
+    print("-" * 80)
+    
+    for model_type in ['custom', 'clip']:
+        for emb_type in ['text', 'image']:
+            key = f"{model_type}_{emb_type}"
+            if key in val_results:
+                metrics = val_results[key]
+                model_name = "Custom Model" if model_type == 'custom' else "CLIP Baseline"
+                print(f"{model_name:<20} {emb_type.capitalize():<10} {metrics['separation_score']:<12.4f} {metrics['accuracy']*100:<10.1f}% {metrics['centroid_accuracy']*100:<12.1f}% {metrics['f1_macro']*100:<10.1f}%")
+
+    if fashion_mnist_results:
+        print("\n👗 FASHION-MNIST TEST DATASET RESULTS:")
+        print(f"Dataset size: {len(df_fashion_mnist)} samples")
+        print(f"{'Model':<20} {'Embedding':<10} {'Sep Score':<12} {'NN Acc':<10} {'Centroid Acc':<12} {'F1 Macro':<10}")
+        print("-" * 80)
+        
+        for model_type in ['custom', 'clip']:
+            for emb_type in ['text', 'image']:
+                key = f"{model_type}_{emb_type}"
+                if key in fashion_mnist_results:
+                    metrics = fashion_mnist_results[key]
+                    model_name = "Custom Model" if model_type == 'custom' else "CLIP Baseline"
+                    print(f"{model_name:<20} {emb_type.capitalize():<10} {metrics['separation_score']:<12.4f} {metrics['accuracy']*100:<10.1f}% {metrics['centroid_accuracy']*100:<12.1f}% {metrics['f1_macro']*100:<10.1f}%")
+    
+    if kagl_results:
+        print("\n🌐 kagl MARQO DATASET RESULTS:")
+        print(f"Dataset size: {len(df_kagl_marqo)} samples")
+        print(f"{'Model':<20} {'Embedding':<10} {'Sep Score':<12} {'NN Acc':<10} {'Centroid Acc':<12} {'F1 Macro':<10}")
+        print("-" * 80)
+        
+        for model_type in ['custom', 'clip']:
+            for emb_type in ['text', 'image']:
+                key = f"{model_type}_{emb_type}"
+                if key in kagl_results:
+                    metrics = kagl_results[key]
+                    model_name = "Custom Model" if model_type == 'custom' else "CLIP Baseline"
+                    print(f"{model_name:<20} {emb_type.capitalize():<10} {metrics['separation_score']:<12.4f} {metrics['accuracy']*100:<10.1f}% {metrics['centroid_accuracy']*100:<12.1f}% {metrics['f1_macro']*100:<10.1f}%")
+    
+    print(f"\n✅ Evaluation completed! Check '{directory}/' for visualization files.")
+    print(f"📊 Custom model hierarchy classes: {len(evaluator.vocab.hierarchy_classes)} classes")
+    print(f"🤗 CLIP baseline comparison included")