refactor: remove overlapping examples and prepare for G-Retriever integration

Author: AJamal27891

examples/llm/relbench_example.py

@@ -26,7 +26,7 @@
 
 # RelBench integration utilities (optional dependencies)
 try:
-    from .relbench import RelBenchDataset, RelBenchProcessor, create_relbench_hetero_data, get_warehouse_task_info  # noqa: F401
+    from .relbench import RelBenchDataset, RelBenchProcessor, create_relbench_hetero_data, get_warehouse_task_info, prepare_for_gretriever  # noqa: F401
     _relbench_available = True
 except ImportError:
     _relbench_available = False
@@ -260,4 +260,5 @@
         "RelBenchProcessor",
         "create_relbench_hetero_data",
         "get_warehouse_task_info",
+        "prepare_for_gretriever",
     ])

examples/relbench/00_llm_gnn_bridge.ipynb

@@ -1,11 +1,9 @@
 """RelBench integration utilities for PyTorch Geometric.
 
 Provides utilities for converting RelBench datasets to PyG HeteroData objects
-with semantic embeddings and graph structure for warehouse applications.
+with semantic embeddings and warehouse-specific enhancements.
 
-TODO: Add subgraph sampling utilities for inference
-TODO: Implement configurable edge weighting schemes
-TODO: Add support for lineage tracking
+Complements examples/rdl.py with G-Retriever preparation and warehouse tasks.
 """
 
 import warnings
@@ -1044,6 +1042,35 @@ def process(self) -> None:
         torch.save((collated_data, slices), self.processed_paths[0])
 
 
+def prepare_for_gretriever(
+        hetero_data: HeteroData) -> Tuple[HeteroData, Dict[str, Any]]:
+    """Prepare RelBench HeteroData for G-Retriever training.
+
+    Enhances HeteroData with G-Retriever-specific attributes and metadata.
+
+    Args:
+        hetero_data: HeteroData object from RelBench integration
+
+    Returns:
+        Tuple of (enhanced_hetero_data, metadata_dict)
+    """
+    metadata = {
+        'embedding_dim': getattr(hetero_data, 'embedding_dim', 384),
+        'node_types': list(hetero_data.node_types),
+        'edge_types': list(hetero_data.edge_types),
+        'warehouse_tasks': ['lineage', 'silo', 'anomaly'],
+        'recommended_qa_pairs': get_warehouse_task_info(),
+        'conversion_ready': True,
+    }
+
+    # Add G-Retriever specific attributes
+    hetero_data.gretriever_ready = True
+    hetero_data.embedding_type = 'sbert'  # Indicates SBERT embeddings
+    hetero_data.warehouse_enhanced = True
+
+    return hetero_data, metadata
+
+
 # Backward compatibility aliases
 RelBenchHeteroDataProcessor = RelBenchProcessor
 create_hetero_data_from_relbench = create_relbench_hetero_data