execute_graph_cora.py

from unittest import loader
import numpy as np
import scipy.sparse as sp
from sklearn.metrics import f1_score
import random
import torch
import torch.nn as nn
import torch.nn.functional as F
from tqdm import tqdm
import argparse
from torch_geometric.datasets import Planetoid, Amazon, Coauthor, Reddit
from torch_geometric.loader import DataLoader
import torch.distributed as dist 
from torch_geometric.data import Batch
from models import LogReg
from preprompt import PrePrompt, pca_compression
import preprompt
from utils import process
import pdb
import aug
import os
from downprompt_graph import downprompt, prefeatureprompt
import csv
from feature_enhancer import FeatureEnhancer
from build_nano_db import load_node_dataset
from search_nano_db import RAGRetriever
from search_motif_db import MotifRetriever
from RAG_enhancer_sample import RAGEnhancer
import pickle
import hashlib

parser = argparse.ArgumentParser("RAG-GFM")
parser.add_argument('--dataset', type=str, default="Cora", help='data')
parser.add_argument('--aug_type', type=str, default="edge", help='aug type: mask or edge')
parser.add_argument('--drop_percent', type=float, default=0.1, help='drop percent')
parser.add_argument('--seed', type=int, default=39, help='seed')
parser.add_argument('--gpu', type=int, default=0, help='gpu (0-3 for multi-GPU support)')
parser.add_argument('--save_name', type=str, default='Cora30.pkl', help='save ckpt name')
parser.add_argument('--val_name', type=str, default='noval_graphcl_BZR.pkl', help='save val')
parser.add_argument('--combinetype', type=str, default='mul', help='the type of text combining')
parser.add_argument('--skip_pretrain', action='store_true', help='skip pretrain')
parser.add_argument('--pretrained_model', type=str, default=None,
                    help='pretrained model path, must be specified when skip_pretrain is True')
parser.add_argument('--use_cache', action='store_true', default=True, help='use cache to accelerate pretrain data preparation')
parser.add_argument('--clear_cache', action='store_true', help='clear pretrain data cache')
args = parser.parse_args()

print('-' * 100)
print("experiment parameters:")
print(args)
print('-' * 100)

if args.skip_pretrain and args.pretrained_model is None:
    raise ValueError("when using --skip_pretrain parameter, must specify --pretrained_model parameter")
if args.skip_pretrain and not os.path.exists(args.pretrained_model):
    raise FileNotFoundError(f"pretrained model file does not exist: {args.pretrained_model}")

CACHE_DIR = "cache"

if args.clear_cache:
    if os.path.exists(CACHE_DIR):
        import glob
        cache_files = glob.glob(os.path.join(CACHE_DIR, "*.pkl"))
        if cache_files:
            for cache_file in cache_files:
                os.remove(cache_file)
                print(f"✓ cleared cache file: {cache_file}")
            print(f"✓ cleared {len(cache_files)} cache files")
        else:
            print("⚠ cache directory does not contain cache files")
    else:
        print("⚠ cache directory does not exist, no need to clear")
    exit(0)

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu) 
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"using device: {device} (GPU {args.gpu})")

seed = args.seed
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  

UNIFIED_RAW_FEATURE_DIM = 50  
TEXT_FEATURE_DIM = 50  

PRETRAIN_DATASETS = ['Citeseer', 'Pubmed', 'wikics', 'home']
DOWNSTREAM_DATASET = 'Cora'
ALL_DB_DATASETS = set(PRETRAIN_DATASETS) | {DOWNSTREAM_DATASET}

def get_cache_key(args, pretrain_datasets):
    key_params = {
        'pretrain_datasets': sorted(pretrain_datasets),  
        'downstream_dataset': args.dataset,  
        'unified_raw_feature_dim': UNIFIED_RAW_FEATURE_DIM,
        'text_feature_dim': TEXT_FEATURE_DIM,
        'sparse': True,  
        'combinetype': args.combinetype,
        'aug_type': args.aug_type,
        'drop_percent': args.drop_percent
    }
    key_string = str(sorted(key_params.items()))
    return hashlib.md5(key_string.encode()).hexdigest()

def get_cache_filename(args, pretrain_datasets):
    pretrain_str = "_".join(sorted(pretrain_datasets))
    downstream_str = args.dataset
    cache_key = get_cache_key(args, pretrain_datasets)
    filename = f"pretrain_{pretrain_str}_downstream_{downstream_str}_{cache_key[:8]}.pkl"
    return os.path.join(CACHE_DIR, filename)

