initial commit

Author: justinlovelace

.gitignore

@@ -0,0 +1,145 @@
+data/
+saved_models/
+bert_ckpts/
+saved_reranking_models/
+pretrained_models/
+pretrained_reranking_models/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/

CONSTANTS.py

@@ -0,0 +1,10 @@
+SNOMED_CORE_DIR = 'data/UMLS/SNOMEDCT_CORE_SUBSET_202008'
+DATA_DIR = {'SNOMED_CT_CORE': 'data/SNOMED-CT-Core/',
+            'FB15K_237': 'data/FB15K-237/',
+            'FB15K_237_SPARSE': 'data/FB15K-237-Sparse/',
+            'CN100K': 'data/CN100K/'}
+COLUMN_NAMES = {'SNOMED_CT_CORE': ('CUI2_id', 'RELA_id', 'CUI1_id'),
+                'FB15K_237': ('entity1_id', 'rel_id', 'entity2_id'),
+                'FB15K_237_SPARSE': ('entity1_id', 'rel_id', 'entity2_id'),
+                'CN100K': ('entity1_id', 'rel_id', 'entity2_id')}
+UMLS_SOURCE_DIR = 'data/UMLS/2020AA/META'

README.md

@@ -1 +1,48 @@
-# robust-kg-completion
+# Robust Knowledge Graph Completion with Stacked Convolutions and a Student Re-Ranking Network
+
+
+This repository contains the implementation for our paper: 
+
+**Robust Knowledge Graph Completion with Stacked Convolutions and a Student Re-Ranking Network** \
+Justin Lovelace, Denis Newman-Griffis, Shikhar Vashishth, Jill Fain Lehman, and Carolyn Penstein Rosé \
+Annual Meeting of the Association for Computational Linguistics and the International Joint Conference on Natural Language Processing
+(**ACL-IJCNLP**) 2021
+
+## Dependencies
+
+Our work was performed with Python 3.8. The dependencies can be installed from `requirements.txt`.
+
+## Data Preparation
+
+- We conduct our work upon the existing FB15K-237 and CN100K datasets. We additionally developed the FB15K-237-Sparse and SNOMED-CT Core datasets for our work.
+- Running `./scripts/prepare_datasets.sh` will unzip the dataset files and process them for use by our models.
+- Because the SNOMED-CT Core dataset was derived from the UMLS, we cannot directly release the dataset files. See [here](snomed_ct_core.md) for full instructions for how to recreate the dataset.
+- The BERT embeddings can be downloaded from [here](https://drive.google.com/drive/folders/1gfbZcJoay69BUzQLQku-qHB6IZ5zxLQw?usp=sharing). The `bert_emb.pt` files should be stored in the corresponding dataset directories, e.g. `data/CN100K/bert_emb.pt`
+
+## Training Ranking Models
+We provide scripts to train our proposed ranking model, denoted as BERT-ResNet in our paper, for all four datasets. 
+
+- FB15K-237: `./scripts/train_resnet_fb15k237.sh`
+- FB15K-237-Sparse: `./scripts/train_resnet_fb15k237_sparse.sh`
+- CN100K: `./scripts/train_resnet_cn100k.sh`
+- SNOMED-CT Core `./scripts/train_resnet_snomed.sh`
+
+## Training Re-Ranking Models
+The re-ranking models can only be trained after the ranking model for the corresponding dataset has already finished training. First, download the BERT checkpoints used for our training from [here](https://drive.google.com/drive/folders/1BsxeWEtFpZuHD_bCQKsIy0zfwl7xi6Bq?usp=sharing). They should be unzipped and stored in `reranking/bert_ckpts`. A re-ranking model can then be trained with the provided scripts similarly to above.
+
+- FB15K-237: `./scripts/train_reranking_fb15k237.sh`
+- FB15K-237-Sparse: `./scripts/train_reranking_fb15k237_sparse.sh`
+- CN100K: `./scripts/train_reranking_cn100k.sh`
+- SNOMED-CT Core `./scripts/train_reranking_snomed.sh`
+
+## Evaluating Pretrained Ranking Models
+Pretrained ranking models can be downloaded from [here](https://drive.google.com/drive/folders/1q20hhUq20wt5OSbHbOWsvAiviFUZ8s8r?usp=sharing). After unzipping them in a `robust-kg-completion/pretrained_models` directory, they can be evaluated by running `./scripts/eval_pretrained_ranking_model.sh {DATASET}` where `{DATASET}` is one of `SNOMED_CT_CORE`, `FB15K_237`, `FB15K_237_SPARSE`, or `CN100K`.
+
+## Evaluating Pretrained Re-Ranking Models 
+Pretrained re-ranking models can be downloaded from [here](https://drive.google.com/drive/folders/1q20hhUq20wt5OSbHbOWsvAiviFUZ8s8r?usp=sharing). After unzipping them in a `robust-kg-completion/reranking/pretrained_reranking_models` directory, they can be evaluated by running the following commands.
+
+- FB15K-237: `./scripts/eval_pretrained_reranking_model.sh FB15K_237 0.75`
+- FB15K-237-Sparse: `./scripts/eval_pretrained_reranking_model.sh FB15K_237_SPARSE 0.75`
+- CN100K: `./scripts/eval_pretrained_reranking_model.sh CN100K 1.0`
+- SNOMED-CT Core `./scripts/eval_pretrained_reranking_model.sh SNOMED_CT_CORE 0.5`
+

dataset.py

@@ -0,0 +1,96 @@
+import pandas as pd
+import numpy as np
+import torch
+from torch.utils.data import Dataset, DataLoader
+import argparse
+import os
+from collections import defaultdict
+import sys
+from CONSTANTS import DATA_DIR, COLUMN_NAMES
+
+
+class KbDataset(Dataset):
+    def __init__(self, args: argparse.ArgumentParser):
+        self.args = args
+        csv_file = os.path.join(DATA_DIR[args.dataset], f'df_train.csv')
+        e1_col, rel_col, e2_col = COLUMN_NAMES[args.dataset]
+        print(f'Loading dataset from {csv_file}')
+        df_kg = pd.read_csv(csv_file)
+        self.pos_samples = df_kg[[e1_col,
+                                  rel_col, e2_col]].to_numpy(np.int64)
+        # print(self.pos_samples.shape)
+        if args.strategy == 'one_to_n':
+            self.gen_one_to_n_data(df_kg)
+        elif args.strategy == 'k_to_n' or args.strategy == 'gen_triplets':
+            self.gen_k_to_n_data(df_kg)
+        elif args.strategy == 'softmax':
+            self.gen_softmax_data(df_kg)
+        else:
+            raise NotImplementedError
+
+    def gen_one_to_n_data(self, df_kg):
+        e1_col, rel_col, e2_col = COLUMN_NAMES[self.args.dataset]
+        self.data = df_kg[[
+            e1_col, rel_col]].to_numpy(np.int64)
+        self.labels = np.zeros(
+            (len(self.data), self.args.num_entities), dtype=np.float32)
+        self.labels[np.arange(start=0, stop=self.labels.shape[0],
+                              step=1), df_kg[e2_col].to_numpy(np.int64)] = 1
+
+    def gen_softmax_data(self, df_kg):
+        e1_col, rel_col, e2_col = COLUMN_NAMES[self.args.dataset]
+        self.data = df_kg[[
+            e1_col, rel_col]].to_numpy(np.int64)
+        self.labels = df_kg[e2_col].to_numpy(np.int64)
+
+    def gen_k_to_n_data(self, df_kg):
+        e1_col, rel_col, e2_col = COLUMN_NAMES[self.args.dataset]
+        self.data = df_kg[[
+            e1_col, rel_col]].drop_duplicates().to_numpy(np.int64)
+        e2_lookup = defaultdict(set)
+        for e1, r, e2 in zip(df_kg[e1_col], df_kg[rel_col], df_kg[e2_col]):
+            e2_lookup[(e1, r)].add(e2)
+        self.labels = np.zeros(
+            (len(self.data), self.args.num_entities), dtype=np.float32)
+        for idx, query in enumerate(self.data):
+            e1, r = query[0], query[1]
+            for e2 in e2_lookup[(e1, r)]:
+                self.labels[idx, e2] = 1
+
+    def __len__(self):
+        return self.data.shape[0]
+
+    def __getitem__(self, idx):
+        return self.data[idx], self.labels[idx]
+
+
+class KbEvalGenerator(Dataset):
+    def __init__(self, eval_split: str, args: argparse.ArgumentParser):
+        self.args = args
+        pos_samples = defaultdict(set)
+        splits = ['train', 'valid', 'test']
+        e1_col, rel_col, e2_col = COLUMN_NAMES[args.dataset]
+        assert eval_split in splits
+        for spl in splits:
+            csv_file = os.path.join(DATA_DIR[args.dataset], f'df_{spl}.csv')
+            df_data = pd.read_csv(csv_file)
+            if spl == eval_split:
+                self.queries = df_data[[
+                    e1_col, rel_col]].to_numpy(np.int64)
+                e2_list = df_data[e2_col].tolist()
+            for e1, r, e2 in zip(df_data[e1_col], df_data[rel_col], df_data[e2_col]):
+                pos_samples[(e1, r)].add(e2)
+        self.labels = np.zeros((self.queries.shape[0], self.args.num_entities))
+        self.filtered_labels = np.zeros(
+            (self.queries.shape[0], self.args.num_entities))
+        for i, query in enumerate(self.queries):
+            e1, r = query[0], query[1]
+            e2 = e2_list[i]
+            self.labels[i, e2] = 1
+            self.filtered_labels[i, list(pos_samples[(e1, r)] - {e2})] = 1
+
+    def __len__(self):
+        return self.queries.shape[0]
+
+    def __getitem__(self, idx):
+        return self.queries[idx], self.labels[idx], self.filtered_labels[idx]