Consolidate embeddings

.dockerignore

@@ -10,12 +10,19 @@ pipeline/
 # OS cruft
 .DS_Store
 
-# Exclude large files but allow model files and embedding files
+# Exclude only large files that aren't essential
 data/*.bin
 data/*.pdf
 
-# Keep embedding files (.pt) for ML inference
-# data/*.pt  <- This line is removed to allow .pt files
+# Exclude individual embedding files (too many for HF Spaces)
+data/embeddings/CLIP_Embeddings/
+data/embeddings/PaintingCLIP_Embeddings/
 
-# Keep model files (.safetensors) for PaintingCLIP LoRA
-# data/*.safetensors  <- This line is removed to allow .safetensors files
+# Allow consolidated embedding files
+!data/embeddings/*_consolidated.pt
+!data/embeddings/*_metadata.json
+
+# Keep essential ML files:
+# - .pt files for embeddings (consolidated)
+# - .safetensors files for models
+# - JSON files for metadata

.gitignore

@@ -4,6 +4,14 @@
 data/*.bin                    # Exclude large binary files
 data/*.pdf                    # Exclude PDF files
 
+# Exclude individual embedding files (too many for HF Spaces)
+data/embeddings/CLIP_Embeddings/
+data/embeddings/PaintingCLIP_Embeddings/
+
+# Allow consolidated embedding files
+!data/embeddings/*_consolidated.pt
+!data/embeddings/*_metadata.json
+
 # Pipeline (exclude entire directory for Phase 1)
 pipeline/
 pipeline/slurm/

backend/runner/inference.py

@@ -446,3 +446,39 @@ def set_model_type(model_type: str) -> None:
     MODEL_TYPE = model_type
     # Clear the cache to force reinitialization
     _initialize_pipeline.cache_clear()
+
+
+def load_consolidated_embeddings(embedding_file: Path, metadata_file: Path):
+    """Load embeddings from consolidated file with metadata"""
+    print(f"Loading consolidated embeddings from {embedding_file}")
+    
+    # Load consolidated data
+    consolidated_data = torch.load(embedding_file, map_location='cpu')
+    embeddings = consolidated_data['embeddings']
+    
+    # Load metadata for file mapping
+    with open(metadata_file, 'r', encoding='utf-8') as f:
+        metadata = json.load(f)
+    
+    # Create filename to index mapping
+    filename_to_index = {item['filename']: item['index'] for item in metadata['file_mapping']}
+    
+    print(f"Loaded {len(embeddings)} embeddings with dimension {embeddings.shape[1]}")
+    
+    return embeddings, filename_to_index
+
+# Update your embedding loading logic
+def load_embeddings_for_model(model_type: str):
+    """Load embeddings for the specified model type"""
+    if model_type == "clip":
+        embedding_file = CLIP_EMBEDDINGS_DIR / "clip_embeddings_consolidated.pt"
+        metadata_file = CLIP_EMBEDDINGS_DIR / "clip_embeddings_metadata.json"
+    else:  # paintingclip
+        embedding_file = PAINTINGCLIP_EMBEDDINGS_DIR / "paintingclip_embeddings_consolidated.pt"
+        metadata_file = PAINTINGCLIP_EMBEDDINGS_DIR / "paintingclip_embeddings_metadata.json"
+    
+    if not embedding_file.exists():
+        print(f"Warning: Consolidated embedding file not found: {embedding_file}")
+        return None, None
+    
+    return load_consolidated_embeddings(embedding_file, metadata_file)

consolidate_embeddings.py

@@ -0,0 +1,156 @@
+#!/usr/bin/env python3
+"""
+Consolidate individual embedding .pt files into larger consolidated files.
+This solves the Hugging Face Spaces 10,000 files per directory limit.
+"""
+
+import torch
+import os
+import json
+from pathlib import Path
+from typing import Dict, List, Tuple
+import argparse
+
+def consolidate_embeddings(
+    input_dir: Path, 
+    output_file: Path, 
+    metadata_file: Path,
+    batch_size: int = 1000
+) -> Dict[str, int]:
+    """
+    Merge individual .pt files into one large tensor file with metadata.
+    
+    Args:
+        input_dir: Directory containing individual .pt files
+        output_file: Path to save consolidated tensor
+        metadata_file: Path to save file mapping metadata
+        batch_size: Process files in batches to manage memory
+    
+    Returns:
+        Dict with statistics about the consolidation
+    """
+    embedding_files = sorted(list(input_dir.glob("*.pt")))
+    
+    if not embedding_files:
+        raise ValueError(f"No .pt files found in {input_dir}")
+    
+    print(f"Found {len(embedding_files)} embedding files in {input_dir}")
+    
+    # Load first file to get embedding dimension
+    print("Loading first embedding to determine dimensions...")
+    first_embedding = torch.load(embedding_files[0])
+    embedding_dim = first_embedding.shape[0]
+    print(f"Embedding dimension: {embedding_dim}")
+    
+    # Pre-allocate tensor
+    all_embeddings = torch.zeros(len(embedding_files), embedding_dim, dtype=first_embedding.dtype)
+    file_mapping = []
+    
+    print(f"Consolidating {len(embedding_files)} embeddings...")
+    
+    for i, file_path in enumerate(embedding_files):
+        if i % 1000 == 0:
+            print(f"Processing {i}/{len(embedding_files)} ({i/len(embedding_files)*100:.1f}%)")
+        
+        try:
+            embedding = torch.load(file_path)
+            all_embeddings[i] = embedding
+            
+            # Store file mapping for later lookup
+            file_mapping.append({
+                'index': i,
+                'filename': file_path.name,
+                'stem': file_path.stem,
+                'file_size': file_path.stat().st_size
+            })
+            
+        except Exception as e:
+            print(f"Error loading {file_path}: {e}")
+            # Fill with zeros if file is corrupted
+            all_embeddings[i] = torch.zeros(embedding_dim, dtype=first_embedding.dtype)
+    
+    # Save consolidated data
+    print(f"Saving consolidated embeddings to {output_file}...")
+    consolidated_data = {
+        'embeddings': all_embeddings,
+        'embedding_dim': embedding_dim,
+        'num_embeddings': len(embedding_files),
+        'dtype': str(first_embedding.dtype)
+    }
+    
+    torch.save(consolidated_data, output_file)
+    
+    # Save metadata for lookup
+    print(f"Saving metadata to {metadata_file}...")
+    metadata = {
+        'input_directory': str(input_dir),
+        'output_file': str(output_file),
+        'num_embeddings': len(embedding_files),
+        'embedding_dim': embedding_dim,
+        'dtype': str(first_embedding.dtype),
+        'file_mapping': file_mapping
+    }
+    
+    with open(metadata_file, 'w', encoding='utf-8') as f:
+        json.dump(metadata, f, indent=2, ensure_ascii=False)
+    
+    # Calculate file sizes
+    original_size = sum(f.stat().st_size for f in embedding_files)
+    consolidated_size = output_file.stat().st_size
+    metadata_size = metadata_file.stat().st_size
+    
+    stats = {
+        'num_files_processed': len(embedding_files),
+        'original_size_mb': original_size / (1024 * 1024),
+        'consolidated_size_mb': consolidated_size / (1024 * 1024),
+        'metadata_size_kb': metadata_size / 1024,
+        'compression_ratio': original_size / consolidated_size if consolidated_size > 0 else 0
+    }
+    
+    print(f"\nConsolidation complete!")
+    print(f"Files processed: {stats['num_files_processed']}")
+    print(f"Original size: {stats['original_size_mb']:.1f} MB")
+    print(f"Consolidated size: {stats['consolidated_size_mb']:.1f} MB")
+    print(f"Metadata size: {stats['metadata_size_kb']:.1f} KB")
+    print(f"Compression ratio: {stats['compression_ratio']:.2f}x")
+    
+    return stats
+
+def main():
+    parser = argparse.ArgumentParser(description='Consolidate embedding files')
+    parser.add_argument('--input-dir', type=str, required=True, 
+                       help='Input directory containing .pt files')
+    parser.add_argument('--output-file', type=str, required=True,
+                       help='Output consolidated .pt file')
+    parser.add_argument('--metadata-file', type=str, required=True,
+                       help='Output metadata JSON file')
+    parser.add_argument('--batch-size', type=int, default=1000,
+                       help='Batch size for processing (default: 1000)')
+    
+    args = parser.parse_args()
+    
+    input_dir = Path(args.input_dir)
+    output_file = Path(args.output_file)
+    metadata_file = Path(args.metadata_file)
+    
+    if not input_dir.exists():
+        print(f"Error: Input directory {input_dir} does not exist")
+        return 1
+    
+    # Create output directory if it doesn't exist
+    output_file.parent.mkdir(parents=True, exist_ok=True)
+    
+    try:
+        stats = consolidate_embeddings(
+            input_dir=input_dir,
+            output_file=output_file,
+            metadata_file=metadata_file,
+            batch_size=args.batch_size
+        )
+        return 0
+    except Exception as e:
+        print(f"Error during consolidation: {e}")
+        return 1
+
+if __name__ == "__main__":
+    exit(main())

data/embeddings/clip_embeddings_consolidated.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ebcce37ea66ceb8417a2e52e80e38e1c2970e7ef6026b5546f57a9a09f2f3c85
+size 60604

data/embeddings/clip_embeddings_metadata.json

@@ -0,0 +1,8 @@
+{
+  "input_directory": "data/embeddings/CLIP_Embeddings",
+  "output_file": "data/embeddings/clip_embeddings_consolidated.pt",
+  "num_embeddings": 14674,
+  "embedding_dim": 1,
+  "dtype": "torch.float32",
+  "file_mapping": []
+}

data/embeddings/paintingclip_embeddings_consolidated.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:303a47cd9c8b24cead5d27b7e6f23e99c2fc400b03028034447374c19f9adfba
+size 60660

data/embeddings/paintingclip_embeddings_metadata.json

@@ -0,0 +1,8 @@
+{
+  "input_directory": "data/embeddings/PaintingCLIP_Embeddings",
+  "output_file": "data/embeddings/paintingclip_embeddings_consolidated.pt",
+  "num_embeddings": 14674,
+  "embedding_dim": 1,
+  "dtype": "torch.float32",
+  "file_mapping": []
+}