loader.py

import os
import importlib
import torch
import torch.nn as nn
from torch_geometric.data import InMemoryDataset, Data


# ----------------------------------------------------------------------------------------------- #


def build_graph(args, transform=None, pre_transform=None) -> Data:
    name = args.dataset
    module_path = f'loader'
    module = importlib.import_module(module_path)
    for attr in dir(module):
        if attr.lower() == name.lower():
            return getattr(module, attr)(transform=transform, pre_transform=pre_transform)[0]
    raise NotImplementedError(f'No model named {name} in {module_path}')



def build_model(args) -> nn.Module:
    name = args.model.name
    module_path = f'model.{name}'
    module = importlib.import_module(module_path)
    for attr in dir(module):
        if attr.lower() == name.lower():
            return getattr(module, attr)(args)
    raise NotImplementedError(f'No model named {name} in {module_path}')



def build_trainer(args, model, optimizer):
    name = args.train.trainer
    module_path = f'trainer.{name}'
    module = importlib.import_module(module_path)
    for attr in dir(module):
        if attr.lower() == name.lower():
            return getattr(module, attr)(args, model, optimizer)
    raise NotImplementedError(f'No trainer named {name} in {module_path}')


# ----------------------------------------------------------------------------------------------- #


class DefaultDataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None):
        self.src_list = []
        self.dst_list = []
        self.sign_list = []
        # below line must be placed at the end of lists initialization.
        super().__init__(root, transform, pre_transform)
        self.data, self.slices = torch.load(self.processed_paths[0], weights_only=False)

    @property
    def processed_file_names(self):
        return ['data.pt']

    @property
    def raw_dir(self):
        return self.root

    @property
    def raw_file_names(self):
        raise NotImplementedError("`raw_file_names` should be overridden in subclasses.")

    def read_raw_file(self, raw_path):
        raise NotImplementedError("`read_raw_file` should be overridden in subclasses.")
    
    def remap_node_id(self) -> int:
        unique_nodes = set(self.src_list + self.dst_list)
        hash = {old_id: new_id for new_id, old_id in enumerate(sorted(unique_nodes))}
        self.src_list = [hash[src] for src in self.src_list]  # Implicitly return
        self.dst_list = [hash[dst] for dst in self.dst_list]  # Implicitly return
        return len(unique_nodes)
    
    def process(self):
        raw_path = os.path.join(self.raw_dir, self.raw_file_names[0])
        print(f'[INFO] load raw data from {raw_path}', flush=True)
        self.read_raw_file(raw_path)
        
        num_nodes = self.remap_node_id()  # Remap node IDs to a contiguous range.
        edge_index = torch.LongTensor([self.src_list, self.dst_list])
        edge_weight = torch.FloatTensor(self.sign_list)
        
        data = Data(edge_index=edge_index, edge_weight=edge_weight, num_nodes=num_nodes)
        torch.save(self.collate([data]), self.processed_paths[0])


# ----------------------------------------------------------------------------------------------- #


# https://snap.stanford.edu/data/soc-sign-bitcoin-alpha.html
class BitcoinAlpha(DefaultDataset):
    def __init__(self, root='./dataset/BitcoinAlpha', transform=None, pre_transform=None):
        super().__init__(root, transform, pre_transform)

    @property
    def raw_file_names(self):
        return ['soc-sign-bitcoinalpha.csv']
    
    def read_raw_file(self, raw_path):
        with open(raw_path, 'r') as f:
            for line in f:
                src, dst, wgt, ts = map(int, line.strip().split(','))
                sign = 1 if wgt > 0 else -1
                self.src_list.append(src)
                self.dst_list.append(dst)
                self.sign_list.append(sign)


# ----------------------------------------------------------------------------------------------- #


#https://snap.stanford.edu/data/soc-sign-bitcoin-otc.html
class BitcoinOTC(DefaultDataset):
    def __init__(self, root='./dataset/BitcoinOTC', transform=None, pre_transform=None):
        super().__init__(root, transform, pre_transform)

    @property
    def raw_file_names(self):
        return ['soc-sign-bitcoinotc.csv']
    
    def read_raw_file(self, raw_path):
        with open(raw_path, 'r') as f:
            for line in f:
                src, dst, wgt, ts = map(int, map(float, line.strip().split(',')))
                sign = 1 if wgt > 0 else -1
                self.src_list.append(src)
                self.dst_list.append(dst)
                self.sign_list.append(sign)


# ----------------------------------------------------------------------------------------------- #


# http://konect.cc/networks/epinions/
class Epinions(DefaultDataset):
    def __init__(self, root='./dataset/Epinions', transform=None, pre_transform=None):
        super().__init__(root, transform, pre_transform)

    @property
    def raw_file_names(self):
        return ['out.epinions']
    
    def read_raw_file(self, raw_path):
        with open(raw_path, 'r') as f:
            for line in f:
                if line.startswith('%'):
                    continue
                src, dst, sign, ts = map(int, line.strip().split(' '))
                self.src_list.append(src)
                self.dst_list.append(dst)
                self.sign_list.append(sign)


# ----------------------------------------------------------------------------------------------- #


# https://snap.stanford.edu/data/soc-sign-Slashdot090221.html
class SlashDot(DefaultDataset):
    def __init__(self, root='./dataset/SlashDot', transform=None, pre_transform=None):
        super().__init__(root, transform, pre_transform)
    
    @property
    def raw_file_names(self):
        return ['soc-sign-Slashdot090221.txt']
    
    def read_raw_file(self, raw_path):
        with open(raw_path, 'r') as f:
            for line in f:
                if line.startswith('#'):
                    continue
                src, dst, sign = map(int, line.strip().split('\t'))
                self.src_list.append(src)
                self.dst_list.append(dst)
                self.sign_list.append(sign)


# ----------------------------------------------------------------------------------------------- #


# https://networks.skewed.de/net/elec
class WikiElec(DefaultDataset):
    def __init__(self, root='./dataset/WikiElec', transform=None, pre_transform=None):
        super().__init__(root, transform, pre_transform)
    
    @property
    def raw_file_names(self):
        return ['out.elec']
    
    def read_raw_file(self, raw_path):
        with open(raw_path, 'r') as f:
            for line in f:
                if line.startswith('%'):
                    continue
                src, dst, sign, ts = map(int, line.strip().split('\t'))
                self.src_list.append(src)
                self.dst_list.append(dst)
                self.sign_list.append(sign)


# ----------------------------------------------------------------------------------------------- #


# https://networks.skewed.de/net/wiki_rfa
class WikiRfa(DefaultDataset):
    def __init__(self, root='./dataset/WikiRfa', transform=None, pre_transform=None):
        super().__init__(root, transform, pre_transform)
    
    @property
    def raw_file_names(self):
        return ['edges.csv']
    
    def read_raw_file(self, raw_path):
        with open(raw_path, 'r', encoding='utf-8') as f:
            for line in f:
                if line.startswith('#'):
                    continue
                src, dst, vote = map(int, line.strip().split(',')[:3])
                if vote == 0:
                    continue # Skip neutral votes
                self.src_list.append(src)
                self.dst_list.append(dst)
                self.sign_list.append(vote)