Try to fix changes

backend/runner/app.py

@@ -98,7 +98,9 @@ from .config import (
     ARTIFACTS_DIR,
     OUTPUTS_DIR,
     JSON_INFO_DIR,
-    MARKER_DIR
+    MARKER_DIR,
+    JSON_DATASETS,
+    EMBEDDINGS_DATASETS
 )
 
 # Import data from config (loaded from HF datasets)

backend/runner/config.py

@@ -5,12 +5,33 @@ All runner modules should import from this module instead of defining their own
 
 import os
 from pathlib import Path
-from datasets import load_dataset
+from typing import Any, Dict, Optional
 
-# HF Dataset IDs
-EMBEDDINGS_DATASET = "samwaugh/artefact-embeddings"
-JSON_DATASET = "samwaugh/artefact-json"
-MARKDOWN_DATASET = "samwaugh/artefact-markdown"
+# Try to import datasets, but handle gracefully if not available
+try:
+    from datasets import load_dataset
+    DATASETS_AVAILABLE = True
+except ImportError:
+    print("⚠️  datasets library not available - HF dataset loading disabled")
+    DATASETS_AVAILABLE = False
+
+# Environment variables for dataset names
+ARTEFACT_JSON_DATASET = os.getenv('ARTEFACT_JSON_DATASET', 'samwaugh/artefact-json')
+ARTEFACT_EMBEDDINGS_DATASET = os.getenv('ARTEFACT_EMBEDDINGS_DATASET', 'samwaugh/artefact-embeddings')
+ARTEFACT_MARKDOWN_DATASET = os.getenv('ARTEFACT_MARKDOWN_DATASET', 'samwaugh/artefact-markdown')
+
+# Legacy path variables for backward compatibility
+JSON_INFO_DIR = "/data/hub/datasets--samwaugh--artefact-json/snapshots/latest"
+EMBEDDINGS_DIR = "/data/hub/datasets--samwaugh--artefact-embeddings/snapshots/latest"
+MARKDOWN_DIR = "/data/hub/datasets--samwaugh--artefact-markdown/snapshots/latest"
+
+# Embedding file paths for backward compatibility
+CLIP_EMBEDDINGS_ST = Path(EMBEDDINGS_DIR) / "clip_embeddings.safetensors"
+PAINTINGCLIP_EMBEDDINGS_ST = Path(EMBEDDINGS_DIR) / "paintingclip_embeddings.safetensors"
+CLIP_SENTENCE_IDS = Path(EMBEDDINGS_DIR) / "clip_embeddings_sentence_ids.json"
+PAINTINGCLIP_SENTENCE_IDS = Path(EMBEDDINGS_DIR) / "paintingclip_embeddings_sentence_ids.json"
+CLIP_EMBEDDINGS_DIR = EMBEDDINGS_DIR
+PAINTINGCLIP_EMBEDDINGS_DIR = EMBEDDINGS_DIR
 
 # READ root (repo data - read-only)
 PROJECT_ROOT = Path(__file__).resolve().parents[2]
@@ -60,50 +81,96 @@ for dir_path in [OUTPUTS_DIR, ARTIFACTS_DIR]:
         print(f"⚠️  Could not create directory {dir_path}: {e}")
 
 # Global data variables (will be populated from HF datasets)
-sentences = {}
-works = {}
-creators = {}
-topics = {}
-topic_names = {}
-
-def load_json_from_hf(dataset_name: str, file_name: str):
-    """Load JSON data from Hugging Face dataset"""
-    try:
-        dataset = load_dataset(dataset_name, split="train")
-        # Access the specific file content
-        return dataset[file_name]
-    except Exception as e:
-        print(f"Failed to load {file_name} from HF: {e}")
-        return None
+sentences: Dict[str, Any] = {}
+works: Dict[str, Any] = {}
+creators: Dict[str, Any] = {}
+topics: Dict[str, Any] = {}
+topic_names: Dict[str, Any] = {}
 