def save_pretrain_cache(cache_data, cache_key, cache_filename):
    os.makedirs(CACHE_DIR, exist_ok=True)
    cache_info = {
        'cache_key': cache_key,
        'data': cache_data
    }
    with open(cache_filename, 'wb') as f:
        pickle.dump(cache_info, f)

    cache_size = os.path.getsize(cache_filename) / (1024 * 1024)  # MB


def load_pretrain_cache(cache_key, cache_filename):
    if not os.path.exists(cache_filename):
        print(f"  - cache file does not exist: {cache_filename}")
        return None
    
    try:
        with open(cache_filename, 'rb') as f:
            cache_info = pickle.load(f)
        
        if cache_info['cache_key'] != cache_key:
            print(f"  ⚠ cache key does not match, will re-generate data")
            return None
        
        cache_size = os.path.getsize(cache_filename) / (1024 * 1024)  # MB
        return cache_info['data']
    except Exception as e:
        print(f"  ⚠ load pretrain data from cache failed: {e}, will re-generate data")
        return None

def get_downstream_cache_key(args, dataset_name):
    key_params = {
        'dataset': dataset_name,
        'unified_raw_feature_dim': UNIFIED_RAW_FEATURE_DIM,
        'text_feature_dim': TEXT_FEATURE_DIM,
        'seed': args.seed
    }
    key_string = str(sorted(key_params.items()))
    return hashlib.md5(key_string.encode()).hexdigest()

def get_downstream_cache_filename(args, dataset_name):
    cache_key = get_downstream_cache_key(args, dataset_name)
    filename = f"downstream_{dataset_name.lower()}_{cache_key[:8]}.pkl"
    return os.path.join(CACHE_DIR, filename)

def save_downstream_cache(cache_data, cache_key, cache_filename):
    os.makedirs(CACHE_DIR, exist_ok=True)
    cache_info = {
        'cache_key': cache_key,
        'data': cache_data
    }
    with open(cache_filename, 'wb') as f:
        pickle.dump(cache_info, f)
    cache_size = os.path.getsize(cache_filename) / (1024 * 1024)

def load_downstream_cache(cache_key, cache_filename):
    if not os.path.exists(cache_filename):
        print(f"  - downstream task cache file does not exist: {cache_filename}")
        return None
    
    try:
        with open(cache_filename, 'rb') as f:
            cache_info = pickle.load(f)
        
        if cache_info['cache_key'] != cache_key:
            print(f"  ⚠ downstream task cache key does not match, will re-generate data")
            return None
        
        return cache_info['data']
    
    except Exception as e:
        print(f"  ⚠ load downstream task cache failed: {e}")
        return None

nb_epochs = 10000
patience = 50
lr = 0.00001
l2_coef = 0.0
hid_units = 256

unify_dim = UNIFIED_RAW_FEATURE_DIM + TEXT_FEATURE_DIM
sparse = True
nonlinearity = 'prelu'

b_xent = nn.BCEWithLogitsLoss()
xent = nn.CrossEntropyLoss()

cache_key = get_cache_key(args, PRETRAIN_DATASETS)
cache_filename = get_cache_filename(args, PRETRAIN_DATASETS)

cached_data = None
if args.use_cache:
    cached_data = load_pretrain_cache(cache_key, cache_filename)
else:
    print("  - cache is disabled, will re-generate pretrain data...")

if cached_data is not None:
    pretrain_features = {name: cached_data[f'features_{name}'] for name in PRETRAIN_DATASETS}
    pretrain_adjs = {name: cached_data[f'adj_{name}'] for name in PRETRAIN_DATASETS}
    pretrain_sp_adjs = {name: cached_data[f'sp_adj_{name}'] for name in PRETRAIN_DATASETS}
    negetive_sample = cached_data['negetive_sample']

else:    
    enhancer = FeatureEnhancer(
        text_encoder='bert',
        text_output_dim=TEXT_FEATURE_DIM
    )
    pretrain_features = {}
    pretrain_adjs = {}
    pretrain_sp_adjs = {}

    for name in PRETRAIN_DATASETS:
        dataset_wrapper = load_node_dataset(root='data', name=name, args=args)
        data = dataset_wrapper.data

        batched_data = Batch.from_data_list([data])

        raw_features_np, adj = process.process_tu(batched_data, batched_data.x.shape[1])

        raw_texts = data.raw_texts

        unified_raw_features_np = pca_compression(raw_features_np, k=UNIFIED_RAW_FEATURE_DIM)

        enhanced_features_tensor = enhancer.enhance(unified_raw_features_np, raw_texts)

        adj_normalized = process.normalize_adj(adj + sp.eye(adj.shape[0]))
        if sparse:
            sp_adj = process.sparse_mx_to_torch_sparse_tensor(adj_normalized)
            pretrain_sp_adjs[name] = sp_adj

        pretrain_features[name] = enhanced_features_tensor
        pretrain_adjs[name] = adj

    combined_adj_for_sampling = process.combine_dataset(*[pretrain_adjs[name] for name in PRETRAIN_DATASETS])
    negetive_sample = preprompt.prompt_pretrain_sample(combined_adj_for_sampling, 50)
    
    if args.use_cache:
        cache_data = {'negetive_sample': negetive_sample}
        for name in PRETRAIN_DATASETS:
            cache_data[f'features_{name}'] = pretrain_features[name]
            cache_data[f'adj_{name}'] = pretrain_adjs[name]
            cache_data[f'sp_adj_{name}'] = pretrain_sp_adjs[name]
        save_pretrain_cache(cache_data, cache_key, cache_filename)

