Use safetensors for consolidated embeddings; load via config; add ids JSON

backend/runner/config.py

@@ -66,3 +66,10 @@ WORKS_JSON = JSON_INFO_DIR / "works.json"
 TOPICS_JSON = JSON_INFO_DIR / "topics.json"
 CREATORS_JSON = JSON_INFO_DIR / "creators.json"
 TOPIC_NAMES_JSON = JSON_INFO_DIR / "topic_names.json"
+
+# Add below existing EMBEDDINGS_DIR constants
+CLIP_EMBEDDINGS_ST = EMBEDDINGS_DIR / "clip_embeddings.safetensors"
+CLIP_SENTENCE_IDS = EMBEDDINGS_DIR / "clip_embeddings_sentence_ids.json"
+
+PAINTINGCLIP_EMBEDDINGS_ST = EMBEDDINGS_DIR / "paintingclip_embeddings.safetensors"
+PAINTINGCLIP_SENTENCE_IDS = EMBEDDINGS_DIR / "paintingclip_embeddings_sentence_ids.json"

backend/runner/inference.py

@@ -25,6 +25,7 @@ import torch.nn.functional as F
 from peft import PeftModel
 from PIL import Image
 from transformers import CLIPModel, CLIPProcessor
+from safetensors.torch import load_file as st_load_file
 
 from .filtering import get_filtered_sentence_ids
 # on-demand Grad-ECLIP & region-aware ranking
@@ -34,7 +35,9 @@ from .config import (
     PAINTINGCLIP_EMBEDDINGS_DIR,
     PAINTINGCLIP_MODEL_DIR,
     SENTENCES_JSON,
-    EMBEDDINGS_DIR  # ← Add this line
+    EMBEDDINGS_DIR,
+    CLIP_EMBEDDINGS_ST, CLIP_SENTENCE_IDS,
+    PAINTINGCLIP_EMBEDDINGS_ST, PAINTINGCLIP_SENTENCE_IDS,
 )
 
 # ─── Configuration ───────────────────────────────────────────────────────────
@@ -456,33 +459,59 @@ def load_consolidated_embeddings(embedding_file: Path, metadata_file: Path):
     
     return embeddings, filename_to_index
 
+def load_consolidated_embeddings_st(embedding_st_file: Path, ids_json_file: Path):
+	print(f"Loading safetensors embeddings from {embedding_st_file}")
+	if not embedding_st_file.exists():
+		raise FileNotFoundError(f"Missing {embedding_st_file}")
+	if not ids_json_file.exists():
+		raise FileNotFoundError(f"Missing {ids_json_file}")
+
+	data = st_load_file(str(embedding_st_file))
+	if "embeddings" not in data:
+		raise KeyError(f"'embeddings' tensor missing in {embedding_st_file}")
+	embeddings = data["embeddings"].to(dtype=torch.float32, device="cpu").contiguous()
+
+	with open(ids_json_file, "r", encoding="utf-8") as f:
+		sentence_ids = json.load(f)
+	if not isinstance(sentence_ids, list):
+		raise ValueError(f"IDs file malformed: {ids_json_file}")
+
+	print(f"Loaded {len(sentence_ids)} embeddings with dim {embeddings.shape[1]}")
+	return embeddings, sentence_ids
+
 # Update your embedding loading logic
 def load_embeddings_for_model(model_type: str):
-    """Load embeddings for the specified model type"""
-    if model_type == "clip":
-        # Consolidated files are in the root embeddings directory
-        embedding_file = EMBEDDINGS_DIR / "clip_embeddings_consolidated.pt"
-        metadata_file = EMBEDDINGS_DIR / "clip_embeddings_metadata.json"
-    else:  # paintingclip
-        # Consolidated files are in the root embeddings directory
-        embedding_file = EMBEDDINGS_DIR / "paintingclip_embeddings_consolidated.pt"
-        metadata_file = EMBEDDINGS_DIR / "paintingclip_embeddings_metadata.json"
-    
-    print(f"🔍 Looking for embeddings at: {embedding_file}")
-    print(f"🔍 Looking for metadata at: {metadata_file}")
-    
-    if not embedding_file.exists():
-        print(f"❌ Consolidated embedding file not found: {embedding_file}")
-        print(f" Available files in embeddings directory:")
-        for file in EMBEDDINGS_DIR.iterdir():
-            print(f"   - {file.name}")
-        return None, None
-    
-    if not metadata_file.exists():
-        print(f"❌ Metadata file not found: {metadata_file}")
-        return None, None
-    
-    print(f"✅ Found embedding file: {embedding_file}")
-    print(f"✅ Found metadata file: {metadata_file}")
-    
-    return load_consolidated_embeddings(embedding_file, metadata_file)
+	"""Load embeddings for the specified model type with safetensors-first strategy."""
+	if model_type == "clip":
+		st_file = CLIP_EMBEDDINGS_ST
+		ids_file = CLIP_SENTENCE_IDS
+		# Legacy PT locations for fallback (if repo still has them)
+		pt_file = EMBEDDINGS_DIR / "clip_embeddings_consolidated.pt"
+		meta_file = EMBEDDINGS_DIR / "clip_embeddings_metadata.json"
+		indiv_dir = CLIP_EMBEDDINGS_DIR
+	else:
+		st_file = PAINTINGCLIP_EMBEDDINGS_ST
+		ids_file = PAINTINGCLIP_SENTENCE_IDS
+		pt_file = EMBEDDINGS_DIR / "paintingclip_embeddings_consolidated.pt"
+		meta_file = EMBEDDINGS_DIR / "paintingclip_embeddings_metadata.json"
+		indiv_dir = PAINTINGCLIP_EMBEDDINGS_DIR
+
+	# 1) safetensors
+	if st_file.exists() and ids_file.exists():
+		try:
+			return load_consolidated_embeddings_st(st_file, ids_file)
+		except Exception as e:
+			print(f"⚠️  Safetensors load failed: {e}")
+
+	# 2) legacy PT (if present)
+	if pt_file.exists() and meta_file.exists():
+		try:
+			return load_consolidated_embeddings(pt_file, meta_file)
+		except Exception as e:
+			print(f"⚠️  Legacy PT load failed: {e}")
+
+	# 3) final fallback: refuse to scan 10k files here (HF Spaces file count limits)
+	print("❌ No valid consolidated embeddings found. Make sure you committed:")
+	print(f"   - {st_file.name}")
+	print(f"   - {ids_file.name}")
+	return None, None

