from utils.utils import *
from models.gcn import GCN
from models.mlp import MLP
from config import CONFIG
cfg = CONFIG()
if len(sys.argv) != 2:
sys.exit("Use: python train.py <dataset>")
datasets = ['20ng', 'R8', 'R52', 'ohsumed', 'mr']
dataset = sys.argv[1]
if dataset not in datasets:
sys.exit("wrong dataset name")
cfg.dataset = dataset
# Set random seed
seed = random.randint(1, 200)
seed = 2019
np.random.seed(seed)
torch.manual_seed(seed)
# if torch.cuda.is_available():
# torch.cuda.manual_seed(seed)
# Settings
# os.environ["CUDA_VISIBLE_DEVICES"] = ""
# Load data
adj, features, y_train, y_val, y_test, train_mask, val_mask, test_mask, train_size, test_size = load_corpus(
cfg.dataset)