-def load_all_data():
-    """Load all data from Hugging Face datasets"""
-    global sentences, works, creators, topics, topic_names
-    
-    print("🔄 Loading data from Hugging Face datasets...")
-    
-    sentences = load_json_from_hf(JSON_DATASET, "sentences.json")
-    works = load_json_from_hf(JSON_DATASET, "works.json")
-    creators = load_json_from_hf(JSON_DATASET, "creators.json")
-    topics = load_json_from_hf(JSON_DATASET, "topics.json")
-    topic_names = load_json_from_hf(JSON_DATASET, "topic_names.json")
-    
-    # Validate data loading
-    if sentences and works and creators and topics and topic_names:
-        print(f"✅ Successfully loaded data from HF:")
+# Load datasets from Hugging Face
+def load_json_datasets() -> Optional[Dict[str, Any]]:
+    """Load all JSON datasets from Hugging Face"""
+    if not DATASETS_AVAILABLE:
+        print("⚠️  datasets library not available - skipping HF dataset loading")
+        return None
+        
+    try:
+        print("🔄 Loading data from Hugging Face datasets...")
+        
+        creators_dataset = load_dataset(ARTEFACT_JSON_DATASET, 'creators.json', split='train')
+        sentences_dataset = load_dataset(ARTEFACT_JSON_DATASET, 'sentences.json', split='train')
+        works_dataset = load_dataset(ARTEFACT_JSON_DATASET, 'works.json', split='train')
+        topics_dataset = load_dataset(ARTEFACT_JSON_DATASET, 'topics.json', split='train')
+        topic_names_dataset = load_dataset(ARTEFACT_JSON_DATASET, 'topic_names.json', split='train')
+        
+        # Convert to dictionaries for backward compatibility
+        global sentences, works, creators, topics, topic_names
+        
+        sentences = {str(i): item for i, item in enumerate(sentences_dataset)}
+        works = {str(i): item for i, item in enumerate(works_dataset)}
+        creators = {str(i): item for i, item in enumerate(creators_dataset)}
+        topics = {str(i): item for i, item in enumerate(topics_dataset)}
+        topic_names = {str(i): item for i, item in enumerate(topic_names_dataset)}
+        
+        print(f"✅ Successfully loaded JSON datasets from HF:")
         print(f"   Sentences: {len(sentences)} entries")
         print(f"   Works: {len(works)} entries")
-        print(f"   Topics: {len(topics)} entries")
         print(f"   Creators: {len(creators)} entries")
-        print(f"   Topic names: {len(topic_names)} entries")
-    else:
-        print("⚠️  Some data failed to load from HF datasets")
-        # Fallback to empty dicts to prevent crashes
-        sentences = sentences or {}
-        works = works or {}
-        creators = creators or {}
-        topics = topics or {}
-        topic_names = topic_names or {}
+        print(f"   Topics: {len(topics)} entries")
+        print(f"   Topic Names: {len(topic_names)} entries")
+        
+        return {
+            'creators': creators_dataset,
+            'sentences': sentences_dataset,
+            'works': works_dataset,
+            'topics': topics_dataset,
+            'topic_names': topic_names_dataset
+        }
+    except Exception as e:
+        print(f"❌ Failed to load JSON datasets from HF: {e}")
+        return None
+
+def load_embeddings_datasets() -> Optional[Dict[str, Any]]:
+    """Load embeddings datasets from Hugging Face"""
+    if not DATASETS_AVAILABLE:
+        print("⚠️  datasets library not available - skipping HF embeddings loading")
+        return None
+        
+    try:
+        clip_embeddings = load_dataset(ARTEFACT_EMBEDDINGS_DATASET, 'clip_embeddings.safetensors', split='train')
+        paintingclip_embeddings = load_dataset(ARTEFACT_EMBEDDINGS_DATASET, 'paintingclip_embeddings.safetensors', split='train')
+        
+        return {
+            'clip': clip_embeddings,
+            'paintingclip': paintingclip_embeddings
+        }
+    except Exception as e:
+        print(f"❌ Failed to load embeddings datasets from HF: {e}")
+        return None
+
+# Initialize datasets
+JSON_DATASETS = load_json_datasets()
+EMBEDDINGS_DATASETS = load_embeddings_datasets()
 
 # Initialize data loading
-load_all_data()
+if JSON_DATASETS is None:
+    print("⚠️  Some data failed to load from HF datasets")
+else:
+    print("✅ All data loaded successfully from HF datasets")
+
+# Add this function for backward compatibility
+def st_load_file(file_path: Path) -> Any:
+    """Load a file using safetensors or other methods"""
+    try:
+        if file_path.suffix == '.safetensors':
+            import safetensors
+            return safetensors.safe_open(str(file_path), framework="pt")
+        else:
+            import torch
+            return torch.load(str(file_path))
+    except ImportError:
+        print(f"⚠️  Required library not available for loading {file_path}")
+        return None
+    except Exception as e:
+        print(f"❌ Error loading {file_path}: {e}")
+        return None

backend/runner/filtering.py

@@ -5,7 +5,10 @@ Filtering logic for sentence selection based on topics and creators.
 from typing import Any, Dict, List, Set
 
 # Import data from config (loaded from HF datasets)
-from .config import sentences, works, creators, topics
+from .config import (
+    JSON_INFO_DIR,
+    JSON_DATASETS
+)
 
 # Data is now loaded from Hugging Face datasets in config.py
 # No need to load from local files anymore