consolidate_embeddings.py

@@ -1,156 +1,81 @@
 #!/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
+import sys
 from pathlib import Path
-from typing import Dict, List, Tuple
-import argparse
+from typing import List, Tuple
+
+import torch
+from safetensors.torch import save_file
+
+ROOT = Path(__file__).resolve().parent
+DATA_DIR = ROOT / "data" / "embeddings"
+CLIP_DIR = DATA_DIR / "CLIP_Embeddings"
+PAINTINGCLIP_DIR = DATA_DIR / "PaintingCLIP_Embeddings"
+
+def load_one(pt_path: Path) -> torch.Tensor:
+	"""Load a single .pt embedding, handling dict-or-tensor variants."""
+	obj = torch.load(pt_path, map_location="cpu", weights_only=True)
+	if isinstance(obj, torch.Tensor):
+		return obj
+	if isinstance(obj, dict):
+		for k in ("embedding", "embeddings", "features"):
+			if k in obj:
+				t = obj[k]
+				if isinstance(t, torch.Tensor):
+					return t
+	raise ValueError(f"Unsupported .pt content in {pt_path}")
 
-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 derive_id_from_filename(stem: str) -> str:
+	"""
+	- CLIP:       Wxxxx_sYYYY_clip → Wxxxx_sYYYY
+	- PaintingCLIP: Wxxxx_sYYYY_painting_clip → Wxxxx_sYYYY
+	"""
+	if stem.endswith("_painting_clip"):
+		return stem[: -len("_painting_clip")]
+	if stem.endswith("_clip"):
+		return stem[: -len("_clip")]
+	return stem  # fallback
+
+def consolidate_dir(indir: Path) -> Tuple[torch.Tensor, List[str]]:
+	pt_files = sorted(indir.glob("*.pt"))
+	if not pt_files:
+		raise RuntimeError(f"No .pt files found under {indir}")
+
+	embs: List[torch.Tensor] = []
+	ids: List[str] = []
+
+	for i, p in enumerate(pt_files, 1):
+		e = load_one(p).float()
+		if e.ndim > 1:
+			e = e.squeeze()
+		if e.ndim != 1:
+			raise ValueError(f"Embedding is not 1D in {p}: shape={tuple(e.shape)}")
+		embs.append(e)
+		ids.append(derive_id_from_filename(p.stem))
+		if i % 1000 == 0:
+			print(f"... processed {i} files from {indir}")
+
+	# Stack to [N, D]
+	embeddings = torch.stack(embs, dim=0).contiguous()
+	return embeddings, ids
+
+def save_as_safetensors(embeddings: torch.Tensor, ids: List[str], out_prefix: Path) -> None:
+	out_st = out_prefix.with_suffix(".safetensors")
+	out_json = out_prefix.with_name(out_prefix.name + "_sentence_ids.json")
+	save_file({"embeddings": embeddings}, str(out_st))
+	with open(out_json, "w", encoding="utf-8") as f:
+		json.dump(ids, f, ensure_ascii=False, indent=2)
+	print(f"Saved embeddings: {out_st} [{tuple(embeddings.shape)}]")
+	print(f"Saved sentence IDs: {out_json} [{len(ids)} ids]")
 
 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
+	print("Consolidating CLIP...")
+	clip_emb, clip_ids = consolidate_dir(CLIP_DIR)
+	save_as_safetensors(clip_emb, clip_ids, DATA_DIR / "clip_embeddings")
+
+	print("Consolidating PaintingCLIP...")
+	pclip_emb, pclip_ids = consolidate_dir(PAINTINGCLIP_DIR)
+	save_as_safetensors(pclip_emb, pclip_ids, DATA_DIR / "paintingclip_embeddings")
 
 if __name__ == "__main__":
-    exit(main())
+	main()

data/embeddings/clip_embeddings.safetensors

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

data/embeddings/paintingclip_embeddings.safetensors

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

requirements.txt

@@ -19,3 +19,5 @@ numpy>=1.24.0
 
 # Optional: GPU acceleration (if available)
 # torchvision>=0.15.0  # Uncomment if you need additional vision models
+
+safetensors>=0.4.0