Allow weights false
Browse files- backend/runner/inference.py +18 -2
backend/runner/inference.py
CHANGED
|
@@ -425,8 +425,24 @@ def load_consolidated_embeddings(embedding_file: Path, metadata_file: Path):
|
|
| 425 |
"""Load embeddings from consolidated file with metadata"""
|
| 426 |
print(f"Loading consolidated embeddings from {embedding_file}")
|
| 427 |
|
| 428 |
-
# Load consolidated data
|
| 429 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 430 |
embeddings = consolidated_data['embeddings']
|
| 431 |
|
| 432 |
# Load metadata for file mapping
|
|
|
|
| 425 |
"""Load embeddings from consolidated file with metadata"""
|
| 426 |
print(f"Loading consolidated embeddings from {embedding_file}")
|
| 427 |
|
| 428 |
+
# Load consolidated data with weights_only=False for compatibility
|
| 429 |
+
# This is safe since we're loading our own pre-computed embeddings
|
| 430 |
+
try:
|
| 431 |
+
consolidated_data = torch.load(embedding_file, map_location='cpu', weights_only=False)
|
| 432 |
+
print(f"✅ Successfully loaded consolidated embeddings")
|
| 433 |
+
except Exception as e:
|
| 434 |
+
print(f"❌ Failed to load with weights_only=False: {e}")
|
| 435 |
+
# Fallback: try with weights_only=True (may fail if file has non-tensor data)
|
| 436 |
+
try:
|
| 437 |
+
print(f"🔍 Trying fallback with weights_only=True...")
|
| 438 |
+
consolidated_data = torch.load(embedding_file, map_location='cpu', weights_only=True)
|
| 439 |
+
print(f"✅ Successfully loaded with weights_only=True")
|
| 440 |
+
except Exception as e2:
|
| 441 |
+
print(f"❌ Both loading methods failed:")
|
| 442 |
+
print(f" weights_only=False: {e}")
|
| 443 |
+
print(f" weights_only=True: {e2}")
|
| 444 |
+
raise RuntimeError(f"Cannot load embedding file with either method: {e2}")
|
| 445 |
+
|
| 446 |
embeddings = consolidated_data['embeddings']
|
| 447 |
|
| 448 |
# Load metadata for file mapping
|