@@ -24,7 +27,7 @@ def get_filtered_sentence_ids(
         Set of sentence IDs that match all filters
     """
     # Start with all sentence IDs
-    valid_sentence_ids = set(sentences.keys())
+    valid_sentence_ids = set(JSON_DATASETS['sentences']['id'])
 
     # If no filters, return all sentences
     if not filter_topics and not filter_creators:
@@ -38,21 +41,21 @@ def get_filtered_sentence_ids(
         # Using topics.json (topic -> works mapping)
         # For each selected topic, get all works that have it
         for topic_id in filter_topics:
-            if topic_id in topics:
+            if topic_id in JSON_DATASETS['topics']:
                 # Add all works that have this topic
-                valid_work_ids.update(topics[topic_id])
+                valid_work_ids.update(JSON_DATASETS['topics'][topic_id])
     else:
         # If no topic filter, all works are valid so far
-        valid_work_ids = set(works.keys())
+        valid_work_ids = set(JSON_DATASETS['works']['id'])
 
     # Apply creator filter
     if filter_creators:
         # Direct lookup in creators.json (more efficient)
         creator_work_ids = set()
         for creator_name in filter_creators:
-            if creator_name in creators:
+            if creator_name in JSON_DATASETS['creators']:
                 # Get all works by this creator directly from creators.json
-                creator_work_ids.update(creators[creator_name])
+                creator_work_ids.update(JSON_DATASETS['creators'][creator_name])
 
         # Intersect with existing valid_work_ids if topics were filtered
         if filter_topics:

backend/runner/inference.py

@@ -31,14 +31,19 @@ from .filtering import get_filtered_sentence_ids
 # on-demand Grad-ECLIP & region-aware ranking
 from .heatmap import generate_heatmap
 from .config import (
+    JSON_INFO_DIR,
+    EMBEDDINGS_DIR,
+    JSON_DATASETS,
+    EMBEDDINGS_DATASETS,
     PAINTINGCLIP_MODEL_DIR,
-    EMBEDDINGS_DATASET,
-    JSON_DATASET,
-    sentences,
-    works,
-    creators,
-    topics,
-    topic_names
+    ARTEFACT_EMBEDDINGS_DATASET,
+    sentences,  # Add this
+    CLIP_EMBEDDINGS_ST,  # Add these for backward compatibility
+    PAINTINGCLIP_EMBEDDINGS_ST,
+    CLIP_SENTENCE_IDS,
+    PAINTINGCLIP_SENTENCE_IDS,
+    CLIP_EMBEDDINGS_DIR,
+    PAINTINGCLIP_EMBEDDINGS_DIR
 )
 
 # ─── Configuration ───────────────────────────────────────────────────────────
@@ -65,8 +70,8 @@ TOP_K = 25  # Number of results to return
 def load_embeddings_from_hf():
     """Load embeddings from HF dataset"""
     try:
-        print(f"🔍 Loading embeddings from {EMBEDDINGS_DATASET}...")
-        dataset = load_dataset(EMBEDDINGS_DATASET, split="train")
+        print(f"🔍 Loading embeddings from {ARTEFACT_EMBEDDINGS_DATASET}...")
+        dataset = load_dataset(ARTEFACT_EMBEDDINGS_DATASET, split="train")
         
         # Load CLIP embeddings
         clip_embeddings = dataset["clip_embeddings"]
@@ -92,6 +97,9 @@ def _load_sentences_metadata() -> Dict[str, Dict[str, Any]]:
     """
     Get sentence metadata from global config (loaded from HF datasets).
     """
+    if not sentences:
+        print("⚠️  No sentence metadata available - check if HF datasets loaded successfully")
+        return {}
     return sentences
 
 @lru_cache(maxsize=1)
@@ -156,7 +164,7 @@ def _initialize_pipeline():
     try:
         embeddings_data = load_embeddings_from_hf()
         if embeddings_data is None:
-            raise ValueError(f"Failed to load embeddings from HF dataset: {EMBEDDINGS_DATASET}")
+            raise ValueError(f"Failed to load embeddings from HF dataset: {ARTEFACT_EMBEDDINGS_DATASET}")
         
         if MODEL_TYPE == "clip":
             embeddings, sentence_ids = embeddings_data["clip"]
@@ -489,3 +497,20 @@ def load_embeddings_for_model(model_type: str):
 	print(f"   - {st_file.name}")
 	print(f"   - {ids_file.name}")
 	return None, None
+
+# Add this function for backward compatibility
+def st_load_file(file_path: Path) -> Any:
+    """Load a file using safetensors or other methods"""
+    try:
+        if file_path.suffix == '.safetensors':
+            import safetensors
+            return safetensors.safe_open(str(file_path), framework="pt")
+        else:
+            import torch
+            return torch.load(str(file_path))
+    except ImportError:
+        print(f"⚠️  Required library not available for loading {file_path}")
+        return None
+    except Exception as e:
+        print(f"❌ Error loading {file_path}: {e}")
+        return None

requirements.txt

@@ -6,7 +6,7 @@ flask-cors
 # Hugging Face ecosystem
 huggingface_hub>=0.20
 hf_transfer>=0.1.4
-datasets>=2.14.0
+datasets>=2.0.0
 
 # Core ML libraries
 torch>=2.0.0