from utils.utils import *
from models.gcn import GCN
from models.gin import GIN
from models.mlp import MLP
from config import CONFIG
cfg = CONFIG()
parser = argparse.ArgumentParser(description='Pytorch implementation of GCN and GIN')
parser.add_argument('data',help='path to dataset')
parser.add_argument('--num_mlp_layers',default=1,type = int)
parser.add_argument('--train_eps',action='store_true')
args = parser.parse_args()
cfg.num_mlp_layers = args.num_mlp_layers
cfg.train_eps = args.train_eps
# if len(sys.argv) != 2:
# sys.exit("Use: python train.py <dataset>")
datasets = ['20ng', 'R8', 'R52', 'ohsumed', 'mr']
dataset = args.data
print('dataset %s'%datasets)
if dataset not in datasets:
sys.exit("wrong dataset name")
cfg.dataset = dataset
# Set random seed
seed = random.randint(1, 200)