model = PrePrompt(unify_dim, hid_units, nonlinearity, negetive_sample, 3, 0.1, args.combinetype)

if not args.skip_pretrain:
    optimiser = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=l2_coef)

if torch.cuda.is_available():
    model = model.cuda()
    for name in PRETRAIN_DATASETS:
        pretrain_features[name] = pretrain_features[name].cuda()
        if sparse:
            pretrain_sp_adjs[name] = pretrain_sp_adjs[name].cuda()

if not args.skip_pretrain:
    best_loss = 1e9
    best_epoch = 0
    epochs_since_best = 0

    for epoch in range(nb_epochs):
        model.train()
        optimiser.zero_grad()

        features_list = [pretrain_features[name] for name in PRETRAIN_DATASETS]
        adjs_list = [pretrain_sp_adjs[name] if sparse else pretrain_adjs[name] for name in PRETRAIN_DATASETS]
        loss = model(*features_list, *adjs_list, sparse, None, None, None)

        loss.backward()
        optimiser.step()

        print(f'Epoch: [{epoch + 1}/{nb_epochs}] | Loss: {loss.item():.4f}')

        if loss.item() < best_loss:
            best_loss = loss.item()
            best_epoch = epoch
            epochs_since_best = 0
            torch.save(model.state_dict(), args.save_name)
            print(f'  -> new best model saved in epoch {epoch + 1}, loss: {best_loss:.4f}')
        else:
            epochs_since_best += 1

        if epochs_since_best == patience:
            print(f'patience reached {patience}, early stopping.')
            break

else:
    model.load_state_dict(torch.load(args.pretrained_model, map_location=device))

rag_retriever = RAGRetriever(
    db_file_path='./nano_db/unified_database.json',
    text_encoder='bert'
)

try:
    motif_retriever = MotifRetriever(
        motif_lib_path='./motif_lib/',
        motif_db_path='./motif_lib/'
    )
except Exception as e:
    print(f"  ⚠ Motif Retriever initialization failed: {e}")
    motif_retriever = None

rag_enhancer = RAGEnhancer(feature_dim=unify_dim)

downstream_task_name = args.dataset.capitalize()
if downstream_task_name not in ALL_DB_DATASETS:
    if downstream_task_name.lower() == 'cs':
        downstream_task_name = 'CS'
    else:
        raise ValueError(f"downstream task data set '{downstream_task_name}' not in known database data sets!")

pretrain_sets_for_this_run = ALL_DB_DATASETS - {downstream_task_name}

downstream_cache_file = ""

if not os.path.exists(downstream_cache_file):
    raise FileNotFoundError(f"cache file does not exist: {downstream_cache_file}")

with open(downstream_cache_file, 'rb') as f:
    cache_info = pickle.load(f)
    cached_downstream_data = cache_info['data']

raw_features_np = cached_downstream_data['raw_features_np']
adj_scipy = cached_downstream_data['adj_scipy']
raw_texts = cached_downstream_data['raw_texts']
unified_raw_features_np = cached_downstream_data['unified_raw_features_np']
base_enhanced_features = cached_downstream_data['base_enhanced_features'].to(device)
adj_normalized = cached_downstream_data['adj_normalized']
sp_adj = process.sparse_mx_to_torch_sparse_tensor(adj_normalized).to(device)
labels = cached_downstream_data['labels'].to(device)
nb_classes = cached_downstream_data['nb_classes']
test_range = cached_downstream_data['test_range']
idx_test = cached_downstream_data['idx_test']
test_lbls = labels[idx_test]
data = cached_downstream_data['data']


if args.skip_pretrain:
    model.load_state_dict(torch.load(args.pretrained_model, map_location=device))
else:
    model.load_state_dict(torch.load(args.save_name, map_location=device))

model.eval() 

testsetsize = test_range
neighboradj = adj_scipy.todense().A

