CTDKG

Source code for 'CTDKG: Towards Uncovering Chemical–Gene–Disease Pathways via Knowledge Graphs' with pytorch and torch_geometric.

Requirements & Installation

The code is written in Python 3.10.0 and required packages as follow:

• torch
• torch_geometric
• torch_scatter (not necessary, need for compgcn implementation)

Please check INSTALL.md for installation details.

Basic Usage

Data download

from module.dataset import LinkPredDataset

# Download and process data at './dataset/cd'
dataset = LinkPredDataset(name = 'cd')
train_triples, valid_triples, test_triples = dataset.get_edge_split()

If you have problems downloading the dataset, refer to the link and manually install the dataset. This will be modified soon.

Data loader

# head loader for train data
data = dataset[0]
train_count, train_true_head, train_true_tail = defaultdict(lambda: 4), defaultdict(list), defaultdict(list)
for i in tqdm(range(len(train_triples['head']))):
    head, relation, tail = train_triples['head'][i].item(), train_triples['relation'][i].item(), train_triples['tail'][i].item()
    head_type, tail_type = train_triples['head_type'][i], train_triples['tail_type'][i]
    train_count[(head, relation, head_type)] += 1
    train_count[(tail, -relation-1, tail_type)] += 1
    train_true_head[(relation, tail)].append(head)
    train_true_tail[(head, relation)].append(tail)

train_dataloader_head = DataLoader(
    TrainDataset(train_triples, nentity, nrelation, 
        args.negative_sample_size, 'head-batch',
        train_count, train_true_head, train_true_tail,
        entity_dict), 
    batch_size=args.batch_size,
    shuffle=True, 
    num_workers=args.num_workers,
    collate_fn=TrainDataset.collate_fn
)

# head loader for validation/test data
valid_dataloader_head = DataLoader(
    TestDataset(
        valid_triples, 
        args, 
        'head-batch',
        random_sampling,
        entity_dict
    ),
    batch_size = args.test_batch_size,
    num_workers = args.num_workers,
    collate_fn = TestDataset.collate_fn
)

The head loader means that head entities of true triples were contaminated $(h', r, t)$. If you want to make a tail loader, substitute 'head-batch' with 'tail-batch.' If you run the code on the local environment, then set the num_workers argument to num_workers=0.

Training & Evaluation

We support eight models and seven datasets.

Datasets: cd, cgd, cgpd, ctd

# for translation models
python main.py --dataset=cd --model=TransE
python main.py --dataset=cd --model=RotatE -num_entity_embedding=2
python main.py --dataset=cd --model=HAKE --num_entity_embedding=2 --num_relation_embedding=3
python main.py --dataset=cd --model=TripleRE --num_relation_embedding=3
python main.py --dataset-cd --model=Rotate4D --num_entity_embedding=4 --num_relation_embedding=4
# python main.py --dataset=cd --model=GIE
# python main.py --dataset=cd --model=HOUSE

# for semantic information models
python main.py --dataset=cd --model=DistMult
python main.py --dataset=cd --model=ComplEx --num_entity_embedding=2 --num_relation_embedding=2
python main.py --dataset=cd --model=QuatRE --num_entity_embedding=4 --num_relation_embedding=4

# for neural network models
# python main_conv.py --dataset=cd --model=conve --learning_rate=0.001 --negative_sample_size=1
# python main_conv.py --dataset=cd --model=convkb --learning_rate=0.005 --negative_sample_size=1
# python main_gnn.py --dataset=cd --model=rgcn --learning_rate=0.001 --negative_sample_size=1