build_all_motif_dbs.py

import os
import torch
import numpy as np
from tqdm import tqdm
import argparse
from typing import Optional, Union, List, Dict, Any
import time
from torch_geometric.data import Data
from nano_vectordb import NanoVectorDB
from subgraph_dataset import SubgraphViewDataset
from centrality_utils import compute_centrality_and_cse
from build_nano_db import load_node_dataset
from train_motif_finder import SubgraphEncoder

def build_motif_vectordb(dataset_name: str, 
                        motif_lib_path: str,
                        output_db_path: str,
                        k: int = 200,
                        ksteps: Optional[List[int]] = None,
                        device: str = 'cuda') -> Dict[str, Any]:

    device = torch.device(device)
    
    if ksteps is None:
        ksteps = list(range(1, 9))  

    try:
        args = argparse.Namespace()
        args.dataset = dataset_name.lower()
        dataset_wrapper = load_node_dataset(root='data', name=dataset_name, args=args)
        full_graph_data = dataset_wrapper.data
    except Exception as e:
        print(f"   ✗ load dataset failed: {e}")
        raise e

    dynamic_k = max(200, int(0.1 * full_graph_data.num_nodes))
    k = dynamic_k  
    try:
        results = compute_centrality_and_cse(
            edge_index=full_graph_data.edge_index,
            num_nodes=full_graph_data.num_nodes,
            ksteps=ksteps,
            k=k
        )
        cse_features = results['cse_encodings']
        top_k_indices = results['top_k_indices']

    except Exception as e:
        print(f"   ✗ CSE calculation failed: {e}")
        cse_features = torch.randn(full_graph_data.num_nodes, len(ksteps))
        top_k_indices = torch.randperm(full_graph_data.num_nodes)[:k]

    try:
        config_path = os.path.join(motif_lib_path, dataset_name, 'config.pth')
        encoder_path = os.path.join(motif_lib_path, dataset_name, 'encoder.pth')
        
        if not os.path.exists(config_path):
            raise FileNotFoundError(f"config file not found: {config_path}")
        if not os.path.exists(encoder_path):
            raise FileNotFoundError(f"encoder weight file not found: {encoder_path}")
        
        config = torch.load(config_path, map_location=device)
        struct_input_dim = config['struct_input_dim']
        hidden_dim = config['hidden_dim']
        output_dim = config['output_dim']

        encoder = SubgraphEncoder(
            input_dim=struct_input_dim,
            hidden_dim=hidden_dim,
            output_dim=output_dim
        )
        
        encoder.load_state_dict(torch.load(encoder_path, map_location=device))
        encoder.to(device)
        encoder.eval()
        
    except Exception as e:
        print(f"   ✗ load encoder failed: {e}")
        raise e
    
    try:
        subgraph_dataset = SubgraphViewDataset(
            full_graph_data=full_graph_data,
            cse_features=cse_features,
            top_k_node_indices=top_k_indices
        )
    except Exception as e:
        print(f"   ✗ create subgraph dataset failed: {e}")
        raise e

    documents = []
    embedding_dim = -1
    
    with torch.no_grad():
        for i in tqdm(range(len(subgraph_dataset)), desc="encode subgraph"):
            struct_view, _ = subgraph_dataset[i]
            
            struct_view.batch = torch.zeros(struct_view.num_nodes, dtype=torch.long, device=device)
            struct_view = struct_view.to(device)
            
            subgraph_embedding = encoder(struct_view)
            
            subgraph_embedding_np = subgraph_embedding.cpu().numpy().flatten()
            
            if embedding_dim == -1:
                embedding_dim = len(subgraph_embedding_np)
            
            doc = {
                "__vector__": subgraph_embedding_np,
                "metadata": {
                    "domain": dataset_name,
                    "center_node_original_idx": top_k_indices[i].item(),
                    "subgraph_size": struct_view.num_nodes,
                    "edge_count": struct_view.edge_index.shape[1]
                }
            }
            documents.append(doc)
    
    try:
        db_dir = os.path.dirname(output_db_path)
        if db_dir and not os.path.exists(db_dir):
            os.makedirs(db_dir)
        
        if os.path.exists(output_db_path):
            os.remove(output_db_path)
        
        db = NanoVectorDB(embedding_dim, storage_file=output_db_path)
        
        db.upsert(documents)

        db.save()
        print(f"   ✓ database saved to: {os.path.abspath(output_db_path)}")
        
    except Exception as e:
        print(f"   ✗ build database failed: {e}")
        raise e
    
    print(f"\n======================================================")
    print(f"✅ {dataset_name} Motif vector database built!")
    print(f"======================================================")
    
    return {
        'database_path': output_db_path,
        'num_documents': len(documents),
        'embedding_dim': embedding_dim,
        'top_k_nodes': len(top_k_indices),
        'success': True
    }


def main():
    parser = argparse.ArgumentParser(description='batch build all motif vector databases')
    parser.add_argument('--datasets', nargs='+', 
                       default=['Cora','Citeseer','Pubmed','home','wikics'],
                       help='list of datasets to process')
    parser.add_argument('--motif_lib_path', type=str, default='./motif_lib/',
                       help='root directory of GNN encoders')
    parser.add_argument('--output_db_path', type=str, default='./motif_lib/',
                       help='root directory of output databases')
    parser.add_argument('--top_k', type=int, default=200,
                       help='number of center nodes selected for each dataset')
    parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu',
                       help='compute device')
    parser.add_argument('--seed', type=int, default=42,
                       help='random seed')
    
    args = parser.parse_args()
    
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)

    os.makedirs(args.output_db_path, exist_ok=True)

    total_start_time = time.time()
    
    for i, dataset_name in enumerate(args.datasets, 1):
        print(f"\n{'='*60}")
        print(f"start to build motif vector database for dataset '{dataset_name}' ({i}/{len(args.datasets)})")
        print(f"{'='*60}")
        
        start_time = time.time()
        
        try:
            config_path = os.path.join(args.motif_lib_path, dataset_name, 'config.pth')
            encoder_path = os.path.join(args.motif_lib_path, dataset_name, 'encoder.pth')
            
            if not os.path.exists(config_path):
                print(f"   ⚠ config file not found: {config_path}")
                print(f"   ⚠ skip dataset '{dataset_name}'")
                continue
                
            if not os.path.exists(encoder_path):
                print(f"   ⚠ encoder weight file not found: {encoder_path}")
                print(f"   ⚠ skip dataset '{dataset_name}'")
                continue
            
            output_db_path = os.path.join(args.output_db_path, dataset_name, 'motif_vectordb.json')
            
            result = build_motif_vectordb(
                dataset_name=dataset_name,
                motif_lib_path=args.motif_lib_path,
                output_db_path=output_db_path,
                k=args.top_k,
                device=args.device
            )
            
            end_time = time.time()
            
        except Exception as e:
            end_time = time.time()
            elapsed_time = end_time - start_time
            
            print(f"\n✗ dataset '{dataset_name}' built failed!")
    
    total_end_time = time.time()
    
if __name__ == "__main__":
    main()