neighborslist = [[] for x in range(testsetsize)]
neighbors_2hoplist = [[] for x in range(testsetsize)]
testindex = [[] for x in range(testsetsize)]
testlist = [[] for x in range(testsetsize)]
for x,y in enumerate(idx_test):
    
    neighborslist[x], neighbors_2hoplist[x]= process.find_2hop_neighbors(neighboradj, y)
    testlist[x] = [y] + neighborslist[x] + neighbors_2hoplist[x]
    testindex[x] = [x] * len(testlist[x])
neighborslist = sum(neighborslist,[])
neighbors_2hoplist = sum(neighbors_2hoplist,[])
testlist = sum(testlist,[])
testindex = sum(testindex,[])
testlist = torch.Tensor(testlist).type(torch.long).cuda()
testindex = torch.Tensor(testindex).type(torch.long).cuda()
with torch.no_grad():
    original_embeds, _ = model.embed(base_enhanced_features, sp_adj, sparse, None, None)
test_embs = original_embeds[0, testlist]

downstreamlr = 0.001
shotnum = 1

all_accs = []

for i in tqdm(range(5), desc="  - Few-shot experiment progress"):
    texttoken_weights = [getattr(model, f'texttoken{idx+1}').weight.detach() for idx in range(len(PRETRAIN_DATASETS))]
    pretext_weights = [getattr(model, f'pretext{idx+1}').weight.detach() for idx in range(len(PRETRAIN_DATASETS))]
    log = downprompt(*texttoken_weights, *pretext_weights, hid_units, nb_classes,
                     args.combinetype, unify_dim).to(device)
    opt = torch.optim.Adam(log.parameters(), lr=downstreamlr)

    idx_train = torch.load(
        f"data/fewshot_{args.dataset.lower()}_graph/{shotnum}-shot_{args.dataset.lower()}_graph/{i}/idx.pt"
    )
    train_batch = torch.load(
        f"data/fewshot_{args.dataset.lower()}_graph/{shotnum}-shot_{args.dataset.lower()}_graph/{i}/batch.pt"
    )
    idx_train = torch.Tensor(idx_train).type(torch.long).to(device)
    train_batch = torch.Tensor(train_batch).type(torch.long).to(device)
    
    train_lbls = torch.load(
        f"data/fewshot_{args.dataset.lower()}_graph/{shotnum}-shot_{args.dataset.lower()}_graph/{i}/labels.pt"
    ).type(torch.long).squeeze().to(device)

    pretrain_weights = {
        name: getattr(model, f'pretext{idx+1}').weight.detach()
        for idx, name in enumerate(PRETRAIN_DATASETS)
    }
    
    rag_enhanced_features = rag_enhancer.enhance_features_for_shot_nodes(
        shot_node_indices=idx_train.tolist(),
        base_enhanced_features=base_enhanced_features,
        raw_texts=raw_texts,
        adj=adj_scipy, 
        pretrain_dataset_names=pretrain_sets_for_this_run,
        pretrain_weights=pretrain_weights,
        pretrain_features=pretrain_features,  
        edge_index=data.edge_index,  
        k=5,
        dataset_name=downstream_task_name,  
        rag_weight=0.0,  
        motif_weight=0.0  
    )
    
    with torch.no_grad():
        rag_enhanced_embeds, _ = model.embed(rag_enhanced_features, sp_adj, sparse, None, None)

        pretrain_embs = rag_enhanced_embeds[0, idx_train]

    best_loss_downstream = 1e9
    patience_downstream = 20  
    cnt_wait_downstream = 0
    for _ in range(400):  
        log.train()
        opt.zero_grad()
        logits = log(rag_enhanced_features, sp_adj, sparse, model.gcn, idx_train, train_batch, pretrain_embs, train_lbls, 1)
        loss = xent(logits, train_lbls)
        loss.backward()
        opt.step()
        if loss.item() < best_loss_downstream:
            best_loss_downstream = loss.item()
            cnt_wait_downstream = 0
        else:
            cnt_wait_downstream += 1
        if cnt_wait_downstream == patience_downstream:
                break
            
    log.eval()
    with torch.no_grad():
        logits = log(base_enhanced_features, sp_adj, sparse, model.gcn, testlist, testindex, test_embs)
    preds = torch.argmax(logits, dim=1)
    test_lbls = labels[idx_test]
    acc = (torch.sum(preds == test_lbls).float() / test_lbls.shape[0]).item()
    all_accs.append(acc * 100)
    
accs_tensor = torch.tensor(all_accs)
mean_acc = accs_tensor.mean().item()
std_acc = accs_tensor.std().item()
print(f"  - average accuracy: {mean_acc:.2f}%")
print(f"  - standard deviation: {std_acc:.2f}")